metrale_model_layers/layers/ops/
gemm_dense.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the W4A16 GEMMs (NVFP4 weights, BF16
4//! activations), and the re-export of the dense BF16 GEMM launchers.
5//!
6//! Owner: model-layers ops.
7//! Invariants: none beyond the types.
8//!
9//! `gemm_dense_tests.rs` reads this file as text: it counts the `.arg_*` calls
10//! of the `_ldb` launchers and reads the body of `w4a16_gemm_n128`.
11
12#![allow(unused_imports)]
13
14use anyhow::Result;
15use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
16use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
17
18use crate::layers::moe;
19use crate::weight_map::{DenseWeight, Fp8DenseWeight, Fp8Weight, QuantizedWeight};
20
21use super::*;
22#[path = "gemm_dense_bf16.rs"]
23mod bf16;
24pub use bf16::{
25    MOE_ROUTER_GEMM_COLS, MOE_ROUTER_GEMM_MAX_PREFILL_ROWS, MOE_ROUTER_GEMM_ROWS,
26    MOE_ROUTER_RT_MIN_ROWS, dense_gemm, dense_gemm_bf16_pipelined, dense_gemm_prefill,
27    dense_gemm_router, dense_gemm_splitk, dense_gemm_tc, dense_gemm_tc_scaled_acc, moe_router_gemm,
28    moe_router_gemm_rt,
29};
30
31/// 2026-09-29: W4A16 GEMM, `C = A @ dequant(B)`: A `[M, K]` BF16, B NVFP4
32/// (packed E2M1, FP8 block scales, FP32 per-tensor scale), C `[M, N]` BF16.
33/// Every launcher here sizes grid.x from the kernel's published N tile
34/// ([`n_tile_blocks`]).
35///
36/// Also launches `w4a16_gemm_t_k64_n64_p3`, which takes the same arguments.
37pub fn w4a16_gemm(
38    gpu: &dyn GpuBackend,
39    kernel: KernelHandle,
40    input: DevicePtr,
41    weight: &QuantizedWeight,
42    output: DevicePtr,
43    m: u32,
44    n: u32,
45    k: u32,
46    stream: u64,
47) -> Result<()> {
48    check_e4m3_activation_range(gpu, kernel, input, m, k, stream)?;
49    KernelLaunch::new(gpu, kernel)
50        .grid([n_tile_blocks(gpu, kernel, n)?, div_ceil(m, 64), 1])
51        .block([128, 1, 1])
52        .arg_ptr(input)
53        .arg_ptr(weight.weight)
54        .arg_ptr(weight.weight_scale)
55        .arg_f32(weight.weight_scale_2)
56        .arg_ptr(output)
57        .arg_u32(m)
58        .arg_u32(n)
59        .arg_u32(k)
60        .launch(stream)
61}
62
63#[allow(clippy::too_many_arguments)]
64/// 2026-09-25: [`w4a16_gemm_n128`] with an explicit row stride `ldb` for the
65/// transposed B, so row `r` starts at `r * ldb`.
66///
67/// The tile kernels load B in 16-byte `cp.async` chunks, which need 16-byte
68/// aligned sources, so an N that is not a multiple of 16 cannot be the
69/// stride. The lm_head's N is the vocab size, which can be odd; its padded
70/// twin is built at `align_up(vocab, 128)` with the pad columns zeroed
71/// (`impl_a1.rs`, `transpose.rs`), and callers pass that as `ldb`.
72pub fn w4a16_gemm_n128_ldb(
73    gpu: &dyn GpuBackend,
74    kernel: KernelHandle,
75    input: DevicePtr,
76    weight: &QuantizedWeight,
77    output: DevicePtr,
78    m: u32,
79    n: u32,
80    k: u32,
81    ldb: u32,
82    stream: u64,
83) -> Result<()> {
84    check_e4m3_activation_range(gpu, kernel, input, m, k, stream)?;
85    KernelLaunch::new(gpu, kernel)
86        .grid([n_tile_blocks(gpu, kernel, n)?, div_ceil(m, 64), 1])
87        .block([128, 1, 1])
88        .arg_ptr(input)
89        .arg_ptr(weight.weight)
90        .arg_ptr(weight.weight_scale)
91        .arg_f32(weight.weight_scale_2)
92        .arg_ptr(output)
93        .arg_u32(m)
94        .arg_u32(n)
95        .arg_u32(k)
96        .arg_u32(ldb)
97        .launch(stream)
98}
99
100/// 2026-09-29: W4A16 GEMM over a packed transposed B, N tiled by the kernel's published
101/// tile: it launches through [`w4a16_gemm_n128_ldb`] with `ldb = n`.
102pub fn w4a16_gemm_n128(
103    gpu: &dyn GpuBackend,
104    kernel: KernelHandle,
105    input: DevicePtr,
106    weight: &QuantizedWeight,
107    output: DevicePtr,
108    m: u32,
109    n: u32,
110    k: u32,
111    stream: u64,
112) -> Result<()> {
113    w4a16_gemm_n128_ldb(gpu, kernel, input, weight, output, m, n, k, n, stream)
114}
115
116/// 2026-09-25: W4A16 GEMM v3 (`w4a16_gemm_t_m128_v3`, only in
117/// `kernels/gb10/minimax-m2-229b/nvfp4/w4a16_gemm_v3.cu`): a 128x128 CTA tile
118/// on 256 threads with a K step of 64.
119#[allow(clippy::too_many_arguments)]
120pub fn w4a16_gemm_n128_m128_v3(
121    gpu: &dyn GpuBackend,
122    kernel: KernelHandle,
123    input: DevicePtr,
124    weight: &QuantizedWeight,
125    output: DevicePtr,
126    m: u32,
127    n: u32,
128    k: u32,
129    stream: u64,
130) -> Result<()> {
131    check_e4m3_activation_range(gpu, kernel, input, m, k, stream)?;
132    KernelLaunch::new(gpu, kernel)
133        .grid([n_tile_blocks(gpu, kernel, n)?, div_ceil(m, 128), 1])
134        .block([256, 1, 1])
135        .arg_ptr(input)
136        .arg_ptr(weight.weight)
137        .arg_ptr(weight.weight_scale)
138        .arg_f32(weight.weight_scale_2)
139        .arg_ptr(output)
140        .arg_u32(m)
141        .arg_u32(n)
142        .arg_u32(k)
143        .launch(stream)
144}
145
146/// 2026-09-25: W4A16 GEMM v2 (`w4a16_gemm_t_m128_v2`, in the minimax-m2-229b
147/// and qwen3.6-27b kernel dirs): the CTA tile of [`w4a16_gemm_n128_m128`]
148/// (M=128, N=128, K step 32) on 256 threads instead of 128. Warps 0-3 compute
149/// the first 64 rows and warps 4-7 the second, so the two chunks' MMAs run in
150/// parallel.
151#[allow(clippy::too_many_arguments)]
152pub fn w4a16_gemm_n128_m128_v2(
153    gpu: &dyn GpuBackend,
154    kernel: KernelHandle,
155    input: DevicePtr,
156    weight: &QuantizedWeight,
157    output: DevicePtr,
158    m: u32,
159    n: u32,
160    k: u32,
161    stream: u64,
162) -> Result<()> {
163    check_e4m3_activation_range(gpu, kernel, input, m, k, stream)?;
164    KernelLaunch::new(gpu, kernel)
165        .grid([n_tile_blocks(gpu, kernel, n)?, div_ceil(m, 128), 1])
166        .block([256, 1, 1])
167        .arg_ptr(input)
168        .arg_ptr(weight.weight)
169        .arg_ptr(weight.weight_scale)
170        .arg_f32(weight.weight_scale_2)
171        .arg_ptr(output)
172        .arg_u32(m)
173        .arg_u32(n)
174        .arg_u32(k)
175        .launch(stream)
176}
177
178/// 2026-09-25: W4A16 GEMM with a 128-row CTA (two 64-row chunks): each CTA
179/// loads a B tile once for both chunks.
180///
181/// Grid contract: N is the fast axis (`blockIdx.x` = N block, `blockIdx.y` =
182/// M block), as every `w4a16_gemm_t_m128` in `kernels/` reads it. Several
183/// layers share this launcher (qwen3_attention, dense_ffn, qwen3_ssm,
184/// nemotron_mamba2, nemotron_moe), so swapping the axes here would mis-map
185/// every CTA of the other kernels with no error. An M-fast order needs its
186/// own kernel and launcher, as `w4a4_gemm_mfast` and `fp8_gemm_t_m128_mfast`
187/// have.
188#[allow(clippy::too_many_arguments)]
189pub fn w4a16_gemm_n128_m128(
190    gpu: &dyn GpuBackend,
191    kernel: KernelHandle,
192    input: DevicePtr,
193    weight: &QuantizedWeight,
194    output: DevicePtr,
195    m: u32,
196    n: u32,
197    k: u32,
198    stream: u64,
199) -> Result<()> {
200    check_e4m3_activation_range(gpu, kernel, input, m, k, stream)?;
201    KernelLaunch::new(gpu, kernel)
202        .grid([n_tile_blocks(gpu, kernel, n)?, div_ceil(m, 128), 1])
203        .block([128, 1, 1])
204        .arg_ptr(input)
205        .arg_ptr(weight.weight)
206        .arg_ptr(weight.weight_scale)
207        .arg_f32(weight.weight_scale_2)
208        .arg_ptr(output)
209        .arg_u32(m)
210        .arg_u32(n)
211        .arg_u32(k)
212        .launch(stream)
213}
214
215#[allow(clippy::too_many_arguments)]
216/// 2026-09-25: [`w4a16_gemm_n128_m128_bf16`] with an explicit transposed-B
217/// row stride, for the same reason as [`w4a16_gemm_n128_ldb`]. This is the
218/// launcher for the 9-parameter `w4a16_gemm_t_m128_bf16_v2`.
219pub fn w4a16_gemm_n128_m128_bf16_ldb(
220    gpu: &dyn GpuBackend,
221    kernel: KernelHandle,
222    input: DevicePtr,
223    weight: &QuantizedWeight,
224    output: DevicePtr,
225    m: u32,
226    n: u32,
227    k: u32,
228    ldb: u32,
229    stream: u64,
230) -> Result<()> {
231    check_e4m3_activation_range(gpu, kernel, input, m, k, stream)?;
232    KernelLaunch::new(gpu, kernel)
233        .grid([n_tile_blocks(gpu, kernel, n)?, div_ceil(m, 128), 1])
234        .block([128, 1, 1])
235        .arg_ptr(input)
236        .arg_ptr(weight.weight)
237        .arg_ptr(weight.weight_scale)
238        .arg_f32(weight.weight_scale_2)
239        .arg_ptr(output)
240        .arg_u32(m)
241        .arg_u32(n)
242        .arg_u32(k)
243        .arg_u32(ldb)
244        .launch(stream)
245}
246
247/// 2026-09-25: W4A16 GEMM launching `w4a16_gemm_t_m128_bf16`, which runs BF16
248/// m16n8k16 MMAs, on the grid, block and transposed NVFP4 weight layout of
249/// [`w4a16_gemm_n128_m128`]. It passes 8 arguments, so it fits only the
250/// 8-parameter v1 kernel; `w4a16_gemm_t_m128_bf16_v2` takes a 9th (`ldb`)
251/// and goes through [`w4a16_gemm_n128_m128_bf16_ldb`] (`ldb = N` for an
252/// unpadded twin).
253pub fn w4a16_gemm_n128_m128_bf16(
254    gpu: &dyn GpuBackend,
255    kernel: KernelHandle,
256    input: DevicePtr,
257    weight: &QuantizedWeight,
258    output: DevicePtr,
259    m: u32,
260    n: u32,
261    k: u32,
262    stream: u64,
263) -> Result<()> {
264    check_e4m3_activation_range(gpu, kernel, input, m, k, stream)?;
265    KernelLaunch::new(gpu, kernel)
266        .grid([n_tile_blocks(gpu, kernel, n)?, div_ceil(m, 128), 1])
267        .block([128, 1, 1])
268        .arg_ptr(input)
269        .arg_ptr(weight.weight)
270        .arg_ptr(weight.weight_scale)
271        .arg_f32(weight.weight_scale_2)
272        .arg_ptr(output)
273        .arg_u32(m)
274        .arg_u32(n)
275        .arg_u32(k)
276        .launch(stream)
277}
278
279#[cfg(test)]
280#[path = "gemm_dense_tests.rs"]
281mod gemm_dense_tests;
metrale_model_layers/layers/ops/gemm_dense.rs

metrale_model_layers/layers/ops/
gemm_dense.rs
