metrale_model_layers/layers/
dense_ffn_w8a8_prefill.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The W8A8 block-scaled arm of the native-FP8 dense FFN: which projections take it, and the GEMM it runs.
4//!
5//! Owner: model-layers (dense FFN).
6//! Invariants:
7//! - `w8a8_prefill_selected` holds only for `4 < m <= max_m`, `Fp8BlockScaled`
8//!   weights, `k` and `n` multiples of 128, the quantizer and GEMM handles
9//!   loaded, `METRALE_FFN_W8A16_ONLY` unset and
10//!   `dispatch.fp8_blockscaled_prefill` on.
11//!
12//! The arm is reachable only on a layer that holds native-FP8 dense-FFN
13//! weights. The one loader that installs them, `qwen35_dense.rs`, does so only
14//! when `ffn_fp8_arm_selected` holds: `METRALE_DENSE_FP8` is exactly `1`,
15//! `config.tp_world_size.max(1) == 1`, the variant is
16//! `Nvfp4Variant::Fp8Dequanted`, and the layer's `mlp.gate_proj` is block-scaled
17//! FP8 on disk (`proj_is_native_fp8`). Without it neither route log in this file
18//! is emitted.
19//!
20//! The activation is quantized to E4M3 with one FP32 scale per token per
21//! 128-wide K group (`per_token_group_quant_fp8`). `w8a8_gemm` then runs
22//! cuBLASLt `fp8_gemm_act_weight_t_blkscaled` when `dispatch.cublas.ffn` holds
23//! (`METRALE_CUBLAS_GEMM` lists `ffn`, or is `1`, `true` or `all`) and its other
24//! clauses pass, and the in-tree `ops::fp8_gemm_t_blockscaled` otherwise.
25//! `examples/native_fp8_ffn_w8a8_microtest.rs` grades W8A8 against the W8A16
26//! reference at cosine >= 0.999 and relative RMS <= 2%.
27//!
28//! The upper M bound is per target: `kernels/<hw>/HARDWARE.toml` `[defaults]
29//! w8a8_prefill_max_m_widening` / `_narrowing`. gb10 declares 64 and 384 beside
30//! its measurement; every other target has no cap (`u32::MAX`).
31
32use anyhow::Result;
33use metrale_gpu_runtime::gpu::{DevicePtr, KernelHandle};
34
35use super::DenseFfnLayer;
36use crate::layer::ForwardContext;
37use crate::layers::ops;
38use crate::weight_map::{Fp8Weight, WeightQuantFormat};
39
40/// 2026-09-25: `METRALE_FFN_W8A16_ONLY`, by presence (any value, including
41/// empty or `0`): no dense-FFN projection takes the W8A8 arm. Read once per
42/// process (`OnceLock`).
43pub fn ffn_w8a16_only() -> bool {
44    static ONLY: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
45    *ONLY.get_or_init(|| std::env::var_os("METRALE_FFN_W8A16_ONLY").is_some())
46}
47
48/// 2026-09-25: The W8A8 upper M bound for a projection of shape `[n, k]`:
49/// `w8a8_prefill_max_m_widening` when `n > k` (gate/up) and
50/// `w8a8_prefill_max_m_narrowing` otherwise (down), from
51/// [`ops::target_defaults::resolved`].
52pub(crate) fn max_m_for(n: u32, k: u32) -> u32 {
53    let levers = ops::target_defaults::resolved();
54    if n > k {
55        levers.w8a8_prefill_max_m_widening.value
56    } else {
57        levers.w8a8_prefill_max_m_narrowing.value
58    }
59}
60
61/// 2026-09-25: The W8A8 selection rule, as a pure function, so the CPU tests
62/// can pin every clause without a `ForwardContext`. `w8a16_only` and `max_m`
63/// are arguments because their production readers are process-global.
64///
65/// Clauses:
66///
67/// * `m > 4`: `m <= 4` belongs to the `w8a16_gemv_batch4` rung.
68/// * `fp8_blockscaled_prefill`: cleared by `METRALE_FP8_SINGLE_SCALE=1`.
69/// * `Fp8BlockScaled`: the GEMM indexes `row_scale` as the `[N/128, K/128]`
70///   grid, which per-row and single-scale weights do not have.
71/// * `k % 128 == 0`: the activation quantizer emits one scale per 128-wide K
72///   group.
73/// * `n % 128 == 0`: the weight scale grid is `[N/128, K/128]`.
74/// * `m <= max_m`: the per-target upper bound. `u32::MAX` is no cap and 0
75///   turns the arm off.
76/// * the quantizer and the GEMM handles are loaded.
77#[allow(clippy::too_many_arguments)]
78pub(crate) fn w8a8_prefill_selected(
79    m: u32,
80    n: u32,
81    k: u32,
82    scale_format: WeightQuantFormat,
83    fp8_blockscaled_prefill: bool,
84    quant_k: ops::Fp8ActQuant,
85    gemm_k: KernelHandle,
86    w8a16_only: bool,
87    max_m: u32,
88) -> bool {
89    !w8a16_only
90        && m > 4
91        && m <= max_m
92        && fp8_blockscaled_prefill
93        && scale_format == WeightQuantFormat::Fp8BlockScaled
94        && k.is_multiple_of(128)
95        && n.is_multiple_of(128)
96        && quant_k.available()
97        && gemm_k.0 != 0
98}
99
100impl DenseFfnLayer {
101    /// 2026-09-25: Whether one dense-FFN projection takes the W8A8 arm.
102    ///
103    /// Beyond [`w8a8_prefill_selected`] this requires the arena's dense-FFN
104    /// activation scratch (`ffn_act_a`, `ffn_act_scale`), which is 0 for MoE
105    /// configs. It is a gate and not an assert: a null scratch would be a launch
106    /// writing to address 0.
107    pub(crate) fn prefill_w8a8_selected(
108        &self,
109        ctx: &ForwardContext,
110        m: u32,
111        n: u32,
112        k: u32,
113        w: &Fp8Weight,
114    ) -> bool {
115        w8a8_prefill_selected(
116            m,
117            n,
118            k,
119            w.scale_format,
120            ctx.dispatch.fp8_blockscaled_prefill,
121            self.per_token_group_quant_fp8_k,
122            self.fp8_gemm_t_blockscaled_k,
123            ffn_w8a16_only(),
124            max_m_for(n, k),
125        ) && ctx.buffers.ffn_act_a().0 != 0
126            && ctx.buffers.ffn_act_scale().0 != 0
127    }
128
129    /// 2026-09-26: Quantize `act[m, k]` BF16 into the arena's dense-FFN scratch as
130    /// FP8 E4M3 plus one FP32 scale per token per 128-wide K group; returns
131    /// `(a_fp8, a_scale)`.
132    ///
133    /// The arena has one scratch pair, so the caller consumes the result before
134    /// the next call. `prefill_fp8` (reached from `forward_prefill_inner`) issues the gate/up GEMMs, the
135    /// SiLU-mul and the down quantization in order on one stream. `k` is `hidden`
136    /// or `intermediate`; the scratch is sized for the larger.
137    pub(crate) fn w8a8_quant_act(
138        &self,
139        ctx: &ForwardContext,
140        act: DevicePtr,
141        m: u32,
142        k: u32,
143        stream: u64,
144    ) -> Result<(DevicePtr, DevicePtr)> {
145        let a_fp8 = ctx.buffers.ffn_act_a();
146        let a_scale = ctx.buffers.ffn_act_scale();
147        // 2026-09-25: Padded extents: the cuBLASLt arm reads `ceil16(m)` rows.
148        let rows = ops::cublas_fp8_m_pad(m) as usize;
149        debug_assert!(rows * k as usize <= ctx.buffers.ffn_act_a_bytes());
150        debug_assert!(rows * (k as usize / 128) * 4 <= ctx.buffers.ffn_act_scale_bytes());
151        ops::per_token_group_quant_fp8(
152            ctx.gpu,
153            self.per_token_group_quant_fp8_k,
154            act,
155            a_fp8,
156            a_scale,
157            m,
158            k,
159            stream,
160        )?;
161        Ok((a_fp8, a_scale))
162    }
163
164    /// 2026-09-25: `out[m, n] = a_fp8[m, k] @ weight[n, k]^T` with both scale sets
165    /// applied in FP32: cuBLASLt when `dispatch.cublas.ffn` holds and the clauses
166    /// below pass, else the in-tree kernel.
167    ///
168    /// `out_capacity_bytes` is the allocated size of `out`. The cuBLASLt helper
169    /// rounds M up to 16 and writes those rows too, so an `out` smaller than the
170    /// padded extent sends the GEMM to the in-tree kernel instead.
171    #[allow(clippy::too_many_arguments)]
172    pub(crate) fn w8a8_gemm(
173        &self,
174        ctx: &ForwardContext,
175        a_fp8: DevicePtr,
176        a_scale: DevicePtr,
177        w: &Fp8Weight,
178        out: DevicePtr,
179        out_capacity_bytes: usize,
180        m: u32,
181        n: u32,
182        k: u32,
183        stream: u64,
184    ) -> Result<()> {
185        let m_pad = ops::cublas_fp8_m_pad(m) as usize;
186        let padded_out_bytes = m_pad * n as usize * 2;
187        // 2026-09-25: The cuBLASLt arm also needs:
188        //
189        // * output room for the padded rows;
190        // * when the k-major scale layout is in use (the default;
191        //   `METRALE_CUBLAS_SCALE_LAYOUT=rowmajor` turns it off), the
192        //   `fp8_act_scale_to_kmajor` kernel and its scratch. cuBLASLt reads the
193        //   activation scales token-contiguous. Measured on H100 2026-09-11 without
194        //   the transpose: rel_rms 7.7e-2 against the in-tree kernel;
195        // * `k / 128` a multiple of 4 (`blk128x128_stride_ok`): the weight scales
196        //   are handed over as the checkpoint's `[N/128, K/128]` grid.
197        let scale_layout_ready = ctx.buffers.ffn_act_scale_kmajor().0 != 0
198            && self.fp8_act_scale_kmajor_k.0 != 0
199            && ctx.buffers.ffn_act_scale_kmajor_bytes() >= m_pad * (k as usize / 128) * 4;
200        let cublas = ctx.dispatch.cublas.ffn
201            && padded_out_bytes <= out_capacity_bytes
202            && metrale_gpu_runtime::cublaslt::scale_layout::blk128x128_stride_ok(k as usize)
203            && (scale_layout_ready || !ops::cublas_scale_layout_kmajor());
204        self.log_w8a8_prefill_route(ctx, cublas);
205        if cublas {
206            return ops::cublas_fp8_proj_prequant(
207                ctx.gpu,
208                self.fp8_act_scale_kmajor_k,
209                a_fp8,
210                a_scale,
211                ctx.buffers.ffn_act_scale_kmajor(),
212                w,
213                out,
214                m,
215                n,
216                k,
217                stream,
218            );
219        }
220        ops::fp8_gemm_t_blockscaled(
221            ctx.gpu,
222            self.fp8_gemm_t_blockscaled_k,
223            a_fp8,
224            a_scale,
225            w.weight,
226            w.row_scale,
227            out,
228            m,
229            n,
230            k,
231            stream,
232        )
233    }
234
235    /// 2026-09-25: Logs once per model (`log:ffn_w8a8_prefill`) which GEMM the
236    /// W8A8 arm used.
237    fn log_w8a8_prefill_route(&self, ctx: &ForwardContext, cublas: bool) {
238        if ctx.stats.once("log:ffn_w8a8_prefill") {
239            let how = if cublas { "cuBLASLt" } else { "kernel" };
240            tracing::info!(
241                "[metrale] dense FFN prefill: W8A8 block-scaled via {how} \
242                 (per-token 1x128 act scales x 128x128 weight scales, FP32 epilogue; \
243                 vLLM-equivalent FP8 numerics). METRALE_FFN_W8A16_ONLY=1 restores W8A16."
244            );
245        }
246    }
247
248    /// 2026-09-25: Logs once per model (`log:ffn_w8a16_prefill`) when a prefill
249    /// takes the W8A8 arm for neither gate/up nor down while the W8A8 arm is
250    /// reachable (neither the batch16 nor a tensor-core tier claims `m`).
251    pub(crate) fn log_w8a16_prefill_route(&self, ctx: &ForwardContext) {
252        if ctx.stats.once("log:ffn_w8a16_prefill") {
253            tracing::info!(
254                "[metrale] dense FFN prefill: W8A16 (BF16 act x FP8 weight). \
255                 W8A8 not selected — see #917/#928."
256            );
257        }
258    }
259}
260
261#[cfg(test)]
262#[path = "dense_ffn_w8a8_prefill_tests.rs"]
263mod tests;
metrale_model_layers/layers/dense_ffn_w8a8_prefill.rs

metrale_model_layers/layers/
dense_ffn_w8a8_prefill.rs
