spark_model/weight_map/
moe.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/// MoE layer weights.
14pub struct MoeWeights {
15    /// Router gate: [hidden_size, num_experts] BF16.
16    pub gate: DenseWeight,
17    /// Shared expert (always active).
18    pub shared_expert: ExpertWeight,
19    /// Shared expert gate sigmoid weight: `[1]` BF16.
20    pub shared_expert_gate: DenseWeight,
21    /// Per-expert weights: 512 experts.
22    pub experts: Vec<ExpertWeight>,
23    /// Optional router pre-normalization weight.
24    /// Set for Gemma-4 MoE where the HF reference applies a pure RMSNorm to
25    /// the router input followed by a per-dim scale multiplication:
26    ///   `router_input = rms_norm(x) * scale * hidden_size^(-0.5)`
27    /// Stored as a BF16 `[hidden_size]` vector containing `scale * root_size`
28    /// so the existing rms_norm kernel (`output = x/rms(x) * weight`) applies
29    /// both steps in one pass. `None` for models that feed the router from
30    /// the raw post-attention residual.
31    pub router_pre_norm: Option<DenseWeight>,
32    /// Optional expert correction bias: `[num_experts]` F32.
33    ///
34    /// Set for models using the DeepSeek-V3 / MiniMax-M2 loss-free-balancing
35    /// routing trick: the bias is added to sigmoid(gate_logits) *only* for
36    /// top-k selection; gathered dispatch weights come from the unbiased
37    /// sigmoid scores. Consumed by `moe_topk_sigmoid` kernel via its `bias`
38    /// argument.
39    ///
40    /// `None` for softmax-routed Qwen/Gemma MoE. Nemotron-H carries its own
41    /// bias in `NemotronMoeWeights::e_score_correction_bias` because its
42    /// MoE is a separate layer type (Mamba-2 interleaved) — those paths
43    /// don't touch this struct.
44    pub correction_bias: Option<DenseWeight>,
45}
46
47impl MoeWeights {
48    /// Create empty MoeWeights for testing (all null pointers).
49    #[cfg(test)]
50    pub fn empty(num_experts: usize) -> Self {
51        let null_dense = DenseWeight {
52            weight: DevicePtr::NULL,
53        };
54        let null_quant = QuantizedWeight {
55            weight: DevicePtr::NULL,
56            weight_scale: DevicePtr::NULL,
57            weight_scale_2: 1.0,
58            input_scale: DevicePtr::NULL,
59            weight_scale_2_vec: DevicePtr::NULL,
60        };
61        let null_expert = ExpertWeight {
62            gate_proj: null_quant,
63            up_proj: null_quant,
64            down_proj: null_quant,
65        };
66        Self {
67            gate: null_dense,
68            shared_expert: null_expert,
69            shared_expert_gate: null_dense,
70            experts: vec![null_expert; num_experts],
71            router_pre_norm: None,
72            correction_bias: None,
73        }
74    }
75}
76
77/// All weights for one transformer layer.
78pub enum LayerWeights {
79    FullAttention {
80        input_norm: DenseWeight,
81        attn: AttentionWeights,
82        post_attn_norm: DenseWeight,
83        moe: MoeWeights,
84    },
85    LinearAttention {
86        input_norm: DenseWeight,
87        ssm: SsmWeights,
88        post_attn_norm: DenseWeight,
89        moe: MoeWeights,
90    },
91}
92
93/// BF16 expert weight (before NVFP4 quantization).
94#[derive(Debug, Clone, Copy)]
95pub struct DenseExpertWeight {
96    pub gate_proj: DenseWeight,
97    pub up_proj: DenseWeight,
98    pub down_proj: DenseWeight,
99}
100
101/// MTP (Multi-Token Prediction) head weights (all BF16 from safetensors).
102///
103/// Single decoder layer + concat projection. All projection weights are BF16
104/// and get quantized to NVFP4 at load time by the weight loader.
105pub struct MtpWeights {
106    /// RMSNorm on token embedding before concat: `[hidden_size]` BF16.
107    pub pre_fc_norm_embedding: DenseWeight,
108    /// RMSNorm on target hidden state before concat: `[hidden_size]` BF16.
109    pub pre_fc_norm_hidden: DenseWeight,
110    /// Concat projection: `[hidden_size, 2*hidden_size]` BF16.
111    pub fc: DenseWeight,
112    /// Input layernorm for the attention layer: `[hidden_size]` BF16.
113    pub input_layernorm: DenseWeight,
114    /// Attention projections (all BF16).
115    pub q_proj: DenseWeight,
116    pub k_proj: DenseWeight,
117    pub v_proj: DenseWeight,
118    pub o_proj: DenseWeight,
119    pub q_norm: DenseWeight,
120    pub k_norm: DenseWeight,
121    /// Post-attention layernorm: `[hidden_size]` BF16.
122    pub post_attn_layernorm: DenseWeight,
123    /// MoE router gate: [num_experts, hidden_size] BF16.
124    /// NULL when `dense_ffn` is `Some` (dense FFN MTP head).
125    pub moe_gate: DenseWeight,
126    /// Shared expert (BF16). NULL fields when `dense_ffn` is `Some`.
127    pub shared_expert: DenseExpertWeight,
128    /// Shared expert gate: [1, hidden_size] BF16.
129    /// NULL when `dense_ffn` is `Some`.
130    pub shared_expert_gate: DenseWeight,
131    /// Per-expert weights (512 experts, BF16).
132    /// Empty when `dense_ffn` is `Some`.
133    pub experts: Vec<DenseExpertWeight>,
134    /// Dense FFN triple (`gate_proj`, `up_proj`, `down_proj`) — used by MTP
135    /// heads bundled with dense (non-MoE) FP8 checkpoints, e.g.
136    /// `Qwen/Qwen3.6-27B-FP8`. When `Some`, the MoE fields above are unused
137    /// and the forward path takes the dense MLP shortcut.
138    pub dense_ffn: Option<DenseExpertWeight>,
139    /// Final output RMSNorm: `[hidden_size]` BF16.
140    pub norm: DenseWeight,
141}
spark_model/weight_map/moe.rs

spark_model/weight_map/
moe.rs
