metrale_model_layers/layers/
dense_ffn_w8a8.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: The declared-W8A8 arm of the dense SiLU FFN (`crate::layers::W8a8Ffn`): the
4//! gate and up projections share one quantized input, and the down projection's input,
5//! `bf16(silu(gate) * up)`, is quantized in the same launch that computes it. Every entry
6//! point (`forward`, `forward_k2`, `forward_k3`, `forward_km`, `forward_prefill`) tries it
7//! first for `1..=ops::W8A8_MAX_ROWS` rows.
8//!
9//! Owner: model-layers (dense FFN).
10//! Invariants:
11//! - On `Ok(true)` the `[rows, hidden]` output is in `ctx.buffers.moe_output()`, as for the
12//!   other arms; `expert_gate_out` / `expert_up_out` hold the gate and up rows.
13//! - `Ok(false)` launches nothing: no W8A8 weights, a LoRA overlay, or a row count outside the
14//!   family.
15
16use anyhow::{Result, ensure};
17use metrale_gpu_runtime::gpu::DevicePtr;
18
19use super::{DenseFfnLayer, FfnActivation};
20use crate::layer::ForwardContext;
21use crate::layers::{W8a8Ffn, ops};
22
23impl DenseFfnLayer {
24    /// 2026-09-28: Install the W8A8 gate, up and down. Refuses a GELU FFN and shapes other than
25    /// `[inter, hidden]` (gate, up) and `[hidden, inter]` (down).
26    pub fn set_w8a8_decode_weights(&mut self, w: W8a8Ffn, hidden: u32, inter: u32) -> Result<()> {
27        ensure!(
28            self.activation == FfnActivation::SiLU,
29            "W8A8 FFN arm fuses SiLU only"
30        );
31        let shapes = [&w.gate, &w.up, &w.down].map(|p| (p.n(), p.k()));
32        ensure!(
33            shapes == [(inter, hidden), (inter, hidden), (hidden, inter)],
34            "W8A8 FFN shapes {shapes:?}, want gate/up [{inter}, {hidden}] and down [{hidden}, {inter}]"
35        );
36        self.w8a8 = Some(w);
37        Ok(())
38    }
39
40    /// 2026-09-28: The W8A8 FFN over `rows` rows of `input` (`[rows, hidden]` BF16).
41    pub(super) fn forward_w8a8(
42        &self,
43        input: DevicePtr,
44        rows: usize,
45        ctx: &ForwardContext,
46        stream: u64,
47    ) -> Result<bool> {
48        let Some(ref w) = self.w8a8 else {
49            return Ok(false);
50        };
51        if self.lora.is_some() || !w.ctx.available(&w.gate, rows) || !w.ctx.available(&w.down, rows)
52        {
53            return Ok(false);
54        }
55        let (h, inter) = (w.gate.k(), w.gate.n());
56        let gate_out = ctx.buffers.expert_gate_out();
57        let up_out = ctx.buffers.expert_up_out();
58        let output = ctx.buffers.moe_output();
59        w.ctx
60            .proj(ctx.gpu, &w.gate, input, h, rows, gate_out, inter, stream)?;
61        // 2026-09-28: `up` reads the activation `gate` quantized into the shared scratch.
62        ops::w8a8_gemv(
63            ctx.gpu,
64            &w.ctx.kernels,
65            &w.up,
66            &w.ctx.scratch,
67            rows,
68            up_out,
69            inter,
70            stream,
71        )?;
72        w.ctx.silu_proj(
73            ctx.gpu, &w.down, gate_out, up_out, inter, rows, output, h, stream,
74        )
75    }
76}
metrale_model_layers/layers/dense_ffn_w8a8.rs

metrale_model_layers/layers/
dense_ffn_w8a8.rs
