metrale_model_layers/layers/ops/
gemm_quant.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the FP8 x FP8 and FP8-weight GEMMs, the per-token
4//! FP8 activation quantizer, the block-scaled W8A8 GEMM, the FP8 weight and
5//! scale transposes, block-scale widening and the M-dispatched BF16 GEMM
6//! [`dense_mm_bf16`]. Re-exports the grouped MoE, W8A16 and GEMV launchers.
7//!
8//! Owner: model-layers ops.
9//! Invariants: none beyond the types.
10
11#![allow(unused_imports)]
12
13use anyhow::{Result, ensure};
14use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
15use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
16
17use crate::layers::moe;
18use crate::weight_map::{DenseWeight, Fp8DenseWeight, Fp8Weight, QuantizedWeight};
19
20use super::*;
21#[path = "gemm_quant_moe.rs"]
22mod moe_grouped;
23pub use moe_grouped::{
24    MOE_E4M3_DN, MOE_E4M3_GU, MoeE4m3Tile, moe_bf16_grouped_gemm, moe_build_tile_worklist,
25    moe_fp8_grouped_gemm, moe_gate_topk_fused, moe_w8a8_gateup_silu_e4m3, moe_w8a8_grouped_gemm,
26    moe_w8a8_grouped_gemm_e4m3, moe_w8a8_grouped_gemm_pm4,
27};
28#[path = "gemm_quant_w8a16.rs"]
29mod w8a16;
30pub use w8a16::{
31    w8a16_gemm, w8a16_gemm_n128_m128, w8a16_gemm_pipelined, w8a16_gemm_t, w8a16_gemm_t_pipelined,
32    w8a16_gemv, w8a16_gemv_row_tiered,
33};
34#[path = "gemm_quant_gemv.rs"]
35mod gemv;
36pub use gemv::{
37    DENSE_GEMV_BATCHM_DECODE_MAX_M, DENSE_GEMV_BATCHM_MAX_M, dense_gemv, dense_gemv_batch2,
38    dense_gemv_batchm, dense_gemv_batchm_split, dense_gemv_fp8w,
39};
40
41/// 2026-09-25: FP8 x FP8 GEMM: A `[M, K]` FP8 E4M3 and B `[N, K]` FP8 E4M3
42/// give C `[M, N]` BF16.
43pub fn fp8_fp8_gemm_n128(
44    gpu: &dyn GpuBackend,
45    kernel: KernelHandle,
46    a_fp8: DevicePtr,
47    b_fp8: DevicePtr,
48    output: DevicePtr,
49    m: u32,
50    n: u32,
51    k: u32,
52    stream: u64,
53) -> Result<()> {
54    KernelLaunch::new(gpu, kernel)
55        .grid([div_ceil(n, 128), div_ceil(m, 64), 1])
56        .block([128, 1, 1])
57        .arg_ptr(a_fp8)
58        .arg_ptr(b_fp8)
59        .arg_ptr(output)
60        .arg_u32(m)
61        .arg_u32(n)
62        .arg_u32(k)
63        .launch(stream)
64}
65
66/// 2026-09-25: FP8-weight GEMM with a 128-row CTA (two 64-row chunks); each
67/// CTA loads a B tile once for both chunks.
68#[allow(clippy::too_many_arguments)]
69pub fn fp8_gemm_n128_m128(
70    gpu: &dyn GpuBackend,
71    kernel: KernelHandle,
72    input: DevicePtr,
73    b_fp8: DevicePtr,
74    output: DevicePtr,
75    m: u32,
76    n: u32,
77    k: u32,
78    stream: u64,
79) -> Result<()> {
80    check_e4m3_activation_range(gpu, kernel, input, m, k, stream)?;
81    KernelLaunch::new(gpu, kernel)
82        .grid([div_ceil(n, 128), div_ceil(m, 128), 1])
83        .block([128, 1, 1])
84        .arg_ptr(input)
85        .arg_ptr(b_fp8)
86        .arg_ptr(output)
87        .arg_u32(m)
88        .arg_u32(n)
89        .arg_u32(k)
90        .launch(stream)
91}
92
93/// 2026-09-25: [`fp8_fp8_gemm_n128`] with a 128-row CTA (two 64-row chunks);
94/// each CTA loads a B tile once for both chunks.
95#[allow(clippy::too_many_arguments)]
96pub fn fp8_fp8_gemm_n128_m128(
97    gpu: &dyn GpuBackend,
98    kernel: KernelHandle,
99    a_fp8: DevicePtr,
100    b_fp8: DevicePtr,
101    output: DevicePtr,
102    m: u32,
103    n: u32,
104    k: u32,
105    stream: u64,
106) -> Result<()> {
107    KernelLaunch::new(gpu, kernel)
108        .grid([div_ceil(n, 128), div_ceil(m, 128), 1])
109        .block([128, 1, 1])
110        .arg_ptr(a_fp8)
111        .arg_ptr(b_fp8)
112        .arg_ptr(output)
113        .arg_u32(m)
114        .arg_u32(n)
115        .arg_u32(k)
116        .launch(stream)
117}
118
119/// 2026-09-25: Per-token FP8 activation quantization over 128-element
120/// K-groups: A_fp8 `[M, K]` FP8 E4M3 and a_scale `[M, K/128]` FP32.
121///
122/// [`Fp8ActQuant::pick`] chooses the shared kernel or the Hopper twin for this
123/// width and returns the kernel and its grid together; both run 128-thread
124/// blocks. This is the only caller of `fp8_quant_log`, which logs the choice
125/// once per branch.
126#[allow(clippy::too_many_arguments)]
127pub fn per_token_group_quant_fp8(
128    gpu: &dyn GpuBackend,
129    quant: Fp8ActQuant,
130    input_bf16: DevicePtr,
131    output_fp8: DevicePtr,
132    a_scale: DevicePtr,
133    m: u32,
134    k: u32,
135    stream: u64,
136) -> Result<()> {
137    let pick = quant.pick(m, k);
138    super::fp8_quant_log(&pick, m, k);
139    KernelLaunch::new(gpu, pick.kernel)
140        .grid(pick.grid)
141        .block([128, 1, 1])
142        .arg_ptr(input_bf16)
143        .arg_ptr(output_fp8)
144        .arg_ptr(a_scale)
145        .arg_u32(m)
146        .arg_u32(k)
147        .launch(stream)
148}
149
150/// 2026-09-25: W8A8 GEMM with per-token activation scales and per-block weight
151/// scales, applied in an FP32 epilogue per 128-element K-group:
152///
153///   C[M, N] = bf16( Σ_g (FP8 MMA over K-group g) × a_scale[M, g] × b_scale[N/128, g] )
154///
155/// Inputs:
156///   - `a_fp8`     [M, K] FP8 E4M3
157///   - `a_scale`   [M, K/128] FP32 (from [`per_token_group_quant_fp8`])
158///   - `b_fp8`     [N, K] FP8 E4M3
159///   - `b_scale`   [N/128, K/128] FP32
160///   - `output`    [M, N] BF16
161#[allow(clippy::too_many_arguments)]
162pub fn fp8_gemm_t_blockscaled(
163    gpu: &dyn GpuBackend,
164    kernel: KernelHandle,
165    a_fp8: DevicePtr,
166    a_scale: DevicePtr,
167    b_fp8: DevicePtr,
168    b_scale: DevicePtr,
169    output: DevicePtr,
170    m: u32,
171    n: u32,
172    k: u32,
173    stream: u64,
174) -> Result<()> {
175    super::log_gemm_shape(gpu, "fp8_gemm_t_blockscaled", m, n, k);
176    if let Some(pipe) = fp8_gemm_pipe_kernel(gpu, m, n, k)? {
177        // 2026-09-27: `fp8_gemm_blockscaled_pipe_128x64`: 128 x 64 tiles, 256 threads, SmemBytes<128, 64, 3>.
178        const PIPE_SMEM: u32 = 3 * (128 + 64) * 64 + 3 * 128 * 4 + 128 * 4;
179        return KernelLaunch::new(gpu, pipe)
180            .grid([n / 64, div_ceil(m, 128), 1])
181            .block([256, 1, 1])
182            .shared_mem(PIPE_SMEM)
183            .arg_ptr(a_fp8)
184            .arg_ptr(a_scale)
185            .arg_ptr(b_fp8)
186            .arg_ptr(b_scale)
187            .arg_ptr(output)
188            .arg_u32(m)
189            .arg_u32(n)
190            .arg_u32(k)
191            .launch(stream);
192    }
193    KernelLaunch::new(gpu, kernel)
194        .grid([div_ceil(n, 128), div_ceil(m, 64), 1])
195        .block([128, 1, 1])
196        .arg_ptr(a_fp8)
197        .arg_ptr(a_scale)
198        .arg_ptr(b_fp8)
199        .arg_ptr(b_scale)
200        .arg_ptr(output)
201        .arg_u32(m)
202        .arg_u32(n)
203        .arg_u32(k)
204        .launch(stream)
205}
206
207/// 2026-09-27: The pipelined twin of `fp8_gemm_t_blockscaled`
208/// (`kernels/gb10/common/fp8_gemm_blockscaled_pipe.cu`), when the backend carries it, the
209/// shape fits (N a multiple of its 64-wide tile, K of 128) and `METRALE_NO_FP8_GEMM_PIPE`
210/// is absent. It performs the same MMAs and scale folds in the same order per output
211/// element, so its output is bit-identical (microbench: 0 differing values at every
212/// 35B projection shape, over every non-NaN E4M3 code), 1.7-3.2x faster on GB10. The
213/// handle is memoized in the backend's `OpCache`.
214fn fp8_gemm_pipe_kernel(
215    gpu: &dyn GpuBackend,
216    m: u32,
217    n: u32,
218    k: u32,
219) -> Result<Option<KernelHandle>> {
220    const MODULE: &str = "fp8_gemm_blockscaled_pipe";
221    static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
222    let off = *OFF.get_or_init(|| std::env::var_os("METRALE_NO_FP8_GEMM_PIPE").is_some());
223    if off
224        || m == 0
225        || n == 0
226        || !n.is_multiple_of(64)
227        || !k.is_multiple_of(128)
228        || !gpu.has_module(MODULE)
229    {
230        return Ok(None);
231    }
232    Ok(Some(gpu.op_cache().kernel(
233        gpu,
234        MODULE,
235        "fp8_gemm_blockscaled_pipe_128x64",
236    )?))
237}
238
239/// 2026-09-25: Transpose an FP8 weight on the GPU: `B[N, K]` to `B_t[K, N]`.
240pub fn transpose_fp8(
241    gpu: &dyn GpuBackend,
242    kernel: KernelHandle,
243    src: DevicePtr,
244    dst: DevicePtr,
245    n: u32,
246    k: u32,
247    stream: u64,
248) -> Result<()> {
249    let total = n as u64 * k as u64;
250    KernelLaunch::new(gpu, kernel)
251        .grid([div_ceil(total as u32, 256), 1, 1])
252        .block([256, 1, 1])
253        .arg_ptr(src)
254        .arg_ptr(dst)
255        .arg_u32(n)
256        .arg_u32(k)
257        .launch(stream)
258}
259
260/// 2026-09-25: Widen an FP8 block-scale tensor to FP32 on the GPU, so the
261/// block-scale kernels read `const float*`. `src` is `[total]` BF16 (dtype 0),
262/// FP32 (1) or E8M0 (2); `dst` is `[total]` FP32. An E8M0 byte `e` becomes
263/// `2^(e - 127)` (the bit pattern `e << 23`), `e = 0` becomes `2^-127` and
264/// `e = 255` becomes 0.0.
265pub fn widen_block_scale_f32(
266    gpu: &dyn GpuBackend,
267    kernel: KernelHandle,
268    src: DevicePtr,
269    dst: DevicePtr,
270    total: u32,
271    input_dtype: u32,
272    stream: u64,
273) -> Result<()> {
274    KernelLaunch::new(gpu, kernel)
275        .grid([div_ceil(total, 256), 1, 1])
276        .block([256, 1, 1])
277        .arg_ptr(src)
278        .arg_ptr(dst)
279        .arg_u32(total)
280        .arg_u32(input_dtype)
281        .launch(stream)
282}
283
284/// 2026-09-25: Transpose FP32 block scales: `[N/128, K/128]` to
285/// `[K/128, N/128]`.
286pub fn transpose_block_scale(
287    gpu: &dyn GpuBackend,
288    kernel: KernelHandle,
289    src: DevicePtr,
290    dst: DevicePtr,
291    n_blocks: u32,
292    k_blocks: u32,
293    stream: u64,
294) -> Result<()> {
295    let total = n_blocks * k_blocks;
296    KernelLaunch::new(gpu, kernel)
297        .grid([div_ceil(total, 256), 1, 1])
298        .block([256, 1, 1])
299        .arg_ptr(src)
300        .arg_ptr(dst)
301        .arg_u32(n_blocks)
302        .arg_u32(k_blocks)
303        .launch(stream)
304}
305
306/// 2026-09-25: The three BF16 dense kernels [`dense_mm_bf16`] chooses from,
307/// resolved once per projection site (the GLM-5.3 mixer and MLP sites bind
308/// one each).
309#[derive(Clone, Copy)]
310pub struct DenseMmKernels {
311    /// 2026-09-25: `dense_gemm_bf16`, the 16x16 tile GEMM: the arm for M above
312    /// [`DENSE_GEMV_BATCHM_MAX_M`], and for any M whose GEMV kernel is missing.
313    pub gemm: KernelHandle,
314    /// 2026-09-25: `dense_gemv_bf16`, for M == 1.
315    pub gemv: KernelHandle,
316    /// 2026-09-25: `dense_gemv_bf16_batchm`, for 2..=[`DENSE_GEMV_BATCHM_MAX_M`]
317    /// rows in one pass over the weight; `KernelHandle(0)` when unavailable.
318    pub batchm: KernelHandle,
319}
320
321/// 2026-09-25: `C[M, N] = A[M, K] @ B[N, K]^T`, BF16 in and out, output row
322/// stride `N`, dispatched on M: `dense_gemv_bf16` at M == 1,
323/// `dense_gemv_bf16_batchm` at 2..=[`DENSE_GEMV_BATCHM_MAX_M`], and the tile
324/// GEMM otherwise or when the chosen GEMV kernel is missing.
325///
326/// `batchm` reads the weight once for all M rows. Each of its rows follows
327/// `dense_gemv_bf16`'s K order and reduction (gb10 kernels), so the two GEMV
328/// arms give the same bits per row; the tile GEMM accumulates in a different
329/// order.
330#[allow(clippy::too_many_arguments)]
331pub fn dense_mm_bf16(
332    gpu: &dyn GpuBackend,
333    k: &DenseMmKernels,
334    a: DevicePtr,
335    b: DevicePtr,
336    c: DevicePtr,
337    m: usize,
338    n: usize,
339    kk: usize,
340    stream: u64,
341) -> Result<()> {
342    // 2026-09-25: Each grid follows its kernel's tile: `N_PER_BLOCK` = 4 outputs
343    // per 256-thread block in both GEMV kernels, `GEMM_TILE` in the tile GEMM.
344    if m == 1 && k.gemv.0 != 0 {
345        return KernelLaunch::new(gpu, k.gemv)
346            .grid([div_ceil(n as u32, 4), 1, 1])
347            .block([256, 1, 1])
348            .arg_ptr(a)
349            .arg_ptr(b)
350            .arg_ptr(c)
351            .arg_u32(n as u32)
352            .arg_u32(kk as u32)
353            .launch(stream);
354    }
355    // 2026-09-25: Without `batchm`, M > 1 goes to the tile GEMM; warn once per
356    // process.
357    if m > 1 && k.batchm.0 == 0 {
358        static ONCE: std::sync::Once = std::sync::Once::new();
359        ONCE.call_once(|| {
360            tracing::warn!(
361                "dense_mm_bf16: no dense_gemv_bf16_batchm on this target -- M>1 sites fall \
362                 back to the tile GEMM (measured 3.6x slower at M<=8)"
363            );
364        });
365    }
366    if (2..=DENSE_GEMV_BATCHM_MAX_M as usize).contains(&m) && k.batchm.0 != 0 {
367        return KernelLaunch::new(gpu, k.batchm)
368            .grid([div_ceil(n as u32, 4), 1, 1])
369            .block([256, 1, 1])
370            .arg_ptr(a)
371            .arg_ptr(b)
372            .arg_ptr(c)
373            .arg_u32(m as u32)
374            .arg_u32(n as u32)
375            .arg_u32(kk as u32)
376            // 2026-09-25: `out_stride = N`: contiguous `[M, N]`, as the tile GEMM
377            // writes.
378            .arg_u32(n as u32)
379            .launch(stream);
380    }
381    const GEMM_TILE: u32 = 16;
382    KernelLaunch::new(gpu, k.gemm)
383        .grid([
384            (n as u32).div_ceil(GEMM_TILE),
385            (m as u32).div_ceil(GEMM_TILE),
386            1,
387        ])
388        .block([GEMM_TILE, GEMM_TILE, 1])
389        .arg_ptr(a)
390        .arg_ptr(b)
391        .arg_ptr(c)
392        .arg_u32(m as u32)
393        .arg_u32(n as u32)
394        .arg_u32(kk as u32)
395        .launch(stream)
396}
metrale_model_layers/layers/ops/gemm_quant.rs

metrale_model_layers/layers/ops/
gemm_quant.rs
