Previously the vanilla reference was a Gemma4ForCausalLM shell wrapping
the language_model (the same construction used inside
gemma4_flex_inference.main for LoRA / state-dict hashing convenience).
That is not plain HF: `Gemma4Model.forward` uses
`create_masks_for_generate(..., mm_token_type_ids, pixel_values)` to
build attention masks, while the shell calls `Gemma4TextModel.forward`
directly, which builds its own per-regime masks via `create_causal_mask`
+ `create_sliding_window_causal_mask`. Both are correct for text-only
input but their mask-bias precision differs enough to produce a
measurable drift.
The script now keeps both references alive and reports three diffs:
shell vs raw, flex vs raw, flex vs shell. On unsloth/gemma-4-E2B-it
bf16 with a 6-token prompt:
shell vs raw max 6.9e-01 mean 3.0e-01 argmax=yes top-10=10/10
flex vs raw max 6.3e-01 mean 2.2e-01 argmax=yes top-10=10/10
flex vs shell max 3.8e-01 mean 8.0e-02 argmax=yes top-10=10/10
Flex is actually closer to raw HF than the shell is. About 0.30 mean of
the flex-vs-shell-and-vs-raw gap comes from the shell's mask
construction alone, not the flex kernel. Either way the bf16 drift
band matches Qwen3 (0.3 / 0.09) and Llama-3.2 (0.13 / 0.02), and
semantic top-1 + top-10 are exact.
One-shot correctness checks for the flex_attention + paged KV path
against a vanilla HF `model(input_ids, use_cache=False)` forward on
the same prompt. Report max / mean abs diff on last-position logits
plus argmax match and top-10 overlap, so numerical drift and semantic
equivalence are both visible.
Shared approach: load the base model, deep-copy it for the flex path
so the attention patching does not mutate the vanilla comparison, run
both on the same tokenized prompt, report diffs.
`verify_gemma4_numerics.py` mirrors `gemma4_flex_inference.py`'s
text-only loader (Gemma4ForConditionalGeneration -> drop vision / audio
towers -> language_model into a Gemma4ForCausalLM shell) before
deep-copying. The softcap is NOT re-applied on the vanilla logits
since `Gemma4ForCausalLM.forward` already applies
`final_logit_softcapping`.
`verify_qwen3_numerics.py` runs `qwen3_flex_inference.FlexInference`
against either Qwen3 or Llama-3.2 via `--model_name`. `fa4_prefill` is
disabled because short prompts hit a CuteDSL sm_100 shape mismatch in
`handle_block_sparse_empty_tile_correction_sm100`.
Results on B200 bf16, 6 to 7 token prompt:
| Model | max abs | mean abs | argmax | top-10 |
|---------------------------------|---------|----------|--------|--------|
| unsloth/Qwen3-4B-Base | 0.313 | 0.088 | yes | 10/10 |
| unsloth/Llama-3.2-3B-Instruct | 0.125 | 0.021 | yes | 10/10 |
| unsloth/gemma-4-E2B-it | 0.375 | 0.080 | yes | 10/10 |
All three land in the same bf16 Triton-flex vs eager-matmul drift band;
Gemma-4's extra per-layer-input path and layer_scalar do not widen the
gap despite 20 of 35 layers going through the shared-KV link.
The sidecar design in the previous cut stored K/V at
[max_batch, n_kv, max_seq, D] layout, but the flex_attention block mask
is built for the paged cache's [1, H, n_pages*page_size, D] layout.
Shared layers running through that mismatch either needed a parallel
block mask (expensive to build per call, per layer) or had to fall back
to SDPA, which breaks the single-CUDA-graph capture story and silently
dropped the sliding-window mask on shared sliding layers.
This commit drops the sidecar entirely. Shared layers now reference the
store layer's `PagedKVCache` directly:
- `patch_gemma4_attention_forwards` allocates a cache on every
non-shared layer, then walks shared layers and points
`shared._paged_cache = store._paged_cache`, plus stashes the store
attention module itself on `shared._store_attn`.
- The store layer keeps its post-rotary `k`, `v` on
`self._last_k_val`, `self._last_v_val` before its own paged update
so shared successors can read the same packed prefill tensors.
- Shared-layer forward reads `_last_k_val` / `_last_v_val` on prefill
(q_len > 1) and `_paged_cache.k_cache` / `.v_cache` on decode. The
block_mask dispatched by `self.layer_type` works for both regimes
uniformly -- one block mask builder, one kernel compile per regime,
one CUDA graph per batch-size bucket.
Also fixes the 4-bit path: `AutoModelForCausalLM.from_pretrained` on
`unsloth/gemma-4-E2B-it-unsloth-bnb-4bit` resolves to
`Gemma4ForConditionalGeneration`, so `.model.embed_tokens` does not
exist. The loader now detects the multimodal wrapper, drops the vision
and audio towers, and moves the language_model into a
`Gemma4ForCausalLM` shell -- mirroring the bf16 path.
Benchmarks on a single B200 (sm_100), CUDA_VISIBLE_DEVICES=2, Gemma-4
E2B-it, n_prompts=64 n_rounds=5 max_new_tokens=512 max_batch_size=64
capture_cudagraph no-fa4_prefill BLOCK_M=32 BLOCK_N=32 (prefill) /
BLOCK_M=16 BLOCK_N=16 (decode):
| Config | Peak GB | Median tok/s | Best tok/s |
|-----------------|---------|--------------|------------|
| bf16 | 14.2 | 2794 | 2797 |
| bf16 + LoRA r32 | 23.0 | 2798 | 2801 |
| 4bit + LoRA r32 | 13.0 | 1659 | 1821 |
Drift verification (10 perturb+refresh cycles, noise_scale=0.01):
`base_bit_identical = true`, `inference_deterministic = true`.
Sample completions are coherent math reasoning ("Let the isosceles
trapezoid be $ABCD$ with bases ...").
The first cut routed every non-shared layer through one causal block mask
and relied on SDPA with is_causal=True for shared layers. That gives the
right semantics for Gemma-4's full_attention layers but silently drops the
sliding_attention window, so sliding layers attend far beyond their
512-token window as soon as the prefix grows.
This commit builds a block mask per attention regime (full_attention is
pure causal; sliding_attention is causal AND q_pos - kv_pos < window) and
passes both into each attention call as a dict, letting the patched
forward select by self.layer_type. For the shared-KV sidecar path, the
SDPA call now receives an explicit attn_mask composed the same way, so
sliding shared layers also respect the window. Strict less-than
comparison matches Unsloth's flex-attention convention for GPT-OSS.
`_causal_blockmask_with_window` and `_prefill_blockmask_with_window` are
local to this file; the shared helpers in `flex_paged_attention.py` stay
untouched. `FlexGemma4Inference` now caches both logical decode masks,
slices both per-row in `_decode_block_mask`, and runs the PageTable's
logical->physical conversion on each before passing them down.
Extends the flex_attention + paged KV + CUDA graphs engine from Qwen3 and
Llama-3.2 to Gemma-4-E2B-it via a new standalone file that imports the
shared helpers (PagedKVCache, PageTable, Sequence, LoRA double-copy,
drift verification, flex_attention_compiled, _apply_rotary) from
qwen3_flex_inference.py. The Qwen3 / Llama path is not modified.
Gemma-4 diverges from Qwen3 / Llama in ways that cannot be folded into a
single hasattr guard:
- KV-sharing layers. E2B has 35 layers; the upper 20 lack k_proj / v_proj
/ k_norm / v_norm and consume the full prefix K/V produced by a store
layer further up the stack. We allocate a sidecar dict of
[max_batch, n_kv, max_seq, head_dim] buffers at fixed device addresses,
populated by store layers during prefill and read by shared layers
through eager SDPA (their layout does not match the paged cache's
block-mask shape).
- Dual attention regimes. full_attention (head_dim=512, rope_theta=1e6)
and sliding_attention (head_dim=256, sliding_window=512) coexist. We
precompute both (cos, sin) pairs via
Gemma4TextRotaryEmbedding(x, position_ids, layer_type) and dispatch on
self.layer_type inside the patched attention forward.
- Per-layer input embeddings. The walker threads the
[B, S, num_layers, hidden_size_per_layer_input] table from
get_per_layer_inputs + project_per_layer_inputs through each layer's
per_layer_input_gate / act_fn / mul / per_layer_projection /
post_per_layer_input_norm path.
- Four norms per block with double residuals. Attn (input_layernorm,
post_attention_layernorm) and MLP (pre_feedforward_layernorm,
post_feedforward_layernorm), plus the per-layer-input residual and
layer_scalar multiply.
- Final logit softcap. tanh(logits / 30.0) * 30.0 on the lm_head output.
Transformers>=5.5.0 is required for the gemma4 module. A _require_gemma4
guard at main() exits with a clear install hint when the module is
missing, so the workspace's Qwen3 / Llama path stays on the existing
transformers install.
The text-only path loads Gemma4ForConditionalGeneration, drops the
vision and audio towers, and moves the language_model into a
Gemma4ForCausalLM shell so LoRA, state-dict hashing, and the double-copy
refresh treat it like any other HF decoder model.
CLI mirrors qwen3_flex_inference.py: --model_name (default
unsloth/gemma-4-E2B-it), --lora_adapter, --load_in_4bit,
--capture_cudagraph, --verify_no_drift, --chat_template {auto,grpo,
native}, --fa4_prefill / --no-fa4_prefill, plus decode / prefill
kernel_options for Triton block tuning.
Smoke-tested on B200 (sm_100) with bf16, batch 2, max_new_tokens 16,
no-fa4_prefill + BLOCK_M=32 / BLOCK_N=32 (Gemma-4 head_dim=256 exceeds
FA4's 128-limit on sm_100): 52 tok/s cold, coherent completions.
scripts/benchmarks/README.md gains a paragraph covering the
transformers>=5.5 dependency, the head_dim=256 constraint, and the
recommended kernel_options for B200.
--chat_template {auto,grpo,native} matches the flag added to
qwen3_flex_inference so cross-engine comparisons can hold the prompt
template constant per model (auto: GRPO for Qwen3, tokenizer native
otherwise). Threaded through build_prompts() and applied to all three
backends (vllm, tpaged, unsloth_fi_false).
--enforce_eager (vLLM backend only) forwards enforce_eager=True to
FastLanguageModel.from_pretrained so the vLLM engine skips torch.compile
+ cudagraph capture. Needed because vLLM 0.19.1 on torch 2.10 raises
`RuntimeError: Tried to erase Node size_1 but it still had 2 users`
inside compilation.backends.split_graph during the first forward; the
eager path still uses PagedAttention + FlashInfer for decode, so the
measurement stays meaningful (just no graph capture).
The flex_attention + paged KV + CUDA graphs inference engine was
Qwen3-specific in a handful of places, but the underlying engine
(PageTable, PagedKVCache, manual forward walker, decode graph capture,
double-copy LoRA rollout, FA4 capability guard) reads only attributes
that LlamaAttention / LlamaModel also expose. This change makes the
engine run on both Qwen3 and Llama-3.2-3B-Instruct.
Attention forward factory:
- make_flex_qwen3_attention_forward -> make_flex_attention_forward
- Guard the per-head QK RMSNorm call behind hasattr(self, "q_norm").
Qwen3 has it, Llama does not. The Qwen3 path is byte-equivalent to
before: RMSNorm on [B, S, H, D] (per-head) then transpose.
- patch_qwen3_model -> patch_model_attention_forwards.
Chat template selection:
- New --chat_template {auto,grpo,native}. auto picks GRPO for Qwen3
and the tokenizer's shipped template otherwise. grpo forces GRPO
(matches prior Qwen3 baselines). native forces the tokenizer's own
template (Llama-3.2-Instruct only produces coherent completions
with its shipped Instruct template).
Stats JSON:
- backend: "qwen3_flex" -> "flex"; adds "model_name" so multi-arch
runs land in a single schema.
README: one paragraph noting Llama-3.2 support + the --chat_template
native flag.
Measured on B200 (sm_100), n_prompts 64, max_new_tokens 512, 5 rounds,
--capture_cudagraph, double-copy LoRA rank 32:
Qwen3-4B-Base bf16 3975 tok/s 44.2 GB
Qwen3-4B-Base bf16 + LoRA 3656 tok/s 52.0 GB
Qwen3-4B-Base 4bit + LoRA 1734 tok/s 40.6 GB
Llama-3.2-3B-Inst bf16 4216 tok/s 34.7 GB
Llama-3.2-3B-Inst bf16+L 4205 tok/s 40.9 GB
Llama-3.2-3B-Inst 4bit+L 1892 tok/s 31.7 GB
--verify_no_drift passes on both arches (base bit-identical across 10
perturb+refresh cycles, inference hash deterministic).
Llama runs the Triton flex_attention backend instead of FA4:
flash-attn-4 b9's sm_100 kernel raises a NoneType in
handle_block_sparse_empty_tile_correction_sm100 on Llama-3.2's head
shapes. Qwen3 is unaffected. Pass --no-fa4_prefill on Llama; auto-FA4
still enables on Qwen3.
Delete two unreferenced drivers (qwen3_grpo_notebook.py, qwen3_grpo_unified.py)
that duplicated the canonical trio. Port the --compile_mode / --compile_dynamic
flags from unified into qwen3_grpo_naive.py and qwen3_grpo_tpaged.py before
deletion so the torch.compile path is preserved on the training-side backends
(vLLM is excluded because it owns its own inference graph).
Extract the 20-line StepTimer TrainerCallback, the per-step stats JSON writer,
the vLLM GuidedDecodingParams shim, and the optional torch.compile wrapper
into unsloth_grpo_common.py so the three canonical drivers
(qwen3_grpo_{vllm,naive,tpaged}.py) share one implementation. Stats schema is
unchanged: backend, train_wall_s, peak_memory_gb, step_wall_s, losses, rewards,
max_prompt_length, max_completion_length, num_generations, max_steps, plus
backend-specific extras (attn_impl, persistent_cb) passed through write_stats's
extra kwarg.
Add a short paragraph to scripts/benchmarks/README.md describing the new
--compile_mode flag.
Verified:
- python -m py_compile on all four modified files.
- --help on all three drivers shows --compile_mode on naive + tpaged only.
- 2-step tpaged smoke (flash_attention_2, num_generations=2, pdb=2) runs to
completion on B200. Stats JSON schema matches the pre-refactor output exactly.
Net: 7 files changed, +235 / -1093, 21 -> 19 benchmark files.
Remove the raw benchmark output JSONs from the PR diff. The writeup
markdowns under scripts/benchmarks/results/*.md keep their filename
anchors as a record of what each table was measured from; anyone who
wants the raw numbers can re-run the benchmark scripts.
41 files removed, 6068 lines deleted.
Removed JSON files under scripts/benchmarks/results/stats/ that no
writeup markdown in scripts/benchmarks/results/*.md referenced.
Intermediate debugging dumps (cb_sync_smoke, cb_tpaged_64_lora_4bit*,
flex_256x512, flex_32x512_eager, flex_32x512_mauto_nocg,
flex_64_lora_4bit*, flex_64_lora_bf16_{fusedmerge,mergeadapter,nomerge},
flex_verify_fusedmerge, unsloth_fi_true_64_lora_4bit). No writeups
edited -- every `results/*.md` reference still resolves.
--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.
Inductor was specialising create_block_mask on the first prefill shape
it saw (warmup with 16 prompts -> small total L). When round-0 prefill
ran at bs=64 with a much larger packed L, the cached triton block-mask
kernels launched with the wrong shape constants and hit CUDA illegal
memory access inside the document_causal mask construction, even
though the fused GEMMs and graph-captured decode path were fine.
torch.compile(create_block_mask, dynamic=True) keeps L as a runtime
arg so the same kernels work across the warmup and full-batch prefill
shapes.
Add the drift-verification and end-to-end rollout stats for the fused
addmm path: base bit-identical across 10 cycles, and bs=64 + LoRA +
capture_cudagraph + decode_kernel_options reaches 5057 tok/s median,
5224 tok/s best on B200 at 52 GB peak -- within noise of the prior
5785 tok/s baseline. Variance is round-0/1 warmup (3298, 3607 tok/s)
rather than steady-state (4940, 5082, 5057 tok/s).
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.
30-step results:
| Backend | Train wall | Median step | Peak mem | % of vLLM |
|----------------|-----------|-------------|----------|-----------|
| vLLM | 215.9 s | 5.14 s | 159 GB | 100 % |
| fi_false | 1165.4 s | 41.30 s | 10.7 GB | 12.4 % |
| cb_paged | 1564.5 s | 39.82 s | 61.9 GB | 12.9 % |
Pairwise diff vs vLLM (30 steps, compare_grpo_runs.py):
| Pair | max |loss| | max |kl| | max |reward| |
|---------------------------------|------------|------------|---------------|
| vLLM vs unsloth_fi_false | 0.39 | **0.015** | 9.25 (noisy) |
| vLLM vs cb_paged | 0.83 | (missing) | 6.25 (noisy) |
KL trajectory match between vLLM and unsloth_fi_false is the load-bearing
equivalence signal: both stay in [0, 0.015] across all 30 steps, so the
policy drift guardrail behaves the same. Reward diffs of ~3-9 are expected
because the rollout backends produce different completions even at
temperature=0.1 (kernel-level non-determinism).
Two caveats documented in results/grpo_equivalence.md:
- cb_paged's StatisticsCallback doesn't capture TRL's kl log entry because
TRL emits kl on a separate log call that doesn't include loss.
- cb_paged's grad_norm (~200-900) is unclipped pre-optimizer, while vLLM's
goes through Unsloth's internal max_grad_norm=1.0. Not a correctness bug,
just not apples-to-apples until cb_paged sets max_grad_norm in GRPOConfig.
Also updates the report with Phase 3 + Phase 4 status (CB sync driver eager
works; CUDA graph capture hangs on output_ids slice pending fix; torch.compile
on training step incompatible with both unsloth_fi_false and cb_paged).
Phase 4 (torch.compile on training step) -- negative result documented:
- unsloth_fi_false + compile_mode=default: crashes immediately with
`PeftModel_fast_forward() got multiple values for argument 'input_ids'`.
Unsloth's monkey-patched forward and Dynamo's argument rebinding don't
compose.
- cb_paged + compile_mode=default: Dynamo emits 700+ graph breaks /
recompiles during the first optimizer step and never makes progress
(killed after 10 minutes at step 0/10). Root cause trail:
`Tensor.requires_grad_()` inside `modeling_utils.make_inputs_require_grads`
triggers GB0125 (no Dynamo support), which propagates up through TRL's
`_compute_loss`. Fixing this would require restructuring GRPO's LoRA
gradient enablement path -- out of scope for this PR.
- vllm is excluded from Phase 4 because vLLM owns its own compile pipeline.
Net: torch.compile on the training step is not a quick win for the non-vLLM
paths in this stack. Phase 3 (CUDA graph on the rollout decode step) remains
the right lever for closing the CB <-> vLLM gap, and Phase 1 already
demonstrates that Unsloth's fast_inference=False path narrows the gap to
~14% of vLLM at 1/7th the peak memory without any compile.
Helper script `scripts/benchmarks/compare_grpo_runs.py`:
- Wraps `torch_debugging_utils.compare_training_runs` (loss/grad-norm diff)
- Adds reward/KL/time pairwise diffs with `max_abs` and `mean_abs`
- Reads the StatisticsCallback-emitted JSON written by
`qwen3_grpo_unified.py`
Sample output on Phase 2 vibe (10 steps):
vllm vs unsloth_fi_false:
max_loss_diff = 0.40, max_kl_diff = 0.009 (both tiny), reward_diff
mean 1.85 (different rollouts expected across backends at temp=0.1)
vllm vs cb_paged:
max_loss_diff = 0.29, max_grad_norm_diff = 715 (cb_paged has no
gradient clipping on the vanilla-HF path; vLLM path is clipped to
1.0 by Unsloth internally -- apples-to-oranges without matching
clipping, tracked for the 30-step run).
Phase 2 results (`scripts/benchmarks/results/grpo_equivalence.md`):
- vLLM: 74.4s train, 4.14s median step, 158 GB peak (100%)
- unsloth_fi_false: 355s train, 23.95s median, 10.7 GB peak (17%)
- cb_paged (via sdpa_paged load): 466s train, 36s median, 55.6 GB peak (11.5%)
Coherence gate passes on all three backends: losses finite, rewards in the
expected early-GRPO range, KL trajectories qualitatively matched between
vLLM and unsloth_fi_false in [0, 0.015]. Memory story is striking:
unsloth_fi_false uses 15x less memory than vLLM.
qwen3_grpo_unified.py fixes:
- Auto-adjust per_device_train_batch_size -> num_generations for vanilla-HF
backends (Unsloth's loader does this automatically; TRL on the HF path
doesn't and crashes on the divisibility check).
- cb_paged now loads with sdpa_paged (not paged_attention). The FA4
paged_attention kernel requires cu_seq_lens_q on every forward, but the
GRPO training forward feeds a dense batch without them. sdpa_paged
gracefully falls back to plain SDPA in that case and still exercises the
paged path during the CB rollout.
cb_sync_driver.py fixes:
- FIFOScheduler no longer accepts manual_eviction in its signature; dropped.
- drive_until_empty used to check has_pending_requests() before calling
prepare_next_batch(), which returned False at startup because nothing had
yet been pulled from the input_queue. Now the loop drains the input queue
first and exits only when both queues + scheduler are empty.
Smoke test on GPU 1 (8 prompts, 64 tokens): eager path produces 512 correct
tokens; CUDA-graph path hangs during first-step capture (PagedAttentionCache
probably allocates on first use). Tracked for the next commit.
Phase 1 results (`scripts/benchmarks/results/lora_rollout_baselines.md`):
- vLLM+LoRA: 4581 decode tok/s, 156 GB peak (gold, 100%)
- unsloth_fi_false+LoRA: 641 tok/s, 15.8 GB peak (14%, 7x lower mem)
- CB paged+FA4 persistent+LoRA: 422 tok/s (9.2%)
- CB sdpa_paged persistent+LoRA: 434 tok/s (9.5%)
Headline finding: Unsloth's `fast_inference=False` path (custom HF inference
kernels with cached fp16 LoRA in `fast_linear_forward`) is 1.5x faster than
CB at 1/7th the peak memory. Phase 2 will include it as a first-class backend.
cb_vs_vllm_generation.py:
- New --lora_adapter flag. vLLM uses LoRARequest; tpaged uses
PeftModel.from_pretrained (no merge_adapter so we measure LoRA-active
inference); unsloth_fi_false copies the adapter weights into Unsloth's
get_peft_model wrapper (with key normalization so PEFT's base_model.model.
prefix and Unsloth's .default. wrapper both match).
- New unsloth_fi_false backend: batched generate across all 32 prompts in
a single call after FastLanguageModel.for_inference(model).
- Exposed sampling knobs (temp/top_p/min_p/top_k); defaults are the
equivalence params.
cb_sync_driver.py (Phase 3 scaffold):
- SyncCBDriver owns PagedAttentionCache + ContinuousBatchProcessor +
FIFOScheduler on the main thread. Never calls manager.start() so there is
no background thread.
- slice_inputs=False => fixed-shape buffer views each step => CUDA graph
replay is safe.
- use_cuda_graph=True path: 2-step eager warmup, then capture one decode
step, then replay. `_is_pure_decode()` keeps prefill out of the graphed
path since those have varying shapes.
- Greedy sampling only (CUDA-graph-safe); stochastic sanity checks stay in
the non-graphed path.
- Standalone benchmark harness at the bottom.
qwen3_grpo_unified.py (Phase 4 scaffold):
- Single entrypoint for vllm / unsloth_fi_false / cb_paged / cb_sdpa /
naive_trl backends sharing dataset, reward funcs, sampling, and the
torch_debugging_utils StatisticsCallback.
- New --compile_mode {default,reduce-overhead,max-autotune-no-cudagraphs}
that compiles `trainer.model.forward` and `trainer.ref_model.forward`
after the trainer is built. CompileDebugger tracks graph breaks and
recompiles. Skipped for vLLM since vLLM owns its own compile pipeline.
- Post-warmup median (skip first 3 steps) is computed and saved alongside
the full per-step logs.
make_lora_adapter.py: writes a canonical PEFT adapter to
outputs/lora_rank32_fresh. Re-initializes lora_B with a tiny gaussian so
the adapter isn't a no-op (PEFT's default zero-init would let LoRA kernels
short-circuit).
qwen3_grpo_notebook.py (Phase 0): notebook-to-script port with
StatisticsCallback and equivalence sampling. 10-step reference reported in
scripts/benchmarks/results/notebook_ref_10.md (median step 5.80s,
peak 158.9 GB).
Phase 0 (canonical reference):
- scripts/benchmarks/qwen3_grpo_notebook.py: notebook-to-script port of
Qwen3_(4B)-GRPO.ipynb with StatisticsCallback from torch_debugging_utils
and equivalence-friendly sampling (temp=0.1, top_p=0.97, min_p=0.5, top_k=5).
- scripts/benchmarks/results/notebook_ref_10.md: 10-step reference run table
(median step post-warmup = 5.80s, peak 158.9 GB).
- scripts/benchmarks/results/stats/notebook_ref_10.json: full per-step logs
for downstream compare_training_runs checks.
Phase 1 (rollout-only LoRA comparison scaffold):
- scripts/benchmarks/make_lora_adapter.py: one-shot that materializes a
rank-32 LoRA at outputs/lora_rank32_fresh. Re-initializes lora_B with a
tiny gaussian so the adapter isn't a no-op (otherwise LoRA kernels can
short-circuit and we'd be measuring the base model).
- scripts/benchmarks/cb_vs_vllm_generation.py: extended with --lora_adapter
for vLLM (LoRARequest) and tpaged (peft.PeftModel.from_pretrained,
no merge_adapter), plus a new unsloth_fi_false backend that exercises the
custom HF inference path (cached fp16 LoRA via fast_linear_forward).
Sampling knobs are exposed and default to equivalence params.
Phase 2 scaffold:
- scripts/benchmarks/qwen3_grpo_unified.py: single entry point for all 5
backends (vllm, unsloth_fi_false, cb_paged, cb_sdpa, naive_trl) sharing
dataset, reward funcs, sampling, and StatisticsCallback. Skips the first
3 steps when reporting median step wall.
No unsloth internals touched.
New scripts under scripts/benchmarks/:
- flash_attn_fa4_shim.py: monkey-patches that let transformers CB dispatch to
Flash Attention 4 on Blackwell (sm_100). CB's ContinuousBatchProcessor
otherwise emits a 4D paged attention mask for flash_attention_2 (which then
breaks _flash_attention_forward's _upad_input branch), and passes
max_seqlen_q instead of max_length_q. The shim skips the mask for FA and
accepts both names.
- persistent_cb.py: replaces model.generate_batch with a version that reuses
one ContinuousBatchingManager across calls, avoiding the per-step
PagedAttentionCache realloc. Wired up behind --persistent_cb on the tpaged
and standalone scripts.
- qwen3_grpo_naive.py: vanilla HF model.generate + TRL GRPOTrainer, no vLLM
and no CB. Mirrors the TRL docs example. Useful as a third column in the
comparison and also as a "will this at least converge" sanity check.
Adds --attn_impl and --persistent_cb flags to the existing generation and
training scripts. No changes to Unsloth internals.
Updated README.md with the FA install recipe (flash-attn-4==4.0.0b9, plus
a small site-packages shim that re-exports FA4's cute.* symbols under the
FA2 flash_attn namespace so transformers' is_flash_attn_2_available() and
_lazy_imports("flash_attention_2") succeed on B200).
Benchmark numbers on a single B200, Qwen3-4B-Base LoRA rank 32, bf16:
Generation microbenchmark (32 prompts, 512 new tokens):
vLLM 7224 decode tok/s (100%)
CB paged|sdpa 527 decode tok/s ( 7.3%)
CB paged|flash_attention_2 (FA4) 709 decode tok/s ( 9.8%)
CB paged|flash_attention_2 persistent 529 decode tok/s ( 7.3%)
GRPO training (max_steps=20, num_generations=2, per_device_batch=2):
vLLM colocated 136.6 s peak 157 GB
naive TRL (HF generate) 910.0 s peak 15 GB
CB SDPA 1521.5 s peak 98 GB (prior run)
CB FA4 1470.7 s peak 82 GB
CB FA4 + persistent 1562.1 s peak 87 GB
CB FA4 + ng=4 persistent 1597.0 s peak 94 GB
FA4 is a real ~1.4x improvement over SDPA for CB decode throughput but the
50% of vLLM target is still not reached. The remaining gap is driven by
CUDA graph capture (which ContinuousBatchingManager still NotImplementedErrors
on) and vLLM's scheduler being more efficient for decode-heavy GRPO rollouts.
Naive TRL generate is the honest small-rig baseline: 6.7x slower than vLLM
at 10% of the VRAM footprint, and ~1.7x faster than CB here.
Adds reproducible scripts under scripts/benchmarks/ that compare vLLM
colocated rollouts against the transformers continuous batching API
(model.generate_batch, paged attention) for GRPO training on Qwen3-4B.
Contents:
- unsloth_grpo_common.py: shared dataset, reward functions, and GRPO
hyperparameters so the two backends differ only in the rollout engine.
- qwen3_grpo_vllm.py: baseline training entry using fast_inference=True
and TRL use_vllm=True, vllm_mode=colocate.
- qwen3_grpo_tpaged.py: candidate using a vanilla HF Qwen3 + PEFT LoRA
with TRL use_transformers_paged=True.
- cb_vs_vllm_generation.py: standalone generation microbenchmark.
- README.md: integration notes, reproduction steps, and observed numbers.
On a single B200 with Unsloth Qwen3-4B-Base at LoRA rank 32 bf16,
transformers continuous batching reaches 7-9 percent of vLLM throughput
on this workload. The README documents the integration sharp edges
(top_k=-1, PagedAttentionCache default upper bounds, Unsloth's
Qwen3Attention_fast_forward bypassing the functional attention
interface, TRL importing GuidedDecodingParams from a newer vLLM that no
longer exports it, and UnslothGRPOTrainer expecting for_training /
for_inference hooks on the model).
The scripts are intentionally self-contained so they are easy to rerun
after either upstream change that could close the throughput gap
(flash-attn availability, CUDA graphs in ContinuousBatchingManager,
persistent paged caches across generate_batch calls, or a
paged-compatible Unsloth attention forward).
