metrale_model_layers/layers/
dense_ffn_gateup_fused.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The fused dense-FFN gate+up decode arm: one W8A8 block-scaled GEMM at `N = 2 * inter`, not two.
4//!
5//! Owner: model-layers (dense FFN).
6//! Invariants:
7//! - The arm is taken only when `gateup_fused_selected` holds: lever on for a
8//!   SiLU layer, the W8A8 arm selected for both gate and up, the fused weight
9//!   installed, the strided SiLU handle loaded, `m` in 5..=`GATEUP_FUSED_MAX_M`,
10//!   and an output buffer of at least `fused_out_bytes(m, inter)`.
11//!
12//! The loader (`qwen35_dense.rs`) builds one `[2 * inter, K]` E4M3 buffer and
13//! one FP32 scale grid, and re-points `gate_proj` and `up_proj` at views inside
14//! them, so every un-fused rung of `w8_gemm!` reads the same bytes. A fused
15//! output row is `[gate | up]`: gate in columns `[0, inter)`, up in
16//! `[inter, 2 * inter)`. `ops::silu_mul_strided` reads the two halves at a row
17//! stride of `2 * inter` and writes the contiguous `[m, inter]` the down
18//! projection reads.
19//!
20//! Output column `j` of the fused GEMM is the dot product of gate column `j`
21//! (`j < inter`) or up column `j - inter`, over the same K with the same
22//! scales. `examples/native_fp8_ffn_gateup_fused_microtest.rs` gates byte
23//! equality of both halves against the un-fused cuBLASLt GEMMs at M = 5, 8 and
24//! 16. The GEMM is `w8a8_gemm`: cuBLASLt when the dispatch's cuBLASLt scope
25//! includes `ffn` and its other clauses pass, the in-tree
26//! `fp8_gemm_t_blockscaled` otherwise.
27//!
28//! The lever is `kernels/<hw>/HARDWARE.toml` `[defaults] ffn_gateup_fused`:
29//! true on hopper, false on gb10, b200 and b300. `silu_mul_strided.cu` exists
30//! only in `kernels/hopper/common`, so on other targets the SiLU handle is 0
31//! and the arm declines even when `METRALE_FFN_GATEUP_FUSED` turns the lever
32//! on.
33
34use anyhow::Result;
35use metrale_gpu_runtime::buffers::GATEUP_FUSED_MAX_M;
36use metrale_gpu_runtime::gpu::DevicePtr;
37
38use super::DenseFfnLayer;
39use crate::layer::ForwardContext;
40use crate::layers::ops;
41use crate::weight_map::Fp8Weight;
42
43/// 2026-09-25: Whether the fused gate+up arm is on for this process:
44/// `[defaults] ffn_gateup_fused`, overridden by `METRALE_FFN_GATEUP_FUSED`
45/// (`0`, `false`, `off` or `no` turns it off; any other value turns it on).
46pub fn ffn_gateup_fused() -> bool {
47    ops::target_defaults::resolved().ffn_gateup_fused.value
48}
49
50/// 2026-09-25: Bytes the fused `[ceil16(m), 2 * inter]` BF16 output occupies:
51/// `cublas_fp8_proj_prequant` hands cuBLASLt `ceil16(m)` rows and writes all
52/// of them.
53pub(crate) fn fused_out_bytes(m: u32, inter: u32) -> usize {
54    ops::cublas_fp8_m_pad(m) as usize * 2 * inter as usize * 2
55}
56
57/// 2026-09-25: The fused-arm selection rule, as a pure function, so the CPU
58/// tests can pin every clause without a `ForwardContext`. `lever` is injected
59/// because [`ffn_gateup_fused`] reads a process-global `OnceLock`.
60///
61/// Clauses:
62///
63/// * `lever`: the target's declaration with the environment override applied.
64/// * `gate_up_w8a8`: the W8A8 arm was selected for both gate and up. The fused
65///   GEMM is that arm at twice the N, so a pair the ladder would have put on a
66///   W8A16 rung must not be fused.
67/// * `5..=GATEUP_FUSED_MAX_M`: the decode band.
68/// * `fused_installed`: the loader built the `[2 * inter, K]` weight.
69/// * `silu_strided_loaded`: the strided SiLU consumer's handle is loaded.
70/// * `out_capacity_bytes`: the buffer holds `fused_out_bytes(m, inter)`, the
71///   padded extent. A short buffer declines instead of being written past its
72///   end.
73#[allow(clippy::too_many_arguments)]
74pub(crate) fn gateup_fused_selected(
75    m: u32,
76    inter: u32,
77    lever: bool,
78    gate_up_w8a8: bool,
79    fused_installed: bool,
80    silu_strided_loaded: bool,
81    out_capacity_bytes: usize,
82) -> bool {
83    lever
84        && gate_up_w8a8
85        && fused_installed
86        && silu_strided_loaded
87        && (5..=GATEUP_FUSED_MAX_M as u32).contains(&m)
88        && fused_out_bytes(m, inter) <= out_capacity_bytes
89}
90
91impl DenseFfnLayer {
92    /// 2026-09-26: Whether this layer fuses gate+up for `m` rows.
93    ///
94    /// `gate_up_w8a8` comes from the caller (`prefill_fp8`), which
95    /// already resolved the W8A8 rule for both projections; asking again here
96    /// would be a second copy of that rule.
97    pub(crate) fn gateup_fused_plan(
98        &self,
99        ctx: &ForwardContext,
100        m: u32,
101        inter: u32,
102        gate_up_w8a8: bool,
103    ) -> Option<&Fp8Weight> {
104        let fused = self.fp8_gate_up_fused.as_ref();
105        // 2026-09-25: SiLU only: the consumer this arm launches is the SiLU-mul,
106        // so any other activation keeps the `w8_gemm!` pair and `self.act_mul`.
107        let silu = self.activation == super::FfnActivation::SiLU;
108        gateup_fused_selected(
109            m,
110            inter,
111            self.gateup_fused && silu,
112            gate_up_w8a8,
113            fused.is_some(),
114            self.silu_mul_strided_k.0 != 0,
115            ctx.buffers.ffn_gate_up_fused_bytes(),
116        )
117        .then_some(fused)
118        .flatten()
119    }
120
121    /// 2026-09-25: gate+up in one GEMM into the arena's `ffn_gate_up_fused`, then
122    /// the strided SiLU-mul into the contiguous `[m, inter]` at `gate_out`.
123    /// `the_fused_arm_allocates_nothing_per_call` checks that the in-tree kernel
124    /// path allocates nothing.
125    #[allow(clippy::too_many_arguments)]
126    pub(crate) fn w8a8_gate_up_fused(
127        &self,
128        ctx: &ForwardContext,
129        a_fp8: DevicePtr,
130        a_scale: DevicePtr,
131        fused_w: &Fp8Weight,
132        gate_out: DevicePtr,
133        m: u32,
134        inter: u32,
135        h: u32,
136        stream: u64,
137    ) -> Result<()> {
138        self.log_gateup_fused_route(ctx, m);
139        let fused_out = ctx.buffers.ffn_gate_up_fused();
140        let cap = ctx.buffers.ffn_gate_up_fused_bytes();
141        debug_assert!(fused_out_bytes(m, inter) <= cap);
142        self.w8a8_gemm(
143            ctx,
144            a_fp8,
145            a_scale,
146            fused_w,
147            fused_out,
148            cap,
149            m,
150            2 * inter,
151            h,
152            stream,
153        )?;
154        // 2026-09-25: `up` starts `inter` BF16 elements into each fused row.
155        const BF16: usize = 2;
156        ops::silu_mul_strided(
157            ctx.gpu,
158            self.silu_mul_strided_k,
159            fused_out,
160            fused_out.offset(inter as usize * BF16),
161            gate_out,
162            m,
163            inter,
164            2 * inter,
165            inter,
166            stream,
167        )
168    }
169
170    /// 2026-09-25: Logs the route once per model (`log:ffn_gateup_fused`).
171    fn log_gateup_fused_route(&self, ctx: &ForwardContext, m: u32) {
172        if ctx.stats.once("log:ffn_gateup_fused") {
173            tracing::info!(
174                "[metrale] dense FFN decode: gate+up FUSED into ONE W8A8 \
175                 block-scaled GEMM at N=2*intermediate (m={m}, band 5..={max}) \
176                 — same weight bytes, one launch instead of two. Round 13 \
177                 priced the un-fused pair at 5 730.5 us/step, 59.4% of HBM, \
178                 against `down`'s 71.4% for the same bytes in one launch. \
179                 Bit-identical per element; METRALE_FFN_GATEUP_FUSED=0 restores \
180                 the two-GEMM arm (#927).",
181                max = GATEUP_FUSED_MAX_M,
182            );
183        }
184    }
185}
186
187#[cfg(test)]
188#[path = "dense_ffn_gateup_fused_tests.rs"]
189mod tests;
metrale_model_layers/layers/dense_ffn_gateup_fused.rs

metrale_model_layers/layers/
dense_ffn_gateup_fused.rs
