metrale_model_layers/layers/moe/forward_fp8_grouped_router.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: The router step of the grouped FP8 MoE decode
4//! (`forward_fp8_grouped_decode_routed`) and its admission per routing.
5//!
6//! The grouped expert kernels reproduce the per-row FP8 expert and blend
7//! kernels bit for bit given the same routing
8//! (`model-arch/examples/fp8_moe_grouped_decode_microtest.rs`). The routing is
9//! what differed: the batched gate GEMM and `moe_topk_softmax_batched` may pick
10//! differently from the per-row router on a near-tie. The two exact routings
11//! repeat the router of the path they replace, so the whole MoE output is the
12//! same bytes:
13//! - `PerRow`: `MoeLayer::forward` once per row. `dense_gemv_bf16_batchm`
14//! (row-for-row `dense_gemv_bf16`, at most 16 rows a launch) and
15//! `moe_topk_softmax_rows` (the `moe_topk_softmax` body per row).
16//! - `PerToken`: `MoeLayer::forward_batched`. The router GEMM of
17//! `batched_gate_logits` (`router_gemm_bf16`, the bits of `dense_gemm_bf16`)
18//! over all rows, then `moe_topk_softmax_rows`.
19//!
20//! Owner: model-layers (MoE).
21//! Invariants: `grouped_route` launches nothing unless
22//! `fp8_grouped_routing_ok` holds for its `m` and `routing`.
23
24use super::*;
25
26/// 2026-09-26: Which router arithmetic `forward_fp8_grouped_decode_routed` runs.
27#[derive(Clone, Copy, Debug, PartialEq, Eq)]
28pub enum GroupedRouting {
29 /// 2026-09-27: The batched gate GEMM (NVFP4 or BF16) and the batched softmax or sigmoid
30 /// top-k: the MTP drafter and the `METRALE_FP8_MOE_GROUPED_DECODE` arm.
31 Batched,
32 /// 2026-09-27: The router of `MoeLayer::forward`, row by row.
33 PerRow,
34 /// 2026-09-27: The router of `MoeLayer::forward_batched` with BF16 logits.
35 PerToken,
36}
37
38/// 2026-09-27: Rows per block row of the per-row router GEMV
39/// (`dense_gemv_batchm_split`); at most `DENSE_GEMV_BATCHM_MAX_M`.
40const ROUTER_ROWS_PER_BLOCK: u32 = 4;
41
42/// 2026-09-29: Block rows of the per-row router GEMV for `n` rows: `ROUTER_ROWS_PER_BLOCK` rows
43/// each up to 16 rows (4 block rows, 256 CTAs), then as many rows per block row as keep 4 block
44/// rows, up to `DENSE_GEMV_BATCHM_MAX_M`. Every block row reads the whole 1 MB router, so at a
45/// 256-row verify 64 block rows read it 64 times (7.8 ms per step on dgx2). A row's result does
46/// not depend on how rows are grouped (`dense_gemv_bf16_batchm.cu`).
47fn router_block_rows(n: u32) -> u32 {
48 let per_block = n
49 .div_ceil(4)
50 .clamp(ROUTER_ROWS_PER_BLOCK, ops::DENSE_GEMV_BATCHM_MAX_M);
51 n.div_ceil(per_block)
52}
53
54impl MoeLayer {
55 /// 2026-09-26: Whether the grouped decode serves `m` rows with `routing`.
56 /// `Batched` is [`Self::fp8_grouped_decode_ok`]. The exact routings also need
57 /// a BF16 softmax router with no correction bias, pre-router norm or FP32
58 /// gate, a layer `forward` does not reroute (DFlash capture), and their
59 /// kernels.
60 pub fn fp8_grouped_routing_ok(
61 &self,
62 m: usize,
63 routing: GroupedRouting,
64 ctx: &ForwardContext,
65 ) -> bool {
66 if !self.fp8_grouped_decode_ok(m, ctx) {
67 return false;
68 }
69 let exact_router = self.gate_nvfp4.is_none()
70 && self.correction_bias_dev.is_none()
71 && self.weights.router_pre_norm.is_none()
72 && !ctx.levers.fp32_gate
73 && !(self.is_dflash_capture_layer && ctx.levers.frankenstein_decode_via_prefill)
74 && self.moe_topk_softmax_rows_k.0 != 0;
75 match routing {
76 GroupedRouting::Batched => true,
77 GroupedRouting::PerToken => exact_router,
78 GroupedRouting::PerRow => {
79 exact_router
80 && self.router_gemv_batchm_k.0 != 0
81 && ctx.config.hidden_size.is_multiple_of(8)
82 }
83 }
84 }
85
86 /// 2026-09-26: Router logits into `gate_logits()` and top-k into
87 /// `indices_dev` / `weights_dev` (`[m * top_k]` u32 and f32) for `m` rows of
88 /// `input`, with the arithmetic `routing` names.
89 #[allow(clippy::too_many_arguments)]
90 pub(super) fn grouped_route(
91 &self,
92 input: DevicePtr,
93 m: usize,
94 routing: GroupedRouting,
95 indices_dev: DevicePtr,
96 weights_dev: DevicePtr,
97 ctx: &ForwardContext,
98 stream: u64,
99 ) -> Result<()> {
100 let h = ctx.config.hidden_size as u32;
101 let num_experts = ctx.config.num_experts as u32;
102 let top_k = ctx.config.num_experts_per_tok as u32;
103 let n = m as u32;
104 let router_in = self.router_input(input, n, h, ctx, stream)?;
105 let gate_logits = ctx.buffers.gate_logits();
106 match routing {
107 GroupedRouting::PerRow => {
108 // 2026-09-27: One launch, ROUTER_ROWS_PER_BLOCK rows per block row: the
109 // router's 256 columns give only 64 blocks per row group.
110 ops::dense_gemv_batchm_split(
111 ctx.gpu,
112 self.router_gemv_batchm_k,
113 router_in,
114 &self.weights.gate,
115 gate_logits,
116 n,
117 router_block_rows(n),
118 num_experts,
119 h,
120 num_experts,
121 stream,
122 )?;
123 }
124 GroupedRouting::PerToken => {
125 self.router_gemm_bf16(router_in, gate_logits, n, num_experts, h, ctx, stream)?;
126 }
127 GroupedRouting::Batched => {
128 return self.grouped_route_batched(
129 router_in,
130 n,
131 indices_dev,
132 weights_dev,
133 ctx,
134 stream,
135 );
136 }
137 }
138 ops::moe_topk_softmax_batched(
139 ctx.gpu,
140 self.moe_topk_softmax_rows_k,
141 gate_logits,
142 indices_dev,
143 weights_dev,
144 num_experts,
145 top_k,
146 ctx.config.norm_topk_prob,
147 n,
148 stream,
149 )
150 }
151
152 /// 2026-09-25: The `Batched` router: `[n, H] x [H, E]` -> `gate_logits [n, E]`
153 /// (NVFP4 or BF16 gate), then the batched sigmoid (with a correction bias)
154 /// or softmax top-k.
155 fn grouped_route_batched(
156 &self,
157 router_in: DevicePtr,
158 n: u32,
159 indices_dev: DevicePtr,
160 weights_dev: DevicePtr,
161 ctx: &ForwardContext,
162 stream: u64,
163 ) -> Result<()> {
164 let h = ctx.config.hidden_size as u32;
165 let num_experts = ctx.config.num_experts as u32;
166 let top_k = ctx.config.num_experts_per_tok as u32;
167 let gate_logits = ctx.buffers.gate_logits();
168 if let Some(ref nvfp4) = self.gate_nvfp4 {
169 ops::w4a16_gemm(
170 ctx.gpu,
171 self.w4a16_gemm,
172 router_in,
173 nvfp4,
174 gate_logits,
175 n,
176 num_experts,
177 h,
178 stream,
179 )?;
180 } else {
181 self.router_gemm_bf16(router_in, gate_logits, n, num_experts, h, ctx, stream)?;
182 }
183 if let Some(bias) = self.correction_bias_dev {
184 ops::moe_topk_sigmoid_batched(
185 ctx.gpu,
186 self.moe_topk_sigmoid_batched_k,
187 gate_logits,
188 bias,
189 indices_dev,
190 weights_dev,
191 num_experts,
192 top_k,
193 ctx.config.norm_topk_prob,
194 ctx.config.routed_scaling_factor as f32,
195 n,
196 stream,
197 )
198 } else {
199 ops::moe_topk_softmax_batched(
200 ctx.gpu,
201 self.moe_topk_batched,
202 gate_logits,
203 indices_dev,
204 weights_dev,
205 num_experts,
206 top_k,
207 ctx.config.norm_topk_prob,
208 n,
209 stream,
210 )
211 }
212 }
213}
214
215#[cfg(test)]
216mod router_block_rows_tests {
217 use super::router_block_rows;
218
219 /// 2026-09-29: Up to 16 rows the grouping is the previous one (4 rows per block row); wider
220 /// batches keep 4 block rows until a block row holds 16, and no block row exceeds 16.
221 #[test]
222 fn narrow_batches_keep_four_rows_per_block_and_wide_ones_cap_at_sixteen() {
223 for n in 1..=16u32 {
224 assert_eq!(router_block_rows(n), n.div_ceil(4), "n={n}");
225 }
226 assert_eq!(router_block_rows(64), 4);
227 assert_eq!(router_block_rows(128), 8);
228 assert_eq!(router_block_rows(256), 16);
229 for n in 1..=512u32 {
230 let y = router_block_rows(n);
231 assert!(
232 n.div_ceil(y) <= 16,
233 "n={n}: {} rows per block row",
234 n.div_ceil(y)
235 );
236 }
237 }
238}