metrale_model_layers/weight_map/
loaders_fp8.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: FP8 block-scaled weight loading, BF16-to-NVFP4 quantization, and the attention, SSM and MoE layer loaders.
4//!
5//! Owner: model-layers (weight loading).
6//! Invariants:
7//! - Every `Fp8Weight` returned here holds its block scale as an FP32 device buffer.
8
9#![allow(unused_imports)]
10
11use anyhow::{Context, Result, bail, ensure};
12use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
13use metrale_model_weights::weights::{WeightDtype, WeightStore};
14
15use super::*;
16
17/// 2026-09-25: Load an FP8 E4M3 block-scaled checkpoint weight as a native [`Fp8Weight`].
18///
19/// `{prefix}.weight` must be a 2D FP8E4M3 tensor `[N, K]`. The block scale is
20/// read from the first of these that exists: `{prefix}.weight_scale_inv`, a 2D
21/// `{prefix}.weight_scale`, or a 2D F8_E8M0 `{prefix}.scale`. The scale must be
22/// BF16, FP32 or F8_E8M0, and it is widened once into an FP32 device buffer.
23/// When none exists, a scalar `{prefix}.weight_scale` is repeated over a
24/// `[ceil(N/128), ceil(K/128)]` FP32 grid. `w8a16_gemv` reads this buffer as
25/// `const float*` and applies `block_scale[n/128, k/128]` to each weight.
26pub fn load_fp8_block_scaled_as_fp8weight(
27    store: &WeightStore,
28    prefix: &str,
29    gpu: &dyn GpuBackend,
30) -> Result<Fp8Weight> {
31    let w = store.get(&format!("{prefix}.weight"))?;
32    ensure!(
33        w.dtype == WeightDtype::FP8E4M3,
34        "Expected FP8E4M3 for {prefix}.weight, got {:?}",
35        w.dtype,
36    );
37    ensure!(
38        w.shape.len() == 2,
39        "Expected 2D weight for {prefix}, got {:?}",
40        w.shape
41    );
42    let n = w.shape[0];
43    let k = w.shape[1];
44    let weight_ptr = w.ptr;
45
46    let scale_inv_key = format!("{prefix}.weight_scale_inv");
47    let plain_scale_key = format!("{prefix}.weight_scale");
48    let e8m0_scale_key = format!("{prefix}.scale");
49    let block_scale_key = if store.contains(&scale_inv_key) {
50        Some(scale_inv_key.clone())
51    } else if store
52        .get(&plain_scale_key)
53        .map(|s| s.shape.len() == 2)
54        .unwrap_or(false)
55    {
56        Some(plain_scale_key.clone())
57    } else if store
58        .get(&e8m0_scale_key)
59        .map(|s| s.shape.len() == 2 && s.dtype == WeightDtype::FP8E8M0)
60        .unwrap_or(false)
61    {
62        Some(e8m0_scale_key.clone())
63    } else {
64        None
65    };
66    let row_scale = if let Some(scale_key) = block_scale_key {
67        let s = store.get(&scale_key)?;
68        ensure!(
69            s.shape.len() == 2,
70            "Expected 2D shape for {scale_key}, got {:?}",
71            s.shape,
72        );
73        ensure!(
74            matches!(
75                s.dtype,
76                WeightDtype::BF16 | WeightDtype::FP32 | WeightDtype::FP8E8M0
77            ),
78            "Expected BF16, FP32, or F8_E8M0 for {scale_key}, got {:?}",
79            s.dtype,
80        );
81
82        tracing::debug!(
83            "FP8 block scales: {prefix} [{n},{k}] scale=[{},{}] dtype={:?} -> FP32",
84            s.shape[0],
85            s.shape[1],
86            s.dtype,
87        );
88
89        // 2026-09-25: Widen on the device so the kernels read one FP32 layout
90        // whatever dtype the checkpoint stored the scale in.
91        let scale_total = s.shape[0] * s.shape[1];
92        let row_scale = gpu.alloc(scale_total * 4)?;
93        let kernel = gpu.kernel("widen_block_scale_f32", "widen_block_scale_f32")?;
94        let stream = gpu.default_stream();
95        let input_dtype = match s.dtype {
96            WeightDtype::BF16 => 0,
97            WeightDtype::FP32 => 1,
98            WeightDtype::FP8E8M0 => 2,
99            _ => unreachable!("validated block-scale dtype"),
100        };
101        crate::layers::ops::widen_block_scale_f32(
102            gpu,
103            kernel,
104            s.ptr,
105            row_scale,
106            scale_total as u32,
107            input_dtype,
108            stream,
109        )?;
110        gpu.synchronize(stream)?;
111        row_scale
112    } else {
113        let scalar_key = plain_scale_key;
114        let scale = scalar_f32(store, &scalar_key, gpu)
115            .with_context(|| format!("Missing {scale_inv_key} or scalar {scalar_key}"))?;
116        let n_blocks = n.div_ceil(128);
117        let k_blocks = k.div_ceil(128);
118        let scale_total = n_blocks * k_blocks;
119        tracing::debug!(
120            "FP8 scalar scale: {prefix} [{n},{k}] scale={scale:.8} -> [{n_blocks},{k_blocks}] FP32"
121        );
122        let mut scale_buf = Vec::with_capacity(scale_total * 4);
123        for _ in 0..scale_total {
124            scale_buf.extend_from_slice(&scale.to_le_bytes());
125        }
126        let ptr = gpu.alloc(scale_buf.len())?;
127        gpu.copy_h2d(&scale_buf, ptr)?;
128        ptr
129    };
130
131    Ok(Fp8Weight {
132        weight: weight_ptr,
133        row_scale,
134        n: n as u32,
135        k: k as u32,
136        scale_format: WeightQuantFormat::Fp8BlockScaled,
137    })
138}
139
140/// 2026-09-25: Quantize a BF16 dense `[n, k]` weight to NVFP4 on the GPU.
141///
142/// Two kernels: `absmax_kernel` finds the global absolute max, then
143/// `quantize_kernel` writes packed E2M1 (`n*k/2` bytes) and one scale byte per
144/// 16 weights. The global scale is `global_max / (6 * 448)`, or 1.0 for an
145/// all-zero weight. It synchronizes `stream` after each kernel, so it belongs
146/// at load time, not in a forward pass.
147pub fn quantize_to_nvfp4(
148    bf16_weight: &DenseWeight,
149    n: usize,
150    k: usize,
151    gpu: &dyn GpuBackend,
152    absmax_kernel: metrale_gpu_runtime::gpu::KernelHandle,
153    quantize_kernel: metrale_gpu_runtime::gpu::KernelHandle,
154    stream: u64,
155) -> Result<QuantizedWeight> {
156    use metrale_gpu_runtime::kernel_args::KernelLaunch;
157    use std::sync::atomic::{AtomicU64, Ordering};
158
159    static T_ALLOC_MAX: AtomicU64 = AtomicU64::new(0);
160    static T_LAUNCH1: AtomicU64 = AtomicU64::new(0);
161    static T_SYNC1: AtomicU64 = AtomicU64::new(0);
162    static T_D2H: AtomicU64 = AtomicU64::new(0);
163    static T_ALLOC_OUT: AtomicU64 = AtomicU64::new(0);
164    static T_LAUNCH2: AtomicU64 = AtomicU64::new(0);
165    static T_SYNC2: AtomicU64 = AtomicU64::new(0);
166    static N_CALLS: AtomicU64 = AtomicU64::new(0);
167
168    let total = n * k;
169
170    let t = std::time::Instant::now();
171    let max_buf = gpu.alloc(4)?;
172    gpu.memset(max_buf, 0, 4)?;
173    T_ALLOC_MAX.fetch_add(t.elapsed().as_nanos() as u64, Ordering::Relaxed);
174
175    let t = std::time::Instant::now();
176    let grid1 = (total / 256).clamp(1, 1024) as u32;
177    KernelLaunch::new(gpu, absmax_kernel)
178        .grid([grid1, 1, 1])
179        .block([256, 1, 1])
180        .arg_ptr(bf16_weight.weight)
181        .arg_ptr(max_buf)
182        .arg_u32(total as u32)
183        .launch(stream)?;
184    T_LAUNCH1.fetch_add(t.elapsed().as_nanos() as u64, Ordering::Relaxed);
185
186    let t = std::time::Instant::now();
187    gpu.synchronize(stream)?;
188    T_SYNC1.fetch_add(t.elapsed().as_nanos() as u64, Ordering::Relaxed);
189    let t = std::time::Instant::now();
190    let mut max_bytes = [0u8; 4];
191    gpu.copy_d2h(max_buf, &mut max_bytes)?;
192    T_D2H.fetch_add(t.elapsed().as_nanos() as u64, Ordering::Relaxed);
193    let global_max = f32::from_le_bytes(max_bytes);
194
195    // 2026-09-25: 6.0 is the E2M1 maximum and 448.0 the E4M3 maximum.
196    let scale2 = if global_max > 0.0 {
197        global_max / (6.0 * 448.0)
198    } else {
199        1.0
200    };
201
202    // 2026-09-25: The first 5 calls on this backend log their absmax
203    // (`OpCache::first_n` counts per backend).
204    if gpu.op_cache().first_n("diag:quantize_nvfp4_absmax", 5) {
205        tracing::info!(
206            "quantize_to_nvfp4: n={n} k={k} total={total} global_max={global_max:.6} scale2={scale2:.8} grid1={grid1}",
207        );
208    }
209
210    let t = std::time::Instant::now();
211    let packed_buf = gpu.alloc(n * k / 2)?;
212    let scale_buf = gpu.alloc(n * k / 16)?;
213    T_ALLOC_OUT.fetch_add(t.elapsed().as_nanos() as u64, Ordering::Relaxed);
214
215    let t = std::time::Instant::now();
216    KernelLaunch::new(gpu, quantize_kernel)
217        .grid([n as u32, 1, 1])
218        .block([256, 1, 1])
219        .arg_ptr(bf16_weight.weight)
220        .arg_ptr(packed_buf)
221        .arg_ptr(scale_buf)
222        .arg_f32(scale2)
223        .arg_u32(n as u32)
224        .arg_u32(k as u32)
225        .launch(stream)?;
226    T_LAUNCH2.fetch_add(t.elapsed().as_nanos() as u64, Ordering::Relaxed);
227
228    let t = std::time::Instant::now();
229    gpu.synchronize(stream)?;
230    T_SYNC2.fetch_add(t.elapsed().as_nanos() as u64, Ordering::Relaxed);
231
232    let c = N_CALLS.fetch_add(1, Ordering::Relaxed) + 1;
233    if c.is_multiple_of(512) {
234        let ms = |a: &AtomicU64| a.load(Ordering::Relaxed) as f64 / 1.0e6;
235        tracing::info!(
236            "quantize_to_nvfp4 PROFILE after {c} calls (ms total): alloc_max={:.1} launch1={:.1} \
237             sync1={:.1} d2h={:.1} alloc_out={:.1} launch2={:.1} sync2={:.1} | sum={:.1} \
238             per_call={:.3}ms",
239            ms(&T_ALLOC_MAX),
240            ms(&T_LAUNCH1),
241            ms(&T_SYNC1),
242            ms(&T_D2H),
243            ms(&T_ALLOC_OUT),
244            ms(&T_LAUNCH2),
245            ms(&T_SYNC2),
246            ms(&T_ALLOC_MAX)
247                + ms(&T_LAUNCH1)
248                + ms(&T_SYNC1)
249                + ms(&T_D2H)
250                + ms(&T_ALLOC_OUT)
251                + ms(&T_LAUNCH2)
252                + ms(&T_SYNC2),
253            (ms(&T_ALLOC_MAX)
254                + ms(&T_LAUNCH1)
255                + ms(&T_SYNC1)
256                + ms(&T_D2H)
257                + ms(&T_ALLOC_OUT)
258                + ms(&T_LAUNCH2)
259                + ms(&T_SYNC2))
260                / c as f64,
261        );
262    }
263
264    Ok(QuantizedWeight {
265        weight: packed_buf,
266        weight_scale: scale_buf,
267        weight_scale_2: scale2,
268        input_scale: DevicePtr::NULL,
269        weight_scale_2_vec: DevicePtr::NULL,
270        act: metrale_config::Nvfp4Act::Unstamped,
271    })
272}
273
274/// 2026-09-25: Load the `{layer_prefix}.self_attn` weights of a full-attention layer.
275pub fn load_attention(
276    store: &WeightStore,
277    layer_prefix: &str,
278    gpu: &dyn GpuBackend,
279    variant: Nvfp4Variant,
280    qctx: QuantizeCtx,
281    config: &metrale_config::ModelConfig,
282) -> Result<AttentionWeights> {
283    let p = format!("{layer_prefix}.self_attn");
284    let (k_scale, v_scale) = load_kv_scales(store, &p, gpu)?;
285    let h = config.hidden_size;
286    let qkv_out = config.num_attention_heads * config.head_dim;
287    // 2026-09-25: q/k/v are `DenseWeight`, so an NVFP4 `.weight_packed`
288    // projection is dequantized to BF16 here; otherwise `.weight` goes through
289    // `dense_auto`. The dims come from the packed tensor `[out, in/2]`, not
290    // from config, because a gated q_proj carries the gate rows as well.
291    let load_qkv = |name: &str| -> Result<DenseWeight> {
292        match store.get(&format!("{p}.{name}.weight_packed")) {
293            Ok(w) => crate::weight_map::dequant_nvfp4_to_bf16(
294                store,
295                &format!("{p}.{name}"),
296                w.shape[0],
297                w.shape[1] * 2,
298                gpu,
299            ),
300            Err(_) => dense_auto(store, &format!("{p}.{name}.weight"), gpu),
301        }
302    };
303    Ok(AttentionWeights {
304        q_proj: load_qkv("q_proj")?,
305        k_proj: load_qkv("k_proj")?,
306        v_proj: load_qkv("v_proj")?,
307        o_proj: quantized_any(
308            store,
309            &format!("{p}.o_proj"),
310            h,
311            qkv_out,
312            gpu,
313            variant,
314            qctx,
315        )?,
316        q_norm: dense(store, &format!("{p}.q_norm.weight"))?,
317        k_norm: dense(store, &format!("{p}.k_norm.weight"))?,
318        q_norm_full: None,
319        k_norm_full: None,
320        k_scale,
321        v_scale,
322    })
323}
324
325/// 2026-09-25: Load the `{layer_prefix}.linear_attn` weights of a linear-attention layer.
326pub fn load_ssm(
327    store: &WeightStore,
328    layer_prefix: &str,
329    gpu: &dyn GpuBackend,
330    variant: Nvfp4Variant,
331    qctx: QuantizeCtx,
332    config: &metrale_config::ModelConfig,
333) -> Result<SsmWeights> {
334    let p = format!("{layer_prefix}.linear_attn");
335    let h = config.hidden_size;
336    let d_inner = config.linear_value_head_dim * config.linear_num_value_heads;
337    Ok(SsmWeights {
338        in_proj_qkvz: dense_auto(store, &format!("{p}.in_proj_qkvz.weight"), gpu)?,
339        in_proj_ba: dense_auto(store, &format!("{p}.in_proj_ba.weight"), gpu)?,
340        conv1d: dense(store, &format!("{p}.conv1d.weight"))?,
341        a_log: dense_keep_f32(store, &format!("{p}.A_log"), gpu)?,
342        dt_bias: dense_keep_f32(store, &format!("{p}.dt_bias"), gpu)?,
343        norm: dense(store, &format!("{p}.norm.weight"))?,
344        out_proj: quantized_any(
345            store,
346            &format!("{p}.out_proj"),
347            h,
348            d_inner,
349            gpu,
350            variant,
351            qctx,
352        )?,
353    })
354}
355
356/// 2026-09-25: Load the `{layer_prefix}.mlp` MoE weights of a layer.
357///
358/// Only experts for which `config.is_local_expert(e)` holds are loaded; the
359/// others are `ExpertWeight::null()`.
360pub fn load_moe(
361    store: &WeightStore,
362    layer_prefix: &str,
363    num_experts: usize,
364    gpu: &dyn GpuBackend,
365    config: &metrale_config::ModelConfig,
366    variant: Nvfp4Variant,
367    qctx: QuantizeCtx,
368) -> Result<MoeWeights> {
369    load_moe_inner(
370        store,
371        layer_prefix,
372        num_experts,
373        gpu,
374        config,
375        variant,
376        qctx,
377        false,
378    )
379}
380
381/// 2026-09-25: [`load_moe`] with every routed expert left as `ExpertWeight::null()`.
382/// The native-FP8 qwen3 loader uses it and loads the routed experts separately.
383pub fn load_moe_skip_experts(
384    store: &WeightStore,
385    layer_prefix: &str,
386    num_experts: usize,
387    gpu: &dyn GpuBackend,
388    config: &metrale_config::ModelConfig,
389    variant: Nvfp4Variant,
390    qctx: QuantizeCtx,
391) -> Result<MoeWeights> {
392    load_moe_inner(
393        store,
394        layer_prefix,
395        num_experts,
396        gpu,
397        config,
398        variant,
399        qctx,
400        true,
401    )
402}
metrale_model_layers/weight_map/loaders_fp8.rs

metrale_model_layers/weight_map/
loaders_fp8.rs
