Module batch16_decode

Module batch16_decode 

Source
Expand description

2026-09-25: The opt-in 5..=32-row native-FP8 decode tier. 2026-09-26: The opt-in 5..=32-row native-FP8 dense-FFN tier, w8a16_gemv_batch16: one weight pass for up to 16 rows, and two launches on row halves for 17..=32.

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

  • batch16_plan returns None unless the tier is armed and the handle resolved, and for every m outside 5..=32.
  • Both launches of a Halves plan cover at most 16 rows, the kernel’s MAX_M.

The w8_gemm! ladder in DenseFfnLayer::prefill_fp8 (reached from forward_prefill_inner) tries this tier after w8a16_gemv_batch4 (m <= 4) and the tensor-core m16_tc tier, and before the W8A8 and tile GEMMs. It selects by row count, not by phase, so a prefill chunk of 5..=32 tokens takes it too. w8a16_gemv_batch16 and w8a16_gemv_batch4 are two MAX_M instantiations of one template (w8a16_gemv_batch4.cu), and examples/native_fp8_ffn_batch16_microtest.rs requires every row to equal the scalar w8a16_gemv byte for byte.

Off by default. Measured 2026-09-11 on 1xH100, Qwen/Qwen3.8-27B-FP8, same binary, 1024x256 at C=16, tier on against off: aggregate 121.4 against 128.0 tok/s, TPOT p50 107.4 against 102.0 ms, and TTFT of a 28-token prompt 150 against 101 ms.

Functions§

ffn_batch16_enabled
2026-09-25: Whether METRALE_FFN_BATCH16 is exactly 1, which arms the tier; any other value, or none, leaves it off. Read once per process; each DenseFfnLayer copies it into batch16_enabled at construction, so the route cannot change between CUDA-graph replays.