metrale_model_layers/layers/ops/
dense_gemm_m16_bf16.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launch wrappers for the tensor-core BF16 decode GEMM with a 16-row M tile
4//! (`dense_gemm_m16_bf16`, `dense_gemm_m16_bf16_n64`), the BF16 LM head's arm for 5 to 16
5//! rows.
6//!
7//! Owner: model-layers ops (kernels in `kernels/hopper/common/dense_gemm_m16_bf16.cu`).
8//! Invariants: a launch has `1 <= m <= 16`, `k` a multiple of 64, row pitches that cover
9//! `k` and `n`, and `a_row_stride` a multiple of 8; `launch` returns an error otherwise.
10//!
11//! Why it exists: measured 2026-09-11 on 1xH100 (Qwen3.8-27B-FP8 with a BF16 LM head,
12//! decode batch 16), the batched scalar GEMV `dense_gemv_bf16_batchm` took 3,571 µs for
13//! the head, about 710 GB/s, against 798 µs at batch 1, so the batched GEMV was FP32-FMA
14//! bound. This kernel feeds each 16-row tile to `mma.sync.m16n8k16` instead.
15//!
16//! Numerics: an MMA sums K products in the tensor core's own order, so outputs can
17//! differ from the scalar GEMV. The accepted difference is
18//! [`crate::layers::dense_ffn::m16_tc::within_m16_tc_budget`]: at most 2 ordinal BF16
19//! ULP, or an absolute error under the accumulation floor. The GPU check is the
20//! model-arch example `native_bf16_lm_head_m16_microtest`. Whether the head uses this
21//! arm is the target's `lm_head_m16_tc` default (on in `kernels/hopper/HARDWARE.toml`,
22//! off for gb10, b200 and b300), which `METRALE_LM_HEAD_M16_TC` overrides
23//! (`target_defaults.rs`); the head's route is in `model-engine lm_head_batched.rs`.
24//!
25//! There is no block scale and no dequant: B is BF16 already. Grid `ceil(N / N_TILE)`,
26//! block 128.
27
28use crate::weight_map::DenseWeight;
29use anyhow::{Result, ensure};
30use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
31use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
32
33/// 2026-09-25: N columns one CTA owns in the default instantiation. It must equal the
34/// kernel's `DGM16_N_TILE`.
35pub const DENSE_GEMM_M16_BF16_N_TILE: u32 = 32;
36
37/// 2026-09-25: N tile of the wide instantiation `dense_gemm_m16_bf16_n64` (kernel
38/// `DGM16_N_TILE_WIDE`), selected by `METRALE_LM_HEAD_M16_TC_NTILE=64`. It launches half
39/// the CTAs for a given N, so the A tile is re-read from L2 half as many times.
40pub const DENSE_GEMM_M16_BF16_N_TILE_WIDE: u32 = 64;
41
42/// 2026-09-25: The kernel's M tile. Rows past it are not computed, so the wrapper
43/// returns an error rather than leave part of the output stale.
44pub const DENSE_GEMM_M16_BF16_MAX_M: u32 = 16;
45
46/// 2026-09-25: K granularity: the kernel's 4-stage cp.async pipeline advances 64
47/// elements per step (`DGM16_K_STEP`), which also keeps every weight row 16-byte
48/// aligned.
49pub const DENSE_GEMM_M16_BF16_K_STEP: u32 = 64;
50
51/// 2026-09-25: The signature both instantiations share, so a caller can pick one and
52/// hold a single function pointer.
53pub type DenseM16Bf16Gemm = fn(
54    &dyn GpuBackend,
55    KernelHandle,
56    DevicePtr,
57    &DenseWeight,
58    DevicePtr,
59    u32,
60    u32,
61    u32,
62    u32,
63    u32,
64    u64,
65) -> Result<()>;
66
67/// 2026-09-25: The default 32-wide CTA. `input` is `[m, a_row_stride]` BF16 with `k`
68/// used, `weight.weight` is the `[n, k]` BF16 weight as loaded, and `output` is
69/// `[m, c_row_stride]` BF16 with `n` used.
70#[allow(clippy::too_many_arguments)]
71pub fn dense_gemm_m16_bf16(
72    gpu: &dyn GpuBackend,
73    kernel: KernelHandle,
74    input: DevicePtr,
75    weight: &DenseWeight,
76    output: DevicePtr,
77    m: u32,
78    n: u32,
79    k: u32,
80    a_row_stride: u32,
81    c_row_stride: u32,
82    stream: u64,
83) -> Result<()> {
84    launch(
85        gpu,
86        kernel,
87        DENSE_GEMM_M16_BF16_N_TILE,
88        "dense_gemm_m16_bf16",
89        input,
90        weight,
91        output,
92        m,
93        n,
94        k,
95        a_row_stride,
96        c_row_stride,
97        stream,
98    )
99}
100
101/// 2026-09-25: The `N_TILE = 64` instantiation of the same kernel template: the same
102/// arguments as [`dense_gemm_m16_bf16`], `ceil(n / 64)` CTAs instead of `ceil(n / 32)`.
103#[allow(clippy::too_many_arguments)]
104pub fn dense_gemm_m16_bf16_n64(
105    gpu: &dyn GpuBackend,
106    kernel: KernelHandle,
107    input: DevicePtr,
108    weight: &DenseWeight,
109    output: DevicePtr,
110    m: u32,
111    n: u32,
112    k: u32,
113    a_row_stride: u32,
114    c_row_stride: u32,
115    stream: u64,
116) -> Result<()> {
117    launch(
118        gpu,
119        kernel,
120        DENSE_GEMM_M16_BF16_N_TILE_WIDE,
121        "dense_gemm_m16_bf16_n64",
122        input,
123        weight,
124        output,
125        m,
126        n,
127        k,
128        a_row_stride,
129        c_row_stride,
130        stream,
131    )
132}
133
134/// 2026-09-25: The checks and launch both instantiations share; only `n_tile` differs.
135#[allow(clippy::too_many_arguments)]
136fn launch(
137    gpu: &dyn GpuBackend,
138    kernel: KernelHandle,
139    n_tile: u32,
140    who: &str,
141    input: DevicePtr,
142    weight: &DenseWeight,
143    output: DevicePtr,
144    m: u32,
145    n: u32,
146    k: u32,
147    a_row_stride: u32,
148    c_row_stride: u32,
149    stream: u64,
150) -> Result<()> {
151    ensure!(
152        (1..=DENSE_GEMM_M16_BF16_MAX_M).contains(&m),
153        "{who}: m={m} outside 1..={DENSE_GEMM_M16_BF16_MAX_M} (kernel M tile; \
154         rows past it are never computed, not a launch failure)"
155    );
156    ensure!(
157        k.is_multiple_of(DENSE_GEMM_M16_BF16_K_STEP),
158        "{who}: K={k} not a multiple of {DENSE_GEMM_M16_BF16_K_STEP} \
159         (cp.async pipeline step, and what keeps each [n, k] weight row 16B-aligned)"
160    );
161    ensure!(
162        a_row_stride >= k && c_row_stride >= n,
163        "{who}: row pitches (a={a_row_stride}, c={c_row_stride}) must cover the \
164         used extents (k={k}, n={n})"
165    );
166    ensure!(
167        a_row_stride.is_multiple_of(8),
168        "{who}: a_row_stride={a_row_stride} must keep rows 16B-aligned \
169         (cp.async stages A in 16-byte chunks)"
170    );
171    KernelLaunch::new(gpu, kernel)
172        .grid([div_ceil(n, n_tile), 1, 1])
173        .block([128, 1, 1])
174        .arg_ptr(input)
175        .arg_ptr(weight.weight)
176        .arg_ptr(output)
177        .arg_u32(m)
178        .arg_u32(n)
179        .arg_u32(k)
180        .arg_u32(a_row_stride)
181        .arg_u32(c_row_stride)
182        .launch(stream)
183}
184
185/// 2026-09-25: Host simulation of the kernel's staging, fragment and store index maps,
186/// checked against a reference GEMM.
187#[cfg(test)]
188#[path = "dense_gemm_m16_bf16_tests.rs"]
189mod tests;
190
191/// 2026-09-25: Host simulation of the accumulation floor that `within_m16_tc_budget`
192/// applies, and its margins.
193#[cfg(test)]
194#[path = "dense_gemm_m16_bf16_floor_tests.rs"]
195mod floor_tests;
metrale_model_layers/layers/ops/dense_gemm_m16_bf16.rs

metrale_model_layers/layers/ops/
dense_gemm_m16_bf16.rs
