metrale_model_layers/layers/
dense_ffn_decode_batch.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: Dense-FFN decode for small verify batches: `forward_k2`, `forward_k3`,
4//! `forward_km` and the packed-Q2 `forward_km_q2`.
5//!
6//! Owner: model-layers (dense FFN).
7//! Invariants:
8//! - On success the `[m, hidden]` output is in `ctx.buffers.moe_output()`.
9//! - A layer with a BF16 or FP8 overlay is handed to `forward_prefill`.
10
11use anyhow::Result;
12use metrale_gpu_runtime::gpu::{DevicePtr, KernelHandle};
13
14use super::{DenseFfnLayer, DenseFfnWeightsQ2, FfnActivation, native_small_batch_uses_prefill};
15use crate::layer::ForwardContext;
16use crate::layers::ops;
17use crate::weight_map::PackedQ2Weight;
18
19impl DenseFfnLayer {
20    /// 2026-09-25: Packed-Q2 FFN for `m` rows, used by `forward_k2` and `forward_k3`:
21    /// `q2_0_gemv_vec_batchm` for gate and up, `act_mul`, then down; SiLU only. Output
22    /// `[m, hidden]` in `moe_output`.
23    fn forward_km_q2(
24        &self,
25        q2w: &DenseFfnWeightsQ2,
26        input: DevicePtr,
27        ctx: &ForwardContext,
28        m: u32,
29        stream: u64,
30    ) -> Result<()> {
31        if self.q2_0_gemv_batchm_k.0 == 0 {
32            anyhow::bail!(
33                "q2_0_gemv_vec_batchm kernel missing in this target build — packed-Q2 \
34                 batched decode (METRALE_GGUF_NATIVE_Q2, C>=2) is unavailable"
35            );
36        }
37        if self.activation != FfnActivation::SiLU {
38            anyhow::bail!(
39                "packed-Q2 FFN batched decode supports SiLU only (got {:?})",
40                self.activation
41            );
42        }
43        let inter = ctx.config.intermediate_size as u32;
44        let gate_out = ctx.buffers.expert_gate_out();
45        let up_out = ctx.buffers.expert_up_out();
46        let batchm = |w: &PackedQ2Weight, inp: DevicePtr, out: DevicePtr| -> Result<()> {
47            ops::q2_0_gemv_vec_batchm(ctx.gpu, self.q2_0_gemv_batchm_k, inp, w, out, m, stream)
48        };
49        batchm(&q2w.gate_proj, input, gate_out)?;
50        batchm(&q2w.up_proj, input, up_out)?;
51        ops::silu_mul(
52            ctx.gpu,
53            self.act_mul,
54            gate_out,
55            up_out,
56            gate_out,
57            m * inter,
58            stream,
59        )?;
60        let output = ctx.buffers.moe_output();
61        batchm(&q2w.down_proj, gate_out, output)?;
62        Ok(())
63    }
64
65    /// 2026-09-25: FFN for 2 rows. Packed-Q2 goes to `forward_km_q2`, FP8/BF16 to
66    /// `forward_prefill`, and to `forward_km` when `small_batch_w4a4` and `can_forward_km(2)`; otherwise NVFP4 runs `w4a16_gemv_dual_batch2`, `act_mul` and `w4a16_gemv_batch2`, plus the
67    /// LoRA deltas.
68    pub fn forward_k2(&self, input: DevicePtr, ctx: &ForwardContext, stream: u64) -> Result<()> {
69        if self.forward_w8a8(input, 2, ctx, stream)? {
70            return Ok(());
71        }
72        if let Some(ref q2w) = self.q2_weights {
73            return self.forward_km_q2(q2w, input, ctx, 2, stream);
74        }
75        if native_small_batch_uses_prefill(self.bf16_weights.is_some(), self.fp8_weights.is_some())
76        {
77            return self.forward_prefill(input, 2, ctx, stream);
78        }
79        // 2026-09-28: The batch2 kernels below are W4A16; `forward_km` routes every projection
80        // through the W4A4 launcher, which follows the weight-quantization tier.
81        if self.small_batch_w4a4() && self.can_forward_km(2) {
82            return self.forward_km(input, 2, ctx, stream);
83        }
84
85        let h = ctx.config.hidden_size as u32;
86        let inter = ctx.config.intermediate_size as u32;
87
88        let gate_out = ctx.buffers.expert_gate_out();
89        let up_out = ctx.buffers.expert_up_out();
90
91        ops::w4a16_gemv_dual_batch2(
92            ctx.gpu,
93            self.w4a16_gemv_dual_batch2,
94            input,
95            &self.weights.gate_proj,
96            gate_out,
97            &self.weights.up_proj,
98            up_out,
99            inter,
100            h,
101            stream,
102        )?;
103        self.apply_lora_gate_up(ctx, input, gate_out, up_out, 2, stream)?;
104        ops::silu_mul(
105            ctx.gpu,
106            self.act_mul,
107            gate_out,
108            up_out,
109            gate_out,
110            2 * inter,
111            stream,
112        )?;
113        let output = ctx.buffers.moe_output();
114        ops::w4a16_gemv_batch2(
115            ctx.gpu,
116            self.w4a16_gemv_batch2,
117            gate_out,
118            &self.weights.down_proj,
119            output,
120            h,
121            inter,
122            stream,
123        )?;
124        self.apply_lora_down(ctx, gate_out, output, 2, stream)?;
125
126        Ok(())
127    }
128
129    /// 2026-09-25: FFN for 3 rows; the same routing as `forward_k2`, with the batch3 kernels.
130    pub fn forward_k3(&self, input: DevicePtr, ctx: &ForwardContext, stream: u64) -> Result<()> {
131        if self.forward_w8a8(input, 3, ctx, stream)? {
132            return Ok(());
133        }
134        if let Some(ref q2w) = self.q2_weights {
135            return self.forward_km_q2(q2w, input, ctx, 3, stream);
136        }
137        if native_small_batch_uses_prefill(self.bf16_weights.is_some(), self.fp8_weights.is_some())
138        {
139            return self.forward_prefill(input, 3, ctx, stream);
140        }
141        // 2026-09-28: As in `forward_k2`.
142        if self.small_batch_w4a4() && self.can_forward_km(3) {
143            return self.forward_km(input, 3, ctx, stream);
144        }
145
146        let h = ctx.config.hidden_size as u32;
147        let inter = ctx.config.intermediate_size as u32;
148
149        let gate_out = ctx.buffers.expert_gate_out();
150        let up_out = ctx.buffers.expert_up_out();
151
152        ops::w4a16_gemv_dual_batch3(
153            ctx.gpu,
154            self.w4a16_gemv_dual_batch3,
155            input,
156            &self.weights.gate_proj,
157            gate_out,
158            &self.weights.up_proj,
159            up_out,
160            inter,
161            h,
162            stream,
163        )?;
164        self.apply_lora_gate_up(ctx, input, gate_out, up_out, 3, stream)?;
165        ops::silu_mul(
166            ctx.gpu,
167            self.act_mul,
168            gate_out,
169            up_out,
170            gate_out,
171            3 * inter,
172            stream,
173        )?;
174        let output = ctx.buffers.moe_output();
175        ops::w4a16_gemv_batch3(
176            ctx.gpu,
177            self.w4a16_gemv_batch3,
178            gate_out,
179            &self.weights.down_proj,
180            output,
181            h,
182            inter,
183            stream,
184        )?;
185        self.apply_lora_down(ctx, gate_out, output, 3, stream)?;
186
187        Ok(())
188    }
189
190    /// 2026-09-25: The `w4a16_gemv_batch{M}` handle for `m` rows, from
191    /// `W4a16BatchmTiers::kernel_for` over the gate weight (so within the W4A4 edge when the
192    /// checkpoint declares FP4 activations); zero when no tier serves `m`.
193    pub(super) fn batchm_kernel(&self, m: u32) -> KernelHandle {
194        self.w4a16_batchm.kernel_for(m, &self.weights.gate_proj)
195    }
196
197    /// 2026-09-28: Any of gate/up/down declares FP4 activations and the W4A4 kernels are
198    /// present (`W4a16BatchmTiers::declares_a4`).
199    fn declares_a4(&self) -> bool {
200        let t = &self.w4a16_batchm;
201        t.declares_a4(&self.weights.gate_proj)
202            || t.declares_a4(&self.weights.up_proj)
203            || t.declares_a4(&self.weights.down_proj)
204    }
205
206    /// 2026-09-28: Whether the 2- and 3-row steps go through `forward_km`, whose projections
207    /// run W4A4 where the tier admits them (`WeightQuantTier::ffn_small_batch_w4a4`).
208    pub fn small_batch_w4a4(&self) -> bool {
209        crate::layers::weight_quantization().ffn_small_batch_w4a4(self.declares_a4())
210    }
211
212    /// 2026-09-28: Whether the single decode row goes through `forward_km`
213    /// (`WeightQuantTier::ffn_single_row_w4a4`).
214    pub fn single_row_w4a4(&self) -> bool {
215        crate::layers::weight_quantization().ffn_single_row_w4a4(self.declares_a4())
216    }
217
218    /// 2026-09-28: Row edge of the dense FFN's narrow arms (`W4a16BatchmTiers::ffn_edge`).
219    pub fn narrow_rows(&self) -> u32 {
220        self.w4a16_batchm.ffn_edge(&self.weights.gate_proj)
221    }
222
223    /// 2026-09-25: Whether `forward_km` can serve `m` rows: a batchm tier resolved, and the layer
224    /// has NVFP4 gate weights or an FP8 overlay.
225    pub fn can_forward_km(&self, m: u32) -> bool {
226        self.batchm_kernel(m).0 != 0
227            && (!self.weights.gate_proj.weight.is_null()
228                // 2026-09-25: An FP8 layer answers true: `forward_km` sends it to
229                // `forward_prefill`, and callers such as
230                // `qwen3_attention/trait_impl/multi_seq/ffn.rs` pick their verify branch by this.
231                || self.fp8_weights.is_some())
232    }
233
234    /// 2026-09-25: FFN for `m` verify rows. FP8/BF16 layers go to `forward_prefill`. NVFP4 runs
235    /// gate, up and down through `ops::w4a4_proj::nvfp4_proj_small_m` with the `batchm_kernel(m)`
236    /// tier (W4A4 where the weight-quantization tier admits it), plus `act_mul` and the LoRA deltas. There is no
237    /// packed-Q2 arm; `can_forward_km` is false for a layer without NVFP4 or FP8 weights.
238    pub fn forward_km(
239        &self,
240        input: DevicePtr,
241        m: u32,
242        ctx: &ForwardContext,
243        stream: u64,
244    ) -> Result<()> {
245        if self.forward_w8a8(input, m as usize, ctx, stream)? {
246            return Ok(());
247        }
248        if native_small_batch_uses_prefill(self.bf16_weights.is_some(), self.fp8_weights.is_some())
249        {
250            return self.forward_prefill(input, m as usize, ctx, stream);
251        }
252        let h = ctx.config.hidden_size as u32;
253        let inter = ctx.config.intermediate_size as u32;
254        let kh = self.batchm_kernel(m);
255
256        let gate_out = ctx.buffers.expert_gate_out();
257        let up_out = ctx.buffers.expert_up_out();
258
259        ops::w4a4_proj::nvfp4_proj_small_m(
260            ctx.gpu,
261            kh,
262            input,
263            &self.weights.gate_proj,
264            gate_out,
265            m,
266            inter,
267            h,
268            stream,
269        )?;
270        // 2026-09-25: Same `input` as gate: the W4A4 path reuses gate's quantisation.
271        ops::w4a4_proj::nvfp4_proj_small_m_same_input(
272            ctx.gpu,
273            kh,
274            input,
275            &self.weights.up_proj,
276            up_out,
277            m,
278            inter,
279            h,
280            stream,
281        )?;
282        self.apply_lora_gate_up(ctx, input, gate_out, up_out, m, stream)?;
283        ops::silu_mul(
284            ctx.gpu,
285            self.act_mul,
286            gate_out,
287            up_out,
288            gate_out,
289            m * inter,
290            stream,
291        )?;
292        let output = ctx.buffers.moe_output();
293        ops::w4a4_proj::nvfp4_proj_small_m(
294            ctx.gpu,
295            kh,
296            gate_out,
297            &self.weights.down_proj,
298            output,
299            m,
300            h,
301            inter,
302            stream,
303        )?;
304        self.apply_lora_down(ctx, gate_out, output, m, stream)?;
305
306        Ok(())
307    }
308}
metrale_model_layers/layers/dense_ffn_decode_batch.rs

metrale_model_layers/layers/
dense_ffn_decode_batch.rs
