spark_model/layers/ops/
dense_gemm_m16_bf16.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Tensor-core DENSE BF16 decode GEMM with a 16-row M tile — the BF16 LM-head
4//! arm (#927/#928).
5//!
6//! WHY. nsys on 1xH100 (round 7, 2026-09-11, Qwen/Qwen3.8-27B-FP8 with
7//! `--lm-head-dtype bf16` and `ATLAS_LM_HEAD_BATCHM_MAX=16`, decode batch 16)
8//! puts `dense_gemv_bf16_batchm` — the LM head — at **3,571 µs in ONE launch,
9//! 8.19% of the 43.6 ms step**, for a single pass over the 2.54 GB
10//! `[248077, 5120]` BF16 vocab weight. That is **~710 GB/s** against ~3,350
11//! GB/s of HBM3. The SAME kernel at C=1 costs 798 µs — 3.2 TB/s-class, at the
12//! roofline — so 16 rows cost 4.47x one row for an identical weight read: the
13//! batched tier is FP32-FMA-bound, not bandwidth-bound.
14//!
15//! `dense_gemm_m16_bf16` replaces the per-row scalar FFMA with one
16//! `mma.sync.m16n8k16` lane slot. The M tile IS 16 rows, so nothing is padded —
17//! the difference from `dense_gemm_tc`'s 16Mx64N tile (which pads the same way
18//! but still reads the weight through a scalar inner loop) and from the tile
19//! GEMMs that pad M to 128. Target: **<= 1.3 ms at M=16 = >= 1,950 GB/s**.
20//!
21//! NUMERICS — REASSOCIATED ON PURPOSE. `dense_gemv_bf16` / `_batchm` reduce
22//! each output in ONE FP32 accumulator in strict K order, and the batched tier
23//! is bit-identical to M serial scalar GEMVs. An MMA reduces 16 K-products in
24//! the tensor core's own order first, so this is NOT. The contract is the
25//! predicate `w8a16_gemm_m16` already answers to —
26//! [`crate::layers::dense_ffn::m16_tc::within_m16_tc_budget`]: <= 2 ordinal
27//! BF16 ULP, or an absolute error under the accumulation floor. Oracle:
28//! `examples/native_bf16_lm_head_m16_microtest.rs`.
29//!
30//! At the LM head a near-tie argmax flip changes the emitted token, which is
31//! why the head arm is behind `ATLAS_LM_HEAD_M16_TC` and defaults OFF
32//! (`model/trait_impl/lm_head_batched.rs`).
33//!
34//! Unlike `w8a16_gemm_m16` there is no block scale and no dequant: B is already
35//! BF16, so a staged weight word IS a B fragment register, the accumulator is
36//! ONE level, and the K constraint is the 64-wide pipeline step rather than the
37//! 128-wide scale block.
38//!
39//! Kernels: `dense_gemm_m16_bf16` / `dense_gemm_m16_bf16_n64` (module
40//! `dense_gemm_m16_bf16`). Grid: (ceil(N/N_TILE), 1, 1)  Block: (128, 1, 1).
41
42use crate::weight_map::DenseWeight;
43use anyhow::{Result, ensure};
44use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
45use spark_runtime::kernel_args::{KernelLaunch, div_ceil};
46
47/// N columns one CTA owns on the DEFAULT instantiation. SSOT for the launch
48/// geometry AND for the dispatch rule's CTA-count reasoning, so the two cannot
49/// drift: the kernel's `DGM16_N_TILE` must equal this.
50pub const DENSE_GEMM_M16_BF16_N_TILE: u32 = 32;
51
52/// The wide instantiation's N tile (`dense_gemm_m16_bf16_n64`, kernel
53/// `DGM16_N_TILE_WIDE`) — opt-in via `ATLAS_LM_HEAD_M16_TC_NTILE=64`. Halves
54/// the CTA count for a given N and halves the L2 traffic the re-read A tile
55/// costs. WHY it exists: `dense_gemm_m16_bf16.cu`.
56pub const DENSE_GEMM_M16_BF16_N_TILE_WIDE: u32 = 64;
57
58/// The kernel's M tile. Rows past it are simply not computed, so the wrapper
59/// REFUSES rather than writing part of the block and leaving the rest stale.
60pub const DENSE_GEMM_M16_BF16_MAX_M: u32 = 16;
61
62/// K granularity: the 4-stage cp.async pipeline advances 64 elements per step,
63/// and 64 BF16 is also what keeps every 16-byte weight-row chunk aligned.
64pub const DENSE_GEMM_M16_BF16_K_STEP: u32 = 64;
65
66/// The shared shape of both instantiations, so a caller that picks between them
67/// (and between them and `dense_gemv_batchm`) can hold ONE function pointer and
68/// the tile stays a dispatch choice rather than a code path.
69pub type DenseM16Bf16Gemm = fn(
70    &dyn GpuBackend,
71    KernelHandle,
72    DevicePtr,
73    &DenseWeight,
74    DevicePtr,
75    u32,
76    u32,
77    u32,
78    u32,
79    u32,
80    u64,
81) -> Result<()>;
82
83/// The default 32-wide CTA. `input` is `[m, a_row_stride]` BF16 with `k` used,
84/// `weight.weight` is the raw `[n, k]` BF16 checkpoint tensor (NOT copied, NOT
85/// quantized — the same pointer `dense_gemv_batchm` reads), and `output` is
86/// `[m, c_row_stride]` BF16 with `n` used.
87#[allow(clippy::too_many_arguments)]
88pub fn dense_gemm_m16_bf16(
89    gpu: &dyn GpuBackend,
90    kernel: KernelHandle,
91    input: DevicePtr,
92    weight: &DenseWeight,
93    output: DevicePtr,
94    m: u32,
95    n: u32,
96    k: u32,
97    a_row_stride: u32,
98    c_row_stride: u32,
99    stream: u64,
100) -> Result<()> {
101    launch(
102        gpu,
103        kernel,
104        DENSE_GEMM_M16_BF16_N_TILE,
105        "dense_gemm_m16_bf16",
106        input,
107        weight,
108        output,
109        m,
110        n,
111        k,
112        a_row_stride,
113        c_row_stride,
114        stream,
115    )
116}
117
118/// `N_TILE=64` twin of [`dense_gemm_m16_bf16`] — identical arguments and
119/// identical per-output arithmetic, `ceil(n/64)` CTAs instead of `ceil(n/32)`.
120#[allow(clippy::too_many_arguments)]
121pub fn dense_gemm_m16_bf16_n64(
122    gpu: &dyn GpuBackend,
123    kernel: KernelHandle,
124    input: DevicePtr,
125    weight: &DenseWeight,
126    output: DevicePtr,
127    m: u32,
128    n: u32,
129    k: u32,
130    a_row_stride: u32,
131    c_row_stride: u32,
132    stream: u64,
133) -> Result<()> {
134    launch(
135        gpu,
136        kernel,
137        DENSE_GEMM_M16_BF16_N_TILE_WIDE,
138        "dense_gemm_m16_bf16_n64",
139        input,
140        weight,
141        output,
142        m,
143        n,
144        k,
145        a_row_stride,
146        c_row_stride,
147        stream,
148    )
149}
150
151/// The guards and the launch both instantiations share — only the CTA width
152/// differs, and it is the ONE thing a reader has to check to tell them apart.
153#[allow(clippy::too_many_arguments)]
154fn launch(
155    gpu: &dyn GpuBackend,
156    kernel: KernelHandle,
157    n_tile: u32,
158    who: &str,
159    input: DevicePtr,
160    weight: &DenseWeight,
161    output: DevicePtr,
162    m: u32,
163    n: u32,
164    k: u32,
165    a_row_stride: u32,
166    c_row_stride: u32,
167    stream: u64,
168) -> Result<()> {
169    ensure!(
170        (1..=DENSE_GEMM_M16_BF16_MAX_M).contains(&m),
171        "{who}: m={m} outside 1..={DENSE_GEMM_M16_BF16_MAX_M} (kernel M tile; \
172         rows past it are never computed, not a launch failure)"
173    );
174    ensure!(
175        k.is_multiple_of(DENSE_GEMM_M16_BF16_K_STEP),
176        "{who}: K={k} not a multiple of {DENSE_GEMM_M16_BF16_K_STEP} \
177         (cp.async pipeline step, and what keeps each [n, k] weight row 16B-aligned)"
178    );
179    ensure!(
180        a_row_stride >= k && c_row_stride >= n,
181        "{who}: row pitches (a={a_row_stride}, c={c_row_stride}) must cover the \
182         used extents (k={k}, n={n})"
183    );
184    ensure!(
185        a_row_stride.is_multiple_of(8),
186        "{who}: a_row_stride={a_row_stride} must keep rows 16B-aligned \
187         (cp.async stages A in 16-byte chunks)"
188    );
189    KernelLaunch::new(gpu, kernel)
190        .grid([div_ceil(n, n_tile), 1, 1])
191        .block([128, 1, 1])
192        .arg_ptr(input)
193        .arg_ptr(weight.weight)
194        .arg_ptr(output)
195        .arg_u32(m)
196        .arg_u32(n)
197        .arg_u32(k)
198        .arg_u32(a_row_stride)
199        .arg_u32(c_row_stride)
200        .launch(stream)
201}
202
203/// Host simulation of the staging / fragment / store index math against a
204/// reference GEMM in the scalar `dense_gemv_bf16`'s reduction order — the test
205/// that says WHERE a mismatch is, without an H100.
206#[cfg(test)]
207#[path = "dense_gemm_m16_bf16_tests.rs"]
208mod tests;
209
210/// Host simulation of the ACCUMULATION FLOOR — the round-9 LM-head red cell,
211/// why round 6's fixed `2^-20 * rms` constant could not carry this tier, and
212/// the margins the K-aware floor keeps. Split from `tests` for the 500-line
213/// cap.
214#[cfg(test)]
215#[path = "dense_gemm_m16_bf16_floor_tests.rs"]
216mod floor_tests;
spark_model/layers/ops/dense_gemm_m16_bf16.rs

spark_model/layers/ops/
dense_gemm_m16_bf16.rs
