metrale_model_layers/layers/ops/
gemv_tc.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Routing of the narrow NVFP4 batched GEMV to its tensor-core kernels (`w4a16_gemv_tc8` / `w4a16_gemv_tc16` in kernels/gb10/common/w4a16_gemv_tc.cu).
4//!
5//! Owner: model-layers ops.
6//! Invariants:
7//! - The tensor-core entries take the same arguments as the CUDA-core `w4a16_gemv_batch*`
8//!   kernels, with their own grid (`ceil(N / cols_per_cta)`) and a 256-thread block.
9//! - A launch routes only when `K % 128 == 0` and `1 <= m <= 16`; otherwise the caller keeps its
10//!   CUDA-core kernel ([`tc_route`]).
11//! - [`super::w4a16_gemv_batchm`] and [`tc_fixed_m`] both route through [`tc_kernel`].
12//! - `METRALE_NO_W4A16_TC` set to any non-empty value, read once per process, turns the
13//!   tensor-core path off ([`tc_enabled`]). The bench crate's `PERF_CONTROLS` discloses it,
14//!   and `record_env_tests` reads [`tc_enabled`] and [`wide_rows_enabled`] from this file.
15
16use std::sync::{Mutex, OnceLock};
17
18use anyhow::Result;
19use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
20use metrale_gpu_runtime::kernel_args::KernelLaunch;
21
22use crate::weight_map::QuantizedWeight;
23
24/// 2026-09-25: Rows the `w4a16_gemv_tc8` entry covers (its `MT` is 8).
25pub const TC8_MAX_M: u32 = 8;
26/// 2026-09-25: Rows the `w4a16_gemv_tc16` entry covers (its `MT` is 16).
27pub const TC16_MAX_M: u32 = 16;
28/// 2026-09-25: Columns per CTA, `8 * NT`: tc8 has `NT = 1`, tc16 has `NT = 2`.
29pub const TC8_COLS_PER_CTA: u32 = 8;
30pub const TC16_COLS_PER_CTA: u32 = 16;
31/// 2026-09-25: Threads per CTA, `TC_WARPS * 32` with `TC_WARPS = 8` in the .cu; the warps split K.
32pub const TC_BLOCK: u32 = 256;
33
34/// 2026-09-25: Which tensor-core entry serves a launch.
35#[derive(Clone, Copy, Debug, PartialEq, Eq)]
36pub enum TcKind {
37    M8,
38    M16,
39}
40
41impl TcKind {
42    pub fn cols_per_cta(self) -> u32 {
43        match self {
44            TcKind::M8 => TC8_COLS_PER_CTA,
45            TcKind::M16 => TC16_COLS_PER_CTA,
46        }
47    }
48}
49
50/// 2026-09-25: The routing decision, without I/O. `None` keeps the caller's CUDA-core kernel.
51///
52/// The kernel walks K in 128-element blocks, so a `K` that is not a multiple of 128 declines.
53/// Any `N` routes: in a partial last tile the kernel loads zeros for columns past `N` and never
54/// stores them.
55/// `tc8` serves `m <= 8`; `tc16` serves `m <= 16`, including `m <= 8` when `tc8` is missing.
56pub fn tc_route(
57    m: u32,
58    n: u32,
59    k: u32,
60    enabled: bool,
61    have8: bool,
62    have16: bool,
63) -> Option<TcKind> {
64    if !enabled || m == 0 || n == 0 || k == 0 || !k.is_multiple_of(128) {
65        return None;
66    }
67    if m <= TC8_MAX_M && have8 {
68        Some(TcKind::M8)
69    } else if m <= TC16_MAX_M && have16 {
70        Some(TcKind::M16)
71    } else {
72        None
73    }
74}
75
76/// 2026-09-25: True when `METRALE_NO_W4A16_TC` is unset or empty. Any non-empty value, `0`
77/// included, turns the tensor-core path off. Read once per process, so every launch, including
78/// one captured in a CUDA graph, sees the same choice.
79pub fn tc_enabled() -> bool {
80    static ON: OnceLock<bool> = OnceLock::new();
81    *ON.get_or_init(|| std::env::var_os("METRALE_NO_W4A16_TC").is_none_or(|v| v.is_empty()))
82}
83
84/// 2026-09-25: Row edges of the narrow-GEMV dispatch arms: `NARROW_MAX_ROWS` by default, and
85/// `WIDE_MAX_ROWS` (the `tc16` reach) when [`wide_rows_enabled`]. Read through
86/// [`narrow_gemv_max_rows`].
87pub const NARROW_MAX_ROWS: u32 = 8;
88pub const WIDE_MAX_ROWS: u32 = TC16_MAX_M;
89
90/// 2026-09-25: The opt-in widening: true when `METRALE_W4A16_TC_WIDE` is set to a non-empty value
91/// and [`tc_enabled`]. Read once per process. It raises [`narrow_gemv_max_rows`] to 16 and
92/// lets [`tc_fixed_m`] route the fixed-M launchers (`w4a16_gemv_batch2/3`,
93/// `w4a16_gemv_dual_batch2/3`) to the tensor-core kernels.
94pub fn wide_rows_enabled() -> bool {
95    static ON: OnceLock<bool> = OnceLock::new();
96    *ON.get_or_init(|| {
97        tc_enabled() && std::env::var_os("METRALE_W4A16_TC_WIDE").is_some_and(|v| !v.is_empty())
98    })
99}
100
101/// 2026-09-25: Row edge of the narrow-GEMV dispatch arms (see [`wide_rows_enabled`]).
102pub fn narrow_gemv_max_rows() -> u32 {
103    if wide_rows_enabled() {
104        WIDE_MAX_ROWS
105    } else {
106        NARROW_MAX_ROWS
107    }
108}
109
110/// 2026-09-25: Resolved handles, cached per backend. A `KernelHandle` names a function in one
111/// backend's loaded module, so the cache is keyed by the backend object's address. A missing
112/// entry is `KernelHandle(0)`.
113#[derive(Clone, Copy)]
114struct TcHandles {
115    tc8: KernelHandle,
116    tc16: KernelHandle,
117}
118
119fn tc_handles(gpu: &dyn GpuBackend) -> TcHandles {
120    static CACHE: OnceLock<Mutex<Vec<(usize, TcHandles)>>> = OnceLock::new();
121    let key = gpu as *const dyn GpuBackend as *const () as usize;
122    let cache = CACHE.get_or_init(|| Mutex::new(Vec::new()));
123    let mut guard = cache.lock().unwrap_or_else(|p| p.into_inner());
124    if let Some((_, h)) = guard.iter().find(|(k, _)| *k == key) {
125        return *h;
126    }
127    let h = TcHandles {
128        tc8: crate::layers::try_kernel(gpu, "w4a16_gemv_tc", "w4a16_gemv_tc8"),
129        tc16: crate::layers::try_kernel(gpu, "w4a16_gemv_tc", "w4a16_gemv_tc16"),
130    };
131    guard.push((key, h));
132    h
133}
134
135/// 2026-09-25: The tensor-core kernel and grid x for this launch, or `None` to keep the
136/// CUDA-core kernel. It does not consult [`wide_rows_enabled`]; the callers' row edges do.
137pub fn tc_kernel(gpu: &dyn GpuBackend, m: u32, n: u32, k: u32) -> Option<(KernelHandle, u32)> {
138    if !tc_enabled() {
139        return None;
140    }
141    let h = tc_handles(gpu);
142    let kind = tc_route(m, n, k, true, h.tc8.0 != 0, h.tc16.0 != 0)?;
143    let handle = match kind {
144        TcKind::M8 => h.tc8,
145        TcKind::M16 => h.tc16,
146    };
147    Some((handle, n.div_ceil(kind.cols_per_cta())))
148}
149
150/// 2026-09-25: The tensor-core route of the fixed-M launchers, taken only when
151/// [`wide_rows_enabled`]: `Ok(true)` when it launched, `Ok(false)` to keep the CUDA-core kernel.
152#[allow(clippy::too_many_arguments)]
153pub fn tc_fixed_m(
154    gpu: &dyn GpuBackend,
155    input: DevicePtr,
156    weight: &QuantizedWeight,
157    output: DevicePtr,
158    m: u32,
159    n: u32,
160    k: u32,
161    stream: u64,
162) -> Result<bool> {
163    if !wide_rows_enabled() {
164        return Ok(false);
165    }
166    let Some((tc, grid_x)) = tc_kernel(gpu, m, n, k) else {
167        return Ok(false);
168    };
169    KernelLaunch::new(gpu, tc)
170        .grid([grid_x, 1, 1])
171        .block([TC_BLOCK, 1, 1])
172        .arg_ptr(input)
173        .arg_ptr(weight.weight)
174        .arg_ptr(weight.weight_scale)
175        .arg_f32(weight.weight_scale_2)
176        .arg_ptr(output)
177        .arg_u32(m)
178        .arg_u32(n)
179        .arg_u32(k)
180        .launch(stream)?;
181    Ok(true)
182}
183
184#[cfg(test)]
185#[path = "gemv_tc_tests.rs"]
186mod gemv_tc_tests;
metrale_model_layers/layers/ops/gemv_tc.rs

metrale_model_layers/layers/ops/
gemv_tc.rs
