metrale_model_layers/layers/ops/
w8a16_gemm_m16.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the tensor-core W8A16 decode GEMM with a 16-row M
4//! tile (`kernels/hopper/common/w8a16_gemm_m16.cu`, module `w8a16_gemm_m16`;
5//! entries `w8a16_gemm_m16`, `w8a16_gemm_m16_n64`, `w8a16_gemm_m16_strided`).
6//!
7//! Where `w8a16_gemv_batch16` runs 16 scalar FFMA per weight byte, this kernel
8//! runs `mma.sync.m16n8k16`, which reduces 16 K-products in the tensor core's
9//! own order. It is therefore not bit-identical to the scalar GEMVs. Its
10//! contract is `within_m16_tc_budget` (`dense_ffn_m16_tc_oracle.rs`): within
11//! `M16_TC_MAX_ULP` BF16 ULP of the reference, or an absolute error under
12//! `m16_tc_acc_floor` for the reduction depth and the row's RMS; the
13//! `native_fp8_ffn_m16_tc_microtest` example applies it. The callers gate it
14//! on the target-declared `ffn_m16_tc` (dense FFN) and `attn_m16_tc`
15//! (attention) levers.
16//!
17//! The 128-K block scale is folded once per block onto an FP32 outer
18//! accumulator, never per element and never into BF16.
19//!
20//! Grid: (ceil(N/N_TILE), 1, 1)  Block: (128, 1, 1).
21//!
22//! Owner: model-layers ops.
23//! Invariants: none beyond the types.
24
25use anyhow::{Result, ensure};
26use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
27use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
28
29/// 2026-09-25: N columns one CTA owns on the default instantiation. The launch
30/// grid and `m16_tc_kernel` (`dense_ffn_m16_tc.rs`) both use it; the kernel's
31/// `M16_N_TILE` must equal it.
32pub const W8A16_GEMM_M16_N_TILE: u32 = 32;
33
34/// 2026-09-25: The wide instantiation's N tile (`w8a16_gemm_m16_n64`, kernel
35/// `M16_N_TILE_WIDE`), selected for the dense FFN by
36/// `METRALE_FFN_M16_TC_NTILE=64`. It halves the CTA count for a given N.
37pub const W8A16_GEMM_M16_N_TILE_WIDE: u32 = 64;
38
39/// 2026-09-25: The shared signature of [`w8a16_gemm_m16`] and
40/// [`w8a16_gemm_m16_n64`], so `m16_tc_kernel` (`dense_ffn_m16_tc.rs`) can
41/// return either as one function pointer.
42pub type ContiguousM16Gemm = fn(
43    &dyn GpuBackend,
44    KernelHandle,
45    DevicePtr,
46    DevicePtr,
47    DevicePtr,
48    DevicePtr,
49    u32,
50    u32,
51    u32,
52    u64,
53) -> Result<()>;
54
55/// 2026-09-25: Contiguous `input` `[m, k]` BF16 and `output` `[m, n]` BF16;
56/// `weight` / `block_scale` are the raw `w8a16_gemv` pointers (`[N, K]` FP8
57/// E4M3 and `[N/128, K/128]` FP32).
58///
59/// Returns an error for `m` outside `1..=16` or `k` not a multiple of 128: the
60/// kernel's M tile is the MMA's 16 rows, and rows past it would be left
61/// unwritten rather than fail. Callers with 17..=32 rows run it twice on
62/// contiguous row halves, as `dense_ffn_m16_tc.rs` does.
63#[allow(clippy::too_many_arguments)]
64pub fn w8a16_gemm_m16(
65    gpu: &dyn GpuBackend,
66    kernel: KernelHandle,
67    input: DevicePtr,
68    weight: DevicePtr,
69    block_scale: DevicePtr,
70    output: DevicePtr,
71    m: u32,
72    n: u32,
73    k: u32,
74    stream: u64,
75) -> Result<()> {
76    ensure!(
77        (1..=16).contains(&m),
78        "w8a16_gemm_m16: m={m} outside 1..=16 (kernel M tile)"
79    );
80    ensure!(
81        k.is_multiple_of(128),
82        "w8a16_gemm_m16: K={k} not a multiple of 128 (block-scale granularity)"
83    );
84    launch_contiguous(
85        gpu,
86        kernel,
87        W8A16_GEMM_M16_N_TILE,
88        input,
89        weight,
90        block_scale,
91        output,
92        m,
93        n,
94        k,
95        stream,
96    )
97}
98
99/// 2026-09-25: The `N_TILE=64` instantiation of [`w8a16_gemm_m16`]: the same
100/// arguments and guards, `ceil(N/64)` CTAs instead of `ceil(N/32)`.
101#[allow(clippy::too_many_arguments)]
102pub fn w8a16_gemm_m16_n64(
103    gpu: &dyn GpuBackend,
104    kernel: KernelHandle,
105    input: DevicePtr,
106    weight: DevicePtr,
107    block_scale: DevicePtr,
108    output: DevicePtr,
109    m: u32,
110    n: u32,
111    k: u32,
112    stream: u64,
113) -> Result<()> {
114    ensure!(
115        (1..=16).contains(&m),
116        "w8a16_gemm_m16_n64: m={m} outside 1..=16 (kernel M tile)"
117    );
118    ensure!(
119        k.is_multiple_of(128),
120        "w8a16_gemm_m16_n64: K={k} not a multiple of 128 (block-scale granularity)"
121    );
122    launch_contiguous(
123        gpu,
124        kernel,
125        W8A16_GEMM_M16_N_TILE_WIDE,
126        input,
127        weight,
128        block_scale,
129        output,
130        m,
131        n,
132        k,
133        stream,
134    )
135}
136
137/// 2026-09-25: The launch both contiguous instantiations share; only the CTA
138/// width differs.
139#[allow(clippy::too_many_arguments)]
140fn launch_contiguous(
141    gpu: &dyn GpuBackend,
142    kernel: KernelHandle,
143    n_tile: u32,
144    input: DevicePtr,
145    weight: DevicePtr,
146    block_scale: DevicePtr,
147    output: DevicePtr,
148    m: u32,
149    n: u32,
150    k: u32,
151    stream: u64,
152) -> Result<()> {
153    KernelLaunch::new(gpu, kernel)
154        .grid([div_ceil(n, n_tile), 1, 1])
155        .block([128, 1, 1])
156        .arg_ptr(input)
157        .arg_ptr(weight)
158        .arg_ptr(block_scale)
159        .arg_ptr(output)
160        .arg_u32(m)
161        .arg_u32(n)
162        .arg_u32(k)
163        .launch(stream)
164}
165
166/// 2026-09-25: Strided [`w8a16_gemm_m16`]: `a_row_stride` / `c_row_stride` are
167/// the A and C row pitches in elements, for callers whose rows are not
168/// contiguous. The multi-seq decode QKV buffer (`qkv_fp8_batch.rs`) is
169/// `[n, per_seq_qkv]` with Q at offset 0, K after Q and V after K in every row,
170/// so one launch per projection writes all `m` rows into their slots. Same
171/// argument order as `w8a16_gemv_batch16_strided`.
172///
173/// `a_row_stride` must keep each activation row 16-byte aligned (a multiple of
174/// 8 BF16): the kernel stages A in 16-byte `cp.async` chunks.
175#[allow(clippy::too_many_arguments)]
176pub fn w8a16_gemm_m16_strided(
177    gpu: &dyn GpuBackend,
178    kernel: KernelHandle,
179    input: DevicePtr,
180    weight: DevicePtr,
181    block_scale: DevicePtr,
182    output: DevicePtr,
183    m: u32,
184    n: u32,
185    k: u32,
186    a_row_stride: u32,
187    c_row_stride: u32,
188    stream: u64,
189) -> Result<()> {
190    ensure!(
191        (1..=16).contains(&m),
192        "w8a16_gemm_m16_strided: m={m} outside 1..=16 (kernel M tile)"
193    );
194    ensure!(
195        k.is_multiple_of(128),
196        "w8a16_gemm_m16_strided: K={k} not a multiple of 128 (block-scale granularity)"
197    );
198    ensure!(
199        a_row_stride >= k && c_row_stride >= n,
200        "w8a16_gemm_m16_strided: row pitches (a={a_row_stride}, c={c_row_stride}) \
201         must cover the used extents (k={k}, n={n})"
202    );
203    ensure!(
204        a_row_stride.is_multiple_of(8),
205        "w8a16_gemm_m16_strided: a_row_stride={a_row_stride} must keep rows \
206         16B-aligned (cp.async stages A in 16-byte chunks)"
207    );
208    KernelLaunch::new(gpu, kernel)
209        .grid([div_ceil(n, W8A16_GEMM_M16_N_TILE), 1, 1])
210        .block([128, 1, 1])
211        .arg_ptr(input)
212        .arg_ptr(weight)
213        .arg_ptr(block_scale)
214        .arg_ptr(output)
215        .arg_u32(m)
216        .arg_u32(n)
217        .arg_u32(k)
218        .arg_u32(a_row_stride)
219        .arg_u32(c_row_stride)
220        .launch(stream)
221}
metrale_model_layers/layers/ops/w8a16_gemm_m16.rs

metrale_model_layers/layers/ops/
w8a16_gemm_m16.rs
