spark_model/weight_map/
loaders_mtp.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Auto-extracted from `weight_map.rs` during refactor wave 4a.
4
5#![allow(unused_imports)]
6
7use anyhow::{Context, Result, bail, ensure};
8use spark_runtime::gpu::{DevicePtr, GpuBackend};
9use spark_runtime::weights::{WeightDtype, WeightStore};
10
11use super::*;
12
13/// Slice a stacked + fused MTP MoE expert layout into per-expert
14/// `DenseExpertWeight`s via DevicePtr offsets (zero-copy).
15///
16/// Expects two BF16 tensors in `store`:
17///   `{mlp}.experts.gate_up_proj` shape `[E, 2*I, H]`
18///       — first `I` rows of axis 1 are gate, next `I` rows are up
19///   `{mlp}.experts.down_proj`    shape `[E, H, I]`
20///
21/// Each expert's `gate`, `up`, `down` are contiguous sub-tensors of the
22/// stacked allocations, so we hand back DenseWeights pointing into the
23/// same underlying GPU memory. The WeightStore retains ownership of the
24/// stacked allocations; the offset pointers are aliases and must NEVER
25/// be passed to `gpu.free()` (the loader doesn't, and ModelWeights drops
26/// the WeightStore last in any case).
27pub(super) fn load_mtp_experts_stacked(
28    store: &WeightStore,
29    mlp: &str,
30    num_experts: usize,
31) -> Result<Vec<DenseExpertWeight>> {
32    let gate_up = store.get(&format!("{mlp}.experts.gate_up_proj"))?;
33    let down = store.get(&format!("{mlp}.experts.down_proj"))?;
34
35    ensure!(
36        gate_up.shape.len() == 3,
37        "MTP stacked experts.gate_up_proj: expected 3D [E,2I,H], got {:?}",
38        gate_up.shape
39    );
40    ensure!(
41        down.shape.len() == 3,
42        "MTP stacked experts.down_proj: expected 3D [E,H,I], got {:?}",
43        down.shape
44    );
45    ensure!(
46        gate_up.shape[0] == num_experts,
47        "MTP stacked experts.gate_up_proj: expert dim {} != num_experts {num_experts}",
48        gate_up.shape[0]
49    );
50    ensure!(
51        down.shape[0] == num_experts,
52        "MTP stacked experts.down_proj: expert dim {} != num_experts {num_experts}",
53        down.shape[0]
54    );
55
56    let two_inter = gate_up.shape[1];
57    let hidden = gate_up.shape[2];
58    ensure!(
59        two_inter % 2 == 0,
60        "MTP stacked experts.gate_up_proj: 2nd dim must be even (gate+up fused), got {two_inter}"
61    );
62    let intermediate = two_inter / 2;
63
64    ensure!(
65        down.shape[1] == hidden,
66        "MTP stacked: gate_up_proj.hidden ({hidden}) != down_proj.hidden ({})",
67        down.shape[1]
68    );
69    ensure!(
70        down.shape[2] == intermediate,
71        "MTP stacked: down_proj.intermediate ({}) != gate_up_proj/2 ({intermediate})",
72        down.shape[2]
73    );
74
75    // Stacked tensors must be BF16 — the per-expert split path also returns
76    // BF16 (norm/gate dense or dequanted projections), so we keep the
77    // contract uniform downstream.
78    ensure!(
79        matches!(gate_up.dtype, WeightDtype::BF16),
80        "MTP stacked experts.gate_up_proj: expected BF16, got {:?}",
81        gate_up.dtype
82    );
83    ensure!(
84        matches!(down.dtype, WeightDtype::BF16),
85        "MTP stacked experts.down_proj: expected BF16, got {:?}",
86        down.dtype
87    );
88
89    let elt = WeightDtype::BF16.byte_size();
90    let half_bytes = intermediate * hidden * elt;
91    let gate_up_stride = two_inter * hidden * elt;
92    let down_stride = hidden * intermediate * elt;
93
94    let mut experts = Vec::with_capacity(num_experts);
95    for e in 0..num_experts {
96        let base_gu = gate_up.ptr.offset(e * gate_up_stride);
97        experts.push(DenseExpertWeight {
98            gate_proj: DenseWeight { weight: base_gu },
99            up_proj: DenseWeight {
100                weight: base_gu.offset(half_bytes),
101            },
102            down_proj: DenseWeight {
103                weight: down.ptr.offset(e * down_stride),
104            },
105        });
106    }
107    Ok(experts)
108}
109
110// ── Qwen3.5-MoE weight loaders ──
111// Two NVFP4 naming conventions exist:
112//   Standard (nvidia/txn545):  weight, weight_scale, weight_scale_2, input_scale
113//   Sehyo (compressed-tensors): weight_packed, weight_scale, weight_global_scale, input_global_scale
114// Additionally, Sehyo quantizes attention/SSM projections; standard keeps them BF16.
115
116/// Weight quantization variant (on-disk format).
117#[derive(Debug, Clone, Copy, PartialEq, Eq)]
118pub enum Nvfp4Variant {
119    /// Standard ModelOpt: weight, weight_scale, weight_scale_2, input_scale.
120    /// Attention/SSM projections are BF16 dense.
121    Standard,
122    /// Sehyo/compressed-tensors: weight_packed, weight_global_scale, input_global_scale.
123    /// Attention/SSM projections are NVFP4 quantized.
124    CompressedTensors,
125    /// FP8 block-scaled (e.g. Qwen/Qwen3.5-35B-A3B-FP8, Qwen/Qwen3.6-35B-A3B-FP8):
126    /// weight (float8_e4m3fn) + weight_scale_inv (BF16 per-`[128,128]`-block).
127    ///
128    /// Loaded **NATIVELY as FP8** in Qwen3 and Qwen3.5/3.6 model families.
129    /// Attention uses `w8a16_gemv` (decode) + `w8a16_gemm` (prefill).
130    /// MoE uses the FP8 fused grouped-GEMM batch1/2/3 path.
131    /// SSM uses `w8a16_gemv` decode + `fp8_gemm_n128` prefill (single-scale).
132    /// No silent FP8→BF16→NVFP4 triple-conversion.
133    ///
134    /// Historical note: the variant name retains "Dequanted" because the
135    /// `Bf16Raw` cousin and the pre-2026-05-24 NVFP4 detour did dequant on
136    /// load. The dispatch tables in `qwen35/load_layers.rs` (`LayerType::
137    /// FullAttention if native_fp8`) and `qwen3.rs` (line 176) now branch
138    /// to native FP8 paths when `quant_format == QuantFormat::Fp8`.
139    Fp8Dequanted,
140    /// Raw BF16/FP16 fine-tunes (e.g. samuelcardillo/Carnice-MoE-35B-A3B):
141    /// only `.weight` tensors exist (no quantization metadata). Runtime-quantize
142    /// from BF16 to NVFP4 at load time. Quality is suboptimal vs. a
143    /// pre-calibrated NVFP4 release — the user gets a warning at startup.
144    Bf16Raw,
145}
spark_model/weight_map/loaders_mtp.rs

spark_model/weight_map/
loaders_mtp.rs
