spark_model/weight_loader/
step3p7.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Step 3.7 Flash weight loader.
4//!
5//! Hybrid of MiniMax M2 and Qwen 3.5 patterns:
6//!   * Sigmoid MoE routing + correction bias (MiniMax M2 pattern)
7//!   * Shared expert per MoE layer (Qwen 3.5 pattern)
8//!   * Attention gate g_proj (Qwen 3.5 pattern)
9//!   * Partial RoPE 0.5 (MiniMax M2 pattern)
10//!   * Per-head q_norm / k_norm
11//!   * Mixed dense FFN (layers 0-2) + MoE (layers 3-44)
12//!   * 3 MTP modules at layers 45-47 (different prefix: `model.layers.`)
13//!
14//! Weight prefix: `model.language_model.layers.{i}` for main layers.
15//! MTP prefix: `model.layers.{45|46|47}` (different namespace!).
16//!
17//! KEY ARCHITECTURAL DIFFERENCE: Step 3.7 stores expert weights as FUSED
18//! tensors — one tensor per projection type containing ALL 288 experts
19//! concatenated. Atlas needs per-expert QuantizedWeight entries, so we
20//! slice by computing byte offsets into the fused GPU allocations.
21//!
22//! NVFP4 format: ModelOpt style with `weight`, `weight_scale`, `weight_scale_2`,
23//! `input_scale` per projection. Shared expert is BF16.
24
25mod load_layers;
26
27use anyhow::Result;
28use atlas_core::config::ModelConfig;
29use spark_runtime::gpu::{DevicePtr, GpuBackend};
30use spark_runtime::kv_cache::KvCacheDtype;
31use spark_runtime::weights::WeightStore;
32
33use super::ModelWeightLoader;
34use crate::layer::TransformerLayer;
35use crate::weight_map::{DenseWeight, MtpWeights, QuantizedWeight, dense};
36
37pub struct Step3p7WeightLoader;
38
39/// Step 3.7 uses shifted RMSNorm: `output = (x / rms) * (weight + 1)`.
40/// The checkpoint stores norm weights centered around 0, not 1.
41/// This function adds 1.0 to each element so the standard RMSNorm kernel
42/// `output = (x / rms) * weight` produces the correct result.
43fn offset_norm_weights_plus_one(
44    weight: &DenseWeight,
45    size: usize,
46    gpu: &dyn GpuBackend,
47) -> Result<()> {
48    let byte_len = size * 2; // BF16 = 2 bytes per element
49    let mut buf = vec![0u8; byte_len];
50    gpu.copy_d2h(weight.weight, &mut buf)?;
51
52    for i in 0..size {
53        let bits = u16::from_le_bytes([buf[i * 2], buf[i * 2 + 1]]);
54        let f32_val = f32::from_bits((bits as u32) << 16);
55        let new_val = f32_val + 1.0;
56        // Round to BF16: add 0x7FFF + bit 16 for round-to-nearest-even
57        let f32_bits = new_val.to_bits();
58        let new_bits = ((f32_bits + 0x7FFF + ((f32_bits >> 16) & 1)) >> 16) as u16;
59        buf[i * 2] = new_bits as u8;
60        buf[i * 2 + 1] = (new_bits >> 8) as u8;
61    }
62
63    gpu.copy_h2d(&buf, weight.weight)?;
64    Ok(())
65}
66
67/// Slice a fused NVFP4 tensor into per-expert QuantizedWeight entries.
68///
69/// Step 3.7's original checkpoint stores all experts in one contiguous
70/// tensor per projection:
71///   weight: [num_experts * n, k] packed NVFP4 (0.5 bytes/element)
72///   weight_scale: [num_experts * n, k/group_size] FP8 per-group scales
73///   input_scale: [num_experts * n] (optional, activation quantization)
74///
75/// This function creates `num_experts` QuantizedWeight entries, each
76/// pointing to a different offset within the fused allocations.
77fn slice_fused_experts(
78    fused_weight: DevicePtr,
79    fused_scale: DevicePtr,
80    fused_input_scale: DevicePtr,
81    global_scale_2: f32,
82    num_experts: usize,
83    n: usize,
84    k: usize,
85) -> Vec<QuantizedWeight> {
86    let group_size = 16usize;
87    let packed_bytes_per_expert = n * k / 2;
88    let scale_bytes_per_expert = n * k.div_ceil(group_size);
89    let input_scale_bytes_per_expert = n * 4;
90
91    (0..num_experts)
92        .map(|e| QuantizedWeight {
93            weight: fused_weight.offset(e * packed_bytes_per_expert),
94            weight_scale: fused_scale.offset(e * scale_bytes_per_expert),
95            weight_scale_2: global_scale_2,
96            input_scale: if fused_input_scale == DevicePtr::NULL {
97                DevicePtr::NULL
98            } else {
99                fused_input_scale.offset(e * input_scale_bytes_per_expert)
100            },
101            weight_scale_2_vec: DevicePtr::NULL,
102        })
103        .collect()
104}
105
106/// Detect whether this checkpoint uses per-expert tensor format.
107fn has_per_expert_tensors(store: &WeightStore, layer_prefix: &str) -> bool {
108    let pattern = format!("{layer_prefix}.moe.experts.");
109    let found = store.names().any(|k| k.starts_with(&pattern));
110    tracing::debug!("has_per_expert_tensors('{layer_prefix}'): pattern='{pattern}', found={found}");
111    found
112}
113
114/// Load a fused NVFP4 tensor from the store (Standard ModelOpt format).
115fn load_fused_nvfp4(
116    store: &WeightStore,
117    prefix: &str,
118    gpu: &dyn GpuBackend,
119) -> Result<(DevicePtr, DevicePtr, DevicePtr, f32)> {
120    let weight = store.get(&format!("{prefix}.weight"))?.ptr;
121    let weight_scale = store.get(&format!("{prefix}.weight_scale"))?.ptr;
122
123    let ws2_key = format!("{prefix}.weight_scale_2");
124    let ws2_ptr = store.get(&ws2_key)?.ptr;
125    let mut ws2_buf = [0u8; 4];
126    gpu.copy_d2h(ws2_ptr, &mut ws2_buf)?;
127    let weight_scale_2 = f32::from_le_bytes(ws2_buf);
128
129    let is_key = format!("{prefix}.input_scale");
130    let input_scale = if store.contains(&is_key) {
131        store.get(&is_key)?.ptr
132    } else {
133        DevicePtr::NULL
134    };
135
136    Ok((weight, weight_scale, input_scale, weight_scale_2))
137}
138
139impl ModelWeightLoader for Step3p7WeightLoader {
140    fn supports_tp(&self) -> bool {
141        false // Single-GPU initial bring-up
142    }
143
144    fn load_layers(
145        &self,
146        store: &WeightStore,
147        config: &ModelConfig,
148        gpu: &dyn GpuBackend,
149        layer_kv_dtypes: &[KvCacheDtype],
150    ) -> Result<Vec<Box<dyn TransformerLayer>>> {
151        load_layers::load_layers(store, config, gpu, layer_kv_dtypes)
152    }
153
154    fn load_embedding(
155        &self,
156        store: &WeightStore,
157        config: &ModelConfig,
158        _gpu: &dyn GpuBackend,
159    ) -> Result<DenseWeight> {
160        let prefix = if config.weight_prefix.is_empty() {
161            "model.language_model"
162        } else {
163            &config.weight_prefix
164        };
165        dense(store, &format!("{prefix}.embed_tokens.weight"))
166    }
167
168    fn load_final_norm(
169        &self,
170        store: &WeightStore,
171        config: &ModelConfig,
172        gpu: &dyn GpuBackend,
173    ) -> Result<DenseWeight> {
174        let prefix = if config.weight_prefix.is_empty() {
175            "model.language_model"
176        } else {
177            &config.weight_prefix
178        };
179        let w = dense(store, &format!("{prefix}.norm.weight"))?;
180        offset_norm_weights_plus_one(&w, config.hidden_size, gpu)?;
181        Ok(w)
182    }
183
184    fn load_lm_head(
185        &self,
186        store: &WeightStore,
187        _config: &ModelConfig,
188        _gpu: &dyn GpuBackend,
189    ) -> Result<DenseWeight> {
190        dense(store, "lm_head.weight")
191    }
192
193    fn load_mtp_weights(
194        &self,
195        _store: &WeightStore,
196        _config: &ModelConfig,
197        _gpu: &dyn GpuBackend,
198    ) -> Result<Option<MtpWeights>> {
199        Ok(None) // Multi-module MTP — use load_mtp_weights_multi
200    }
201
202    fn load_mtp_weights_multi(
203        &self,
204        store: &WeightStore,
205        config: &ModelConfig,
206        _gpu: &dyn GpuBackend,
207    ) -> Result<Vec<MtpWeights>> {
208        let first_mtp_idx = config.num_hidden_layers;
209        let probe = format!("model.layers.{first_mtp_idx}.input_layernorm.weight");
210        if !store.contains(&probe) {
211            tracing::info!(
212                "step3p7: no MTP module weights found \
213                 (expected at layer {first_mtp_idx}); MTP disabled"
214            );
215            return Ok(Vec::new());
216        }
217
218        tracing::info!(
219            "step3p7: MTP module weights detected at layers {}-{} but MTP loader \
220             not yet implemented. Run with --speculative 0 for non-MTP decode.",
221            first_mtp_idx,
222            first_mtp_idx + config.mtp_num_hidden_layers - 1,
223        );
224        Ok(Vec::new())
225    }
226}
spark_model/weight_loader/step3p7.rs

spark_model/weight_loader/
step3p7.rs
