fix(fmha_v2): fix FP8 V-scratch pipeline and varlen scheduler on SM90 (#3276)
Forked and modified from @bobboli's Original PR: https://github.com/flashinfer-ai/flashinfer/pull/2957 ### Root cause On Hopper (SM90), the warp-specialized FP8 FMHAv2 kernel runs a dedicated DMA warpgroup that explicitly transposes the V tile (via STSM) from `smem_v_scratch` into `smem_v` before BMM2, because FP8 QGMMA can't accept a transposed B-operand. The V-scratch FIFO between TMA load and STSM transpose was hardcoded to **depth 1** (`V_SCRATCH_BUFFERS = DMA_GROUP_TRANSPOSE_V ? 1 : 0`). At depth 1, producer (`_wptr`) and consumer (`_rptr`) share one physical slot and the phase bit alone distinguishes epochs N and N+1. Any momentary ordering inversion between the 128 transposer threads and DMA thread 0 corrupts the phase and deadlocks at the next `bar.sync` — observable as a hang in FP8 prefill at `head_dim=256` after enabling FP8 output coverage. The fix is to raise the V-scratch FIFO to **depth `KV_BUFFERS`** (≥ 2 for FP8 head_size ≤ 128). Raising the depth unmasks two latent bugs that depth-1 was hiding, and exposes a third race the kernel didn't previously have a window for. All three need to land together for `KV_BUFFERS > 1` to be correct, and a fourth fix is needed for the persistent scheduler on the FP8 path. ### Changes `kernel_traits.h` — raise V-scratch FIFO depth from 1 to `KV_BUFFERS`. Also tighten the V-transposer UNROLL gate from `STEP_KV > 128` to `STEP_KV >= 128` to avoid register spilling at the deeper buffer state. `dma.h` — three coupled fixes the new depth requires: - `smem_v_scratch[v_scratch_barrier_id]` → `smem_v_scratch[v_scratch_barrier_id * TILE_SIZE_V]`. At depth 1 `v_scratch_barrier_id` is always 0 so the missing stride was dormant; at depth > 1 it would corrupt the transposed V output. - Thread the captured `v_scratch_barrier_id` into `cbr_v_scratch.peek/wait/pop()` instead of letting the reader use its internal `_rptr`. Required because `_rptr` walks independently of the slot the writer reserved when the V tile was loaded. - Insert `named_barrier_wait(SYNC_BARRIER, NUM_THREADS_IN_DMA_GROUP)` immediately after `threadReserve()` in `transpose_v_tile`. Without it, DMA thread 0 can race ahead into iteration N+1 (advancing the consumed-barrier phase) while warps 1–3 still poll iteration N's expected phase, deadlocking the warpgroup. Canonical leader-ahead fix, same pattern `push_with_sync` uses. `circular_buffer.h` — add `ptr`-taking overloads of `peek/wait/advance/pop` so callers can wait on a specific slot rather than the reader's internal `_rptr`. Required by the `dma.h` fix above. FP8 persistent scheduler for ragged q-tiles — `decode_exact_dynamic_tile_id` walks `cu_q_seqlens` and computes each batch element's q-tile count on the fly. The old scheduler assumed a uniform `num_tiles_per_head` across the batch and relied on downstream skip-tile logic to ignore invalid tiles. That worked for FP16/BF16, but on the FP8 path the DMA warpgroup does real work for those ghost tiles (TMA + transpose) and the barrier machinery doesn't unwind cleanly. Gated on `DMA_GROUP_TRANSPOSE_V` so FP16/BF16 keeps its tuned scheduling. FP8-output coverage + smem budget — `fmha_library.py` now drops `kv_tile_buffers` to 1 for FP8-output `head_dim=256` (FP8→FP8 adds two output staging buffers, pushing past H100's 228KB cap). Depth-1 is safe at this config because the new `named_barrier_wait` after `threadReserve` keeps the DMA warpgroup synchronized across iterations. ### Test - Enables FP8 output prefill coverage in `tests/attention/test_fmha_v2_prefill.py`. - Removes the FP8 sliding-window test markers that were masking the original hang. ## Test Results: ``` (py312) root@82a6e7db19d1:/workspace/flashinfer# pytest tests/attention/test_fmha_v2_prefill.py ======================================================= test session starts ======================================================= platform linux -- Python 3.12.12, pytest-9.0.2, pluggy-1.6.0 rootdir: /workspace/flashinfer configfile: pytest.ini collected 2345 items tests/attention/test_fmha_v2_prefill.py sssssssss.......................................................................... [ 3%] ..............................................................ssssssssssssssssssssssss..................................... [ 8%] ........................................................................................................................... [ 14%] ........................................................................................................................... [ 19%] .............ssssssssssssssssssssssss...................................................................................... [ 24%] ........................................................................................................................... [ 29%] .......................................................................................ssssssssssssssssssssssss............ [ 35%] ........................................................................................................................... [ 40%] ........................................................................................................................... [ 45%] ......................ssssssssssssssssssssssssssssssssssssssss............................................................. [ 50%] ........................................................................................................................... [ 55%] ................................................................................................sssssssssssssssssssssssssss [ 61%] sssssssssssss.............................................................................................................. [ 66%] ........................................................................................................................... [ 71%] ...............................................ssssssssssssssssssssssssssssssssssssssss.................................... [ 76%] ........................................................................................................................... [ 82%] .................ssssssss................ssssssss................ssssssss................ssssssss................ssssssss.. [ 87%] ..............ssssssss..................................................................................................... [ 92%] ........................................................................................................................... [ 97%] ................................................ [100%] ========================================= 2096 passed, 249 skipped in 1824.87s (0:30:24) ========================================== ``` --------- Co-authored-by: Bo Li <22713281+bobboli@users.noreply.github.com>
J
Jimmy Zhou committed
d625ed4d09b0a3d23d5d4fa3d08b94f2a8cbeeee
Parent: ee31370
Committed by GitHub <noreply@github.com>
on 5/14/2026, 11:41:20 PM