metrale_model_layers/layers/ops/
gemv_sw.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Decode NVFP4 W4A16 GEMV launchers whose grid is tied to the kernel's `N_PER_BLOCK` / `N_PER_BLOCK_SW` defines.
4//!
5//! Owner: model-layers ops.
6//! Invariants:
7//! - `w4a16_gemv` computes 4 outputs per 256-thread block (64 threads each) and
8//!   `w4a16_gemv_sw` computes 8 (one warp each), so a launch that swaps the kernel must also
9//!   swap the grid. [`w4a16_decode_gemv`] picks both together.
10//! - `gemv_sw_tests::cuda_n_per_block_matches_rust_ssot` checks that every compiled copy of the
11//!   kernel sources defines the two constants below.
12
13use anyhow::Result;
14use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
15use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
16
17use crate::weight_map::QuantizedWeight;
18
19/// 2026-09-25: Outputs per 256-thread block of `w4a16_gemv`; equals `#define N_PER_BLOCK` in
20/// kernels/gb10/common/w4a16_gemv.cu.
21pub const W4A16_GEMV_OUTS_PER_BLOCK: u32 = 4;
22
23/// 2026-09-25: Outputs per 256-thread block of the single-warp kernels (`w4a16_gemv_sw` and the
24/// `_sw` kernels in w4a16_gemv_fused.cu); equals `#define N_PER_BLOCK_SW`.
25pub const W4A16_GEMV_SW_OUTS_PER_BLOCK: u32 = 8;
26
27pub fn w4a16_gemv_grid_x(n: u32) -> u32 {
28    div_ceil(n, W4A16_GEMV_OUTS_PER_BLOCK)
29}
30
31pub fn w4a16_gemv_sw_grid_x(n: u32) -> u32 {
32    div_ceil(n, W4A16_GEMV_SW_OUTS_PER_BLOCK)
33}
34
35/// 2026-09-25: Polarity of the `METRALE_NO_GEMV_SW` kill switch, given the variable's value: the
36/// single-warp GEMV is on unless the value is exactly `"1"`, so `=0` and an empty value leave
37/// it on.
38pub fn gemv_sw_from(no_gemv_sw: Option<&str>) -> bool {
39    no_gemv_sw != Some("1")
40}
41
42/// 2026-09-25: Use the single-warp kernel only when the lever is on and its handle resolved
43/// (a missing kernel is `KernelHandle(0)`).
44pub fn use_gemv_sw(lever: bool, sw_handle: KernelHandle) -> bool {
45    lever && sw_handle.0 != 0
46}
47
48/// 2026-09-25: Single-warp-per-output W4A16 GEMV for M=1, grid `ceil(N/8)`.
49#[allow(clippy::too_many_arguments)]
50pub fn w4a16_gemv_sw(
51    gpu: &dyn GpuBackend,
52    kernel: KernelHandle,
53    input: DevicePtr,
54    weight: &QuantizedWeight,
55    output: DevicePtr,
56    n: u32,
57    k: u32,
58    stream: u64,
59) -> Result<()> {
60    KernelLaunch::new(gpu, kernel)
61        .grid([w4a16_gemv_sw_grid_x(n), 1, 1])
62        .block([256, 1, 1])
63        .arg_ptr(input)
64        .arg_ptr(weight.weight)
65        .arg_ptr(weight.weight_scale)
66        .arg_f32(weight.weight_scale_2)
67        .arg_ptr(output)
68        .arg_u32(n)
69        .arg_u32(k)
70        .launch(stream)
71}
72
73/// 2026-09-25: The launch of [`w4a16_gemv_sw`] for callers that hold the NVFP4 packed weight,
74/// scale and global scale as separate values rather than a [`QuantizedWeight`] (the GLM-5.3
75/// MLP in `model-arch/src/glm5next_mlp/forward.rs`). It keeps the grid on
76/// [`w4a16_gemv_sw_grid_x`].
77#[allow(clippy::too_many_arguments)]
78pub fn w4a16_gemv_sw_raw(
79    gpu: &dyn GpuBackend,
80    kernel: KernelHandle,
81    input: DevicePtr,
82    packed: DevicePtr,
83    scale: DevicePtr,
84    scale_2: f32,
85    output: DevicePtr,
86    n: u32,
87    k: u32,
88    stream: u64,
89) -> Result<()> {
90    KernelLaunch::new(gpu, kernel)
91        .grid([w4a16_gemv_sw_grid_x(n), 1, 1])
92        .block([256, 1, 1])
93        .arg_ptr(input)
94        .arg_ptr(packed)
95        .arg_ptr(scale)
96        .arg_f32(scale_2)
97        .arg_ptr(output)
98        .arg_u32(n)
99        .arg_u32(k)
100        .launch(stream)
101}
102
103/// 2026-09-25: Decode GEMV: the single-warp kernel when [`use_gemv_sw`] allows it, otherwise
104/// the base `w4a16_gemv`, each with its own grid.
105#[allow(clippy::too_many_arguments)]
106pub fn w4a16_decode_gemv(
107    gpu: &dyn GpuBackend,
108    gemv: KernelHandle,
109    gemv_sw: KernelHandle,
110    use_sw: bool,
111    input: DevicePtr,
112    weight: &QuantizedWeight,
113    output: DevicePtr,
114    n: u32,
115    k: u32,
116    stream: u64,
117) -> Result<()> {
118    if use_gemv_sw(use_sw, gemv_sw) {
119        w4a16_gemv_sw(gpu, gemv_sw, input, weight, output, n, k, stream)
120    } else {
121        super::quant_dispatch::w4a16_gemv(gpu, gemv, input, weight, output, n, k, stream)
122    }
123}
124
125#[cfg(test)]
126#[path = "gemv_sw_tests.rs"]
127mod tests;
metrale_model_layers/layers/ops/gemv_sw.rs

metrale_model_layers/layers/ops/
gemv_sw.rs
