metrale_model_layers/layers/
dense_ffn_fp8_down.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Chooses the arm that runs the native-FP8 decode FFN for one token
4//! (`Fp8DownArm`), as a pure function `DenseFfnLayer::forward` calls.
5//!
6//! Owner: model-layers (dense FFN).
7//! Invariants:
8//! - `fp8_down_arm` returns `PerProjection` for a non-SiLU layer or a missing
9//!   `w8a16_gemv_dual`, and returns `SplitSilu` or `FusedSilu` only when every kernel
10//!   that arm launches resolved.
11//!
12//! Why `SplitSilu` is the preferred arm: `w8a16_gemv_silu_input` gives each of the 4
13//! outputs of a 256-thread CTA a 64-lane team that walks the whole K and evaluates
14//! `(g / (1 + __expf(-g))) * u` for every element, so one launch computes the SwiGLU once
15//! per output rather than once per element. `SplitSilu` computes it once with
16//! `moe_silu_mul` and then runs the plain `w8a16_gemv`. Measured 2026-09-11 on 1xH100,
17//! Qwen/Qwen3.8-27B-FP8, C=1: the fused down projection took 6.65 ms (30.4%) of a
18//! 21.891 ms decode step at 858 GB/s, against 1,979 GB/s for `w8a16_gemv_dual`.
19
20/// 2026-09-25: Which kernels the native-FP8 decode FFN runs for one token.
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub(crate) enum Fp8DownArm {
23    /// 2026-09-25: `w8a16_gemv_dual`, then `moe_silu_mul` writes `silu(gate)*up` over
24    /// `gate_out`, then the plain `w8a16_gemv` for down. The default:
25    /// `ModelLevers::decode_split_silu` comes from `[defaults] decode_split_silu`, true in
26    /// every `kernels/*/HARDWARE.toml` that declares it, and `METRALE_NO_DECODE_SPLIT_SILU`
27    /// turns it off.
28    ///
29    /// Not bit-identical to `FusedSilu`: `moe_silu_mul` rounds `g * (1 / (1 + e^-g)) * u`
30    /// to BF16 before the GEMV reads it, while the fused kernel keeps
31    /// `(g / (1 + e^-g)) * u` in FP32 inside the dot product. The non-fused FP8 prefill
32    /// path runs the same `moe_silu_mul`.
33    SplitSilu,
34    /// 2026-09-25: `w8a16_gemv_dual`, then the fused `w8a16_gemv_silu_input`. Taken
35    /// when the split-SiLU lever is off or `act_mul` / `w8a16_gemv` did not resolve.
36    FusedSilu,
37    /// 2026-09-25: Neither fused path is usable: four launches, `w8a16_gemv` for gate and
38    /// up, `act_mul`, then `w8a16_gemv` for down.
39    PerProjection,
40}
41
42/// 2026-09-25: The arm rule `DenseFfnLayer::forward` uses, pure so the CPU tests can
43/// cover every combination without a GPU.
44///
45/// `is_silu` is the layer's activation; `split_silu_lever` is
46/// `ModelLevers::decode_split_silu`; every other `bool` is "this handle is nonzero".
47pub(crate) fn fp8_down_arm(
48    is_silu: bool,
49    dual: bool,
50    fused_silu: bool,
51    act_mul: bool,
52    plain_gemv: bool,
53    split_silu_lever: bool,
54) -> Fp8DownArm {
55    // 2026-09-25: Both fused arms read the gate/up pair that `w8a16_gemv_dual` writes
56    // to `gate_out`/`up_out`, so without it neither can run.
57    if !is_silu || !dual {
58        return Fp8DownArm::PerProjection;
59    }
60    if split_silu_lever && act_mul && plain_gemv {
61        Fp8DownArm::SplitSilu
62    } else if fused_silu {
63        Fp8DownArm::FusedSilu
64    } else {
65        Fp8DownArm::PerProjection
66    }
67}
metrale_model_layers/layers/dense_ffn_fp8_down.rs

metrale_model_layers/layers/
dense_ffn_fp8_down.rs
