metrale_model_arch/nemotron_moe/
shared_e4m3.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-29: The `--nemotron-shared-expert-e4m3` switch, published once by the serve.
4//!
5//! Owner: model-arch (Nemotron-H).
6//! Invariants: read by the weight loader only after the serve published it.
7
8use std::sync::OnceLock;
9
10static E4M3: OnceLock<bool> = OnceLock::new();
11
12/// 2026-09-29: Publish the command line's value; returns the value in force, which differs
13/// when something read it first.
14pub fn set_shared_expert_e4m3_from_cli(on: bool) -> bool {
15    let _ = E4M3.set(on);
16    *E4M3.get().expect("just set")
17}
18
19/// 2026-09-29: Whether the Nemotron-H shared expert's prefill runs the E4M3 tile GEMM. Off
20/// unless the serve published otherwise: the checkpoint's precision (NVFP4 weights, 16-bit
21/// activations on the BF16-MMA `w4a16_gemm`) is the default.
22pub fn shared_expert_e4m3() -> bool {
23    *E4M3.get_or_init(|| false)
24}