spark_model/layers/ops/
gemm_fp4.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Native FP4 (NVFP4 / mxf4nvf4) prefill launchers: activation quantization
4//! and the W4A4 tensor-core GEMMs. Split from `gemm_dense.rs` (500-LoC cap).
5
6#![allow(unused_imports)]
7
8use anyhow::Result;
9use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
10use spark_runtime::kernel_args::{KernelLaunch, div_ceil};
11
12use crate::weight_map::QuantizedWeight;
13
14use super::*;
15
16/// Quantize a BF16 [M, K] matrix to NVFP4 (single-level, scale2=1.0): packed E2M1
17/// `[M, K/2]` + per-group-16 E4M3 scales `[M, K/16]`. Prepares W4A4 prefill
18/// activations. Grid = M rows (one block/row), block 128 (threads stride groups).
19#[allow(clippy::too_many_arguments)]
20pub fn quantize_bf16_to_nvfp4(
21    gpu: &dyn GpuBackend,
22    kernel: KernelHandle,
23    input: DevicePtr,
24    packed_out: DevicePtr,
25    scale_out: DevicePtr,
26    m: u32,
27    k: u32,
28    stream: u64,
29) -> Result<()> {
30    KernelLaunch::new(gpu, kernel)
31        .grid([m, 1, 1])
32        .block([128, 1, 1])
33        .arg_ptr(input)
34        .arg_ptr(packed_out)
35        .arg_ptr(scale_out)
36        .arg_f32(1.0) // scale2 = 1.0 (single-level; activation range fits E4M3 group scales)
37        .arg_u32(m) // kernel's N param = rows = tokens
38        .arg_u32(k)
39        .launch(stream)
40}
41
42/// W4A4 NVFP4 prefill GEMM (native FP4 tensor cores, sm_121a). Activation is
43/// pre-quantized NVFP4 (`a_packed`/`a_scale`, scale2=1.0); weight is the native
44/// NVFP4 `QuantizedWeight`. Output BF16 [M, N]. See kernels/.../w4a4_gemm.cu.
45/// Grid: (ceil(N/128), ceil(M/128), 1)  Block: (256, 1, 1).
46#[allow(clippy::too_many_arguments)]
47pub fn w4a4_gemm(
48    gpu: &dyn GpuBackend,
49    kernel: KernelHandle,
50    a_packed: DevicePtr,
51    a_scale: DevicePtr,
52    weight: &QuantizedWeight,
53    output: DevicePtr,
54    m: u32,
55    n: u32,
56    k: u32,
57    stream: u64,
58) -> Result<()> {
59    KernelLaunch::new(gpu, kernel)
60        .grid([div_ceil(n, 128), div_ceil(m, 128), 1])
61        .block([256, 1, 1])
62        .arg_ptr(a_packed)
63        .arg_ptr(a_scale)
64        .arg_ptr(weight.weight)
65        .arg_ptr(weight.weight_scale)
66        .arg_ptr(output)
67        .arg_f32(1.0) // scaleA2 (activation single-level)
68        .arg_f32(weight.weight_scale_2) // scaleB2 (weight per-tensor)
69        .arg_u32(m)
70        .arg_u32(n)
71        .arg_u32(k)
72        .launch(stream)
73}
74
75/// W4A16 GEMM with N_TILE=128: same kernel signature, wider N tile.
76///
77/// Grid: (ceil(N/128), ceil(M/64), 1)  Block: (128, 1, 1)
78#[allow(clippy::too_many_arguments)]
79/// `w4a4_gemm_mfast`: same W4A4 GEMM with M on the fast grid axis, so the
80/// M-blocks sharing a B panel are co-resident and B streams from DRAM once.
81pub fn w4a4_gemm_mfast(
82    gpu: &dyn GpuBackend,
83    kernel: KernelHandle,
84    a_packed: DevicePtr,
85    a_scale: DevicePtr,
86    weight: &QuantizedWeight,
87    output: DevicePtr,
88    m: u32,
89    n: u32,
90    k: u32,
91    stream: u64,
92) -> Result<()> {
93    KernelLaunch::new(gpu, kernel)
94        .grid([div_ceil(m, 128), div_ceil(n, 128), 1])
95        .block([128, 1, 1])
96        .arg_ptr(a_packed)
97        .arg_ptr(a_scale)
98        .arg_ptr(weight.weight)
99        .arg_ptr(weight.weight_scale)
100        .arg_ptr(output)
101        .arg_f32(1.0) // scaleA2 (activation single-level)
102        .arg_f32(weight.weight_scale_2) // scaleB2 (weight per-tensor)
103        .arg_u32(m)
104        .arg_u32(n)
105        .arg_u32(k)
106        .launch(stream)
107}
spark_model/layers/ops/gemm_fp4.rs

spark_model/layers/ops/
gemm_fp4.rs
