metrale_model_layers/layers/
dense_ffn_fixed.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-30: The dense FFN under a fixed `--activation-quantization` for the `ffn` family: one
4//! activation format at every row count, so a row's output bits do not depend on how many rows
5//! share the launch. Decode sites (the one-row decode, the multi-sequence arms of both layer
6//! kinds, the batched verify's GDN arm) ask [`DenseFfnLayer::fixed_ok`] first.
7//!
8//! Owner: model-layers (dense FFN).
9//! Invariants:
10//! - `fixed_ok` is false for `adaptive` row counts, so those take today's arms unchanged.
11//! - `forward_fixed` launches only when `fixed_ok` holds, and writes `moe_output` rows `0..m`.
12
13use anyhow::Result;
14use metrale_config::{ActQuantFormat, ProjFamily};
15use metrale_gpu_runtime::gpu::DevicePtr;
16
17use super::DenseFfnLayer;
18use crate::layer::ForwardContext;
19use crate::layers::ops;
20
21/// 2026-09-30: How a layer runs a fixed format, if it can.
22#[derive(Clone, Copy, Debug, PartialEq, Eq)]
23enum Fixed {
24    /// 2026-09-30: The declared W8A8 layer (`dense_ffn_w8a8.rs`): per-row FP8 at every row count.
25    W8a8,
26    /// 2026-09-30: NVFP4 weights, NVFP4 activations (`ops::w4a4_proj::nvfp4_proj_mx`).
27    W4a4,
28}
29
30impl DenseFfnLayer {
31    fn fixed_route(
32        &self,
33        m: usize,
34        gpu: &dyn metrale_gpu_runtime::gpu::GpuBackend,
35    ) -> Option<Fixed> {
36        let format = crate::layers::fixed_act(ProjFamily::Ffn, m)?;
37        if self.lora.is_some() {
38            return None;
39        }
40        if self.w8a8.is_some() {
41            return matches!(format, ActQuantFormat::Declared | ActQuantFormat::Fp8)
42                .then_some(Fixed::W8a8);
43        }
44        let nvfp4 = !self.weights.gate_proj.weight.is_null()
45            && ops::w4a4_proj::max_rows(gpu) > 0
46            && self.activation == super::FfnActivation::SiLU;
47        let wants = match format {
48            ActQuantFormat::Nvfp4 => true,
49            ActQuantFormat::Declared => self.w4a16_batchm.declares_a4(&self.weights.gate_proj),
50            _ => false,
51        };
52        (nvfp4 && wants).then_some(Fixed::W4a4)
53    }
54
55    /// 2026-09-30: Whether `m` rows run a fixed format on this layer.
56    pub fn fixed_ok(&self, m: usize, ctx: &ForwardContext) -> bool {
57        self.fixed_route(m, ctx.gpu).is_some()
58    }
59
60    /// 2026-09-30: The FFN of `m` rows at `input` under the fixed format, into `moe_output`.
61    pub fn forward_fixed(
62        &self,
63        input: DevicePtr,
64        m: usize,
65        ctx: &ForwardContext,
66        stream: u64,
67    ) -> Result<()> {
68        match self.fixed_route(m, ctx.gpu) {
69            Some(Fixed::W8a8) => {
70                anyhow::ensure!(
71                    self.forward_w8a8(input, m, ctx, stream)?,
72                    "ffn: the declared W8A8 arm declined {m} rows"
73                );
74                Ok(())
75            }
76            Some(Fixed::W4a4) => {
77                let h = ctx.config.hidden_size as u32;
78                let inter = ctx.config.intermediate_size as u32;
79                let (gate, up) = (ctx.buffers.expert_gate_out(), ctx.buffers.expert_up_out());
80                let w = &self.weights;
81                ops::w4a4_proj::nvfp4_proj_mx(
82                    ctx.gpu,
83                    input,
84                    &w.gate_proj,
85                    gate,
86                    m as u32,
87                    inter,
88                    h,
89                    stream,
90                )?;
91                ops::w4a4_proj::nvfp4_proj_mx(
92                    ctx.gpu, input, &w.up_proj, up, m as u32, inter, h, stream,
93                )?;
94                ops::silu_mul(
95                    ctx.gpu,
96                    self.act_mul,
97                    gate,
98                    up,
99                    gate,
100                    m as u32 * inter,
101                    stream,
102                )?;
103                ops::w4a4_proj::nvfp4_proj_mx(
104                    ctx.gpu,
105                    gate,
106                    &w.down_proj,
107                    ctx.buffers.moe_output(),
108                    m as u32,
109                    h,
110                    inter,
111                    stream,
112                )
113            }
114            None => anyhow::bail!("ffn: no fixed format for {m} rows on this layer"),
115        }
116    }
117}
metrale_model_layers/layers/dense_ffn_fixed.rs

metrale_model_layers/layers/
dense_ffn_fixed.rs
