metrale_model_layers/layers/
row_tiers.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-27: Row-invariant tier policies: whether a row's arithmetic may depend
4//! on how many rows share its launch.
5//!
6//! By default several ops pick a kernel by row count R, and the kernels sum in
7//! different orders: the W8A16 projections (scalar-order GEMV up to 16 rows, a
8//! tensor-core tile above), the NVFP4 LM head (GEMV below 9 rows, a tile GEMM
9//! above), and the MoE router and experts (two- and three-row kernels, the
10//! pairwise walk, the grouped paths). A row's logits then depend on R, which
11//! depends on the draft count and on which sequences are active, so greedy
12//! output varies across `num_drafts` and across runs at a fixed concurrency.
13//! Two policies make every row's bits independent of R:
14//! - `METRALE_ROW_EXACT_TIERS` (presence, off by default): the W8A16
15//!   projections keep the scalar GEMV's order at every R (16-row batch16
16//!   chunks above 16 rows).
17//! - Canonical: the W8A16 projections take the tensor-core tile family's order
18//!   at every R (`w8a16_gemm_pipelined` and its 32/64-row twins, which sum
19//!   identically). 2026-09-27: The default for FP8 MoE checkpoints;
20//!   `--no-canonical-tiers` opts out, and `METRALE_CANONICAL_TIERS` (presence)
21//!   turns it on for any checkpoint.
22//!
23//! Either one also makes the NVFP4 LM head take its tile GEMM at every R and
24//! the FP8 MoE take the grouped kernels with the per-row router at every R of
25//! two or more (`GroupedRouting::PerRow`, the bits of `MoeLayer::forward`).
26//!
27//! Owner: model-layers.
28//! Invariants:
29//! - The serve publishes the policy of each model before building it
30//!   ([`publish_row_tiers`]); a process that never publishes (a test, an
31//!   example) reads the two levers once, and runs by rows without them.
32//! - `--no-canonical-tiers` wins over `METRALE_CANONICAL_TIERS`, which wins over
33//!   `METRALE_ROW_EXACT_TIERS`.
34
35use std::sync::atomic::{AtomicU8, Ordering};
36
37/// 2026-09-27: The projection order a process runs.
38#[derive(Clone, Copy, Debug, PartialEq, Eq)]
39pub enum RowTiers {
40    /// 2026-09-27: Pick by row count (the order before canonical tiers).
41    ByRows,
42    /// 2026-09-27: Scalar-GEMV order at every row count.
43    Exact,
44    /// 2026-09-27: Tensor-core tile order at every row count.
45    Canonical,
46}
47
48/// 2026-09-27: `RowTiers` from the presence of the two levers.
49pub fn row_tiers_from(exact: bool, canonical: bool) -> RowTiers {
50    if canonical {
51        RowTiers::Canonical
52    } else if exact {
53        RowTiers::Exact
54    } else {
55        RowTiers::ByRows
56    }
57}
58
59/// 2026-09-27: A model's policy: `--no-canonical-tiers` keeps the scalar order
60/// if `METRALE_ROW_EXACT_TIERS` is present and the by-rows tiers otherwise; else
61/// a present lever decides; else FP8 MoE checkpoints run canonical and the
62/// others by rows. The dense mixed FP8/NVFP4 checkpoints (Qwen3.8-27B) keep
63/// their tiers: canonical tiers cost them 2-4% of decode throughput.
64pub fn resolve_row_tiers(
65    no_canonical: bool,
66    exact_env: bool,
67    canonical_env: bool,
68    fp8_moe_checkpoint: bool,
69) -> RowTiers {
70    if no_canonical {
71        row_tiers_from(exact_env, false)
72    } else if canonical_env || exact_env {
73        row_tiers_from(exact_env, canonical_env)
74    } else if fp8_moe_checkpoint {
75        RowTiers::Canonical
76    } else {
77        RowTiers::ByRows
78    }
79}
80
81const UNPUBLISHED: u8 = 0;
82
83fn encode(t: RowTiers) -> u8 {
84    match t {
85        RowTiers::ByRows => 1,
86        RowTiers::Exact => 2,
87        RowTiers::Canonical => 3,
88    }
89}
90
91/// 2026-09-27: The policy the serve published for the model it builds.
92static PUBLISHED: AtomicU8 = AtomicU8::new(UNPUBLISHED);
93
94/// 2026-09-27: Publish the policy of the model about to be built. A model swap
95/// publishes again before the next build.
96pub fn publish_row_tiers(t: RowTiers) {
97    PUBLISHED.store(encode(t), Ordering::Relaxed);
98    tracing::info!(
99        "row tiers: {t:?}{}",
100        if t == RowTiers::ByRows {
101            " (kernels picked by row count)"
102        } else {
103            " (one summation order per op at every row count)"
104        }
105    );
106}
107
108/// 2026-09-27: This process's `RowTiers`: the published policy, else the two
109/// levers, read once.
110pub fn row_tiers() -> RowTiers {
111    match PUBLISHED.load(Ordering::Relaxed) {
112        1 => RowTiers::ByRows,
113        2 => RowTiers::Exact,
114        3 => RowTiers::Canonical,
115        _ => {
116            static FALLBACK: std::sync::OnceLock<RowTiers> = std::sync::OnceLock::new();
117            *FALLBACK.get_or_init(|| {
118                row_tiers_from(
119                    std::env::var_os("METRALE_ROW_EXACT_TIERS").is_some(),
120                    std::env::var_os("METRALE_CANONICAL_TIERS").is_some(),
121                )
122            })
123        }
124    }
125}
126
127/// 2026-09-27: Whether this process runs a row-invariant policy.
128pub fn row_invariant() -> bool {
129    row_tiers() != RowTiers::ByRows
130}
131
132/// 2026-10-01: Whether a prefill pass picks its kernels without regard to its row count,
133/// so a row's bits are the same in a long pass and in a short pass over the same rows (a
134/// warm request replaying a prompt's tail). On under the canonical tiers: the FP8 MoE then
135/// takes the grouped prefill kernels at every row count (`MoeLayer::forward_prefill`),
136/// where it took the per-row decode kernels up to 64 rows. The row-count switches left in
137/// a prefill (W8A16 `pipe128`, the BA-gates twin, the router's register tile, the gate/up
138/// `_w1`/`_w2` entries) choose between kernels that give the same bits.
139pub fn prefill_row_invariant() -> bool {
140    row_tiers() == RowTiers::Canonical
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146
147    /// 2026-09-27: Canonical wins over exact; neither lever keeps the by-rows tiers.
148    #[test]
149    fn canonical_wins_and_absence_keeps_by_rows() {
150        assert_eq!(row_tiers_from(false, false), RowTiers::ByRows);
151        assert_eq!(row_tiers_from(true, false), RowTiers::Exact);
152        assert_eq!(row_tiers_from(false, true), RowTiers::Canonical);
153        assert_eq!(row_tiers_from(true, true), RowTiers::Canonical);
154    }
155
156    /// 2026-09-27: FP8 MoE checkpoints default to canonical and the others to by rows;
157    /// the opt-out beats the canonical lever but keeps an exact one; a lever
158    /// decides for any checkpoint.
159    #[test]
160    fn fp8_defaults_to_canonical_and_the_opt_out_wins() {
161        for fp8 in [false, true] {
162            let default = resolve_row_tiers(false, false, false, fp8);
163            assert_eq!(
164                default,
165                if fp8 {
166                    RowTiers::Canonical
167                } else {
168                    RowTiers::ByRows
169                }
170            );
171            assert_eq!(resolve_row_tiers(true, false, false, fp8), RowTiers::ByRows);
172            assert_eq!(resolve_row_tiers(true, false, true, fp8), RowTiers::ByRows);
173            assert_eq!(resolve_row_tiers(true, true, true, fp8), RowTiers::Exact);
174            assert_eq!(resolve_row_tiers(false, true, false, fp8), RowTiers::Exact);
175            assert_eq!(
176                resolve_row_tiers(false, false, true, fp8),
177                RowTiers::Canonical
178            );
179            assert_eq!(
180                resolve_row_tiers(false, true, true, fp8),
181                RowTiers::Canonical
182            );
183        }
184    }
185}
metrale_model_layers/layers/row_tiers.rs

metrale_model_layers/layers/
row_tiers.rs
