spark_model/weight_map/
loaders_fp8.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/// Load an FP8 E4M3 block-scaled checkpoint weight as a native [`Fp8Weight`].
14///
15/// The FP8 checkpoint stores:
16///   - `{prefix}.weight`: FP8E4M3 tensor [N, K]
17///   - `{prefix}.weight_scale_inv`: BF16 (Qwen/DeepSeek) or FP32 (MiniMax)
18///     tensor [N/block, K/block]
19///
20/// The `w8a16_gemv` kernel uses 2D block scales directly:
21///   `dequant[i,j] = E4M3_LUT[fp8[i,j]] * block_scale[i/BS, j/BS]`
22/// No per-row max reduction needed — the kernel loads the correct block
23/// scale for each 128-element K chunk.
24///
25/// **Scale precision (block-FP8 numerics):** the block scale is *widened to a
26/// genuine FP32 device buffer here, once*, so it is applied in full FP32 in the
27/// W8A8/W8A16 GEMM epilogues — matching vLLM / DeepGEMM / HF block-FP8 (which
28/// also accumulate the scale in FP32). The checkpoint may store the scale as
29/// BF16 (lossless widen), FP32 (straight copy), or F8_E8M0 (exact power-of-two
30/// widen); in every case `row_scale` ends up an FP32 `[N/BS, K/BS]` buffer.
31/// Every FP8 block-scale kernel reads
32/// `const float*` — see `kernels/gb10/common/w8a16_gemv.cu` et al.
33pub fn load_fp8_block_scaled_as_fp8weight(
34    store: &WeightStore,
35    prefix: &str,
36    gpu: &dyn GpuBackend,
37) -> Result<Fp8Weight> {
38    let w = store.get(&format!("{prefix}.weight"))?;
39    ensure!(
40        w.dtype == WeightDtype::FP8E4M3,
41        "Expected FP8E4M3 for {prefix}.weight, got {:?}",
42        w.dtype,
43    );
44    ensure!(
45        w.shape.len() == 2,
46        "Expected 2D weight for {prefix}, got {:?}",
47        w.shape
48    );
49    let n = w.shape[0];
50    let k = w.shape[1];
51    let weight_ptr = w.ptr;
52
53    // Load block scale [N/BS, K/BS] — already on GPU from safetensors. The
54    // tensor name varies by producer: DeepSeek/Qwen-native FP8 ships
55    // `weight_scale_inv` (2D); compressed-tensors `float-quantized` (e.g.
56    // Hcompany/Holo-3.1-*-FP8) ships a 2D `weight_scale`; DeepSeek-V4 ships
57    // a 2D F8_E8M0 `.scale`; ModelOpt
58    // MIXED_PRECISION ships a *scalar* `weight_scale` (expanded to the block
59    // matrix shape below). All three are the per-block FP8 dequant multiplier
60    // the W8A16 kernels apply in FP32. Prefer whichever 2D block scale exists.
61    let scale_inv_key = format!("{prefix}.weight_scale_inv");
62    let plain_scale_key = format!("{prefix}.weight_scale");
63    let e8m0_scale_key = format!("{prefix}.scale");
64    let block_scale_key = if store.contains(&scale_inv_key) {
65        Some(scale_inv_key.clone())
66    } else if store
67        .get(&plain_scale_key)
68        .map(|s| s.shape.len() == 2)
69        .unwrap_or(false)
70    {
71        Some(plain_scale_key.clone())
72    } else if store
73        .get(&e8m0_scale_key)
74        .map(|s| s.shape.len() == 2 && s.dtype == WeightDtype::FP8E8M0)
75        .unwrap_or(false)
76    {
77        Some(e8m0_scale_key.clone())
78    } else {
79        None
80    };
81    let row_scale = if let Some(scale_key) = block_scale_key {
82        let s = store.get(&scale_key)?;
83        ensure!(
84            s.shape.len() == 2,
85            "Expected 2D shape for {scale_key}, got {:?}",
86            s.shape,
87        );
88        ensure!(
89            matches!(
90                s.dtype,
91                WeightDtype::BF16 | WeightDtype::FP32 | WeightDtype::FP8E8M0
92            ),
93            "Expected BF16, FP32, or F8_E8M0 for {scale_key}, got {:?}",
94            s.dtype,
95        );
96
97        tracing::debug!(
98            "FP8 block scales: {prefix} [{n},{k}] scale=[{},{}] dtype={:?} -> FP32",
99            s.shape[0],
100            s.shape[1],
101            s.dtype,
102        );
103
104        // Widen the block scale to a genuine FP32 device buffer (lossless from
105        // BF16, straight copy from FP32). The W8A8/W8A16 kernels apply this scale
106        // in FP32; reading the checkpoint BF16 directly would clamp it to BF16
107        // precision (and an FP32-scale checkpoint would be misread as BF16).
108        let scale_total = s.shape[0] * s.shape[1];
109        let row_scale = gpu.alloc(scale_total * 4)?;
110        let kernel = gpu.kernel("widen_block_scale_f32", "widen_block_scale_f32")?;
111        let stream = gpu.default_stream();
112        let input_dtype = match s.dtype {
113            WeightDtype::BF16 => 0,
114            WeightDtype::FP32 => 1,
115            WeightDtype::FP8E8M0 => 2,
116            _ => unreachable!("validated block-scale dtype"),
117        };
118        crate::layers::ops::widen_block_scale_f32(
119            gpu,
120            kernel,
121            s.ptr,
122            row_scale,
123            scale_total as u32,
124            input_dtype,
125            stream,
126        )?;
127        gpu.synchronize(stream)?;
128        row_scale
129    } else {
130        let scalar_key = plain_scale_key;
131        let scale = scalar_f32(store, &scalar_key, gpu)
132            .with_context(|| format!("Missing {scale_inv_key} or scalar {scalar_key}"))?;
133        let n_blocks = n.div_ceil(128);
134        let k_blocks = k.div_ceil(128);
135        let scale_total = n_blocks * k_blocks;
136        tracing::debug!(
137            "FP8 scalar scale: {prefix} [{n},{k}] scale={scale:.8} -> [{n_blocks},{k_blocks}] FP32"
138        );
139        let mut scale_buf = Vec::with_capacity(scale_total * 4);
140        for _ in 0..scale_total {
141            scale_buf.extend_from_slice(&scale.to_le_bytes());
142        }
143        let ptr = gpu.alloc(scale_buf.len())?;
144        gpu.copy_h2d(&scale_buf, ptr)?;
145        ptr
146    };
147
148    Ok(Fp8Weight {
149        weight: weight_ptr,
150        row_scale, // FP32 [N/BS, K/BS] block scales on GPU
151        n: n as u32,
152        k: k as u32,
153        scale_format: WeightQuantFormat::Fp8BlockScaled,
154    })
155}
156
157/// Quantize a BF16 dense weight to NVFP4 on GPU.
158///
159/// Two-phase: (1) find global max, (2) per-group E2M1 quantization.
160/// Halves weight bandwidth vs FP8 (0.5 bytes/weight + group scales vs 1 byte/weight).
161/// Called once at model load time (not on the hot path).
162pub(crate) fn quantize_to_nvfp4(
163    bf16_weight: &DenseWeight,
164    n: usize,
165    k: usize,
166    gpu: &dyn GpuBackend,
167    absmax_kernel: spark_runtime::gpu::KernelHandle,
168    quantize_kernel: spark_runtime::gpu::KernelHandle,
169    stream: u64,
170) -> Result<QuantizedWeight> {
171    use spark_runtime::kernel_args::KernelLaunch;
172    use std::sync::atomic::{AtomicU64, Ordering};
173
174    static T_ALLOC_MAX: AtomicU64 = AtomicU64::new(0);
175    static T_LAUNCH1: AtomicU64 = AtomicU64::new(0);
176    static T_SYNC1: AtomicU64 = AtomicU64::new(0);
177    static T_D2H: AtomicU64 = AtomicU64::new(0);
178    static T_ALLOC_OUT: AtomicU64 = AtomicU64::new(0);
179    static T_LAUNCH2: AtomicU64 = AtomicU64::new(0);
180    static T_SYNC2: AtomicU64 = AtomicU64::new(0);
181    static N_CALLS: AtomicU64 = AtomicU64::new(0);
182
183    let total = n * k;
184
185    // Phase 1: Find global absolute max
186    let t = std::time::Instant::now();
187    let max_buf = gpu.alloc(4)?;
188    gpu.memset(max_buf, 0, 4)?;
189    T_ALLOC_MAX.fetch_add(t.elapsed().as_nanos() as u64, Ordering::Relaxed);
190
191    let t = std::time::Instant::now();
192    let grid1 = (total / 256).clamp(1, 1024) as u32;
193    KernelLaunch::new(gpu, absmax_kernel)
194        .grid([grid1, 1, 1])
195        .block([256, 1, 1])
196        .arg_ptr(bf16_weight.weight)
197        .arg_ptr(max_buf)
198        .arg_u32(total as u32)
199        .launch(stream)?;
200    T_LAUNCH1.fetch_add(t.elapsed().as_nanos() as u64, Ordering::Relaxed);
201
202    let t = std::time::Instant::now();
203    gpu.synchronize(stream)?;
204    T_SYNC1.fetch_add(t.elapsed().as_nanos() as u64, Ordering::Relaxed);
205    let t = std::time::Instant::now();
206    let mut max_bytes = [0u8; 4];
207    gpu.copy_d2h(max_buf, &mut max_bytes)?;
208    T_D2H.fetch_add(t.elapsed().as_nanos() as u64, Ordering::Relaxed);
209    let global_max = f32::from_le_bytes(max_bytes);
210
211    // scale2 = global_max / (6.0 * 448.0)  [FP8 E4M3 max = 448]
212    let scale2 = if global_max > 0.0 {
213        global_max / (6.0 * 448.0)
214    } else {
215        1.0
216    };
217
218    // Diagnostic: the first few quantizations report their absmax. Counted on
219    // the BACKEND, so a second model loaded into the process reports its own
220    // instead of inheriting a spent counter.
221    if gpu.op_cache().first_n("diag:quantize_nvfp4_absmax", 5) {
222        tracing::info!(
223            "quantize_to_nvfp4: n={n} k={k} total={total} global_max={global_max:.6} scale2={scale2:.8} grid1={grid1}",
224        );
225    }
226
227    // Phase 2: Quantize
228    let t = std::time::Instant::now();
229    let packed_buf = gpu.alloc(n * k / 2)?;
230    let scale_buf = gpu.alloc(n * k / 16)?;
231    T_ALLOC_OUT.fetch_add(t.elapsed().as_nanos() as u64, Ordering::Relaxed);
232
233    let t = std::time::Instant::now();
234    KernelLaunch::new(gpu, quantize_kernel)
235        .grid([n as u32, 1, 1])
236        .block([256, 1, 1])
237        .arg_ptr(bf16_weight.weight)
238        .arg_ptr(packed_buf)
239        .arg_ptr(scale_buf)
240        .arg_f32(scale2)
241        .arg_u32(n as u32)
242        .arg_u32(k as u32)
243        .launch(stream)?;
244    T_LAUNCH2.fetch_add(t.elapsed().as_nanos() as u64, Ordering::Relaxed);
245
246    let t = std::time::Instant::now();
247    gpu.synchronize(stream)?;
248    T_SYNC2.fetch_add(t.elapsed().as_nanos() as u64, Ordering::Relaxed);
249
250    let c = N_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
251    if c.is_multiple_of(512) {
252        let ms = |a: &AtomicU64| a.load(Ordering::Relaxed) as f64 / 1.0e6;
253        tracing::info!(
254            "quantize_to_nvfp4 PROFILE after {c} calls (ms total): alloc_max={:.1} launch1={:.1} \
255             sync1={:.1} d2h={:.1} alloc_out={:.1} launch2={:.1} sync2={:.1} | sum={:.1} \
256             per_call={:.3}ms",
257            ms(&T_ALLOC_MAX),
258            ms(&T_LAUNCH1),
259            ms(&T_SYNC1),
260            ms(&T_D2H),
261            ms(&T_ALLOC_OUT),
262            ms(&T_LAUNCH2),
263            ms(&T_SYNC2),
264            ms(&T_ALLOC_MAX)
265                + ms(&T_LAUNCH1)
266                + ms(&T_SYNC1)
267                + ms(&T_D2H)
268                + ms(&T_ALLOC_OUT)
269                + ms(&T_LAUNCH2)
270                + ms(&T_SYNC2),
271            (ms(&T_ALLOC_MAX)
272                + ms(&T_LAUNCH1)
273                + ms(&T_SYNC1)
274                + ms(&T_D2H)
275                + ms(&T_ALLOC_OUT)
276                + ms(&T_LAUNCH2)
277                + ms(&T_SYNC2))
278                / c as f64,
279        );
280    }
281
282    Ok(QuantizedWeight {
283        weight: packed_buf,
284        weight_scale: scale_buf,
285        weight_scale_2: scale2,
286        input_scale: DevicePtr::NULL,
287        weight_scale_2_vec: DevicePtr::NULL,
288    })
289}
290
291/// Load attention weights for a full_attention layer.
292pub(crate) fn load_attention(
293    store: &WeightStore,
294    layer_prefix: &str,
295    gpu: &dyn GpuBackend,
296    variant: Nvfp4Variant,
297    qctx: QuantizeCtx,
298    config: &atlas_core::config::ModelConfig,
299) -> Result<AttentionWeights> {
300    let p = format!("{layer_prefix}.self_attn");
301    let (k_scale, v_scale) = load_kv_scales(store, &p, gpu);
302    let h = config.hidden_size;
303    let qkv_out = config.num_attention_heads * config.head_dim;
304    // q/k/v may ship either as dense BF16/FP8 (`.weight`, kept in the quant
305    // ignore list) OR as compressed-tensors NVFP4 (`.weight_packed`) — e.g.
306    // RedHatAI/Qwen3-Coder-Next-NVFP4, which quantizes the attention
307    // projections too. This attention path consumes q/k/v as dense BF16
308    // (`AttentionWeights.{q,k,v}_proj: DenseWeight`), so dequant the NVFP4
309    // case to BF16 at load, mirroring the gemma4 loader; the dense case is
310    // untouched. Dims come from the packed tensor itself, not config: under
311    // `attn_output_gate` q_proj's row count is 2×(heads·head_dim), so a
312    // config-derived `n` would run the dequant kernel off the end. weight_packed
313    // is [out_features, in_features/2] (2 fp4 nibbles per byte). Without this,
314    // packed-q/k/v checkpoints die on `self_attn.q_proj.weight not found` at the
315    // first full_attention layer (issue #299 follow-on — the reported
316    // shared_expert half already loads).
317    let load_qkv = |name: &str| -> Result<DenseWeight> {
318        match store.get(&format!("{p}.{name}.weight_packed")) {
319            Ok(w) => crate::weight_map::dequant_nvfp4_to_bf16(
320                store,
321                &format!("{p}.{name}"),
322                w.shape[0],
323                w.shape[1] * 2,
324                gpu,
325            ),
326            Err(_) => dense_auto(store, &format!("{p}.{name}.weight"), gpu),
327        }
328    };
329    Ok(AttentionWeights {
330        q_proj: load_qkv("q_proj")?,
331        k_proj: load_qkv("k_proj")?,
332        v_proj: load_qkv("v_proj")?,
333        o_proj: quantized_any(
334            store,
335            &format!("{p}.o_proj"),
336            h,
337            qkv_out,
338            gpu,
339            variant,
340            qctx,
341        )?,
342        q_norm: dense(store, &format!("{p}.q_norm.weight"))?,
343        k_norm: dense(store, &format!("{p}.k_norm.weight"))?,
344        q_norm_full: None,
345        k_norm_full: None,
346        k_scale,
347        v_scale,
348    })
349}
350
351/// Load SSM weights for a linear_attention layer.
352pub(crate) fn load_ssm(
353    store: &WeightStore,
354    layer_prefix: &str,
355    gpu: &dyn GpuBackend,
356    variant: Nvfp4Variant,
357    qctx: QuantizeCtx,
358    config: &atlas_core::config::ModelConfig,
359) -> Result<SsmWeights> {
360    let p = format!("{layer_prefix}.linear_attn");
361    let h = config.hidden_size;
362    // out_proj is [hidden_size, d_inner] where d_inner = linear_value_head_dim * linear_num_value_heads
363    let d_inner = config.linear_value_head_dim * config.linear_num_value_heads;
364    Ok(SsmWeights {
365        in_proj_qkvz: dense_auto(store, &format!("{p}.in_proj_qkvz.weight"), gpu)?,
366        in_proj_ba: dense_auto(store, &format!("{p}.in_proj_ba.weight"), gpu)?,
367        conv1d: dense(store, &format!("{p}.conv1d.weight"))?,
368        a_log: dense_keep_f32(store, &format!("{p}.A_log"), gpu)?,
369        dt_bias: dense_keep_f32(store, &format!("{p}.dt_bias"), gpu)?,
370        norm: dense(store, &format!("{p}.norm.weight"))?,
371        out_proj: quantized_any(
372            store,
373            &format!("{p}.out_proj"),
374            h,
375            d_inner,
376            gpu,
377            variant,
378            qctx,
379        )?,
380    })
381}
382
383/// Load MoE weights for a layer.
384///
385/// Under EP (ep_world_size > 1), only local experts are loaded from the store.
386/// Remote experts get NULL pointers — kernels detect NULL and write zero output.
387pub(crate) fn load_moe(
388    store: &WeightStore,
389    layer_prefix: &str,
390    num_experts: usize,
391    gpu: &dyn GpuBackend,
392    config: &atlas_core::config::ModelConfig,
393    variant: Nvfp4Variant,
394    qctx: QuantizeCtx,
395) -> Result<MoeWeights> {
396    load_moe_inner(
397        store,
398        layer_prefix,
399        num_experts,
400        gpu,
401        config,
402        variant,
403        qctx,
404        false,
405    )
406}
407
408/// Load MoE with option to skip routed experts (native FP8 loads them separately).
409pub(crate) fn load_moe_skip_experts(
410    store: &WeightStore,
411    layer_prefix: &str,
412    num_experts: usize,
413    gpu: &dyn GpuBackend,
414    config: &atlas_core::config::ModelConfig,
415    variant: Nvfp4Variant,
416    qctx: QuantizeCtx,
417) -> Result<MoeWeights> {
418    load_moe_inner(
419        store,
420        layer_prefix,
421        num_experts,
422        gpu,
423        config,
424        variant,
425        qctx,
426        true,
427    )
428}
spark_model/weight_map/loaders_fp8.rs

spark_model/weight_map/
loaders_fp8.rs
