metrale_model_layers/layers/ops/
kquant_mmq.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2// provenance-id: 526f6e616c6420522e205374657369616b
3
4//! 2026-09-25: Launchers for the DeepSeek-V4.1 K-quant kernels (`kquant_moe` module, kernels/gb10/deepseek-v4-flash/nvfp4/kquant_moe.cu), which read raw Q2_K / Q3_K / Q6_K blocks in place.
5//!
6//! Owner: model-layers ops.
7//! Invariants:
8//! - Every GEMV launcher refuses `m` outside `1..=8` (the kernels' `KQ_MAX_M`) and a `k` that is
9//!   not a multiple of `QK_K` (256).
10//! - Decode quantizes the activation with [`kquant_q8_1_rows`] and runs a `kquant_mmvq*`
11//!   GEMV; prefill quantizes it into the type's MMQ q8_1 layout through
12//!   [`super::quantize_act_q8_1`] (D2S6 for Q2_K, D4 for Q3_K) and runs [`kquant_mmq_gemm`].
13
14use anyhow::Result;
15use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
16use metrale_gpu_runtime::kernel_args::KernelLaunch;
17
18use super::q4k_mmq::QK_K;
19
20pub const KQUANT_MODULE: &str = "kquant_moe";
21/// 2026-09-25: Bytes per 256-value super-block: `block_q2_K`, `block_q3_K`, `block_q6_K`.
22pub const Q2K_BLOCK_BYTES: usize = 84;
23pub const Q3K_BLOCK_BYTES: usize = 110;
24pub const Q6K_BLOCK_BYTES: usize = 210;
25
26/// 2026-09-25: A resident `[N, K]` row-major projection from the GGUF path: expanded bf16, or the
27/// raw `Q2_K` / `Q3_K` blocks (`WeightDtype::Q2K` / `Q3K`) that the K-quant GEMV and MMQ
28/// kernels read.
29#[derive(Clone, Copy, Debug)]
30pub enum ResidentMat {
31    Bf16(DevicePtr),
32    Q2K(DevicePtr),
33    Q3K(DevicePtr),
34}
35
36impl ResidentMat {
37    /// 2026-09-25: The sub-matrix starting `rows` rows in (each row `k` weights long).
38    pub fn at_rows(self, rows: usize, k: usize) -> ResidentMat {
39        let at = |p: DevicePtr, off: usize| DevicePtr(p.0 + off as u64);
40        match self {
41            ResidentMat::Bf16(p) => ResidentMat::Bf16(at(p, rows * k * 2)),
42            ResidentMat::Q2K(p) => ResidentMat::Q2K(at(p, rows * (k / 256) * Q2K_BLOCK_BYTES)),
43            ResidentMat::Q3K(p) => ResidentMat::Q3K(at(p, rows * (k / 256) * Q3K_BLOCK_BYTES)),
44        }
45    }
46}
47/// 2026-09-25: Dynamic shared memory of the 128x128 MMQ tile on the MMA path, as the vendored
48/// `mmq_get_nbytes_shared` computes it with `MMQ_TILE_NE_K = 32`: ids 512 B + tile_x (128 rows x
49/// tile_x_k ints) + tile_y (128 x 144 B). `MMQ_MMA_TILE_X_K_Q2_K` = 2*32 + 32 + 4 = 100;
50/// `MMQ_MMA_TILE_X_K_Q3_K` = 2*32 + 16 + 4 = 84.
51pub const Q2K_MMQ_SMEM: u32 = 512 + 128 * 100 * 4 + 128 * 144;
52pub const Q3K_MMQ_SMEM: u32 = 512 + 128 * 84 * 4 + 128 * 144;
53/// 2026-09-25: `block_q8_1`: 32 int8 + half2 (d, sum).
54pub const Q8_1_BLOCK_BYTES: usize = 36;
55
56fn div_ceil(a: u32, b: u32) -> u32 {
57    a.div_ceil(b)
58}
59
60/// 2026-09-25: Raw bytes of an `[n, k]` K-quant weight.
61pub fn kquant_weight_bytes(n: u32, k: u32, block_bytes: usize) -> usize {
62    (n as usize) * (k as usize / QK_K as usize) * block_bytes
63}
64
65/// 2026-09-25: Bytes of the plain `block_q8_1` buffer for `[m, k]` activations.
66pub fn kquant_q8_1_rows_bytes(m: u32, k: u32) -> usize {
67    (m as usize) * (k as usize / 32) * Q8_1_BLOCK_BYTES
68}
69
70/// 2026-09-25: Bytes of the MMQ-layout q8_1 buffer for `[m, k]` activations: 144 B per 128
71/// values, with `m` rounded up to the 128-row tile and `k` to 256.
72pub fn kquant_mmq_act_bytes(m: u32, k: u32) -> usize {
73    let mpad = div_ceil(m, 128) * 128;
74    let kpad = div_ceil(k, QK_K) * QK_K;
75    (mpad as usize) * (kpad as usize / 128) * 144
76}
77
78/// 2026-09-25: bf16 `[m, k]` -> `block_q8_1 [m, k/32]` (`kquant_q8_1_rows_bf16`), one warp per
79/// block. Fails unless `k % 32 == 0`.
80pub fn kquant_q8_1_rows(
81    gpu: &dyn GpuBackend,
82    kernel: KernelHandle,
83    x_bf16: DevicePtr,
84    out_q8: DevicePtr,
85    m: u32,
86    k: u32,
87    stream: u64,
88) -> Result<()> {
89    anyhow::ensure!(
90        k.is_multiple_of(32),
91        "kquant_q8_1_rows: k={k} is not a multiple of 32"
92    );
93    let warps = m * (k / 32);
94    KernelLaunch::new(gpu, kernel)
95        .grid([div_ceil(warps * 32, 128), 1, 1])
96        .block([128, 1, 1])
97        .arg_ptr(x_bf16)
98        .arg_ptr(out_q8)
99        .arg_u32(k)
100        .arg_u32(m)
101        .launch(stream)
102}
103
104/// 2026-09-25: `moe_v41_swiglu` and [`kquant_q8_1_rows`] in one launch
105/// (`kquant_swiglu_q8_1_rows_bf16`): `h[rows, inter] = bf16(swiglu(gate, up) * w[row])`, then
106/// the q8_1 blocks of those bf16 rows. `kquant_fold_tests` compares both outputs byte for byte
107/// with the two launches. `w` may be null. Fails unless `inter % 32 == 0`.
108#[allow(clippy::too_many_arguments)]
109pub fn kquant_swiglu_q8_1_rows(
110    gpu: &dyn GpuBackend,
111    kernel: KernelHandle,
112    gate: DevicePtr,
113    up: DevicePtr,
114    w: DevicePtr,
115    h_bf16: DevicePtr,
116    out_q8: DevicePtr,
117    rows: u32,
118    inter: u32,
119    limit: f32,
120    stream: u64,
121) -> Result<()> {
122    anyhow::ensure!(
123        inter.is_multiple_of(32),
124        "kquant_swiglu_q8_1_rows: inter={inter} is not a multiple of 32"
125    );
126    let warps = rows * (inter / 32);
127    KernelLaunch::new(gpu, kernel)
128        .grid([div_ceil(warps * 32, 128), 1, 1])
129        .block([128, 1, 1])
130        .arg_ptr(gate)
131        .arg_ptr(up)
132        .arg_ptr(w)
133        .arg_ptr(h_bf16)
134        .arg_ptr(out_q8)
135        .arg_u32(rows)
136        .arg_u32(inter)
137        .arg_f32(limit)
138        .launch(stream)
139}
140
141/// 2026-09-25: Decode GEMV on raw blocks: `out[m][n] = sum_k xq8[m][k] * W[n][k]`
142/// (`kquant_mmvq_q2_k` / `kquant_mmvq_q3_k`), one 4-warp block per output row. `w_blocks` is
143/// `[n][k/256]` raw super-blocks; `y_q8` is the `block_q8_1 [m][k/32]` buffer; `out_bf16` is
144/// `[m][n]`.
145pub fn kquant_mmvq(
146    gpu: &dyn GpuBackend,
147    kernel: KernelHandle,
148    w_blocks: DevicePtr,
149    y_q8: DevicePtr,
150    out_bf16: DevicePtr,
151    n: u32,
152    k: u32,
153    m: u32,
154    stream: u64,
155) -> Result<()> {
156    anyhow::ensure!((1..=8).contains(&m), "kquant_mmvq: m={m} outside 1..=8");
157    anyhow::ensure!(
158        k.is_multiple_of(QK_K),
159        "kquant_mmvq: k={k} is not a multiple of {QK_K}"
160    );
161    KernelLaunch::new(gpu, kernel)
162        .grid([n, 1, 1])
163        .block([32, 4, 1])
164        .arg_ptr(w_blocks)
165        .arg_ptr(y_q8)
166        .arg_ptr(out_bf16)
167        .arg_u32(k)
168        .arg_u32(n)
169        .arg_u32(m)
170        .launch(stream)
171}
172
173#[allow(clippy::too_many_arguments)]
174/// 2026-09-25: `kquant_mmvq_*_experts`: one launch over `n_experts` experts whose block
175/// pointers sit in the device table `vxs`. The activation for expert `e` is
176/// `y_q8 + e * y_stride_bytes` (0 = shared), its output `out_bf16 + e * m * n`. Each expert's
177/// rows go through the same device function as [`kquant_mmvq`].
178pub fn kquant_mmvq_experts(
179    gpu: &dyn GpuBackend,
180    kernel: KernelHandle,
181    vxs: DevicePtr,
182    y_q8: DevicePtr,
183    out_bf16: DevicePtr,
184    n: u32,
185    k: u32,
186    m: u32,
187    n_experts: u32,
188    y_stride_bytes: u32,
189    stream: u64,
190) -> Result<()> {
191    anyhow::ensure!(
192        (1..=8).contains(&m),
193        "kquant_mmvq_experts: m={m} outside 1..=8"
194    );
195    anyhow::ensure!(
196        k.is_multiple_of(QK_K),
197        "kquant_mmvq_experts: k={k} is not a multiple of {QK_K}"
198    );
199    KernelLaunch::new(gpu, kernel)
200        .grid([n, n_experts, 1])
201        .block([32, 4, 1])
202        .arg_ptr(vxs)
203        .arg_ptr(y_q8)
204        .arg_ptr(out_bf16)
205        .arg_u32(k)
206        .arg_u32(n)
207        .arg_u32(m)
208        .arg_u32(y_stride_bytes)
209        .launch(stream)
210}
211
212/// 2026-09-25: `kquant_mmvq_*_w`: the warp-per-row GEMV (four rows a block, reduced by warp
213/// shuffle), with the same arguments as [`kquant_mmvq`].
214pub fn kquant_mmvq_w(
215    gpu: &dyn GpuBackend,
216    kernel: KernelHandle,
217    w_blocks: DevicePtr,
218    y_q8: DevicePtr,
219    out_bf16: DevicePtr,
220    n: u32,
221    k: u32,
222    m: u32,
223    stream: u64,
224) -> Result<()> {
225    anyhow::ensure!((1..=8).contains(&m), "kquant_mmvq_w: m={m} outside 1..=8");
226    anyhow::ensure!(
227        k.is_multiple_of(QK_K),
228        "kquant_mmvq_w: k={k} is not a multiple of {QK_K}"
229    );
230    KernelLaunch::new(gpu, kernel)
231        .grid([div_ceil(n, 4), 1, 1])
232        .block([32, 4, 1])
233        .arg_ptr(w_blocks)
234        .arg_ptr(y_q8)
235        .arg_ptr(out_bf16)
236        .arg_u32(k)
237        .arg_u32(n)
238        .arg_u32(m)
239        .launch(stream)
240}
241
242/// 2026-09-25: `kquant_mmvq_q2_k_groups_w`: `n_groups` independent `[n, k]` GEMVs in one launch.
243/// Group `g` takes weight rows `[g*n, (g+1)*n)` (`w_blocks` is `[n_groups * n][k/256]`) against
244/// activation columns `[g*k, (g+1)*k)` of the `block_q8_1 [m][n_groups * k / 32]` buffer
245/// `y_q8`, and writes columns `[g*n, (g+1)*n)` of `out_bf16` (`[m][n_groups * n]`). Each group
246/// runs the per-row code of [`kquant_mmvq_w`] with offsets.
247#[allow(clippy::too_many_arguments)]
248pub fn kquant_mmvq_groups_w(
249    gpu: &dyn GpuBackend,
250    kernel: KernelHandle,
251    w_blocks: DevicePtr,
252    y_q8: DevicePtr,
253    out_bf16: DevicePtr,
254    n: u32,
255    k: u32,
256    m: u32,
257    n_groups: u32,
258    stream: u64,
259) -> Result<()> {
260    anyhow::ensure!(
261        (1..=8).contains(&m),
262        "kquant_mmvq_groups_w: m={m} outside 1..=8"
263    );
264    anyhow::ensure!(
265        k.is_multiple_of(QK_K),
266        "kquant_mmvq_groups_w: k={k} is not a multiple of {QK_K}"
267    );
268    anyhow::ensure!(n_groups >= 1, "kquant_mmvq_groups_w: no groups");
269    KernelLaunch::new(gpu, kernel)
270        .grid([div_ceil(n, 4), n_groups, 1])
271        .block([32, 4, 1])
272        .arg_ptr(w_blocks)
273        .arg_ptr(y_q8)
274        .arg_ptr(out_bf16)
275        .arg_u32(k)
276        .arg_u32(n)
277        .arg_u32(m)
278        .arg_u32(n_groups)
279        .launch(stream)
280}
281
282/// 2026-09-25: `kquant_mmvq_q2_k_pair_w`: two `[n_i, k]` Q2_K projections of one activation
283/// (`y_q8`, `block_q8_1 [m][k/32]`) in one launch, `out_i` = `[m][n_i]` bf16. Grid y picks the
284/// projection, and each runs the same device function as [`kquant_mmvq_w`].
285#[allow(clippy::too_many_arguments)]
286pub fn kquant_mmvq_pair_w(
287    gpu: &dyn GpuBackend,
288    kernel: KernelHandle,
289    (w0, out0, n0): (DevicePtr, DevicePtr, u32),
290    (w1, out1, n1): (DevicePtr, DevicePtr, u32),
291    y_q8: DevicePtr,
292    k: u32,
293    m: u32,
294    stream: u64,
295) -> Result<()> {
296    anyhow::ensure!(
297        (1..=8).contains(&m),
298        "kquant_mmvq_pair_w: m={m} outside 1..=8"
299    );
300    anyhow::ensure!(
301        k.is_multiple_of(QK_K),
302        "kquant_mmvq_pair_w: k={k} is not a multiple of {QK_K}"
303    );
304    anyhow::ensure!(n0 >= 1 && n1 >= 1, "kquant_mmvq_pair_w: empty projection");
305    KernelLaunch::new(gpu, kernel)
306        .grid([div_ceil(n0.max(n1), 4), 2, 1])
307        .block([32, 4, 1])
308        .arg_ptr(w0)
309        .arg_ptr(out0)
310        .arg_u32(n0)
311        .arg_ptr(w1)
312        .arg_ptr(out1)
313        .arg_u32(n1)
314        .arg_ptr(y_q8)
315        .arg_u32(k)
316        .arg_u32(m)
317        .launch(stream)
318}
319
320/// 2026-09-25: [`kquant_mmvq_experts_w`] with `nwarps` rows a block (the `_w2` / `_w8` entries
321/// for 2 / 8; the plain `_w` entry is 4). One warp computes each row in the same way for any
322/// `nwarps`. Fails unless `nwarps` is 2, 4 or 8.
323#[allow(clippy::too_many_arguments)]
324pub fn kquant_mmvq_experts_wn(
325    gpu: &dyn GpuBackend,
326    kernel: KernelHandle,
327    w_tables: DevicePtr,
328    y_q8: DevicePtr,
329    out_bf16: DevicePtr,
330    n: u32,
331    k: u32,
332    m: u32,
333    n_experts: u32,
334    y_stride_bytes: u32,
335    nwarps: u32,
336    stream: u64,
337) -> Result<()> {
338    anyhow::ensure!(
339        (1..=8).contains(&m),
340        "kquant_mmvq_experts_wn: m={m} outside 1..=8"
341    );
342    anyhow::ensure!(
343        k.is_multiple_of(QK_K),
344        "kquant_mmvq_experts_wn: k={k} is not a multiple of {QK_K}"
345    );
346    anyhow::ensure!(
347        matches!(nwarps, 2 | 4 | 8),
348        "kquant_mmvq_experts_wn: nwarps={nwarps} not in 2/4/8"
349    );
350    KernelLaunch::new(gpu, kernel)
351        .grid([div_ceil(n, nwarps), n_experts, 1])
352        .block([32, nwarps, 1])
353        .arg_ptr(w_tables)
354        .arg_ptr(y_q8)
355        .arg_ptr(out_bf16)
356        .arg_u32(k)
357        .arg_u32(n)
358        .arg_u32(m)
359        .arg_u32(y_stride_bytes)
360        .launch(stream)
361}
362
363/// 2026-09-25: `kquant_mmvq_*_experts_w`: the warp-per-row expert batch, with the same arguments
364/// as [`kquant_mmvq_experts`].
365pub fn kquant_mmvq_experts_w(
366    gpu: &dyn GpuBackend,
367    kernel: KernelHandle,
368    vxs: DevicePtr,
369    y_q8: DevicePtr,
370    out_bf16: DevicePtr,
371    n: u32,
372    k: u32,
373    m: u32,
374    n_experts: u32,
375    y_stride_bytes: u32,
376    stream: u64,
377) -> Result<()> {
378    anyhow::ensure!(
379        (1..=8).contains(&m),
380        "kquant_mmvq_experts_w: m={m} outside 1..=8"
381    );
382    anyhow::ensure!(
383        k.is_multiple_of(QK_K),
384        "kquant_mmvq_experts_w: k={k} is not a multiple of {QK_K}"
385    );
386    KernelLaunch::new(gpu, kernel)
387        .grid([div_ceil(n, 4), n_experts, 1])
388        .block([32, 4, 1])
389        .arg_ptr(vxs)
390        .arg_ptr(y_q8)
391        .arg_ptr(out_bf16)
392        .arg_u32(k)
393        .arg_u32(n)
394        .arg_u32(m)
395        .arg_u32(y_stride_bytes)
396        .launch(stream)
397}
398
399/// 2026-09-25: Prefill GEMM on raw blocks through the vendored MMQ engine:
400/// `out[m][n]` (bf16) = `A_q8[m][k]` x `W[n][k]`. `a_q8` must be in the type's MMQ q8_1 layout
401/// (D2S6 for Q2_K, D4 for Q3_K), and `smem` is the matching `*_MMQ_SMEM`. `kernel_nc` serves
402/// `n % 128 == 0`, and `kernel_wc` (bounds-checked) any other `n`. The kernel's arguments are
403/// `(x = weights, y = activations, dst, nrows_x = n, ncols_dst = m, ncols_x = k,
404/// stride_row_x = k / 256, ncols_y = m, stride_col_dst = n)`.
405pub fn kquant_mmq_gemm(
406    gpu: &dyn GpuBackend,
407    kernel_nc: KernelHandle,
408    kernel_wc: KernelHandle,
409    a_q8: DevicePtr,
410    w_blocks: DevicePtr,
411    out_bf16: DevicePtr,
412    m: u32,
413    n: u32,
414    k: u32,
415    smem: u32,
416    stream: u64,
417) -> Result<()> {
418    anyhow::ensure!(
419        k.is_multiple_of(QK_K),
420        "kquant_mmq_gemm: k={k} is not a multiple of {QK_K}"
421    );
422    let kernel = if n.is_multiple_of(128) {
423        kernel_nc
424    } else {
425        kernel_wc
426    };
427    KernelLaunch::new(gpu, kernel)
428        .grid([div_ceil(n, 128), div_ceil(m, 128), 1])
429        .block([32, 8, 1])
430        .shared_mem(smem)
431        .arg_ptr(w_blocks)
432        .arg_ptr(a_q8)
433        .arg_ptr(out_bf16)
434        .arg_u32(n)
435        .arg_u32(m)
436        .arg_u32(k)
437        .arg_u32(k / QK_K)
438        .arg_u32(m)
439        .arg_u32(n)
440        .launch(stream)
441}
442
443#[cfg(test)]
444mod tests {
445    use super::*;
446
447    #[test]
448    fn sizes_match_the_block_geometry() {
449        // 2026-09-25: A DeepSeek-V4.1 expert's gate (2304 x 5120, Q2_K) and down (5120 x 2304,
450        // Q3_K) projections (`moe_intermediate_size`, `hidden_dim` in its MODEL.toml).
451        assert_eq!(kquant_weight_bytes(2304, 5120, Q2K_BLOCK_BYTES), 3_870_720);
452        assert_eq!(kquant_weight_bytes(5120, 2304, Q3K_BLOCK_BYTES), 5_068_800);
453        assert_eq!(kquant_q8_1_rows_bytes(1, 5120), 160 * 36);
454        assert_eq!(kquant_mmq_act_bytes(200, 5120), 256 * 40 * 144);
455        assert_eq!((Q2K_MMQ_SMEM, Q3K_MMQ_SMEM), (70_144, 61_952));
456    }
457}
metrale_model_layers/layers/ops/kquant_mmq.rs

metrale_model_layers/layers/ops/
kquant_mmq.rs
