spark_model/weight_map/
model_a.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/// All model weights organized by layer.
14pub struct ModelWeights {
15    /// Embedding table: [vocab_size, hidden_size] BF16.
16    pub embed_tokens: DenseWeight,
17    /// Final RMS norm: `[hidden_size]` BF16.
18    pub final_norm: DenseWeight,
19    /// LM head: [hidden_size, vocab_size] BF16.
20    pub lm_head: DenseWeight,
21    /// Per-layer weights.
22    pub layers: Vec<LayerWeights>,
23}
24
25/// Extract a DevicePtr from the weight store by name.
26pub(crate) fn ptr(store: &WeightStore, name: &str) -> Result<DevicePtr> {
27    Ok(store.get(name)?.ptr)
28}
29
30/// Extract a scalar f32 from a single-element weight tensor via D2H copy.
31pub(crate) fn scalar_f32(store: &WeightStore, name: &str, gpu: &dyn GpuBackend) -> Result<f32> {
32    let w = store.get(name)?;
33    ensure!(
34        w.dtype == WeightDtype::FP32,
35        "Expected FP32 for {name}, got {:?}",
36        w.dtype
37    );
38    ensure!(
39        w.num_elements() == 1,
40        "Expected scalar for {name}, got {} elements",
41        w.num_elements()
42    );
43    let mut buf = [0u8; 4];
44    gpu.copy_d2h(w.ptr, &mut buf)?;
45    Ok(f32::from_le_bytes(buf))
46}
47
48/// Load FP8 KV cache quantization scales from a checkpoint.
49///
50/// Searches for `{attn_prefix}.k_proj.k_scale` and `{attn_prefix}.v_proj.v_scale`
51/// tensors in the weight store. When found, returns the scalar f32 values from the
52/// checkpoint (calibrated per-tensor scales for FP8 KV cache quantization).
53///
54/// When not found, returns (1.0, 1.0) with a debug-level log. This fallback is
55/// correct for BF16 KV cache (scales are unused) and acceptable as uncalibrated
56/// default for FP8 KV cache (Workstream 4C will add proper calibration).
57pub(crate) fn load_kv_scales(
58    store: &WeightStore,
59    attn_prefix: &str,
60    gpu: &dyn GpuBackend,
61) -> (f32, f32) {
62    let k_key = format!("{attn_prefix}.k_proj.k_scale");
63    let v_key = format!("{attn_prefix}.v_proj.v_scale");
64
65    let k_scale = if store.contains(&k_key) {
66        match scalar_f32(store, &k_key, gpu) {
67            Ok(v) => {
68                tracing::debug!("Loaded k_scale={v:.6} from {k_key}");
69                v
70            }
71            Err(e) => {
72                tracing::warn!("Failed to load {k_key}: {e:#}, using 1.0");
73                1.0
74            }
75        }
76    } else {
77        tracing::debug!("No {k_key} in checkpoint, using k_scale=1.0");
78        1.0
79    };
80
81    let v_scale = if store.contains(&v_key) {
82        match scalar_f32(store, &v_key, gpu) {
83            Ok(v) => {
84                tracing::debug!("Loaded v_scale={v:.6} from {v_key}");
85                v
86            }
87            Err(e) => {
88                tracing::warn!("Failed to load {v_key}: {e:#}, using 1.0");
89                1.0
90            }
91        }
92    } else {
93        tracing::debug!("No {v_key} in checkpoint, using v_scale=1.0");
94        1.0
95    };
96
97    (k_scale, v_scale)
98}
99
100/// Build a QuantizedWeight from the store using the standard NVFP4 naming.
101///
102/// `weight_scale_2` is a single FP32 scalar — extracted from GPU via D2H copy.
103pub(crate) fn quantized(
104    store: &WeightStore,
105    prefix: &str,
106    gpu: &dyn GpuBackend,
107) -> Result<QuantizedWeight> {
108    let input_scale_key = format!("{prefix}.input_scale");
109    Ok(QuantizedWeight {
110        weight: ptr(store, &format!("{prefix}.weight"))?,
111        weight_scale: ptr(store, &format!("{prefix}.weight_scale"))?,
112        weight_scale_2: scalar_f32(store, &format!("{prefix}.weight_scale_2"), gpu)?,
113        input_scale: if store.contains(&input_scale_key) {
114            ptr(store, &input_scale_key)?
115        } else {
116            DevicePtr::NULL
117        },
118        weight_scale_2_vec: DevicePtr::NULL,
119    })
120}
121
122/// Native MXFP4 routed expert: land the on-disk bytes device-resident
123/// **UNCHANGED** — no dequant, no re-quantize, no dtype coercion (the
124/// transcode-free path, contrast `quantized_from_fp8` / the old E8M0
125/// `dequant→quantize_to_nvfp4` arm that cost TWO lossy 4-bit conversions).
126///
127/// On-disk layout (DeepSeek-V4-Flash ORIGINAL routed format):
128///   - `.weight` = 4-bit E2M1 nibbles, 2 per byte, stored U8/I8, shape `[n, k/2]`
129///   - `.scale`  = F8_E8M0 per-block (biased exponent; scale = `2^(byte-127)`),
130///     `GROUP_SIZE=32`, NO per-tensor global.
131///
132/// The buffer is tagged `WeightQuantFormat::Mxfp4E8m0` at the MoE-layer level
133/// (`MoeWeights::experts_scale_kind`); the E8M0 kernel variants (Phase-K)
134/// consume `weight_scale` as E8M0 bytes and ignore `weight_scale_2`. Asserts
135/// the inferred block size is 32 so a non-MX checkpoint can't slip through.
136pub(crate) fn quantized_mxfp4_e8m0(store: &WeightStore, prefix: &str) -> Result<QuantizedWeight> {
137    let w = store.get(&format!("{prefix}.weight"))?;
138    let n = w.shape[0];
139    let k_packed = w.shape[1];
140    let total_nibbles = n * k_packed * 2;
141    let scale_t = store.get(&format!("{prefix}.scale"))?;
142    let num_groups = scale_t.num_elements();
143    ensure!(
144        num_groups > 0 && total_nibbles.is_multiple_of(num_groups),
145        "{prefix}: MXFP4 weight nibbles {total_nibbles} not divisible by E8M0 scale groups {num_groups}"
146    );
147    let block = total_nibbles / num_groups;
148    ensure!(
149        block == 32,
150        "{prefix}: native MXFP4 expects GROUP_SIZE=32, inferred {block} (scale groups {num_groups}) \
151         — refusing to land a non-MX checkpoint on the transcode-free path"
152    );
153    Ok(QuantizedWeight {
154        weight: ptr(store, &format!("{prefix}.weight"))?,
155        weight_scale: ptr(store, &format!("{prefix}.scale"))?,
156        weight_scale_2: 1.0, // native MXFP4 has no per-tensor global
157        input_scale: DevicePtr::NULL,
158        // native MXFP4 uses the scalar `weight_scale_2` (E8M0 per-group), not
159        // the per-output-row scale2 vector added by #257 → NULL.
160        weight_scale_2_vec: DevicePtr::NULL,
161    })
162}
163
164pub(crate) fn dense(store: &WeightStore, name: &str) -> Result<DenseWeight> {
165    let w = store.get(name)?;
166    Ok(DenseWeight { weight: w.ptr })
167}
168
169// REMOVED: `dense_minus_one` — the offset-from-1 norm loader.
170//
171// It pre-subtracted 1.0 and stored `bf16(w - 1)` so the offset-from-1 `rms_norm`
172// kernel would recover `1 + (w - 1) = w`. That round-trip is only lossless when
173// `w ≈ 1`. DeepSeek-V4 — its ONLY caller — ships HF-vanilla norm weights of
174// ≈ 0.03, so it stored ≈ −0.97 and BF16's rounding error there (~1.9e-3 absolute)
175// became a 1.8–3.4 % RELATIVE error on the weight once 1 was added back
176// (catastrophic cancellation; up to 19 % on `q_norm`, 100 % with sign flips on the
177// compressor norms — measured over all 249 V4 norm tensors, 2026-07-13).
178//
179// V4 now loads its norm weights exactly (`dense_auto`) and dispatches
180// `rms_norm_vanilla`. See `crate::ships_vanilla_norm_weights`. Models whose norm
181// weights are genuinely stored as an offset (Qwen3-Next, init 0) use `dense` and
182// keep the offset kernel — they never needed this function.
183
184/// Load a weight, auto-dequanting FP8 block-scaled to BF16 when needed.
185///
186/// Used for models with mixed-precision layers — Qwen3.6's ViT, for
187/// example, keeps the first four blocks in BF16 but stores the rest as
188/// FP8. Callers pass a prefix (e.g. `"model.visual.blocks.5.attn.qkv"`)
189/// and get back a BF16 GPU buffer either way.
190pub(crate) fn dense_auto_fp8_or_bf16(
191    store: &WeightStore,
192    prefix: &str,
193    gpu: &dyn GpuBackend,
194) -> Result<DenseWeight> {
195    let w = store.get(&format!("{prefix}.weight"))?;
196    match w.dtype {
197        WeightDtype::BF16 => Ok(DenseWeight { weight: w.ptr }),
198        WeightDtype::FP8E4M3 => dequant_fp8_blockscaled_to_bf16(store, prefix, gpu),
199        other => anyhow::bail!(
200            "dense_auto_fp8_or_bf16: unsupported dtype {:?} for {prefix}.weight",
201            other
202        ),
203    }
204}
205
206/// Load a dense weight, converting FP32 → BF16 if needed. Used for norm weights
207/// that may be FP32 in some checkpoints (e.g. Qwen FP8 `linear_attn.norm.weight`).
208pub(crate) fn dense_f32_safe(
209    store: &WeightStore,
210    name: &str,
211    gpu: &dyn GpuBackend,
212) -> Result<DenseWeight> {
213    let w = store.get(name)?;
214    if w.dtype == WeightDtype::FP32 {
215        // On-device FP32→BF16 truncation — ONE async kernel on the load stream,
216        // no D2H/CPU/H2D round-trip. The old path did copy_d2h→CPU-truncate→
217        // copy_h2d per weight, each with 2 cuStreamSynchronize on the busy load
218        // stream → ~104s across 635 FP32 weights (the dominant cold-load cost).
219        // The kernel reads the high 2 bytes of each f32 → bit-identical to the
220        // prior CPU truncation. Ordered after the weight's upload (same stream),
221        // so no sync needed.
222        let n = w.num_elements();
223        let ptr = gpu.alloc(n * 2)?;
224        let trunc = gpu.kernel("quantize_nvfp4", "f32_to_bf16_trunc")?;
225        let blocks = (n.div_ceil(256) as u32).max(1);
226        spark_runtime::kernel_args::KernelLaunch::new(gpu, trunc)
227            .grid([blocks, 1, 1])
228            .block([256, 1, 1])
229            .arg_ptr(w.ptr)
230            .arg_ptr(ptr)
231            .arg_u32(n as u32)
232            .launch(gpu.default_stream())?;
233        Ok(DenseWeight { weight: ptr })
234    } else {
235        Ok(DenseWeight { weight: w.ptr })
236    }
237}
238
239/// Load a weight and ensure it's FP32 on GPU, regardless of source dtype.
240///
241/// Used for SSM gate parameters (A_log, dt_bias) where BF16 precision
242/// causes exponential error amplification in the recurrent state at
243/// long context (8k+ tokens). A 1-ULP BF16 error in the decay gate
244/// produces (g_correct/g_error)^8000 ≈ 3000x magnitude divergence.
245///
246/// - FP32 in safetensors: keep as-is (no conversion)
247/// - BF16 in safetensors: convert BF16 → FP32 via zero-extension
248pub(crate) fn dense_keep_f32(
249    store: &WeightStore,
250    name: &str,
251    gpu: &dyn GpuBackend,
252) -> Result<DenseWeight> {
253    let w = store.get(name)?;
254    match w.dtype {
255        WeightDtype::FP32 => {
256            // Already FP32 — use directly, no conversion needed
257            Ok(DenseWeight { weight: w.ptr })
258        }
259        WeightDtype::BF16 => {
260            // Convert BF16 → FP32 to preserve precision
261            tracing::info!(
262                "dense_keep_f32: promoting {name} from BF16 to FP32 ({:?})",
263                w.shape
264            );
265            let n = w.num_elements();
266            let mut bf16_buf = vec![0u8; n * 2];
267            gpu.copy_d2h(w.ptr, &mut bf16_buf)?;
268            let f32_buf: Vec<u8> = bf16_buf
269                .chunks_exact(2)
270                .flat_map(|c| {
271                    let bits = u16::from_le_bytes([c[0], c[1]]);
272                    let f32_bits = (bits as u32) << 16;
273                    f32_bits.to_le_bytes()
274                })
275                .collect();
276            let ptr = gpu.alloc(f32_buf.len())?;
277            gpu.copy_h2d(&f32_buf, ptr)?;
278            Ok(DenseWeight { weight: ptr })
279        }
280        other => {
281            bail!("dense_keep_f32: unsupported dtype {:?} for {name}", other);
282        }
283    }
284}
285
286/// Load a BF16 tensor and convert to F32 on-device via CPU roundtrip.
287///
288/// Used for Nemotron-H SSM parameters (A_log, D, dt_bias, conv1d.bias)
289/// which are stored as BF16 in safetensors but consumed as F32 by CUDA kernels.
290pub(crate) fn dense_bf16_as_f32(
291    store: &WeightStore,
292    name: &str,
293    gpu: &dyn GpuBackend,
294) -> Result<DenseWeight> {
295    let w = store.get(name)?;
296    ensure!(
297        w.dtype == WeightDtype::BF16,
298        "Expected BF16 for {name}, got {:?}",
299        w.dtype
300    );
301    let n = w.num_elements();
302    let mut bf16_buf = vec![0u8; n * 2];
303    gpu.copy_d2h(w.ptr, &mut bf16_buf)?;
304    let f32_buf: Vec<u8> = bf16_buf
305        .chunks_exact(2)
306        .flat_map(|c| {
307            let bits = u16::from_le_bytes([c[0], c[1]]);
308            let f32_bits = (bits as u32) << 16;
309            f32_bits.to_le_bytes()
310        })
311        .collect();
312    let ptr = gpu.alloc(f32_buf.len())?;
313    gpu.copy_h2d(&f32_buf, ptr)?;
314    Ok(DenseWeight { weight: ptr })
315}
316
317/// Load an F32 tensor and convert to BF16 on-device via CPU roundtrip.
318///
319/// Used for Nemotron-H gate weights (F32 in safetensors, consumed as BF16).
320pub(crate) fn dense_f32_as_bf16(
321    store: &WeightStore,
322    name: &str,
323    gpu: &dyn GpuBackend,
324) -> Result<DenseWeight> {
325    let w = store.get(name)?;
326    ensure!(
327        w.dtype == WeightDtype::FP32,
328        "Expected FP32 for {name}, got {:?}",
329        w.dtype
330    );
331    let n = w.num_elements();
332    let mut f32_buf = vec![0u8; n * 4];
333    gpu.copy_d2h(w.ptr, &mut f32_buf)?;
334    let bf16_buf: Vec<u8> = f32_buf
335        .chunks_exact(4)
336        .flat_map(|c| {
337            let val = f32::from_le_bytes([c[0], c[1], c[2], c[3]]);
338            f32_to_bf16(val).to_le_bytes()
339        })
340        .collect();
341    let ptr = gpu.alloc(bf16_buf.len())?;
342    gpu.copy_h2d(&bf16_buf, ptr)?;
343    Ok(DenseWeight { weight: ptr })
344}
spark_model/weight_map/model_a.rs

spark_model/weight_map/
model_a.rs
