metrale_model_layers/layers/
dense_ffn_batch16_decode.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-26: The opt-in 5..=32-row native-FP8 dense-FFN tier, `w8a16_gemv_batch16`:
4//! one weight pass for up to 16 rows, and two launches on row halves for 17..=32.
5//!
6//! Owner: model-layers (dense FFN).
7//! Invariants:
8//! - `batch16_plan` returns `None` unless the tier is armed and the handle resolved, and
9//!   for every `m` outside 5..=32.
10//! - Both launches of a `Halves` plan cover at most 16 rows, the kernel's `MAX_M`.
11//!
12//! The `w8_gemm!` ladder in `DenseFfnLayer::prefill_fp8` (reached from `forward_prefill_inner`) tries this tier after
13//! `w8a16_gemv_batch4` (m <= 4) and the tensor-core `m16_tc` tier, and before the W8A8
14//! and tile GEMMs. It selects by row count, not by phase, so a prefill chunk of 5..=32
15//! tokens takes it too. `w8a16_gemv_batch16` and `w8a16_gemv_batch4` are two `MAX_M`
16//! instantiations of one template (`w8a16_gemv_batch4.cu`), and
17//! `examples/native_fp8_ffn_batch16_microtest.rs` requires every row to equal the scalar
18//! `w8a16_gemv` byte for byte.
19//!
20//! Off by default. Measured 2026-09-11 on 1xH100, Qwen/Qwen3.8-27B-FP8, same binary,
21//! 1024x256 at C=16, tier on against off: aggregate 121.4 against 128.0 tok/s, TPOT p50
22//! 107.4 against 102.0 ms, and TTFT of a 28-token prompt 150 against 101 ms.
23
24use anyhow::Result;
25use metrale_gpu_runtime::gpu::DevicePtr;
26
27use super::DenseFfnLayer;
28use crate::layer::ForwardContext;
29use crate::layers::ops;
30use crate::weight_map::Fp8Weight;
31
32/// 2026-09-25: Whether `METRALE_FFN_BATCH16` is exactly `1`, which arms the tier; any
33/// other value, or none, leaves it off. Read once per process; each `DenseFfnLayer`
34/// copies it into `batch16_enabled` at construction, so the route cannot change between
35/// CUDA-graph replays.
36pub fn ffn_batch16_enabled() -> bool {
37    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
38    *ON.get_or_init(|| std::env::var("METRALE_FFN_BATCH16").as_deref() == Ok("1"))
39}
40
41/// 2026-09-25: How the batch16 tier serves `m` rows.
42#[derive(Debug, Clone, Copy, PartialEq, Eq)]
43pub(crate) enum Batch16Plan {
44    /// 2026-09-25: One launch covering rows `0..m` (m <= 16).
45    Single,
46    /// 2026-09-25: Two launches on contiguous row halves: rows `0..first`, then
47    /// `first..m`. `first` is `ceil(m/2)`, so both halves are <= 16 for every m <= 32 and
48    /// the first half is the wider one (m=17 -> 9 + 8).
49    Halves { first: u32 },
50}
51
52/// 2026-09-25: The batch16 selection rule, pure over the row count, the handle's
53/// presence and the opt-in, so the CPU tests can cover it without a `ForwardContext`.
54/// `enabled` is a parameter because the process-wide `OnceLock` behind
55/// `ffn_batch16_enabled` cannot be toggled per test.
56pub(crate) fn batch16_plan(m: u32, batch16_loaded: bool, enabled: bool) -> Option<Batch16Plan> {
57    if !enabled || !batch16_loaded {
58        return None;
59    }
60    match m {
61        5..=16 => Some(Batch16Plan::Single),
62        // 2026-09-25: Both halves must be <= 16, the kernel's `MAX_M`; `div_ceil` puts
63        // the odd row in the first half.
64        17..=32 => Some(Batch16Plan::Halves {
65            first: m.div_ceil(2),
66        }),
67        _ => None,
68    }
69}
70
71impl DenseFfnLayer {
72    /// 2026-09-25: The plan for `m` rows on this layer: its handle and the opt-in it
73    /// copied at construction.
74    pub(crate) fn ffn_batch16_plan(&self, m: u32) -> Option<Batch16Plan> {
75        batch16_plan(m, self.w8a16_gemv_batch16_k.0 != 0, self.batch16_enabled)
76    }
77
78    /// 2026-09-25: Run one dense-FFN projection through `w8a16_gemv_batch16`.
79    ///
80    /// `input` is `[m, k]` BF16 and `out` is `[m, n]` BF16, both contiguous, so the
81    /// `Halves` plan is two launches at byte offsets. Rows that are not contiguous need
82    /// `ops::w8a16_gemv_batch16_strided`, which the attention QKV path uses.
83    #[allow(clippy::too_many_arguments)]
84    pub(crate) fn w8a16_batch16_proj(
85        &self,
86        ctx: &ForwardContext,
87        plan: Batch16Plan,
88        w: &Fp8Weight,
89        input: DevicePtr,
90        out: DevicePtr,
91        m: u32,
92        n: u32,
93        k: u32,
94        stream: u64,
95    ) -> Result<()> {
96        self.log_batch16_decode_route(ctx, plan);
97        const BF16: usize = 2;
98        let launch = |rows: u32, first: u32| {
99            ops::w8a16_gemv_batch16(
100                ctx.gpu,
101                self.w8a16_gemv_batch16_k,
102                input.offset(first as usize * k as usize * BF16),
103                w.weight,
104                w.row_scale,
105                out.offset(first as usize * n as usize * BF16),
106                rows,
107                n,
108                k,
109                stream,
110            )
111        };
112        match plan {
113            Batch16Plan::Single => launch(m, 0),
114            Batch16Plan::Halves { first } => {
115                launch(first, 0)?;
116                launch(m - first, first)
117            }
118        }
119    }
120
121    /// 2026-09-25: Logs the first batch16 launch once per `ModelStats`
122    /// (`log:ffn_batch16_decode`), so a serve's log shows that the opt-in tier ran.
123    fn log_batch16_decode_route(&self, ctx: &ForwardContext, plan: Batch16Plan) {
124        if ctx.stats.once("log:ffn_batch16_decode") {
125            let how = match plan {
126                Batch16Plan::Single => "one launch",
127                Batch16Plan::Halves { .. } => "two launches on contiguous row halves",
128            };
129            tracing::info!(
130                "[metrale] dense FFN decode: native FP8 w8a16_gemv_batch16 ({how}) \
131                 for 5..=32 rows — one weight pass, bit-identical per row to the \
132                 M=1 w8a16_gemv. ARMED BY METRALE_FFN_BATCH16=1, off by default: it \
133                 measured -5.4% aggregate and +50 ms TTFT on H100, and has never \
134                 been measured on GB10 (#927)."
135            );
136        }
137    }
138}
139
140#[cfg(test)]
141#[path = "dense_ffn_batch16_decode_tests.rs"]
142mod tests;
metrale_model_layers/layers/dense_ffn_batch16_decode.rs

metrale_model_layers/layers/
dense_ffn_batch16_decode.rs
