metrale_model_layers/layers/ops/
w8a16_tc_rows.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: `w8a16_tc_rows_{16,32,64}` (`kernels/gb10/common/w8a16_tc_rows.cu`): the
4//! block-scaled W8A16 projection of 1..=64 decode rows with the rows as the tensor-core MMA's
5//! N columns, and the rule for when it stands in for the 32- and 64-row tile twins.
6//!
7//! Under `RowTiers::Canonical` every W8A16 projection of 1..=64 rows goes through those twins
8//! ([`super::w8a16_gemm_pipelined_m32_strided`], [`super::w8a16_gemm_pipelined_m64_strided`]);
9//! they launch this kernel instead (`METRALE_NO_W8A16_TC_ROWS` keeps the twins). Same precision (FP8 weights,
10//! BF16 activations, FP32 accumulation, the 128 x 128 block scales), another summation order,
11//! so the output bits differ from the twins'; a row's bits still do not depend on the row
12//! count, so the canonical policy holds from 1 to 64 rows. 2026-09-30: above 64 rows
13//! `w8a16_tc_rows_64c` runs the same 64-row body per 64-row chunk, so the policy holds at every
14//! row count (the 128-row tile it replaces there sums in another order).
15//!
16//! Owner: model-layers ops.
17//! Invariants: [`w8a16_tc_rows_launch`] launches only when [`w8a16_tc_rows_shape_ok`] holds.
18
19use anyhow::Result;
20use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
21use metrale_gpu_runtime::kernel_args::KernelLaunch;
22
23/// 2026-09-28: Output columns per CTA (`TR_COLS`).
24pub const W8A16_TC_ROWS_COLS: u32 = 64;
25
26/// 2026-09-28: Widest row count of one row chunk (`8 * NT` of `w8a16_tc_rows_64`); more rows
27/// run as chunks of it (`w8a16_tc_rows_64c`).
28pub const W8A16_TC_ROWS_MAX_M: u32 = 64;
29
30const MODULE: &str = "w8a16_tc_rows";
31
32/// 2026-09-28: On unless `METRALE_NO_W8A16_TC_ROWS` is present (a debugging kill switch: the
33/// tile twins). Read once per process.
34fn w8a16_tc_rows_enabled() -> bool {
35    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
36    *ON.get_or_init(|| std::env::var_os("METRALE_NO_W8A16_TC_ROWS").is_none())
37}
38
39/// 2026-09-28: The kernel's shape contract, without a GPU: at least one row, N a positive
40/// multiple of 128 (whole scale blocks, whole CTAs), K a positive multiple of 128, and an A
41/// pitch that keeps rows 16-byte aligned and covers K; the C pitch covers N.
42pub fn w8a16_tc_rows_shape_ok(
43    m: u32,
44    n: u32,
45    k: u32,
46    a_row_stride: u32,
47    c_row_stride: u32,
48) -> bool {
49    m >= 1
50        && n > 0
51        && n.is_multiple_of(128)
52        && k > 0
53        && k.is_multiple_of(128)
54        && a_row_stride >= k
55        && a_row_stride.is_multiple_of(8)
56        && c_row_stride >= n
57}
58
59/// 2026-09-28: Launches `[m, k] x [n, k]^T` through `w8a16_tc_rows_{16,32,64}` and returns true
60/// when the lever is on, the process runs `RowTiers::Canonical`, the shape passes
61/// [`w8a16_tc_rows_shape_ok`] and the target ships the module; otherwise returns false and
62/// launches nothing.
63#[allow(clippy::too_many_arguments)]
64pub fn w8a16_tc_rows_launch(
65    gpu: &dyn GpuBackend,
66    input: DevicePtr,
67    weight: DevicePtr,
68    block_scale: DevicePtr,
69    output: DevicePtr,
70    m: u32,
71    n: u32,
72    k: u32,
73    a_row_stride: u32,
74    c_row_stride: u32,
75    stream: u64,
76) -> Result<bool> {
77    if !w8a16_tc_rows_enabled()
78        || crate::layers::row_tiers() != crate::layers::RowTiers::Canonical
79        || !w8a16_tc_rows_shape_ok(m, n, k, a_row_stride, c_row_stride)
80        || !gpu.has_module(MODULE)
81    {
82        return Ok(false);
83    }
84    w8a16_tc_rows(
85        gpu,
86        input,
87        weight,
88        block_scale,
89        output,
90        m,
91        n,
92        k,
93        a_row_stride,
94        c_row_stride,
95        stream,
96    )?;
97    Ok(true)
98}
99
100/// 2026-09-28: The launch itself, with no policy check: `input` `[m, a_row_stride]` BF16,
101/// `weight` `[n, k]` E4M3, `block_scale` `[n / 128, k / 128]` FP32, `output` `[m,
102/// c_row_stride]` BF16. Refuses a shape [`w8a16_tc_rows_shape_ok`] refuses.
103#[allow(clippy::too_many_arguments)]
104pub fn w8a16_tc_rows(
105    gpu: &dyn GpuBackend,
106    input: DevicePtr,
107    weight: DevicePtr,
108    block_scale: DevicePtr,
109    output: DevicePtr,
110    m: u32,
111    n: u32,
112    k: u32,
113    a_row_stride: u32,
114    c_row_stride: u32,
115    stream: u64,
116) -> Result<()> {
117    anyhow::ensure!(
118        w8a16_tc_rows_shape_ok(m, n, k, a_row_stride, c_row_stride),
119        "w8a16_tc_rows: m={m} n={n} k={k} lda={a_row_stride} ldc={c_row_stride} outside the kernel's contract"
120    );
121    let (entry, chunks) = if m <= 16 {
122        ("w8a16_tc_rows_16", 1)
123    } else if m <= 32 {
124        ("w8a16_tc_rows_32", 1)
125    } else if m <= W8A16_TC_ROWS_MAX_M {
126        ("w8a16_tc_rows_64", 1)
127    } else {
128        ("w8a16_tc_rows_64c", m.div_ceil(W8A16_TC_ROWS_MAX_M))
129    };
130    let kernel = gpu.op_cache().kernel(gpu, MODULE, entry)?;
131    KernelLaunch::new(gpu, kernel)
132        .grid([n / W8A16_TC_ROWS_COLS * chunks, 1, 1])
133        .block([128, 1, 1])
134        .arg_ptr(input)
135        .arg_ptr(weight)
136        .arg_ptr(block_scale)
137        .arg_ptr(output)
138        .arg_u32(m)
139        .arg_u32(n)
140        .arg_u32(k)
141        .arg_u32(a_row_stride)
142        .arg_u32(c_row_stride)
143        .launch(stream)
144}
145
146#[cfg(test)]
147mod tests {
148    use super::*;
149
150    const CU: &str = include_str!("../../../../../kernels/gb10/common/w8a16_tc_rows.cu");
151    const ROWS: &str = include_str!("../../../../../kernels/gb10/common/tc_rows.cuh");
152
153    /// 2026-09-28: The grid covers `n / W8A16_TC_ROWS_COLS` CTAs of 128 threads and the widest
154    /// entry point takes 64 rows: both must match the kernel (`TR_WARPS` warps of 16 columns,
155    /// `tr_block<8, ...>` row tiles of 8), or columns or rows go unwritten.
156    #[test]
157    fn launch_matches_the_kernel() {
158        assert!(
159            ROWS.contains("#define TR_WARPS 4\n")
160                && ROWS.contains("#define TR_COLS (TR_WARPS * 16)\n")
161        );
162        assert_eq!(W8A16_TC_ROWS_COLS, 4 * 16);
163        assert!(CU.contains(
164            "tr_block<Fp8Block128, 8, 2>(A, {B, block_scale}, C, M, N, K, lda, ldc, blockIdx.x);"
165        ));
166        // 2026-09-30: The chunked entry runs the 64-row body per chunk of `W8A16_TC_ROWS_MAX_M`.
167        assert!(CU.contains("const unsigned int chunks = (M + 63) / 64;"));
168        assert!(
169            CU.contains("tr_block<Fp8Block128, 8, 2>(A + (unsigned long long)chunk * 64 * lda")
170        );
171        assert_eq!(W8A16_TC_ROWS_MAX_M, 8 * 8);
172        for entry in [
173            "w8a16_tc_rows_16(",
174            "w8a16_tc_rows_32(",
175            "w8a16_tc_rows_64(",
176            "w8a16_tc_rows_64c(",
177        ] {
178            assert!(CU.contains(entry), "{entry} missing from the kernel");
179        }
180    }
181
182    /// 2026-09-28: The shape contract refuses each bound it names and admits the 35B shapes.
183    #[test]
184    fn shape_contract() {
185        assert!(w8a16_tc_rows_shape_ok(1, 12288, 2048, 2048, 12288));
186        assert!(w8a16_tc_rows_shape_ok(64, 512, 2048, 2048, 512));
187        assert!(!w8a16_tc_rows_shape_ok(0, 512, 2048, 2048, 512));
188        assert!(w8a16_tc_rows_shape_ok(65, 512, 2048, 2048, 512));
189        assert!(w8a16_tc_rows_shape_ok(256, 12288, 2048, 2048, 12288));
190        assert!(!w8a16_tc_rows_shape_ok(8, 576, 2048, 2048, 576));
191        assert!(!w8a16_tc_rows_shape_ok(8, 512, 2000, 2048, 512));
192        assert!(!w8a16_tc_rows_shape_ok(8, 512, 2048, 2044, 512));
193        assert!(!w8a16_tc_rows_shape_ok(8, 512, 2048, 2052, 512));
194        assert!(!w8a16_tc_rows_shape_ok(8, 512, 2048, 2048, 511));
195    }
196}
metrale_model_layers/layers/ops/w8a16_tc_rows.rs

metrale_model_layers/layers/ops/
w8a16_tc_rows.rs
