metrale_model_layers/layers/ops/
nvfp4_mmq.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the NVFP4 W4A4 MMQ GEMM that the dense FFN uses in
4//! prefill, built on the vendored llama.cpp MMQ headers
5//! (`kernels/gb10/qwen3.6-27b/nvfp4/nvfp4_mmq.cu`, `q4k_vendor/`), plus its
6//! weight repack, activation quantizers and scale2 folds.
7//!
8//! The MMA applies only the per-16 E4M3 scales, so each GEMM output lacks the
9//! per-tensor FP32 `weight_scale_2`. `DenseFfn` folds it in: into the SiLU-mul
10//! for gate and up ([`nvfp4_silu_mul_scaled`] or [`nvfp4_silu_mul_quant`]),
11//! and with [`nvfp4_scale_bf16`] for down. The arm runs when
12//! `ModelLevers::ffn_nvfp4_mmq` is on (unless `METRALE_NO_FFN_NVFP4_MMQ` is set)
13//! and its kernels resolved.
14//!
15//! Owner: model-layers ops.
16//! Invariants: none beyond the types.
17
18use anyhow::Result;
19use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
20use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
21
22/// 2026-09-25: Weights per repacked NVFP4 block: 36 bytes, 32 of E2M1 nibbles and
23/// 4 UE4M3 scales.
24pub const QK_NVFP4: u32 = 64;
25pub const NVFP4_BLOCK_BYTES: usize = 36;
26/// 2026-09-25: `block_fp4_mmq`, the activation format: 256 values in 144 bytes,
27/// the size of `block_q8_1_mmq` (`q4k_vendor/mmq.cuh`).
28const FP4_MMQ_Y_BLOCK_VALS: u32 = 256;
29const FP4_MMQ_Y_BLOCK_BYTES: usize = 144;
30/// 2026-09-25: Dynamic shared memory of the 128-wide M tile.
31pub const NVFP4_MMQ_SMEM: u32 = nvfp4_mmq_smem(128);
32
33/// 2026-09-25: Dynamic shared memory for M tile `mmq_x`, in bytes: `mmq_x` ids,
34/// a y tile of `mmq_x * 36` ints padded to 256 ints, and an x tile of `128 * 76`
35/// ints. This is `mmq_get_nbytes_shared` in `q4k_vendor/mmq.cuh` with
36/// `mmq_y = 128`, `MMQ_MMA_TILE_X_K_FP4 = 76` and 8 warps of 32.
37pub const fn nvfp4_mmq_smem(mmq_x: u32) -> u32 {
38    let y = mmq_x * 36;
39    let y_padded = y.div_ceil(256) * 256;
40    4 * (mmq_x + y_padded + 128 * 76)
41}
42const QUANT_BLOCK_THREADS: u32 = 128;
43
44/// 2026-09-25: Bytes of the repacked form of an `[n, k]` weight; `k` must be a
45/// multiple of 64.
46pub fn nvfp4_mmq_weight_bytes(n: u32, k: u32) -> usize {
47    (n as usize) * (k as usize / QK_NVFP4 as usize) * NVFP4_BLOCK_BYTES
48}
49
50/// 2026-09-25: `block_fp4_mmq` activation bytes for `[m, k]`, with `k` rounded up
51/// to 256, plus 1 MiB of slack. It never exceeds
52/// [`q8_1_scratch_bytes`](super::q8_1_scratch_bytes)`(m, k)`, the size of the
53/// `ffn_act_q8` scratch it is written into.
54pub fn fp4_act_scratch_bytes(m: u32, k: u32) -> usize {
55    let bpc = div_ceil(k, FP4_MMQ_Y_BLOCK_VALS) as usize;
56    (m as usize) * bpc * FP4_MMQ_Y_BLOCK_BYTES + (1 << 20)
57}
58
59/// 2026-09-25: Repack a checkpoint NVFP4 weight (E2M1 `[n, k/2]`, low nibble =
60/// even k, and E4M3 `[n, k/16]` scales) for the MMQ kernel. Each row holds
61/// `k/64` blocks as all nibble bytes (32 per block) followed by all scale bytes
62/// (4 per block). Codes and scale bytes are copied, not requantized.
63pub fn nvfp4_mmq_repack(
64    gpu: &dyn GpuBackend,
65    kernel: KernelHandle,
66    packed: DevicePtr,
67    scales: DevicePtr,
68    out_blocks: DevicePtr,
69    n: u32,
70    k: u32,
71    stream: u64,
72) -> Result<()> {
73    let nblocks = (n as u64) * (k as u64 / QK_NVFP4 as u64);
74    KernelLaunch::new(gpu, kernel)
75        .grid([div_ceil(nblocks as u32, 256), 1, 1])
76        .block([256, 1, 1])
77        .arg_ptr(packed)
78        .arg_ptr(scales)
79        .arg_ptr(out_blocks)
80        .arg_u32(n)
81        .arg_u32(k)
82        .launch(stream)
83}
84
85/// 2026-09-25: Quantize BF16 activations `[m, k]` to `block_fp4_mmq` in `out_y`:
86/// E2M1 values with one UE4M3 scale per 16, `k` padded to 256, one thread per
87/// 16-value group.
88pub fn nvfp4_mmq_quantize_act(
89    gpu: &dyn GpuBackend,
90    kernel: KernelHandle,
91    input_bf16: DevicePtr,
92    out_y: DevicePtr,
93    m: u32,
94    k: u32,
95    stream: u64,
96) -> Result<()> {
97    let kpad = div_ceil(k, FP4_MMQ_Y_BLOCK_VALS) * FP4_MMQ_Y_BLOCK_VALS;
98    let grid_y = div_ceil(kpad, 16 * QUANT_BLOCK_THREADS);
99    KernelLaunch::new(gpu, kernel)
100        .grid([m, grid_y, 1])
101        .block([QUANT_BLOCK_THREADS, 1, 1])
102        .arg_ptr(input_bf16)
103        .arg_ptr(out_y)
104        .arg_u64(k as u64)
105        .arg_u64(k as u64)
106        .arg_u64(kpad as u64)
107        .arg_u32(m)
108        .launch(stream)
109}
110
111/// 2026-09-25: `C[m, n] = A[m, k] x W[n, k]` in BF16, without `weight_scale_2`,
112/// on the 128-wide M tile. `a_fp4` is `block_fp4_mmq`, `w_nvfp4` the
113/// [`nvfp4_mmq_repack`] output. `kernel_wc` is used when `n` is not a multiple
114/// of 128.
115pub fn nvfp4_mmq_gemm(
116    gpu: &dyn GpuBackend,
117    kernel_nc: KernelHandle,
118    kernel_wc: KernelHandle,
119    a_fp4: DevicePtr,
120    w_nvfp4: DevicePtr,
121    out_bf16: DevicePtr,
122    m: u32,
123    n: u32,
124    k: u32,
125    stream: u64,
126) -> Result<()> {
127    let kernel = if !n.is_multiple_of(128) {
128        kernel_wc
129    } else {
130        kernel_nc
131    };
132    KernelLaunch::new(gpu, kernel)
133        .grid([div_ceil(n, 128), div_ceil(m, 128), 1])
134        .block([32, 8, 1])
135        .shared_mem(NVFP4_MMQ_SMEM)
136        .arg_ptr(w_nvfp4)
137        .arg_ptr(a_fp4)
138        .arg_ptr(out_bf16)
139        .arg_u32(n)
140        .arg_u32(m)
141        .arg_u32(k)
142        .arg_u32(k / QK_NVFP4)
143        .arg_u32(m)
144        .arg_u32(n)
145        .launch(stream)
146}
147
148/// 2026-09-25: [`nvfp4_mmq_gemm`] with the M tile `mmq_x` chosen by the caller.
149/// 2026-09-28: With `out_scale`, the pipelined 128-row tile applies it in its store and
150/// the call returns `true`; otherwise the caller still owes [`nvfp4_scale_bf16`].
151/// The tile choice:
152/// the kernels must be that tile's entries. `nvfp4_mmq.cu` instantiates 16, 32,
153/// 64 and 128. A 128-wide tile issues MMAs for all 128 columns whatever `m` is,
154/// so a small `m` wants a small tile. `grid.y = ceil(m / mmq_x)`, and each M tile
155/// streams the whole weight again, so `m` must not exceed `mmq_x`
156/// (debug-asserted).
157#[allow(clippy::too_many_arguments)]
158pub fn nvfp4_mmq_gemm_tiled(
159    gpu: &dyn GpuBackend,
160    kernel_nc: KernelHandle,
161    kernel_wc: KernelHandle,
162    mmq_x: u32,
163    a_fp4: DevicePtr,
164    w_nvfp4: DevicePtr,
165    out_bf16: DevicePtr,
166    m: u32,
167    n: u32,
168    k: u32,
169    out_scale: Option<f32>,
170    stream: u64,
171) -> Result<bool> {
172    // 2026-09-28: The 128-wide tile runs on `metrale_nvfp4_gemm_pipe`, bit-identical to
173    // `metrale_nvfp4_mmq128_*` and 2.5-2.8x faster on the dense FFN prefill shapes. It is
174    // exported by the same module, so it resolves wherever the MMQ family does.
175    if mmq_x == 128 && k.is_multiple_of(256) {
176        let pipe = gpu
177            .op_cache()
178            .kernel(gpu, "nvfp4_mmq", "metrale_nvfp4_gemm_pipe")?;
179        nvfp4_pipe_gemm(
180            gpu, pipe, a_fp4, w_nvfp4, out_bf16, m, n, k, out_scale, stream,
181        )?;
182        return Ok(out_scale.is_some());
183    }
184    debug_assert!(
185        matches!(mmq_x, 16 | 32 | 64 | 128),
186        "mmq_x must be an instantiated tile"
187    );
188    debug_assert!(
189        m <= mmq_x,
190        "grid.y>1 would re-stream the weights per M-tile"
191    );
192    let kernel = if !n.is_multiple_of(128) {
193        kernel_wc
194    } else {
195        kernel_nc
196    };
197    KernelLaunch::new(gpu, kernel)
198        .grid([div_ceil(n, 128), div_ceil(m, mmq_x), 1])
199        .block([32, 8, 1])
200        .shared_mem(nvfp4_mmq_smem(mmq_x))
201        .arg_ptr(w_nvfp4)
202        .arg_ptr(a_fp4)
203        .arg_ptr(out_bf16)
204        .arg_u32(n)
205        .arg_u32(m)
206        .arg_u32(k)
207        .arg_u32(k / QK_NVFP4)
208        .arg_u32(m)
209        .arg_u32(n)
210        .launch(stream)?;
211    Ok(false)
212}
213/// 2026-09-28: Dynamic shared memory of `metrale_nvfp4_gemm_pipe`: two stages of
214/// 128 weight rows and 192 activation rows, 144 bytes each.
215pub const NVFP4_PIPE_SMEM: u32 = 2 * (128 + 192) * 144;
216
217/// 2026-09-28: `C[m, n] = A[m, k] x W[n, k]` in BF16, without `weight_scale_2`, through
218/// `metrale_nvfp4_gemm_pipe` (same operands and output as [`nvfp4_mmq_gemm`], bit-identical).
219/// `k` must be a multiple of 256. Grid of 128-channel x 192-token tiles, 256 threads. With `out_scale`,
220/// each output is multiplied by it as [`nvfp4_scale_bf16`] does in place, bit for bit.
221#[allow(clippy::too_many_arguments)]
222pub fn nvfp4_pipe_gemm(
223    gpu: &dyn GpuBackend,
224    kernel: KernelHandle,
225    a_fp4: DevicePtr,
226    w_nvfp4: DevicePtr,
227    out_bf16: DevicePtr,
228    m: u32,
229    n: u32,
230    k: u32,
231    out_scale: Option<f32>,
232    stream: u64,
233) -> Result<()> {
234    debug_assert!(k.is_multiple_of(256), "nvfp4_pipe_gemm needs K % 256 == 0");
235    KernelLaunch::new(gpu, kernel)
236        .grid([div_ceil(m, 192) * div_ceil(n, 128), 1, 1])
237        .block([256, 1, 1])
238        .shared_mem(NVFP4_PIPE_SMEM)
239        .arg_ptr(w_nvfp4)
240        .arg_ptr(a_fp4)
241        .arg_ptr(out_bf16)
242        .arg_u32(n)
243        .arg_u32(m)
244        .arg_u32(k)
245        .arg_f32(out_scale.unwrap_or(1.0))
246        .arg_u32(u32::from(out_scale.is_some()))
247        .launch(stream)
248}
249
250/// 2026-09-25: `silu(gate * gate_scale) * (up * up_scale)` quantized straight to
251/// `block_fp4_mmq` in `out_y` for the down GEMM, without writing a BF16
252/// intermediate. `k` is the intermediate width, padded to 256. The kernel
253/// applies no clamp.
254#[allow(clippy::too_many_arguments)]
255pub fn nvfp4_silu_mul_quant(
256    gpu: &dyn GpuBackend,
257    kernel: KernelHandle,
258    gate: DevicePtr,
259    up: DevicePtr,
260    out_y: DevicePtr,
261    gate_scale: f32,
262    up_scale: f32,
263    m: u32,
264    k: u32,
265    stream: u64,
266) -> Result<()> {
267    let kpad = div_ceil(k, FP4_MMQ_Y_BLOCK_VALS) * FP4_MMQ_Y_BLOCK_VALS;
268    let grid_y = div_ceil(kpad, 16 * QUANT_BLOCK_THREADS);
269    KernelLaunch::new(gpu, kernel)
270        .grid([m, grid_y, 1])
271        .block([QUANT_BLOCK_THREADS, 1, 1])
272        .arg_ptr(gate)
273        .arg_ptr(up)
274        .arg_ptr(out_y)
275        .arg_f32(gate_scale)
276        .arg_f32(up_scale)
277        .arg_u64(k as u64)
278        .arg_u64(kpad as u64)
279        .arg_u32(m)
280        .launch(stream)
281}
282
283/// 2026-09-25: `data[i] *= scale` in place over `total` BF16 values; the caller
284/// passes the down projection's `weight_scale_2`.
285pub fn nvfp4_scale_bf16(
286    gpu: &dyn GpuBackend,
287    kernel: KernelHandle,
288    data: DevicePtr,
289    scale: f32,
290    total: u32,
291    stream: u64,
292) -> Result<()> {
293    KernelLaunch::new(gpu, kernel)
294        .grid([div_ceil(total, 256), 1, 1])
295        .block([256, 1, 1])
296        .arg_ptr(data)
297        .arg_f32(scale)
298        .arg_u32(total)
299        .launch(stream)
300}
301
302/// 2026-09-25: `out[i] = silu(gate[i] * gate_scale) * (up[i] * up_scale)` over
303/// `total` BF16 values. The kernel applies no clamp.
304#[allow(clippy::too_many_arguments)]
305pub fn nvfp4_silu_mul_scaled(
306    gpu: &dyn GpuBackend,
307    kernel: KernelHandle,
308    gate: DevicePtr,
309    up: DevicePtr,
310    out: DevicePtr,
311    gate_scale: f32,
312    up_scale: f32,
313    total: u32,
314    stream: u64,
315) -> Result<()> {
316    KernelLaunch::new(gpu, kernel)
317        .grid([div_ceil(total, 256), 1, 1])
318        .block([256, 1, 1])
319        .arg_ptr(gate)
320        .arg_ptr(up)
321        .arg_ptr(out)
322        .arg_f32(gate_scale)
323        .arg_f32(up_scale)
324        .arg_u32(total)
325        .launch(stream)
326}
metrale_model_layers/layers/ops/nvfp4_mmq.rs

metrale_model_layers/layers/ops/
nvfp4_mmq.rs
