--fa4_prefill now accepts three states: True (force on, warn + fall
back on sub-Hopper), False (force off), None / default (auto-enable
where supported). Argparse switches to BooleanOptionalAction so both
--fa4_prefill and --no-fa4_prefill work, with the default being
auto-detect from torch.cuda.get_device_capability.
Adds a cu13 / cu12 install section and a per-GPU support matrix to
scripts/benchmarks/README.md.
Adds tests/test_fa4_capability_guard.py covering the nine
combinations of (explicit-on / auto / explicit-off) x (sm_80 / sm_90
/ sm_100 / sm_120). Monkey-patches get_device_capability and stubs
PageTable / patch_qwen3_model so it runs without CUDA.
Replace the copy+merge_adapter pair in refresh_lora_merge_from_pristine
with one torch.addmm(pristine, B, A, alpha=scaling, out=W_inf) per
LoraLayer, then set merged_adapters directly so PEFT's forward
short-circuits to base_layer(x).
Previously each refresh did two passes per weight: a bf16 copy from
pristine, then PEFT merge_adapter which materialises a full [out, in]
fp32 delta via get_delta_weight and in-place adds it back. The fused
path skips the transient delta allocation and runs one cuBLAS GEMM
instead. cuBLAS accumulates the bf16 matmul in fp32 internally, so the
numerical result stays within 1 bf16 ULP of PEFT's path (verified on
the rank-32 Qwen3-4B adapter: max abs diff 1.22e-04).
DoRA, fan_in_fan_out, and lora_bias=True layers fall back to PEFT's
get_delta_weight/merge path via a single trailing merge_adapter call
after restoring their base_layer.weight from pristine. rslora is not a
fallback -- PEFT folds alpha/sqrt(r) into module.scaling[adapter], so
the fused addmm picks it up transparently via alpha=.
Drift verification still passes: base bit-identical across 10
perturb+refresh cycles, inference state deterministic after LoRA
restore.
PEFT's merge/unmerge pair is asymmetric at bf16 and leaks ~1 ULP per
cycle onto base_layer.weight. Across hundreds of GRPO refreshes the
base drifts, so the adapter trains against a moving target.
Keep a pristine base_model on GPU and a deep-copied inference_model
wrapped by PEFT. Before each rollout, restore the inference copy's
LoRA-target base_layer weights in-place from pristine and call
merge_adapter fresh. Never call unmerge_adapter.
Adds --verify_no_drift which hashes base params before/after N
perturb+refresh cycles and asserts bit-identical, and checks that the
merged inference state is deterministic after restoring the LoRA.
Update flex_vs_vllm.md with the double-copy row and memory cost.
Prior default was `peft_model.merge_and_unload()` which bakes LoRA into
the base and destroys the adapter. Same inference speed, but the adapter
is unrecoverable so you can't train on it for the next rollout -- which
GRPO explicitly needs.
Switch default to `peft_model.merge_adapter()`, which:
- Folds LoRA into `base_layer.weight` non-destructively.
- Keeps `lora_A` / `lora_B` parameters intact.
- Flips a `merged` flag inside each `LoraLayer` so its forward
short-circuits to just `base_layer(x)`, giving identical inference speed
to the destructive merge.
- Is fully reversible via `unmerge_adapter()` (bf16 round-trip error ~6e-5).
Measured end-to-end at batch 64 + LoRA rank 32:
- merge_adapter: 5785 tok/s best (was 5744 with merge_and_unload)
- merge+unmerge cycle: ~48 ms total for the 36-layer 7-target adapter,
which is <1 % of a ~5-7 s rollout -- fully amortizable per iteration.
This is *the* rollout path GRPO should use. vLLM's LoRARequest achieves
the same outcome via double-copy (pristine base + materialized base+LoRA
copy) or Punica-style fused kernels, but from a throughput standpoint
both get you to "near-merged speed with adapter separable for training".
Reframes the writeup: removes the previous panic correction that claimed
flex was 35 % of vLLM. The 35 % row is what you'd get with a naive PEFT
wrapper (3 matmuls per projection) -- a path nobody should actually
use. The real headline is still flex reaches 74 % of vLLM at 3.5 x less
memory under proper LoRA semantics.
`--no_merge_lora` flag preserved for the unmerged-PEFT path; documented as
reference only. 4-bit still uses the unmerged path (bnb merging is
unsupported).
Previous bf16 + LoRA rank 32 runs called `peft_model.merge_and_unload()`,
which bakes LoRA into the base weights and destroys the adapter. Every
subsequent forward is then plain bf16 with no LoRA-active cost -- one
matmul per projection. vLLM's LoRARequest path keeps LoRA dynamic (base
matmul + rank-r adapter matmuls + add), which is what GRPO actually needs
because the adapter has to be updateable between rollouts and training
steps.
Adds `--no_merge_lora` flag and runs the honest comparison:
| Backend | tok/s best | vs vLLM |
|--------------------------------|-----------:|--------:|
| vLLM (LoRARequest) | 7775 | 100 % |
| flex -- LoRA merged (prior) | 5744 | 74 % |
| flex -- LoRA active (no merge) | 2683 | 35 % |
The 74 % number in the earlier writeup was only meaningful if you can eat
the merge/unmerge cost between rollouts and training steps (which is not
free). The real flex-vs-vLLM gap under GRPO semantics is ~35 %, not 72 %.
vLLM wins its dynamic-LoRA number via Punica-style fused kernels that
avoid the extra matmul roundtrip. flex has no equivalent and runs
base + LoRA_A + LoRA_B as three separate matmuls per projection.
Writeup updated.
Adds --load_in_4bit (+ --model_name_4bit override) to both the flex
benchmark script and the vllm/tpaged benchmark script. When set, loads the
pre-quantized Unsloth bnb-4bit shard (e.g.
unsloth/Qwen3-4B-Base-unsloth-bnb-4bit) and keeps the LoRA adapter as a
PEFT wrapper instead of merging, because merging into 4-bit weights is not
supported.
Ties lm_head.weight to model.embed_tokens.weight post-load in both scripts,
because the bnb-4bit shards ship without an lm_head parameter even though
tie_word_embeddings is True in the config, so transformers leaves it
randomly initialised otherwise (garbage generations).
Results at batch 64 + LoRA rank 32:
| Backend | tok/s | peak mem | output |
|-------------------------------|------:|---------:|-----------|
| Unsloth fast_inference (vLLM) | 4515 | 159 GB | coherent |
| flex (this PR) | 1738 | 40.6 GB | coherent |
| transformers CB (sdpa) | 504 | 124 GB | gibberish |
4-bit costs ~40 % throughput on the vLLM path vs bf16 and ~70 % on flex.
flex regresses worse because PEFT-without-merge doubles the matmuls per
projection (base + LoRA add) on top of bnb dequant, whereas bf16 flex
merges LoRA into the base. Peak memory barely moves for vLLM because KV
cache at gpu_memory_utilization=0.8 dominates regardless of base size.
transformers CB (generate_batch) at 4-bit + LoRA produces garbage even
with lm_head tied. Likely PEFT-over-bnb + batched CB interaction; not
debugged further -- it was always the 10 % reference path.
Writeup updated in scripts/benchmarks/results/flex_vs_vllm.md with a new
"Same workload at load_in_4bit=True" section.
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.
