Module gateup_fused

Module gateup_fused 

Source
Expand description

2026-09-25: The fused FP8 gate+up GEMM (ffn_gateup_fused). 2026-09-25: The fused dense-FFN gate+up decode arm: one W8A8 block-scaled GEMM at N = 2 * inter, not two.

Owner: model-layers (dense FFN). Invariants:

  • The arm is taken only when gateup_fused_selected holds: lever on for a SiLU layer, the W8A8 arm selected for both gate and up, the fused weight installed, the strided SiLU handle loaded, m in 5..=GATEUP_FUSED_MAX_M, and an output buffer of at least fused_out_bytes(m, inter).

The loader (qwen35_dense.rs) builds one [2 * inter, K] E4M3 buffer and one FP32 scale grid, and re-points gate_proj and up_proj at views inside them, so every un-fused rung of w8_gemm! reads the same bytes. A fused output row is [gate | up]: gate in columns [0, inter), up in [inter, 2 * inter). ops::silu_mul_strided reads the two halves at a row stride of 2 * inter and writes the contiguous [m, inter] the down projection reads.

Output column j of the fused GEMM is the dot product of gate column j (j < inter) or up column j - inter, over the same K with the same scales. examples/native_fp8_ffn_gateup_fused_microtest.rs gates byte equality of both halves against the un-fused cuBLASLt GEMMs at M = 5, 8 and 16. The GEMM is w8a8_gemm: cuBLASLt when the dispatch’s cuBLASLt scope includes ffn and its other clauses pass, the in-tree fp8_gemm_t_blockscaled otherwise.

The lever is kernels/<hw>/HARDWARE.toml [defaults] ffn_gateup_fused: true on hopper, false on gb10, b200 and b300. silu_mul_strided.cu exists only in kernels/hopper/common, so on other targets the SiLU handle is 0 and the arm declines even when METRALE_FFN_GATEUP_FUSED turns the lever on.

Functions§

ffn_gateup_fused
2026-09-25: Whether the fused gate+up arm is on for this process: [defaults] ffn_gateup_fused, overridden by METRALE_FFN_GATEUP_FUSED (0, false, off or no turns it off; any other value turns it on).