metrale_model_arch/nemotron_moe/
prefill_weights.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: `NemotronMoeLayer` prefill-side weight preparation (the
4//! transposed-NVFP4 expert and shared-expert copies, and the FP8 shared-expert
5//! and latent-projection copies) and the dense-GEMM prefill dispatcher.
6//!
7//! Owner: model-arch (Nemotron-H).
8//! Invariants:
9//! - `prepare_prefill_weights` never fails: a copy that cannot be built
10//!   leaves its field `None` and prefill uses the next arm.
11
12use anyhow::Result;
13use metrale_config::ModelConfig;
14use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
15
16use super::NemotronMoeLayer;
17use super::build_ptr_table_from_weights;
18use metrale_model_layers::layers::ops;
19use metrale_model_layers::weight_map::DenseWeight;
20
21impl NemotronMoeLayer {
22    /// 2026-09-25: Dense BF16 GEMM for prefill: `dense_gemm_bf16_pipelined`
23    /// when it resolved, else `dense_gemm_bf16`. Used for the gate GEMM and,
24    /// without FP8 copies, the fc1 / fc2 latent GEMMs.
25    #[allow(clippy::too_many_arguments)]
26    pub(super) fn dense_gemm_prefill(
27        &self,
28        gpu: &dyn GpuBackend,
29        input: DevicePtr,
30        weight: &DenseWeight,
31        output: DevicePtr,
32        m: u32,
33        n: u32,
34        k: u32,
35        stream: u64,
36    ) -> Result<()> {
37        if self.dense_gemm_pipelined_k.0 != 0 {
38            ops::dense_gemm_bf16_pipelined(
39                gpu,
40                self.dense_gemm_pipelined_k,
41                input,
42                weight,
43                output,
44                m,
45                n,
46                k,
47                stream,
48            )
49        } else {
50            ops::dense_gemm(
51                gpu,
52                self.dense_gemm_k,
53                input,
54                weight,
55                output,
56                m,
57                n,
58                k,
59                stream,
60            )
61        }
62    }
63
64    /// 2026-09-29: Build the prefill-only weight copies; the loader calls it
65    /// after construction. Routed experts are transposed only when
66    /// `moe_latent_size == 0`. `shared_e4m3` (`--nemotron-shared-expert-e4m3`) builds the
67    /// shared expert's transposed copies, which the E4M3 tile GEMM reads; without it the
68    /// shared expert runs the BF16-MMA `w4a16_gemm` on the checkpoint's NVFP4 weights.
69    pub fn prepare_prefill_weights(
70        &mut self,
71        gpu: &dyn GpuBackend,
72        config: &ModelConfig,
73        shared_e4m3: bool,
74    ) {
75        let h = config.hidden_size;
76        let inter = self.moe_inter;
77        let shared_inter = config.shared_expert_intermediate_size;
78
79        // 2026-09-25: A table is installed only when every expert transposed;
80        // without it the sorted path uses the base grouped GEMM.
81        if self.moe_latent_size == 0 {
82            let expert_k = h;
83            let mut up_t = Vec::new();
84            let mut down_t = Vec::new();
85            for expert in &self.weights.experts {
86                if let Ok(ut) = expert.up_proj.transpose_for_gemm(gpu, inter, expert_k) {
87                    up_t.push(ut);
88                }
89                if let Ok(dt) = expert.down_proj.transpose_for_gemm(gpu, expert_k, inter) {
90                    down_t.push(dt);
91                }
92            }
93            if up_t.len() == self.weights.experts.len()
94                && let Ok(ptrs) = build_ptr_table_from_weights(&up_t, gpu)
95            {
96                self.up_ptrs_t = Some(ptrs);
97            }
98            if down_t.len() == self.weights.experts.len()
99                && let Ok(ptrs) = build_ptr_table_from_weights(&down_t, gpu)
100            {
101                self.down_ptrs_t = Some(ptrs);
102            }
103        }
104
105        // 2026-09-29: The shared expert's copies, in every layer shape: with
106        // `METRALE_SHARED_FP8_PREFILL` set (any value) and
107        // `fp8_gemm_t_m128_mfast` resolved, pre-dequantized FP8 copies; otherwise,
108        // or if either failed, transposed NVFP4 copies, and only under `shared_e4m3`:
109        // the tile GEMM that reads them casts activations and dequantized weights to
110        // E4M3 unscaled (saturating at 448; 1-30% local error measured on
111        // Nemotron-3-Nano), so it is opt-in. Under native FP8 the NVFP4 shared weights
112        // are `QuantizedWeight::null()` and every copy here is derived from them, so
113        // none is built.
114        let native_shared =
115            self.weights.shared_up_fp8.is_some() || self.weights.shared_down_fp8.is_some();
116        let fp8_prefill = !native_shared && std::env::var("METRALE_SHARED_FP8_PREFILL").is_ok();
117        if fp8_prefill
118            && self.fp8_gemm_m128_k.0 != 0
119            && let Ok(pdq_k) = gpu.kernel("w4a16", "predequant_nvfp4_to_fp8")
120        {
121            self.shared_up_pd_fp8 = self
122                .weights
123                .shared_up
124                .predequant_to_fp8(gpu, pdq_k, shared_inter, h, 0)
125                .ok();
126            self.shared_down_pd_fp8 = self
127                .weights
128                .shared_down
129                .predequant_to_fp8(gpu, pdq_k, h, shared_inter, 0)
130                .ok();
131        }
132        if !native_shared
133            && shared_e4m3
134            && (self.shared_up_pd_fp8.is_none() || self.shared_down_pd_fp8.is_none())
135        {
136            self.shared_up_t = self
137                .weights
138                .shared_up
139                .transpose_for_gemm(gpu, shared_inter, h)
140                .ok();
141            self.shared_down_t = self
142                .weights
143                .shared_down
144                .transpose_for_gemm(gpu, h, shared_inter)
145                .ok();
146        }
147
148        // 2026-09-25: Under the same `fp8_prefill` gate, FP8 E4M3 copies of the
149        // BF16 fc1 / fc2 latent projections for `fp8_gemm_t_m128_mfast`.
150        let lat = self.moe_latent_size;
151        if lat > 0
152            && fp8_prefill
153            && self.fp8_gemm_m128_k.0 != 0
154            && let Ok(b2f) = gpu.kernel("w4a16", "bf16_to_fp8")
155        {
156            let conv = |w: &DenseWeight, n: usize, k: usize| -> Option<DevicePtr> {
157                let dst = gpu.alloc(n * k).ok()?;
158                metrale_model_layers::layers::ops::bf16_to_fp8(
159                    gpu,
160                    b2f,
161                    w.weight,
162                    dst,
163                    (n * k) as u32,
164                    0,
165                )
166                .ok()?;
167                gpu.synchronize(0).ok()?;
168                Some(dst)
169            };
170            self.fc1_pd_fp8 = self
171                .weights
172                .fc1_latent_proj
173                .as_ref()
174                .and_then(|w| conv(w, lat, h));
175            self.fc2_pd_fp8 = self
176                .weights
177                .fc2_latent_proj
178                .as_ref()
179                .and_then(|w| conv(w, h, lat));
180        }
181    }
182}
183
184#[cfg(test)]
185#[path = "prefill_weights_tests.rs"]
186mod tests;
metrale_model_arch/nemotron_moe/prefill_weights.rs

metrale_model_arch/nemotron_moe/
prefill_weights.rs
