Module batch16_decode

Module batch16_decode 

Source
Expand description

The 5..=32-row native-FP8 DECODE tier (#927) — same child-module reason as w8a8_prefill above: it reads this layer’s private kernel handles, and the arm’s rule, its opt-in and the measurements that made it opt-in do not fit in this file’s budget. The 5..=32-row native-FP8 dense-FFN DECODE tier — w8a16_gemv_batch16.

WHY (#927). Measured on 1xH100, 2026-09-11, Qwen/Qwen3.8-27B-FP8, tip fbbe70767: the decode step cost 44 ms at 4 active rows and 224 ms at 16 (TPOT), so raising the batch cap from 4 to 16 made C=16 aggregate throughput FALL from 76 to 62 tok/s. Five extra rows cost 5x the step.

The cliff is a dispatch gap, not a kernel one. dense_ffn.rs’s w8_gemm! claimed only (1..=4) for w8a16_gemv_batch4; at m = 5..16 it fell to the transposed w8a16_gemm_n128_m128 / w8a16_gemm_pipelined tile GEMMs. Those pad M to a 128-row MMA tile, so at M=16 seven eighths of every tile is padding and the kernel turns 5-12 TFLOP/s while the FFN at decode widths is purely weight-bandwidth bound. w8a16_gemv_batch16 — the MAX_M=16 instantiation of the SAME template as w8a16_gemv_batch4, already in kernels/gb10/common/w8a16_gemv_batch4.cu — makes ONE pass over the FP8 weight for up to 16 rows.

NUMERICS. Every row is bit-identical to the scalar w8a16_gemv that M=1 decode runs: same K-iteration order, same per-row reduction tree, the accumulators are independent and M appears in no row’s operand sequence. H100 receipt on #932: M=8 and M=16 both unequal_bf16=0 max_abs=0. So this moves widths 5..=32 from a REASSOCIATING tile GEMM onto the bits decode already produces at M=1 — the direction that removes a numerics seam rather than adding one.

17..=32 runs the same kernel TWICE on contiguous row halves. The FFN activations and outputs are contiguous [m, k] / [m, n], so a half is a plain byte offset — no staging, no strided variant. Two weight passes still beat one M-padded MMA tile at these widths, and it means a max_batch_size of 32 never reaches the tile GEMMs at decode either.

ARM ORDER in w8_gemm! (see dense_ffn.rs) is deliberate and this module owns the 4th and 5th rungs:

  1. m <= 4 -> w8a16_gemv_batch4
  2. m 5..=16 -> w8a16_gemm_m16 (ATLAS_FFN_M16_TC only)
  3. m 17..=32 -> w8a16_gemm_m16 x2 halves (ATLAS_FFN_M16_TC only)
  4. m 5..=16 -> w8a16_gemv_batch16 (here)
  5. m 17..=32 -> w8a16_gemv_batch16 x2 halves (here)
  6. W8A8 block-scaled prefill (#917/#928)
  7. transposed / pipelined / base W8A16 tile GEMMs

Rungs 2-3 (dense_ffn_m16_tc.rs) are OFF by default and, when an operator sets the lever, they take these same widths onto a tensor-core MMA that REASSOCIATES the K reduction. This module’s bit-exactness claim below is about the arm THIS module owns; with the lever set, the FFN’s 5..=32 output is the MMA’s, within 2 BF16 ULP of the scalar rather than equal to it.

🪤 CONSEQUENCE, stated because it is a real boundary move: the W8A8 prefill arm’s own rule (dense_ffn_w8a8_prefill.rs) starts at m > 4, so with rungs 2-3 ahead of it the W8A8 path begins at m > 32 in practice. 5..=32 are decode widths where one weight pass beats any MMA tile — but a prefill of 5..=32 tokens (a very short prompt, or the TAIL CHUNK of a chunked prefill) takes the GEMV too. That consequence is what the serving A/B below caught, and it is why this tier ships disarmed.

🔴 DEFAULT OFF — OPT-IN VIA ATLAS_FFN_BATCH16=1. The cliff above is real and this kernel is the right shape for it, but on the one target where the tier has been A/B’d end to end it is a net LOSS in serving. H100 round 5, 2026-09-11, Qwen/Qwen3.8-27B-FP8, single variable — same binary, same 16-way burst, the tier the only difference:

1024x256, C=16tier ONtier OFF
aggregate tok/s121.4128.0
TPOT p50107.4 ms102.0 ms
28-token smoke TTFT150 ms101 ms

Per-phase at n=16 (ATLAS_MS_PROFILE=1, so eager — read the ratios): with the tier OFF the step goes 86.80 -> 82.32 ms, ssm 63.31 -> 59.91 ms and attn 19.90 -> 18.82 ms (-5.2% to -5.4% each); head does not move. ssm per layer returns to 1248 us against a pre-#927 1252 us — the tier’s cost is the whole of the regression it introduced, not part of it.

WHY it loses although the kernel wins at 16 rows: it dispatches by ROW COUNT, not by phase, so a chunked prefill’s tail chunk lands in the band — a 1193-token prompt splits 1168 + 25, and the 25-row tail takes the GEMV. That is a FIXED ~35 ms TTFT cost per request (49 ms on the 28-token smoke), which no decode-rate gain at these widths pays back.

🚨 AND IT HAS NEVER BEEN MEASURED ON GB10. w8a16_gemv_batch16 is an instantiation in w8a16_gemv_batch4.cu, so the handle resolves on every target that carries that module — GB10 included. A default-ON tier would ship an unmeasured routing change to the target this repo serves, on the strength of an H100 number that came out negative. Opt-in is the honest default until a GB10 A/B exists; if one wins there, the lever to flip is this file’s, not the caller’s.

WHAT STAYS DEFAULT-ON, and why it is a different lever: the attention o_proj groups-of-16 arm, the QKV band widening and the SSM MTP-verify arms key off their OWN kernel handles and never read this switch. They were ON in BOTH arms of the A/B above, so none of the movement in that table is theirs to claim or to blame — including the attn phase’s -5.4%, which moved while they were untouched. Each is bit-identical per row to the M=1 w8a16_gemv it replaces, which is a numerics improvement that does not depend on the FFN result either way.

Functions§

ffn_batch16_enabled
ATLAS_FFN_BATCH16 opt-in: the value 1 — and only 1 — arms the 5..=32-row tier. Anything else, absence included, leaves those widths on the pre-#927 arms.