metrale_model_layers/layers/ops/
gemm_quant_w8a16.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the FP8 E4M3 weight x BF16 activation (W8A16) GEMV and GEMMs, row-major and transposed.
4//!
5//! Owner: model-layers ops.
6//! Invariants:
7//! - Every launcher passes the arguments in the order `(A, B, scale, C, [M,] N, K)` that its
8//!   kernel's `extern "C"` signature declares.
9//! - Each grid and block is derived from the tile constants of the kernel it launches; the
10//!   kernel files are named on each function.
11
12use super::*;
13
14/// 2026-09-25: W8A16 GEMV for M=1 decode: `C[1,N] = A[1,K] @ dequant(B[N,K])^T`.
15///
16/// `row_scale` is the 2D block scale `[ceil(N/128), ceil(K/128)]` FP32, one scale per
17/// 128x128 weight block, despite the parameter name. Each 256-thread block computes 4 outputs
18/// with 64 threads per output, so the grid is `ceil(N/4)` (kernels/gb10/common/w8a16_gemv.cu,
19/// and the same geometry in kernels/hopper/common/w8a16_gemv.cu).
20#[allow(clippy::too_many_arguments)]
21pub fn w8a16_gemv(
22    gpu: &dyn GpuBackend,
23    kernel: KernelHandle,
24    input: DevicePtr,
25    weight: DevicePtr,
26    row_scale: DevicePtr,
27    output: DevicePtr,
28    n: u32,
29    k: u32,
30    stream: u64,
31) -> Result<()> {
32    KernelLaunch::new(gpu, kernel)
33        .grid([div_ceil(n, 4), 1, 1])
34        .block([256, 1, 1])
35        .arg_ptr(input)
36        .arg_ptr(weight)
37        .arg_ptr(row_scale)
38        .arg_ptr(output)
39        .arg_u32(n)
40        .arg_u32(k)
41        .launch(stream)
42}
43
44/// 2026-09-25: W8A16 GEMM for M>1: `C[M,N] = A[M,K] @ dequant(B[N,K])^T`, with the 2D block
45/// scale `[ceil(N/128), ceil(K/128)]` FP32.
46#[allow(clippy::too_many_arguments)]
47pub fn w8a16_gemm(
48    gpu: &dyn GpuBackend,
49    kernel: KernelHandle,
50    input: DevicePtr,
51    weight: DevicePtr,
52    block_scale: DevicePtr,
53    output: DevicePtr,
54    m: u32,
55    n: u32,
56    k: u32,
57    stream: u64,
58) -> Result<()> {
59    // 2026-09-25: The geometry is per target because the kernel source is. The native-HIP
60    // kernel (kernels/strix-hip/common/w8a16_gemm.cu) is a 256x128 MxN tile with 512 threads;
61    // every other target builds kernels/gb10/common/w8a16_gemm.cu, a 64x64 tile with 128
62    // threads. Keep both arms equal to their kernel's `M_TILE`/`N_TILE`/`THREADS`.
63    #[cfg(metrale_hip)]
64    let (grid, block) = ([div_ceil(n, 128), div_ceil(m, 256), 1], [512, 1, 1]);
65    #[cfg(not(metrale_hip))]
66    let (grid, block) = ([div_ceil(n, 64), div_ceil(m, 64), 1], [128, 1, 1]);
67    KernelLaunch::new(gpu, kernel)
68        .grid(grid)
69        .block(block)
70        .arg_ptr(input)
71        .arg_ptr(weight)
72        .arg_ptr(block_scale)
73        .arg_ptr(output)
74        .arg_u32(m)
75        .arg_u32(n)
76        .arg_u32(k)
77        .launch(stream)
78}
79
80/// 2026-09-25: Pipelined W8A16 GEMM (kernel `w8a16_gemm_pipelined`), with the same arguments
81/// and block-scale layout as [`w8a16_gemm`]. A 128x32 MxN tile with 256 threads and a
82/// cp.async prefetch (`PM_M_TILE`, `PM_N_TILE`, `PM_THREADS` in
83/// kernels/gb10/common/w8a16_gemm_pipelined.cu).
84#[allow(clippy::too_many_arguments)]
85pub fn w8a16_gemm_pipelined(
86    gpu: &dyn GpuBackend,
87    kernel: KernelHandle,
88    input: DevicePtr,
89    weight: DevicePtr,
90    block_scale: DevicePtr,
91    output: DevicePtr,
92    m: u32,
93    n: u32,
94    k: u32,
95    stream: u64,
96) -> Result<()> {
97    if let Some(twin) = w8a16_pipe128_kernel(gpu, m, n, k)? {
98        // 2026-09-28: `w8a16_gemm_pipe128`: 128 x 128 tiles, 256 threads, 54,272 B of dynamic
99        // shared memory (3 x (8 KiB A + 4 KiB weight bytes) + 2 x 8 KiB decoded + 1 KiB LUT).
100        return KernelLaunch::new(gpu, twin)
101            .grid([n / 128, div_ceil(m, 128), 1])
102            .block([256, 1, 1])
103            .shared_mem(3 * (8192 + 4096) + 2 * 8192 + 1024)
104            .arg_ptr(input)
105            .arg_ptr(weight)
106            .arg_ptr(block_scale)
107            .arg_ptr(output)
108            .arg_u32(m)
109            .arg_u32(n)
110            .arg_u32(k)
111            .launch(stream);
112    }
113    KernelLaunch::new(gpu, kernel)
114        .grid([div_ceil(n, 32), div_ceil(m, 128), 1])
115        .block([256, 1, 1])
116        .arg_ptr(input)
117        .arg_ptr(weight)
118        .arg_ptr(block_scale)
119        .arg_ptr(output)
120        .arg_u32(m)
121        .arg_u32(n)
122        .arg_u32(k)
123        .launch(stream)
124}
125
126/// 2026-09-28: Fewest rows for which [`w8a16_gemm_pipelined`] launches its 128x128 twin;
127/// decode batches (at most 64 rows at the concurrency the MoE ladder runs) keep the original.
128pub const W8A16_PIPE128_MIN_ROWS: u32 = 256;
129
130/// 2026-09-28: `w8a16_gemm_pipe128` (`kernels/gb10/common/w8a16_gemm_pipe128.cu`), the
131/// bit-identical 128x128-tile twin of `w8a16_gemm_pipelined` (same LUT decode, same BF16 MMAs
132/// in the same K order, same per-128-K scale fold; microtest `w8a16_gemm_pipe128_microtest`),
133/// 1.5-1.9x faster on GB10: when the backend carries it, `m >= W8A16_PIPE128_MIN_ROWS`, N and
134/// K are multiples of 128 and `METRALE_NO_W8A16_PIPE128` is absent. Memoized in the
135/// backend's `OpCache`.
136fn w8a16_pipe128_kernel(
137    gpu: &dyn GpuBackend,
138    m: u32,
139    n: u32,
140    k: u32,
141) -> Result<Option<KernelHandle>> {
142    const MODULE: &str = "w8a16_gemm_pipe128";
143    static OFF: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
144    let off = *OFF.get_or_init(|| std::env::var_os("METRALE_NO_W8A16_PIPE128").is_some());
145    if off
146        || m < W8A16_PIPE128_MIN_ROWS
147        || n == 0
148        || !n.is_multiple_of(128)
149        || k == 0
150        || !k.is_multiple_of(128)
151        || !gpu.has_module(MODULE)
152    {
153        return Ok(None);
154    }
155    Ok(Some(gpu.op_cache().kernel(
156        gpu,
157        MODULE,
158        "w8a16_gemm_pipe128",
159    )?))
160}
161
162/// 2026-09-25: Transposed W8A16 GEMM: `C[M,N] = A[M,K] @ dequant(B_t[K,N])`, with the
163/// transposed FP32 block scale `block_scale_t[K/128, N/128]`, so weight reads are contiguous
164/// along N. Each CTA covers 64 rows and 64 columns with 128 threads
165/// (kernels/gb10/common/w8a16_gemm_t.cu, kernels/strix-hip/common/w8a16_gemm_t.cu).
166#[allow(clippy::too_many_arguments)]
167pub fn w8a16_gemm_t(
168    gpu: &dyn GpuBackend,
169    kernel: KernelHandle,
170    input: DevicePtr,
171    weight_t: DevicePtr,      // 2026-09-25: [K, N] FP8 transposed
172    block_scale_t: DevicePtr, // 2026-09-25: [K/128, N/128] FP32 transposed
173    output: DevicePtr,
174    m: u32,
175    n: u32,
176    k: u32,
177    stream: u64,
178) -> Result<()> {
179    KernelLaunch::new(gpu, kernel)
180        .grid([div_ceil(n, 64), div_ceil(m, 64), 1])
181        .block([128, 1, 1])
182        .arg_ptr(input)
183        .arg_ptr(weight_t)
184        .arg_ptr(block_scale_t)
185        .arg_ptr(output)
186        .arg_u32(m)
187        .arg_u32(n)
188        .arg_u32(k)
189        .launch(stream)
190}
191
192/// 2026-09-25: Transposed W8A16 GEMM with a 128x128 MxN tile (kernel `w8a16_gemm_t_m128`):
193/// two 64-row chunks, 8 warps, `m16n8k16` BF16 MMA, and the block scale folded into an FP32
194/// accumulator once per 128-K block (kernels/gb10/common/w8a16_gemm_t_m128.cu). Same
195/// arguments and layout as [`w8a16_gemm_t`], so it takes the output of the `transpose_fp8` /
196/// `transpose_block_scale` kernels unchanged.
197#[allow(clippy::too_many_arguments)]
198pub fn w8a16_gemm_n128_m128(
199    gpu: &dyn GpuBackend,
200    kernel: KernelHandle,
201    input: DevicePtr,
202    weight_t: DevicePtr,      // 2026-09-25: [K, N] FP8 transposed
203    block_scale_t: DevicePtr, // 2026-09-25: [K/128, N/128] FP32 transposed
204    output: DevicePtr,
205    m: u32,
206    n: u32,
207    k: u32,
208    stream: u64,
209) -> Result<()> {
210    super::super::log_gemm_shape(gpu, "w8a16_gemm_t_m128", m, n, k);
211    KernelLaunch::new(gpu, kernel)
212        .grid([div_ceil(n, 128), div_ceil(m, 128), 1])
213        .block([256, 1, 1])
214        .arg_ptr(input)
215        .arg_ptr(weight_t)
216        .arg_ptr(block_scale_t)
217        .arg_ptr(output)
218        .arg_u32(m)
219        .arg_u32(n)
220        .arg_u32(k)
221        .launch(stream)
222}
223
224/// 2026-09-25: Pipelined transposed W8A16 GEMM (kernel `w8a16_gemm_t_pipelined`), with the
225/// same arguments as [`w8a16_gemm_t`]: a 128x32 MxN tile with 256 threads (`PT_M_TILE`,
226/// `PT_N_TILE`, `PT_THREADS` in kernels/gb10/common/w8a16_gemm_t.cu).
227#[allow(clippy::too_many_arguments)]
228pub fn w8a16_gemm_t_pipelined(
229    gpu: &dyn GpuBackend,
230    kernel: KernelHandle,
231    input: DevicePtr,
232    weight_t: DevicePtr,
233    block_scale_t: DevicePtr,
234    output: DevicePtr,
235    m: u32,
236    n: u32,
237    k: u32,
238    stream: u64,
239) -> Result<()> {
240    super::super::log_gemm_shape(gpu, "w8a16_gemm_t_pipelined", m, n, k);
241    KernelLaunch::new(gpu, kernel)
242        .grid([div_ceil(n, 32), div_ceil(m, 128), 1])
243        .block([256, 1, 1])
244        .arg_ptr(input)
245        .arg_ptr(weight_t)
246        .arg_ptr(block_scale_t)
247        .arg_ptr(output)
248        .arg_u32(m)
249        .arg_u32(n)
250        .arg_u32(k)
251        .launch(stream)
252}
253
254/// 2026-09-27: One row's block-scaled W8A16 projection `[1, k] x [n, k]^T` for
255/// a single-row decode site: under `RowTiers::Canonical`, when `m32_kernel` is
256/// linked and `k` is whole 128-wide scale blocks, the 32-row tile (the order
257/// every other row count takes under that policy); otherwise the scalar
258/// [`w8a16_gemv`].
259#[allow(clippy::too_many_arguments)]
260pub fn w8a16_gemv_row_tiered(
261    gpu: &dyn GpuBackend,
262    gemv_kernel: KernelHandle,
263    m32_kernel: KernelHandle,
264    input: DevicePtr,
265    weight: DevicePtr,
266    block_scale: DevicePtr,
267    output: DevicePtr,
268    n: u32,
269    k: u32,
270    stream: u64,
271) -> Result<()> {
272    if crate::layers::row_tiers() == crate::layers::RowTiers::Canonical
273        && super::w8a16_pipelined_prefers_m32(1, k, m32_kernel)
274    {
275        return super::w8a16_gemm_pipelined_m32(
276            gpu,
277            m32_kernel,
278            input,
279            weight,
280            block_scale,
281            output,
282            1,
283            n,
284            k,
285            stream,
286        );
287    }
288    w8a16_gemv(
289        gpu,
290        gemv_kernel,
291        input,
292        weight,
293        block_scale,
294        output,
295        n,
296        k,
297        stream,
298    )
299}
metrale_model_layers/layers/ops/gemm_quant_w8a16.rs

metrale_model_layers/layers/ops/
gemm_quant_w8a16.rs
