metrale_model_layers/layers/ops/
gemm_dense_bf16.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the dense BF16 GEMMs: scalar, tensor-core,
4//! scaled-accumulate, split-K, router, pipelined and prefill. The W4A16
5//! launchers are in `gemm_dense.rs`, which the arity tests read.
6//!
7//! Owner: model-layers ops.
8//! Invariants: none beyond the types.
9
10use super::*;
11
12/// 2026-09-25: Tensor-core BF16 GEMM, `C = A @ B^T`, with m16n8k16 MMAs
13/// (kernel `dense_gemm_tc`).
14pub fn dense_gemm_tc(
15    gpu: &dyn GpuBackend,
16    kernel: KernelHandle,
17    input: DevicePtr,
18    weight: &DenseWeight,
19    output: DevicePtr,
20    m: u32,
21    n: u32,
22    k: u32,
23    stream: u64,
24) -> Result<()> {
25    KernelLaunch::new(gpu, kernel)
26        .grid([div_ceil(n, 64), div_ceil(m, 16), 1])
27        .block([128, 1, 1])
28        .arg_ptr(input)
29        .arg_ptr(weight.weight)
30        .arg_ptr(output)
31        .arg_u32(m)
32        .arg_u32(n)
33        .arg_u32(k)
34        .launch(stream)
35}
36
37/// 2026-09-25: `output[m, n] += scale * bf16(input[m, k] @ weight[n, k]^T)` in
38/// one pass: [`dense_gemm_tc`] with an accumulating epilogue, used for the
39/// LoRA expand+fold (`lora_delta.rs`). It saves the `[m, n]` scratch that a
40/// GEMM followed by `bf16_scaled_add` writes and reads back.
41///
42/// Same result as that pair: the kernel rounds each dot product to BF16
43/// before it applies `scale` and adds, as `bf16_scaled_add` does with a BF16
44/// scratch.
45#[allow(clippy::too_many_arguments)]
46pub fn dense_gemm_tc_scaled_acc(
47    gpu: &dyn GpuBackend,
48    kernel: KernelHandle,
49    input: DevicePtr,
50    weight: &DenseWeight,
51    output: DevicePtr,
52    m: u32,
53    n: u32,
54    k: u32,
55    scale: f32,
56    stream: u64,
57) -> Result<()> {
58    KernelLaunch::new(gpu, kernel)
59        .grid([div_ceil(n, 64), div_ceil(m, 16), 1])
60        .block([128, 1, 1])
61        .arg_ptr(input)
62        .arg_ptr(weight.weight)
63        .arg_ptr(output)
64        .arg_u32(m)
65        .arg_u32(n)
66        .arg_u32(k)
67        .arg_f32(scale)
68        .launch(stream)
69}
70
71/// 2026-09-25: Split-K BF16 GEMM: `partial_kernel` writes FP32 partial
72/// products for `k_splits` K chunks into `workspace`, and `reduce_kernel` sums
73/// them into BF16 `output`. `workspace` holds `k_splits * m * n` FP32 values.
74#[allow(clippy::too_many_arguments)]
75pub fn dense_gemm_splitk(
76    gpu: &dyn GpuBackend,
77    partial_kernel: KernelHandle,
78    reduce_kernel: KernelHandle,
79    input: DevicePtr,
80    weight: &DenseWeight,
81    output: DevicePtr,
82    workspace: DevicePtr,
83    m: u32,
84    n: u32,
85    k: u32,
86    k_splits: u32,
87    stream: u64,
88) -> Result<()> {
89    KernelLaunch::new(gpu, partial_kernel)
90        .grid([div_ceil(n, 16), div_ceil(m, 16), k_splits])
91        .block([16, 16, 1])
92        .arg_ptr(input)
93        .arg_ptr(weight.weight)
94        .arg_ptr(workspace)
95        .arg_u32(m)
96        .arg_u32(n)
97        .arg_u32(k)
98        .arg_u32(k_splits)
99        .launch(stream)?;
100    KernelLaunch::new(gpu, reduce_kernel)
101        .grid([div_ceil(n, 256), m, 1])
102        .block([256, 1, 1])
103        .arg_ptr(workspace)
104        .arg_ptr(output)
105        .arg_u32(m)
106        .arg_u32(n)
107        .arg_u32(k_splits)
108        .launch(stream)
109}
110
111/// 2026-09-25: Scalar BF16 GEMM, `C = A @ B^T` (kernel
112/// `dense_gemm_bf16(A, B, C, M, N, K)`): A `[M, K]` row-major, B `[N, K]`
113/// row-major (HuggingFace layout), C `[M, N]` row-major.
114pub fn dense_gemm(
115    gpu: &dyn GpuBackend,
116    kernel: KernelHandle,
117    input: DevicePtr,
118    weight: &DenseWeight,
119    output: DevicePtr,
120    m: u32,
121    n: u32,
122    k: u32,
123    stream: u64,
124) -> Result<()> {
125    KernelLaunch::new(gpu, kernel)
126        .grid([div_ceil(n, 16), div_ceil(m, 16), 1])
127        .block([16, 16, 1])
128        .arg_ptr(input)
129        .arg_ptr(weight.weight)
130        .arg_ptr(output)
131        .arg_u32(m)
132        .arg_u32(n)
133        .arg_u32(k)
134        .launch(stream)
135}
136
137/// 2026-09-26: Output columns per block of `moe_router_gemm_bf16`
138/// (`RC_COLS` in `kernels/gb10/common/moe_router_gemm.cu`).
139pub const MOE_ROUTER_GEMM_COLS: u32 = 4;
140/// 2026-09-26: Activation rows per block of `moe_router_gemm_bf16` (`RC_ROWS`).
141pub const MOE_ROUTER_GEMM_ROWS: u32 = 16;
142
143/// 2026-09-26: The MoE router gate GEMM, `C = A @ B^T` with the operands of
144/// [`dense_gemm`] and its bits (kernel `moe_router_gemm_bf16`): each output
145/// keeps one FP32 accumulator over k = 0..K-1 in order, but a block stages a
146/// 512-wide K slice of 4 weight rows and up to 16 activation rows in shared
147/// memory at once instead of waiting on every 16-wide tile. `k` must be a
148/// multiple of 16 (below that [`dense_gemm`] adds zero padding terms).
149#[allow(clippy::too_many_arguments)]
150pub fn moe_router_gemm(
151    gpu: &dyn GpuBackend,
152    kernel: KernelHandle,
153    input: DevicePtr,
154    weight: &DenseWeight,
155    output: DevicePtr,
156    m: u32,
157    n: u32,
158    k: u32,
159    stream: u64,
160) -> Result<()> {
161    anyhow::ensure!(
162        k.is_multiple_of(16),
163        "moe_router_gemm: k={k} is not a multiple of 16"
164    );
165    KernelLaunch::new(gpu, kernel)
166        .grid([
167            div_ceil(n, MOE_ROUTER_GEMM_COLS),
168            div_ceil(m, MOE_ROUTER_GEMM_ROWS),
169            1,
170        ])
171        .block([256, 1, 1])
172        .arg_ptr(input)
173        .arg_ptr(weight.weight)
174        .arg_ptr(output)
175        .arg_u32(m)
176        .arg_u32(n)
177        .arg_u32(k)
178        .launch(stream)
179}
180
181/// 2026-09-25: Register-blocked BF16 GEMM (kernel `dense_gemm_bf16_router`)
182/// that keeps the scalar `dense_gemm_bf16`'s per-output FP32 accumulation
183/// order, strict k = 0..K-1, and with it the scalar kernel's results under the
184/// `--fmad=false` build (`kernels/gb10/common/KERNEL.toml`).
185/// `router_gate_gemm_dense` uses it for the MoE router gate GEMM.
186pub fn dense_gemm_router(
187    gpu: &dyn GpuBackend,
188    kernel: KernelHandle,
189    input: DevicePtr,
190    weight: &DenseWeight,
191    output: DevicePtr,
192    m: u32,
193    n: u32,
194    k: u32,
195    stream: u64,
196) -> Result<()> {
197    KernelLaunch::new(gpu, kernel)
198        .grid([div_ceil(n, 64), div_ceil(m, 16), 1])
199        .block([16, 16, 1])
200        .arg_ptr(input)
201        .arg_ptr(weight.weight)
202        .arg_ptr(output)
203        .arg_u32(m)
204        .arg_u32(n)
205        .arg_u32(k)
206        .launch(stream)
207}
208
209/// 2026-10-01: Most rows for which the MoE prefill router takes [`moe_router_gemm`] (the
210/// decode routers' kernel, the same bits as [`dense_gemm_router`]); the decode routers run it
211/// at 2 to 64 rows.
212pub const MOE_ROUTER_GEMM_MAX_PREFILL_ROWS: u32 = 64;
213
214/// 2026-09-28: Fewest rows for which [`moe_router_gemm_rt`] beats [`dense_gemm_router`] on
215/// GB10 (measured: 0.82x at 300 rows, 2.5x at 1000, 3.0-3.5x at 4k-32k).
216pub const MOE_ROUTER_RT_MIN_ROWS: u32 = 1024;
217
218/// 2026-09-28: [`dense_gemm_router`] on 64x128 tiles with a 4x8 register tile per thread
219/// (kernel `moe_router_gemm_rt`, `kernels/gb10/common/moe_router_gemm_prefill.cu`): the
220/// same per-output FP32 accumulation in ascending k, so the same bits. Needs `k % 16 == 0`.
221#[allow(clippy::too_many_arguments)]
222pub fn moe_router_gemm_rt(
223    gpu: &dyn GpuBackend,
224    kernel: KernelHandle,
225    input: DevicePtr,
226    weight: &DenseWeight,
227    output: DevicePtr,
228    m: u32,
229    n: u32,
230    k: u32,
231    stream: u64,
232) -> Result<()> {
233    anyhow::ensure!(
234        k.is_multiple_of(16),
235        "moe_router_gemm_rt: K={k} is not a multiple of 16"
236    );
237    KernelLaunch::new(gpu, kernel)
238        .grid([div_ceil(n, 128), div_ceil(m, 64), 1])
239        .block([256, 1, 1])
240        .arg_ptr(input)
241        .arg_ptr(weight.weight)
242        .arg_ptr(output)
243        .arg_u32(m)
244        .arg_u32(n)
245        .arg_u32(k)
246        .launch(stream)
247}
248
249/// 2026-09-25: Tensor-core BF16 GEMM (kernel `dense_gemm_bf16_pipelined`):
250/// m16n8k16 MMAs fed by a 2-stage `cp.async` pipeline on a 128x128 tile, with
251/// the inputs and output layout of [`dense_gemm`].
252#[allow(clippy::too_many_arguments)]
253pub fn dense_gemm_bf16_pipelined(
254    gpu: &dyn GpuBackend,
255    kernel: KernelHandle,
256    input: DevicePtr,
257    weight: &DenseWeight,
258    output: DevicePtr,
259    m: u32,
260    n: u32,
261    k: u32,
262    stream: u64,
263) -> Result<()> {
264    KernelLaunch::new(gpu, kernel)
265        .grid([div_ceil(n, 128), div_ceil(m, 128), 1])
266        .block([256, 1, 1])
267        .arg_ptr(input)
268        .arg_ptr(weight.weight)
269        .arg_ptr(output)
270        .arg_u32(m)
271        .arg_u32(n)
272        .arg_u32(k)
273        .launch(stream)
274}
275
276/// 2026-09-25: Dense BF16 prefill GEMM: [`dense_gemm_bf16_pipelined`] when
277/// `pipelined_kernel` is loaded, otherwise [`dense_gemm`] with
278/// `fallback_kernel`.
279#[allow(clippy::too_many_arguments)]
280pub fn dense_gemm_prefill(
281    gpu: &dyn GpuBackend,
282    fallback_kernel: KernelHandle,
283    pipelined_kernel: KernelHandle,
284    input: DevicePtr,
285    weight: &DenseWeight,
286    output: DevicePtr,
287    m: u32,
288    n: u32,
289    k: u32,
290    stream: u64,
291) -> Result<()> {
292    if pipelined_kernel.0 != 0 {
293        dense_gemm_bf16_pipelined(
294            gpu,
295            pipelined_kernel,
296            input,
297            weight,
298            output,
299            m,
300            n,
301            k,
302            stream,
303        )
304    } else {
305        dense_gemm(gpu, fallback_kernel, input, weight, output, m, n, k, stream)
306    }
307}
metrale_model_layers/layers/ops/gemm_dense_bf16.rs

metrale_model_layers/layers/ops/
gemm_dense_bf16.rs
