Module gateup_fused

Module gateup_fused 

Source
Expand description

The FUSED gate+up DECODE GEMM (ffn_gateup_fused, #927) — a child module because it reads this layer’s private kernel handles, and the round-13 receipt, the layout decision and the residency-neutrality argument need room this file does not have. The FUSED dense-FFN gate+up DECODE GEMM — one block-scaled FP8 cuBLASLt call at N = 2 * intermediate in place of two at N = intermediate.

§WHY (#927)

nsys --cuda-graph-trace=node, 1xH100 80GB HBM3, Qwen/Qwen3.8-27B-FP8 @ 3717cb05e, round 13 cell V, median n = 16 decode step 19.887 ms of kernel busy (h100-r13-attribution.md §§C.2–C.4). Resolved by grid shape, the dense FFN’s gate and up projections are 128 graph nodes, 5 730.5 µs = 44.77 µs/node, K = 5120 N = 17408 each. At 89.1 MB of E4M3 weight per node that is 1 991 GB/s = 59.4 % of HBM.

In the SAME step, on the SAME arm, moving the SAME bytes:

projectionKNnodesµs/nodeGB/s% HBM
FFN gate + up51201740812844.771 99159.4 %
FFN down1740851206437.252 39371.4 %
SSM in_proj_qkvz5120163844834.212 45373.2 %

down reads the same 89.1 MB as one of the gate/up nodes and is 7.5 µs faster; the difference between the 59.4 % arm and the 71–73 % arms is that the first issues two launches per layer for one weight pass. The weight bytes are read once either way — this is not a traffic saving. One launch of twice the N halves the per-launch fixed cost and doubles the tile count per wave, which is what the 71.4 % row already demonstrates on this card. At an 80 % target the pair costs 11.41 GB / (0.8 × 3.35 TB/s) = 4.26 ms against a measured 5 730.5 µs → 1 476 µs/step, 7.4 % of the step and the largest single decode kernel item in the round-13 table.

§Numerics: a bit claim, not a tolerance

The fused weight is the two [inter, K] E4M3 blocks appended along N, and its [N/128, K/128] FP32 block-scale grid is the two grids appended along N/128. Splitting N therefore produces independent output columns over the same K with the same scales: output element (m, j) of the fused GEMM is the same dot product, in the same order, as element (m, j) of gate (for j < inter) or (m, j - inter) of up. Same cuBLASLt op, same epilogue. examples/native_fp8_ffn_gateup_fused_microtest.rs asserts byte equality of both halves at M ∈ {5, 8, 16} rather than a cosine.

§Layout, and why it is N-concatenation rather than an interleave

The fused output is [m, 2*inter] with gate in columns [0, inter) and up in [inter, 2*inter) — a row is [gate | up]. Three things stay simple that a tile-interleave would complicate: the loader’s fused weight is a straight device-to-device append, its scale grid is the same append one row-block wider, and each half remains addressable as an un-fused Fp8Weight VIEW, so every other rung of dense_ffn.rs’s w8_gemm! ladder keeps working on the same bytes with no change at all. The consumer pays a row stride instead of a flat index (ops::silu_mul_strided), and coalescing survives it: a row half is inter contiguous BF16 — 34 816 B at these shapes — so every warp’s 128-byte segments are whole and only the jump BETWEEN rows differs.

§Residency: net zero, by construction

A second copy of gate+up is 178.3 MB × 64 layers = 11.4 GB, which would not fit beside the bs32 KV budget. So the loader does not make one: it builds the fused buffer, re-points gate_proj and up_proj at VIEWS inside it, and Qwen35DenseWeightLoader::prune_after_load releases the two source store tensors the copy consumed. Steady-state delta is zero, and predicted_residency prices it as zero for the preflight ring fit. The load-time transient is one layer’s 178.3 MB at a time against the store tensors that have not been pruned yet — the same shape, and the same precedent, as the SSM [QKV|Z] concat that has shipped since #915.

§Band

5..=spark_runtime::buffers::GATEUP_FUSED_MAX_M rows.

  • Below 5 rung 1 of w8_gemm! (w8a16_gemv_batch4) owns the width and already makes ONE pass over each weight; there is no second launch to fuse, and the W8A8 rule this arm rides on starts at m > 4 anyway.
  • Above 16 the arm stops because the saving does. At the prefill widths these same two GEMMs run at 68.6 % of FP8 PEAK (M=4576, h100-r13-attribution.md §A.4) — compute-bound, where a launch buys nothing measurable — and the attribution’s own advice is to keep the lever scoped to decode until a prefill microtest says otherwise.

The lever is [defaults] ffn_gateup_fused: hopper true, gb10 and b200 false (no receipt, and both declare cublas_gemm_scope = "off", so the arm this changes is not even armed there). The strided SiLU consumer lives in kernels/hopper/common/silu_mul_strided.cu — HOPPER-OWNED ([kernels] overrides, an addition), so the two other NVIDIA targets do not compile a kernel they can never launch and no new cross-hardware symlink is created. ATLAS_FFN_GATEUP_FUSED=0 kills it; on gb10/b200 =1 arms a lever whose kernel lookup returns 0 and the arm declines.

Functions§

ffn_gateup_fused
Whether the compiled target arms the fused gate+up decode GEMM.