metrale_model_layers/layers/qwen3_ssm/
init_fp8.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: FP8 weight-install setters and the NVFP4→FP8 `out_proj`
4//! pre-dequant for [`Qwen3SsmLayer`].
5//!
6//! Owner: model-layers (qwen3_ssm).
7//! Invariants:
8//! - `set_fp8_decode_weights` panics unless each weight is tagged
9//!   `Fp8BlockScaled`, and `set_fp8_rowwise_prefill_weights` unless each is
10//!   tagged `Fp8PerRow`, so neither FP8 layout reaches the other's readers.
11
12use anyhow::Result;
13use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
14
15use super::Qwen3SsmLayer;
16use crate::weight_map::Fp8Weight;
17
18impl Qwen3SsmLayer {
19    /// 2026-09-25: Install block-scaled FP8 weights as `qkvz_fp8w` /
20    /// `out_proj_fp8w`. Panics unless each is tagged `Fp8BlockScaled`: their
21    /// readers include `w8a16_gemv`, which computes
22    /// `sum_k A[k] * E4M3_LUT[B[n,k]] * block_scale[n/128, k/128]` with FP32
23    /// scales (`kernels/gb10/common/w8a16_gemv.cu`).
24    ///
25    /// It does not touch `qkvz_fp8` / `out_proj_fp8`: those feed
26    /// `fp8_gemm_n128`, which takes no scale argument, so block-scaled bytes
27    /// there would be read unscaled. `set_fp8_prefill_only_weights` installs
28    /// them.
29    pub fn set_fp8_decode_weights(&mut self, qkvz: Option<Fp8Weight>, out_proj: Option<Fp8Weight>) {
30        if let Some(ref w) = qkvz {
31            w.scale_format.expect(
32                crate::weight_map::WeightQuantFormat::Fp8BlockScaled,
33                "set_fp8_decode_weights::qkvz (w8a16_gemv expects [N/BS,K/BS] BF16 block scales)",
34            );
35        }
36        if let Some(ref w) = out_proj {
37            w.scale_format.expect(
38                crate::weight_map::WeightQuantFormat::Fp8BlockScaled,
39                "set_fp8_decode_weights::out_proj (w8a16_gemv expects [N/BS,K/BS] BF16 block scales)",
40            );
41        }
42        self.qkvz_fp8w = qkvz;
43        self.out_proj_fp8w = out_proj;
44    }
45
46    /// 2026-09-25: Install per-row FP8 weights as `qkvz_fp8w_rowwise` /
47    /// `out_proj_fp8w_rowwise`. Panics unless each is tagged `Fp8PerRow`. Only
48    /// the row-wise prefill arms read them (`trait_prefill_proj.rs`,
49    /// `trait_prefill_helper.rs`); decode never does.
50    pub fn set_fp8_rowwise_prefill_weights(
51        &mut self,
52        qkvz: Option<Fp8Weight>,
53        out_proj: Option<Fp8Weight>,
54    ) {
55        for (w, what) in [(&qkvz, "qkvz"), (&out_proj, "out_proj")] {
56            if let Some(w) = w {
57                w.scale_format.expect(
58                    crate::weight_map::WeightQuantFormat::Fp8PerRow,
59                    "set_fp8_rowwise_prefill_weights (cuBLASLt row-wise expects [N] f32)",
60                );
61                let _ = what;
62            }
63        }
64        self.qkvz_fp8w_rowwise = qkvz;
65        self.out_proj_fp8w_rowwise = out_proj;
66    }
67
68    /// 2026-09-25: Install unscaled FP8 weights as `qkvz_fp8` / `out_proj_fp8`,
69    /// the operands of the `fp8_gemm_n128` arms in prefill and in the batched
70    /// decode/verify projections (`trait_decode_batched.rs`). A `None`
71    /// argument leaves that field as it was.
72    pub fn set_fp8_prefill_only_weights(
73        &mut self,
74        qkvz_fp8: Option<DevicePtr>,
75        out_proj_fp8: Option<DevicePtr>,
76    ) {
77        if qkvz_fp8.is_some() {
78            self.qkvz_fp8 = qkvz_fp8;
79        }
80        if out_proj_fp8.is_some() {
81            self.out_proj_fp8 = out_proj_fp8;
82        }
83    }
84
85    /// 2026-09-25: When the layer holds `out_proj_nvfp4_t`, dequantize the NVFP4
86    /// `out_proj` to FP8 into a new allocation and store it as `out_proj_fp8`.
87    /// QKVZ is not converted. Synchronizes `stream`.
88    pub fn predequant_for_prefill(
89        &mut self,
90        gpu: &dyn GpuBackend,
91        config: &metrale_config::ModelConfig,
92        stream: u64,
93    ) -> Result<()> {
94        let predequant_k = gpu.kernel("w4a16", "predequant_nvfp4_to_fp8")?;
95        let h = config.hidden_size;
96        let qkvz_size = config.ssm_qkvz_size();
97        let value_dim = config.linear_num_value_heads * config.linear_value_head_dim;
98
99        let _ = qkvz_size;
100        // 2026-09-25: `predequant_nvfp4_to_fp8` reads a packed `[N, K/2]` weight,
101        // so this passes the untransposed `ssm.out_proj`.
102        if self.out_proj_nvfp4_t.is_some() {
103            self.out_proj_fp8 = Some(self.ssm.out_proj.predequant_to_fp8(
104                gpu,
105                predequant_k,
106                h,
107                value_dim,
108                stream,
109            )?);
110        }
111        Ok(())
112    }
113}
metrale_model_layers/layers/qwen3_ssm/init_fp8.rs

metrale_model_layers/layers/qwen3_ssm/
init_fp8.rs
