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