metrale_model_layers/layers/moe/
forward_nvfp4_grouped_decode.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-27: `MoeLayer::forward_nvfp4_grouped_decode`: the NVFP4 MoE of `m` rows in one
4//! routed+shared expert dispatch, under an NVFP4 `--expert-quantization` tier.
5//!
6//! The steps are those of the grouped FP8 decode (`forward_fp8_grouped_decode.rs`): the
7//! per-row router (`GroupedRouting::PerRow`), `moe_fp8_grouped_sort` (the slot sort and the
8//! active-expert list), gate+up and SiLU, down, and the grouped blend. The expert kernels
9//! (`ops/nvfp4_moe_grouped.rs`) read each active expert's NVFP4 weights once per pass for all
10//! the rows routed to it. When the layer keeps its FP8 experts too (the qwen35 loader does,
11//! for a native-FP8 checkpoint), the shared expert runs in FP8 through the grouped FP8 kernels
12//! (with no active experts), the NVFP4 kernels take only the routed experts, and prefill keeps
13//! the FP8 experts; under `nvfp4-gate-up` every down projection, routed and shared, runs
14//! through the grouped FP8 down kernel instead. Every step computes a row independently of
15//! the other rows, so a row's output bits do not depend on `m`: the path serves every width
16//! from one row up, and a row-invariant tier policy (`row_tiers.rs`) holds with it on.
17//! `model-arch/examples/nvfp4_moe_grouped_microtest.rs` checks this.
18//!
19//! Owner: model-layers (MoE).
20//! Invariants: `forward_nvfp4_grouped_decode` launches nothing unless
21//! `nvfp4_grouped_decode_ok(m, ctx)` holds.
22
23use super::*;
24
25/// 2026-09-27: Widest row count this path admits on the CUDA-core expert kernels, the grouped
26/// FP8 decode's.
27pub const NVFP4_GROUPED_DECODE_MAX_ROWS: usize =
28    super::forward_fp8_grouped_decode::FP8_GROUPED_DECODE_MAX_ROWS;
29
30/// 2026-10-02: Widest row count on the tensor-core expert kernels, the grouped FP8 tensor-core
31/// decode's: one verify of 128 sequences at one draft.
32pub const NVFP4_GROUPED_DECODE_TC_MAX_ROWS: usize =
33    super::forward_fp8_grouped_decode::FP8_GROUPED_DECODE_TC_MAX_ROWS;
34
35/// 2026-09-27: The two expert kernels of this path, looked up with `try_kernel`; a zero handle
36/// declines it. The sort, router and blend are the grouped FP8 decode's.
37/// 2026-10-02: `gate_up_tc` / `down_tc` are the tensor-core twins (`moe_nvfp4_grouped_tc.cu`),
38/// taken when both resolved unless `METRALE_NO_MOE_NVFP4_TC` is present. `declared_experts`: the
39/// layer's routed and shared experts are the checkpoint's own NVFP4 (declared W4A16, not a
40/// requantized copy), so decode takes this path at every width whatever the
41/// `--expert-quantization` tier (which governs FP8 checkpoints only); set by the qwen35 loader
42/// under `--weight-quantization declared` ([`MoeLayer::set_declared_nvfp4_experts`]).
43pub(super) struct Nvfp4GroupedKernels {
44    pub gate_up: KernelHandle,
45    pub down: KernelHandle,
46    pub gate_up_tc: KernelHandle,
47    pub down_tc: KernelHandle,
48    pub declared_experts: bool,
49    /// 2026-10-04: The lean point's pair (`Nvfp4G16Lean`) and its load-time repack; `lean` is set
50    /// once this layer's tables were repacked (`nvfp4_lean.rs`), after which only that pair reads them.
51    pub gate_up_tc_lean: KernelHandle,
52    pub down_tc_lean: KernelHandle,
53    pub lean_repack: KernelHandle,
54    pub lean: bool,
55    /// 2026-10-02: The BF16 point's pair (`moe_bf16_grouped_tc.cu`, `forward_bf16_grouped_decode.rs`).
56    pub bf16_gate_up_tc: KernelHandle,
57    pub bf16_down_tc: KernelHandle,
58}
59
60impl Nvfp4GroupedKernels {
61    /// 2026-09-27: One direct `try_kernel` call per kernel (`#[track_caller]` audit lines).
62    pub(super) fn resolve(gpu: &dyn GpuBackend) -> Self {
63        use super::super::try_kernel;
64        const MODULE: &str = "moe_nvfp4_grouped";
65        const TC: &str = "moe_nvfp4_grouped_tc";
66        Self {
67            gate_up: try_kernel(gpu, MODULE, "moe_expert_gate_up_act_nvfp4_grouped"),
68            down: try_kernel(gpu, MODULE, "moe_expert_down_act_nvfp4_grouped"),
69            gate_up_tc: try_kernel(gpu, TC, "moe_expert_gate_up_act_nvfp4_grouped_tc"),
70            down_tc: try_kernel(gpu, TC, "moe_expert_down_act_nvfp4_grouped_tc"),
71            declared_experts: false,
72            gate_up_tc_lean: try_kernel(gpu, TC, "moe_expert_gate_up_act_nvfp4_grouped_tc_lean"),
73            down_tc_lean: try_kernel(gpu, TC, "moe_expert_down_act_nvfp4_grouped_tc_lean"),
74            lean_repack: try_kernel(gpu, TC, "nvfp4_tc_lean_repack"),
75            lean: false,
76            bf16_gate_up_tc: try_kernel(
77                gpu,
78                "moe_bf16_grouped_tc",
79                "moe_expert_gate_up_act_bf16_grouped_tc",
80            ),
81            bf16_down_tc: try_kernel(
82                gpu,
83                "moe_bf16_grouped_tc",
84                "moe_expert_down_act_bf16_grouped_tc",
85            ),
86        }
87    }
88
89    /// 2026-10-02: The gate+up and down launches for an `inter` x `hidden` expert: the
90    /// tensor-core twins when on and the shape fits them, else the CUDA-core kernels.
91    /// 2026-10-04: Always the lean pair once the tables are lean (the repack checked the shape).
92    fn select(&self, hidden: u32, inter: u32) -> Nvfp4GroupedLaunch {
93        if let Some(lean) = self.lean_launch() {
94            return lean;
95        }
96        if nvfp4_grouped_tc_enabled()
97            && self.gate_up_tc.0 != 0
98            && self.down_tc.0 != 0
99            && ops::nvfp4_grouped_tc_shape_ok(inter, hidden, ops::NVFP4_GROUPED_GATE_UP_TC)
100            && ops::nvfp4_grouped_tc_shape_ok(hidden, inter, ops::NVFP4_GROUPED_DOWN_TC)
101        {
102            Nvfp4GroupedLaunch {
103                gate_up: self.gate_up_tc,
104                gate_up_geometry: ops::NVFP4_GROUPED_GATE_UP_TC,
105                down: self.down_tc,
106                down_geometry: ops::NVFP4_GROUPED_DOWN_TC,
107                max_rows: NVFP4_GROUPED_DECODE_TC_MAX_ROWS,
108            }
109        } else {
110            Nvfp4GroupedLaunch {
111                gate_up: self.gate_up,
112                gate_up_geometry: ops::NVFP4_GROUPED_GATE_UP_SCALAR,
113                down: self.down,
114                down_geometry: ops::NVFP4_GROUPED_DOWN_SCALAR,
115                max_rows: NVFP4_GROUPED_DECODE_MAX_ROWS,
116            }
117        }
118    }
119}
120
121/// 2026-10-02: The expert kernels one grouped NVFP4 decode launches, and the widest row count
122/// they admit.
123pub(super) struct Nvfp4GroupedLaunch {
124    pub(super) gate_up: KernelHandle,
125    pub(super) gate_up_geometry: ops::Fp8GroupedGeometry,
126    pub(super) down: KernelHandle,
127    pub(super) down_geometry: ops::Fp8GroupedGeometry,
128    pub(super) max_rows: usize,
129}
130
131/// 2026-10-02: The grouped NVFP4 decode takes the tensor-core expert kernels unless
132/// `METRALE_NO_MOE_NVFP4_TC` is present (a debugging kill switch: the CUDA-core kernels). Read
133/// once per process. The two pairs differ in summation order, and in the SiLU product's
134/// carrier (FP32 against BF16 hi + lo), so their bits differ; each is row-invariant.
135pub(super) fn nvfp4_grouped_tc_enabled() -> bool {
136    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
137    *ON.get_or_init(|| std::env::var_os("METRALE_NO_MOE_NVFP4_TC").is_none())
138}
139
140/// 2026-09-27: Shape admission without a GPU: `m` in `1..=max_rows` (2026-10-02: the selected
141/// expert kernels' widest, `NVFP4_GROUPED_DECODE_MAX_ROWS` or `NVFP4_GROUPED_DECODE_TC_MAX_ROWS`),
142/// `hidden % 32 == 0` (gate+up reads 32-element chunks) and `inter % 16 == 0` (down reads
143/// 16-element blocks), and a shared expert as wide as a routed one (the shared SiLU product
144/// shares the routed layout).
145pub fn nvfp4_grouped_decode_shape_ok(
146    m: usize,
147    max_rows: usize,
148    hidden: usize,
149    inter: usize,
150    shared_inter: usize,
151) -> bool {
152    (1..=max_rows).contains(&m)
153        && hidden >= 32
154        && hidden.is_multiple_of(32)
155        && inter >= 16
156        && inter.is_multiple_of(16)
157        && shared_inter == inter
158}
159
160impl MoeLayer {
161    /// 2026-09-27: Whether `forward_nvfp4_grouped_decode` serves `m` rows on this layer: an
162    /// NVFP4 `--expert-quantization` tier is in force or (2026-10-02) the experts are the
163    /// checkpoint's declared NVFP4 (`set_declared_nvfp4_experts`), the routed projections that tier decodes
164    /// as NVFP4 are present in the row-major decode layout, the shared expert (and, under
165    /// `nvfp4-gate-up`, the routed down projections) are there in FP8 or NVFP4 as the tier
166    /// reads them, the per-row router applies (BF16 softmax gate, no correction bias,
167    /// pre-router norm, FP32 gate or FP32 routing, no DFlash reroute), the kernels resolved,
168    /// the shape is admitted, the arena is wide enough, and there is no LoRA, pre-expert norm,
169    /// hash routing or expert parallelism.
170    pub fn nvfp4_grouped_decode_ok(&self, m: usize, ctx: &ForwardContext) -> bool {
171        let cfg = ctx.config;
172        let (h, inter) = (cfg.hidden_size, cfg.moe_intermediate_size);
173        let need = super::forward_fp8_grouped_decode::grouped_decode_buffer_need(
174            m,
175            h,
176            inter,
177            cfg.num_experts,
178            cfg.num_experts_per_tok,
179        );
180        let b = ctx.buffers;
181        let tier = crate::layers::expert_quantization();
182        let fp8_shared_ok = self.nvfp4_decode_fp8_shared().is_some()
183            && self.moe_expert_gate_up_act_fp8_grouped_k.0 != 0
184            && self.moe_expert_down_act_fp8_grouped_k.0 != 0;
185        let nvfp4_shared_ok = !self.weights.shared_expert.gate_proj.weight.is_null()
186            && !self.weights.shared_expert.up_proj.weight.is_null()
187            && !self.weights.shared_expert.down_proj.weight.is_null();
188        // 2026-09-27: The routed NVFP4 projections must have been built: a layer loaded without
189        // them (the MTP head's native-FP8 MoE) has pointer tables of NULL entries, which the
190        // kernels would read as zero experts. Expert parallelism is refused below, so every
191        // expert is local and the first stands for all.
192        let routed = self.weights.experts.first();
193        let gate_up_ok =
194            routed.is_some_and(|e| !e.gate_proj.weight.is_null() && !e.up_proj.weight.is_null());
195        let native = self.nvfp4_grouped.declared_experts;
196        let launch = self.nvfp4_grouped.select(h as u32, inter as u32);
197        let down_ok = if tier.nvfp4_down() || native {
198            !self.down_ptrs.packed_ptrs.is_null()
199                && routed.is_some_and(|e| !e.down_proj.weight.is_null())
200        } else {
201            fp8_shared_ok && self.fp8_down_weight_ptrs.is_some()
202        };
203        (tier.nvfp4_decode() || native)
204            && nvfp4_grouped_decode_shape_ok(
205                m,
206                launch.max_rows,
207                h,
208                inter,
209                cfg.shared_expert_intermediate_size,
210            )
211            && launch.gate_up.0 != 0
212            && launch.down.0 != 0
213            && gate_up_ok
214            && down_ok
215            && (fp8_shared_ok || nvfp4_shared_ok)
216            && self.moe_weighted_sum_blend_fp8_grouped_k.0 != 0
217            && self.moe_fp8_grouped_sort_k.0 != 0
218            && cfg.num_experts <= ops::FP8_GROUPED_SORT_MAX_EXPERTS as usize
219            && self.moe_topk_softmax_rows_k.0 != 0
220            && self.router_gemv_batchm_k.0 != 0
221            && self.experts_scale_kind == crate::weight_map::WeightQuantFormat::Nvfp4
222            && self.shared_experts_scale_kind == crate::weight_map::WeightQuantFormat::Nvfp4
223            && !self.gate_ptrs.packed_ptrs.is_null()
224            && !self.up_ptrs.packed_ptrs.is_null()
225            && self.bf16_gate_weight_ptrs.is_none()
226            && self.bf16_shared_expert.is_none()
227            && !self.use_t_layout_for_decode()
228            && self.lora.is_none()
229            && self.pre_expert_norm.is_none()
230            && self.tid2eid_dev.is_none()
231            && self.router_logits_n as usize == cfg.num_experts
232            && self.gate_nvfp4.is_none()
233            && self.correction_bias_dev.is_none()
234            && self.weights.router_pre_norm.is_none()
235            && !ctx.levers.fp32_gate
236            && !self.fp32_routing_active(ctx.levers)
237            && !(self.is_dflash_capture_layer && ctx.levers.frankenstein_decode_via_prefill)
238            && h.is_multiple_of(8)
239            && !(ctx.comm.is_some() && cfg.ep_world_size > 1)
240            && b.scratch_bytes() >= need.scratch
241            && b.gate_logits_bytes() >= need.gate_logits
242            && b.expert_gate_out_bytes() >= need.expert_gate_out
243            && b.expert_down_out_bytes() >= need.expert_down_out
244            && b.logits_bytes() >= need.shared_act
245            && b.attn_output_bytes() >= need.row_hidden
246            && b.moe_output_bytes() >= need.row_hidden
247    }
248
249    /// 2026-10-02: Mark this layer's routed and shared experts as the checkpoint's declared
250    /// NVFP4, so `nvfp4_grouped_decode_ok` admits it under every `--expert-quantization` tier.
251    /// Refused for a layer that also holds FP8 experts: the tier decides between those.
252    pub fn set_declared_nvfp4_experts(&mut self) -> Result<()> {
253        self.set_draft_nvfp4_experts()
254    }
255
256    /// 2026-10-02: The same mark for a drafter whose experts were requantized to NVFP4 for drafting
257    /// (`MtpHead::new_nvfp4_draft_moe`): the grouped NVFP4 decode serves them at every width.
258    pub fn set_draft_nvfp4_experts(&mut self) -> Result<()> {
259        anyhow::ensure!(
260            self.fp8_shared_expert.is_none() && self.fp8_down_weight_ptrs.is_none(),
261            "set_declared_nvfp4_experts: the layer holds FP8 experts"
262        );
263        self.nvfp4_grouped.declared_experts = true;
264        Ok(())
265    }
266
267    /// 2026-09-27: The FP8 shared expert this path runs instead of the NVFP4 one, when the
268    /// layer keeps its FP8 experts (`set_fp8_experts`).
269    fn nvfp4_decode_fp8_shared(&self) -> Option<&Fp8ExpertWeight> {
270        self.fp8_shared_expert.as_ref()
271    }
272
273    /// 2026-09-27: The NVFP4 MoE of `m` rows: `input` is `[m, H]` BF16, the output lands in
274    /// rows `0..m` of `moe_output()`. Returns an error, launching nothing, when
275    /// `nvfp4_grouped_decode_ok(m, ctx)` is false.
276    pub fn forward_nvfp4_grouped_decode(
277        &self,
278        input: DevicePtr,
279        m: usize,
280        ctx: &ForwardContext,
281        stream: u64,
282    ) -> Result<()> {
283        anyhow::ensure!(
284            self.nvfp4_grouped_decode_ok(m, ctx),
285            "forward_nvfp4_grouped_decode: predicate false for m={m} (caller must gate on it)"
286        );
287        let h = ctx.config.hidden_size as u32;
288        let inter = ctx.config.moe_intermediate_size as u32;
289        let num_experts = ctx.config.num_experts as u32;
290        let top_k = ctx.config.num_experts_per_tok as u32;
291        let n = m as u32;
292        let te = m * top_k as usize;
293        if ctx.stats.once("log:moe_nvfp4_grouped_decode") {
294            tracing::info!(
295                "MoE NVFP4 grouped decode active (first use M={m}, top_k={top_k}, \
296                 experts={num_experts}; one-time log)"
297            );
298        }
299
300        let scratch = ctx.buffers.scratch();
301        let indices_dev = scratch;
302        let weights_dev = scratch.offset(te * 4);
303        self.grouped_route(
304            input,
305            m,
306            GroupedRouting::PerRow,
307            indices_dev,
308            weights_dev,
309            ctx,
310            stream,
311        )?;
312
313        // 2026-09-27: The sort scratch reuses `gate_logits`, which top-k has read earlier on
314        // this stream (the layout of `grouped_decode_buffer_need`).
315        let ne = num_experts as usize;
316        let gate_logits = ctx.buffers.gate_logits();
317        let sorted_token_ids = gate_logits;
318        let sorted_expert_ids = gate_logits.offset(te * 4);
319        let expert_offsets = gate_logits.offset(te * 4 * 2);
320        let token_to_perm = gate_logits.offset(te * 4 * 2 + (ne + 1) * 4);
321        let cap = ops::fp8_grouped_active_cap(n, top_k, num_experts);
322        let active_experts = token_to_perm.offset(te * 4);
323        let active_count = active_experts.offset(cap as usize * 4);
324        ops::moe_fp8_grouped_sort(
325            ctx.gpu,
326            self.moe_fp8_grouped_sort_k,
327            ops::Fp8GroupedSortOut {
328                sorted_token_ids,
329                sorted_expert_ids,
330                expert_offsets,
331                token_to_perm,
332                active_experts,
333                active_count,
334            },
335            indices_dev,
336            te as u32,
337            num_experts,
338            top_k,
339            stream,
340        )?;
341
342        let act = ctx.buffers.expert_gate_out();
343        let expert_down_out = ctx.buffers.expert_down_out();
344        let shared_act = ctx.buffers.logits();
345        let shared_out = ctx.buffers.attn_output();
346        let tables = |t: &ExpertPtrTable| ops::Nvfp4ExpertTables {
347            packed_ptrs: t.packed_ptrs,
348            scale_ptrs: t.scale_ptrs,
349            scale2_vals: t.scale2_vals,
350        };
351        let sh = &self.weights.shared_expert;
352        // 2026-09-27: With the layer's FP8 experts kept, the grouped FP8 gate+up kernel runs the
353        // FP8 shared expert alone (no active experts) and the NVFP4 one only the routed experts
354        // (no shared rows). The down projections: under `nvfp4` the grouped FP8 down kernel runs
355        // the shared expert alone and the NVFP4 one the routed experts; under `nvfp4-gate-up`
356        // the grouped FP8 down kernel runs all of them from the FP8 experts.
357        let fp8_shared = self.nvfp4_decode_fp8_shared();
358        let fp8_down = if crate::layers::expert_quantization().nvfp4_down() {
359            None
360        } else {
361            self.fp8_down_weight_ptrs.as_ref().zip(fp8_shared)
362        };
363        if let Some(fsh) = fp8_shared {
364            ops::moe_expert_gate_up_act_fp8_grouped(
365                ctx.gpu,
366                self.moe_expert_gate_up_act_fp8_grouped_k,
367                ops::FP8_GROUPED_GATE_UP_SCALAR,
368                input,
369                DevicePtr::NULL,
370                DevicePtr::NULL,
371                DevicePtr::NULL,
372                DevicePtr::NULL,
373                act,
374                expert_offsets,
375                sorted_token_ids,
376                active_experts,
377                active_count,
378                &fsh.gate_proj,
379                &fsh.up_proj,
380                shared_act,
381                inter,
382                h,
383                0,
384                n,
385                stream,
386            )?;
387        }
388        if let (Some(fsh), None) = (fp8_shared, fp8_down) {
389            ops::moe_expert_down_act_fp8_grouped(
390                ctx.gpu,
391                self.moe_expert_down_act_fp8_grouped_k,
392                ops::FP8_GROUPED_DOWN_SCALAR,
393                act,
394                DevicePtr::NULL,
395                DevicePtr::NULL,
396                expert_down_out,
397                expert_offsets,
398                active_experts,
399                active_count,
400                shared_act,
401                &fsh.down_proj,
402                shared_out,
403                h,
404                inter,
405                0,
406                n,
407                stream,
408            )?;
409        }
410        let nvfp4_shared_rows = if fp8_shared.is_some() { 0 } else { n };
411        let launch = self.nvfp4_grouped.select(h, inter);
412        ops::moe_expert_gate_up_act_nvfp4_grouped(
413            ctx.gpu,
414            launch.gate_up,
415            launch.gate_up_geometry,
416            input,
417            tables(&self.gate_ptrs),
418            tables(&self.up_ptrs),
419            act,
420            expert_offsets,
421            sorted_token_ids,
422            active_experts,
423            active_count,
424            &sh.gate_proj,
425            &sh.up_proj,
426            shared_act,
427            inter,
428            h,
429            cap,
430            nvfp4_shared_rows,
431            stream,
432        )?;
433        if let Some((dp, fsh)) = fp8_down {
434            ops::moe_expert_down_act_fp8_grouped(
435                ctx.gpu,
436                self.moe_expert_down_act_fp8_grouped_k,
437                ops::FP8_GROUPED_DOWN_SCALAR,
438                act,
439                dp.weight_ptrs,
440                dp.scale_ptrs,
441                expert_down_out,
442                expert_offsets,
443                active_experts,
444                active_count,
445                shared_act,
446                &fsh.down_proj,
447                shared_out,
448                h,
449                inter,
450                cap,
451                n,
452                stream,
453            )?;
454        } else {
455            ops::moe_expert_down_act_nvfp4_grouped(
456                ctx.gpu,
457                launch.down,
458                launch.down_geometry,
459                act,
460                tables(&self.down_ptrs),
461                expert_down_out,
462                expert_offsets,
463                active_experts,
464                active_count,
465                shared_act,
466                &sh.down_proj,
467                shared_out,
468                h,
469                inter,
470                cap,
471                nvfp4_shared_rows,
472                stream,
473            )?;
474        }
475        ops::moe_weighted_sum_blend_fp8_grouped(
476            ctx.gpu,
477            self.moe_weighted_sum_blend_fp8_grouped_k,
478            ctx.buffers.moe_output(),
479            expert_down_out,
480            weights_dev,
481            token_to_perm,
482            shared_out,
483            input,
484            self.weights.shared_expert_gate.weight,
485            h,
486            top_k,
487            h,
488            n,
489            stream,
490        )
491    }
492}
493
494#[cfg(test)]
495#[path = "forward_nvfp4_grouped_decode_tests.rs"]
496mod tests;
metrale_model_layers/layers/moe/forward_nvfp4_grouped_decode.rs

metrale_model_layers/layers/moe/
forward_nvfp4_grouped_decode.rs
