metrale_model_layers/layers/moe/
nvfp4_lean.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-10-04: The lean point of the tensor-core grouped NVFP4 decode (`Nvfp4G16Lean`,
4//! `kernels/gb10/common/tc_weight_formats.cuh`): the layer's declared NVFP4 experts, routed and
5//! shared, are repacked IN PLACE once at load (`nvfp4_tc_lean_repack`: same bytes, same size, no
6//! extra memory) into a tile-contiguous layout whose E2M1 bits a fragment extracts with rotates
7//! and masks and whose scale bytes become one IMAD each. The `_lean` expert kernels then return
8//! the row-major tensor-core kernels' output bytes exactly (`nvfp4_moe_grouped_microtest`,
9//! `all-nvfp4-tc-lean`) at 10-17 % less expert-kernel energy at equal time.
10//!
11//! Why in place: the row-major tables are read by nothing else once the prefill copies exist,
12//! and a second copy of every expert would cost the layer's expert memory again.
13//!
14//! Owner: model-layers (MoE).
15//! Invariants:
16//! - A layer is repacked only when its prefill reads copies (the `*_t` tables built, no CUTLASS
17//!   tables, no unified layout), its experts are the checkpoint's declared NVFP4 and all local,
18//!   the tensor-core decode and its lean pair resolved for its shape, and no lever sends its
19//!   decode to the legacy kernels (FP32 gate or routing, grouped decode off).
20//! - After the repack `Nvfp4GroupedKernels::lean` holds and only the lean pair reads the tables;
21//!   every arm that would read them in the row-major layout calls `refuse_lean_layout` first.
22
23use anyhow::Context;
24
25use super::forward_nvfp4_grouped_decode::{Nvfp4GroupedKernels, Nvfp4GroupedLaunch};
26use super::*;
27
28impl Nvfp4GroupedKernels {
29    /// 2026-10-04: The lean pair, once this layer's tables are lean.
30    pub(super) fn lean_launch(&self) -> Option<Nvfp4GroupedLaunch> {
31        self.lean.then_some(Nvfp4GroupedLaunch {
32            gate_up: self.gate_up_tc_lean,
33            gate_up_geometry: ops::NVFP4_GROUPED_GATE_UP_TC,
34            down: self.down_tc_lean,
35            down_geometry: ops::NVFP4_GROUPED_DOWN_TC,
36            max_rows: super::NVFP4_GROUPED_DECODE_TC_MAX_ROWS,
37        })
38    }
39}
40
41impl MoeLayer {
42    /// 2026-10-04: Whether this layer's NVFP4 expert tables are in the lean layout.
43    pub fn nvfp4_experts_lean(&self) -> bool {
44        self.nvfp4_grouped.lean
45    }
46
47    /// 2026-10-04: Fails an arm that would read this layer's NVFP4 expert tables in the
48    /// row-major layout once they are lean (the bytes would decode as other values).
49    pub(super) fn refuse_lean_layout(&self, arm: &str) -> Result<()> {
50        anyhow::ensure!(
51            !self.nvfp4_grouped.lean,
52            "{arm}: this layer's NVFP4 experts were repacked for the lean tensor-core decode at \
53             load; only `forward_nvfp4_grouped_decode` reads them (set METRALE_NO_MOE_NVFP4_TC to \
54             keep the row-major layout)"
55        );
56        Ok(())
57    }
58
59    /// 2026-10-04: Why this layer keeps the row-major layout, or `None` when it qualifies for the
60    /// lean one (module invariants). Call after the load-time prefill passes.
61    pub fn nvfp4_lean_refusal(&self, config: &metrale_config::ModelConfig) -> Option<&'static str> {
62        let (h, inter) = (
63            config.hidden_size as u32,
64            config.moe_intermediate_size as u32,
65        );
66        let k = &self.nvfp4_grouped;
67        let levers = ops::ModelLevers::get();
68        let local = |q: &QuantizedWeight| !q.weight.is_null() && !q.weight_scale.is_null();
69        let checks: [(bool, &'static str); 12] = [
70            (
71                k.declared_experts
72                    && self.fp8_shared_expert.is_none()
73                    && self.fp8_down_weight_ptrs.is_none(),
74                "experts are not the checkpoint's declared NVFP4 alone (the FP8 kernels would read \
75                 the lifted SiLU rows)",
76            ),
77            (
78                super::forward_nvfp4_grouped_decode::nvfp4_grouped_tc_enabled(),
79                "METRALE_NO_MOE_NVFP4_TC",
80            ),
81            (
82                k.gate_up_tc_lean.0 != 0 && k.down_tc_lean.0 != 0 && k.lean_repack.0 != 0,
83                "lean kernels not in the kernel set",
84            ),
85            (
86                ops::nvfp4_grouped_tc_shape_ok(inter, h, ops::NVFP4_GROUPED_GATE_UP_TC)
87                    && ops::nvfp4_grouped_tc_shape_ok(h, inter, ops::NVFP4_GROUPED_DOWN_TC)
88                    && ops::nvfp4_lean_repack_shape_ok(inter, h)
89                    && ops::nvfp4_lean_repack_shape_ok(h, inter)
90                    && config.shared_expert_intermediate_size == config.moe_intermediate_size,
91                "shape",
92            ),
93            (
94                self.gate_ptrs_t.is_some()
95                    && self.up_ptrs_t.is_some()
96                    && self.down_ptrs_t.is_some()
97                    && self.shared_gate_t.is_some()
98                    && self.shared_up_t.is_some()
99                    && self.shared_down_t.is_some(),
100                "prefill reads the row-major tables (no transposed copies)",
101            ),
102            (
103                self.cutlass_grouped_host.is_none() && !self.unified_layout,
104                "CUTLASS grouped prefill or the unified layout",
105            ),
106            (
107                !self.weights.experts.is_empty()
108                    && self
109                        .weights
110                        .experts
111                        .iter()
112                        .all(|e| local(&e.gate_proj) && local(&e.up_proj) && local(&e.down_proj))
113                    && config.ep_world_size <= 1,
114                "experts not all local",
115            ),
116            (
117                local(&self.weights.shared_expert.gate_proj)
118                    && local(&self.weights.shared_expert.up_proj)
119                    && local(&self.weights.shared_expert.down_proj),
120                "no NVFP4 shared expert",
121            ),
122            (
123                self.experts_scale_kind == crate::weight_map::WeightQuantFormat::Nvfp4
124                    && self.shared_experts_scale_kind
125                        == crate::weight_map::WeightQuantFormat::Nvfp4,
126                "scale format",
127            ),
128            (
129                !levers.fp32_gate && !self.fp32_routing_active(levers),
130                "FP32 gate or routing",
131            ),
132            (
133                crate::layers::moe_grouped_decode_enabled(),
134                "METRALE_NO_MOE_GROUPED_DECODE",
135            ),
136            (self.lora.is_none(), "MoE LoRA"),
137        ];
138        checks.iter().find(|(ok, _)| !ok).map(|(_, why)| *why)
139    }
140
141    /// 2026-10-04: Repacks this layer's routed and shared NVFP4 experts in place for the lean
142    /// decode when it qualifies (`nvfp4_lean_refusal`), on `stream`, and synchronizes it. Returns
143    /// whether it did. Fails the load on a negative or NaN scale.
144    pub fn repack_nvfp4_experts_lean(
145        &mut self,
146        gpu: &dyn GpuBackend,
147        config: &metrale_config::ModelConfig,
148        stream: u64,
149    ) -> Result<bool> {
150        if let Some(why) = self.nvfp4_lean_refusal(config) {
151            tracing::debug!("NVFP4 experts keep the row-major layout: {why}");
152            return Ok(false);
153        }
154        let (h, inter) = (
155            config.hidden_size as u32,
156            config.moe_intermediate_size as u32,
157        );
158        let kernel = self.nvfp4_grouped.lean_repack;
159        let bad = gpu.alloc(16)?;
160        gpu.copy_h2d(&[0u8; 4], bad)?;
161        let sh = &self.weights.shared_expert;
162        let mats = self
163            .weights
164            .experts
165            .iter()
166            .chain(std::iter::once(sh))
167            .flat_map(|e| {
168                [
169                    (&e.gate_proj, inter, h),
170                    (&e.up_proj, inter, h),
171                    (&e.down_proj, h, inter),
172                ]
173            });
174        let launched = mats
175            .map(|(q, n, k)| {
176                ops::nvfp4_tc_lean_repack(gpu, kernel, q.weight, q.weight_scale, n, k, bad, stream)
177            })
178            .collect::<Result<Vec<()>>>();
179        let synced = launched.and_then(|_| gpu.synchronize(stream));
180        let mut flag = [0u8; 4];
181        let read = synced.and_then(|()| gpu.copy_d2h(bad, &mut flag));
182        gpu.free(bad)?;
183        read.context("nvfp4_tc_lean_repack")?;
184        anyhow::ensure!(
185            flag == [0u8; 4],
186            "nvfp4_tc_lean_repack: a negative or NaN E4M3 block scale in this layer's NVFP4 experts"
187        );
188        self.nvfp4_grouped.lean = true;
189        Ok(true)
190    }
191}
metrale_model_layers/layers/moe/nvfp4_lean.rs

metrale_model_layers/layers/moe/
nvfp4_lean.rs
