metrale_model_layers/layers/ops/
dispatch_proj.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: cuBLASLt projection routes (block-scaled FP8, BF16) and the
4//! FP8-to-BF16 weight dequant.
5//!
6//! Owner: model-layers ops.
7//! Invariants:
8//! - The block-scaled FP8 route hands cuBLASLt `cublas_fp8_m_pad(m)` rows, and
9//!   zeroes the pad rows' FP8 bytes and activation scales before the GEMM.
10
11#![allow(unused_imports)]
12
13use super::*;
14
15// 2026-09-25: `ops.rs` loads this file with `#[path]`, so a plain
16// `mod cutlass;` would resolve in `ops/`, not `ops/dispatch_proj/`; hence the
17// explicit `#[path]`.
18#[path = "dispatch_proj/cutlass.rs"]
19mod cutlass;
20pub use cutlass::*;
21
22/// 2026-09-25: Which VEC128 activation-scale layout the cuBLASLt block-scaled
23/// arm passes, from `METRALE_CUBLAS_SCALE_LAYOUT`, read once per process.
24///
25/// * `true` (unset, or any value but `rowmajor`): `[K/128, ceil16(M)]`, tokens
26///   contiguous, the layout `metrale_gpu_runtime::cublaslt::scale_layout`
27///   quotes from the cuBLAS manual.
28/// * `false` (`rowmajor`): the quantizer's `[M, K/128]` passed as is. It is a
29///   measurement control: measured 2026-09-11 on H100, it gave rel_rms 7.7e-2
30///   against the in-tree kernel.
31pub fn cublas_scale_layout_kmajor() -> bool {
32    static KMAJOR: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
33    *KMAJOR.get_or_init(|| {
34        !matches!(
35            std::env::var("METRALE_CUBLAS_SCALE_LAYOUT").as_deref(),
36            Ok("rowmajor")
37        )
38    })
39}
40
41/// 2026-09-25: Launch `fp8_act_scale_to_kmajor`: copy the quantizer's
42/// row-major `[M, K/128]` FP32 activation scales into `[K/128, M_pad]`, with
43/// zeros in the `M..M_pad` pad rows. `a_scale` is only read.
44///
45/// `metrale_gpu_runtime::cublaslt::scale_layout` holds the same index math and
46/// its CPU tests.
47pub fn fp8_act_scale_to_kmajor(
48    gpu: &dyn metrale_gpu_runtime::gpu::GpuBackend,
49    kernel: metrale_gpu_runtime::gpu::KernelHandle,
50    a_scale: metrale_gpu_runtime::gpu::DevicePtr,
51    a_scale_kmajor: metrale_gpu_runtime::gpu::DevicePtr,
52    m: u32,
53    m_pad: u32,
54    k: u32,
55    stream: u64,
56) -> anyhow::Result<()> {
57    use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
58    let l = k / 128;
59    KernelLaunch::new(gpu, kernel)
60        .grid([div_ceil(m_pad, 256), l, 1])
61        .block([256, 1, 1])
62        .arg_ptr(a_scale)
63        .arg_ptr(a_scale_kmajor)
64        .arg_u32(m)
65        .arg_u32(m_pad)
66        .arg_u32(l)
67        .launch(stream)
68}
69
70/// 2026-09-25: Quantize the activation with `per_token_group_quant_fp8` (FP8
71/// plus one scale per token and 128 of K), then run
72/// [`cublas_fp8_proj_prequant`]. The FP8 weight and its 128x128 block scales
73/// go to cuBLASLt as they are, with no dequant.
74///
75/// The scratch buffers need the extents listed on
76/// [`cublas_fp8_proj_prequant`].
77#[allow(clippy::too_many_arguments)]
78pub fn cublas_fp8_proj(
79    gpu: &dyn metrale_gpu_runtime::gpu::GpuBackend,
80    ptg_quant_k: Fp8ActQuant,
81    scale_kmajor_k: metrale_gpu_runtime::gpu::KernelHandle,
82    act_bf16: metrale_gpu_runtime::gpu::DevicePtr,
83    act_fp8_scratch: metrale_gpu_runtime::gpu::DevicePtr,
84    act_scale_scratch: metrale_gpu_runtime::gpu::DevicePtr,
85    act_scale_kmajor_scratch: metrale_gpu_runtime::gpu::DevicePtr,
86    fp8w: &crate::weight_map::Fp8Weight,
87    out: metrale_gpu_runtime::gpu::DevicePtr,
88    m: u32,
89    n: u32,
90    k: u32,
91    stream: u64,
92) -> anyhow::Result<()> {
93    per_token_group_quant_fp8(
94        gpu,
95        ptg_quant_k,
96        act_bf16,
97        act_fp8_scratch,
98        act_scale_scratch,
99        m,
100        k,
101        stream,
102    )?;
103    cublas_fp8_proj_prequant(
104        gpu,
105        scale_kmajor_k,
106        act_fp8_scratch,
107        act_scale_scratch,
108        act_scale_kmajor_scratch,
109        fp8w,
110        out,
111        m,
112        n,
113        k,
114        stream,
115    )
116}
117
118/// 2026-09-25: [`cublas_fp8_proj`] for an activation the caller already ran
119/// through `per_token_group_quant_fp8`, so one quantization can feed several
120/// projections (the dense FFN's gate and up share one).
121///
122/// With [`cublas_scale_layout_kmajor`] on, the scales are first transposed
123/// into `act_scale_kmajor`; a zero kernel handle or scratch pointer is an
124/// error.
125///
126/// cuBLASLt is handed `ceil16(M)` rows ([`cublas_fp8_m_pad`]), so:
127///
128/// * `out` must hold `ceil16(M) * N` BF16 elements; the pad rows are written.
129/// * `act_fp8` must hold `ceil16(M) * K` bytes; the pad rows are zeroed here.
130/// * `act_scale_kmajor` must hold `ceil16(M) * (K/128)` f32, and `act_scale`
131///   `M * (K/128)` f32, or `ceil16(M) * (K/128)` under `rowmajor`, where the
132///   pad rows are zeroed in place.
133#[allow(clippy::too_many_arguments)]
134pub fn cublas_fp8_proj_prequant(
135    gpu: &dyn metrale_gpu_runtime::gpu::GpuBackend,
136    scale_kmajor_k: metrale_gpu_runtime::gpu::KernelHandle,
137    act_fp8: metrale_gpu_runtime::gpu::DevicePtr,
138    act_scale: metrale_gpu_runtime::gpu::DevicePtr,
139    act_scale_kmajor: metrale_gpu_runtime::gpu::DevicePtr,
140    fp8w: &crate::weight_map::Fp8Weight,
141    out: metrale_gpu_runtime::gpu::DevicePtr,
142    m: u32,
143    n: u32,
144    k: u32,
145    stream: u64,
146) -> anyhow::Result<()> {
147    let m_pad = cublas_fp8_m_pad(m);
148    let kg = (k / 128) as usize;
149    if m_pad > m {
150        // 2026-09-25: A zero scale alone is not enough: the FP8 dot product
151        // still reads the pad bytes, and `NaN * 0.0` is `NaN`.
152        gpu.memset_async(
153            act_fp8.offset(m as usize * k as usize),
154            0,
155            (m_pad - m) as usize * k as usize,
156            stream,
157        )?;
158    }
159    let b_scale = if cublas_scale_layout_kmajor() {
160        if scale_kmajor_k.0 == 0 || act_scale_kmajor.0 == 0 {
161            anyhow::bail!(
162                "cuBLASLt block-scaled FP8 needs the fp8_act_scale_to_kmajor adapter \
163                 (kernel={:#x}, scratch={:#x}) — see cublas_scale_layout_kmajor()",
164                scale_kmajor_k.0,
165                act_scale_kmajor.0
166            );
167        }
168        // 2026-09-25: The adapter writes every `[K/128, m_pad]` slot, pad rows
169        // included, so the scale pad needs no memset.
170        fp8_act_scale_to_kmajor(
171            gpu,
172            scale_kmajor_k,
173            act_scale,
174            act_scale_kmajor,
175            m,
176            m_pad,
177            k,
178            stream,
179        )?;
180        act_scale_kmajor
181    } else {
182        // 2026-09-25: In the `rowmajor` layout the pad rows are a contiguous
183        // tail, so they are zeroed in place.
184        if m_pad > m {
185            gpu.memset_async(
186                act_scale.offset(m as usize * kg * 4),
187                0,
188                (m_pad - m) as usize * kg * 4,
189                stream,
190            )?;
191        }
192        act_scale
193    };
194    metrale_gpu_runtime::cublaslt::fp8_gemm_act_weight_t_blkscaled(
195        act_fp8.0,
196        b_scale.0,
197        fp8w.weight.0,
198        fp8w.row_scale.0,
199        out.0,
200        m_pad,
201        n,
202        k,
203        stream,
204    )
205}
206
207/// 2026-09-25: The M extent [`cublas_fp8_proj_prequant`] hands cuBLASLt; callers
208/// bound their output buffers against it.
209pub fn cublas_fp8_m_pad(m: u32) -> u32 {
210    m.div_ceil(16) * 16
211}
212
213/// 2026-09-25: Dequantize a block-scaled or per-row FP8 weight `[N,K]` to BF16
214/// into `dst`, a caller-owned buffer of [`dequant_fp8_bf16_bytes`] bytes.
215///
216/// One kernel serves both scale layouts; the block geometry passed to it
217/// selects between them:
218///
219///   block-scaled   block_n = block_k = 128, sk = K/128
220///   per-row        block_n = 1, block_k = K, sk = 1
221pub fn dequant_fp8_bf16_into(
222    gpu: &dyn metrale_gpu_runtime::gpu::GpuBackend,
223    fp8w: &crate::weight_map::Fp8Weight,
224    dst: metrale_gpu_runtime::gpu::DevicePtr,
225    stream: u64,
226) -> anyhow::Result<()> {
227    use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
228    let (n, kk) = (fp8w.n, fp8w.k);
229    let per_row = fp8w.scale_format == crate::weight_map::WeightQuantFormat::Fp8PerRow;
230    let (block_n, block_k, sk) = if per_row {
231        (1u32, kk, 1u32)
232    } else {
233        (128u32, 128u32, kk / 128)
234    };
235    let kernel = gpu.kernel(
236        "dequant_fp8_blockscaled_bf16",
237        "dequant_fp8_blockscaled_bf16",
238    )?;
239    KernelLaunch::new(gpu, kernel)
240        .grid([div_ceil(kk, 64), div_ceil(n, 4), 1])
241        .block([64, 4, 1])
242        .arg_ptr(fp8w.weight)
243        .arg_ptr(fp8w.row_scale)
244        .arg_ptr(dst)
245        .arg_u32(n)
246        .arg_u32(kk)
247        .arg_u32(block_n)
248        .arg_u32(block_k)
249        .arg_u32(sk)
250        .arg_u32(1)
251        .launch(stream)
252}
253
254/// 2026-09-25: BF16 bytes [`dequant_fp8_bf16_into`] writes for `fp8w`.
255pub fn dequant_fp8_bf16_bytes(fp8w: &crate::weight_map::Fp8Weight) -> usize {
256    fp8w.n as usize * fp8w.k as usize * 2
257}
258
259/// 2026-09-25: Route a projection `out[M,N] = act[M,K] @ weightᵀ` through
260/// cuBLASLt BF16 for a weight that is already BF16 `[N,K]`.
261pub fn cublas_bf16_proj_dense(
262    act: metrale_gpu_runtime::gpu::DevicePtr,
263    weight_bf16: metrale_gpu_runtime::gpu::DevicePtr,
264    out: metrale_gpu_runtime::gpu::DevicePtr,
265    m: u32,
266    n: u32,
267    k: u32,
268    stream: u64,
269) -> anyhow::Result<()> {
270    metrale_gpu_runtime::cublaslt::bf16_gemm_act_weight_t(
271        act.0,
272        weight_bf16.0,
273        out.0,
274        m,
275        n,
276        k,
277        stream,
278    )
279}
280
281/// 2026-09-25: [`cublas_bf16_proj_dense`] writing an FP32 output buffer, for
282/// consumers whose next kernel reads FP32.
283pub fn cublas_bf16_proj_dense_f32_out(
284    act: metrale_gpu_runtime::gpu::DevicePtr,
285    weight_bf16: metrale_gpu_runtime::gpu::DevicePtr,
286    out: metrale_gpu_runtime::gpu::DevicePtr,
287    m: u32,
288    n: u32,
289    k: u32,
290    stream: u64,
291) -> anyhow::Result<()> {
292    metrale_gpu_runtime::cublaslt::bf16_gemm_act_weight_t_f32_out(
293        act.0,
294        weight_bf16.0,
295        out.0,
296        m,
297        n,
298        k,
299        stream,
300    )
301}
metrale_model_layers/layers/ops/dispatch_proj.rs

metrale_model_layers/layers/ops/
dispatch_proj.rs
