metrale_model_layers/layers/
expert_quantization.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-27: `--expert-quantization`: the precision a native-FP8 checkpoint's routed MoE
4//! experts are decoded at. `fp8` (the default) decodes the checkpoint's FP8 experts. The two
5//! NVFP4 tiers add an NVFP4 copy of routed-expert projections at load and run every MoE decode
6//! of 1 to 64 rows through the grouped decode (`moe/forward_nvfp4_grouped_decode.rs`), whose
7//! output for a row does not depend on how many rows share the call; the shared expert and
8//! prefill keep the FP8 weights under both.
9//!
10//! It lowers the routed experts' precision in decode, so it is `fp8` unless the serve command
11//! line asks otherwise; there is no environment fallback. The loader reads it
12//! (`qwen35/load_layers.rs`) beside the diagnostic `METRALE_FORCE_NVFP4_MOE`, which loads
13//! NVFP4 experts only and keeps today's NVFP4 arms.
14//!
15//! Owner: model-layers (MoE).
16//! Invariants:
17//! - The first publication or read wins (`OnceLock`); the serve publishes before the model is
18//!   built, and anything that reads first fixes it at `fp8`.
19//! - Under an NVFP4 tier a row's MoE output bits do not depend on how many rows share the
20//!   call (`FfnComponent::nvfp4_grouped_ok` admits every width the decode arms pass it).
21
22use std::sync::OnceLock;
23
24use anyhow::Result;
25use metrale_gpu_runtime::gpu::DevicePtr;
26
27use crate::layer::ForwardContext;
28
29/// 2026-09-27: The expert precision tiers. Measured on GB10 (Qwen3.6-35B-A3B-FP8, canonical
30/// tiers, `--mtp-gate force`):
31///
32/// | tier | BFCL echolp N=1004 overall/normalized | agentic-webserver | C16 tok/s, J/tok (dgx1) |
33/// |---|---|---|---|
34/// | fp8 | 84.96/86.03 | pass, 535-546 s | 315.5, 0.227 |
35/// | nvfp4-gate-up | 84.86/85.33 | pass 3/3, 504-533 s | 359.6, 0.211 |
36/// | nvfp4 | 85.46/86.57 | FAIL: 168 turns, 774 s over the 700 s ceiling | 383.1, 0.201 |
37///
38/// Weights resident before the KV cache: fp8 38.4 GB, nvfp4-gate-up 49.6 GB, nvfp4 55.2 GB.
39#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
40pub enum ExpertQuantization {
41    /// 2026-09-27: The checkpoint's FP8 experts: most stable, still fast.
42    #[default]
43    Fp8,
44    /// 2026-09-27: Routed gate and up projections NVFP4 in decode; routed down, the shared expert and
45    /// prefill FP8. The stable speed lever.
46    Nvfp4GateUp,
47    /// 2026-09-27: Every routed projection NVFP4 in decode; the shared expert and prefill FP8. Fastest,
48    /// and it changes the agent's trajectories most.
49    Nvfp4,
50}
51
52impl ExpertQuantization {
53    /// 2026-09-27: Every tier, in the order the flag lists them.
54    pub const ALL: [Self; 3] = [Self::Fp8, Self::Nvfp4GateUp, Self::Nvfp4];
55
56    /// 2026-09-27: The tier's flag value and recipe value.
57    pub fn name(self) -> &'static str {
58        match self {
59            Self::Fp8 => "fp8",
60            Self::Nvfp4GateUp => "nvfp4-gate-up",
61            Self::Nvfp4 => "nvfp4",
62        }
63    }
64
65    /// 2026-09-27: Whether MoE decode takes the grouped NVFP4 path.
66    pub fn nvfp4_decode(self) -> bool {
67        self != Self::Fp8
68    }
69
70    /// 2026-09-27: Whether the routed down projections decode as NVFP4, so their NVFP4 copy
71    /// is built at load; under `nvfp4-gate-up` they decode from the FP8 experts.
72    pub fn nvfp4_down(self) -> bool {
73        self == Self::Nvfp4
74    }
75}
76
77static EXPERT_QUANTIZATION: OnceLock<ExpertQuantization> = OnceLock::new();
78
79/// 2026-09-27: Publish `--expert-quantization`. Returns the tier in force; a caller that gets a
80/// different one should warn.
81pub fn set_expert_quantization_from_cli(q: ExpertQuantization) -> ExpertQuantization {
82    let _ = EXPERT_QUANTIZATION.set(q);
83    *EXPERT_QUANTIZATION.get().expect("just set")
84}
85
86/// 2026-09-27: The `--expert-quantization` tier in force: `fp8` unless the serve published
87/// another.
88pub fn expert_quantization() -> ExpertQuantization {
89    *EXPERT_QUANTIZATION.get_or_init(ExpertQuantization::default)
90}
91
92impl super::FfnComponent {
93    /// 2026-09-27: Whether this is a MoE whose grouped NVFP4 decode serves `m` rows
94    /// (`MoeLayer::nvfp4_grouped_decode_ok`).
95    pub fn nvfp4_grouped_ok(&self, m: usize, ctx: &ForwardContext) -> bool {
96        matches!(self, Self::Moe(moe) if moe.nvfp4_grouped_decode_ok(m, ctx))
97    }
98
99    /// 2026-09-27: The grouped NVFP4 MoE decode over `[m, H]` rows into `moe_output()`.
100    /// Errors for dense and none; callers gate on [`Self::nvfp4_grouped_ok`] first.
101    pub fn forward_nvfp4_grouped(
102        &self,
103        input: DevicePtr,
104        m: usize,
105        ctx: &ForwardContext,
106        stream: u64,
107    ) -> Result<()> {
108        match self {
109            Self::Moe(moe) => moe.forward_nvfp4_grouped_decode(input, m, ctx, stream),
110            _ => anyhow::bail!("forward_nvfp4_grouped is MoE-only (m={m})"),
111        }
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use super::ExpertQuantization as Q;
118
119    /// 2026-09-27: The names are distinct and each tier's decode and down precision.
120    #[test]
121    fn tiers_name_and_split_precision() {
122        let names: Vec<_> = Q::ALL.iter().map(|q| q.name()).collect();
123        assert_eq!(names, ["fp8", "nvfp4-gate-up", "nvfp4"]);
124        assert_eq!(Q::default(), Q::Fp8);
125        assert!(!Q::Fp8.nvfp4_decode() && !Q::Fp8.nvfp4_down());
126        assert!(Q::Nvfp4GateUp.nvfp4_decode() && !Q::Nvfp4GateUp.nvfp4_down());
127        assert!(Q::Nvfp4.nvfp4_decode() && Q::Nvfp4.nvfp4_down());
128    }
129}
metrale_model_layers/layers/expert_quantization.rs

metrale_model_layers/layers/
expert_quantization.rs
