metrale_model_layers/layers/ops/
dispatch_proj_decode.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: W8A8 block-scaled cuBLASLt route for the 5..=16-row decode
4//! projections (GDN `in_proj_qkvz` and `out_proj`, attention Q/K/V and O), and
5//! the strided-output extent those need.
6//!
7//! On this route a projection runs W8A8 instead of the W8A16 GEMV: an E4M3
8//! activation with one FP32 scale per token and 128 of K, against the
9//! checkpoint's E4M3 weight and its 128x128 FP32 block scales.
10//!
11//! Owner: model-layers ops.
12//! Invariants:
13//! - [`decode_w8a8_selected`] reads the environment only through values cached
14//!   once per process ([`w8a8_decode_proj_disabled`],
15//!   `cublas_scale_layout_kmajor`), so for the same padded rows, shape, weight
16//!   format and buffers it gives the same answer on every step, CUDA-graph
17//!   capture included.
18
19use anyhow::Result;
20use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
21
22use super::{
23    Fp8ActQuant, cublas_fp8_m_pad, cublas_scale_layout_kmajor, fp8_act_scale_to_kmajor,
24    per_token_group_quant_fp8,
25};
26
27/// 2026-09-25: Kill switch: `METRALE_NO_W8A8_DECODE_PROJ` set to any value,
28/// empty and `0` included, turns this route off. Read once, so it is constant
29/// for the process and safe to branch on under CUDA-graph capture.
30pub fn w8a8_decode_proj_disabled() -> bool {
31    static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
32    *OFF.get_or_init(|| std::env::var_os("METRALE_NO_W8A8_DECODE_PROJ").is_some())
33}
34
35/// 2026-09-25: The padded decode row counts this route takes. Any other count
36/// stays on the caller's own kernels.
37pub const DECODE_W8A8_ROWS: std::ops::RangeInclusive<usize> = 5..=16;
38
39/// 2026-09-25: Elements from the start of the output to one past the last
40/// element a cuBLASLt GEMM writes when its `m_pad` output rows of `n` elements
41/// are `ldc` elements apart: `(m_pad - 1) * ldc + n`. Pass the padded row
42/// count, because the pad rows are written too. The strided call sites bound
43/// their buffers with it.
44pub fn strided_out_extent_elems(m_pad: u32, ldc: u32, n: u32) -> usize {
45    debug_assert!(m_pad >= 1);
46    (m_pad as usize - 1) * ldc as usize + n as usize
47}
48
49/// 2026-09-25: The shape and output of one decode projection, as
50/// [`decode_w8a8_selected`] judges it. The route is taken only when all of
51/// these hold:
52///
53/// * `family_armed`: the caller's [`super::CublasScope`] family (`ssm` or
54///   `attn`) is on.
55/// * `disabled` is false ([`w8a8_decode_proj_disabled`]).
56/// * `rows` is in [`DECODE_W8A8_ROWS`].
57/// * The weight is `Fp8BlockScaled`, since cuBLASLt is given its scales as a
58///   `[N/128, K/128]` grid.
59/// * `n` and `k` are multiples of 128, and `blk128x128_stride_ok(k)` holds
60///   (`k` a multiple of 512), the weight-scale stride rule that
61///   `metrale_gpu_runtime::cublaslt::scale_layout` quotes.
62/// * `ldc >= n`.
63/// * The output holds the padded write extent ([`strided_out_extent_elems`]).
64/// * The quantizer resolved, and the scratch holds `m_pad` rows in the active
65///   scale layout, adapter kernel and buffer included when it is K-major.
66#[derive(Clone, Copy, Debug)]
67pub struct DecodeW8a8Plan {
68    /// 2026-09-25: Padded decode rows (the forward context's `n`).
69    pub rows: usize,
70    /// 2026-09-25: Output width of this projection.
71    pub n: u32,
72    /// 2026-09-25: Contract width of this projection.
73    pub k: u32,
74    /// 2026-09-25: Output row pitch in BF16 elements (`n` for a contiguous
75    /// output).
76    pub ldc: u32,
77    /// 2026-09-25: Allocated size of the output buffer, in bytes.
78    pub out_capacity_bytes: usize,
79}
80
81impl DecodeW8a8Plan {
82    /// 2026-09-25: A contiguous `[rows, n]` output (GDN `in_proj_qkvz` and
83    /// `out_proj`, attention `o_proj`).
84    pub fn contiguous(rows: usize, n: u32, k: u32, out_capacity_bytes: usize) -> Self {
85        Self {
86            rows,
87            n,
88            k,
89            ldc: n,
90            out_capacity_bytes,
91        }
92    }
93
94    /// 2026-09-25: A strided output: rows `ldc` BF16 elements apart (attention
95    /// Q/K/V into the `[n, per_seq_qkv]` multi-sequence QKV buffer).
96    pub fn strided(rows: usize, n: u32, k: u32, ldc: u32, out_capacity_bytes: usize) -> Self {
97        Self {
98            rows,
99            n,
100            k,
101            ldc,
102            out_capacity_bytes,
103        }
104    }
105
106    /// 2026-09-25: The padded M cuBLASLt is handed.
107    pub fn m_pad(&self) -> u32 {
108        cublas_fp8_m_pad(self.rows as u32)
109    }
110
111    /// 2026-09-25: Bytes of `out` this projection may write, pad rows included.
112    pub fn write_extent_bytes(&self) -> usize {
113        strided_out_extent_elems(self.m_pad(), self.ldc, self.n) * 2
114    }
115}
116
117/// 2026-09-25: The activation-quant scratch buffers with their capacities, and
118/// the quantizer and scale-layout adapter kernels.
119#[derive(Clone, Copy, Debug)]
120pub struct DecodeW8a8Scratch {
121    pub act_fp8: DevicePtr,
122    pub act_fp8_bytes: usize,
123    pub act_scale: DevicePtr,
124    pub act_scale_bytes: usize,
125    pub act_scale_kmajor: DevicePtr,
126    pub act_scale_kmajor_bytes: usize,
127    pub quant_k: Fp8ActQuant,
128    pub scale_kmajor_k: KernelHandle,
129}
130
131impl DecodeW8a8Scratch {
132    /// 2026-09-25: Whether the quantizer resolved and the scratch holds `m_pad`
133    /// rows of a K-wide activation in the active scale layout.
134    fn fits(&self, m_pad: u32, k: u32) -> bool {
135        let rows = m_pad as usize;
136        let kg = k as usize / 128;
137        self.act_fp8.0 != 0
138            && self.act_scale.0 != 0
139            && self.quant_k.available()
140            && self.act_fp8_bytes >= rows * k as usize
141            && self.act_scale_bytes >= rows * kg * 4
142            && (!cublas_scale_layout_kmajor()
143                || (self.scale_kmajor_k.0 != 0
144                    && self.act_scale_kmajor.0 != 0
145                    && self.act_scale_kmajor_bytes >= rows * kg * 4))
146    }
147}
148
149/// 2026-09-25: Whether one decode projection takes the W8A8 cuBLASLt route;
150/// [`DecodeW8a8Plan`] lists the conditions.
151pub fn decode_w8a8_selected(
152    family_armed: bool,
153    disabled: bool,
154    plan: &DecodeW8a8Plan,
155    scale_format: crate::weight_map::WeightQuantFormat,
156    scratch: &DecodeW8a8Scratch,
157) -> bool {
158    let m_pad = plan.m_pad();
159    family_armed
160        && !disabled
161        && DECODE_W8A8_ROWS.contains(&plan.rows)
162        && scale_format == crate::weight_map::WeightQuantFormat::Fp8BlockScaled
163        && plan.n.is_multiple_of(128)
164        && plan.k.is_multiple_of(128)
165        && metrale_gpu_runtime::cublaslt::scale_layout::blk128x128_stride_ok(plan.k as usize)
166        && plan.ldc >= plan.n
167        && plan.write_extent_bytes() <= plan.out_capacity_bytes
168        && scratch.fits(m_pad, plan.k)
169}
170
171/// 2026-09-25: Quantize BF16 `act[rows, k]` into the scratch: FP8 E4M3 bytes
172/// and one FP32 scale per token and 128 of K, with the pad rows
173/// `rows..ceil16(rows)` zeroed, and the scales transposed for cuBLASLt when
174/// the layout is K-major.
175///
176/// Separate from the GEMM so projections over one activation share one
177/// quantization; the attention Q/K/V do.
178///
179/// The pad rows' FP8 bytes are zeroed as well as their scales: a zero scale
180/// alone is not enough, because `NaN * 0.0` is `NaN`.
181pub fn decode_w8a8_quant_act(
182    gpu: &dyn GpuBackend,
183    scratch: &DecodeW8a8Scratch,
184    act_bf16: DevicePtr,
185    rows: u32,
186    k: u32,
187    stream: u64,
188) -> Result<()> {
189    per_token_group_quant_fp8(
190        gpu,
191        scratch.quant_k,
192        act_bf16,
193        scratch.act_fp8,
194        scratch.act_scale,
195        rows,
196        k,
197        stream,
198    )?;
199    let m_pad = cublas_fp8_m_pad(rows);
200    if m_pad > rows {
201        gpu.memset_async(
202            scratch.act_fp8.offset(rows as usize * k as usize),
203            0,
204            (m_pad - rows) as usize * k as usize,
205            stream,
206        )?;
207    }
208    if cublas_scale_layout_kmajor() {
209        // 2026-09-25: Writes every `[K/128, m_pad]` slot, pad rows included.
210        fp8_act_scale_to_kmajor(
211            gpu,
212            scratch.scale_kmajor_k,
213            scratch.act_scale,
214            scratch.act_scale_kmajor,
215            rows,
216            m_pad,
217            k,
218            stream,
219        )?;
220    } else {
221        // 2026-09-25: In the `rowmajor` layout the pad rows are a contiguous
222        // tail, so they are zeroed in place.
223        let kg = k as usize / 128;
224        if m_pad > rows {
225            gpu.memset_async(
226                scratch.act_scale.offset(rows as usize * kg * 4),
227                0,
228                (m_pad - rows) as usize * kg * 4,
229                stream,
230            )?;
231        }
232    }
233    Ok(())
234}
235
236/// 2026-09-25: `out[m_pad, n] = act_fp8[m_pad, k] @ weight[n, k]ᵀ` at row pitch
237/// `plan.ldc`, with both sets of scales applied. The activation must already
238/// be through [`decode_w8a8_quant_act`].
239pub fn decode_w8a8_gemm(
240    scratch: &DecodeW8a8Scratch,
241    fp8w: &crate::weight_map::Fp8Weight,
242    out: DevicePtr,
243    plan: &DecodeW8a8Plan,
244    stream: u64,
245) -> Result<()> {
246    let b_scale = if cublas_scale_layout_kmajor() {
247        scratch.act_scale_kmajor
248    } else {
249        scratch.act_scale
250    };
251    metrale_gpu_runtime::cublaslt::fp8_gemm_act_weight_t_blkscaled_ldc(
252        scratch.act_fp8.0,
253        b_scale.0,
254        fp8w.weight.0,
255        fp8w.row_scale.0,
256        out.0,
257        plan.m_pad(),
258        plan.n,
259        plan.k,
260        plan.ldc,
261        stream,
262    )
263}
264
265#[cfg(test)]
266#[path = "dispatch_proj_decode_tests.rs"]
267mod tests;
metrale_model_layers/layers/ops/dispatch_proj_decode.rs

metrale_model_layers/layers/ops/
dispatch_proj_decode.rs
