spark_model/layers/ops/
gemv_sw.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Decode W4A16 GEMV launchers whose grid is coupled to the CUDA
4//! `N_PER_BLOCK` / `N_PER_BLOCK_SW` defines.
5//!
6//! The single-warp kernel (`w4a16_gemv_sw`) is bit-identical to the 64-thread
7//! base (`examples/w4a16_gemv_sw_microtest.rs`). Shipping it as the default
8//! decode GEMV is a free occupancy win — **if and only if** the launch grid
9//! stays coupled: 8 outputs/block vs the base kernel's 4. Swapping the kernel
10//! without swapping the grid writes the wrong outputs.
11
12use anyhow::Result;
13use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
14use spark_runtime::kernel_args::{KernelLaunch, div_ceil};
15
16use crate::weight_map::QuantizedWeight;
17
18/// Base `w4a16_gemv`: 4 outputs / 256-thread block.
19/// SSOT with `kernels/**/w4a16_gemv.cu` `#define N_PER_BLOCK 4`.
20pub const W4A16_GEMV_OUTS_PER_BLOCK: u32 = 4;
21
22/// Single-warp `w4a16_gemv_sw`: 8 outputs / 256-thread block.
23/// SSOT with `#define N_PER_BLOCK_SW 8`.
24pub const W4A16_GEMV_SW_OUTS_PER_BLOCK: u32 = 8;
25
26pub fn w4a16_gemv_grid_x(n: u32) -> u32 {
27    div_ceil(n, W4A16_GEMV_OUTS_PER_BLOCK)
28}
29
30pub fn w4a16_gemv_sw_grid_x(n: u32) -> u32 {
31    div_ceil(n, W4A16_GEMV_SW_OUTS_PER_BLOCK)
32}
33
34/// Kill-switch polarity for lossless SW GEMV. ON unless `ATLAS_NO_GEMV_SW` is
35/// exactly `"1"`. `=0` does **not** disable (same `== "1"` reading as
36/// `ATLAS_NO_LM_HEAD_BATCH_GEMV`).
37pub fn gemv_sw_from(no_gemv_sw: Option<&str>) -> bool {
38    no_gemv_sw != Some("1")
39}
40
41/// SW kernel when the model lever is on **and** the handle resolved.
42pub fn use_gemv_sw(lever: bool, sw_handle: KernelHandle) -> bool {
43    lever && sw_handle.0 != 0
44}
45
46/// Single-warp-per-output W4A16 GEMV (M=1). Grid: `(ceil(N/8), 1, 1)`.
47#[allow(clippy::too_many_arguments)]
48pub fn w4a16_gemv_sw(
49    gpu: &dyn GpuBackend,
50    kernel: KernelHandle,
51    input: DevicePtr,
52    weight: &QuantizedWeight,
53    output: DevicePtr,
54    n: u32,
55    k: u32,
56    stream: u64,
57) -> Result<()> {
58    KernelLaunch::new(gpu, kernel)
59        .grid([w4a16_gemv_sw_grid_x(n), 1, 1])
60        .block([256, 1, 1])
61        .arg_ptr(input)
62        .arg_ptr(weight.weight)
63        .arg_ptr(weight.weight_scale)
64        .arg_f32(weight.weight_scale_2)
65        .arg_ptr(output)
66        .arg_u32(n)
67        .arg_u32(k)
68        .launch(stream)
69}
70
71/// Same launch as [`w4a16_gemv_sw`] for callers that hold the NVFP4 operand triple as
72/// loose pointers rather than a [`QuantizedWeight`] (GLM-5.3's `Nvfp4Proj`).
73///
74/// Exists so the `ceil(N/8)` grid stays in this file — the one place that is SSOT with the
75/// kernel's `N_PER_BLOCK_SW`.
76#[allow(clippy::too_many_arguments)]
77pub fn w4a16_gemv_sw_raw(
78    gpu: &dyn GpuBackend,
79    kernel: KernelHandle,
80    input: DevicePtr,
81    packed: DevicePtr,
82    scale: DevicePtr,
83    scale_2: f32,
84    output: DevicePtr,
85    n: u32,
86    k: u32,
87    stream: u64,
88) -> Result<()> {
89    KernelLaunch::new(gpu, kernel)
90        .grid([w4a16_gemv_sw_grid_x(n), 1, 1])
91        .block([256, 1, 1])
92        .arg_ptr(input)
93        .arg_ptr(packed)
94        .arg_ptr(scale)
95        .arg_f32(scale_2)
96        .arg_ptr(output)
97        .arg_u32(n)
98        .arg_u32(k)
99        .launch(stream)
100}
101
102/// Decode GEMV: software-pipelined single-warp when the lever and handle agree.
103#[allow(clippy::too_many_arguments)]
104pub fn w4a16_decode_gemv(
105    gpu: &dyn GpuBackend,
106    gemv: KernelHandle,
107    gemv_sw: KernelHandle,
108    use_sw: bool,
109    input: DevicePtr,
110    weight: &QuantizedWeight,
111    output: DevicePtr,
112    n: u32,
113    k: u32,
114    stream: u64,
115) -> Result<()> {
116    if use_gemv_sw(use_sw, gemv_sw) {
117        w4a16_gemv_sw(gpu, gemv_sw, input, weight, output, n, k, stream)
118    } else {
119        super::quant_dispatch::w4a16_gemv(gpu, gemv, input, weight, output, n, k, stream)
120    }
121}
122
123#[cfg(test)]
124mod tests {
125    use super::*;
126    use spark_runtime::gpu::KernelHandle;
127    use spark_runtime::gpu::mock::MockGpuBackend;
128    use std::fs;
129    use std::path::{Path, PathBuf};
130
131    #[test]
132    fn gemv_sw_ships_on_and_only_the_one_value_kills() {
133        assert!(gemv_sw_from(None), "unset → ON");
134        assert!(gemv_sw_from(Some("0")), "`=0` is NOT off");
135        assert!(gemv_sw_from(Some("")), "empty is NOT off");
136        assert!(!gemv_sw_from(Some("1")), "`=1` is the kill");
137    }
138
139    #[test]
140    fn sw_requires_both_the_lever_and_a_live_handle() {
141        assert!(use_gemv_sw(true, KernelHandle(1)));
142        assert!(
143            !use_gemv_sw(true, KernelHandle(0)),
144            "missing kernel falls back"
145        );
146        assert!(!use_gemv_sw(false, KernelHandle(1)), "kill switch wins");
147        assert!(!use_gemv_sw(false, KernelHandle(0)));
148    }
149
150    #[test]
151    fn sw_grid_covers_every_output_and_is_half_base_when_n_divisible_by_8() {
152        for n in 1..=64 {
153            assert!(w4a16_gemv_sw_grid_x(n) * W4A16_GEMV_SW_OUTS_PER_BLOCK >= n);
154            assert!(w4a16_gemv_grid_x(n) * W4A16_GEMV_OUTS_PER_BLOCK >= n);
155        }
156        for n in [8u32, 16, 256, 5120, 14336] {
157            assert_eq!(
158                w4a16_gemv_sw_grid_x(n) * 2,
159                w4a16_gemv_grid_x(n),
160                "N={n}: SW is 8 outs/block, base is 4 — grid_x must be half"
161            );
162        }
163    }
164
165    #[test]
166    fn decode_dispatch_uses_the_selected_handle_and_matching_grid() {
167        for (lever, sw_handle, expected_handle, expected_grid_x) in [
168            (true, KernelHandle(22), 22, 2),
169            (false, KernelHandle(22), 11, 3),
170            (true, KernelHandle(0), 11, 3),
171        ] {
172            let gpu = MockGpuBackend::new();
173            w4a16_decode_gemv(
174                &gpu,
175                KernelHandle(11),
176                sw_handle,
177                lever,
178                DevicePtr::NULL,
179                &QuantizedWeight::null(),
180                DevicePtr::NULL,
181                9,
182                128,
183                0,
184            )
185            .unwrap();
186            let launches = gpu.launches_snapshot();
187            assert_eq!(launches.len(), 1);
188            assert_eq!(launches[0].func, expected_handle);
189            assert_eq!(launches[0].grid, [expected_grid_x, 1, 1]);
190            assert_eq!(launches[0].block, [256, 1, 1]);
191        }
192    }
193
194    fn kernel_root() -> PathBuf {
195        Path::new(env!("CARGO_MANIFEST_DIR")).join("../../kernels")
196    }
197
198    fn named_cu(file_name: &str) -> Vec<PathBuf> {
199        fn visit(d: &Path, name: &str, out: &mut Vec<PathBuf>) {
200            let Ok(rd) = std::fs::read_dir(d) else { return };
201            for e in rd.flatten() {
202                let p = e.path();
203                if p.is_dir() {
204                    visit(&p, name, out);
205                } else if p.file_name().is_some_and(|n| n == name) {
206                    out.push(p);
207                }
208            }
209        }
210        let root = kernel_root();
211        let mut files = Vec::new();
212        visit(&root, file_name, &mut files);
213        files.sort();
214        files
215    }
216
217    /// POSITIVE: every copy of the GEMV sources pins the same occupancy
218    /// constants the Rust launchers use. A new backend copy that changes
219    /// `N_PER_BLOCK_SW` without updating the launcher writes the wrong N
220    /// slice — silent, not a CUDA error.
221    ///
222    /// PROVEN BY: changing either `#define` in one `.cu` copy turns this red.
223    #[test]
224    fn cuda_n_per_block_matches_rust_ssot() {
225        let gemv = named_cu("w4a16_gemv.cu");
226        assert!(
227            gemv.len() >= 3,
228            "expected gb10 + strix + strix-hip copies, got {gemv:?}"
229        );
230        let want_base = format!("#define N_PER_BLOCK {W4A16_GEMV_OUTS_PER_BLOCK}");
231        let want_sw = format!("#define N_PER_BLOCK_SW {W4A16_GEMV_SW_OUTS_PER_BLOCK}");
232        for p in &gemv {
233            let src = fs::read_to_string(p).unwrap();
234            assert!(
235                src.contains(&want_base),
236                "{} missing {want_base}",
237                p.display()
238            );
239            assert!(src.contains(&want_sw), "{} missing {want_sw}", p.display());
240        }
241        let fused = named_cu("w4a16_gemv_fused.cu");
242        assert!(
243            !fused.is_empty(),
244            "dual_sw / silu_input_sw live in w4a16_gemv_fused.cu"
245        );
246        for p in &fused {
247            let src = fs::read_to_string(p).unwrap();
248            assert!(src.contains(&want_sw), "{} missing {want_sw}", p.display());
249        }
250    }
251
252    /// POSITIVE: SW GEMV must share the 2-chunk K16 pipeline with the 64-thread
253    /// kernel. A stride-64 sequential `acc += a*w` copy was 1 ULP lossy on GB10
254    /// (`w4a16_gemv_sw_microtest`: gdn in_proj 99.992%, K-tail 99.976%).
255    ///
256    /// PROVEN BY: restoring `k16 += 64u` in `w4a16_gemv.cu` or dropping
257    /// `orig_lane * 2u` from `w4a16_gemv_partial` turns this red.
258    #[test]
259    fn sw_partial_shares_pipelined_k16_loop() {
260        for p in named_cu("w4a16_gemv.cu") {
261            let src = fs::read_to_string(&p).unwrap();
262            assert!(
263                src.contains("orig_lane * 2u"),
264                "{}: w4a16_gemv_partial must start k16 at orig_lane*2",
265                p.display()
266            );
267            assert!(
268                src.contains("k16 < K16 + 1u"),
269                "{}: pipelined K16+1 bound missing",
270                p.display()
271            );
272            assert!(
273                !src.contains("k16 += 64u"),
274                "{}: stride-64 sequential loop drifted back in",
275                p.display()
276            );
277        }
278        for p in named_cu("w4a16_gemv_fused.cu") {
279            let src = fs::read_to_string(&p).unwrap();
280            assert!(
281                src.contains("w4a16_dual_partial"),
282                "{}: dual and dual_sw must share w4a16_dual_partial",
283                p.display()
284            );
285            assert!(
286                src.contains("orig_lane * 2u"),
287                "{}: dual_partial must start k16 at orig_lane*2",
288                p.display()
289            );
290        }
291    }
292
293    /// Split the declaration starting at `sig` into (parameter list, body).
294    /// The body is brace-matched, so nested blocks are kept and the next
295    /// function is not swept in.
296    fn fn_signature_and_body<'a>(src: &'a str, sig: &str) -> (&'a str, &'a str) {
297        let start = src
298            .find(sig)
299            .unwrap_or_else(|| panic!("signature `{sig}` not found"));
300        let open = start
301            + src[start..]
302                .find('{')
303                .expect("no body brace after signature");
304        let mut depth = 0usize;
305        for (i, c) in src[open..].char_indices() {
306            match c {
307                '{' => depth += 1,
308                '}' => {
309                    depth -= 1;
310                    if depth == 0 {
311                        return (&src[start..open], &src[open..=open + i]);
312                    }
313                }
314                _ => {}
315            }
316        }
317        panic!("unbalanced braces after `{sig}`");
318    }
319
320    fn fn_body<'a>(src: &'a str, sig: &str) -> &'a str {
321        fn_signature_and_body(src, sig).1
322    }
323
324    /// (file, partial signature, `__constant__` table it must NOT index,
325    ///  callers that must hand it a shared-staged copy)
326    const DECODE_PARTIALS: &[(&str, &str, &str, &[(&str, &str)])] = &[
327        (
328            "w4a16_gemv.cu",
329            "__device__ __forceinline__ float w4a16_gemv_partial(",
330            "E2M1_LUT",
331            &[("w4a16_gemv", "s_lut"), ("w4a16_gemv_sw", "warp_lut")],
332        ),
333        (
334            "w4a16_gemv_fused.cu",
335            "__device__ __forceinline__ float w4a16_dual_partial(",
336            "E2M1_LUT_FUSED_W4",
337            &[
338                ("w4a16_gemv_dual", "s_lut"),
339                ("w4a16_gemv_dual_sw", "warp_lut"),
340            ],
341        ),
342        (
343            "w4a16_gemv_fused.cu",
344            "__device__ __forceinline__ float w4a16_silu_partial(",
345            "E2M1_LUT_FUSED_W4",
346            &[("w4a16_gemv_silu_input_sw", "warp_lut")],
347        ),
348    ];
349
350    /// POSITIVE: every M=1 decode GEMV partial must index a SHARED-staged copy
351    /// of the E2M1 table, never `__constant__` memory directly.
352    ///
353    /// WHY (this is the PR #479 regression, not style): the table index is a
354    /// data-dependent weight nibble. `__constant__` is a BROADCAST cache — a
355    /// warp request replays once per distinct address, and 32 lanes over NVFP4
356    /// weights cover ~14 of the 16 entries, so one lookup costs ~14
357    /// transactions. There is exactly one lookup per weight element, i.e. on
358    /// every FMA of the decode inner loop. Shared memory answers all 16
359    /// indices from 16 distinct banks in one conflict-free transaction.
360    ///
361    /// Staging is numerically inert — `s_lut[i]` is a bit-exact FP32 copy — so
362    /// this coexists with `sw_partial_shares_pipelined_k16_loop`, which pins
363    /// the association order that buys base/SW bit-identity.
364    ///
365    /// PROVEN BY: swapping any `lut[byte_val ...]` back to
366    /// `E2M1_LUT[byte_val ...]` turns this red (SASS: 1 LDS + 32 indexed
367    /// `LDC c[0x3][R]` instead of 33 LDS + 1 LDC in `w4a16_gemv`).
368    #[test]
369    fn decode_gemv_partials_index_a_shared_staged_lut() {
370        for &(file, sig, table, callers) in DECODE_PARTIALS {
371            let partial = sig
372                .strip_suffix('(')
373                .and_then(|sig| sig.split_whitespace().last())
374                .expect("partial signature ends in a function name");
375            let paths = named_cu(file);
376            assert!(paths.len() >= 3, "{file}: expected 3 backend copies");
377            for path in paths {
378                let src = fs::read_to_string(&path).unwrap();
379                let where_ = format!("{}::{sig}", path.display());
380                let (params, body) = fn_signature_and_body(&src, sig);
381                assert!(
382                    !body.contains(&format!("{table}[byte_val")),
383                    "{where_}: data-dependent index into __constant__ {table} \
384                     serializes the warp — take the staged `lut` instead"
385                );
386                assert!(
387                    body.contains("lut[byte_val"),
388                    "{where_}: must dequant through the staged `lut` parameter"
389                );
390                assert!(
391                    params.contains("const float* __restrict__ lut"),
392                    "{where_}: must accept the staged table as `const float* __restrict__ lut`"
393                );
394                for &(caller, expected_lut) in callers {
395                    let cb = fn_body(&src, &format!("void {caller}("));
396                    assert!(
397                        cb.contains("__shared__ float s_lut"),
398                        "{}::{caller}: must stage the E2M1 table in shared memory",
399                        path.display()
400                    );
401                    let calls: Vec<_> = cb
402                        .lines()
403                        .filter(|line| line.contains(&format!("{partial}(")))
404                        .collect();
405                    assert!(
406                        !calls.is_empty(),
407                        "{}::{caller}: no {partial} call",
408                        path.display()
409                    );
410                    for call in calls {
411                        assert!(
412                            call.contains(&format!(", {expected_lut})")),
413                            "{}::{caller}: `{partial}` must receive {expected_lut}, got `{}`",
414                            path.display(),
415                            call.trim()
416                        );
417                        assert!(
418                            !call.contains(table),
419                            "{}::{caller}: `{partial}` received constant-memory {table}",
420                            path.display()
421                        );
422                    }
423                }
424            }
425        }
426    }
427
428    /// STRUCTURAL: the single-warp kernels must stage the LUT PER WARP and
429    /// stay free of block barriers.
430    ///
431    /// Two invariants the shared staging leans on, both silently breakable:
432    ///   1. `w4a16_gemv_sw` / `_dual_sw` / `_silu_input_sw` early-return on
433    ///      `n >= N`, which is warp-uniform (`n = blockIdx.x*8 + tid/32`) but
434    ///      NOT block-uniform. A `__syncthreads()` after that return is a
435    ///      divergent barrier — undefined behaviour, not a compile error. So
436    ///      the staging must publish with `__syncwarp()`, and these kernels
437    ///      must contain no `__syncthreads()` at all (the documented
438    ///      "no smem, no __syncthreads in the reduction" property).
439    ///   2. One private 16-float row per warp, so the row count must track
440    ///      `N_PER_BLOCK_SW`. A hardcoded 8 would index out of bounds the day
441    ///      that define moves. 8 rows = 512 B/block, ~50x under the smem that
442    ///      would cap occupancy at 8 blocks/SM, so it is occupancy-neutral.
443    ///
444    /// PROVEN BY: replacing `__syncwarp()` with `__syncthreads()`, or writing
445    /// `s_lut[8][16]`, turns this red.
446    #[test]
447    fn sw_gemv_stages_the_lut_per_warp_without_a_block_barrier() {
448        let want_rows = "__shared__ float s_lut[N_PER_BLOCK_SW][16]";
449        for (file, kernels, helper) in [
450            (
451                "w4a16_gemv.cu",
452                &["w4a16_gemv_sw"][..],
453                "stage_e2m1_lut_warp",
454            ),
455            (
456                "w4a16_gemv_fused.cu",
457                &["w4a16_gemv_dual_sw", "w4a16_gemv_silu_input_sw"][..],
458                "stage_e2m1_lut_fused_warp",
459            ),
460        ] {
461            for path in named_cu(file) {
462                let src = fs::read_to_string(&path).unwrap();
463                let hb = fn_body(&src, &format!("void {helper}("));
464                assert!(
465                    hb.contains("__syncwarp()"),
466                    "{}::{helper}: warp-scoped staging must publish with __syncwarp()",
467                    path.display()
468                );
469                for k in kernels {
470                    let kb = fn_body(&src, &format!("void {k}("));
471                    assert!(
472                        kb.contains(want_rows),
473                        "{}::{k}: per-warp LUT rows must be sized by N_PER_BLOCK_SW",
474                        path.display()
475                    );
476                    assert!(
477                        kb.contains(&format!("{helper}(s_lut[local_out], lane)")),
478                        "{}::{k}: must stage its own warp row",
479                        path.display()
480                    );
481                    assert!(
482                        !kb.contains("__syncthreads()"),
483                        "{}::{k}: block barrier after a warp-uniform early return is \
484                         divergent UB — and it would undo the barrier-free reduction",
485                        path.display()
486                    );
487                }
488            }
489        }
490    }
491
492    /// NEGATIVE: attention decode must not launch the base GEMV directly.
493    /// A new `ops::w4a16_gemv(` site there ships the 64-thread kernel on
494    /// the default path even though `nvfp4_decode_gemv` exists.
495    ///
496    /// PROVEN BY: restoring any of the pre-PR call sites turns this red.
497    #[test]
498    fn attention_decode_does_not_call_base_w4a16_gemv() {
499        let attn = Path::new(env!("CARGO_MANIFEST_DIR")).join("src/layers/qwen3_attention");
500        let mut offenders = Vec::new();
501        for rel in [
502            "decode/attention_forward.rs",
503            "decode/attention_forward_v4.rs",
504            "decode/attention_forward_oproj.rs",
505            "decode/attention_forward_mla.rs",
506            "decode/attention_forward_kv.rs",
507            "trait_impl/multi_seq/qkv.rs",
508            "trait_impl/multi_seq/attn.rs",
509            "trait_impl/multi_seq/attn/o_proj.rs",
510            "trait_impl/multi_seq/mla.rs",
511        ] {
512            let src = fs::read_to_string(attn.join(rel)).unwrap();
513            if src.contains("w4a16_gemv(") {
514                offenders.push(rel);
515            }
516        }
517        assert!(
518            offenders.is_empty(),
519            "use nvfp4_decode_gemv (N/8 grid) not ops::w4a16_gemv: {offenders:?}"
520        );
521    }
522}
spark_model/layers/ops/gemv_sw.rs

spark_model/layers/ops/
gemv_sw.rs
