metrale_model_layers/weight_map/
ssm_qwen35.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Qwen3.5-family loaders: the separate-projection linear-attention weights and three MoE loaders (NVFP4, native FP8 experts, no shared expert).
4//!
5//! Owner: model-layers (weight loading).
6//! Invariants: none beyond the types.
7
8#![allow(unused_imports)]
9
10use anyhow::{Context, Result, bail, ensure};
11use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
12use metrale_model_weights::weights::{WeightDtype, WeightStore};
13
14use super::*;
15
16#[path = "ssm_qwen35/dequant_fp8.rs"]
17mod dequant_fp8;
18use dequant_fp8::dequant_fp8_block_slice_bf16;
19
20/// 2026-09-25: Qwen3.5 linear-attention weights with separate projections. The
21/// projections, `conv1d` and `norm` are BF16; `a_log` and `dt_bias` are FP32.
22pub struct SsmWeightsQwen35 {
23    /// 2026-09-25: Q, K and V projections in one weight (no Z).
24    pub in_proj_qkv: DenseWeight,
25    pub in_proj_z: DenseWeight,
26    pub in_proj_a: DenseWeight,
27    pub in_proj_b: DenseWeight,
28    pub conv1d: DenseWeight,
29    pub a_log: DenseWeight,
30    pub dt_bias: DenseWeight,
31    pub norm: DenseWeight,
32    pub out_proj: DenseWeight,
33}
34
35/// 2026-09-25: Load the Qwen3.5 `{layer_prefix}.linear_attn` weights.
36pub fn load_ssm_qwen35(
37    store: &WeightStore,
38    layer_prefix: &str,
39    gpu: &dyn GpuBackend,
40    // 2026-09-25: Unused: each projection is loaded by its own on-disk dtype.
41    _variant: Nvfp4Variant,
42) -> Result<SsmWeightsQwen35> {
43    let p = format!("{layer_prefix}.linear_attn");
44
45    // 2026-09-25: A projection with `weight_packed` (NVFP4, 2 values per byte) is
46    // dequantized to BF16 with dims from the packed shape; any other goes
47    // through `dense_auto`.
48    let load_proj = |prefix: &str| -> Result<DenseWeight> {
49        if store.contains(&format!("{prefix}.weight_packed")) {
50            let shape = store.get(&format!("{prefix}.weight_packed"))?.shape.clone();
51            dequant_nvfp4_to_bf16(store, prefix, shape[0], shape[1] * 2, gpu)
52        } else {
53            dense_auto(store, &format!("{prefix}.weight"), gpu)
54        }
55    };
56
57    Ok(SsmWeightsQwen35 {
58        in_proj_qkv: load_proj(&format!("{p}.in_proj_qkv"))?,
59        in_proj_z: load_proj(&format!("{p}.in_proj_z"))?,
60        in_proj_a: load_proj(&format!("{p}.in_proj_a"))?,
61        in_proj_b: load_proj(&format!("{p}.in_proj_b"))?,
62        conv1d: dense_auto(store, &format!("{p}.conv1d.weight"), gpu)?,
63        a_log: dense_keep_f32(store, &format!("{p}.A_log"), gpu)?,
64        dt_bias: dense_keep_f32(store, &format!("{p}.dt_bias"), gpu)?,
65        norm: dense_f32_safe(store, &format!("{p}.norm.weight"), gpu)?,
66        out_proj: load_proj(&format!("{p}.out_proj"))?,
67    })
68}
69
70/// 2026-09-27: Which NVFP4 copies `load_moe_qwen35` builds. A projection it does not build is
71/// NULL (`QuantizedWeight::null()`).
72#[derive(Clone, Copy, Debug, PartialEq, Eq)]
73pub struct Nvfp4MoeCopies {
74    /// 2026-09-27: The routed experts' gate and up projections.
75    pub routed_gate_up: bool,
76    /// 2026-09-27: The routed experts' down projections.
77    pub routed_down: bool,
78    /// 2026-09-27: The shared expert.
79    pub shared: bool,
80}
81
82impl Nvfp4MoeCopies {
83    /// 2026-09-27: Every projection.
84    pub const ALL: Self = Self {
85        routed_gate_up: true,
86        routed_down: true,
87        shared: true,
88    };
89    /// 2026-09-27: The shared expert only; every routed expert is `ExpertWeight::null()`.
90    pub const SHARED_ONLY: Self = Self {
91        routed_gate_up: false,
92        routed_down: false,
93        shared: true,
94    };
95}
96
97/// 2026-09-25: Load a Qwen3.5 MoE block (`{layer_prefix}.mlp`) as NVFP4.
98///
99/// 2026-09-27: `copies` names the projections to build; routed experts that this rank does not
100/// hold are `ExpertWeight::null()`. The fused layout builds every projection (it errors on
101/// anything but `Nvfp4MoeCopies::ALL`).
102pub fn load_moe_qwen35(
103    store: &WeightStore,
104    layer_prefix: &str,
105    num_experts: usize,
106    gpu: &dyn GpuBackend,
107    config: &metrale_config::ModelConfig,
108    variant: Nvfp4Variant,
109    absmax_k: metrale_gpu_runtime::gpu::KernelHandle,
110    quantize_k: metrale_gpu_runtime::gpu::KernelHandle,
111    stream: u64,
112    copies: Nvfp4MoeCopies,
113) -> Result<MoeWeights> {
114    let p = format!("{layer_prefix}.mlp");
115
116    let gate = dense_auto(store, &format!("{p}.gate.weight"), gpu)?;
117    let shared_expert_gate = dense_auto(store, &format!("{p}.shared_expert_gate.weight"), gpu)?;
118
119    let inter = config.moe_intermediate_size;
120    let h = config.hidden_size;
121
122    let qctx = QuantizeCtx {
123        absmax_k,
124        quantize_k,
125        stream,
126    };
127
128    let fused_gate_up_key = format!("{p}.experts.gate_up_proj");
129    let fused_down_key = format!("{p}.experts.down_proj");
130    // 2026-09-25: Fused layout: one `experts.gate_up_proj` `[E, 2*inter, h]` and
131    // one `experts.down_proj` `[E, h, inter]` per layer, sliced per expert and
132    // quantized to NVFP4. A BF16 slice is quantized directly; an FP8E4M3 slice
133    // is first dequantized to BF16 with its `{key}_scale_inv` block scales
134    // `[E, sn, sk]`. The dtype decides, not `variant`.
135    let is_fused = store.contains(&fused_gate_up_key) && store.contains(&fused_down_key);
136    anyhow::ensure!(
137        !is_fused || copies == Nvfp4MoeCopies::ALL,
138        "load_moe_qwen35: the fused expert layout builds every NVFP4 projection ({copies:?} asked)"
139    );
140    let fused_is_fp8 = is_fused
141        && store
142            .get(&fused_gate_up_key)
143            .map(|w| w.dtype == WeightDtype::FP8E4M3)
144            .unwrap_or(false);
145
146    let load_expert_fused = |expert_idx: usize| -> Result<ExpertWeight> {
147        let fused_gu = store.get(&fused_gate_up_key)?;
148        let fused_d = store.get(&fused_down_key)?;
149        if fused_is_fp8 {
150            let gu_s = store.get(&format!("{fused_gate_up_key}_scale_inv"))?;
151            let d_s = store.get(&format!("{fused_down_key}_scale_inv"))?;
152            let (gu_sn, gu_sk) = (gu_s.shape[1], gu_s.shape[2]);
153            let (d_sn, d_sk) = (d_s.shape[1], d_s.shape[2]);
154            let gu_s_f32 = gu_s.dtype == WeightDtype::FP32;
155            let d_s_f32 = d_s.dtype == WeightDtype::FP32;
156            let gu_w_stride = 2 * inter * h;
157            let d_w_stride = h * inter;
158            let gu_s_elem = if gu_s_f32 { 4 } else { 2 };
159            let d_s_elem = if d_s_f32 { 4 } else { 2 };
160            let gu_s_stride = gu_sn * gu_sk * gu_s_elem;
161            let d_s_stride = d_sn * d_sk * d_s_elem;
162            // 2026-09-25: Dequantize all of `gate_up[e]`, then take gate as rows
163            // `0..inter` and up as rows `inter..2*inter`.
164            let gu_bf16 = dequant_fp8_block_slice_bf16(
165                gpu,
166                fused_gu.ptr.offset(expert_idx * gu_w_stride),
167                gu_s.ptr.offset(expert_idx * gu_s_stride),
168                2 * inter,
169                h,
170                gu_sn,
171                gu_sk,
172                gu_s_f32,
173            )?;
174            let down_bf16 = dequant_fp8_block_slice_bf16(
175                gpu,
176                fused_d.ptr.offset(expert_idx * d_w_stride),
177                d_s.ptr.offset(expert_idx * d_s_stride),
178                h,
179                inter,
180                d_sn,
181                d_sk,
182                d_s_f32,
183            )?;
184            let gate_dw = DenseWeight { weight: gu_bf16 };
185            let up_dw = DenseWeight {
186                weight: gu_bf16.offset(inter * h * 2),
187            };
188            let down_dw = DenseWeight { weight: down_bf16 };
189            let out = ExpertWeight {
190                gate_proj: quantize_to_nvfp4(
191                    &gate_dw, inter, h, gpu, absmax_k, quantize_k, stream,
192                )?,
193                up_proj: quantize_to_nvfp4(&up_dw, inter, h, gpu, absmax_k, quantize_k, stream)?,
194                down_proj: quantize_to_nvfp4(
195                    &down_dw, h, inter, gpu, absmax_k, quantize_k, stream,
196                )?,
197            };
198            gpu.free(gu_bf16)?;
199            gpu.free(down_bf16)?;
200            Ok(out)
201        } else {
202            let bf16 = 2usize;
203            let gu_per_expert_bytes = 2 * inter * h * bf16;
204            let d_per_expert_bytes = h * inter * bf16;
205            let gate_off = expert_idx * gu_per_expert_bytes;
206            let up_off = gate_off + inter * h * bf16;
207            let down_off = expert_idx * d_per_expert_bytes;
208            let gate_dw = DenseWeight {
209                weight: fused_gu.ptr.offset(gate_off),
210            };
211            let up_dw = DenseWeight {
212                weight: fused_gu.ptr.offset(up_off),
213            };
214            let down_dw = DenseWeight {
215                weight: fused_d.ptr.offset(down_off),
216            };
217            Ok(ExpertWeight {
218                gate_proj: quantize_to_nvfp4(
219                    &gate_dw, inter, h, gpu, absmax_k, quantize_k, stream,
220                )?,
221                up_proj: quantize_to_nvfp4(&up_dw, inter, h, gpu, absmax_k, quantize_k, stream)?,
222                down_proj: quantize_to_nvfp4(
223                    &down_dw, h, inter, gpu, absmax_k, quantize_k, stream,
224                )?,
225            })
226        }
227    };
228
229    // 2026-09-25: `quantized_any` picks the format per key, so a BF16 expert in
230    // an FP8 or NVFP4 checkpoint still loads.
231    let load_expert = |prefix: &str, gate_up: bool, down: bool| -> Result<ExpertWeight> {
232        let proj = |name: &str, n: usize, k: usize, on: bool| -> Result<QuantizedWeight> {
233            if on {
234                quantized_any(store, &format!("{prefix}.{name}"), n, k, gpu, variant, qctx)
235            } else {
236                Ok(QuantizedWeight::null())
237            }
238        };
239        Ok(ExpertWeight {
240            gate_proj: proj("gate_proj", inter, h, gate_up)?,
241            up_proj: proj("up_proj", inter, h, gate_up)?,
242            down_proj: proj("down_proj", h, inter, down)?,
243        })
244    };
245
246    let shared_expert = if copies.shared {
247        load_expert(&format!("{p}.shared_expert"), true, true)?
248    } else {
249        ExpertWeight::null()
250    };
251    let routed_any = copies.routed_gate_up || copies.routed_down;
252
253    let mut experts = Vec::with_capacity(num_experts);
254    for e in 0..num_experts {
255        if !routed_any || !config.is_local_expert(e) {
256            experts.push(ExpertWeight::null());
257        } else if is_fused {
258            experts.push(load_expert_fused(e)?);
259        } else {
260            experts.push(load_expert(
261                &format!("{p}.experts.{e}"),
262                copies.routed_gate_up,
263                copies.routed_down,
264            )?);
265        }
266    }
267
268    // 2026-09-25: Every loaded expert has its own NVFP4 copy here, so the fused source
269    // tensors are freed; a failed free is ignored. The store keeps the freed
270    // pointers, and nothing may read these keys afterwards.
271    if is_fused {
272        if let Ok(w) = store.get(&fused_gate_up_key) {
273            let _ = gpu.free(w.ptr);
274        }
275        if let Ok(w) = store.get(&fused_down_key) {
276            let _ = gpu.free(w.ptr);
277        }
278    }
279
280    Ok(MoeWeights {
281        gate,
282        shared_expert,
283        shared_expert_gate,
284        experts,
285        router_pre_norm: None,
286        correction_bias: None,
287    })
288}
289
290/// 2026-09-25: Load the routed experts of `{layer_prefix}.mlp` as native FP8
291/// block-scaled tables, one entry per expert id. Experts this rank does not
292/// hold are NULL entries tagged `Fp8BlockScaled`.
293pub fn load_moe_qwen35_fp8_experts(
294    store: &WeightStore,
295    layer_prefix: &str,
296    num_experts: usize,
297    gpu: &dyn GpuBackend,
298    config: &metrale_config::ModelConfig,
299) -> Result<Vec<Fp8ExpertWeight>> {
300    let p = format!("{layer_prefix}.mlp");
301    let mut fp8_experts = Vec::with_capacity(num_experts);
302
303    for e in 0..num_experts {
304        if config.is_local_expert(e) {
305            let ep = format!("{p}.experts.{e}");
306            fp8_experts.push(Fp8ExpertWeight {
307                gate_proj: load_fp8_block_scaled_as_fp8weight(
308                    store,
309                    &format!("{ep}.gate_proj"),
310                    gpu,
311                )?,
312                up_proj: load_fp8_block_scaled_as_fp8weight(store, &format!("{ep}.up_proj"), gpu)?,
313                down_proj: load_fp8_block_scaled_as_fp8weight(
314                    store,
315                    &format!("{ep}.down_proj"),
316                    gpu,
317                )?,
318            });
319        } else {
320            let null_block = Fp8Weight {
321                weight: DevicePtr::NULL,
322                row_scale: DevicePtr::NULL,
323                n: 0,
324                k: 0,
325                scale_format: WeightQuantFormat::Fp8BlockScaled,
326            };
327            fp8_experts.push(Fp8ExpertWeight {
328                gate_proj: null_block,
329                up_proj: null_block,
330                down_proj: null_block,
331            });
332        }
333    }
334
335    // 2026-09-25: The shared expert's FP8 tables are loaded (so a missing or
336    // malformed one is an error) and then dropped: `_shared_fp8` is unused.
337    let shared_prefix = format!("{p}.shared_expert");
338    let _shared_fp8 = Fp8ExpertWeight {
339        gate_proj: load_fp8_block_scaled_as_fp8weight(
340            store,
341            &format!("{shared_prefix}.gate_proj"),
342            gpu,
343        )?,
344        up_proj: load_fp8_block_scaled_as_fp8weight(
345            store,
346            &format!("{shared_prefix}.up_proj"),
347            gpu,
348        )?,
349        down_proj: load_fp8_block_scaled_as_fp8weight(
350            store,
351            &format!("{shared_prefix}.down_proj"),
352            gpu,
353        )?,
354    };
355
356    Ok(fp8_experts)
357}
358
359/// 2026-09-25: Load a MoE block that has no shared expert. A zero-filled shared
360/// expert of real expert size stands in: its packed codes, scales and
361/// `weight_scale_2` are all 0, so it contributes 0. Routed experts go through
362/// `quantized_auto`, which panics for `Fp8Dequanted` and `Bf16Raw`.
363pub fn load_moe_no_shared(
364    store: &WeightStore,
365    layer_prefix: &str,
366    num_experts: usize,
367    gpu: &dyn GpuBackend,
368    config: &metrale_config::ModelConfig,
369    variant: Nvfp4Variant,
370) -> Result<MoeWeights> {
371    let p = format!("{layer_prefix}.mlp");
372
373    let gate = dense(store, &format!("{p}.gate.weight"))?;
374
375    let h = config.hidden_size;
376    let inter = config.moe_intermediate_size;
377    let group_size = 16usize;
378
379    let gu_packed_bytes = inter * h / 2;
380    let gu_scale_bytes = inter * (h / group_size);
381    let d_packed_bytes = h * inter / 2;
382    let d_scale_bytes = h * (inter / group_size);
383
384    let alloc_zero = |size: usize| -> Result<DevicePtr> {
385        let ptr = gpu.alloc(size)?;
386        gpu.memset(ptr, 0, size)?;
387        Ok(ptr)
388    };
389
390    let mk_zero_quant = |packed_sz: usize, scale_sz: usize| -> Result<QuantizedWeight> {
391        Ok(QuantizedWeight {
392            weight: alloc_zero(packed_sz)?,
393            weight_scale: alloc_zero(scale_sz)?,
394            weight_scale_2: 0.0,
395            input_scale: DevicePtr::NULL,
396            weight_scale_2_vec: DevicePtr::NULL,
397            act: metrale_config::Nvfp4Act::Unstamped,
398        })
399    };
400
401    let shared_expert = ExpertWeight {
402        gate_proj: mk_zero_quant(gu_packed_bytes, gu_scale_bytes)?,
403        up_proj: mk_zero_quant(gu_packed_bytes, gu_scale_bytes)?,
404        down_proj: mk_zero_quant(d_packed_bytes, d_scale_bytes)?,
405    };
406    // 2026-09-25: A zero gate gives sigmoid(0) = 0.5, times a zero shared output.
407    let shared_expert_gate = DenseWeight {
408        weight: alloc_zero(h * 2)?,
409    };
410
411    let mut experts = Vec::with_capacity(num_experts);
412    for e in 0..num_experts {
413        if config.is_local_expert(e) {
414            experts.push(ExpertWeight {
415                gate_proj: quantized_auto(
416                    store,
417                    &format!("{p}.experts.{e}.gate_proj"),
418                    gpu,
419                    variant,
420                )?,
421                up_proj: quantized_auto(store, &format!("{p}.experts.{e}.up_proj"), gpu, variant)?,
422                down_proj: quantized_auto(
423                    store,
424                    &format!("{p}.experts.{e}.down_proj"),
425                    gpu,
426                    variant,
427                )?,
428            });
429        } else {
430            experts.push(ExpertWeight::null());
431        }
432    }
433
434    Ok(MoeWeights {
435        gate,
436        shared_expert,
437        shared_expert_gate,
438        experts,
439        router_pre_norm: None,
440        correction_bias: None,
441    })
442}
metrale_model_layers/weight_map/ssm_qwen35.rs

metrale_model_layers/weight_map/
ssm_qwen35.rs
