metrale_model_layers/layers/dense_ffn_decode.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: Single-token dense-FFN decode, `DenseFfnLayer::forward`.
4//!
5//! Owner: model-layers (dense FFN).
6//! Invariants:
7//! - On success the `[1, hidden]` output is in `ctx.buffers.moe_output()`, and `forward`
8//! returns that pointer.
9
10use anyhow::Result;
11use metrale_gpu_runtime::gpu::DevicePtr;
12
13use super::{DenseFfnLayer, FfnActivation, fp8_down};
14
15/// 2026-09-28: The FFN's circuit binding (`circuit_bind`); a child of this module so it reads
16/// the layer's private fields without growing `dense_ffn.rs`.
17#[path = "dense_ffn_circuit.rs"]
18mod circuit;
19use crate::layer::ForwardContext;
20use crate::layers::ops;
21use crate::weight_map::QuantizedWeight;
22
23impl DenseFfnLayer {
24 /// 2026-09-25: Decode for one token; returns `moe_output` holding `[1, hidden]`. Packed-Q2, FP8
25 /// and BF16 layers take their own branches. NVFP4 that `single_row_w4a4` admits runs
26 /// `forward_km` at one row (W4A4, 2026-09-28). Other NVFP4 runs the fused gate+up GEMV, then either
27 /// `act_mul` plus the down GEMV (split SiLU: the default, and always with a LoRA adapter or
28 /// GELU) or the SiLU-fused down GEMV. `METRALE_DECODE_FFN_VIA_GEMM=1` runs NVFP4 through
29 /// `w4a16_prefill_gemm` instead.
30 pub fn forward(
31 &self,
32 input: DevicePtr,
33 ctx: &ForwardContext,
34 stream: u64,
35 ) -> Result<DevicePtr> {
36 // 2026-09-30: A fixed `--activation-quantization` for the FFN (`dense_ffn_fixed.rs`).
37 if self.fixed_ok(1, ctx) {
38 self.forward_fixed(input, 1, ctx, stream)?;
39 return Ok(ctx.buffers.moe_output());
40 }
41 // 2026-09-28: The declared-W8A8 arm first (`dense_ffn_w8a8.rs`).
42 if self.forward_w8a8(input, 1, ctx, stream)? {
43 return Ok(ctx.buffers.moe_output());
44 }
45 let h = ctx.config.hidden_size as u32;
46 let inter = ctx.config.intermediate_size as u32;
47
48 let gate_out = ctx.buffers.expert_gate_out();
49 let up_out = ctx.buffers.expert_up_out();
50
51 // 2026-09-25: Packed-Q2: gate, up, `act_mul`, down (four launches); SiLU only.
52 if let Some(ref q2w) = self.q2_weights {
53 if self.q2_0_gemv_k.0 == 0 {
54 anyhow::bail!(
55 "q2_0_gemv kernel missing in this target build — packed-Q2 decode \
56 (METRALE_GGUF_NATIVE_Q2) is unavailable"
57 );
58 }
59 if self.activation != FfnActivation::SiLU {
60 anyhow::bail!(
61 "packed-Q2 FFN decode supports SiLU only (got {:?})",
62 self.activation
63 );
64 }
65 let output = ctx.buffers.moe_output();
66 ops::q2_0_gemv_vec(
67 ctx.gpu,
68 self.q2_0_gemv_k,
69 input,
70 &q2w.gate_proj,
71 gate_out,
72 stream,
73 )?;
74 ops::q2_0_gemv_vec(
75 ctx.gpu,
76 self.q2_0_gemv_k,
77 input,
78 &q2w.up_proj,
79 up_out,
80 stream,
81 )?;
82 ops::silu_mul(
83 ctx.gpu,
84 self.act_mul,
85 gate_out,
86 up_out,
87 gate_out,
88 inter,
89 stream,
90 )?;
91 ops::q2_0_gemv_vec(
92 ctx.gpu,
93 self.q2_0_gemv_k,
94 gate_out,
95 &q2w.down_proj,
96 output,
97 stream,
98 )?;
99 return Ok(output);
100 }
101
102 // 2026-09-25: FP8: `fp8_down::fp8_down_arm` picks the arm. The default `SplitSilu` stages
103 // `silu(gate)*up` once with `act_mul` before a plain `w8a16_gemv`, because the fused
104 // `w8a16_gemv_silu_input` recomputes it for every output. Numerics and the measurement
105 // are on `Fp8DownArm`.
106 if let Some(ref fp8w) = self.fp8_weights {
107 let output = ctx.buffers.moe_output();
108 let arm = fp8_down::fp8_down_arm(
109 self.activation == FfnActivation::SiLU,
110 self.w8a16_gemv_dual_k.0 != 0,
111 self.w8a16_gemv_silu_input_k.0 != 0,
112 self.act_mul.0 != 0,
113 self.w8a16_gemv_k.0 != 0,
114 ctx.levers.decode_split_silu,
115 );
116 if arm != fp8_down::Fp8DownArm::PerProjection {
117 ops::w8a16_gemv_dual(
118 ctx.gpu,
119 self.w8a16_gemv_dual_k,
120 input,
121 fp8w.gate_proj.weight,
122 fp8w.gate_proj.row_scale,
123 gate_out,
124 fp8w.up_proj.weight,
125 fp8w.up_proj.row_scale,
126 up_out,
127 inter,
128 h,
129 stream,
130 )?;
131 if arm == fp8_down::Fp8DownArm::SplitSilu {
132 ops::silu_mul(
133 ctx.gpu,
134 self.act_mul,
135 gate_out,
136 up_out,
137 gate_out,
138 inter,
139 stream,
140 )?;
141 ops::w8a16_gemv(
142 ctx.gpu,
143 self.w8a16_gemv_k,
144 gate_out,
145 fp8w.down_proj.weight,
146 fp8w.down_proj.row_scale,
147 output,
148 h,
149 inter,
150 stream,
151 )?;
152 } else {
153 ops::w8a16_gemv_silu_input(
154 ctx.gpu,
155 self.w8a16_gemv_silu_input_k,
156 gate_out,
157 up_out,
158 fp8w.down_proj.weight,
159 fp8w.down_proj.row_scale,
160 output,
161 h,
162 inter,
163 stream,
164 )?;
165 }
166 return Ok(output);
167 }
168 ops::w8a16_gemv(
169 ctx.gpu,
170 self.w8a16_gemv_k,
171 input,
172 fp8w.gate_proj.weight,
173 fp8w.gate_proj.row_scale,
174 gate_out,
175 inter,
176 h,
177 stream,
178 )?;
179 ops::w8a16_gemv(
180 ctx.gpu,
181 self.w8a16_gemv_k,
182 input,
183 fp8w.up_proj.weight,
184 fp8w.up_proj.row_scale,
185 up_out,
186 inter,
187 h,
188 stream,
189 )?;
190 ops::silu_mul(
191 ctx.gpu,
192 self.act_mul,
193 gate_out,
194 up_out,
195 gate_out,
196 inter,
197 stream,
198 )?;
199 ops::w8a16_gemv(
200 ctx.gpu,
201 self.w8a16_gemv_k,
202 gate_out,
203 fp8w.down_proj.weight,
204 fp8w.down_proj.row_scale,
205 output,
206 h,
207 inter,
208 stream,
209 )?;
210 return Ok(output);
211 }
212
213 if let Some(ref bf16w) = self.bf16_weights {
214 ops::dense_gemv(
215 ctx.gpu,
216 self.dense_gemv_bf16_k,
217 input,
218 &bf16w.gate_proj,
219 gate_out,
220 inter,
221 h,
222 stream,
223 )?;
224 ops::dense_gemv(
225 ctx.gpu,
226 self.dense_gemv_bf16_k,
227 input,
228 &bf16w.up_proj,
229 up_out,
230 inter,
231 h,
232 stream,
233 )?;
234 ops::silu_mul(
235 ctx.gpu,
236 self.act_mul,
237 gate_out,
238 up_out,
239 gate_out,
240 inter,
241 stream,
242 )?;
243 let output = ctx.buffers.moe_output();
244 ops::dense_gemv(
245 ctx.gpu,
246 self.dense_gemv_bf16_k,
247 gate_out,
248 &bf16w.down_proj,
249 output,
250 h,
251 inter,
252 stream,
253 )?;
254 return Ok(output);
255 }
256
257 // 2026-09-28: Under `--weight-quantization declared`, an FFN that declares FP4
258 // activations sends its single row through `forward_km`, whose projections run the
259 // W4A4 FP4 MMA as at 2..=32 rows. The GEMVs below are W4A16.
260 if self.single_row_w4a4()
261 && self.activation == FfnActivation::SiLU
262 && self.can_forward_km(1)
263 {
264 self.forward_km(input, 1, ctx, stream)?;
265 return Ok(ctx.buffers.moe_output());
266 }
267
268 // 2026-09-25: `METRALE_DECODE_FFN_VIA_GEMM=1` (SiLU layers): run the three NVFP4
269 // projections through `w4a16_prefill_gemm` at m=1 instead of the decode GEMVs. It needs the
270 // gate and up `_t` copies, which the `finalize_*_load` functions may have freed; without
271 // them it logs a warning and uses the GEMVs. This branch applies no LoRA delta.
272 if ctx.levers.decode_ffn_via_gemm
273 && self.activation == FfnActivation::SiLU
274 && self.act_mul.0 != 0
275 {
276 let wt_alive =
277 |w: &Option<QuantizedWeight>| w.as_ref().is_some_and(|w| !w.weight.is_null());
278 if wt_alive(&self.weights.gate_proj_t) && wt_alive(&self.weights.up_proj_t) {
279 if ctx.stats.once("log:decode_ffn_via_gemm") {
280 tracing::info!(target: "metrale_model_layers::layers::dense_ffn", "decode FFN via verify GEMM path (METRALE_DECODE_FFN_VIA_GEMM=1): \
281 gate/up/down through w4a16_prefill_gemm at M=1"
282 );
283 }
284 self.w4a16_prefill_gemm(
285 ctx,
286 &self.weights.gate_proj,
287 self.weights.gate_proj_t.as_ref(),
288 input,
289 gate_out,
290 1,
291 inter,
292 h,
293 stream,
294 )?;
295 self.w4a16_prefill_gemm(
296 ctx,
297 &self.weights.up_proj,
298 self.weights.up_proj_t.as_ref(),
299 input,
300 up_out,
301 1,
302 inter,
303 h,
304 stream,
305 )?;
306 ops::silu_mul(
307 ctx.gpu,
308 self.act_mul,
309 gate_out,
310 up_out,
311 gate_out,
312 inter,
313 stream,
314 )?;
315 let output = ctx.buffers.moe_output();
316 self.w4a16_prefill_gemm(
317 ctx,
318 &self.weights.down_proj,
319 self.weights.down_proj_t.as_ref(),
320 gate_out,
321 output,
322 1,
323 h,
324 inter,
325 stream,
326 )?;
327 return Ok(output);
328 }
329 if ctx.stats.once("log:decode_ffn_no_twins") {
330 tracing::warn!(target: "metrale_model_layers::layers::dense_ffn", "METRALE_DECODE_FFN_VIA_GEMM=1 requested but transposed FFN copies \
331 are freed/absent (NVFP4-MMQ prefill arm?) — falling back to GEMV; \
332 the unification experiment is NOT active"
333 );
334 }
335 }
336
337 // 2026-09-25: Fused gate+up GEMV. The single-warp choice is made separately for the dual
338 // and the silu-input kernels, so a missing `w4a16_gemv_silu_input_sw` does not disable
339 // `w4a16_gemv_dual_sw`.
340 let use_dual_sw = ops::use_gemv_sw(ctx.levers.gemv_sw, self.w4a16_gemv_dual_sw);
341 let use_silu_sw = ops::use_gemv_sw(ctx.levers.gemv_sw, self.w4a16_gemv_silu_input_sw);
342 if use_dual_sw {
343 ops::w4a16_gemv_dual_sw(
344 ctx.gpu,
345 self.w4a16_gemv_dual_sw,
346 input,
347 &self.weights.gate_proj,
348 gate_out,
349 &self.weights.up_proj,
350 up_out,
351 inter,
352 h,
353 stream,
354 )?;
355 } else {
356 ops::w4a16_gemv_dual(
357 ctx.gpu,
358 self.w4a16_gemv_dual,
359 input,
360 &self.weights.gate_proj,
361 gate_out,
362 &self.weights.up_proj,
363 up_out,
364 inter,
365 h,
366 stream,
367 )?;
368 }
369
370 let output = ctx.buffers.moe_output();
371 // 2026-09-25: Split SiLU (default; off with `METRALE_NO_DECODE_SPLIT_SILU`): stage
372 // `silu(gate)*up` once with `act_mul`, then run the down GEMV. The fused
373 // `w4a16_gemv_silu_input` recomputes the activation for every output row. An installed LoRA
374 // adapter forces this path, because the down delta needs the materialised activation;
375 // `set_lora_weights` refuses an adapter on a layer that cannot take it.
376 let split_silu = self.activation == FfnActivation::SiLU
377 && self.act_mul.0 != 0
378 && self.w4a16_gemv.0 != 0
379 && (ctx.levers.decode_split_silu || self.lora.is_some());
380 if split_silu {
381 self.apply_lora_gate_up(ctx, input, gate_out, up_out, 1, stream)?;
382 ops::silu_mul(
383 ctx.gpu,
384 self.act_mul,
385 gate_out,
386 up_out,
387 gate_out,
388 inter,
389 stream,
390 )?;
391 ops::w4a16_decode_gemv(
392 ctx.gpu,
393 self.w4a16_gemv,
394 self.w4a16_gemv_sw,
395 ctx.levers.gemv_sw,
396 gate_out,
397 &self.weights.down_proj,
398 output,
399 h,
400 inter,
401 stream,
402 )?;
403 self.apply_lora_down(ctx, gate_out, output, 1, stream)?;
404 return Ok(output);
405 }
406 debug_assert!(
407 self.lora.is_none(),
408 "LoRA installed but decode took the fused silu_input path, which \
409 never materialises the activation the down delta contracts over; \
410 set_lora_weights is supposed to make this unreachable"
411 );
412 match self.activation {
413 FfnActivation::SiLU => {
414 if use_silu_sw {
415 ops::w4a16_gemv_silu_input_sw(
416 ctx.gpu,
417 self.w4a16_gemv_silu_input_sw,
418 gate_out,
419 up_out,
420 &self.weights.down_proj,
421 output,
422 h,
423 inter,
424 stream,
425 )?;
426 } else {
427 ops::w4a16_gemv_silu_input(
428 ctx.gpu,
429 self.w4a16_gemv_silu_input,
430 gate_out,
431 up_out,
432 &self.weights.down_proj,
433 output,
434 h,
435 inter,
436 stream,
437 )?;
438 }
439 }
440 FfnActivation::GeLU => {
441 ops::silu_mul(
442 ctx.gpu,
443 self.act_mul,
444 gate_out,
445 up_out,
446 gate_out,
447 inter,
448 stream,
449 )?;
450 ops::w4a16_decode_gemv(
451 ctx.gpu,
452 self.w4a16_gemv,
453 self.w4a16_gemv_sw,
454 ctx.levers.gemv_sw,
455 gate_out,
456 &self.weights.down_proj,
457 output,
458 h,
459 inter,
460 stream,
461 )?;
462 }
463 }
464
465 Ok(output)
466 }
467}