metrale_model_layers/layers/dense_ffn_prefill.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: Dense-FFN prefill: `forward_prefill`, its dispatch over the installed weights,
4//! the packed-Q2 branch, and the NVFP4 projection GEMM `w4a16_prefill_gemm`.
5//!
6//! Owner: model-layers (dense FFN).
7//! Invariants:
8//! - On success the `[num_tokens, hidden]` output is in `ctx.buffers.moe_output()`.
9//! - The weights are tried in the order packed Q2, FP8, BF16, NVFP4; the first installed
10//! one serves the call.
11
12use anyhow::Result;
13use metrale_gpu_runtime::gpu::DevicePtr;
14
15use super::{DenseFfnLayer, DenseFfnWeightsQ2, FfnActivation};
16use crate::layer::ForwardContext;
17use crate::layers::ops;
18use crate::weight_map::{DenseWeight, PackedQ2Weight, QuantizedWeight};
19
20impl DenseFfnLayer {
21 /// 2026-09-25: One NVFP4 projection GEMM for `m` rows. With a transposed copy `wt`, m <= 64,
22 /// k % 32 == 0 and `ModelLevers::ffn_small_m` (off with `METRALE_FFN_SMALLM=0`), it runs
23 /// the deep-K `w4a16_gemm_t_k64` when k >= `w4a16_k64_min_k()` and k % 64 == 0, else
24 /// `w4a16_gemm_t`. Otherwise, with `wt`, it runs `w4a16_gemm_t_m128_v2` or
25 /// `w4a16_gemm_t_m128`; without `wt` or those handles, the base `w4a16_gemm`.
26 #[allow(clippy::too_many_arguments)]
27 pub(super) fn w4a16_prefill_gemm(
28 &self,
29 ctx: &ForwardContext,
30 w: &QuantizedWeight,
31 wt: Option<&QuantizedWeight>,
32 input: DevicePtr,
33 output: DevicePtr,
34 m: u32,
35 n: u32,
36 k: u32,
37 stream: u64,
38 ) -> Result<()> {
39 if let Some(wt) = wt {
40 if m <= 64 && k.is_multiple_of(32) && ctx.levers.ffn_small_m {
41 if k >= crate::layers::w4a16_k64_min_k()
42 && k.is_multiple_of(64)
43 && self.w4a16_gemm_t_k64_k.0 != 0
44 {
45 return ops::w4a16_gemm_n128(
46 ctx.gpu,
47 self.w4a16_gemm_t_k64_k,
48 input,
49 wt,
50 output,
51 m,
52 n,
53 k,
54 stream,
55 );
56 }
57 if self.w4a16_gemm_t_k.0 != 0 {
58 return ops::w4a16_gemm_n128(
59 ctx.gpu,
60 self.w4a16_gemm_t_k,
61 input,
62 wt,
63 output,
64 m,
65 n,
66 k,
67 stream,
68 );
69 }
70 }
71 if self.w4a16_gemm_t_m128_v2_k.0 != 0 {
72 return ops::w4a16_gemm_n128_m128_v2(
73 ctx.gpu,
74 self.w4a16_gemm_t_m128_v2_k,
75 input,
76 wt,
77 output,
78 m,
79 n,
80 k,
81 stream,
82 );
83 }
84 if self.w4a16_gemm_t_m128_k.0 != 0 {
85 return ops::w4a16_gemm_n128_m128(
86 ctx.gpu,
87 self.w4a16_gemm_t_m128_k,
88 input,
89 wt,
90 output,
91 m,
92 n,
93 k,
94 stream,
95 );
96 }
97 }
98 ops::w4a16_gemm(ctx.gpu, self.w4a16_gemm, input, w, output, m, n, k, stream)
99 }
100
101 /// 2026-09-25: Dense-FFN prefill of `num_tokens` rows of `input`; the output lands in
102 /// `moe_output`. Under `ctx.profile` it synchronizes `stream` and logs
103 /// `FFN prefill [dense_total] N=<n>: <us>µs`.
104 pub fn forward_prefill(
105 &self,
106 input: DevicePtr,
107 num_tokens: usize,
108 ctx: &ForwardContext,
109 stream: u64,
110 ) -> Result<()> {
111 if !ctx.profile {
112 return self.forward_prefill_inner(input, num_tokens, ctx, stream);
113 }
114 let t0 = std::time::Instant::now();
115 let r = self.forward_prefill_inner(input, num_tokens, ctx, stream);
116 // 2026-09-25: Synchronize so the time covers the kernels, not just their launch.
117 ctx.gpu.synchronize(stream)?;
118 tracing::info!(target: "metrale_model_layers::layers::dense_ffn", " FFN prefill [dense_total] N={}: {}µs",
119 num_tokens,
120 t0.elapsed().as_micros()
121 );
122 r
123 }
124
125 fn forward_prefill_inner(
126 &self,
127 input: DevicePtr,
128 num_tokens: usize,
129 ctx: &ForwardContext,
130 stream: u64,
131 ) -> Result<()> {
132 // 2026-09-28: The declared-W8A8 arm for up to `ops::W8A8_MAX_ROWS` rows
133 // (`dense_ffn_w8a8.rs`): the wide decode batches that land here, and short prefill tails.
134 if self.forward_w8a8(input, num_tokens, ctx, stream)? {
135 return Ok(());
136 }
137 let h = ctx.config.hidden_size as u32;
138 let inter = ctx.config.intermediate_size as u32;
139 let m = num_tokens as u32;
140
141 let gate_out = ctx.buffers.expert_gate_out();
142 let up_out = ctx.buffers.expert_up_out();
143
144 // 2026-09-25: Packed-Q2 prefill, SiLU only: the Q2_0 MMQ GEMM when it can run, else
145 // dequantize each projection to BF16 and run the BF16 GEMM.
146 if let Some(ref q2w) = self.q2_weights {
147 return self.prefill_q2(q2w, input, ctx, m, h, inter, gate_out, up_out, stream);
148 }
149
150 // 2026-09-25: FP8. `w8_gemm!` takes, per projection, the first arm that matches, each
151 // only when its kernel resolved:
152 // 1. m <= 4: `w8a16_gemv_batch4`.
153 // 2. m 5..=32 with `m16_tc` on: tensor-core `w8a16_gemm_m16`, two launches above 16
154 // rows. It reorders the K reduction; its tests compare it with `w8a16_gemv` within
155 // `m16_tc::within_m16_tc_budget`, not bit for bit.
156 // 3. m 5..=32 with `METRALE_FFN_BATCH16=1`: `w8a16_gemv_batch16`.
157 // 4. The W8A8 block-scaled GEMM, when `prefill_w8a8_selected` held.
158 // 5. `w8a16_gemm_pipelined`, else the base `w8a16_gemm`.
159 // The transposed `w8a16_gemm_t_m128` arm never matches: no transposed FP8 copy is passed.
160 // Every `[defaults]` table in `kernels/*/HARDWARE.toml` sets `ffn_m16_tc = false`, so with
161 // no variables set, m 5..=32 reaches arm 4 or 5.
162 if let Some(ref fp8w) = self.fp8_weights {
163 return self.prefill_fp8(fp8w, input, ctx, m, h, inter, gate_out, up_out, stream);
164 }
165
166 if let Some(ref bf16w) = self.bf16_weights {
167 return self.prefill_bf16(bf16w, input, ctx, m, h, inter, gate_out, up_out, stream);
168 }
169
170 self.prefill_nvfp4(
171 input, num_tokens, ctx, m, h, inter, gate_out, up_out, stream,
172 )
173 }
174
175 /// 2026-09-26: The packed-Q2 branch of `forward_prefill_inner`, SiLU only: the Q2_0 MMQ
176 /// GEMM when it can run, else a BF16 dequant of each projection and the BF16 GEMM.
177 fn prefill_q2(
178 &self,
179 q2w: &DenseFfnWeightsQ2,
180 input: DevicePtr,
181 ctx: &ForwardContext,
182 m: u32,
183 h: u32,
184 inter: u32,
185 gate_out: DevicePtr,
186 up_out: DevicePtr,
187 stream: u64,
188 ) -> Result<()> {
189 if self.activation != FfnActivation::SiLU {
190 anyhow::bail!(
191 "packed-Q2 FFN prefill supports SiLU only (got {:?})",
192 self.activation
193 );
194 }
195
196 // 2026-09-25: Q2_0 MMQ, when `METRALE_GGUF_NATIVE_Q2_MMQ=1`, `q2_0_mmq_nc_k` and
197 // `q4k_quant_act_k` resolved and all three weights use 128-element groups. Gate and
198 // up share one q8_1 quantization of `input`; down quantizes the activation.
199 let q2_mmq = self.q2_0_mmq_nc_k.0 != 0
200 && self.q4k_quant_act_k.0 != 0
201 && ops::native_q2_mmq_enabled()
202 && q2w.gate_proj.group == 128
203 && q2w.up_proj.group == 128
204 && q2w.down_proj.group == 128;
205 if q2_mmq {
206 static Q2MMQ_LOG: std::sync::Once = std::sync::Once::new();
207 Q2MMQ_LOG.call_once(|| {
208 eprintln!(
209 "[metrale] METRALE_GGUF_NATIVE_Q2_MMQ=1: dense-FFN prefill via native packed Q2_0 MMQ (W2A8, keep-packed)"
210 );
211 });
212 let a_q8 = ctx.buffers.q2_act_q8();
213 let mmq = |w: &PackedQ2Weight, out: DevicePtr| -> Result<()> {
214 ops::q2_0_mmq_gemm(
215 ctx.gpu,
216 self.q2_0_mmq_nc_k,
217 self.q2_0_mmq_wc_k,
218 a_q8,
219 w.weight,
220 out,
221 m,
222 w.n,
223 w.k,
224 stream,
225 )
226 };
227 ops::quantize_act_q8_1(ctx.gpu, self.q4k_quant_act_k, input, a_q8, m, h, stream)?;
228 mmq(&q2w.gate_proj, gate_out)?;
229 mmq(&q2w.up_proj, up_out)?;
230 ops::silu_mul(
231 ctx.gpu,
232 self.act_mul,
233 gate_out,
234 up_out,
235 gate_out,
236 m * inter,
237 stream,
238 )?;
239 let output = ctx.buffers.moe_output();
240 ops::quantize_act_q8_1(
241 ctx.gpu,
242 self.q4k_quant_act_k,
243 gate_out,
244 a_q8,
245 m,
246 inter,
247 stream,
248 )?;
249 mmq(&q2w.down_proj, output)?;
250 return Ok(());
251 }
252
253 if self.dequant_q2_0_gn_k.0 == 0 {
254 anyhow::bail!(
255 "dequant_q2_0_gn_to_bf16 kernel missing in this target build — \
256 packed-Q2 (METRALE_GGUF_NATIVE_Q2) prefill is unavailable"
257 );
258 }
259 let tc = self.dense_gemm_tc_k.0 != 0;
260 // 2026-09-25: The three projections share the arena's `q2_dequant_scratch`. That is
261 // safe because each dequant and its GEMM run on `stream` before the next projection's
262 // dequant.
263 let scratch = ctx.buffers.q2_dequant_scratch();
264 let q2_gemm = |w: &PackedQ2Weight, input: DevicePtr, out: DevicePtr| -> Result<()> {
265 let (n, k) = (w.n, w.k);
266 debug_assert!(
267 (n as usize) * (k as usize) * 2 <= ctx.buffers.q2_dequant_scratch_bytes(),
268 "packed-Q2 FFN dequant scratch too small for [{n},{k}] BF16"
269 );
270 ops::dequant_q2_0_gn_to_bf16(
271 ctx.gpu,
272 self.dequant_q2_0_gn_k,
273 w.weight,
274 scratch,
275 n,
276 k,
277 w.group as u32,
278 stream,
279 )?;
280 let dw = DenseWeight { weight: scratch };
281 if tc {
282 ops::dense_gemm_tc(
283 ctx.gpu,
284 self.dense_gemm_tc_k,
285 input,
286 &dw,
287 out,
288 m,
289 n,
290 k,
291 stream,
292 )?;
293 } else {
294 ops::dense_gemm(
295 ctx.gpu,
296 self.dense_gemm_bf16_k,
297 input,
298 &dw,
299 out,
300 m,
301 n,
302 k,
303 stream,
304 )?;
305 }
306 Ok(())
307 };
308 q2_gemm(&q2w.gate_proj, input, gate_out)?;
309 q2_gemm(&q2w.up_proj, input, up_out)?;
310 ops::silu_mul(
311 ctx.gpu,
312 self.act_mul,
313 gate_out,
314 up_out,
315 gate_out,
316 m * inter,
317 stream,
318 )?;
319 let output = ctx.buffers.moe_output();
320 q2_gemm(&q2w.down_proj, gate_out, output)?;
321 Ok(())
322 }
323}