SIGN IN SIGN UP

FlashInfer: Kernel Library for LLM Serving

0 0 188 Python

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