metrale_model_layers/layers/
moe_grouped_decode.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Row threshold and switches for the grouped-GEMM MoE decode arm.
4//!
5//! Owner: model-layers (MoE).
6//! Invariants:
7//! - `moe_grouped_decode_for` is false at every width while
8//!   `METRALE_NO_MOE_GROUPED_DECODE` is set, even when the arm is forced.
9
10/// 2026-09-25: Minimum rows for the grouped-GEMM MoE decode arm. The SSM stack
11/// (`qwen3_ssm::trait_decode_multi_seq`) and the attention layers' multi-seq FFN
12/// both decide through `moe_grouped_decode_for`, so they share this width. The
13/// arm reads each routed expert once for all rows but first sorts and permutes
14/// the rows, a fixed cost that the threshold keeps off narrow batches.
15pub fn moe_grouped_decode_min_rows() -> usize {
16    16
17}
18
19/// 2026-09-25: False when `METRALE_NO_MOE_GROUPED_DECODE` is set to any value,
20/// `0` included. Read once per process.
21pub fn moe_grouped_decode_enabled() -> bool {
22    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
23    *ON.get_or_init(|| std::env::var_os("METRALE_NO_MOE_GROUPED_DECODE").is_none())
24}
25
26/// 2026-09-25: `METRALE_MOE_GROUPED_DECODE=1` runs the arm below
27/// `moe_grouped_decode_min_rows()` too, for measuring narrow widths. Read once
28/// per process.
29pub fn moe_grouped_decode_forced() -> bool {
30    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
31    *ON.get_or_init(|| std::env::var("METRALE_MOE_GROUPED_DECODE").as_deref() == Ok("1"))
32}
33
34/// 2026-09-25: Whether the arm runs for `n` rows, given the two switches. It
35/// reads no env, so tests do not depend on the latched `OnceLock`s above.
36pub fn moe_grouped_decode_decide(n: usize, enabled: bool, forced: bool) -> bool {
37    enabled && (n >= moe_grouped_decode_min_rows() || forced)
38}
39
40/// 2026-09-25: Whether the arm runs for `n` rows under this process's switches.
41pub fn moe_grouped_decode_for(n: usize) -> bool {
42    moe_grouped_decode_decide(n, moe_grouped_decode_enabled(), moe_grouped_decode_forced())
43}
44
45#[cfg(test)]
46#[path = "moe_grouped_decode_tests.rs"]
47mod moe_grouped_decode_tests;
metrale_model_layers/layers/moe_grouped_decode.rs

metrale_model_layers/layers/
moe_grouped_decode.rs
