metrale_model_layers/layers/ops/
gemm_fp8_prefill.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the FP8-weight prefill GEMMs, the NVFP4-to-FP8
4//! weight pre-dequant, the BF16-to-FP8 activation cast and the row-scaled
5//! BF16-to-FP8 weight quantizer.
6//!
7//! Owner: model-layers ops.
8//! Invariants: none beyond the types.
9
10#![allow(unused_imports)]
11
12use anyhow::Result;
13use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
14use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
15
16use crate::weight_map::{Fp8DenseWeight, QuantizedWeight};
17
18use super::*;
19
20/// 2026-09-25: FP8-weight prefill GEMM, `C = A @ B_fp8^T`: A `[M, K]` BF16,
21/// B_fp8 `[N, K]` FP8 E4M3 (from [`predequant_nvfp4_to_fp8`]), C `[M, N]`
22/// BF16.
23///
24/// When `k` is a multiple of 32 and `METRALE_FP8_LDMAB` is not `0`, it casts A
25/// to FP8 into a scratch buffer ([`bf16_to_fp8`]) and launches
26/// `fp8_fp8_gemm_ldmab` (module `w4a16_fp8_ldmab`, 128x256 tiles, two 128-byte K
27/// stages) instead of `kernel`.
28#[allow(clippy::too_many_arguments)]
29pub fn fp8_gemm_n128(
30    gpu: &dyn GpuBackend,
31    kernel: KernelHandle,
32    input: DevicePtr,
33    b_fp8: DevicePtr,
34    output: DevicePtr,
35    m: u32,
36    n: u32,
37    k: u32,
38    stream: u64,
39) -> Result<()> {
40    check_e4m3_activation_range(gpu, kernel, input, m, k, stream)?;
41    if k.is_multiple_of(32) && std::env::var("METRALE_FP8_LDMAB").as_deref() != Ok("0") {
42        return fp8_act_ldmab_gemm(gpu, input, b_fp8, output, m, n, k, stream);
43    }
44    KernelLaunch::new(gpu, kernel)
45        .grid([div_ceil(n, 128), div_ceil(m, 64), 1])
46        .block([128, 1, 1])
47        .arg_ptr(input)
48        .arg_ptr(b_fp8)
49        .arg_ptr(output)
50        .arg_u32(m)
51        .arg_u32(n)
52        .arg_u32(k)
53        .launch(stream)
54}
55
56/// 2026-09-28: `C = e4m3(A) @ B_fp8^T` through `bf16_to_fp8` and `fp8_fp8_gemm_ldmab`
57/// (module `w4a16_fp8_ldmab`). The handles and the activation scratch come from the
58/// backend's op cache, not from statics: they belong to this backend's modules and
59/// allocations, which a process-wide cache would outlive. `k` must be a multiple of 32.
60#[allow(clippy::too_many_arguments)]
61fn fp8_act_ldmab_gemm(
62    gpu: &dyn GpuBackend,
63    input: DevicePtr,
64    b_fp8: DevicePtr,
65    output: DevicePtr,
66    m: u32,
67    n: u32,
68    k: u32,
69    stream: u64,
70) -> Result<()> {
71    let cache = gpu.op_cache();
72    let qk = cache.kernel(gpu, "w4a16", "bf16_to_fp8")?;
73    let lk = cache.kernel(gpu, "w4a16_fp8_ldmab", "fp8_fp8_gemm_ldmab")?;
74    let a8 = cache.scratch(gpu, "fp8_prefill_activation", (m as usize) * (k as usize))?;
75    bf16_to_fp8(gpu, qk, input, a8, m * k, stream)?;
76    // 2026-09-28: 1-D grid of 128x256 tiles (the kernel orders them in groups of 8 M
77    // tiles), 256 threads, two 48 KiB stages of dynamic shared memory.
78    KernelLaunch::new(gpu, lk)
79        .grid([div_ceil(m, 128) * div_ceil(n, 256), 1, 1])
80        .block([256, 1, 1])
81        .shared_mem(2 * (128 + 256) * 128)
82        .arg_ptr(a8)
83        .arg_ptr(b_fp8)
84        .arg_ptr(output)
85        .arg_u32(m)
86        .arg_u32(n)
87        .arg_u32(k)
88        .launch(stream)
89}
90
91/// 2026-09-28: Rows at or above which [`w4a16_t_via_fp8_ldmab`] beats
92/// `w4a16_gemm_t_m128` on the dense 27B attention projections (dgx2: x1.23-1.35 at 1024
93/// rows, x0.69 at 300, x2.15 at 8192): below it the per-call weight dequant dominates.
94pub const W4A16_VIA_FP8_MIN_M: u32 = 1024;
95
96/// 2026-09-28: `w4a16_gemm_t_m128`'s product on the multistage FP8 GEMM: the transposed
97/// NVFP4 weight (`[K/2, N]` codes, `[K/16, N]` E4M3 scales, scalar `weight_scale_2`) is
98/// dequantized to E4M3 `[N, K]` with that kernel's arithmetic (`fp8_predequant_nvfp4_t`),
99/// then `fp8_act_ldmab_gemm` runs. Bit-identical to `w4a16_gemm_t_m128` (kbench
100/// `w4a16_bench.cu`, nine shapes). `k` must be a multiple of 32.
101#[allow(clippy::too_many_arguments)]
102pub fn w4a16_t_via_fp8_ldmab(
103    gpu: &dyn GpuBackend,
104    input: DevicePtr,
105    weight: &QuantizedWeight,
106    output: DevicePtr,
107    m: u32,
108    n: u32,
109    k: u32,
110    stream: u64,
111) -> Result<()> {
112    debug_assert!(
113        k.is_multiple_of(32),
114        "w4a16_t_via_fp8_ldmab needs K % 32 == 0"
115    );
116    let cache = gpu.op_cache();
117    let pk = cache.kernel(gpu, "w4a16_fp8_ldmab", "fp8_predequant_nvfp4_t")?;
118    let b8 = cache.scratch(gpu, "w4a16_t_fp8_weight", (n as usize) * (k as usize))?;
119    KernelLaunch::new(gpu, pk)
120        .grid([div_ceil(n, 256), k / 16, 1])
121        .block([256, 1, 1])
122        .arg_ptr(weight.weight)
123        .arg_ptr(weight.weight_scale)
124        .arg_f32(weight.weight_scale_2)
125        .arg_ptr(b8)
126        .arg_u32(n)
127        .arg_u32(k)
128        .launch(stream)?;
129    fp8_act_ldmab_gemm(gpu, input, b8, output, m, n, k, stream)
130}
131
132#[allow(clippy::too_many_arguments)]
133/// 2026-09-25: FP8-weight GEMM through the `fp8_gemm_t_mfast` kernel, with M
134/// on the fast grid axis (`blockIdx.x` = M block), so consecutive CTAs share a
135/// B panel.
136pub fn fp8_gemm_n128_mfast(
137    gpu: &dyn GpuBackend,
138    kernel: KernelHandle,
139    input: DevicePtr,
140    b_fp8: DevicePtr,
141    output: DevicePtr,
142    m: u32,
143    n: u32,
144    k: u32,
145    stream: u64,
146) -> Result<()> {
147    check_e4m3_activation_range(gpu, kernel, input, m, k, stream)?;
148    KernelLaunch::new(gpu, kernel)
149        .grid([div_ceil(m, 64), div_ceil(n, 128), 1])
150        .block([128, 1, 1])
151        .arg_ptr(input)
152        .arg_ptr(b_fp8)
153        .arg_ptr(output)
154        .arg_u32(m)
155        .arg_u32(n)
156        .arg_u32(k)
157        .launch(stream)
158}
159
160/// 2026-09-25: FP8-weight GEMM through `fp8_gemm_t_m128_mfast`: a 128-row M
161/// tile (two 64-row chunks per CTA) with M on the fast grid axis.
162pub fn fp8_gemm_m128_mfast(
163    gpu: &dyn GpuBackend,
164    kernel: KernelHandle,
165    input: DevicePtr,
166    b_fp8: DevicePtr,
167    output: DevicePtr,
168    m: u32,
169    n: u32,
170    k: u32,
171    stream: u64,
172) -> Result<()> {
173    check_e4m3_activation_range(gpu, kernel, input, m, k, stream)?;
174    KernelLaunch::new(gpu, kernel)
175        .grid([div_ceil(m, 128), div_ceil(n, 128), 1])
176        .block([128, 1, 1])
177        .arg_ptr(input)
178        .arg_ptr(b_fp8)
179        .arg_ptr(output)
180        .arg_u32(m)
181        .arg_u32(n)
182        .arg_u32(k)
183        .launch(stream)
184}
185
186/// 2026-09-25: FP8 x FP8 GEMM through `fp8_fp8_gemm_t_m128_mfast`: a 128-row
187/// M tile with M on the fast grid axis. A must already be FP8 E4M3
188/// ([`bf16_to_fp8`]).
189pub fn fp8_fp8_gemm_m128_mfast(
190    gpu: &dyn GpuBackend,
191    kernel: KernelHandle,
192    input: DevicePtr,
193    b_fp8: DevicePtr,
194    output: DevicePtr,
195    m: u32,
196    n: u32,
197    k: u32,
198    stream: u64,
199) -> Result<()> {
200    KernelLaunch::new(gpu, kernel)
201        .grid([div_ceil(m, 128), div_ceil(n, 128), 1])
202        .block([128, 1, 1])
203        .arg_ptr(input)
204        .arg_ptr(b_fp8)
205        .arg_ptr(output)
206        .arg_u32(m)
207        .arg_u32(n)
208        .arg_u32(k)
209        .launch(stream)
210}
211
212/// 2026-09-25: Dequantize an NVFP4 weight to FP8 E4M3: reads
213/// `B_packed[N, K/2]`, `B_scale[N, K/16]` and `scale2`, writes `B_fp8[N, K]`.
214/// `QuantizedWeight::predequant_to_fp8` calls it.
215pub fn predequant_nvfp4_to_fp8(
216    gpu: &dyn GpuBackend,
217    kernel: KernelHandle,
218    b_packed: DevicePtr,
219    b_scale: DevicePtr,
220    scale2: f32,
221    b_fp8: DevicePtr,
222    n: u32,
223    k: u32,
224    stream: u64,
225) -> Result<()> {
226    let total = n * k / 2;
227    KernelLaunch::new(gpu, kernel)
228        .grid([div_ceil(total, 256), 1, 1])
229        .block([256, 1, 1])
230        .arg_ptr(b_packed)
231        .arg_ptr(b_scale)
232        .arg_f32(scale2)
233        .arg_ptr(b_fp8)
234        .arg_u32(n)
235        .arg_u32(k)
236        .launch(stream)
237}
238
239/// 2026-09-25: Cast `total_elements` BF16 values to FP8 E4M3, two per thread.
240/// `total_elements` must be even: the kernel reads and writes pairs.
241///
242/// 2026-09-29: The cast has no scale, so a value past 448 saturates; debug builds check
243/// the source range first ([`check_e4m3_range`]).
244pub fn bf16_to_fp8(
245    gpu: &dyn GpuBackend,
246    kernel: KernelHandle,
247    src: DevicePtr,
248    dst: DevicePtr,
249    total_elements: u32,
250    stream: u64,
251) -> Result<()> {
252    check_e4m3_range(gpu, src, total_elements as usize, stream)?;
253    let threads_needed = total_elements / 2;
254    KernelLaunch::new(gpu, kernel)
255        .grid([div_ceil(threads_needed, 256), 1, 1])
256        .block([256, 1, 1])
257        .arg_ptr(src)
258        .arg_ptr(dst)
259        .arg_u32(total_elements)
260        .launch(stream)
261}
262
263/// 2026-09-25: Quantize a BF16 weight `[N, K]` to FP8 E4M3 `[N, K]` with one
264/// f32 scale per row `[N]`, one CTA per row (kernel `quantize_bf16_to_fp8` in
265/// `kernels/gb10/common/dense_gemv_fp8w.cu`). `DenseWeight::quantize_to_fp8`
266/// calls it; the resulting `Fp8DenseWeight` feeds the row-scaled FP8 kernels
267/// such as [`fp8_gemm_n128_row_scaled`].
268#[allow(clippy::too_many_arguments)]
269pub fn quantize_bf16_to_fp8(
270    gpu: &dyn GpuBackend,
271    kernel: KernelHandle,
272    input: DevicePtr,
273    output: DevicePtr,
274    row_scales: DevicePtr,
275    n: u32,
276    k: u32,
277    stream: u64,
278) -> Result<()> {
279    KernelLaunch::new(gpu, kernel)
280        .grid([n, 1, 1])
281        .block([256, 1, 1])
282        .arg_ptr(input)
283        .arg_ptr(output)
284        .arg_ptr(row_scales)
285        .arg_u32(n)
286        .arg_u32(k)
287        .launch(stream)
288}
289
290/// 2026-09-25: [`fp8_gemm_n128_row_scaled`] through the single-warp kernel
291/// `fp8_gemm_t_row_scaled_m16`, whose M tile is 16 rows. The grid has no M
292/// dimension, so `m` must be at most 16: rows past the tile are not computed.
293#[allow(clippy::too_many_arguments)]
294pub fn fp8_gemm_n128_row_scaled_m16(
295    gpu: &dyn GpuBackend,
296    kernel: KernelHandle,
297    input: DevicePtr,
298    weight: &Fp8DenseWeight,
299    output: DevicePtr,
300    m: u32,
301    n: u32,
302    k: u32,
303    stream: u64,
304) -> Result<()> {
305    check_e4m3_activation_range(gpu, kernel, input, m, k, stream)?;
306    KernelLaunch::new(gpu, kernel)
307        .grid([div_ceil(n, 128), 1, 1])
308        .block([32, 1, 1])
309        .arg_ptr(input)
310        .arg_ptr(weight.weight)
311        .arg_ptr(weight.row_scale)
312        .arg_ptr(output)
313        .arg_u32(m)
314        .arg_u32(n)
315        .arg_u32(k)
316        .launch(stream)
317}
318
319/// 2026-09-25: Row-scaled FP8 GEMM,
320/// `C[M, N] = A[M, K] @ (dequant(B_fp8[N, K]) * row_scale[N])^T`; the scale
321/// multiplies each output column before the BF16 store. Takes the
322/// `Fp8DenseWeight` from [`crate::weight_map::DenseWeight::quantize_to_fp8`].
323/// The DFlash drafter calls it when its quantization is `Fp8Weights`
324/// (`dflash_head/forward_block_layer_paged.rs`).
325///
326/// Kernel: `fp8_gemm_t_row_scaled` (per-model `w4a16_gemm.cu`).
327#[allow(clippy::too_many_arguments)]
328pub fn fp8_gemm_n128_row_scaled(
329    gpu: &dyn GpuBackend,
330    kernel: KernelHandle,
331    input: DevicePtr,
332    weight: &Fp8DenseWeight,
333    output: DevicePtr,
334    m: u32,
335    n: u32,
336    k: u32,
337    stream: u64,
338) -> Result<()> {
339    check_e4m3_activation_range(gpu, kernel, input, m, k, stream)?;
340    KernelLaunch::new(gpu, kernel)
341        .grid([div_ceil(n, 128), div_ceil(m, 64), 1])
342        .block([128, 1, 1])
343        .arg_ptr(input)
344        .arg_ptr(weight.weight)
345        .arg_ptr(weight.row_scale)
346        .arg_ptr(output)
347        .arg_u32(m)
348        .arg_u32(n)
349        .arg_u32(k)
350        .launch(stream)
351}
metrale_model_layers/layers/ops/gemm_fp8_prefill.rs

metrale_model_layers/layers/ops/
gemm_fp8_prefill.rs
