metrale_model_layers/layers/moe/helpers_c.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Expert precision setup (FP8 and BF16 expert tables, the BF16 shared
4//! expert, FP8 prefill copies), the BF16 shared-expert pass, the router GEMM for
5//! a BF16 gate, and the router pre-norm.
6//!
7//! Owner: model-layers (MoE).
8//! Invariants: none beyond the types.
9
10use super::*;
11
12impl MoeLayer {
13 /// 2026-09-25: Make FP8 copies of the NVFP4 router gate and of the NVFP4 shared
14 /// expert for the prefill FP8 GEMMs. Routed experts are untouched.
15 pub fn predequant_for_prefill(
16 &mut self,
17 gpu: &dyn GpuBackend,
18 config: &metrale_config::ModelConfig,
19 stream: u64,
20 ) -> Result<()> {
21 let h = config.hidden_size;
22 let shared_inter = config.shared_expert_intermediate_size;
23 let num_experts = config.num_experts;
24 let predequant_k = gpu.kernel("w4a16", "predequant_nvfp4_to_fp8")?;
25
26 if let Some(ref nvfp4) = self.gate_nvfp4 {
27 self.gate_fp8 =
28 Some(nvfp4.predequant_to_fp8(gpu, predequant_k, num_experts, h, stream)?);
29 }
30
31 // 2026-09-25: An installed BF16 shared expert is used as is; no FP8 copy is
32 // made of the NVFP4 placeholder beside it. 2026-10-02: Nor of a checkpoint's own NVFP4
33 // shared expert under `--weight-quantization declared` (`set_declared_nvfp4_experts`):
34 // `fp8_gemm_n128` casts the activations to E4M3 unscaled, below the declared W4A16, so
35 // its prefill runs `w4a16_gemm_n128` on the transposed copies instead.
36 if self.bf16_shared_expert.is_none()
37 && !self.nvfp4_grouped.declared_experts
38 && !self.weights.shared_expert.gate_proj.is_null()
39 && shared_inter > 0
40 {
41 self.shared_gate_fp8 = Some(self.weights.shared_expert.gate_proj.predequant_to_fp8(
42 gpu,
43 predequant_k,
44 shared_inter,
45 h,
46 stream,
47 )?);
48 self.shared_up_fp8 = Some(self.weights.shared_expert.up_proj.predequant_to_fp8(
49 gpu,
50 predequant_k,
51 shared_inter,
52 h,
53 stream,
54 )?);
55 self.shared_down_fp8 = Some(self.weights.shared_expert.down_proj.predequant_to_fp8(
56 gpu,
57 predequant_k,
58 h,
59 shared_inter,
60 stream,
61 )?);
62 }
63
64 Ok(())
65 }
66
67 /// 2026-09-25: Install FP8 routed experts: device pointer tables for gate, up
68 /// and down, indexed by expert id in the FP8 kernels, plus the FP8 shared
69 /// expert.
70 pub fn set_fp8_experts(
71 &mut self,
72 experts: &[Fp8ExpertWeight],
73 shared_expert: Fp8ExpertWeight,
74 gpu: &dyn GpuBackend,
75 ) -> Result<()> {
76 self.fp8_gate_weight_ptrs = Some(build_fp8_ptr_table(experts, |e| &e.gate_proj, gpu)?);
77 self.fp8_up_weight_ptrs = Some(build_fp8_ptr_table(experts, |e| &e.up_proj, gpu)?);
78 self.fp8_down_weight_ptrs = Some(build_fp8_ptr_table(experts, |e| &e.down_proj, gpu)?);
79 self.fp8_shared_expert = Some(shared_expert);
80 Ok(())
81 }
82
83 /// 2026-09-25: Install BF16 routed experts (FP8 experts dequantized at load;
84 /// the qwen35 loader does this under METRALE_FP8_DEQUANT_MOE_TO_BF16=1).
85 ///
86 /// `shared_*` are the shared expert's BF16 gate/up/down pointers; all three
87 /// NULL means no shared expert.
88 pub fn set_bf16_experts(
89 &mut self,
90 gate_experts: &[crate::weight_map::DenseWeight],
91 up_experts: &[crate::weight_map::DenseWeight],
92 down_experts: &[crate::weight_map::DenseWeight],
93 shared_gate: DevicePtr,
94 shared_up: DevicePtr,
95 shared_down: DevicePtr,
96 gpu: &dyn GpuBackend,
97 ) -> Result<()> {
98 use super::build_bf16_ptr_table;
99 self.bf16_gate_weight_ptrs = Some(build_bf16_ptr_table(gate_experts, gpu)?);
100 self.bf16_up_weight_ptrs = Some(build_bf16_ptr_table(up_experts, gpu)?);
101 self.bf16_down_weight_ptrs = Some(build_bf16_ptr_table(down_experts, gpu)?);
102 if shared_gate.is_null() && shared_up.is_null() && shared_down.is_null() {
103 self.bf16_shared_expert = None;
104 } else {
105 self.set_bf16_shared_expert(
106 DenseWeight {
107 weight: shared_gate,
108 },
109 DenseWeight { weight: shared_up },
110 DenseWeight {
111 weight: shared_down,
112 },
113 )?;
114 }
115 Ok(())
116 }
117
118 /// 2026-09-25: Install a BF16 shared expert, whatever the routed experts'
119 /// precision.
120 pub fn set_bf16_shared_expert(
121 &mut self,
122 gate_proj: DenseWeight,
123 up_proj: DenseWeight,
124 down_proj: DenseWeight,
125 ) -> Result<()> {
126 self.bf16_shared_expert = Some(Bf16SharedExpert::new(gate_proj, up_proj, down_proj)?);
127 Ok(())
128 }
129
130 /// 2026-09-25: True when a BF16 shared expert sits beside non-BF16 routed
131 /// experts, so the quantized fused kernels cannot compute the shared term.
132 pub(super) fn has_mixed_bf16_shared_expert(&self) -> bool {
133 self.bf16_shared_expert.is_some() && self.bf16_gate_weight_ptrs.is_none()
134 }
135
136 /// 2026-09-25: Run the installed BF16 shared expert (gate, up, `moe_act_mul`,
137 /// down) for `num_tokens` rows into `down_out`. Callers pass the scratch
138 /// buffers because the free aliases differ between decode and prefill.
139 ///
140 /// Returns `Ok(false)` without launching when no BF16 shared expert is
141 /// installed. Errors if `num_tokens` or `shared_intermediate` is 0.
142 #[allow(clippy::too_many_arguments)]
143 pub(super) fn run_bf16_shared_expert(
144 &self,
145 input: DevicePtr,
146 num_tokens: u32,
147 hidden_size: u32,
148 shared_intermediate: u32,
149 gate_out: DevicePtr,
150 up_out: DevicePtr,
151 down_out: DevicePtr,
152 ctx: &ForwardContext,
153 stream: u64,
154 ) -> Result<bool> {
155 let Some(shared) = self.bf16_shared_expert else {
156 return Ok(false);
157 };
158 anyhow::ensure!(
159 num_tokens > 0 && shared_intermediate > 0,
160 "BF16 shared expert requires non-zero token and intermediate dimensions"
161 );
162
163 let project = |activation: DevicePtr,
164 weight: &DenseWeight,
165 output: DevicePtr,
166 n: u32,
167 k: u32|
168 -> Result<()> {
169 if num_tokens == 1 {
170 ops::dense_gemv(
171 ctx.gpu,
172 self.dense_gemv,
173 activation,
174 weight,
175 output,
176 n,
177 k,
178 stream,
179 )
180 } else if ctx.dispatch.cublas.ffn {
181 ops::cublas_bf16_proj_dense(
182 activation,
183 weight.weight,
184 output,
185 num_tokens,
186 n,
187 k,
188 stream,
189 )
190 } else {
191 ops::dense_gemm_prefill(
192 ctx.gpu,
193 self.dense_gemm,
194 self.dense_gemm_pipelined,
195 activation,
196 weight,
197 output,
198 num_tokens,
199 n,
200 k,
201 stream,
202 )
203 }
204 };
205
206 project(
207 input,
208 &shared.gate_proj,
209 gate_out,
210 shared_intermediate,
211 hidden_size,
212 )?;
213 project(
214 input,
215 &shared.up_proj,
216 up_out,
217 shared_intermediate,
218 hidden_size,
219 )?;
220 ops::silu_mul(
221 ctx.gpu,
222 self.moe_act_mul,
223 gate_out,
224 up_out,
225 gate_out,
226 num_tokens * shared_intermediate,
227 stream,
228 )?;
229 project(
230 gate_out,
231 &shared.down_proj,
232 down_out,
233 hidden_size,
234 shared_intermediate,
235 )?;
236 Ok(true)
237 }
238
239 /// 2026-09-25: Router GEMM for a BF16 gate weight:
240 /// `gate_logits[num_tokens, num_experts] = router_in @ gate^T`.
241 ///
242 /// It keeps the scalar kernel's accumulation order: `dense_gemm_bf16_router`
243 /// when the target has it (same per-output order as `dense_gemm_bf16`, see
244 /// dense_gemm_bf16.cu), else `dense_gemm_bf16` itself. Router logits select
245 /// experts after a BF16 store, so an order change can flip near-tied
246 /// selections. Measured 2026-08-12: routing this GEMM to
247 /// `dense_gemm_bf16_pipelined` moved BFCL on the FP8 MoE flagship from 86.55
248 /// to 84.76.
249 #[allow(clippy::too_many_arguments)]
250 pub(super) fn router_gate_gemm_dense(
251 &self,
252 router_in: DevicePtr,
253 gate_logits: DevicePtr,
254 num_tokens: u32,
255 num_experts: u32,
256 hidden_size: u32,
257 ctx: &ForwardContext,
258 stream: u64,
259 ) -> Result<()> {
260 // 2026-10-01: Up to 64 rows (a warm request's tail replay under
261 // `prefill_row_invariant`), `router_gemm_bf16`: `moe_router_gemm_bf16` has the same
262 // bits (dense_gemm_bf16's ascending-k FP32 chain, --fmad=false) and is ~5x faster
263 // there (GB10, 27 rows: 23 vs 110 us per layer).
264 if num_tokens <= ops::MOE_ROUTER_GEMM_MAX_PREFILL_ROWS
265 && self.moe_router_gemm_k.0 != 0
266 && hidden_size.is_multiple_of(16)
267 {
268 return self.router_gemm_bf16(
269 router_in,
270 gate_logits,
271 num_tokens,
272 num_experts,
273 hidden_size,
274 ctx,
275 stream,
276 );
277 }
278 if self.moe_router_rt_k.0 != 0
279 && num_tokens >= ops::MOE_ROUTER_RT_MIN_ROWS
280 && hidden_size.is_multiple_of(16)
281 {
282 return ops::moe_router_gemm_rt(
283 ctx.gpu,
284 self.moe_router_rt_k,
285 router_in,
286 &self.weights.gate,
287 gate_logits,
288 num_tokens,
289 num_experts,
290 hidden_size,
291 stream,
292 );
293 }
294 if self.dense_gemm_router.0 != 0 {
295 return ops::dense_gemm_router(
296 ctx.gpu,
297 self.dense_gemm_router,
298 router_in,
299 &self.weights.gate,
300 gate_logits,
301 num_tokens,
302 num_experts,
303 hidden_size,
304 stream,
305 );
306 }
307 ops::dense_gemm(
308 ctx.gpu,
309 self.dense_gemm,
310 router_in,
311 &self.weights.gate,
312 gate_logits,
313 num_tokens,
314 num_experts,
315 hidden_size,
316 stream,
317 )
318 }
319
320 /// 2026-09-26: The BF16 router gate GEMM `[m, k] x [n, k]^T`, with the bits of
321 /// `dense_gemm_bf16`: `moe_router_gemm_bf16` when this target ships it and `k`
322 /// is a multiple of 16, else `dense_gemm_bf16` itself. The decode routers
323 /// (`forward_k2`, `forward_k3`, `batched_gate_logits`, the grouped decode) run
324 /// it at 2 to 64 rows, where `dense_gemm_bf16` waits on a global load per
325 /// 16-wide K tile across 16 blocks.
326 #[allow(clippy::too_many_arguments)]
327 pub(super) fn router_gemm_bf16(
328 &self,
329 router_in: DevicePtr,
330 gate_logits: DevicePtr,
331 m: u32,
332 n: u32,
333 k: u32,
334 ctx: &ForwardContext,
335 stream: u64,
336 ) -> Result<()> {
337 if self.moe_router_gemm_k.0 != 0 && k.is_multiple_of(16) {
338 return ops::moe_router_gemm(
339 ctx.gpu,
340 self.moe_router_gemm_k,
341 router_in,
342 &self.weights.gate,
343 gate_logits,
344 m,
345 n,
346 k,
347 stream,
348 );
349 }
350 ops::dense_gemm(
351 ctx.gpu,
352 self.dense_gemm,
353 router_in,
354 &self.weights.gate,
355 gate_logits,
356 m,
357 n,
358 k,
359 stream,
360 )
361 }
362
363 /// 2026-09-25: The router input. With `router_pre_norm` (the Gemma-4 weight map
364 /// stores `scale * hidden_size^-0.5` in it), rms-norm `input` with that
365 /// weight into `ctx.buffers.qkv_output()` and return it; without, return
366 /// `input`. The caller must not need `qkv_output` across the MoE block.
367 pub(super) fn router_input(
368 &self,
369 input: DevicePtr,
370 num_tokens: u32,
371 h: u32,
372 ctx: &ForwardContext,
373 stream: u64,
374 ) -> Result<DevicePtr> {
375 let Some(ref weight) = self.weights.router_pre_norm else {
376 return Ok(input);
377 };
378 let eps = ctx.config.rms_norm_eps as f32;
379 let normed = ctx.buffers.qkv_output();
380 ops::rms_norm(
381 ctx.gpu,
382 self.pre_expert_norm_k,
383 input,
384 weight,
385 normed,
386 num_tokens,
387 h,
388 eps,
389 stream,
390 )?;
391 Ok(normed)
392 }
393}