metrale_model_layers/layers/ops/
dense_gemv_tc.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Routes the MTP drafter's BF16 small-M GEMM rows to the
4//! tensor-core GEMV entries of `dense_gemv_bf16_tc.cu`.
5//!
6//! Owner: model-layers ops.
7//! Invariants:
8//! - [`try_dense_gemv_tc`] launches only when `METRALE_NO_MTP_TC` is unset or
9//!   empty, `MIN_M <= m <= TC32_MAX_M`, `n > 0`, `k % K_STEP == 0` and a
10//!   covering entry resolved; otherwise it returns `Ok(false)` and launches
11//!   nothing.
12//! - `METRALE_NO_MTP_TC` is read once per process. Any non-empty value, `0`
13//!   included, turns the path off.
14//!
15//! The entries take `dense_gemv_bf16_batchm`'s arguments in the same order:
16//! A `[M,K]` contiguous, W `[N,K]`, output rows `out_stride` elements apart.
17//! The warps of a CTA reduce through shared memory in a fixed order with no
18//! atomics, so the output is the same run to run.
19
20use std::sync::{Mutex, OnceLock};
21
22use anyhow::Result;
23use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
24use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
25
26use crate::weight_map::DenseWeight;
27
28/// 2026-09-25: Widest M the `tc8` entry covers: one 8-row token tile. The
29/// `tc16` and `tc32` entries run 2 and 4 tiles.
30pub const TC8_MAX_M: u32 = 8;
31pub const TC16_MAX_M: u32 = 16;
32pub const TC32_MAX_M: u32 = 32;
33/// 2026-09-25: Weight rows per CTA: `16 * NT`, and every entry uses `NT = 1`.
34pub const ROWS_PER_CTA: u32 = 16;
35/// 2026-09-25: Threads per CTA: `DTC_WARPS` (8) warps that split K.
36pub const BLOCK: u32 = 256;
37/// 2026-09-25: K per warp step (`DTC_KB`); a K that is not a multiple of it
38/// does not route.
39pub const K_STEP: u32 = 64;
40/// 2026-09-25: Narrowest M that routes; at M=1 the caller keeps its own kernel.
41pub const MIN_M: u32 = 2;
42
43/// 2026-09-25: Which tensor-core entry serves a launch.
44#[derive(Clone, Copy, Debug, PartialEq, Eq)]
45pub enum DtcKind {
46    M8,
47    M16,
48    M32,
49}
50
51impl DtcKind {
52    pub fn entry(self) -> &'static str {
53        match self {
54            DtcKind::M8 => "dense_gemv_bf16_tc8",
55            DtcKind::M16 => "dense_gemv_bf16_tc16",
56            DtcKind::M32 => "dense_gemv_bf16_tc32",
57        }
58    }
59}
60
61/// 2026-09-25: Pure routing decision; `None` keeps the caller's kernel.
62///
63/// `have` says which of `[tc8, tc16, tc32]` resolved. The narrowest resolved
64/// entry that covers `m` wins; a wider one is also correct, because a token
65/// tile past M is skipped. Declines when `enabled` is false, `m < MIN_M`,
66/// `n == 0`, `k == 0` or `k` is not a multiple of `K_STEP`. Any other N
67/// routes: the kernel loads zeros for weight rows past N and never stores
68/// them.
69pub fn route(m: u32, n: u32, k: u32, enabled: bool, have: [bool; 3]) -> Option<DtcKind> {
70    if !enabled || m < MIN_M || n == 0 || k == 0 || !k.is_multiple_of(K_STEP) {
71        return None;
72    }
73    [
74        (TC8_MAX_M, DtcKind::M8),
75        (TC16_MAX_M, DtcKind::M16),
76        (TC32_MAX_M, DtcKind::M32),
77    ]
78    .into_iter()
79    .zip(have)
80    .find(|&((cap, _), ok)| ok && m <= cap)
81    .map(|((_, kind), _)| kind)
82}
83
84/// 2026-09-25: The `METRALE_NO_MTP_TC` rule over the looked-up value: the
85/// path is on unless the variable holds a non-empty value.
86pub fn mtp_tc_from(kill: Option<&std::ffi::OsStr>) -> bool {
87    kill.is_none_or(|v| v.is_empty())
88}
89
90/// 2026-09-25: Whether the path is on (`METRALE_NO_MTP_TC` unset or empty).
91/// Read once, so every launch in a process sees the same choice.
92pub fn mtp_tc_enabled() -> bool {
93    static ON: OnceLock<bool> = OnceLock::new();
94    *ON.get_or_init(|| mtp_tc_from(std::env::var_os("METRALE_NO_MTP_TC").as_deref()))
95}
96
97/// 2026-09-25: Handles `[tc8, tc16, tc32]`, cached per backend, since a
98/// `KernelHandle` belongs to one backend's loaded module. A missing entry is
99/// the zero handle.
100fn handles(gpu: &dyn GpuBackend) -> [KernelHandle; 3] {
101    static CACHE: OnceLock<Mutex<Vec<(usize, [KernelHandle; 3])>>> = OnceLock::new();
102    let key = gpu as *const dyn GpuBackend as *const () as usize;
103    let cache = CACHE.get_or_init(|| Mutex::new(Vec::new()));
104    let mut guard = cache.lock().unwrap_or_else(|p| p.into_inner());
105    if let Some((_, h)) = guard.iter().find(|(k, _)| *k == key) {
106        return *h;
107    }
108    let h = [DtcKind::M8, DtcKind::M16, DtcKind::M32]
109        .map(|kind| crate::layers::try_kernel(gpu, "dense_gemv_bf16_tc", kind.entry()));
110    guard.push((key, h));
111    h
112}
113
114/// 2026-09-25: The tensor-core kernel and grid-x for this launch when the path
115/// is on and an entry routes, else `None`.
116pub fn kernel_for(gpu: &dyn GpuBackend, m: u32, n: u32, k: u32) -> Option<(KernelHandle, u32)> {
117    if !mtp_tc_enabled() {
118        return None;
119    }
120    let h = handles(gpu);
121    let kind = route(m, n, k, true, h.map(|x| x.0 != 0))?;
122    let handle = match kind {
123        DtcKind::M8 => h[0],
124        DtcKind::M16 => h[1],
125        DtcKind::M32 => h[2],
126    };
127    Some((handle, div_ceil(n, ROWS_PER_CTA)))
128}
129
130/// 2026-09-25: Launch the given entry with the given grid-x, without routing.
131/// [`try_dense_gemv_tc`] is the routed entry point.
132#[allow(clippy::too_many_arguments)]
133pub fn launch(
134    gpu: &dyn GpuBackend,
135    kernel: KernelHandle,
136    grid_x: u32,
137    input: DevicePtr,
138    weight: &DenseWeight,
139    output: DevicePtr,
140    m: u32,
141    n: u32,
142    k: u32,
143    out_stride: u32,
144    stream: u64,
145) -> Result<()> {
146    KernelLaunch::new(gpu, kernel)
147        .grid([grid_x, 1, 1])
148        .block([BLOCK, 1, 1])
149        .arg_ptr(input)
150        .arg_ptr(weight.weight)
151        .arg_ptr(output)
152        .arg_u32(m)
153        .arg_u32(n)
154        .arg_u32(k)
155        .arg_u32(out_stride)
156        .launch(stream)
157}
158
159/// 2026-09-25: Run `C[t] = A[t] @ W^T` on a tensor-core entry if the path is
160/// on and the shape routes. `Ok(false)` means nothing was launched and the
161/// caller must run its own kernel.
162#[allow(clippy::too_many_arguments)]
163pub fn try_dense_gemv_tc(
164    gpu: &dyn GpuBackend,
165    input: DevicePtr,
166    weight: &DenseWeight,
167    output: DevicePtr,
168    m: u32,
169    n: u32,
170    k: u32,
171    out_stride: u32,
172    stream: u64,
173) -> Result<bool> {
174    let Some((kernel, grid_x)) = kernel_for(gpu, m, n, k) else {
175        return Ok(false);
176    };
177    launch(
178        gpu, kernel, grid_x, input, weight, output, m, n, k, out_stride, stream,
179    )?;
180    Ok(true)
181}
182
183#[cfg(test)]
184#[path = "dense_gemv_tc_tests.rs"]
185mod dense_gemv_tc_tests;
metrale_model_layers/layers/ops/dense_gemv_tc.rs

metrale_model_layers/layers/ops/
dense_gemv_tc.rs
