metrale_model_layers/layers/mtp_head/
new_bf16_moe.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-10-02: Construction of the MTP drafter's MoE layer on the checkpoint's BF16 experts (an
4//! MTP head excluded from quantization, as nvidia/Qwen3.6-35B-A3B-NVFP4 ships it), so the batched
5//! propose runs the drafter's MoE grouped (`MoeLayer::forward_bf16_grouped_decode`) instead of
6//! per-expert GEMVs per sequence (`moe_forward_generic`).
7//!
8//! `--mtp-experts-nvfp4` instead requantizes those experts to NVFP4 at load, for drafting only.
9//!
10//! Owner: model-layers (MTP head).
11//! Invariants: the switch is read by `MtpHead::new` only after the serve published it.
12
13use std::sync::OnceLock;
14
15use anyhow::Result;
16use metrale_gpu_runtime::gpu::GpuBackend;
17
18use super::MtpHead;
19use crate::layers::MoeLayer;
20use crate::weight_map::{DenseWeight, ExpertWeight, MoeWeights, MtpWeights, quantize_to_nvfp4};
21
22static MTP_EXPERTS_NVFP4: OnceLock<bool> = OnceLock::new();
23
24/// 2026-10-02: Publish `--mtp-experts-nvfp4`; returns the value in force, which differs when
25/// something read it first.
26pub fn set_mtp_experts_nvfp4_from_cli(on: bool) -> bool {
27    let _ = MTP_EXPERTS_NVFP4.set(on);
28    *MTP_EXPERTS_NVFP4.get().expect("just set")
29}
30
31/// 2026-10-02: Whether a BF16 MoE MTP head drafts on NVFP4 copies of its experts. Off unless the
32/// serve published otherwise: the head's declared BF16 is the default.
33pub fn mtp_experts_nvfp4() -> bool {
34    *MTP_EXPERTS_NVFP4.get_or_init(|| false)
35}
36
37impl MtpHead {
38    /// 2026-10-02: The drafter's MoE on the checkpoint's BF16 tensors: null NVFP4 routed and shared
39    /// slots, the BF16 router (`gate_nvfp4 = None`), then `MoeLayer::set_bf16_experts` with the
40    /// routed experts' and the shared expert's BF16 weights (no copy). Errors when the expert count
41    /// differs from the config's.
42    pub(super) fn new_bf16_moe(
43        weights: &MtpWeights,
44        config: &metrale_config::ModelConfig,
45        gpu: &dyn GpuBackend,
46    ) -> Result<MoeLayer> {
47        anyhow::ensure!(
48            weights.experts.len() == config.num_experts,
49            "MTP BF16 experts: {} loaded, config says {}",
50            weights.experts.len(),
51            config.num_experts
52        );
53        let moe_weights = MoeWeights {
54            gate: weights.moe_gate,
55            shared_expert: ExpertWeight::null(),
56            shared_expert_gate: weights.shared_expert_gate,
57            experts: vec![ExpertWeight::null(); config.num_experts],
58            router_pre_norm: None,
59            correction_bias: None,
60        };
61        let mut moe = MoeLayer::new(moe_weights, config.num_experts, None, gpu, config)?;
62        let col =
63            |f: fn(&crate::weight_map::DenseExpertWeight) -> DenseWeight| -> Vec<DenseWeight> {
64                weights.experts.iter().map(f).collect()
65            };
66        moe.set_bf16_experts(
67            &col(|e| e.gate_proj),
68            &col(|e| e.up_proj),
69            &col(|e| e.down_proj),
70            weights.shared_expert.gate_proj.weight,
71            weights.shared_expert.up_proj.weight,
72            weights.shared_expert.down_proj.weight,
73            gpu,
74        )?;
75        Ok(moe)
76    }
77}
78
79impl MtpHead {
80    /// 2026-10-02: The drafter's MoE with its BF16 experts requantized to NVFP4 at load (a
81    /// draft-only precision: the drafter proposes, the target verifies, so output tokens do not
82    /// change, only how many drafts are accepted). The router stays BF16. The layer runs the grouped
83    /// NVFP4 tensor-core decode at every width (`set_draft_nvfp4_experts`), reading 3.5x fewer
84    /// bytes per expert than the BF16 point.
85    pub(super) fn new_nvfp4_draft_moe(
86        weights: &MtpWeights,
87        config: &metrale_config::ModelConfig,
88        gpu: &dyn GpuBackend,
89        absmax_k: metrale_gpu_runtime::gpu::KernelHandle,
90        nvfp4_k: metrale_gpu_runtime::gpu::KernelHandle,
91        stream: u64,
92    ) -> Result<MoeLayer> {
93        let (h, inter) = (config.hidden_size, config.moe_intermediate_size);
94        let shared_inter = config.shared_expert_intermediate_size;
95        let q = |w: &DenseWeight, n: usize, k: usize| {
96            quantize_to_nvfp4(w, n, k, gpu, absmax_k, nvfp4_k, stream)
97        };
98        let experts = weights
99            .experts
100            .iter()
101            .map(|e| -> Result<ExpertWeight> {
102                Ok(ExpertWeight {
103                    gate_proj: q(&e.gate_proj, inter, h)?,
104                    up_proj: q(&e.up_proj, inter, h)?,
105                    down_proj: q(&e.down_proj, h, inter)?,
106                })
107            })
108            .collect::<Result<Vec<_>>>()?;
109        let sh = &weights.shared_expert;
110        let moe_weights = MoeWeights {
111            gate: weights.moe_gate,
112            shared_expert: ExpertWeight {
113                gate_proj: q(&sh.gate_proj, shared_inter, h)?,
114                up_proj: q(&sh.up_proj, shared_inter, h)?,
115                down_proj: q(&sh.down_proj, h, shared_inter)?,
116            },
117            shared_expert_gate: weights.shared_expert_gate,
118            experts,
119            router_pre_norm: None,
120            correction_bias: None,
121        };
122        let mut moe = MoeLayer::new(moe_weights, config.num_experts, None, gpu, config)?;
123        moe.set_draft_nvfp4_experts()?;
124        tracing::info!(
125            "--mtp-experts-nvfp4: the MTP head's {} routed experts and shared expert draft as NVFP4 \
126             (requantized from BF16, draft-only); router BF16",
127            weights.experts.len()
128        );
129        Ok(moe)
130    }
131}
metrale_model_layers/layers/mtp_head/new_bf16_moe.rs

metrale_model_layers/layers/mtp_head/
new_bf16_moe.rs
