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.
cab1bcf576