metrale_model_layers/layers/
moe_expert_tables.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-10-02: Whether a MoE's transposed NVFP4 prefill tables (the routed and shared experts'
4//! K-major copies `MoeLayer::transpose_for_prefill` builds) are built. The serve decides once,
5//! before the model loads, from its memory plan (`met circuit memory`'s model at the serve's
6//! declared `--max-seq-len` and batch): the tables are built only when the plan fits the util
7//! budget with them. Without them MoE prefill runs the grouped GEMM on the row-major experts.
8//!
9//! Owner: model-layers (MoE).
10//! Invariants:
11//! - Published once (`OnceLock`); a second publication does not change it.
12//! - There is no default: a loader that would build the tables and finds nothing published
13//!   refuses to load ([`moe_expert_tables`] is `None`).
14
15use std::sync::OnceLock;
16
17/// 2026-10-02: The decision.
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub enum MoeExpertTables {
20    /// 2026-10-02: The plan fits with the tables.
21    Build,
22    /// 2026-10-02: The plan does not fit with them.
23    Skip,
24}
25
26impl MoeExpertTables {
27    /// 2026-10-02: `build` or `skip`, the spelling the memory plan's copy rules and a gate
28    /// record's `serve_resolved` use.
29    pub fn name(self) -> &'static str {
30        match self {
31            Self::Build => "build",
32            Self::Skip => "skip",
33        }
34    }
35}
36
37static VALUE: OnceLock<MoeExpertTables> = OnceLock::new();
38
39/// 2026-10-02: Publish the serve's decision. Returns the decision in force; a caller that gets a
40/// different one should warn.
41pub fn set_moe_expert_tables_from_plan(t: MoeExpertTables) -> MoeExpertTables {
42    let _ = VALUE.set(t);
43    *VALUE.get().expect("just set")
44}
45
46/// 2026-10-02: The published decision, if any.
47pub fn moe_expert_tables() -> Option<MoeExpertTables> {
48    VALUE.get().copied()
49}
metrale_model_layers/layers/moe_expert_tables.rs

metrale_model_layers/layers/
moe_expert_tables.rs
