metrale_model_layers/layers/ops/
fp8_act_quant_floor.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: When the Hopper FP8 activation-quant twin takes a launch: the
4//! CTA-count floor, the lever that requests it, and one log line per branch.
5//!
6//! Owner: model-layers ops.
7//! Invariants:
8//! - With the twin present and no shared kernel in the pair,
9//!   [`fp8_act_quant_hopper_reject`] returns `None` (the twin runs), whatever
10//!   the lever and the floor say.
11//!
12//! The twin packs 8 K-groups into a CTA, so its grid (`M x ceil(K/128 / 8)`)
13//! has about an eighth of the shared kernel's CTAs. It takes a launch only
14//! when that grid has at least [`FP8_QUANT_MIN_CTAS_PER_SM`] x `sm_count`
15//! CTAs. The floor counts CTAs, not tokens, because the same M makes a
16//! different grid at a different K. With `sm_count = 132`
17//! (`kernels/hopper/HARDWARE.toml`) the floor is 264 CTAs, and the smallest M
18//! the twin accepts is:
19//!
20//! | K | K/128 | grid Y | min M |
21//! |---:|---:|---:|---:|
22//! | 5120 | 40 | 5 | 53 |
23//! | 6144 | 48 | 6 | 44 |
24//! | 17408 | 136 | 17 | 16 |
25//!
26//! `fp8_act_quant_tests.rs` pins these thresholds.
27
28use metrale_gpu_runtime::gpu::KernelHandle;
29
30/// 2026-09-25: CTAs per SM the twin's grid must reach before it takes a
31/// launch. The same value as `ssm_ba_gates_hopper::MIN_CTAS_PER_SM`.
32pub const FP8_QUANT_MIN_CTAS_PER_SM: u32 = 2;
33
34/// 2026-09-25: The smallest twin grid, in CTAs, the floor accepts for this SM
35/// count; an `sm_count` of 0 counts as 1.
36pub fn fp8_quant_min_ctas(sm_count: u32) -> u32 {
37    FP8_QUANT_MIN_CTAS_PER_SM.saturating_mul(sm_count.max(1))
38}
39
40/// 2026-09-25: CTAs the twin's grid launches for `(m, k)`: the product of
41/// [`super::fp8_quant_grid`]'s Hopper arm, so the floor measures the grid the
42/// launch uses.
43pub fn fp8_quant_hopper_ctas(m: u32, k: u32) -> u32 {
44    let [x, y, z] = super::fp8_quant_grid(true, m, k);
45    x.saturating_mul(y).saturating_mul(z)
46}
47
48/// 2026-09-25: The smallest `m` the twin accepts at this `k`: the table in
49/// this module's header, as a function.
50pub fn fp8_quant_hopper_min_m(k: u32, sm_count: u32) -> u32 {
51    let [_, y, _] = super::fp8_quant_grid(true, 1, k);
52    fp8_quant_min_ctas(sm_count).div_ceil(y.max(1))
53}
54
55/// 2026-09-25: Is the twin requested? The compiled target's `[defaults]
56/// fp8_act_quant_hopper`, overridden by `METRALE_FP8_ACT_QUANT_HOPPER` when
57/// it is set ([`super::target_defaults`]).
58pub fn fp8_act_quant_hopper_enabled() -> bool {
59    super::target_defaults::resolved()
60        .fp8_act_quant_hopper
61        .value
62}
63
64/// 2026-09-25: The floor's refusal reason.
65pub const FP8_QUANT_TOO_FEW_CTAS: &str = "too few CTAs to fill the device at 8 K-groups per CTA";
66
67/// 2026-09-25: Every reason [`fp8_act_quant_hopper_reject`] can return.
68/// [`fp8_quant_log`] keeps one once-flag per entry, by index.
69pub const FP8_QUANT_REJECTS: [&str; 3] = [
70    "not requested",
71    "kernel absent from this image (kernels/hopper only)",
72    FP8_QUANT_TOO_FEW_CTAS,
73];
74
75/// 2026-09-25: Why the twin is not running; `None` means it runs.
76///
77/// Pure, so the rule is testable without a GPU or the process environment.
78///
79/// When the twin is present and the pair has no shared kernel, it returns
80/// `None` before it reads the lever or the floor: a refusal hands the launch
81/// to `Fp8ActQuant::shared`, and there is none to hand it to.
82/// `native_fp8_act_quant_hopper_microtest` builds such a pair
83/// (`shared: KernelHandle(0)`) to run the twin at every M.
84pub fn fp8_act_quant_hopper_reject(
85    requested: bool,
86    twin_present: bool,
87    parent_present: bool,
88    m: u32,
89    k: u32,
90    sm_count: u32,
91) -> Option<&'static str> {
92    if !twin_present {
93        Some(FP8_QUANT_REJECTS[1])
94    } else if !parent_present {
95        None
96    } else if !requested {
97        Some(FP8_QUANT_REJECTS[0])
98    } else if fp8_quant_hopper_ctas(m, k) < fp8_quant_min_ctas(sm_count) {
99        Some(FP8_QUANT_TOO_FEW_CTAS)
100    } else {
101        None
102    }
103}
104
105/// 2026-09-25: Which kernel a launch runs, on what grid, and why the other did
106/// not. [`super::Fp8ActQuant::pick_with`] sets the handle and the grid from one
107/// decision.
108///
109/// No `PartialEq`, because `KernelHandle` does not implement it; tests compare
110/// `twin`, `reject` and `grid`.
111#[derive(Debug, Clone, Copy)]
112pub struct Fp8QuantPick {
113    pub kernel: KernelHandle,
114    pub grid: [u32; 3],
115    /// 2026-09-25: `true` when `kernel` is the Hopper twin.
116    pub twin: bool,
117    /// 2026-09-25: `None` when the twin runs; otherwise the refusal reason.
118    pub reject: Option<&'static str>,
119    /// 2026-09-25: The lever value the pick was made with.
120    pub requested: bool,
121}
122
123/// 2026-09-25: Which once-flag [`fp8_quant_log`] uses for a verdict. There is
124/// one per branch because the verdict changes with M, so one process can log
125/// both the twin line and a refusal.
126#[derive(Debug, Clone, Copy, PartialEq, Eq)]
127pub enum Fp8QuantLogSlot {
128    /// 2026-09-25: The twin took the launch.
129    Twin,
130    /// 2026-09-25: The shared kernel took it, for the reason at this index of
131    /// [`FP8_QUANT_REJECTS`].
132    Reject(usize),
133}
134
135/// 2026-09-25: Once-flags [`fp8_quant_log`] keeps: one per refusal reason,
136/// plus the twin's.
137pub const FP8_QUANT_LOG_SLOTS: usize = FP8_QUANT_REJECTS.len() + 1;
138
139/// 2026-09-25: The slot a verdict belongs to, or `None` when it is not logged.
140/// Pure, so tests replay a call order on the CPU.
141pub fn fp8_quant_log_slot(pick: &Fp8QuantPick) -> Option<Fp8QuantLogSlot> {
142    match pick.reject {
143        None => Some(Fp8QuantLogSlot::Twin),
144        // 2026-09-25: The twin was not requested, so the shared kernel is the
145        // expected outcome and nothing is logged.
146        Some(_) if !pick.requested => None,
147        Some(why) => FP8_QUANT_REJECTS
148            .iter()
149            .position(|r| *r == why)
150            .map(Fp8QuantLogSlot::Reject),
151    }
152}
153
154/// 2026-09-25: Log which quantizer runs and, when the lever requested the twin
155/// and did not get it, which reason refused it: once per process per slot.
156/// `per_token_group_quant_fp8` calls this on every launch.
157pub fn fp8_quant_log(pick: &Fp8QuantPick, m: u32, k: u32) {
158    static SAID: [std::sync::Once; FP8_QUANT_LOG_SLOTS] =
159        [const { std::sync::Once::new() }; FP8_QUANT_LOG_SLOTS];
160    let Some(slot) = fp8_quant_log_slot(pick) else {
161        return;
162    };
163    let idx = match slot {
164        Fp8QuantLogSlot::Twin => FP8_QUANT_LOG_SLOTS - 1,
165        Fp8QuantLogSlot::Reject(i) => i,
166    };
167    let [gx, gy, _] = pick.grid;
168    SAID[idx].call_once(|| match pick.reject {
169        Some(why) => tracing::info!(
170            "FP8 act-quant: the Hopper twin is NOT running at M={m} K={k}: {why} \
171             (METRALE_FP8_ACT_QUANT_HOPPER)"
172        ),
173        None => tracing::info!(
174            "FP8 act-quant: per_token_group_quant_fp8_hopper \
175             (METRALE_FP8_ACT_QUANT_HOPPER) M={m} K={k} grid=({gx},{gy},1) block=128"
176        ),
177    });
178}
metrale_model_layers/layers/ops/fp8_act_quant_floor.rs

metrale_model_layers/layers/ops/
fp8_act_quant_floor.rs
