metrale_model_layers/layers/ops/w4a4_proj/
mx_plan.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Which `w4a4_gemv_mx*` entry serves an `m`-row projection, and
4//! with what launch shape: the `METRALE_W4A4_MX_NT` / `METRALE_W4A4_MX_PS`
5//! levers and the pure planning in `mx_plan`.
6//!
7//! Owner: model-layers ops.
8//! Invariants: none beyond the types.
9
10use std::sync::OnceLock;
11
12use metrale_gpu_runtime::gpu::KernelHandle;
13
14use super::{W4A4_MAX_M, W4a4State};
15
16/// 2026-09-25: `METRALE_W4A4_MX_NT` = 4 (default, also unset or empty) | 1; any
17/// other value panics. At 4 the 9..=64-row W4A4 GEMV runs the activation-reuse
18/// twins: 17..=32 rows take four 16-row weight tiles per CTA, 9..=16 and
19/// 33..=64 rows take two, and each warp feeds its activation fragments to
20/// every tile from registers, so the activation matrix is read from L2 that
21/// many times less. The twins keep the one-tile kernels' per-warp chunk order
22/// and warp-order reduction. At 1 every width runs the one-tile kernels.
23pub(super) fn mx_nt() -> u32 {
24    static NT: OnceLock<u32> = OnceLock::new();
25    *NT.get_or_init(
26        || match std::env::var("METRALE_W4A4_MX_NT").ok().as_deref() {
27            None | Some("") | Some("4") => 4,
28            Some("1") => 1,
29            Some(v) => panic!("METRALE_W4A4_MX_NT={v}: expected 1 or 4"),
30        },
31    )
32}
33
34/// 2026-09-25: `METRALE_W4A4_MX_PS` = 1 (default, also unset or empty) | 0; any
35/// other value panics. At 1, a 9..=32-row launch whose whole activation stripe
36/// fits in [`PS_SMEM_MAX`] and whose weight has at least
37/// [`PS_MIN_TILES_PER_SM`] 16-row tiles per SM runs the persistent
38/// activation-staged entry (`w4a4_gemv_mx{16,32}_ps`): one CTA per SM pulls
39/// 16-row tiles from a counter and reads the activations from shared memory.
40/// 0 routes those launches back to the [`mx_nt`] kernels. It acts only when the
41/// tile factor is not 1, so `METRALE_W4A4_MX_NT=1` selects the one-tile kernels
42/// at every width.
43pub(super) fn mx_ps() -> bool {
44    static PS: OnceLock<bool> = OnceLock::new();
45    *PS.get_or_init(
46        || match std::env::var("METRALE_W4A4_MX_PS").ok().as_deref() {
47            None | Some("") | Some("1") => true,
48            Some("0") => false,
49            Some(v) => panic!("METRALE_W4A4_MX_PS={v}: expected 0 or 1"),
50        },
51    )
52}
53
54/// 2026-09-25: Largest dynamic shared memory, in bytes, one CTA may opt in to
55/// on GB10 (sm_121).
56pub const PS_SMEM_MAX: u32 = 101_376;
57/// 2026-09-25: The persistent entries run only when the weight has at least
58/// this many 16-row tiles per SM; below it `mx_plan` uses the tiled kernels.
59pub const PS_MIN_TILES_PER_SM: u32 = 8;
60
61/// 2026-09-25: Pure: 8-token column blocks (`MB`) of the persistent entry
62/// serving `m` rows (`w4a4_gemv_mx16_ps` or `w4a4_gemv_mx32_ps`).
63pub fn ps_column_blocks(m: u32) -> u32 {
64    if m <= 16 { 2 } else { 4 }
65}
66
67/// 2026-09-25: Pure: the host side of the persistent entries' launch contract
68/// (`w4a4_gemv_mx_ps.cuh`): dynamic shared memory for `mb` column blocks with
69/// `sst` k128 chunks staged per warp. 8 warps x sst x mb x (512 B of
70/// fragments + 64 B of scales), plus the reduction buffer of `RJ = 2` 4 KiB
71/// blocks that both entries instantiate.
72pub fn ps_smem_bytes(mb: u32, sst: u32) -> u32 {
73    8 * sst * mb * 576 + 2 * 4096
74}
75
76/// 2026-09-25: Pure: k128 chunks in one warp's stripe (chunks c = warp mod 8).
77pub fn ps_stripe_chunks(k: u32) -> u32 {
78    (k / 128).div_ceil(8)
79}
80
81/// 2026-09-25: How one W4A4 GEMV launch is shaped.
82#[derive(Clone, Copy, Debug)]
83pub(super) enum MxLaunch {
84    /// 2026-09-25: Grid ceil(N / rows_per_cta).
85    Tiles {
86        kernel: KernelHandle,
87        rows_per_cta: u32,
88    },
89    /// 2026-09-25: Grid #SMs, `smem` bytes of dynamic shared memory, `sst`
90    /// staged chunks.
91    Persistent {
92        kernel: KernelHandle,
93        sst: u32,
94        smem: u32,
95    },
96}
97
98/// 2026-09-25: Pure: the launch for an `m`-row `[n, k]` projection.
99pub(super) fn mx_plan(s: &W4a4State, m: u32, n: u32, k: u32, nt: u32, ps: bool) -> MxLaunch {
100    if ps && nt != 1 && m > 8 && m <= W4A4_MAX_M {
101        let kernel = if m <= 16 { s.mx16_ps } else { s.mx32_ps };
102        let sst = ps_stripe_chunks(k);
103        let smem = ps_smem_bytes(ps_column_blocks(m), sst);
104        if smem <= PS_SMEM_MAX && n.div_ceil(16) >= PS_MIN_TILES_PER_SM * s.sms {
105            return MxLaunch::Persistent { kernel, sst, smem };
106        }
107    }
108    let (kernel, rows_per_cta) = mx_pick(s, m, nt);
109    MxLaunch::Tiles {
110        kernel,
111        rows_per_cta,
112    }
113}
114
115/// 2026-09-25: Pure: (kernel, rows per CTA) for an `m`-row launch at tile
116/// factor `nt`.
117pub(super) fn mx_pick(s: &W4a4State, m: u32, nt: u32) -> (KernelHandle, u32) {
118    match (m, nt) {
119        (0..=8, _) => (s.mx8, 16),
120        (9..=16, 1) => (s.mx16, 16),
121        (9..=16, _) => (s.mx16_nt2, 32),
122        (17..=32, 1) => (s.mx32, 16),
123        (17..=32, _) => (s.mx32_nt4, 64),
124        (_, 1) => (s.mx64, 16),
125        _ => (s.mx64_nt2, 32),
126    }
127}
metrale_model_layers/layers/ops/w4a4_proj/mx_plan.rs

metrale_model_layers/layers/ops/w4a4_proj/
mx_plan.rs
