metrale_model_layers/layers/ops/
w4a16_tc_rows.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-10-02: `w4a16_tc_rows_{16,32,64}` (`kernels/gb10/common/w4a16_tc_rows.cu`): the NVFP4
4//! W4A16 projection of 1..=64 rows with the rows as the tensor-core MMA's N columns, the NVFP4
5//! point of the row-tile family whose FP8 point is `w8a16_tc_rows`. A row's bits do not depend
6//! on the row count or the entry point, so callers chunk wider row counts in 64-row calls.
7//!
8//! Owner: model-layers ops.
9//! Invariants: [`w4a16_tc_rows`] launches only when [`w4a16_tc_rows_shape_ok`] holds.
10
11use anyhow::Result;
12use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
13use metrale_gpu_runtime::kernel_args::KernelLaunch;
14
15use crate::weight_map::QuantizedWeight;
16
17/// 2026-10-02: Output columns per CTA (`TR_COLS`).
18pub const W4A16_TC_ROWS_COLS: u32 = 64;
19
20/// 2026-10-02: Widest row count of one launch (`8 * NT` of `w4a16_tc_rows_64`).
21pub const W4A16_TC_ROWS_MAX_M: u32 = 64;
22
23/// 2026-10-02: The kernel module.
24pub const W4A16_TC_ROWS_MODULE: &str = "w4a16_tc_rows";
25
26/// 2026-10-02: The kernel's shape contract, without a GPU: 1..=64 rows, any positive N (the
27/// entry points are ragged: a partial last CTA loads zero weight rows and stores nothing past N),
28/// K a positive multiple of 256 (whole load groups of every entry point), and an A pitch that
29/// keeps rows 16-byte aligned and covers K; the C pitch covers N.
30pub fn w4a16_tc_rows_shape_ok(m: u32, n: u32, k: u32, lda: u32, ldc: u32) -> bool {
31    (1..=W4A16_TC_ROWS_MAX_M).contains(&m)
32        && n > 0
33        && k > 0
34        && k.is_multiple_of(256)
35        && lda >= k
36        && lda.is_multiple_of(8)
37        && ldc >= n
38}
39
40/// 2026-10-02: `output [m, ldc] = input [m, lda] x W^T` for the row-major NVFP4 `weight` `[n, k]`
41/// (packed E2M1, E4M3 scales of 16, `weight_scale_2`). Refuses a shape
42/// [`w4a16_tc_rows_shape_ok`] refuses.
43#[allow(clippy::too_many_arguments)]
44pub fn w4a16_tc_rows(
45    gpu: &dyn GpuBackend,
46    input: DevicePtr,
47    weight: &QuantizedWeight,
48    output: DevicePtr,
49    m: u32,
50    n: u32,
51    k: u32,
52    lda: u32,
53    ldc: u32,
54    stream: u64,
55) -> Result<()> {
56    anyhow::ensure!(
57        w4a16_tc_rows_shape_ok(m, n, k, lda, ldc),
58        "w4a16_tc_rows: m={m} n={n} k={k} lda={lda} ldc={ldc} outside the kernel's contract"
59    );
60    let entry = if m <= 16 {
61        "w4a16_tc_rows_16"
62    } else if m <= 32 {
63        "w4a16_tc_rows_32"
64    } else {
65        "w4a16_tc_rows_64"
66    };
67    let kernel = gpu.op_cache().kernel(gpu, W4A16_TC_ROWS_MODULE, entry)?;
68    KernelLaunch::new(gpu, kernel)
69        .grid([n.div_ceil(W4A16_TC_ROWS_COLS), 1, 1])
70        .block([128, 1, 1])
71        .arg_ptr(input)
72        .arg_ptr(weight.weight)
73        .arg_ptr(weight.weight_scale)
74        .arg_f32(weight.weight_scale_2)
75        .arg_ptr(output)
76        .arg_u32(m)
77        .arg_u32(n)
78        .arg_u32(k)
79        .arg_u32(lda)
80        .arg_u32(ldc)
81        .launch(stream)
82}
83
84#[cfg(test)]
85mod tests {
86    use super::*;
87
88    const CU: &str = include_str!("../../../../../kernels/gb10/common/w4a16_tc_rows.cu");
89    const ROWS: &str = include_str!("../../../../../kernels/gb10/common/tc_rows.cuh");
90
91    /// 2026-10-02: The grid and the widest entry point match the kernel (`TR_WARPS` warps of 16
92    /// columns, `tr_block<Nvfp4G16, 8, ...>` row tiles of 8), or columns or rows go unwritten.
93    #[test]
94    fn launch_matches_the_kernel() {
95        assert!(
96            ROWS.contains("#define TR_WARPS 4\n")
97                && ROWS.contains("#define TR_COLS (TR_WARPS * 16)\n")
98        );
99        assert_eq!(W4A16_TC_ROWS_COLS, 4 * 16);
100        assert!(CU.contains(
101            "tr_block<Nvfp4G16, 8, 1, true>(A, {packed, scale, s2}, C, M, N, K, lda, ldc, blockIdx.x);"
102        ));
103        assert_eq!(W4A16_TC_ROWS_MAX_M, 8 * 8);
104        for entry in [
105            "w4a16_tc_rows_16(",
106            "w4a16_tc_rows_32(",
107            "w4a16_tc_rows_64(",
108        ] {
109            assert!(CU.contains(entry), "{entry} missing from the kernel");
110        }
111    }
112
113    /// 2026-10-02: The shape contract refuses each bound it names and admits the 35B head.
114    #[test]
115    fn shape_contract() {
116        assert!(w4a16_tc_rows_shape_ok(1, 248320, 2048, 2048, 248320));
117        assert!(w4a16_tc_rows_shape_ok(64, 248320, 2048, 2048, 248320));
118        assert!(!w4a16_tc_rows_shape_ok(0, 248320, 2048, 2048, 248320));
119        assert!(!w4a16_tc_rows_shape_ok(65, 248320, 2048, 2048, 248320));
120        // 2026-10-02: Ragged N: the checkpoint's 248070-entry vocab.
121        assert!(w4a16_tc_rows_shape_ok(8, 248070, 2048, 2048, 248070));
122        assert!(!w4a16_tc_rows_shape_ok(8, 248320, 2048 + 128, 2176, 248320));
123        assert!(!w4a16_tc_rows_shape_ok(8, 248320, 2048, 2044, 248320));
124        assert!(!w4a16_tc_rows_shape_ok(8, 248320, 2048, 2048, 248319));
125    }
126}
metrale_model_layers/layers/ops/w4a16_tc_rows.rs

metrale_model_layers/layers/ops/
w4a16_tc_rows.rs
