spark_model/layers/ops/w8a16_gemm_m16.rs
1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Tensor-core W8A16 decode GEMM with a 16-row M tile (#927).
4//!
5//! WHY. `w8a16_gemv_batch16` streams the FP8 weight once for up to 16 rows and
6//! is bit-exact, but at M=16 it is FP32-FMA-bound, not bandwidth-bound: on
7//! 1xH100 with Qwen/Qwen3.8-27B-FP8 it measured 0.260 ms / **342 GB/s** for
8//! gate/up (N=17408, K=5120) and 0.330 ms / **270 GB/s** for down (N=5120,
9//! K=17408), against ~3,000 GB/s of HBM3 — an 89 MB weight matrix should
10//! stream in ~30 us. Its inner loop spends ~37 ALU ops per weight BYTE (16
11//! scalar FFMA, 16 BF16->FP32 converts, a LUT lookup, a scale multiply);
12//! `w8a16_gemm_m16` turns those 16 FFMA into one `mma.sync.m16n8k16` lane-slot
13//! and the dequant into ~2 instructions per byte, so the shape becomes
14//! weight-bandwidth bound.
15//!
16//! NUMERICS — REASSOCIATED ON PURPOSE. The MMA reduces 16 K-products in the
17//! tensor core's own order before they reach the FP32 accumulator, so this is
18//! NOT bit-identical to the scalar `w8a16_gemv` the way the batched GEMVs are.
19//! The contract is <= 2 BF16 ULP per element (oracle:
20//! `examples/native_fp8_ffn_m16_tc_microtest.rs`). This is not a new seam: the
21//! arm the FFN used at these widths BEFORE #927 (`w8a16_gemm_n128_m128` /
22//! `w8a16_gemm_pipelined`) reassociates identically. It is why every call site
23//! sits behind `ATLAS_FFN_M16_TC`, default OFF.
24//!
25//! The 128-K block scale is folded ONCE per block onto an FP32 outer
26//! accumulator (two-level fold, preserved exactly from `w8a16_gemm_pipelined`),
27//! never per element and never into BF16.
28//!
29//! Kernels: `w8a16_gemm_m16` / `w8a16_gemm_m16_strided` (module
30//! `w8a16_gemm_m16`). Grid: (ceil(N/32), 1, 1) Block: (128, 1, 1).
31
32use anyhow::{Result, ensure};
33use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
34use spark_runtime::kernel_args::{KernelLaunch, div_ceil};
35
36/// N columns one CTA owns on the DEFAULT instantiation. SSOT for the launch
37/// geometry AND for the dispatch rule's "does this shape have enough CTAs"
38/// reasoning, so the two cannot drift: the kernel's `M16_N_TILE` must equal
39/// this.
40pub const W8A16_GEMM_M16_N_TILE: u32 = 32;
41
42/// The wide instantiation's N tile (`w8a16_gemm_m16_n64`, kernel
43/// `M16_N_TILE_WIDE`) — opt-in via `ATLAS_FFN_M16_TC_NTILE=64`. Halves the CTA
44/// count for a given N and doubles the reuse of each staged A fragment. WHY it
45/// exists and what it is meant to settle: `dense_ffn_m16_tc.rs`.
46pub const W8A16_GEMM_M16_N_TILE_WIDE: u32 = 64;
47
48/// The shared shape of [`w8a16_gemm_m16`], so a caller that picks between it
49/// and `w8a16_gemv_batch16` can hold one function pointer.
50pub type ContiguousM16Gemm = fn(
51 &dyn GpuBackend,
52 KernelHandle,
53 DevicePtr,
54 DevicePtr,
55 DevicePtr,
56 DevicePtr,
57 u32,
58 u32,
59 u32,
60 u64,
61) -> Result<()>;
62
63/// Contiguous `input` `[m, k]` BF16 and `output` `[m, n]` BF16; `weight` /
64/// `block_scale` are the raw `w8a16_gemv` pointers (`[N, K]` FP8 E4M3 and
65/// `[N/128, K/128]` FP32).
66///
67/// REFUSES m > 16: the kernel's M tile IS the MMA's 16 rows, and rows past it
68/// are simply not computed (stale output, not a launch failure). Callers with
69/// 17..=32 rows run it twice on contiguous row halves, the way
70/// `dense_ffn_m16_tc.rs` does.
71#[allow(clippy::too_many_arguments)]
72pub fn w8a16_gemm_m16(
73 gpu: &dyn GpuBackend,
74 kernel: KernelHandle,
75 input: DevicePtr,
76 weight: DevicePtr,
77 block_scale: DevicePtr,
78 output: DevicePtr,
79 m: u32,
80 n: u32,
81 k: u32,
82 stream: u64,
83) -> Result<()> {
84 ensure!(
85 (1..=16).contains(&m),
86 "w8a16_gemm_m16: m={m} outside 1..=16 (kernel M tile)"
87 );
88 ensure!(
89 k.is_multiple_of(128),
90 "w8a16_gemm_m16: K={k} not a multiple of 128 (block-scale granularity)"
91 );
92 launch_contiguous(
93 gpu,
94 kernel,
95 W8A16_GEMM_M16_N_TILE,
96 input,
97 weight,
98 block_scale,
99 output,
100 m,
101 n,
102 k,
103 stream,
104 )
105}
106
107/// `N_TILE=64` twin of [`w8a16_gemm_m16`] — identical arguments and identical
108/// per-output arithmetic, `ceil(N/64)` CTAs instead of `ceil(N/32)`. Same
109/// signature, so a caller holds ONE [`ContiguousM16Gemm`] pointer and the tile
110/// is a dispatch choice, not a code path.
111#[allow(clippy::too_many_arguments)]
112pub fn w8a16_gemm_m16_n64(
113 gpu: &dyn GpuBackend,
114 kernel: KernelHandle,
115 input: DevicePtr,
116 weight: DevicePtr,
117 block_scale: DevicePtr,
118 output: DevicePtr,
119 m: u32,
120 n: u32,
121 k: u32,
122 stream: u64,
123) -> Result<()> {
124 ensure!(
125 (1..=16).contains(&m),
126 "w8a16_gemm_m16_n64: m={m} outside 1..=16 (kernel M tile)"
127 );
128 ensure!(
129 k.is_multiple_of(128),
130 "w8a16_gemm_m16_n64: K={k} not a multiple of 128 (block-scale granularity)"
131 );
132 launch_contiguous(
133 gpu,
134 kernel,
135 W8A16_GEMM_M16_N_TILE_WIDE,
136 input,
137 weight,
138 block_scale,
139 output,
140 m,
141 n,
142 k,
143 stream,
144 )
145}
146
147/// The launch both contiguous instantiations share — only the CTA width
148/// differs, and it is the ONE thing a reader has to check to tell them apart.
149#[allow(clippy::too_many_arguments)]
150fn launch_contiguous(
151 gpu: &dyn GpuBackend,
152 kernel: KernelHandle,
153 n_tile: u32,
154 input: DevicePtr,
155 weight: DevicePtr,
156 block_scale: DevicePtr,
157 output: DevicePtr,
158 m: u32,
159 n: u32,
160 k: u32,
161 stream: u64,
162) -> Result<()> {
163 KernelLaunch::new(gpu, kernel)
164 .grid([div_ceil(n, n_tile), 1, 1])
165 .block([128, 1, 1])
166 .arg_ptr(input)
167 .arg_ptr(weight)
168 .arg_ptr(block_scale)
169 .arg_ptr(output)
170 .arg_u32(m)
171 .arg_u32(n)
172 .arg_u32(k)
173 .launch(stream)
174}
175
176/// Strided sibling of [`w8a16_gemm_m16`]: `a_row_stride` / `c_row_stride` are
177/// the A and C row pitches in ELEMENTS, for callers whose rows are not
178/// contiguous. The multi-seq decode QKV buffer is `[n, per_seq_qkv]` with Q at
179/// offset 0, K after Q and V after K inside every row, so one launch per
180/// projection writes all `m` rows straight into their slots — the same reason
181/// `w8a16_gemv_batch16_strided` exists, and the same argument order.
182///
183/// `a_row_stride` must keep each activation row 16-byte aligned (a multiple of
184/// 8 BF16): the kernel stages A with 16-byte `cp.async` chunks.
185#[allow(clippy::too_many_arguments)]
186pub fn w8a16_gemm_m16_strided(
187 gpu: &dyn GpuBackend,
188 kernel: KernelHandle,
189 input: DevicePtr,
190 weight: DevicePtr,
191 block_scale: DevicePtr,
192 output: DevicePtr,
193 m: u32,
194 n: u32,
195 k: u32,
196 a_row_stride: u32,
197 c_row_stride: u32,
198 stream: u64,
199) -> Result<()> {
200 ensure!(
201 (1..=16).contains(&m),
202 "w8a16_gemm_m16_strided: m={m} outside 1..=16 (kernel M tile)"
203 );
204 ensure!(
205 k.is_multiple_of(128),
206 "w8a16_gemm_m16_strided: K={k} not a multiple of 128 (block-scale granularity)"
207 );
208 ensure!(
209 a_row_stride >= k && c_row_stride >= n,
210 "w8a16_gemm_m16_strided: row pitches (a={a_row_stride}, c={c_row_stride}) \
211 must cover the used extents (k={k}, n={n})"
212 );
213 ensure!(
214 a_row_stride.is_multiple_of(8),
215 "w8a16_gemm_m16_strided: a_row_stride={a_row_stride} must keep rows \
216 16B-aligned (cp.async stages A in 16-byte chunks)"
217 );
218 KernelLaunch::new(gpu, kernel)
219 .grid([div_ceil(n, W8A16_GEMM_M16_N_TILE), 1, 1])
220 .block([128, 1, 1])
221 .arg_ptr(input)
222 .arg_ptr(weight)
223 .arg_ptr(block_scale)
224 .arg_ptr(output)
225 .arg_u32(m)
226 .arg_u32(n)
227 .arg_u32(k)
228 .arg_u32(a_row_stride)
229 .arg_u32(c_row_stride)
230 .launch(stream)
231}