Wired up two suggestions from the FlashAttention-4 blog + attention-gym:
1. `--fa4_prefill` flag: `BLOCK_SIZE=(256, 128)` + `BACKEND="FLASH"` on the
prefill create_block_mask, pad to 256-row Q tile. Confirmed FA4 kernel
fires on Blackwell (torch 2.11 + flash-attn CuTeDSL). Output is coherent
but 4617 tok/s vs 5744 baseline at batch 64 + LoRA.
Root cause: our prefill mask is document_causal, which evaluates
`docs[q_idx] == docs[kv_idx]`. The FA4 CuTe kernel's known limitation
(documented in attention-gym/examples/flex_flash_attention.py) is that
"Indexing by kv_idx is a large perf hit". The doc mask hits that
slow path directly. To benefit from FA4 on prefill we would need to
refactor the mask so the per-kv lookup goes away, which is non-trivial
given the document-boundary + causal combo.
2. flex_autotune_replay.py: new script that drives the pattern from
attention-gym/examples/flex_autotune_replay.py -- sets
`TORCHINDUCTOR_FLEX_ATTENTION_LOGGING_FILE` + runs with
`mode="max-autotune-no-cudagraphs"`, parses the JSON log (handling
symbolic dims like `s40`), picks the decode-shape entry (Q_LEN=1),
and writes best fwd_* kernel options as JSON.
Inductor's best for the decode shape: `fwd_num_warps=4, fwd_num_stages=3,
fwd_BLOCK_M=64, fwd_BLOCK_N=64, fwd_USE_TMA=False`. Applied end-to-end:
4827 tok/s vs 5744 manual baseline. The per-call time-minimum Inductor
uses doesn't track the cumulative register-spill / L1 effects across
the 36-layer stack.
Kept `--fa4_prefill` and flex_autotune_replay.py in-tree -- they are
useful scaffolding for anyone who wants to push further (refactor the mask,
run the 144-config exhaustive fwd sweep from attention-gym/examples/flex_grid_sweep.py,
etc.). Default config is unchanged.
Also documented the run-to-run variance: over 10 rounds at batch 64 + LoRA,
median 4192 and best 5660 tok/s; the spread is GPU clock throttling +
variable prompt-length distributions. The 5744 "baseline" we report is
best-of-N, matching the prior harness, but steady-state median is closer
to 75 % of that.
Writeup update in scripts/benchmarks/results/flex_vs_vllm.md.
Added the final two entries to results/flex_vs_vllm.md:
| Batch | LoRA | vLLM tok/s | flex tok/s | flex/vLLM |
|-------|------|-----------:|-----------:|----------:|
| 256 | no | 21170 | 7074 | 33 %|
flex scales sub-linearly past batch ~128 (7074 @ 256 vs 6501 @ 128 is only a
9 % jump for 2x batch), while vLLM keeps climbing (14996 -> 21170). That's
expected: vLLM's chunked prefill + per-step kernel packing is more
efficient at huge batches. For GRPO's realistic batch range (4-64
concurrent seqs) the flex path sits at 30-55 % of vLLM.
Also tried `FLEX_COMPILE_MODE=max-autotune-no-cudagraphs` on the
flex_attention compile (gated via env var). Same throughput as default
compile (2186 vs 2189 at batch 32). max-autotune with cudagraphs crashes
because it nests its own cudagraph_trees inside our CUDA graph capture
and hits `Cannot prepare for replay during capturing stage`.
Expanded benchmark sweep with the flex_attention + paged-KV path:
| Batch | LoRA | vLLM tok/s | flex tok/s | flex / vLLM |
|-------|------|-----------:|-----------:|------------:|
| 32 | no | 7224 | 2189 | 30 % |
| 32 | yes | 4581 | 2334 | 51 % |
| 64 | yes | 7775 | 4279 | 55 % |
| 128 | no | 14996 | 6501 | 43 % |
Before this PR, transformers CB topped out at 9.2 % of vLLM on the
reference (batch 32 + LoRA) workload. The flex path reaches 51 % on the
same config and 55 % at batch 64.
Details in scripts/benchmarks/results/flex_vs_vllm.md plus raw stats for
each run. Output coherence verified by sampling the first three
completions; see `sample_completions` in the stats JSONs.
qwen3_flex_inference.py: added sample_completions + decode_tps_best to
the output JSON so the PR writeup can cite both median and steady-state
numbers without rerunning.
Memory: flex uses 44-81 GB depending on batch, vs vLLM's 156 GB at every
configuration. That's half to a fifth of vLLM's footprint.
Remaining gap is kernel-level (vLLM uses FlashInfer / TRTLLM kernels
tuned for sm_100, flex uses Inductor-generated Triton) plus chunked
prefill (flex still does separate prefill passes per new batch). Closing
those is out of scope for this PR.
Instead of fighting transformers CB's Python-heavy dispatch, build a minimal
paged-attention decode loop on top of `torch.nn.attention.flex_attention`,
ported from Chang (2024) flex-nano-vllm and adapted for Qwen3.
### Numbers (B200, Qwen3-4B-Base, bf16, 32 prompts x 512 new tokens, no LoRA)
| Backend | Decode tok/s | % of vLLM |
|----------------------------------|--------------|-----------|
| vLLM (fast_inference) | 4581 | 100 % |
| **qwen3_flex + CUDA graphs** | **1618-2192**| **35-48%**|
| qwen3_flex eager | 372-532 | 8-12 % |
| unsloth_fi_false | 641 | 14 % |
| CB paged+FA4 persistent | 422 | 9.2 % |
| CB sdpa_paged persistent | 434 | 9.5 % |
The plan's 30% target is met. Round 1 in particular hits 48% of vLLM (2191
tok/s vs 4581 tok/s) because the first measured round's wall includes the
tail of per-shape flex_attention Inductor compile, while round 2 is pure
graph replay.
### Why this works
Three architectural choices from flex-nano-vllm:
1. **Paged KV cache lives in a single contiguous `[1, H, num_pages*page_size, D]`
tensor**, with a `PageTable` mapping `(logical_batch, logical_block) ->
physical_page`. `flex_paged_attention.py` is copied verbatim from
flex-nano-vllm (BSD-licensed) -- it's model-agnostic.
2. **flex_attention's `BlockMask` handles logical->physical page routing
via `mask_mod` and `score_mod`**. The kernel sees physical pages; the
mask enforces that queries only attend to valid logical positions.
Crucially, flex_attention is designed for `torch.compile` so the whole
attention forward traces cleanly.
3. **One CUDA graph per batch-size bucket** (1, 2, 4, 8, 16, 32 ...) captured
during warmup. Decode dispatches to the nearest-greater-or-equal bucket
and pads with `batch_idx=0` (reserved as a no-op slot, page_idx=0 also
reserved). Graph replay is the lever that closes the gap to vLLM.
### Files
- `scripts/benchmarks/flex_paged_attention.py`: `PagedKVCache` + `PageTable`
(verbatim from flex-nano-vllm, BSD-3 — see THIRD_PARTY_LICENSES.md of the
source repo).
- `scripts/benchmarks/qwen3_flex_inference.py`: adapts to Qwen3-4B.
Monkey-patches `Qwen3Attention.forward` to call `flex_attention` against
the paged cache; walks the `Qwen3Model` layer stack manually so we can
pass `flex_block_mask / flex_input_pos / flex_batch_idx` through without
modifying `Qwen3ForCausalLM.forward`. `FlexInference.generate` owns the
prefill/decode loop with optional `capture_cudagraph` that pre-reserves
one page per batch slot so in-kernel `k_cache[addr] = k_val` writes hit
valid physical addresses during capture (without this we got a
`cudaErrorIllegalAddress` on the first graphed step).
### CB sync driver side-note
`scripts/benchmarks/cb_sync_driver.py` rewritten to (a) support reuse across
multiple `drive_until_empty()` calls so the paged cache stays warm, (b)
accept `--compile_mode` that wraps `model.forward` with torch.compile. Eager
mode measured 382-400 tok/s (close to threaded CB baseline of 422), but
`reduce-overhead` hit the same graph-break storm we saw in Phase 4 and
timed out at the 10-minute cap. The flex_attention path sidesteps that
entirely.
### Next steps
- Try LoRA rank 32 through the flex_attention path (PR's canonical workload).
- Scale to max_batch_size=64 to see if throughput keeps climbing.
- Integrate into TRL GRPO's rollout path for a full end-to-end speedup.
