spark_model/layers/ops/
q4k_mmq.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//
3// Launchers for the vendored llama Q4_K MMQ FFN prefill GEMM (ATLAS_FFN_MMQ).
4// Kernels in kernels/gb10/qwen3.6-27b/nvfp4/q4k_mmq.cu + q4k_quantize.cu (verified
5// 54.9/53.7 TFLOP/s gate/up·down, +25%/+10% vs faith2, rel_err 6-7e-3).
6// Pipeline: weights NVFP4 -> dequant_nvfp4_to_bf16 -> q4k_quantize (at load); per-prefill
7// activation bf16 -> q8_1_mmq, then MMQ -> bf16 (fused store, no cast).
8use anyhow::Result;
9use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
10use spark_runtime::kernel_args::{KernelLaunch, div_ceil};
11
12/// Q4_K block size (256 weights -> 144-byte block_q4_K).
13pub const QK_K: u32 = 256;
14/// sizeof(block_q4_K) bytes.
15pub const Q4K_BLOCK_BYTES: usize = 144;
16/// Dynamic shared memory for the Q4_K MMQ kernel (mmq_x=mmq_y=128, GB10). >48KB -> registry sets attr.
17pub const Q4K_MMQ_SMEM: u32 = 57856;
18const CUDA_QUANTIZE_BLOCK_SIZE_MMQ: u32 = 128;
19
20/// Bytes for the Q4_K-quantized form of an [nrows, n_per_row] weight (n_per_row % 256 == 0).
21pub fn q4k_weight_bytes(nrows: u32, n_per_row: u32) -> usize {
22    (nrows as usize) * (n_per_row as usize / QK_K as usize) * Q4K_BLOCK_BYTES
23}
24
25/// q8_1_mmq activation scratch bytes for [m, k]; generous (kpad rounded to 256).
26pub fn q8_1_scratch_bytes(m: u32, k: u32) -> usize {
27    let kpad = div_ceil(k, QK_K) * QK_K;
28    (m as usize) * (kpad as usize) * 4 + (1 << 20)
29}
30
31/// Dequantize NVFP4 weight [n, k] (packed E2M1 + E4M3 group scales + per-tensor scale2) -> bf16 [n, k].
32pub fn dequant_nvfp4_to_bf16(
33    gpu: &dyn GpuBackend,
34    kernel: KernelHandle,
35    packed: DevicePtr,
36    scales: DevicePtr,
37    out_bf16: DevicePtr,
38    scale2: f32,
39    n: u32,
40    k: u32,
41    stream: u64,
42) -> Result<()> {
43    KernelLaunch::new(gpu, kernel)
44        .grid([n, 1, 1])
45        .block([256, 1, 1])
46        .arg_ptr(packed)
47        .arg_ptr(scales)
48        .arg_ptr(out_bf16)
49        .arg_f32(scale2)
50        .arg_u32(n)
51        .arg_u32(k)
52        .launch(stream)
53}
54
55/// Quantize bf16 weights [nrows, n_per_row] -> GGML block_q4_K (at model load).
56pub fn quantize_weight_q4k(
57    gpu: &dyn GpuBackend,
58    kernel: KernelHandle,
59    input_bf16: DevicePtr,
60    out_q4k: DevicePtr,
61    nrows: u32,
62    n_per_row: u32,
63    stream: u64,
64) -> Result<()> {
65    let total_sb = (nrows as u64) * (n_per_row as u64 / QK_K as u64);
66    let grid_x = div_ceil(total_sb as u32, 128);
67    KernelLaunch::new(gpu, kernel)
68        .grid([grid_x, 1, 1])
69        .block([128, 1, 1])
70        .arg_ptr(input_bf16)
71        .arg_ptr(out_q4k)
72        .arg_u32(nrows)
73        .arg_u32(n_per_row)
74        .launch(stream)
75}
76
77/// Quantize bf16 activations [m, k] -> q8_1_mmq (DS4 layout) into `out_q8`.
78pub fn quantize_act_q8_1(
79    gpu: &dyn GpuBackend,
80    kernel: KernelHandle, // atlas_q8_1_quantize_ds4_bf16
81    input_bf16: DevicePtr,
82    out_q8: DevicePtr,
83    m: u32,
84    k: u32,
85    stream: u64,
86) -> Result<()> {
87    let kpad = div_ceil(k, QK_K) * QK_K;
88    let grid_y = div_ceil(kpad, 4 * CUDA_QUANTIZE_BLOCK_SIZE_MMQ);
89    KernelLaunch::new(gpu, kernel)
90        .grid([m, grid_y, 1])
91        .block([CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1])
92        .arg_ptr(input_bf16)
93        .arg_ptr(out_q8)
94        .arg_u64(k as u64) // ne00
95        .arg_u64(k as u64) // s01 (contiguous rows)
96        .arg_u64(kpad as u64) // ne0
97        .arg_u32(m) // ne1
98        .launch(stream)
99}
100
101/// Q4_K MMQ GEMM: C\[m,n\] (bf16) = A_q8\[m,k\] x W_q4k\[n,k\]. Fused bf16 store.
102pub fn q4k_mmq_gemm(
103    gpu: &dyn GpuBackend,
104    kernel_nc: KernelHandle, // atlas_q4k_mmq128_nc
105    kernel_wc: KernelHandle, // atlas_q4k_mmq128_wc
106    a_q8: DevicePtr,         // q8_1_mmq activations
107    w_q4k: DevicePtr,        // block_q4_K weights [n, k]
108    out_bf16: DevicePtr,
109    m: u32,
110    n: u32,
111    k: u32,
112    stream: u64,
113) -> Result<()> {
114    let kernel = if !n.is_multiple_of(128) {
115        kernel_wc
116    } else {
117        kernel_nc
118    };
119    KernelLaunch::new(gpu, kernel)
120        .grid([div_ceil(n, 128), div_ceil(m, 128), 1])
121        .block([32, 8, 1])
122        .shared_mem(Q4K_MMQ_SMEM)
123        .arg_ptr(w_q4k) // x = weights
124        .arg_ptr(a_q8) // y = q8_1 activations
125        .arg_ptr(out_bf16) // dst
126        .arg_u32(n) // nrows_x
127        .arg_u32(m) // ncols_dst
128        .arg_u32(k) // ncols_x
129        .arg_u32(k / QK_K) // stride_row_x = K/256
130        .arg_u32(m) // ncols_y
131        .arg_u32(n) // stride_col_dst
132        .launch(stream)
133}
spark_model/layers/ops/q4k_mmq.rs

spark_model/layers/ops/
q4k_mmq.rs
