metrale_model_layers/layers/ops/
q4k_mmq.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the Q4_K MMQ GEMM that the dense FFN uses in
4//! prefill under `METRALE_FFN_MMQ`, built on the vendored llama.cpp MMQ headers
5//! (`kernels/gb10/qwen3.6-27b/nvfp4/q4k_mmq.cu`, `q4k_quantize.cu`), plus the
6//! weight conversion and the q8_1 activation quantizer.
7//!
8//! `DenseFfn::ensure_q4k_weight` converts each NVFP4 weight once (dequantize
9//! to BF16, then quantize to `block_q4_K`) and caches the result. Each prefill
10//! quantizes its activations to q8_1 and runs the GEMM with a BF16 store.
11//!
12//! Owner: model-layers ops.
13//! Invariants: none beyond the types.
14
15use anyhow::Result;
16use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
17use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
18
19/// 2026-09-25: Weights per `block_q4_K`.
20pub const QK_K: u32 = 256;
21/// 2026-09-25: `sizeof(block_q4_K)`.
22pub const Q4K_BLOCK_BYTES: usize = 144;
23/// 2026-09-25: Dynamic shared memory of the Q4_K MMQ kernel at `mmq_x = mmq_y =
24/// 128`: `mmq_get_nbytes_shared` in `q4k_vendor/mmq.cuh` with an x tile of 76
25/// ints per row. Above 48 KiB, so the registry raises the kernel's dynamic
26/// shared memory limit (`registry.rs`).
27pub const Q4K_MMQ_SMEM: u32 = 57856;
28const CUDA_QUANTIZE_BLOCK_SIZE_MMQ: u32 = 128;
29
30/// 2026-09-25: Bytes of the Q4_K form of an `[nrows, n_per_row]` weight;
31/// `n_per_row` must be a multiple of 256.
32pub fn q4k_weight_bytes(nrows: u32, n_per_row: u32) -> usize {
33    (nrows as usize) * (n_per_row as usize / QK_K as usize) * Q4K_BLOCK_BYTES
34}
35
36/// 2026-09-25: q8_1 activation scratch bytes for `[m, k]`: 4 bytes per value with
37/// `k` rounded up to 256, plus 1 MiB. `ffn_act_q8` in
38/// `gpu-runtime/src/buffers/sizes.rs` uses the same expression with `k` the
39/// larger of the hidden and intermediate sizes.
40pub fn q8_1_scratch_bytes(m: u32, k: u32) -> usize {
41    let kpad = div_ceil(k, QK_K) * QK_K;
42    (m as usize) * (kpad as usize) * 4 + (1 << 20)
43}
44
45/// 2026-09-25: Dequantize an NVFP4 weight `[n, k]` (packed E2M1, E4M3 group
46/// scales and the per-tensor `scale2`) to BF16 `[n, k]`, one block per row.
47pub fn dequant_nvfp4_to_bf16(
48    gpu: &dyn GpuBackend,
49    kernel: KernelHandle,
50    packed: DevicePtr,
51    scales: DevicePtr,
52    out_bf16: DevicePtr,
53    scale2: f32,
54    n: u32,
55    k: u32,
56    stream: u64,
57) -> Result<()> {
58    KernelLaunch::new(gpu, kernel)
59        .grid([n, 1, 1])
60        .block([256, 1, 1])
61        .arg_ptr(packed)
62        .arg_ptr(scales)
63        .arg_ptr(out_bf16)
64        .arg_f32(scale2)
65        .arg_u32(n)
66        .arg_u32(k)
67        .launch(stream)
68}
69
70/// 2026-09-25: Quantize BF16 weights `[nrows, n_per_row]` to `block_q4_K`, one
71/// thread per 256-value block.
72pub fn quantize_weight_q4k(
73    gpu: &dyn GpuBackend,
74    kernel: KernelHandle,
75    input_bf16: DevicePtr,
76    out_q4k: DevicePtr,
77    nrows: u32,
78    n_per_row: u32,
79    stream: u64,
80) -> Result<()> {
81    let total_sb = (nrows as u64) * (n_per_row as u64 / QK_K as u64);
82    let grid_x = div_ceil(total_sb as u32, 128);
83    KernelLaunch::new(gpu, kernel)
84        .grid([grid_x, 1, 1])
85        .block([128, 1, 1])
86        .arg_ptr(input_bf16)
87        .arg_ptr(out_q4k)
88        .arg_u32(nrows)
89        .arg_u32(n_per_row)
90        .launch(stream)
91}
92
93/// 2026-09-25: Quantize BF16 activations `[m, k]` to q8_1 in `out_q8`, with `k`
94/// padded to 256. The scale layout (DS4, D4 or D2S6) is the kernel's.
95pub fn quantize_act_q8_1(
96    gpu: &dyn GpuBackend,
97    kernel: KernelHandle,
98    input_bf16: DevicePtr,
99    out_q8: DevicePtr,
100    m: u32,
101    k: u32,
102    stream: u64,
103) -> Result<()> {
104    let kpad = div_ceil(k, QK_K) * QK_K;
105    let grid_y = div_ceil(kpad, 4 * CUDA_QUANTIZE_BLOCK_SIZE_MMQ);
106    KernelLaunch::new(gpu, kernel)
107        .grid([m, grid_y, 1])
108        .block([CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1])
109        .arg_ptr(input_bf16)
110        .arg_ptr(out_q8)
111        .arg_u64(k as u64)
112        .arg_u64(k as u64)
113        .arg_u64(kpad as u64)
114        .arg_u32(m)
115        .launch(stream)
116}
117
118/// 2026-09-25: `C[m, n] = A[m, k] x W[n, k]` in BF16, with `a_q8` in q8_1 and
119/// `w_q4k` in `block_q4_K`, on a 128 x 128 tile. `kernel_wc` is used when `n`
120/// is not a multiple of 128.
121pub fn q4k_mmq_gemm(
122    gpu: &dyn GpuBackend,
123    kernel_nc: KernelHandle,
124    kernel_wc: KernelHandle,
125    a_q8: DevicePtr,
126    w_q4k: DevicePtr,
127    out_bf16: DevicePtr,
128    m: u32,
129    n: u32,
130    k: u32,
131    stream: u64,
132) -> Result<()> {
133    let kernel = if !n.is_multiple_of(128) {
134        kernel_wc
135    } else {
136        kernel_nc
137    };
138    KernelLaunch::new(gpu, kernel)
139        .grid([div_ceil(n, 128), div_ceil(m, 128), 1])
140        .block([32, 8, 1])
141        .shared_mem(Q4K_MMQ_SMEM)
142        .arg_ptr(w_q4k)
143        .arg_ptr(a_q8)
144        .arg_ptr(out_bf16)
145        .arg_u32(n)
146        .arg_u32(m)
147        .arg_u32(k)
148        .arg_u32(k / QK_K)
149        .arg_u32(m)
150        .arg_u32(n)
151        .launch(stream)
152}
metrale_model_layers/layers/ops/q4k_mmq.rs

metrale_model_layers/layers/ops/
q4k_mmq.rs
