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    quantized_mxfp4_e8m0_pair(
138        store,
139        &format!("{prefix}.weight"),
140        &format!("{prefix}.scale"),
141    )
142}
143
144/// Shared native MXFP4 lander for DeepSeek and K3 checkpoint key conventions.
145pub(crate) fn quantized_mxfp4_e8m0_pair(
146    store: &WeightStore,
147    weight_key: &str,
148    scale_key: &str,
149) -> Result<QuantizedWeight> {
150    let w = store.get(weight_key)?;
151    ensure!(
152        w.shape.len() == 2,
153        "{weight_key}: expected packed rank-2 matrix"
154    );
155    let n = w.shape[0];
156    let k_packed = w.shape[1];
157    let total_nibbles = n * k_packed * 2;
158    let scale_t = store.get(scale_key)?;
159    let num_groups = scale_t.num_elements();
160    ensure!(
161        num_groups > 0 && total_nibbles.is_multiple_of(num_groups),
162        "{weight_key}: MXFP4 weight nibbles {total_nibbles} not divisible by E8M0 scale groups {num_groups}"
163    );
164    let block = total_nibbles / num_groups;
165    ensure!(
166        block == 32,
167        "{weight_key}: native MXFP4 expects GROUP_SIZE=32, inferred {block} (scale groups {num_groups}) \
168         — refusing to land a non-MX checkpoint on the transcode-free path"
169    );
170    Ok(QuantizedWeight {
171        weight: ptr(store, weight_key)?,
172        weight_scale: ptr(store, scale_key)?,
173        weight_scale_2: 1.0, // native MXFP4 has no per-tensor global
174        input_scale: DevicePtr::NULL,
175        // native MXFP4 uses the scalar `weight_scale_2` (E8M0 per-group), not
176        // the per-output-row scale2 vector added by #257 → NULL.
177        weight_scale_2_vec: DevicePtr::NULL,
178    })
179}
180
181pub(crate) fn dense(store: &WeightStore, name: &str) -> Result<DenseWeight> {
182    let w = store.get(name)?;
183    Ok(DenseWeight { weight: w.ptr })
184}
185
186// REMOVED: `dense_minus_one` — the offset-from-1 norm loader.
187//
188// It pre-subtracted 1.0 and stored `bf16(w - 1)` so the offset-from-1 `rms_norm`
189// kernel would recover `1 + (w - 1) = w`. That round-trip is only lossless when
190// `w ≈ 1`. DeepSeek-V4 — its ONLY caller — ships HF-vanilla norm weights of
191// ≈ 0.03, so it stored ≈ −0.97 and BF16's rounding error there (~1.9e-3 absolute)
192// became a 1.8–3.4 % RELATIVE error on the weight once 1 was added back
193// (catastrophic cancellation; up to 19 % on `q_norm`, 100 % with sign flips on the
194// compressor norms — measured over all 249 V4 norm tensors, 2026-07-13).
195//
196// V4 now loads its norm weights exactly (`dense_auto`) and dispatches
197// `rms_norm_vanilla`. See `crate::ships_vanilla_norm_weights`. Models whose norm
198// weights are genuinely stored as an offset (Qwen3-Next, init 0) use `dense` and
199// keep the offset kernel — they never needed this function.
200
201/// Load a weight, auto-dequanting FP8 block-scaled to BF16 when needed.
202///
203/// Used for models with mixed-precision layers — Qwen3.6's ViT, for
204/// example, keeps the first four blocks in BF16 but stores the rest as
205/// FP8. Callers pass a prefix (e.g. `"model.visual.blocks.5.attn.qkv"`)
206/// and get back a BF16 GPU buffer either way.
207pub(crate) fn dense_auto_fp8_or_bf16(
208    store: &WeightStore,
209    prefix: &str,
210    gpu: &dyn GpuBackend,
211) -> Result<DenseWeight> {
212    let w = store.get(&format!("{prefix}.weight"))?;
213    match w.dtype {
214        WeightDtype::BF16 => Ok(DenseWeight { weight: w.ptr }),
215        WeightDtype::FP8E4M3 => dequant_fp8_blockscaled_to_bf16(store, prefix, gpu),
216        other => anyhow::bail!(
217            "dense_auto_fp8_or_bf16: unsupported dtype {:?} for {prefix}.weight",
218            other
219        ),
220    }
221}
222
223/// Load a dense weight, converting FP32 → BF16 if needed. Used for norm weights
224/// that may be FP32 in some checkpoints (e.g. Qwen FP8 `linear_attn.norm.weight`).
225pub(crate) fn dense_f32_safe(
226    store: &WeightStore,
227    name: &str,
228    gpu: &dyn GpuBackend,
229) -> Result<DenseWeight> {
230    let w = store.get(name)?;
231    if w.dtype == WeightDtype::FP32 {
232        // On-device FP32→BF16 truncation — ONE async kernel on the load stream,
233        // no D2H/CPU/H2D round-trip. The old path did copy_d2h→CPU-truncate→
234        // copy_h2d per weight, each with 2 cuStreamSynchronize on the busy load
235        // stream → ~104s across 635 FP32 weights (the dominant cold-load cost).
236        // The kernel reads the high 2 bytes of each f32 → bit-identical to the
237        // prior CPU truncation. Ordered after the weight's upload (same stream),
238        // so no sync needed.
239        let n = w.num_elements();
240        let ptr = gpu.alloc(n * 2)?;
241        let trunc = gpu.kernel("quantize_nvfp4", "f32_to_bf16_trunc")?;
242        let blocks = (n.div_ceil(256) as u32).max(1);
243        spark_runtime::kernel_args::KernelLaunch::new(gpu, trunc)
244            .grid([blocks, 1, 1])
245            .block([256, 1, 1])
246            .arg_ptr(w.ptr)
247            .arg_ptr(ptr)
248            .arg_u32(n as u32)
249            .launch(gpu.default_stream())?;
250        Ok(DenseWeight { weight: ptr })
251    } else {
252        Ok(DenseWeight { weight: w.ptr })
253    }
254}
255
256/// Load a weight and ensure it's FP32 on GPU, regardless of source dtype.
257///
258/// Used for SSM gate parameters (A_log, dt_bias) where BF16 precision
259/// causes exponential error amplification in the recurrent state at
260/// long context (8k+ tokens). A 1-ULP BF16 error in the decay gate
261/// produces (g_correct/g_error)^8000 ≈ 3000x magnitude divergence.
262///
263/// - FP32 in safetensors: keep as-is (no conversion)
264/// - BF16 in safetensors: convert BF16 → FP32 via zero-extension
265pub(crate) fn dense_keep_f32(
266    store: &WeightStore,
267    name: &str,
268    gpu: &dyn GpuBackend,
269) -> Result<DenseWeight> {
270    let w = store.get(name)?;
271    match w.dtype {
272        WeightDtype::FP32 => {
273            // Already FP32 — use directly, no conversion needed
274            Ok(DenseWeight { weight: w.ptr })
275        }
276        WeightDtype::BF16 => {
277            // Convert BF16 → FP32 to preserve precision
278            tracing::info!(
279                "dense_keep_f32: promoting {name} from BF16 to FP32 ({:?})",
280                w.shape
281            );
282            let n = w.num_elements();
283            let mut bf16_buf = vec![0u8; n * 2];
284            gpu.copy_d2h(w.ptr, &mut bf16_buf)?;
285            let f32_buf: Vec<u8> = bf16_buf
286                .chunks_exact(2)
287                .flat_map(|c| {
288                    let bits = u16::from_le_bytes([c[0], c[1]]);
289                    let f32_bits = (bits as u32) << 16;
290                    f32_bits.to_le_bytes()
291                })
292                .collect();
293            let ptr = gpu.alloc(f32_buf.len())?;
294            gpu.copy_h2d(&f32_buf, ptr)?;
295            Ok(DenseWeight { weight: ptr })
296        }
297        other => {
298            bail!("dense_keep_f32: unsupported dtype {:?} for {name}", other);
299        }
300    }
301}
302
303/// Load a BF16 tensor and convert to F32 on-device via CPU roundtrip.
304///
305/// Used for Nemotron-H SSM parameters (A_log, D, dt_bias, conv1d.bias)
306/// which are stored as BF16 in safetensors but consumed as F32 by CUDA kernels.
307pub(crate) fn dense_bf16_as_f32(
308    store: &WeightStore,
309    name: &str,
310    gpu: &dyn GpuBackend,
311) -> Result<DenseWeight> {
312    let w = store.get(name)?;
313    ensure!(
314        w.dtype == WeightDtype::BF16,
315        "Expected BF16 for {name}, got {:?}",
316        w.dtype
317    );
318    let n = w.num_elements();
319    let mut bf16_buf = vec![0u8; n * 2];
320    gpu.copy_d2h(w.ptr, &mut bf16_buf)?;
321    let f32_buf: Vec<u8> = bf16_buf
322        .chunks_exact(2)
323        .flat_map(|c| {
324            let bits = u16::from_le_bytes([c[0], c[1]]);
325            let f32_bits = (bits as u32) << 16;
326            f32_bits.to_le_bytes()
327        })
328        .collect();
329    let ptr = gpu.alloc(f32_buf.len())?;
330    gpu.copy_h2d(&f32_buf, ptr)?;
331    Ok(DenseWeight { weight: ptr })
332}
333
334/// Load an F32 tensor and convert to BF16 on-device via CPU roundtrip.
335///
336/// Used for Nemotron-H gate weights (F32 in safetensors, consumed as BF16).
337pub(crate) fn dense_f32_as_bf16(
338    store: &WeightStore,
339    name: &str,
340    gpu: &dyn GpuBackend,
341) -> Result<DenseWeight> {
342    let w = store.get(name)?;
343    ensure!(
344        w.dtype == WeightDtype::FP32,
345        "Expected FP32 for {name}, got {:?}",
346        w.dtype
347    );
348    let n = w.num_elements();
349    let mut f32_buf = vec![0u8; n * 4];
350    gpu.copy_d2h(w.ptr, &mut f32_buf)?;
351    let bf16_buf: Vec<u8> = f32_buf
352        .chunks_exact(4)
353        .flat_map(|c| {
354            let val = f32::from_le_bytes([c[0], c[1], c[2], c[3]]);
355            f32_to_bf16(val).to_le_bytes()
356        })
357        .collect();
358    let ptr = gpu.alloc(bf16_buf.len())?;
359    gpu.copy_h2d(&bf16_buf, ptr)?;
360    Ok(DenseWeight { weight: ptr })
361}
spark_model/weight_map/model_a.rs

spark_model/weight_map/
model_a.rs
