metrale_model_layers/layers/ops/
gemm_fp4.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Launchers for the W4A4 prefill path: BF16-to-NVFP4 activation
4//! quantization and the NVFP4 x NVFP4 tensor-core GEMMs.
5//!
6//! Owner: model-layers ops.
7//! Invariants: none beyond the types.
8
9#![allow(unused_imports)]
10
11use anyhow::Result;
12use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
13use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
14
15use crate::weight_map::QuantizedWeight;
16
17use super::*;
18
19/// 2026-09-25: Quantize a BF16 `[M, K]` matrix to NVFP4 with `scale2 = 1.0`:
20/// packed E2M1 `[M, K/2]` and one E4M3 scale per 16 elements `[M, K/16]`. One
21/// block per row; the kernel's `N` parameter is the row count. Produces the
22/// W4A4 prefill activations.
23#[allow(clippy::too_many_arguments)]
24pub fn quantize_bf16_to_nvfp4(
25    gpu: &dyn GpuBackend,
26    kernel: KernelHandle,
27    input: DevicePtr,
28    packed_out: DevicePtr,
29    scale_out: DevicePtr,
30    m: u32,
31    k: u32,
32    stream: u64,
33) -> Result<()> {
34    KernelLaunch::new(gpu, kernel)
35        .grid([m, 1, 1])
36        .block([128, 1, 1])
37        .arg_ptr(input)
38        .arg_ptr(packed_out)
39        .arg_ptr(scale_out)
40        .arg_f32(1.0)
41        .arg_u32(m)
42        .arg_u32(k)
43        .launch(stream)
44}
45
46/// 2026-09-25: W4A4 GEMM on FP4 tensor cores (kernel `w4a4_gemm`). The
47/// activation comes from [`quantize_bf16_to_nvfp4`] and is launched with
48/// `scaleA2 = 1.0`; the weight is an NVFP4 `QuantizedWeight`; the output is
49/// BF16 `[M, N]`.
50#[allow(clippy::too_many_arguments)]
51pub fn w4a4_gemm(
52    gpu: &dyn GpuBackend,
53    kernel: KernelHandle,
54    a_packed: DevicePtr,
55    a_scale: DevicePtr,
56    weight: &QuantizedWeight,
57    output: DevicePtr,
58    m: u32,
59    n: u32,
60    k: u32,
61    stream: u64,
62) -> Result<()> {
63    KernelLaunch::new(gpu, kernel)
64        .grid([div_ceil(n, 128), div_ceil(m, 128), 1])
65        .block([256, 1, 1])
66        .arg_ptr(a_packed)
67        .arg_ptr(a_scale)
68        .arg_ptr(weight.weight)
69        .arg_ptr(weight.weight_scale)
70        .arg_ptr(output)
71        .arg_f32(1.0)
72        .arg_f32(weight.weight_scale_2)
73        .arg_u32(m)
74        .arg_u32(n)
75        .arg_u32(k)
76        .launch(stream)
77}
78
79#[allow(clippy::too_many_arguments)]
80/// 2026-09-25: [`w4a4_gemm`] through the `w4a4_gemm_mfast` kernel, with M on
81/// the fast grid axis (`blockIdx.x` = M block), so consecutive CTAs share a B
82/// panel.
83pub fn w4a4_gemm_mfast(
84    gpu: &dyn GpuBackend,
85    kernel: KernelHandle,
86    a_packed: DevicePtr,
87    a_scale: DevicePtr,
88    weight: &QuantizedWeight,
89    output: DevicePtr,
90    m: u32,
91    n: u32,
92    k: u32,
93    stream: u64,
94) -> Result<()> {
95    KernelLaunch::new(gpu, kernel)
96        .grid([div_ceil(m, 128), div_ceil(n, 128), 1])
97        .block([128, 1, 1])
98        .arg_ptr(a_packed)
99        .arg_ptr(a_scale)
100        .arg_ptr(weight.weight)
101        .arg_ptr(weight.weight_scale)
102        .arg_ptr(output)
103        .arg_f32(1.0)
104        .arg_f32(weight.weight_scale_2)
105        .arg_u32(m)
106        .arg_u32(n)
107        .arg_u32(k)
108        .launch(stream)
109}
metrale_model_layers/layers/ops/gemm_fp4.rs

metrale_model_layers/layers/ops/
gemm_fp4.rs
