[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

scripts/benchmarks/README.md

@@ -0,0 +1,252 @@
+# Qwen3-4B GRPO rollout engine benchmarks
+
+This directory holds the reproducible scripts behind the experiment
+documented in the accompanying PR: can Hugging Face transformers'
+continuous-batching API (`model.generate_batch`, backed by
+`PagedAttentionCache`) serve as a drop-in replacement for vLLM during GRPO
+rollouts on the Qwen3-4B notebook?
+
+The short answer on a single NVIDIA B200 with Unsloth Qwen3-4B-Base, LoRA
+rank 32, bf16: transformers continuous batching is functionally correct and
+integrates with TRL's `use_transformers_paged=True` path, but end-to-end
+throughput lands at around 7 to 10 percent of vLLM colocated even after
+wiring in Flash Attention 4 on Blackwell. Full numbers and per-step timings
+are in the PR description.
+
+## Files
+
+| File | Purpose |
+|---|---|
+| `unsloth_grpo_common.py` | Shared dataset loading, reward functions, and GRPO hyperparameters |
+| `qwen3_grpo_vllm.py` | vLLM baseline training entry (`fast_inference=True`, `use_vllm=True`, `vllm_mode="colocate"`) |
+| `qwen3_grpo_naive.py` | Naive TRL path (vanilla HF `model.generate`, no vLLM, no CB) matching https://huggingface.co/docs/trl/grpo_trainer |
+| `qwen3_grpo_tpaged.py` | Continuous-batching candidate (`fast_inference=False`, `use_transformers_paged=True`, vanilla HF + PEFT). Supports `--persistent_cb` |
+| `cb_vs_vllm_generation.py` | Standalone generation microbenchmark across both engines; supports `--attn_impl`, `--persistent_cb` |
+| `flash_attn_fa4_shim.py` | Installs two monkey-patches that let CB dispatch to FA4 when `--attn_impl flash_attention_2` is selected |
+| `persistent_cb.py` | Replaces `model.generate_batch` with a version that reuses a single `ContinuousBatchingManager` |
+
+## Flash Attention 4 on Blackwell (sm_100)
+
+The CB code path has three attention implementations: `eager_paged`,
+`sdpa_paged`, `flash_attention_2`. The last one requires the legacy
+`flash_attn` Python package, which does not install cleanly on B200 today:
+
+- `flash_attn==2.8.3+cu12torch2.8cxx11abiTRUE-cp313` from the Dao-AILab
+  releases hits `undefined symbol: _ZNK3c106SymInt6sym_neERKS0_` on torch
+  2.9.1 (ABI drift between torch 2.8 and 2.9).
+- `flash_attn_3-3.0.0-cp39-abi3-manylinux_2_28_x86_64.whl` from the PyTorch
+  wheel index installs but was built for sm_80 and sm_90a only. B200 is
+  sm_100. The kernel call fails with "no kernel image is available for
+  execution on the device".
+- `flash-attn-4==4.0.0b9` (pure Python CuTeDSL, Dao-AILab) works on B200.
+  It exposes `flash_attn.cute.flash_attn_varlen_func`.
+
+The recipe this repo uses:
+
+```bash
+uv pip install --no-deps flash-attn-4==4.0.0b9
+```
+
+plus a tiny site-packages shim that re-exports FA4 symbols under the
+FA2 `flash_attn` namespace so transformers' `is_flash_attn_2_available()` and
+`_lazy_imports("flash_attention_2")` succeed. The shim lives out of tree in
+`lib/python3.13/site-packages/flash_attn/__init__.py` +
+`flash_attn/bert_padding.py` + a `flash_attn-2.8.3.dist-info/` directory
+with enough metadata to satisfy `importlib.metadata.version("flash_attn")`.
+
+On top of that, `flash_attn_fa4_shim.py` monkey-patches two rough edges in
+the CB to FA integration that are unrelated to which FA version you use:
+
+1. `ContinuousBatchProcessor.return_attention_mask` returns `False` for
+   `flash_attention_2` / `flash_attention_3` so CB does not emit a 4-D paged
+   attention mask that breaks `_flash_attention_forward`'s `_upad_input`
+   branch.
+2. `_flash_attention_forward` accepts `max_seqlen_q` / `max_seqlen_k`
+   as aliases for `max_length_q` / `max_length_k`. Without this rename CB's
+   model kwargs never bind and FA is called with `max_seqlen_q=None`.
+
+## Install
+
+Prereqs: `torch >= 2.5` for `torch.nn.attention.flex_attention`. The
+Triton backend that flex_attention uses by default runs on Ampere,
+Hopper, and Blackwell -- no separate install.
+
+FA4 (CuTeDSL) targets Hopper and Blackwell only. `qwen3_flex_inference.py`
+auto-enables FA4 on supported GPUs and falls back to the Triton
+`flex_attention` backend elsewhere. `--fa4_prefill` forces on (warns +
+falls back if the GPU does not support it); `--no-fa4_prefill` forces off.
+
+CUDA 13 (recommended, used on B200 / RTX 50xx):
+
+    pip install --index-url https://download.pytorch.org/whl/cu130 torch
+    pip install "flash-attn-4[cu13]"
+
+CUDA 12 (H100 boxes still on cu12):
+
+    pip install torch  # default index is cu12
+    pip install flash-attn-4
+
+Pin `flash-attn-4==4.0.0b9` to match this benchmark. The `[cu13]`
+extra pulls in `nvidia-cutlass-dsl` built for CUDA 13.
+
+`qwen3_flex_inference.py` runs on both Qwen3 and Llama-3.2 (the only
+arch-specific branch is Qwen3's per-head QK RMSNorm; the rest of the
+flex_attention + paged KV + CUDA graphs stack is identical). Pass
+`--model_name unsloth/Llama-3.2-3B-Instruct` to target Llama, along with
+`--chat_template native` to use Llama's shipped Instruct template instead
+of the Qwen3 GRPO template.
+
+`gemma4_flex_inference.py` extends the engine to `unsloth/gemma-4-E2B-it`.
+Gemma-4 is not a drop-in: its text backbone has KV-sharing layers
+(layers 15-34 consume full-sequence K/V produced by a store layer), two
+attention regimes (`full_attention` with `head_dim=512` / rope_theta=1e6
+and `sliding_attention` with `head_dim=256` / sliding_window=512),
+per-layer input embeddings, four norms per block with double residuals,
+and a final logit softcap. The new file keeps the shared helpers
+(`PagedKVCache`, `PageTable`, LoRA double-copy, drift verification)
+imported from `qwen3_flex_inference.py` and adds:
+
+- a KV-sharing sidecar dict, sized `[max_batch, n_kv, max_seq, head_dim]`
+  per store layer, populated at prefill and read by the paired shared
+  layers through eager SDPA (shared layers don't fit the paged-cache
+  block-mask shape);
+- dual RoPE precomputation — `rotary_emb(x, pos, layer_type)` called once
+  per unique layer type, indexed by `self.layer_type`;
+- a walker that threads `per_layer_inputs` from
+  `get_per_layer_inputs` + `project_per_layer_inputs` into each layer
+  and applies the `layer_scalar` multiply at layer end;
+- `tanh(logits / final_logit_softcapping) * final_logit_softcapping`
+  applied on the lm_head output.
+
+Requires `transformers>=5.5.0` for the `gemma4` module; if absent the
+script exits with a clear install hint. Gemma-4 head_dim=256 exceeds FA4
+on sm_100 (B200), so pass `--no-fa4_prefill` and small Triton blocks:
+
+```bash
+CUDA_VISIBLE_DEVICES=6 python scripts/benchmarks/gemma4_flex_inference.py \
+    --model_name unsloth/gemma-4-E2B-it \
+    --n_prompts 64 --max_new_tokens 512 --capture_cudagraph \
+    --no-fa4_prefill \
+    --prefill_kernel_options '{"FORCE_USE_FLEX_ATTENTION": true, "BLOCK_M": 32, "BLOCK_N": 32}' \
+    --decode_kernel_options '{"BLOCK_M": 16, "BLOCK_N": 16}' \
+    --stats_path logs/flex_gemma4_bf16.json
+```
+
+| GPU          | arch      | sm    | Auto FA4 | Triton flex_attention |
+|--------------|-----------|-------|----------|------------------------|
+| A100         | Ampere    | sm_80 | off (uses Triton) | Works |
+| H100 / H200  | Hopper    | sm_90 | on | Works |
+| RTX 50xx     | Blackwell | sm_120 | on | Works |
+| B200 / GB200 | Blackwell | sm_100 | on | Works |
+
+The transformers continuous-batching path's FA4 wiring (the
+`flash_attn_fa4_shim.py` monkey-patches and the
+`site-packages/flash_attn/__init__.py` namespace shim that makes FA4
+visible under the FA2 import name) is covered below under "Known
+integration notes".
+
+## Reproduce
+
+```bash
+pip install unsloth "transformers>=4.57" "trl>=0.25" peft vllm
+
+# Generation microbenchmark (32 prompts, 512 new tokens each)
+CUDA_VISIBLE_DEVICES=2 python scripts/benchmarks/cb_vs_vllm_generation.py \
+    --backend vllm --stats_path logs/vllm_gen.json \
+    --n_prompts 32 --n_rounds 2 --max_new_tokens 512 \
+    --gpu_memory_utilization 0.6
+
+# CB variants
+CUDA_VISIBLE_DEVICES=6 python scripts/benchmarks/cb_vs_vllm_generation.py \
+    --backend tpaged --attn_impl sdpa \
+    --stats_path logs/cb_gen_sdpa.json \
+    --n_prompts 32 --n_rounds 2 --max_new_tokens 512
+
+CUDA_VISIBLE_DEVICES=6 python scripts/benchmarks/cb_vs_vllm_generation.py \
+    --backend tpaged --attn_impl flash_attention_2 \
+    --stats_path logs/cb_gen_fa.json \
+    --n_prompts 32 --n_rounds 2 --max_new_tokens 512
+
+# Full GRPO training (20 steps)
+CUDA_VISIBLE_DEVICES=2 python scripts/benchmarks/qwen3_grpo_vllm.py \
+    --max_steps 20 --num_generations 2 --per_device_train_batch_size 2 \
+    --output_dir outputs/grpo_vllm --stats_path logs/vllm_stats.json \
+    --gpu_memory_utilization 0.6
+
+CUDA_VISIBLE_DEVICES=7 python scripts/benchmarks/qwen3_grpo_naive.py \
+    --max_steps 20 --num_generations 2 --per_device_train_batch_size 2 \
+    --output_dir outputs/grpo_naive --stats_path logs/naive_stats.json
+
+CUDA_VISIBLE_DEVICES=6 python scripts/benchmarks/qwen3_grpo_tpaged.py \
+    --max_steps 20 --num_generations 2 --per_device_train_batch_size 2 \
+    --attn_impl flash_attention_2 \
+    --output_dir outputs/grpo_tpaged_fa --stats_path logs/tpaged_stats_fa.json \
+    --max_batch_tokens 16384 --num_blocks 16384
+```
+
+`qwen3_grpo_naive.py` and `qwen3_grpo_tpaged.py` accept an optional
+`--compile_mode {default,reduce-overhead,max-autotune-no-cudagraphs}` flag.
+When set, `trainer.model.forward` (and `trainer.ref_model.forward`, if present)
+are wrapped with `torch.compile` after trainer construction. The vLLM driver
+has no such flag because vLLM owns its own inference graph. `--compile_dynamic`
+(default on) toggles dynamic-shape compilation.
+
+## Known integration notes for transformers continuous batching + TRL + Unsloth
+
+These are the sharp edges you hit going down the continuous-batching path and
+how `qwen3_grpo_tpaged.py` handles them:
+
+1. **`top_k=-1` is not a valid value for transformers.** vLLM treats `-1` as
+   "disabled", but `TopKLogitsWarper` raises
+   `ValueError: top_k has to be a strictly positive integer`. The script
+   rewrites `top_k=None` on the shared GRPOConfig before handing it to
+   `GRPOConfig(use_transformers_paged=True, ...)`.
+
+2. **`PagedAttentionCache` default upper bounds are extremely conservative.**
+   `_upper_bound_max_batch_tokens=256` and `_upper_bound_num_blocks=4096`
+   choke decode throughput. The script passes
+   `generation_kwargs={"max_batch_tokens": 16384, "num_blocks": 16384}` which
+   TRL forwards to `GenerationConfig`, and the CB manager reads them when
+   sizing the paged cache.
+
+3. **Unsloth's `Qwen3Attention_fast_forward` bypasses the functional
+   attention interface.** Calling `model.generate_batch` on an
+   Unsloth-patched Qwen3 model fails inside
+   `unsloth.utils.attention_dispatch.run_attention` because Unsloth routes
+   through its own dispatcher rather than reading
+   `config._attn_implementation`. The benchmark script works around this by
+   loading a vanilla HF Qwen3 with PEFT LoRA for the tpaged and naive paths.
+   This costs the Unsloth training kernels but keeps the comparison clean.
+   A proper upstream fix is to detect `config._attn_implementation` being
+   `flash_attention_2` / `sdpa_paged` / `eager_paged` and delegate to the
+   stock transformers forward.
+
+4. **TRL imports `GuidedDecodingParams` from `vllm.sampling_params`.**
+   Newer vLLM releases (>= 0.13) have moved or removed that symbol, so
+   `trl.trainer.grpo_trainer` fails to import on a fresh vLLM install even
+   if you are not using vLLM. `qwen3_grpo_tpaged.py` installs a minimal
+   shim before importing TRL.
+
+5. **`UnslothGRPOTrainer` calls `model.for_training()` /
+   `for_inference()`.** Importing `unsloth` replaces
+   `trl.GRPOTrainer` with `UnslothGRPOTrainer`, which assumes the model has
+   these hooks. A vanilla HF model does not, so
+   `qwen3_grpo_tpaged.py` does not `import unsloth` at all.
+
+## Why continuous batching is still slower than vLLM on this workload
+
+- `ContinuousBatchingManager` does not yet implement CUDA graphs
+  (`use_cuda_graph=True` raises `NotImplementedError`). vLLM captures 100+
+  mixed prefill-decode and decode graphs during warmup.
+- CB re-allocates a fresh `PagedAttentionCache` on every `generate_batch`
+  call. For GRPO that is once per step. `--persistent_cb` (via
+  `persistent_cb.py`) keeps the cache warm across steps.
+- FA4 is a CuTeDSL package: first call per shape pays a one-time
+  JIT-compile.
+- vLLM uses its own colocated attention + FlashInfer / TRTLLM kernels
+  tuned for decode, which currently outperform everything a generic CB
+  path can do.
+
+These are all upstream transformers issues, not Unsloth issues. The scripts
+in this directory are intentionally simple so they are easy to port into a
+future upstream fix.

scripts/benchmarks/cb_sync_driver.py

@@ -0,0 +1,357 @@
+"""Main-thread synchronous driver for `ContinuousBatchProcessor`.
+
+`ContinuousBatchingManager.start()` spawns a background thread that owns the
+decode loop. That thread conflicts with two things we want to enable here:
+
+1. `torch.compile(mode="reduce-overhead")` which uses `cudagraph_trees` and
+   requires main-thread TLS.
+2. Raw `torch.cuda.CUDAGraph` capture / replay on the decode forward.
+
+The manager's dead `warmup()` path suggests CB was supposed to grow CUDA
+graph support upstream but `init_continuous_batching` raises
+`NotImplementedError` on `use_cuda_graph=True`. This driver side-steps the
+whole manager thread, so the cudagraph integration point is now available.
+
+Key fixed-shape invariant: with `slice_inputs=False`, the full pre-allocated
+tensor buffers (input_ids, position_ids, cu_seq_lens_*, attention_mask,
+read_index / write_index) are returned as views of the same storage every
+step, so their shapes are constant across iterations. That is the
+precondition for graph replay / cudagraph_trees to be safe.
+
+Greedy sampling only (`do_sample=False`). `torch.multinomial` is not
+graph-friendly.
+
+Usage:
+    driver = SyncCBDriver(model, gen_config, CBSyncConfig(compile_mode="reduce-overhead"))
+    # Reuse across many rollouts (cache / compiled forward stay warm):
+    for batch in batches:
+        driver.add_requests(batch)
+        out = driver.drive_until_empty()
+"""
+
+from __future__ import annotations
+
+import queue
+import time
+from dataclasses import dataclass, field
+from typing import Optional
+
+import torch
+from transformers.generation.configuration_utils import GenerationConfig
+from transformers.generation.continuous_batching import (
+    PagedAttentionCache,
+    RequestStatus,
+)
+from transformers.generation.continuous_batching.continuous_api import (
+    ContinuousBatchProcessor,
+    ContinuousBatchingManager,
+)
+from transformers.generation.continuous_batching.scheduler import FIFOScheduler
+
+
+@dataclass
+class CBSyncConfig:
+    """Tunables for the sync driver."""
+
+    max_new_tokens: int = 512
+    # torch.compile mode for the model forward. None = eager.
+    # "reduce-overhead" triggers cudagraph_trees, which captures a CUDA
+    # graph per unique input shape and replays it afterwards.
+    compile_mode: Optional[str] = None
+    do_sample: bool = False  # greedy only (graph-safe)
+    eos_token_id: Optional[int] = None
+    pad_token_id: Optional[int] = None
+    # `slice_inputs=False` would give fixed shapes every step (graph-friendly)
+    # but forces every decode step to do `max_batch_tokens` tokens of work,
+    # which at max_batch_tokens=8192 is ~256x more than the real decode batch.
+    # Prefer `slice_inputs=True` (natural shapes) and let torch.compile bucket
+    # per shape. Steady-state decode has one shape so most steps replay the
+    # same graph anyway.
+    slice_inputs: bool = True
+    # Paged cache upper bounds.
+    max_batch_tokens: int = 8192
+    num_blocks: int = 8192
+    # `torch._dynamo.config.cache_size_limit`: raise when varying shapes.
+    dynamo_cache_size_limit: int = 256
+    on_step: Optional[callable] = field(default = None)
+
+
+class SyncCBDriver:
+    """Main-thread driver that owns the PagedAttentionCache,
+    ContinuousBatchProcessor, and optionally a `torch.compile`-compiled
+    forward. Reusable across multiple `drive_until_empty` calls -- the cache
+    and compiled forward stay warm between rounds.
+    """
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        generation_config: GenerationConfig,
+        cfg: CBSyncConfig,
+    ):
+        self.model = model.eval()
+        self.cfg = cfg
+        gc = GenerationConfig.from_dict(generation_config.to_dict())
+        gc.do_sample = cfg.do_sample
+        if cfg.max_new_tokens:
+            gc.max_new_tokens = cfg.max_new_tokens
+        if cfg.eos_token_id is not None:
+            gc.eos_token_id = cfg.eos_token_id
+        if cfg.pad_token_id is not None:
+            gc.pad_token_id = cfg.pad_token_id
+        gc.max_batch_tokens = cfg.max_batch_tokens
+        gc.num_blocks = cfg.num_blocks
+        self.generation_config = gc
+
+        self.manager = ContinuousBatchingManager(
+            model = self.model,
+            generation_config = gc,
+            manual_eviction = False,
+            streaming = False,
+            slice_inputs = cfg.slice_inputs,
+        )
+
+        self.cache = PagedAttentionCache(
+            self.model.config,
+            gc,
+            self.model.device,
+            self.model.dtype,
+            tp_size = getattr(self.model, "_tp_size", None),
+        )
+        self.batch_processor = ContinuousBatchProcessor(
+            self.cache,
+            self.model.config,
+            gc,
+            self.manager.input_queue,
+            self.manager.output_queue,
+            self.manager.stop_event,
+            self.model.device,
+            self.model.dtype,
+            FIFOScheduler(self.cache),
+            streaming = False,
+            manual_eviction = False,
+            slice_inputs = cfg.slice_inputs,
+        )
+        self.manager.batch_processor = self.batch_processor
+
+        # torch.compile on the model forward. With slice_inputs=True the
+        # forward sees varying shapes (prefill bursts + decode steady state);
+        # `dynamic=True` lets Inductor bucket per shape without re-tracing
+        # every call, and `mode="reduce-overhead"` wraps each bucket in a
+        # CUDA graph replay path.
+        if cfg.compile_mode:
+            import torch._dynamo
+
+            torch._dynamo.config.cache_size_limit = cfg.dynamo_cache_size_limit
+            # GRPO's `requires_grad_` issue doesn't apply here (eval mode).
+            try:
+                torch._dynamo.config.allow_unspec_int_on_nn_module = True
+            except AttributeError:
+                pass
+            print(
+                f"[cb_sync] torch.compile(model, mode='{cfg.compile_mode}', "
+                f"dynamic=True)"
+            )
+            self.model.forward = torch.compile(
+                self.model.forward,
+                mode = cfg.compile_mode,
+                dynamic = True,
+                fullgraph = False,
+            )
+        self._step_count = 0
+
+    def add_requests(self, prompt_ids_list: list[list[int]]) -> list[str]:
+        return [self.manager.add_request(ids) for ids in prompt_ids_list]
+
+    def drive_until_empty(self) -> dict[str, list[int]]:
+        """Run the decode loop until every request finishes. Returns a dict
+        {request_id: generated_token_ids}.
+
+        Reusable across calls on the same driver -- the paged cache and the
+        compiled forward stay warm.
+        """
+        results: dict[str, list[int]] = {}
+        while True:
+            if (
+                self.manager.input_queue.empty()
+                and not self.batch_processor.has_pending_requests()
+            ):
+                break
+            if not self.batch_processor.prepare_next_batch():
+                break
+            # With `reduce-overhead`, Inductor captures a CUDA graph on the
+            # first call with a given shape signature and replays it after.
+            # Our shapes are constant (slice_inputs=False), so the second call
+            # is already replaying. No per-step torch.cuda.synchronize() --
+            # the replay itself fences appropriately.
+            self.manager._generation_step(self.batch_processor)
+            self.batch_processor.update_batch()
+            self._step_count += 1
+            if self.cfg.on_step is not None:
+                self.cfg.on_step(self._step_count, 0)
+            # Drain output_queue as requests finish.
+            while True:
+                try:
+                    out = self.manager.output_queue.get_nowait()
+                except queue.Empty:
+                    break
+                if out.status == RequestStatus.FINISHED:
+                    results[out.request_id] = out.generated_tokens
+        while True:
+            try:
+                out = self.manager.output_queue.get_nowait()
+            except queue.Empty:
+                break
+            if out.status == RequestStatus.FINISHED:
+                results[out.request_id] = out.generated_tokens
+        return results
+
+    def close(self):
+        self.cache = None
+        self.batch_processor = None
+        self.manager.batch_processor = None
+
+
+# Simple microbench harness so the file is runnable standalone.
+if __name__ == "__main__":
+    import argparse
+    import json
+    import os
+    import sys
+    from pathlib import Path
+
+    HERE = Path(__file__).resolve().parent
+    sys.path.insert(0, str(HERE))
+
+    import flash_attn_fa4_shim  # noqa: E402
+
+    flash_attn_fa4_shim.apply()
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_name", default = "unsloth/Qwen3-4B-Base")
+    parser.add_argument("--n_prompts", type = int, default = 32)
+    parser.add_argument("--n_rounds", type = int, default = 2)
+    parser.add_argument("--max_new_tokens", type = int, default = 512)
+    parser.add_argument("--attn_impl", default = "paged_attention")
+    parser.add_argument(
+        "--compile_mode",
+        default = None,
+        choices = [
+            None,
+            "default",
+            "reduce-overhead",
+            "max-autotune",
+            "max-autotune-no-cudagraphs",
+        ],
+    )
+    parser.add_argument("--max_batch_tokens", type = int, default = 8192)
+    parser.add_argument("--num_blocks", type = int, default = 8192)
+    parser.add_argument("--lora_adapter", default = None)
+    parser.add_argument("--stats_path", required = True)
+    args = parser.parse_args()
+
+    from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
+
+    tok = AutoTokenizer.from_pretrained(args.model_name)
+    if tok.pad_token is None:
+        tok.pad_token = tok.eos_token
+    model = AutoModelForCausalLM.from_pretrained(
+        args.model_name,
+        dtype = torch.bfloat16,
+        attn_implementation = args.attn_impl,
+    ).to("cuda")
+    model.eval()
+
+    if args.lora_adapter:
+        from peft import PeftModel
+
+        model = PeftModel.from_pretrained(
+            model, str(Path(args.lora_adapter).resolve()), is_trainable = False
+        )
+        model.eval()
+
+    from unsloth_grpo_common import (
+        SYSTEM_PROMPT,
+        apply_chat_template_to_tokenizer,
+    )
+    from datasets import load_dataset
+
+    apply_chat_template_to_tokenizer(tok)
+    ds = load_dataset("open-r1/DAPO-Math-17k-Processed", "en", split = "train")
+    ds = ds.shuffle(seed = 3407).select(range(args.n_prompts))
+    messages = [
+        [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": x["prompt"]},
+        ]
+        for x in ds
+    ]
+    prompt_ids = [
+        tok.apply_chat_template(m, add_generation_prompt = True, tokenize = True)
+        for m in messages
+    ]
+
+    gc_cfg = GenerationConfig(
+        max_new_tokens = args.max_new_tokens,
+        do_sample = False,
+        pad_token_id = tok.pad_token_id,
+        bos_token_id = tok.bos_token_id,
+        eos_token_id = tok.eos_token_id,
+        use_cache = True,
+    )
+
+    cfg = CBSyncConfig(
+        max_new_tokens = args.max_new_tokens,
+        compile_mode = args.compile_mode,
+        max_batch_tokens = args.max_batch_tokens,
+        num_blocks = args.num_blocks,
+        eos_token_id = tok.eos_token_id,
+        pad_token_id = tok.pad_token_id or tok.eos_token_id,
+    )
+
+    torch.cuda.reset_peak_memory_stats()
+
+    # One driver, multiple rounds -- cache + compiled forward stay warm.
+    driver = SyncCBDriver(model, gc_cfg, cfg)
+
+    # Warmup (first 16 prompts). With compile, this amortizes the capture.
+    print("[cb_sync] warmup...")
+    driver.add_requests(prompt_ids[:16])
+    _ = driver.drive_until_empty()
+    torch.cuda.synchronize()
+
+    wall_times = []
+    total_decoded = 0
+    for r in range(args.n_rounds):
+        torch.cuda.synchronize()
+        t0 = time.perf_counter()
+        driver.add_requests(prompt_ids)
+        results = driver.drive_until_empty()
+        torch.cuda.synchronize()
+        wall_times.append(time.perf_counter() - t0)
+        total_decoded = sum(len(v) for v in results.values())
+        print(
+            f"[cb_sync] round {r}: {wall_times[-1]:.2f}s, {total_decoded} tokens, "
+            f"{total_decoded / wall_times[-1]:.1f} tok/s"
+        )
+
+    med = sorted(wall_times)[len(wall_times) // 2]
+    out = {
+        "backend": "cb_sync_driver",
+        "compile_mode": args.compile_mode,
+        "attn_impl": args.attn_impl,
+        "lora_adapter": args.lora_adapter,
+        "n_prompts": args.n_prompts,
+        "n_decoded_tokens": total_decoded,
+        "wall_times_s": wall_times,
+        "median_wall_s": med,
+        "decode_tps": total_decoded / med if med else 0,
+        "max_new_tokens": args.max_new_tokens,
+        "peak_memory_gb": torch.cuda.max_memory_allocated() / 1024**3,
+    }
+    os.makedirs(os.path.dirname(os.path.abspath(args.stats_path)) or ".", exist_ok = True)
+    with open(args.stats_path, "w") as f:
+        json.dump(out, f, indent = 2)
+    print(json.dumps(out, indent = 2))
+    driver.close()
+    os._exit(0)

scripts/benchmarks/cb_vs_vllm_generation.py

@@ -0,0 +1,535 @@
+"""Standalone generation microbenchmark: vLLM vs transformers CB vs Unsloth.
+
+Backends (one per process; all engines are GPU-greedy):
+- `vllm`             : Unsloth `fast_inference=True` (vLLM colocated).
+- `tpaged`           : `model.generate_batch` on paged HF + `--attn_impl`.
+- `unsloth_fi_false` : Unsloth `fast_inference=False` with the custom HF
+                       inference kernels (cached fp16 LoRA).
+
+LoRA: pass `--lora_adapter PATH` to activate a PEFT-style rank-32 adapter on
+both vLLM (`LoRARequest`) and the HF paths (`peft.PeftModel.from_pretrained`,
+or for `unsloth_fi_false` `FastLanguageModel.get_peft_model` pointed at the
+same weights).
+
+Equivalence-friendly sampling defaults (`--temperature 0.1 --top_p 0.97
+--min_p 0.5 --top_k 5`) keep rollouts comparable across backends for the KL /
+reward diff checks done in Phase 2.
+
+Usage:
+    CUDA_VISIBLE_DEVICES=6 python scripts/benchmarks/cb_vs_vllm_generation.py \
+        --backend vllm --stats_path logs/lora_vllm_gen.json \
+        --lora_adapter outputs/lora_rank32_fresh
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+import time
+from pathlib import Path
+
+HERE = Path(__file__).resolve().parent
+sys.path.insert(0, str(HERE))
+
+import torch  # noqa: E402
+
+# FA4 shim for the `tpaged` backend with paged_attention. No-op for vLLM /
+# unsloth_fi_false.
+import flash_attn_fa4_shim  # noqa: E402
+
+flash_attn_fa4_shim.apply()
+
+
+def build_prompts(tokenizer, n_prompts, chat_template = "auto", model_type_name = None):
+    from unsloth_grpo_common import (
+        apply_chat_template_to_tokenizer,
+        SYSTEM_PROMPT,
+    )
+    from datasets import load_dataset
+
+    if chat_template == "auto":
+        use_grpo = (model_type_name or "").startswith("Qwen3")
+    elif chat_template == "grpo":
+        use_grpo = True
+    else:  # "native"
+        use_grpo = False
+    if use_grpo:
+        apply_chat_template_to_tokenizer(tokenizer)
+        print("[bench] chat_template: GRPO")
+    else:
+        print("[bench] chat_template: tokenizer native")
+    ds = load_dataset("open-r1/DAPO-Math-17k-Processed", "en", split = "train")
+    ds = ds.shuffle(seed = 3407).select(range(n_prompts))
+    messages = [
+        [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": x["prompt"]},
+        ]
+        for x in ds
+    ]
+    prompts_text = [
+        tokenizer.apply_chat_template(m, add_generation_prompt = True, tokenize = False)
+        for m in messages
+    ]
+    prompt_ids = [
+        tokenizer.apply_chat_template(m, add_generation_prompt = True, tokenize = True)
+        for m in messages
+    ]
+    return prompts_text, prompt_ids
+
+
+def run_vllm(args):
+    os.environ.setdefault("UNSLOTH_VLLM_STANDBY", "1")
+    from unsloth import FastLanguageModel
+
+    fi_kwargs = dict(
+        model_name = args.model_name,
+        max_seq_length = args.max_seq_length,
+        load_in_4bit = args.load_in_4bit,
+        fast_inference = True,
+        max_lora_rank = 32,
+        gpu_memory_utilization = args.gpu_memory_utilization,
+    )
+    if args.enforce_eager:
+        # vLLM 0.19 + torch 2.10 hits `RuntimeError: Tried to erase Node
+        # size_1 but it still had 2 users` during split_graph. enforce_eager
+        # skips vLLM's torch.compile path entirely (still PagedAttention +
+        # FlashInfer decode, just no graph capture).
+        fi_kwargs["enforce_eager"] = True
+    model, tokenizer = FastLanguageModel.from_pretrained(**fi_kwargs)
+    prompts_text, prompt_ids = build_prompts(
+        tokenizer,
+        args.n_prompts,
+        chat_template = args.chat_template,
+        model_type_name = type(getattr(model, "model", model)).__name__,
+    )
+
+    lora_request = None
+    if args.lora_adapter:
+        from vllm.lora.request import LoRARequest
+
+        lora_request = LoRARequest("fresh", 1, str(Path(args.lora_adapter).resolve()))
+
+    from vllm import SamplingParams
+
+    sp = SamplingParams(
+        temperature = args.temperature,
+        top_p = args.top_p,
+        min_p = args.min_p,
+        top_k = args.top_k,
+        seed = 3407,
+        max_tokens = args.max_new_tokens,
+        stop = [tokenizer.eos_token],
+        include_stop_str_in_output = True,
+    )
+
+    # Warmup on 16 prompts then discard.
+    warmup_text = prompts_text[:16]
+    _ = model.fast_generate(warmup_text, sampling_params = sp, lora_request = lora_request)
+    torch.cuda.synchronize()
+
+    n_prompt_tokens = sum(len(p) for p in prompt_ids)
+    wall_times = []
+    total_decoded = None
+    last_outputs = None
+    for _ in range(args.n_rounds):
+        torch.cuda.synchronize()
+        t0 = time.perf_counter()
+        outputs = model.fast_generate(
+            prompts_text, sampling_params = sp, lora_request = lora_request
+        )
+        torch.cuda.synchronize()
+        wall_times.append(time.perf_counter() - t0)
+        total_decoded = sum(len(o.outputs[0].token_ids) for o in outputs)
+        last_outputs = outputs
+
+    med = sorted(wall_times)[len(wall_times) // 2]
+    sample_texts = (
+        [o.outputs[0].text[:200] for o in (last_outputs[:3] or [])]
+        if last_outputs
+        else []
+    )
+    return {
+        "backend": "vllm",
+        "lora_adapter": args.lora_adapter,
+        "n_prompts": args.n_prompts,
+        "n_prompt_tokens": n_prompt_tokens,
+        "n_decoded_tokens": total_decoded,
+        "wall_times_s": wall_times,
+        "median_wall_s": med,
+        "prompt_tps": n_prompt_tokens / med,
+        "decode_tps": (total_decoded or 0) / med,
+        "max_new_tokens": args.max_new_tokens,
+        "sample_completions": sample_texts,
+    }
+
+
+def run_tpaged(args):
+    """Vanilla HF + paged cache.
+
+    Unsloth's Qwen3Attention monkey-patch does not compose with the
+    `paged|<impl>` functional attention interface, so we use plain HF.
+    """
+    from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    if args.load_in_4bit:
+        bnb_model_name = args.model_name_4bit or f"{args.model_name}-unsloth-bnb-4bit"
+        print(f"[tpaged] loading 4-bit base: {bnb_model_name}")
+        model = AutoModelForCausalLM.from_pretrained(
+            bnb_model_name,
+            attn_implementation = args.attn_impl,
+            device_map = "cuda:0",
+        )
+        # HF transformers logs "lm_head.weight newly initialized" for
+        # bnb-4bit shards of tied-embedding models. tie_word_embeddings is
+        # True in the config but the dequant path leaves lm_head unbound.
+        # Tie manually so we don't generate gibberish.
+        if getattr(model.config, "tie_word_embeddings", False):
+            model.lm_head.weight = model.model.embed_tokens.weight
+    else:
+        model = AutoModelForCausalLM.from_pretrained(
+            args.model_name,
+            dtype = torch.bfloat16,
+            attn_implementation = args.attn_impl,
+        ).to("cuda")
+    model.eval()
+
+    if args.lora_adapter:
+        from peft import PeftModel
+
+        # NOTE: no merge_adapter -- we measure LoRA-active inference.
+        model = PeftModel.from_pretrained(
+            model, str(Path(args.lora_adapter).resolve()), is_trainable = False
+        )
+        model.eval()
+
+    if args.persistent_cb:
+        from persistent_cb import install_for_model  # noqa: WPS433
+    prompts_text, prompt_ids = build_prompts(
+        tokenizer,
+        args.n_prompts,
+        chat_template = args.chat_template,
+        model_type_name = type(model).__name__,
+    )
+
+    gen_config = GenerationConfig(
+        max_new_tokens = args.max_new_tokens,
+        do_sample = True,
+        temperature = args.temperature,
+        top_p = args.top_p,
+        min_p = args.min_p,
+        top_k = args.top_k,
+        pad_token_id = tokenizer.pad_token_id or tokenizer.eos_token_id,
+        bos_token_id = tokenizer.bos_token_id,
+        eos_token_id = tokenizer.eos_token_id,
+        use_cache = True,
+    )
+    gen_config.max_batch_tokens = args.max_batch_tokens
+    gen_config.num_blocks = args.num_blocks
+
+    if args.persistent_cb:
+        install_for_model(model, gen_config)
+
+    warmup_ids = prompt_ids[:16]
+    with torch.inference_mode():
+        _ = model.generate_batch(
+            warmup_ids, generation_config = gen_config, progress_bar = False
+        )
+    torch.cuda.synchronize()
+
+    n_prompt_tokens = sum(len(p) for p in prompt_ids)
+    wall_times = []
+    total_decoded = None
+    last_outputs = None
+    for _ in range(args.n_rounds):
+        torch.cuda.synchronize()
+        t0 = time.perf_counter()
+        with torch.inference_mode():
+            outputs = model.generate_batch(
+                prompt_ids, generation_config = gen_config, progress_bar = False
+            )
+        torch.cuda.synchronize()
+        wall_times.append(time.perf_counter() - t0)
+        total_decoded = sum(len(v.generated_tokens) for v in outputs.values())
+        last_outputs = outputs
+
+    # Sample completions for coherence sanity check.
+    sample_texts = []
+    if last_outputs is not None:
+        for k in list(last_outputs.keys())[:3]:
+            toks = last_outputs[k].generated_tokens
+            sample_texts.append(tokenizer.decode(toks, skip_special_tokens = False)[:200])
+
+    med = sorted(wall_times)[len(wall_times) // 2]
+    return {
+        "backend": "tpaged",
+        "lora_adapter": args.lora_adapter,
+        "attn_impl": args.attn_impl,
+        "persistent_cb": args.persistent_cb,
+        "n_prompts": args.n_prompts,
+        "n_prompt_tokens": n_prompt_tokens,
+        "n_decoded_tokens": total_decoded,
+        "wall_times_s": wall_times,
+        "median_wall_s": med,
+        "prompt_tps": n_prompt_tokens / med,
+        "decode_tps": (total_decoded or 0) / med,
+        "max_new_tokens": args.max_new_tokens,
+        "sample_completions": sample_texts,
+    }
+
+
+def run_unsloth_fi_false(args):
+    """Unsloth `fast_inference=False` path with custom HF inference kernels.
+
+    This is the path that backs regular Unsloth training's sampling loop
+    (Triton RMSNorm/RoPE, cached fp16 LoRA copies in `fast_linear_forward`).
+    Previously only exercised through full GRPO runs -- isolating it lets us
+    compare it head-to-head against vLLM on the same workload.
+
+    LoRA is attached via `FastLanguageModel.get_peft_model`; if a PEFT adapter
+    path is provided we re-load its weights into the Unsloth-wrapped model so
+    every backend uses the *same* weights.
+    """
+    os.environ.setdefault("UNSLOTH_VLLM_STANDBY", "1")
+    from unsloth import FastLanguageModel
+
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name = args.model_name,
+        max_seq_length = args.max_seq_length,
+        load_in_4bit = False,
+        fast_inference = False,
+        max_lora_rank = 32,
+    )
+    # Attach LoRA rank 32 the same way the GRPO notebook does.
+    model = FastLanguageModel.get_peft_model(
+        model,
+        r = 32,
+        target_modules = [
+            "q_proj",
+            "k_proj",
+            "v_proj",
+            "o_proj",
+            "gate_proj",
+            "up_proj",
+            "down_proj",
+        ],
+        lora_alpha = 64,
+        use_gradient_checkpointing = "unsloth",
+        random_state = 3407,
+    )
+
+    # Optional: overlay a shared adapter so weights match other backends.
+    if args.lora_adapter:
+        from safetensors import safe_open
+
+        adapter_file = Path(args.lora_adapter).resolve() / "adapter_model.safetensors"
+        loaded_tensors = {}
+        with safe_open(str(adapter_file), framework = "pt") as f:
+            for key in f.keys():
+                loaded_tensors[key] = f.get_tensor(key)
+
+        # Both PEFT and Unsloth's `get_peft_model` produce parameter names with
+        # `base_model.model.` prefix plus `.lora_{A,B}.default.weight`. Build a
+        # normalized (core-path) -> param map, then match by core path only.
+        def _core(name: str) -> str:
+            n = name
+            for pref in ("base_model.model.", "model."):
+                if n.startswith(pref):
+                    n = n[len(pref) :]
+            n = n.replace(".lora_A.default.", ".lora_A.").replace(
+                ".lora_B.default.", ".lora_B."
+            )
+            return n
+
+        own_by_core = {}
+        for n, p in model.named_parameters():
+            if "lora_" in n:
+                own_by_core.setdefault(_core(n), []).append(p)
+        matched = 0
+        with torch.no_grad():
+            for name, tensor in loaded_tensors.items():
+                core = _core(name)
+                for own in own_by_core.get(core, []):
+                    if own.shape == tensor.shape:
+                        own.data.copy_(tensor.to(own.device, own.dtype))
+                        matched += 1
+                        break
+        print(
+            f"[unsloth_fi_false] LoRA weight sync matched {matched} tensors "
+            f"(out of {len(loaded_tensors)} adapter entries)."
+        )
+
+    FastLanguageModel.for_inference(model)
+
+    prompts_text, prompt_ids = build_prompts(
+        tokenizer,
+        args.n_prompts,
+        chat_template = args.chat_template,
+        model_type_name = type(model).__name__,
+    )
+
+    # `model.generate` accepts batched input_ids; pad to max length.
+    from transformers import GenerationConfig
+
+    if tokenizer.padding_side != "left":
+        tokenizer.padding_side = "left"  # decoder needs left padding
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    gen_config = GenerationConfig(
+        max_new_tokens = args.max_new_tokens,
+        do_sample = True,
+        temperature = args.temperature,
+        top_p = args.top_p,
+        min_p = args.min_p,
+        top_k = args.top_k,
+        pad_token_id = tokenizer.pad_token_id,
+        bos_token_id = tokenizer.bos_token_id,
+        eos_token_id = tokenizer.eos_token_id,
+        use_cache = True,
+    )
+
+    def _batched_generate(texts):
+        batch = tokenizer(texts, return_tensors = "pt", padding = True).to("cuda")
+        with torch.inference_mode():
+            out = model.generate(**batch, generation_config = gen_config)
+        prompt_len = batch["input_ids"].shape[1]
+        return out, prompt_len
+
+    # Warmup on 16 prompts.
+    _ = _batched_generate(prompts_text[:16])
+    torch.cuda.synchronize()
+
+    n_prompt_tokens = sum(len(p) for p in prompt_ids)
+    wall_times = []
+    total_decoded = None
+    last_out_ids = None
+    last_prompt_len = None
+    for _ in range(args.n_rounds):
+        torch.cuda.synchronize()
+        t0 = time.perf_counter()
+        out_ids, prompt_len = _batched_generate(prompts_text)
+        torch.cuda.synchronize()
+        wall_times.append(time.perf_counter() - t0)
+        # Count generated tokens past prompt_len per sequence (subtract any
+        # trailing pad-only tail by comparing against EOS).
+        total_decoded = int(
+            (out_ids[:, prompt_len:] != tokenizer.pad_token_id).sum().item()
+        )
+        last_out_ids = out_ids
+        last_prompt_len = prompt_len
+
+    med = sorted(wall_times)[len(wall_times) // 2]
+    sample_texts = []
+    if last_out_ids is not None:
+        for i in range(min(3, last_out_ids.shape[0])):
+            sample_texts.append(
+                tokenizer.decode(
+                    last_out_ids[i, last_prompt_len:], skip_special_tokens = False
+                )[:200]
+            )
+
+    return {
+        "backend": "unsloth_fi_false",
+        "lora_adapter": args.lora_adapter,
+        "n_prompts": args.n_prompts,
+        "n_prompt_tokens": n_prompt_tokens,
+        "n_decoded_tokens": total_decoded,
+        "wall_times_s": wall_times,
+        "median_wall_s": med,
+        "prompt_tps": n_prompt_tokens / med,
+        "decode_tps": (total_decoded or 0) / med,
+        "max_new_tokens": args.max_new_tokens,
+        "sample_completions": sample_texts,
+    }
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument(
+        "--backend", choices = ["vllm", "tpaged", "unsloth_fi_false"], required = True
+    )
+    p.add_argument("--model_name", default = "unsloth/Qwen3-4B-Base")
+    p.add_argument("--max_seq_length", type = int, default = 2048)
+    p.add_argument("--n_prompts", type = int, default = 32)
+    p.add_argument("--n_rounds", type = int, default = 2)
+    p.add_argument("--max_new_tokens", type = int, default = 512)
+    p.add_argument("--gpu_memory_utilization", type = float, default = 0.8)
+    p.add_argument("--attn_impl", default = "sdpa")
+    p.add_argument("--max_batch_tokens", type = int, default = 8192)
+    p.add_argument("--num_blocks", type = int, default = 16384)
+    p.add_argument("--persistent_cb", action = "store_true")
+    p.add_argument(
+        "--lora_adapter",
+        default = None,
+        help = "Path to a PEFT adapter (rank 32) applied in every backend.",
+    )
+    p.add_argument(
+        "--load_in_4bit",
+        action = "store_true",
+        help = "Load base as bitsandbytes 4-bit (Unsloth shard).",
+    )
+    p.add_argument(
+        "--model_name_4bit",
+        default = None,
+        help = "Override 4-bit shard name. Default `{model_name}-unsloth-bnb-4bit`.",
+    )
+    p.add_argument("--temperature", type = float, default = 0.1)
+    p.add_argument("--top_p", type = float, default = 0.97)
+    p.add_argument("--min_p", type = float, default = 0.5)
+    p.add_argument("--top_k", type = int, default = 5)
+    p.add_argument("--stats_path", required = True)
+    p.add_argument(
+        "--chat_template",
+        choices = ["auto", "grpo", "native"],
+        default = "auto",
+        help = (
+            "`auto`: GRPO for Qwen3, tokenizer native otherwise. "
+            "`grpo`: force GRPO template. `native`: force tokenizer's "
+            "built-in Instruct template (Llama-3.2-Instruct)."
+        ),
+    )
+    p.add_argument(
+        "--enforce_eager",
+        action = "store_true",
+        help = (
+            "vLLM only: skip the torch.compile + cudagraph path and run "
+            "eager. Useful when vLLM's compile regresses on the local "
+            "torch build."
+        ),
+    )
+    return p.parse_args()
+
+
+def main():
+    args = parse_args()
+    os.makedirs(os.path.dirname(os.path.abspath(args.stats_path)) or ".", exist_ok = True)
+
+    torch.cuda.reset_peak_memory_stats()
+    if args.backend == "vllm":
+        out = run_vllm(args)
+    elif args.backend == "unsloth_fi_false":
+        out = run_unsloth_fi_false(args)
+    else:
+        out = run_tpaged(args)
+
+    out["peak_memory_gb"] = torch.cuda.max_memory_allocated() / 1024**3
+    out["sampling"] = {
+        "temperature": args.temperature,
+        "top_p": args.top_p,
+        "min_p": args.min_p,
+        "top_k": args.top_k,
+    }
+    with open(args.stats_path, "w") as f:
+        json.dump(out, f, indent = 2)
+    print(json.dumps(out, indent = 2))
+    os._exit(0)
+
+
+if __name__ == "__main__":
+    main()

scripts/benchmarks/compare_grpo_runs.py

@@ -0,0 +1,80 @@
+"""Pairwise equivalence diff for GRPO backend runs.
+
+Runs `torch_debugging_utils.compare_training_runs` over the StatisticsCallback
+JSONs produced by `qwen3_grpo_unified.py`, plus reward / KL diffs which the
+base util doesn't track (it's loss/grad-focused).
+
+Usage:
+    python scripts/benchmarks/compare_grpo_runs.py \
+        --ref logs/grpo_vllm_30.json \
+        --candidate logs/grpo_unsloth_fi_false_30.json
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+HERE = Path(__file__).resolve().parent
+WORKSPACE_ROOT = Path("/mnt/disks/unslothai/ubuntu/workspace_31")
+for p in (HERE, WORKSPACE_ROOT):
+    sys.path.insert(0, str(p))
+
+
+def _arrays(path: str):
+    with open(path) as f:
+        logs = json.load(f)
+    return {
+        "loss": [l.get("loss") for l in logs if "loss" in l],
+        "reward": [l.get("reward") for l in logs if "reward" in l],
+        "kl": [l.get("kl") for l in logs if "kl" in l],
+        "grad_norm": [l.get("grad_norm") for l in logs if "grad_norm" in l],
+        "time_ms": [l.get("time_ms") for l in logs if "time_ms" in l],
+    }
+
+
+def _diff(a, b):
+    if not a or not b:
+        return None
+    n = min(len(a), len(b))
+    diffs = [
+        abs(a[i] - b[i]) for i in range(n) if a[i] is not None and b[i] is not None
+    ]
+    if not diffs:
+        return None
+    return {
+        "n_compared": len(diffs),
+        "max_abs": max(diffs),
+        "mean_abs": sum(diffs) / len(diffs),
+    }
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--ref", required = True)
+    p.add_argument("--candidate", required = True)
+    args = p.parse_args()
+
+    from torch_debugging_utils import compare_training_runs
+
+    base = compare_training_runs(args.ref, args.candidate, loss_tol = 1e-3, grad_tol = 1e-3)
+
+    ref_a = _arrays(args.ref)
+    cand_a = _arrays(args.candidate)
+    extras = {k: _diff(ref_a[k], cand_a[k]) for k in ("reward", "kl", "time_ms")}
+
+    out = {
+        "ref": args.ref,
+        "candidate": args.candidate,
+        "compare_training_runs": base,
+        "reward_diff": extras["reward"],
+        "kl_diff": extras["kl"],
+        "time_diff_ms": extras["time_ms"],
+    }
+    print(json.dumps(out, indent = 2))
+
+
+if __name__ == "__main__":
+    main()

scripts/benchmarks/flash_attn_fa4_shim.py

@@ -0,0 +1,96 @@
+"""Make transformers' continuous batching dispatch to Flash Attention 4 on B200.
+
+Two integration gaps between transformers' continuous batching (CB) and the
+FA2 varlen path make `attn_implementation="flash_attention_2"` fail out of the
+box in transformers 4.57 even with a working FA varlen kernel:
+
+1. CB's `ContinuousBatchProcessor` creates a 4D paged attention mask of shape
+   `[1, 1, q_len, k_len]` for every attention implementation except
+   `"paged_attention"`. `_flash_attention_forward` then enters the
+   `if attention_mask is not None:` branch and calls `_upad_input`, which
+   expects a 2D mask and fails (the FA4 kernel ultimately asserts on
+   `cu_seqlens_k.shape`).
+
+2. CB passes `max_seqlen_q`/`max_seqlen_k` as model kwargs while
+   `_flash_attention_forward` names the parameters `max_length_q`/
+   `max_length_k`. The former therefore never bind and the varlen branch
+   inside `_flash_attention_forward` invokes the FA kernel with
+   `max_seqlen_q=None`.
+
+This module patches around both, and (via the sibling
+`site-packages/flash_attn/__init__.py` shim) points the FA2 varlen dispatch
+at FA4's Blackwell-capable kernel. Call `apply()` once, before
+`model.generate_batch` is invoked.
+"""
+
+from __future__ import annotations
+
+import functools
+
+
+_APPLIED = False
+
+
+def apply() -> None:
+    global _APPLIED
+    if _APPLIED:
+        return
+
+    import transformers  # noqa: F401  - force load
+    from transformers.generation.continuous_batching import continuous_api as _cb
+    from transformers import modeling_flash_attention_utils as _fa_utils
+
+    _patch_return_attention_mask(_cb)
+    _patch_flash_attention_forward(_fa_utils)
+
+    _APPLIED = True
+
+
+def _patch_return_attention_mask(cb_module) -> None:
+    """Don't materialise a 4D attention mask when the kernel is FA varlen.
+
+    The existing `return_attention_mask` only skips the mask for
+    `"paged_attention"`. We extend the skip set to `"flash_attention_2"`
+    (and `"flash_attention_3"` for future-proofing) because the varlen path
+    relies on `cu_seq_lens_*` and reads no mask.
+    """
+    _SKIP_MASK_IMPLS = {
+        "paged_attention",
+        "flash_attention_2",
+        "flash_attention_3",
+    }
+
+    def return_attention_mask(self) -> bool:
+        return self.config._attn_implementation not in _SKIP_MASK_IMPLS
+
+    cb_module.ContinuousBatchProcessor.return_attention_mask = return_attention_mask
+
+
+def _patch_flash_attention_forward(fa_utils_module) -> None:
+    """Accept CB's `max_seqlen_q`/`max_seqlen_k` kwargs as aliases.
+
+    transformers names the parameters `max_length_q`/`max_length_k`, but CB
+    (and most downstream call sites) name them `max_seqlen_q`/
+    `max_seqlen_k`. We rename at the boundary so callers on either side work
+    unchanged.
+    """
+    original = fa_utils_module._flash_attention_forward
+
+    @functools.wraps(original)
+    def wrapper(*args, **kwargs):
+        if kwargs.get("max_length_q") is None and "max_seqlen_q" in kwargs:
+            kwargs["max_length_q"] = kwargs.pop("max_seqlen_q")
+        if kwargs.get("max_length_k") is None and "max_seqlen_k" in kwargs:
+            kwargs["max_length_k"] = kwargs.pop("max_seqlen_k")
+        return original(*args, **kwargs)
+
+    fa_utils_module._flash_attention_forward = wrapper
+
+    # Some integration modules imported the function by name before we
+    # patched. Re-bind the most common consumers so they pick up the wrapper.
+    try:
+        from transformers.integrations import flash_attention as _flash_integration
+
+        _flash_integration._flash_attention_forward = wrapper
+    except Exception:
+        pass

scripts/benchmarks/flex_autotune_replay.py

@@ -0,0 +1,221 @@
+"""Autotune replay for flex_attention decode.
+
+Pattern from attention-gym/examples/flex_autotune_replay.py:
+
+1. Run once with `mode="max-autotune-no-cudagraphs"` and
+   `TORCHINDUCTOR_FLEX_ATTENTION_LOGGING_FILE` set. Inductor writes a JSON
+   log of every kernel config it tried, sorted by wall time per shape.
+2. Parse the log for the decode-shape entry (Q_LEN small, large KV).
+3. Emit the best fwd_* options as a JSON string that the main flex script
+   can accept via --decode_kernel_options.
+
+Usage:
+    CUDA_VISIBLE_DEVICES=7 python scripts/benchmarks/flex_autotune_replay.py \
+        --log_file logs/flex_autotune.json \
+        --max_batch_size 64 \
+        --n_prompts 16 \
+        --max_new_tokens 64
+
+Writes best decode kernel options to --output_opts (JSON), prints to stdout.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+
+HERE = Path(__file__).resolve().parent
+
+
+def run_autotune_pass(log_file: str, args) -> None:
+    env = os.environ.copy()
+    env["FLEX_COMPILE_MODE"] = "max-autotune-no-cudagraphs"
+    # Inductor appends `.json` to this env var value.
+    env["TORCHINDUCTOR_FLEX_ATTENTION_LOGGING_FILE"] = log_file.replace(".json", "")
+
+    cmd = [
+        sys.executable,
+        "-u",
+        str(HERE / "qwen3_flex_inference.py"),
+        "--n_prompts",
+        str(args.n_prompts),
+        "--n_rounds",
+        "1",
+        "--max_new_tokens",
+        str(args.max_new_tokens),
+        "--max_batch_size",
+        str(args.max_batch_size),
+        # NB: autotune in `max-autotune-no-cudagraphs` mode is incompatible with
+        # our raw CUDA graph capture path, so we skip --capture_cudagraph here.
+        # Goal is only to produce the log, not to benchmark.
+        "--stats_path",
+        str(HERE / "logs" / "flex_autotune_stats.json"),
+    ]
+    if args.lora_adapter:
+        cmd += ["--lora_adapter", args.lora_adapter]
+    print("[autotune] running:", " ".join(cmd))
+    print(f"[autotune] logging to {log_file}")
+    subprocess.run(cmd, env = env, check = True)
+
+
+class _SymStub:
+    """Pretend-symbolic value so eval() can handle SymPy-ish free vars like `s40`."""
+
+    def __init__(self, name):
+        self.name = name
+
+    def __repr__(self):
+        return f"Sym({self.name})"
+
+
+class _SymNamespace(dict):
+    """Any unknown name becomes a _SymStub instead of NameError."""
+
+    def __getitem__(self, key):
+        if key in self:
+            return super().__getitem__(key)
+        # Don't catch obvious builtins.
+        if key in ("True", "False", "None"):
+            return eval(key)
+        return _SymStub(key)
+
+    def __contains__(self, key):
+        return True  # satisfies eval's name resolution
+
+
+def parse_log(log_file: str) -> list[tuple[tuple, dict]]:
+    """Return list of (shape_tuple, best_fwd_options_dict) per shape entry."""
+    if not Path(log_file).exists():
+        raise FileNotFoundError(
+            f"Inductor log file missing: {log_file}. "
+            f"Did `TORCHINDUCTOR_FLEX_ATTENTION_LOGGING_FILE` fire?"
+        )
+    with open(log_file) as f:
+        data = json.load(f)
+
+    ns = _SymNamespace()
+    shapes = []
+    for entry in data:
+        key, choices = next(iter(entry.items()))
+        try:
+            parsed = eval(key, {"__builtins__": {}}, ns)
+        except Exception:
+            parsed = (key,)
+        kernel_type = None
+        if isinstance(parsed, (list, tuple)) and len(parsed) > 0:
+            first = parsed[0]
+            if isinstance(first, str):
+                kernel_type = first
+        best = choices[0]
+        opts = {k: v for k, v in best.items() if k not in ("type", "time")}
+        shapes.append((parsed, opts, best.get("time"), kernel_type, key))
+    return shapes
+
+
+def pick_decode_shape(shapes):
+    """Pick the decode-shape entry.
+
+    Decode has Q_LEN=1. Prefill has Q_LEN large. The shape tuple is
+    `('forward', B, H_q, H_kv, Q_LEN, KV_LEN, D_q, D_v)` — so Q_LEN is at
+    index 4. When B/KV_LEN are symbolic (`s0`, `s40`), the raw key string
+    is the form `('forward', s40, 32, 8, 1, s0, 128, 128)`.
+    """
+    import re
+
+    def q_len_of(shape_key, parsed):
+        # If we successfully parsed and there's a real int at index 4, use it.
+        if (
+            isinstance(parsed, (list, tuple))
+            and len(parsed) > 4
+            and isinstance(parsed[4], int)
+        ):
+            return parsed[4]
+        # Else extract from the raw string form, which is always
+        # `('forward', <B>, 32, 8, <Q_LEN>, <KV_LEN>, 128, 128)`.
+        m = re.match(r"\('forward',\s*[^,]+,\s*[^,]+,\s*[^,]+,\s*(\d+)", shape_key)
+        if m:
+            return int(m.group(1))
+        return 10**9
+
+    # (parsed, opts, time, kernel_type) -> plus we need the raw key string.
+    # Pass shape_key via _SymNamespace too — actually we'll redo parse_log to
+    # include the raw key. Simpler: re-read the file.
+    return min(shapes, key = lambda s: q_len_of(s[4] if len(s) > 4 else "", s[0]))
+
+
+def format_best_opts(best_opts: dict) -> dict:
+    """Filter Inductor log keys to those acceptable to FlexKernelOptions as
+    fwd_* prefix."""
+    from torch.nn.attention.flex_attention import FlexKernelOptions
+
+    annotations = FlexKernelOptions.__annotations__
+    out = {}
+    for k, v in best_opts.items():
+        if k in annotations:
+            out[f"fwd_{k}"] = v
+    return out
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument(
+        "--log_file",
+        default = "logs/flex_autotune.json",
+        help = "Inductor writes the autotune log here. Will have .json appended.",
+    )
+    p.add_argument(
+        "--output_opts",
+        default = "logs/flex_best_decode_opts.json",
+        help = "Extracted best kernel options go here.",
+    )
+    p.add_argument("--n_prompts", type = int, default = 16)
+    p.add_argument("--max_batch_size", type = int, default = 64)
+    p.add_argument("--max_new_tokens", type = int, default = 64)
+    p.add_argument("--lora_adapter", default = None)
+    p.add_argument(
+        "--skip_autotune",
+        action = "store_true",
+        help = "Skip autotune pass and just parse existing log.",
+    )
+    args = p.parse_args()
+
+    log_file = args.log_file
+    if not log_file.endswith(".json"):
+        log_file = log_file + ".json"
+
+    if not args.skip_autotune:
+        run_autotune_pass(log_file, args)
+
+    shapes = parse_log(log_file)
+    print(f"[autotune] parsed {len(shapes)} shapes from log:")
+    for shape, opts, t, kt, key in shapes:
+        print(f"  kernel_type={kt!r}  shape={shape}  time={t!r}  opts={opts}")
+        print(f"    raw key: {key}")
+
+    if not shapes:
+        raise SystemExit("no shapes found in autotune log")
+
+    decode_shape, decode_opts, decode_time, _, _ = pick_decode_shape(shapes)
+    print("\n[autotune] selected decode-ish shape:", decode_shape)
+    print("[autotune] best decode options:", decode_opts, f"(time={decode_time!r})")
+
+    best = format_best_opts(decode_opts)
+    # Always add tuned knobs we already confirmed helpful.
+    best.setdefault("PRESCALE_QK", True)
+    best.setdefault("USE_TMA", True)
+    best.setdefault("BLOCKS_ARE_CONTIGUOUS", True)
+    print("\n[autotune] final decode kernel_options:", json.dumps(best, indent = 2))
+
+    Path(args.output_opts).parent.mkdir(parents = True, exist_ok = True)
+    with open(args.output_opts, "w") as f:
+        json.dump(best, f, indent = 2)
+    print(f"[autotune] wrote {args.output_opts}")
+
+
+if __name__ == "__main__":
+    main()

scripts/benchmarks/flex_paged_attention.py

@@ -0,0 +1,502 @@
+# Adapted from attention-gym
+# Original source: https://github.com/pytorch-labs/attention-gym
+# License: BSD 3-Clause (see THIRD_PARTY_LICENSES.md)
+# Copyright (c) 2023, Driss Guessous
+
+# the original implementation has some bugs and has some feature that lives outside of the PageTable class
+
+from typing import Optional
+import torch
+from torch import Tensor
+from torch.nn.attention.flex_attention import (
+    _identity,
+    _mask_mod_signature,
+    _score_mod_signature,
+    BlockMask,
+    noop_mask,
+    create_block_mask,
+)
+
+create_block_mask = torch.compile(create_block_mask, dynamic = True)
+
+
+def _cdiv(x: int | float | torch.Tensor, multiple: int | float | torch.Tensor):
+    return (x + multiple - 1) // multiple
+
+
+class PagedKVCache(torch.nn.Module):
+    def __init__(self, page_table, n_heads, head_dim, dtype):
+        super().__init__()
+        cache_shape = (1, n_heads, page_table.n_pages * page_table.page_size, head_dim)
+        self.register_buffer("k_cache", torch.zeros(cache_shape, dtype = dtype))
+        self.register_buffer("v_cache", torch.zeros(cache_shape, dtype = dtype))
+
+        self.page_table = page_table
+
+    def update(self, input_pos, k_val, v_val, batch_idx = None):
+        assert (
+            batch_idx is not None
+        ), "batch_idx is required for paged kv cache, are you using non-paged attention?"
+
+        if batch_idx.ndim == 1:
+            # batch_idx should be [B] (decode)
+            return self.page_table.assign(
+                batch_idx, input_pos, k_val, v_val, self.k_cache, self.v_cache
+            )
+        else:
+            assert batch_idx.ndim == 2, "batch_idx must be 1D or 2D"
+            # batch_idx should be [1, L] (batch prefill)
+            return self.page_table.assign_prefill_no_paging(
+                batch_idx, input_pos, k_val, v_val, self.k_cache, self.v_cache
+            )
+
+
+class PageTable:
+    """
+    PageTable is a modified version of PagedAttention from attention-gym.
+
+    PageTable improves it by:
+    - maintaining a cpu copy of the page table, to avoid device-to-host transfers
+    - support batch prefill
+    - fix the bug in the original code in mask_mod and score_mod by mapping physical batch index to logical batch index
+    - subsuming the free_batch_idx into the page table, so we don't need to maintain it separately
+    """
+
+    def __init__(
+        self,
+        n_pages: int,
+        page_size: int,
+        max_batch_size: int,
+        device: str = "cuda",
+    ):
+        self.n_pages = n_pages
+        self.page_size = page_size
+        self.max_batch_size = max_batch_size
+        self.device = device
+
+        # page table: [logical_batch_idx, logical_block_idx] -> physical_page_idx
+        self.page_table = -torch.ones(
+            (max_batch_size, self.n_pages), dtype = torch.int64, device = device
+        )
+        self.page_table[0, :] = (
+            0  # page 0 is reserved for simpler code in assign_prefill_no_paging
+        )
+        self.page_table_cpu = [[] for _ in range(max_batch_size)]
+
+        self.capacity = [
+            0 for _ in range(max_batch_size)
+        ]  # capacity: batch_idx -> number of pages allocated * page size
+        self.free_pages = list(
+            reversed(range(1, n_pages))
+        )  # page 0 is reserved for simpler code in assign_prefill_no_paging
+        self.free_batch_idx = list(
+            reversed(range(1, max_batch_size))
+        )  # batch_idx 0 is reserved for no-op
+
+        # [logical_batch_idx, physical_page_idx] -> logical_page_idx
+        self.physical_to_logical = -torch.ones(
+            (max_batch_size, n_pages), dtype = torch.int64, device = device
+        )
+
+    def can_reserve(self, size: int, batch_idx_int: int | None = None) -> bool:
+        """check if we can reserve new pages for an existing request or a new request, without gpu operations"""
+        if batch_idx_int is None:
+            # check if we can schedule a new request
+            return (
+                self.pages_available * self.page_size >= size
+                and len(self.free_batch_idx) > 0
+            )
+        else:
+            # check if we can reserve new pages for an existing request
+            return self.reserve(batch_idx_int, None, size, dry_run = True)
+
+    def allocate(self) -> int:
+        """allocate a new batch"""
+        batch_idx = self.free_batch_idx.pop()
+
+        self.capacity[batch_idx] = 0
+        self.physical_to_logical[batch_idx, :] = -1
+        self.page_table[batch_idx, :] = -1
+        return batch_idx
+
+    @property
+    def pages_available(self) -> int:
+        return len(self.free_pages)
+
+    def reserve(
+        self,
+        batch_idx_int: int,
+        batch_idx: torch.Tensor,
+        seq_len: int,
+        dry_run: bool = False,
+    ) -> bool:
+        """
+        Requests the capacity of a given batch to be at least enough to
+        hold `seq_len` elements.
+
+        Args:
+            batch_idx_int (int): batch index to be reserved;
+            batch_idx (Tensor): batch index to be reserved; shape :math:`(1)`.
+            seq_len (Tensor): minimum capacity for the given batch; shape :math:`(1)`.
+
+        Returns:
+            bool: True if the reservation was successful, False if the reservation was not successful (no space, and in this case, no update is done)
+        """
+
+        if seq_len <= self.capacity[batch_idx_int]:
+            return True
+
+        num_pages_to_allocate = _cdiv(
+            seq_len - self.capacity[batch_idx_int], self.page_size
+        )
+
+        can_allocate = num_pages_to_allocate <= self.pages_available
+        if dry_run:
+            return can_allocate
+
+        if not can_allocate:
+            raise RuntimeError(
+                f"Cannot reserve {num_pages_to_allocate} pages for a sequence of length {seq_len} "
+                f"in batch {batch_idx_int}. Only {self.pages_available} pages available. "
+                f"Current capacity is {self.capacity[batch_idx_int]} tokens."
+            )
+
+        start_page_idx = self.capacity[batch_idx_int] // self.page_size
+        end_page_idx = start_page_idx + num_pages_to_allocate
+
+        # find empty physical pages
+        allocated_pages_list = self.free_pages[-num_pages_to_allocate:]
+        allocated_pages = torch.tensor(allocated_pages_list, device = self.device)
+        # update page table
+        self.page_table[batch_idx, start_page_idx:end_page_idx] = allocated_pages
+
+        # update metadata
+        self.physical_to_logical[batch_idx, allocated_pages] = torch.arange(
+            start_page_idx,
+            end_page_idx,
+            device = self.device,
+        )
+        # update cpu side metadata
+        self.page_table_cpu[batch_idx_int] += allocated_pages_list
+        self.free_pages = self.free_pages[:-num_pages_to_allocate]
+        self.capacity[batch_idx_int] += num_pages_to_allocate * self.page_size
+        return True
+
+    def erase(self, batch_idx: int) -> None:
+        """
+        Removes a single batch from paged attention.
+
+        Args:
+            batch_idx (int): batch index to be removed;
+        """
+        # NOTE: the GPU side data will only be reset/overwritten when we allocate it for a new batch
+        self.free_batch_idx.append(batch_idx)
+        allocated_pages_cpu = self.page_table_cpu[batch_idx]
+        self.free_pages.extend(reversed(allocated_pages_cpu))
+        self.page_table_cpu[batch_idx] = []
+
+    def assign(
+        self,
+        batch_idx: torch.Tensor,
+        input_pos: torch.Tensor,
+        k_val: torch.Tensor,
+        v_val: torch.Tensor,
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+    ) -> None:
+        """
+        Assigns new contents `val` to the storage `cache` at the location
+        `batch_idx` and `input_pos`.
+
+        Args:
+            batch_idx (Tensor): batch index; shape :math:`(B)`.
+            input_pos (Tensor): input positions to be assigned for the given batch; shape :math:`(B, S)`.
+            val (Tensor): value to be assigned; shape :math:`(B, H, S, D)`
+            cache (Tensor): the cache to store the values; shape:`(1, H, MAX_S, D)`
+        """
+        if k_val.requires_grad:
+            raise RuntimeError("val must not require gradient")
+
+        B, H, S, K_D = k_val.shape
+        _, H_cache, MAX_S, D_cache = k_cache.shape
+        assert H_cache == H, "number of heads must match"
+        assert MAX_S >= S, "cache must have enough space"
+        assert D_cache == K_D, "hidden dim must match"
+        assert input_pos.shape == (B, S), "input_pos must have the same shape as val"
+        assert batch_idx.shape == (B,), "batch_idx must have one dimension only"
+
+        V_D = v_val.shape[3]
+        if B != batch_idx.shape[0]:
+            raise RuntimeError(
+                f"Expect val and batch_idx have the same batch size but got B={B} and B={batch_idx.shape[0]}."
+            )
+        if H != k_cache.shape[1]:
+            raise RuntimeError(
+                f"Expect val and cache has the same number of heads but got H={H} and H={k_cache.shape[1]}."
+            )
+        if S != input_pos.shape[1]:
+            raise RuntimeError(
+                f"Expect val and input_pos has the same length but got S={S} and S={input_pos.shape[0]}."
+            )
+        if K_D != k_cache.shape[3]:
+            raise RuntimeError(
+                f"Expect k_val and k_cache has the same hidden dim but got D={K_D} and D={k_cache.shape[3]}."
+            )
+        if V_D != v_cache.shape[3]:
+            raise RuntimeError(
+                f"Expect v_val and v_cache has the same hidden dim but got D={V_D} and D={v_cache.shape[3]}."
+            )
+
+        # find address
+        logical_block_idx = input_pos // self.page_size  # [B, S]
+        logical_block_offset = input_pos % self.page_size  # [B, S]
+
+        # NOTE: this code path is only used for decoding. For batch prefill, use assign_prefill_no_paging() instead
+        physical_block_idx = torch.gather(
+            self.page_table[batch_idx], 1, logical_block_idx.to(torch.int64)
+        ).to(torch.int32)  # [B, S]
+
+        addr = (physical_block_idx * self.page_size + logical_block_offset).view(
+            -1
+        )  # [B*S]
+
+        k_val = k_val.permute(1, 0, 2, 3).contiguous().view(1, H, B * S, K_D)
+        v_val = v_val.permute(1, 0, 2, 3).contiguous().view(1, H, B * S, V_D)
+
+        k_cache[:, :, addr, :] = k_val
+        v_cache[:, :, addr, :] = v_val
+
+        return k_cache, v_cache
+
+    def convert_logical_block_mask(
+        self,
+        block_mask: BlockMask,
+        batch_idx: Optional[torch.Tensor] = None,
+    ) -> BlockMask:
+        """
+        Converts a logical block mask by mapping its logical kv indices to the corresponding
+        physical kv indices.
+
+        Args:
+            block_mask (BlockMask): logical block mask;
+                kv_indices shape :math:`(B, H, ROWS, MAX_BLOCKS_IN_COL)`.
+            batch_idx (Tensor): batch index corresponding to the block_mask
+                batch dimension. This provides flexibility to convert a
+                block mask with smaller batch size than the page table;
+                shape :math:`(B)`.
+        """
+        B, H, ROWS, MAX_BLOCKS_IN_COL = block_mask.kv_indices.shape
+
+        if block_mask.BLOCK_SIZE[1] != self.page_size:
+            raise RuntimeError(
+                f"Expect block_mask has the same column block size as page_sizebut got size={block_mask.BLOCK_SIZE[1]} and size={self.page_size}"
+            )
+
+        device = block_mask.kv_num_blocks.device
+
+        if batch_idx is None:
+            batch_idx = torch.arange(B, device = device)
+
+        assert batch_idx.ndim == 1, "batch_idx must be a 1D tensor"
+        assert (
+            batch_idx.shape[0] == B
+        ), "batch_idx must have the same shape as block_mask"
+        assert (
+            B <= self.max_batch_size
+        ), "batch_idx must be less than or equal to max_batch_size"
+
+        page_table = self.page_table[batch_idx]
+
+        def transform(num_blocks, indices):
+            """
+            transform the block mask from [B, H, num_q_blocks, num_logical_kv_blocks]
+            to [B, H, num_q_blocks, num_physical_kv_blocks]
+
+            kv_num_blocks: [B, H, num_q_blocks] -> unchanged
+            kv_indices: [B, H, num_q_blocks, num_logical_kv_blocks] -> [B, H, num_q_blocks, num_physical_kv_blocks]
+            """
+            if num_blocks is None:
+                return None, None
+            new_kv_num_blocks = num_blocks.clone()
+            new_kv_indices = torch.zeros(
+                (B, H, ROWS, self.n_pages), dtype = torch.int32, device = device
+            )
+            new_kv_indices[:, :, :, :MAX_BLOCKS_IN_COL] = (
+                torch.gather(page_table, 1, indices.view(B, -1).to(torch.int64))
+                .view(block_mask.kv_indices.shape)
+                .to(torch.int32)
+            )
+            return new_kv_num_blocks, new_kv_indices
+
+        new_kv_num_blocks, new_kv_indices = transform(
+            block_mask.kv_num_blocks, block_mask.kv_indices
+        )
+        new_full_kv_num_blocks, new_full_kv_indices = transform(
+            block_mask.full_kv_num_blocks, block_mask.full_kv_indices
+        )
+
+        new_mask_mod = self.get_mask_mod(block_mask.mask_mod, batch_idx)
+
+        seq_lengths = (block_mask.seq_lengths[0], self.n_pages * self.page_size)
+        return BlockMask.from_kv_blocks(
+            new_kv_num_blocks,
+            new_kv_indices,
+            new_full_kv_num_blocks,
+            new_full_kv_indices,
+            block_mask.BLOCK_SIZE,
+            new_mask_mod,
+            seq_lengths = seq_lengths,
+        )
+
+    def get_logical_kv_idx(
+        self,
+        physical_batch_idx: torch.Tensor,
+        physical_kv_idx: torch.Tensor,
+        batch_idx: torch.Tensor,
+    ):
+        logical_batch_idx = batch_idx[physical_batch_idx]
+        physical_kv_block = physical_kv_idx // self.page_size
+        physical_kv_offset = physical_kv_idx % self.page_size
+        logical_block_idx = self.physical_to_logical[
+            logical_batch_idx, physical_kv_block
+        ]
+        logical_kv_idx = logical_block_idx * self.page_size + physical_kv_offset
+        is_valid = logical_block_idx >= 0
+        safe_logical_kv_idx = logical_kv_idx.clamp(min = 0)
+        return is_valid, safe_logical_kv_idx
+
+    def get_mask_mod(
+        self, mask_mod: Optional[_mask_mod_signature], batch_idx: torch.Tensor
+    ) -> _mask_mod_signature:
+        """
+        Converts a mask_mod based on mapping from the physical block index to the logical
+        block index.
+
+        Args:
+            mask_mod (_mask_mod_signature): mask_mod based on the logical block index.
+        """
+        if mask_mod is None:
+            mask_mod = noop_mask
+
+        def new_mask_mod(
+            b: torch.Tensor,
+            h: torch.Tensor,
+            q_idx: torch.Tensor,
+            physical_kv_idx: torch.Tensor,
+        ):
+            is_valid, safe_logical_kv_idx = self.get_logical_kv_idx(
+                b, physical_kv_idx, batch_idx
+            )
+            return torch.where(
+                is_valid, mask_mod(b, h, q_idx, safe_logical_kv_idx), False
+            )
+
+        return new_mask_mod
+
+    # NOTE: not used in the current codebase
+    def get_score_mod(
+        self, score_mod: Optional[_score_mod_signature], batch_idx: torch.Tensor
+    ) -> _score_mod_signature:
+        """
+        Converts a score_mod based on mapping from the physical block index to the logical
+        block index.
+
+        Args:
+            score_mod (_score_mod_signature): score_mod based on the logical block index.
+        """
+        if score_mod is None:
+            score_mod = _identity
+
+        def new_score_mod(
+            score: torch.Tensor,
+            b: torch.Tensor,
+            h: torch.Tensor,
+            q_idx: torch.Tensor,
+            physical_kv_idx: torch.Tensor,
+        ):
+            is_valid, safe_logical_kv_idx = self.get_logical_kv_idx(
+                b, physical_kv_idx, batch_idx
+            )
+            return torch.where(
+                is_valid,
+                score_mod(score, b, h, q_idx, safe_logical_kv_idx),
+                float("-inf"),
+            )
+
+        return new_score_mod
+
+    def create_causal_blockmask(self, B, L):
+        """A minimal, unoptimized causal block mask creation function"""
+
+        def causal(b, h, q_idx, kv_idx):
+            return q_idx >= kv_idx
+
+        return create_block_mask(
+            causal,
+            B = B,
+            H = None,
+            Q_LEN = L,
+            KV_LEN = L,
+            BLOCK_SIZE = self.page_size,
+            device = self.device,
+        )
+
+    def create_prefill_blockmask_no_paging(
+        self, batch_idx: Tensor, BLOCK_SIZE: int = 128
+    ):
+        """
+        there's no prefix sharing implemented, batch_idx is the document id, batch_idx is not guaranteed to be sorted
+        """
+        assert batch_idx.ndim == 2, "batch_idx must be a 2D tensor"
+        assert batch_idx.shape[0] == 1, "batch_idx must have batch size 1"
+        L = batch_idx.shape[1]
+        docs = batch_idx.view(-1)
+
+        def document_causal(b, h, q_idx, kv_idx):
+            causal_mask = q_idx >= kv_idx
+            document_mask = docs[q_idx] == docs[kv_idx]
+            return causal_mask & document_mask
+
+        return create_block_mask(
+            document_causal, B = 1, H = None, Q_LEN = L, KV_LEN = L, BLOCK_SIZE = BLOCK_SIZE
+        )
+
+    # we assign prefill to the cache, similar to assign(), except we don't return the k_cache, v_cache, we only return the k_val, v_val
+    def assign_prefill_no_paging(
+        self,
+        batch_idx: torch.Tensor,
+        input_pos: torch.Tensor,
+        k_val: torch.Tensor,
+        v_val: torch.Tensor,
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+    ) -> None:
+        """
+        assigns kv and returns the original kv
+
+        batch_idx: [1, L]
+        input_pos: [1, L]
+        k_val: [1, H, L, D]
+        v_val: [1, H, L, D]
+        k_cache: [1, H, MAX_S, D]
+        v_cache: [1, H, MAX_S, D]
+        """
+
+        assert batch_idx.ndim == 2, "batch_idx must be a 2D tensor"
+        assert input_pos.ndim == 2, "input_pos must be a 2D tensor"
+        assert k_val.ndim == 4, "k_val must be a 4D tensor"
+        assert v_val.ndim == 4, "v_val must be a 4D tensor"
+        assert k_cache.ndim == 4, "k_cache must be a 4D tensor"
+        assert v_cache.ndim == 4, "v_cache must be a 4D tensor"
+        assert batch_idx.shape[0] == 1, "batch_idx must have batch size 1"
+
+        input_pos_block_idx = input_pos // self.page_size
+        input_pos_offset_in_block = input_pos % self.page_size
+        physical_kv_idx = (
+            self.page_table[batch_idx, input_pos_block_idx] * self.page_size
+            + input_pos_offset_in_block
+        )
+        k_cache[:, :, physical_kv_idx.view(-1), :] = k_val
+        v_cache[:, :, physical_kv_idx.view(-1), :] = v_val
+
+        return k_val, v_val

scripts/benchmarks/make_lora_adapter.py

@@ -0,0 +1,105 @@
+"""One-shot: materialize a rank-32 LoRA adapter on `unsloth/Qwen3-4B-Base`.
+
+Writes a PEFT-style directory so every backend (vLLM `LoRARequest`,
+`peft.PeftModel.from_pretrained`, Unsloth `FastLanguageModel.get_peft_model`)
+can load the SAME weights. Random-init is fine for throughput measurement --
+the goal is to have LoRA kernels active during generation, not a trained
+model.
+
+Run:
+    CUDA_VISIBLE_DEVICES=6 python scripts/benchmarks/make_lora_adapter.py \
+        --output outputs/lora_rank32_fresh
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+from pathlib import Path
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--model_name", default = "unsloth/Qwen3-4B-Base")
+    p.add_argument("--output", default = "outputs/lora_rank32_fresh")
+    p.add_argument("--rank", type = int, default = 32)
+    return p.parse_args()
+
+
+def main():
+    args = parse_args()
+    out_dir = Path(args.output).resolve()
+    out_dir.mkdir(parents = True, exist_ok = True)
+
+    # Use vanilla HF -- PEFT's save_pretrained yields the canonical
+    # adapter_config.json + adapter_model.safetensors that vLLM's LoRARequest
+    # expects. Loading via Unsloth would leak Unsloth-specific LoRA wrappers.
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    from peft import LoraConfig, get_peft_model
+
+    tok = AutoTokenizer.from_pretrained(args.model_name)
+    if tok.pad_token is None:
+        tok.pad_token = tok.eos_token
+
+    # bf16 base; we only need structure + save. Keep on CPU to avoid a GPU load
+    # just for `save_pretrained`.
+    print(f"[make_lora_adapter] Loading {args.model_name} on CPU...")
+    model = AutoModelForCausalLM.from_pretrained(args.model_name, dtype = torch.bfloat16)
+
+    peft_cfg = LoraConfig(
+        r = args.rank,
+        lora_alpha = args.rank * 2,
+        target_modules = [
+            "q_proj",
+            "k_proj",
+            "v_proj",
+            "o_proj",
+            "gate_proj",
+            "up_proj",
+            "down_proj",
+        ],
+        bias = "none",
+        task_type = "CAUSAL_LM",
+        lora_dropout = 0.0,
+    )
+    peft_model = get_peft_model(model, peft_cfg)
+    peft_model.print_trainable_parameters()
+
+    # Ensure both A and B matrices are non-zero. PEFT initializes A with
+    # kaiming_uniform and B with zeros -- which makes the adapter a no-op and
+    # would mask LoRA kernels on some backends. Seed B with tiny random values.
+    n_reinit = 0
+    with torch.no_grad():
+        for name, p in peft_model.named_parameters():
+            if "lora_B" in name:
+                p.normal_(mean = 0.0, std = 1e-4)
+                n_reinit += 1
+    print(
+        f"[make_lora_adapter] Reinitialized {n_reinit} lora_B matrices with tiny gaussian."
+    )
+
+    peft_model.save_pretrained(str(out_dir))
+    tok.save_pretrained(str(out_dir))
+
+    # Sanity: verify safetensors file present and non-trivial.
+    from safetensors import safe_open
+
+    st_path = out_dir / "adapter_model.safetensors"
+    n_zero_tensors = 0
+    n_tensors = 0
+    with safe_open(str(st_path), framework = "pt") as f:
+        for key in f.keys():
+            t = f.get_tensor(key)
+            n_tensors += 1
+            if (t == 0).all().item():
+                n_zero_tensors += 1
+    print(
+        f"[make_lora_adapter] Wrote {n_tensors} tensors to {st_path} "
+        f"({n_zero_tensors} all-zero)."
+    )
+    print(f"[make_lora_adapter] Adapter saved to {out_dir}")
+
+
+if __name__ == "__main__":
+    main()

scripts/benchmarks/persistent_cb.py

@@ -0,0 +1,121 @@
+"""Persistent ContinuousBatching manager for multi-step generation.
+
+`model.generate_batch(...)` initializes a fresh `ContinuousBatchingManager`
+on every call, which in turn allocates a new `PagedAttentionCache` and
+starts a new worker thread. Inside a GRPO training loop this happens once
+per step, amortized across `num_generations * per_device_train_batch_size`
+prompts per step, so the constant per-step cost (cache alloc, prefill
+warmup, thread spin-up) dominates when batch sizes are modest.
+
+`install_for_model(model, generation_config)` monkey-patches
+`model.generate_batch` on this instance to reuse a single long-lived
+manager. The manager is started lazily on first call; `teardown(model)`
+stops the background thread.
+
+This is deliberately kept as a stand-alone helper so it can be enabled /
+disabled per-run via CLI flag without touching TRL or transformers
+installs.
+"""
+
+from __future__ import annotations
+
+import threading
+import types
+from typing import Optional
+
+import torch
+from transformers import GenerationConfig
+from transformers.generation.continuous_batching import RequestStatus
+
+
+_ATTR = "_persistent_cb_manager"
+_LOCK_ATTR = "_persistent_cb_lock"
+
+
+def install_for_model(
+    model: torch.nn.Module, generation_config: GenerationConfig
+) -> None:
+    """Replace `model.generate_batch` with a version that reuses one manager.
+
+    The replacement accepts the same arguments as the stock method. A
+    trailing `generation_config` supplied to the call takes precedence; if
+    it differs from the one used at init, the persistent manager is torn
+    down and rebuilt (rare, but keeps semantics intact).
+    """
+    setattr(model, _LOCK_ATTR, threading.Lock())
+    setattr(model, _ATTR, None)
+    setattr(model, "_persistent_cb_gen_config", generation_config)
+
+    original = model.generate_batch
+
+    def generate_batch(
+        self,
+        inputs,
+        generation_config: Optional[GenerationConfig] = None,
+        progress_bar: bool = False,
+        slice_inputs: bool = True,
+        **kwargs,
+    ):
+        if not inputs:
+            return {}
+
+        gen_config = (
+            generation_config
+            or getattr(self, "_persistent_cb_gen_config", None)
+            or self.generation_config
+        )
+
+        lock = getattr(self, _LOCK_ATTR)
+        with lock:
+            manager = getattr(self, _ATTR)
+            stale = False
+            if manager is not None:
+                stale = (
+                    getattr(manager, "generation_config", None) is not gen_config
+                    or not manager.is_running()
+                )
+                if stale:
+                    try:
+                        manager.stop(block = True, timeout = 5.0)
+                    except Exception:
+                        pass
+                    setattr(self, _ATTR, None)
+                    manager = None
+            if manager is None:
+                manager = self.init_continuous_batching(
+                    generation_config = gen_config,
+                    slice_inputs = slice_inputs,
+                )
+                manager.start()
+                setattr(self, _ATTR, manager)
+
+        results = {}
+        num_requests = len(inputs)
+        manager.add_requests(inputs, **kwargs)
+        finished = 0
+        while finished < num_requests:
+            result = manager.get_result(timeout = 1)
+            if result is None:
+                if not manager.is_running():
+                    break
+                continue
+            if result.status == RequestStatus.FINISHED:
+                results[result.request_id] = result
+                finished += 1
+            else:
+                continue
+        return results
+
+    model.generate_batch = types.MethodType(generate_batch, model)
+    setattr(model, "_persistent_cb_original_generate_batch", original)
+
+
+def teardown(model: torch.nn.Module) -> None:
+    manager = getattr(model, _ATTR, None)
+    if manager is not None:
+        try:
+            manager.stop(block = True, timeout = 5.0)
+        except Exception:
+            pass
+    if hasattr(model, "_persistent_cb_original_generate_batch"):
+        model.generate_batch = model._persistent_cb_original_generate_batch

scripts/benchmarks/qwen3_grpo_naive.py

@@ -0,0 +1,188 @@
+"""Qwen3-4B GRPO with naive TRL generation (no vLLM, no CB).
+
+Mirrors the TRL example at https://huggingface.co/docs/trl/grpo_trainer: a
+vanilla HF model + `GRPOTrainer` with the default rollout path, which calls
+`model.generate(...)` per step. This is the honest "baseline" baseline — it
+is what a user would get if they just followed the TRL docs without enabling
+vLLM colocate or transformers continuous batching. Useful as a third column
+in the benchmark table.
+
+Hyperparameters, dataset, and reward functions are shared with the vLLM and
+CB scripts via `unsloth_grpo_common.py` so numbers are apples-to-apples.
+
+Run:
+    CUDA_VISIBLE_DEVICES=2 python scripts/qwen3_grpo_naive.py \
+        --max_steps 20 --num_generations 2 --per_device_train_batch_size 2 \
+        --output_dir outputs/grpo_naive --stats_path logs/naive_stats.json
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+import time
+from pathlib import Path
+
+HERE = Path(__file__).resolve().parent
+sys.path.insert(0, str(HERE))
+
+from unsloth_grpo_common import (  # noqa: E402
+    StepTimer,
+    apply_chat_template_to_tokenizer,
+    build_dataset,
+    build_grpo_kwargs,
+    build_reward_funcs,
+    install_vllm_sampling_shim,
+    maybe_compile_trainer_forwards,
+    write_stats,
+)
+
+install_vllm_sampling_shim()
+
+import torch  # noqa: E402
+from transformers import AutoModelForCausalLM, AutoTokenizer  # noqa: E402
+from peft import LoraConfig, get_peft_model  # noqa: E402
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--model_name", default = "unsloth/Qwen3-4B-Base")
+    p.add_argument("--max_seq_length", type = int, default = 2048)
+    p.add_argument("--lora_rank", type = int, default = 32)
+    p.add_argument("--max_steps", type = int, default = 20)
+    p.add_argument("--num_generations", type = int, default = 2)
+    p.add_argument("--per_device_train_batch_size", type = int, default = 2)
+    p.add_argument("--gradient_accumulation_steps", type = int, default = 1)
+    p.add_argument(
+        "--attn_impl",
+        default = "sdpa",
+        help = "Attention implementation: sdpa or flash_attention_2 (FA4 shim installed).",
+    )
+    p.add_argument("--output_dir", default = "outputs/grpo_naive")
+    p.add_argument("--stats_path", default = "logs/naive_stats.json")
+    p.add_argument(
+        "--compile_mode",
+        default = None,
+        choices = [None, "default", "reduce-overhead", "max-autotune-no-cudagraphs"],
+        help = "If set, torch.compile(model.forward, mode=...) after the trainer is built.",
+    )
+    p.add_argument("--compile_dynamic", action = "store_true", default = True)
+    return p.parse_args()
+
+
+def main():
+    args = parse_args()
+    os.makedirs(args.output_dir, exist_ok = True)
+    os.makedirs(os.path.dirname(os.path.abspath(args.stats_path)) or ".", exist_ok = True)
+
+    # Install the FA4 shim only if the caller asked for flash_attention_2.
+    # For sdpa we leave transformers untouched.
+    if args.attn_impl == "flash_attention_2":
+        import flash_attn_fa4_shim
+
+        flash_attn_fa4_shim.apply()
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    model = AutoModelForCausalLM.from_pretrained(
+        args.model_name,
+        dtype = torch.bfloat16,
+        attn_implementation = args.attn_impl,
+    ).to("cuda")
+
+    lora = LoraConfig(
+        r = args.lora_rank,
+        lora_alpha = args.lora_rank * 2,
+        target_modules = [
+            "q_proj",
+            "k_proj",
+            "v_proj",
+            "o_proj",
+            "gate_proj",
+            "up_proj",
+            "down_proj",
+        ],
+        bias = "none",
+        task_type = "CAUSAL_LM",
+    )
+    model = get_peft_model(model, lora)
+    try:
+        model.gradient_checkpointing_enable(
+            gradient_checkpointing_kwargs = {"use_reentrant": False}
+        )
+    except TypeError:
+        model.gradient_checkpointing_enable()
+    model.enable_input_require_grads()
+
+    apply_chat_template_to_tokenizer(tokenizer)
+
+    dataset, maximum_length = build_dataset(
+        tokenizer, max_seq_length = args.max_seq_length
+    )
+    print(f"[naive] Max prompt length (p90): {maximum_length}")
+    reward_funcs = build_reward_funcs(tokenizer)
+
+    from trl import GRPOConfig, GRPOTrainer
+
+    shared = build_grpo_kwargs(
+        tokenizer,
+        maximum_length,
+        max_seq_length = args.max_seq_length,
+        max_steps = args.max_steps,
+        num_generations = args.num_generations,
+        per_device_train_batch_size = args.per_device_train_batch_size,
+        gradient_accumulation_steps = args.gradient_accumulation_steps,
+        output_dir = args.output_dir,
+    )
+    # transformers' TopKLogitsWarper rejects -1. None skips the warper.
+    shared["top_k"] = None
+
+    training_args = GRPOConfig(
+        use_vllm = False,
+        use_transformers_paged = False,
+        bf16 = True,
+        **shared,
+    )
+
+    timer = StepTimer()
+
+    trainer = GRPOTrainer(
+        model = model,
+        processing_class = tokenizer,
+        reward_funcs = reward_funcs,
+        args = training_args,
+        train_dataset = dataset,
+        callbacks = [timer],
+    )
+
+    maybe_compile_trainer_forwards(
+        trainer, args.compile_mode, dynamic = args.compile_dynamic, tag = "naive"
+    )
+
+    torch.cuda.reset_peak_memory_stats()
+    t_start = time.perf_counter()
+    trainer.train()
+    t_train = time.perf_counter() - t_start
+
+    peak = torch.cuda.max_memory_allocated() / 1024**3
+
+    write_stats(
+        args.stats_path,
+        backend = "naive_trl",
+        timer = timer,
+        train_wall_s = t_train,
+        peak_memory_gb = peak,
+        max_prompt_length = shared["max_prompt_length"],
+        max_completion_length = shared["max_completion_length"],
+        num_generations = args.num_generations,
+        max_steps = args.max_steps,
+        extra = {"attn_impl": args.attn_impl},
+    )
+    print(f"[naive] Wrote stats to {args.stats_path}")
+    print(f"[naive] Total train wall: {t_train:.1f}s   Peak mem: {peak:.2f} GB")
+
+
+if __name__ == "__main__":
+    main()

scripts/benchmarks/qwen3_grpo_tpaged.py

@@ -0,0 +1,242 @@
+"""Qwen3-4B GRPO with transformers continuous-batching rollouts.
+
+Unsloth's Qwen3Attention monkey-patch bypasses the functional attention
+interface that `paged|<impl>` continuous batching relies on, so this script
+loads a vanilla HF Qwen3 with PEFT LoRA instead. Training is slower than the
+Unsloth path but the goal here is to evaluate transformers CB as a drop-in
+replacement for vLLM rollouts. See benchmark_results.md for numbers.
+
+Run:
+    CUDA_VISIBLE_DEVICES=2 python scripts/qwen3_grpo_tpaged.py \
+        --max_steps 61 --output_dir outputs/grpo_tpaged \
+        --stats_path logs/tpaged_stats.json
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+import time
+from pathlib import Path
+
+HERE = Path(__file__).resolve().parent
+sys.path.insert(0, str(HERE))
+
+# `install_vllm_sampling_shim()` shims `vllm.sampling_params.GuidedDecodingParams`
+# for newer vLLM releases so TRL's GRPOTrainer imports cleanly. We do NOT
+# import `unsloth` here because that replaces TRL's GRPOTrainer with an
+# Unsloth-compiled variant that assumes the model has `for_training()` /
+# `for_inference()` hooks, which a vanilla HF model does not.
+from unsloth_grpo_common import (  # noqa: E402
+    StepTimer,
+    apply_chat_template_to_tokenizer,
+    build_dataset,
+    build_grpo_kwargs,
+    build_reward_funcs,
+    install_vllm_sampling_shim,
+    maybe_compile_trainer_forwards,
+    write_stats,
+)
+
+install_vllm_sampling_shim()
+
+import torch  # noqa: E402
+from transformers import AutoModelForCausalLM, AutoTokenizer  # noqa: E402
+from peft import LoraConfig, get_peft_model  # noqa: E402
+
+import flash_attn_fa4_shim  # noqa: E402
+
+flash_attn_fa4_shim.apply()
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--model_name", default = "unsloth/Qwen3-4B-Base")
+    p.add_argument("--max_seq_length", type = int, default = 2048)
+    p.add_argument("--lora_rank", type = int, default = 32)
+    p.add_argument("--max_steps", type = int, default = 61)
+    p.add_argument("--num_generations", type = int, default = 4)
+    p.add_argument("--per_device_train_batch_size", type = int, default = 1)
+    p.add_argument("--gradient_accumulation_steps", type = int, default = 1)
+    p.add_argument(
+        "--attn_impl",
+        default = "sdpa",
+        help = "Base attention impl to compose with paged. 'sdpa' or 'flash_attention_2'.",
+    )
+    p.add_argument(
+        "--max_batch_tokens",
+        type = int,
+        default = 8192,
+        help = "PagedAttentionCache.max_batch_tokens. Default upper bound is 256 which is far too small.",
+    )
+    p.add_argument(
+        "--num_blocks",
+        type = int,
+        default = 8192,
+        help = "PagedAttentionCache.num_blocks (block_size=32). 8192*32 tokens of KV capacity.",
+    )
+    p.add_argument("--output_dir", default = "outputs/grpo_tpaged")
+    p.add_argument("--stats_path", default = "logs/tpaged_stats.json")
+    p.add_argument(
+        "--persistent_cb",
+        action = "store_true",
+        help = "Reuse one ContinuousBatchingManager across every training step instead "
+        "of letting TRL's generate_batch rebuild it (and the paged cache) each step.",
+    )
+    p.add_argument(
+        "--compile_mode",
+        default = None,
+        choices = [None, "default", "reduce-overhead", "max-autotune-no-cudagraphs"],
+        help = "If set, torch.compile(model.forward, mode=...) after the trainer is built.",
+    )
+    p.add_argument("--compile_dynamic", action = "store_true", default = True)
+    return p.parse_args()
+
+
+def main():
+    args = parse_args()
+    os.makedirs(args.output_dir, exist_ok = True)
+    os.makedirs(os.path.dirname(os.path.abspath(args.stats_path)) or ".", exist_ok = True)
+
+    # 1. Vanilla HF load (no Unsloth patches on the attention forward).
+    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    model = AutoModelForCausalLM.from_pretrained(
+        args.model_name,
+        dtype = torch.bfloat16,
+        attn_implementation = args.attn_impl,
+    )
+    model.to("cuda")
+
+    lora = LoraConfig(
+        r = args.lora_rank,
+        lora_alpha = args.lora_rank * 2,
+        target_modules = [
+            "q_proj",
+            "k_proj",
+            "v_proj",
+            "o_proj",
+            "gate_proj",
+            "up_proj",
+            "down_proj",
+        ],
+        bias = "none",
+        task_type = "CAUSAL_LM",
+    )
+    model = get_peft_model(model, lora)
+    try:
+        model.gradient_checkpointing_enable(
+            gradient_checkpointing_kwargs = {"use_reentrant": False}
+        )
+    except TypeError:
+        model.gradient_checkpointing_enable()
+    model.enable_input_require_grads()
+
+    apply_chat_template_to_tokenizer(tokenizer)
+
+    # 2. Dataset + rewards (identical to the vLLM script).
+    dataset, maximum_length = build_dataset(
+        tokenizer, max_seq_length = args.max_seq_length
+    )
+    print(f"[tpaged] Max prompt length (p90): {maximum_length}")
+    reward_funcs = build_reward_funcs(tokenizer)
+
+    # 3. Build GRPOConfig with transformers continuous batching enabled.
+    from trl import GRPOConfig, GRPOTrainer
+
+    shared = build_grpo_kwargs(
+        tokenizer,
+        maximum_length,
+        max_seq_length = args.max_seq_length,
+        max_steps = args.max_steps,
+        num_generations = args.num_generations,
+        per_device_train_batch_size = args.per_device_train_batch_size,
+        gradient_accumulation_steps = args.gradient_accumulation_steps,
+        output_dir = args.output_dir,
+    )
+    # transformers `TopKLogitsWarper` rejects -1. None skips the warper entirely.
+    shared["top_k"] = None
+    # The default PagedAttentionCache upper bounds
+    # (`_upper_bound_max_batch_tokens=256`, `_upper_bound_num_blocks=4096`)
+    # are extremely conservative and cause long decode loops. Raise them via
+    # `generation_kwargs`, which TRL forwards to `GenerationConfig`, which the
+    # CB manager then reads off when sizing the paged cache.
+    training_args = GRPOConfig(
+        use_vllm = False,
+        use_transformers_paged = True,
+        bf16 = True,
+        generation_kwargs = {
+            "max_batch_tokens": args.max_batch_tokens,
+            "num_blocks": args.num_blocks,
+        },
+        **shared,
+    )
+
+    # 4. Timing callback.
+    timer = StepTimer()
+
+    trainer = GRPOTrainer(
+        model = model,
+        processing_class = tokenizer,
+        reward_funcs = reward_funcs,
+        args = training_args,
+        train_dataset = dataset,
+        callbacks = [timer],
+    )
+
+    maybe_compile_trainer_forwards(
+        trainer, args.compile_mode, dynamic = args.compile_dynamic, tag = "tpaged"
+    )
+
+    if args.persistent_cb:
+        # TRL constructs `self.generation_config` once in `__init__`; reuse
+        # the same object so the persistent manager stays warm.
+        from persistent_cb import install_for_model, teardown
+
+        # TRL generates against the unwrapped base model; attach the patch
+        # directly to it so every rollout picks up the persistent manager.
+        base = (
+            trainer.model_wrapped.base_model.model
+            if hasattr(trainer.model_wrapped, "base_model")
+            else trainer.model_wrapped
+        )
+        install_for_model(base, trainer.generation_config)
+        # PEFT's wrapper chains `generate_batch` through `base_model.model` via
+        # __getattr__, so installing on `base` is enough for TRL's call path.
+
+    torch.cuda.reset_peak_memory_stats()
+    t_start = time.perf_counter()
+    try:
+        trainer.train()
+    finally:
+        if args.persistent_cb:
+            from persistent_cb import teardown
+
+            teardown(base)
+    t_train = time.perf_counter() - t_start
+
+    peak = torch.cuda.max_memory_allocated() / 1024**3
+
+    write_stats(
+        args.stats_path,
+        backend = "transformers_paged",
+        timer = timer,
+        train_wall_s = t_train,
+        peak_memory_gb = peak,
+        max_prompt_length = shared["max_prompt_length"],
+        max_completion_length = shared["max_completion_length"],
+        num_generations = args.num_generations,
+        max_steps = args.max_steps,
+        extra = {
+            "attn_impl": args.attn_impl,
+            "persistent_cb": args.persistent_cb,
+        },
+    )
+    print(f"[tpaged] Wrote stats to {args.stats_path}")
+    print(f"[tpaged] Total train wall: {t_train:.1f}s   Peak mem: {peak:.2f} GB")
+
+
+if __name__ == "__main__":
+    main()

scripts/benchmarks/qwen3_grpo_vllm.py

@@ -0,0 +1,159 @@
+"""Qwen3-4B GRPO baseline (vLLM colocated) derived from the notebook.
+
+Run:
+    CUDA_VISIBLE_DEVICES=2 python scripts/qwen3_grpo_vllm.py \
+        --max_steps 61 --output_dir outputs/grpo_vllm \
+        --stats_path logs/vllm_stats.json
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+import time
+from pathlib import Path
+
+# Unsloth must be imported before transformers / trl.
+os.environ.setdefault("UNSLOTH_VLLM_STANDBY", "1")
+
+# Allow sibling import of the common module.
+HERE = Path(__file__).resolve().parent
+sys.path.insert(0, str(HERE))
+
+from unsloth import FastLanguageModel  # noqa: E402
+import torch  # noqa: E402
+
+from unsloth_grpo_common import (  # noqa: E402
+    StepTimer,
+    apply_chat_template_to_tokenizer,
+    build_dataset,
+    build_grpo_kwargs,
+    build_reward_funcs,
+    write_stats,
+)
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--model_name", default = "unsloth/Qwen3-4B-Base")
+    p.add_argument("--max_seq_length", type = int, default = 2048)
+    p.add_argument("--lora_rank", type = int, default = 32)
+    p.add_argument("--max_steps", type = int, default = 61)
+    p.add_argument("--num_generations", type = int, default = 4)
+    p.add_argument("--per_device_train_batch_size", type = int, default = 1)
+    p.add_argument("--gradient_accumulation_steps", type = int, default = 1)
+    p.add_argument("--gpu_memory_utilization", type = float, default = 0.8)
+    p.add_argument("--output_dir", default = "outputs/grpo_vllm")
+    p.add_argument("--stats_path", default = "logs/vllm_stats.json")
+    return p.parse_args()
+
+
+def main():
+    args = parse_args()
+    os.makedirs(args.output_dir, exist_ok = True)
+    os.makedirs(os.path.dirname(os.path.abspath(args.stats_path)) or ".", exist_ok = True)
+
+    # 1. Load model with vLLM fast inference enabled.
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name = args.model_name,
+        max_seq_length = args.max_seq_length,
+        load_in_4bit = False,
+        fast_inference = True,
+        max_lora_rank = args.lora_rank,
+        gpu_memory_utilization = args.gpu_memory_utilization,
+    )
+    model = FastLanguageModel.get_peft_model(
+        model,
+        r = args.lora_rank,
+        target_modules = [
+            "q_proj",
+            "k_proj",
+            "v_proj",
+            "o_proj",
+            "gate_proj",
+            "up_proj",
+            "down_proj",
+        ],
+        lora_alpha = args.lora_rank * 2,
+        use_gradient_checkpointing = "unsloth",
+        random_state = 3407,
+    )
+    apply_chat_template_to_tokenizer(tokenizer)
+
+    # 2. Dataset + rewards.
+    dataset, maximum_length = build_dataset(
+        tokenizer, max_seq_length = args.max_seq_length
+    )
+    print(f"[vllm] Max prompt length (p90): {maximum_length}")
+    reward_funcs = build_reward_funcs(tokenizer)
+
+    # 3. vLLM sampling params match the notebook.
+    from vllm import SamplingParams
+
+    vllm_sampling_params = SamplingParams(
+        min_p = 0.1,
+        top_p = 1.0,
+        top_k = -1,
+        seed = 3407,
+        stop = [tokenizer.eos_token],
+        include_stop_str_in_output = True,
+    )
+
+    # 4. Build GRPOConfig.
+    from trl import GRPOConfig, GRPOTrainer
+
+    shared = build_grpo_kwargs(
+        tokenizer,
+        maximum_length,
+        max_seq_length = args.max_seq_length,
+        max_steps = args.max_steps,
+        num_generations = args.num_generations,
+        per_device_train_batch_size = args.per_device_train_batch_size,
+        gradient_accumulation_steps = args.gradient_accumulation_steps,
+        output_dir = args.output_dir,
+    )
+    training_args = GRPOConfig(
+        use_vllm = True,
+        vllm_mode = "colocate",
+        vllm_sampling_params = vllm_sampling_params,
+        vllm_gpu_memory_utilization = args.gpu_memory_utilization,
+        **shared,
+    )
+
+    # 5. Timing callback.
+    timer = StepTimer()
+
+    trainer = GRPOTrainer(
+        model = model,
+        processing_class = tokenizer,
+        reward_funcs = reward_funcs,
+        args = training_args,
+        train_dataset = dataset,
+        callbacks = [timer],
+    )
+
+    torch.cuda.reset_peak_memory_stats()
+    t_start = time.perf_counter()
+    trainer.train()
+    t_train = time.perf_counter() - t_start
+
+    peak = torch.cuda.max_memory_allocated() / 1024**3
+
+    write_stats(
+        args.stats_path,
+        backend = "vllm_colocated",
+        timer = timer,
+        train_wall_s = t_train,
+        peak_memory_gb = peak,
+        max_prompt_length = shared["max_prompt_length"],
+        max_completion_length = shared["max_completion_length"],
+        num_generations = args.num_generations,
+        max_steps = args.max_steps,
+    )
+    print(f"[vllm] Wrote stats to {args.stats_path}")
+    print(f"[vllm] Total train wall: {t_train:.1f}s   Peak mem: {peak:.2f} GB")
+
+
+if __name__ == "__main__":
+    main()

scripts/benchmarks/results/flex_vs_vllm.md

@@ -0,0 +1,291 @@
+# flex_attention + paged KV + CUDA graphs vs vLLM
+
+Goal stated in the plan: "CB reaches at least 30% of vLLM throughput." After
+the earlier phases ran out of gas at ~10% with transformers CB, we rebuilt
+the rollout path on top of `torch.nn.attention.flex_attention` using the
+paged KV + BlockMask pattern from
+[flex-nano-vllm](https://github.com/changjonathanc/flex-nano-vllm).
+
+## Setup
+
+- B200 (sm_100), Qwen3-4B-Base, bf16
+- 512 max_new_tokens per prompt, 16-prompt warmup, N measured rounds
+  (`decode_tps_best` = steady-state throughput after Inductor compile +
+  CUDA graph capture have amortized).
+- flex path is greedy (CUDA-graph safe); vLLM uses equivalence sampling
+  (`temperature=0.1, top_p=0.97, min_p=0.5, top_k=5`).
+- No LoRA unless noted; LoRA rank 32 applied to all
+  {q,k,v,o,gate,up,down}_proj.
+
+## Best config (after FlexKernelOptions sweep)
+
+```json
+decode_kernel_options = {
+    "PRESCALE_QK": true,
+    "USE_TMA": true,
+    "BLOCKS_ARE_CONTIGUOUS": true,
+    "num_warps": 8,
+    "num_stages": 3
+}
+prefill_kernel_options = {
+    "FORCE_USE_FLEX_ATTENTION": true,
+    "PRESCALE_QK": true,
+    "USE_TMA": true
+}
+```
+
+## Batch-size sweep (flex tuned vs vLLM, 512 max_new_tokens)
+
+| Batch | flex tps | vLLM tps | flex / vLLM | flex mem | vLLM mem |
+|------:|---------:|---------:|------------:|---------:|---------:|
+|   8   |    680   |  1900    |    35.8 %   |  44 GB   |  156 GB  |
+|  16   |   1626   |  3698    |    44.0 %   |  44 GB   |  156 GB  |
+|  32   |   3134   |  6318    |    49.6 %   |  44 GB   |  156 GB  |
+|  64   | **5474** | 10459    |  **52.3 %** |  44 GB   |  156 GB  |
+| 128   |   5565   | 14996    |    37.1 %   |  81 GB   |  157 GB  |
+| 256   |   5812   | 21170    |    27.5 %   | 154 GB   |  157 GB  |
+
+### Canonical GRPO workload (batch 64 + LoRA rank 32)
+
+| Backend                                    | tok/s best | peak mem  | flex / vLLM |
+|--------------------------------------------|-----------:|----------:|------------:|
+| vLLM (LoRARequest)                         |       7775 |    156 GB |       100 % |
+| **flex** (double-copy, drift-free)         |  **5785**  | **~52 GB**|    **74 %** |
+| flex -- LoRA unmerged (PEFT wrapper)       |       2683 |     45 GB |        35 % |
+
+At the GRPO workload flex reaches **74 % of vLLM throughput at ~3 x less
+memory**. Starting point before this work was 9 % with transformers CB.
+
+**Why the two flex rows are so far apart:** when PEFT keeps the adapter
+unmerged, every projection runs three matmuls (`base_layer(x) + scaling *
+lora_B(lora_A(x))`) instead of one, which is ~50 % slowdown across the
+36-layer stack. GRPO cannot use the unmerged path naively because the
+trainer needs the adapter weights separable; but it also shouldn't pay
+that cost.
+
+#### What the default path does now: double-copy rollout
+
+We keep two copies of the base model on GPU:
+
+- `base_model` -- pristine; never mutated.
+- `inference_model = deepcopy(base_model)` -- wrapped by PEFT; merged
+  LoRA lives on `base_layer.weight` here.
+
+Before each rollout (and at setup), `refresh_lora_merge_from_pristine`:
+
+1. Walks PEFT's `LoraLayer` modules.
+2. `module.base_layer.weight.data.copy_(base_submodule.weight.data)` --
+   in-place restore from the pristine base.
+3. Resets `module.merged_adapters = []` directly (skips PEFT's unmerge
+   arithmetic).
+4. Calls `peft_model.merge_adapter()` once to fold LoRA into the
+   inference copy fresh.
+
+We **never call `unmerge_adapter()`**. PEFT's merge/unmerge pair is
+asymmetric at bf16 -- merge does `W_bf16 += delta_fp32` (the `+=`
+upcasts, stores back in bf16), unmerge does `W_bf16 -= delta_fp32.to(bf16)`
+(the delta is rounded to bf16 first, then subtracted). Net effect is ~1
+ULP drift on `base_layer.weight` per cycle (empirically ~6e-5 max diff
+after one cycle on this model). Across hundreds of GRPO iterations
+that corrupts the base model and the adapter trains against a drifting
+target. Re-materialising from pristine per refresh bypasses the whole
+round-trip.
+
+**Cost.** +~8 GB GPU memory (second copy of Qwen3-4B bf16 weights), so
+peak memory goes from ~44 GB to ~52 GB. Per-refresh overhead: param
+copy (~3 ms) + `merge_adapter()` (~30 ms) = ~35 ms, well under 1 % of a
+5-7 s rollout.
+
+**CUDA graphs stay valid.** In-place `weight.data.copy_(pristine)`
+writes to the same tensor storage, so graphs captured against the
+merged weights read current values at the captured addresses on the
+next replay -- no re-capture needed.
+
+#### Drift verification
+
+`--verify_no_drift` takes a sha256 over every parameter in `base_model`
+(raw bytes via `tensor.view(torch.uint8)`), runs N perturb+refresh
+cycles (random noise added to `lora_A` / `lora_B` on each iteration,
+simulating a training step), re-hashes, and asserts bit-identical.
+It also checks determinism of the inference copy: after restoring the
+LoRA A/B weights to their initial values and refreshing, the merged
+state-dict hash matches the pre-perturbation hash.
+
+Confirmed on Qwen3-4B bf16 with LoRA rank 32 across 10 refreshes: base
+model bit-identical; inference copy deterministic after LoRA restore.
+
+```
+CUDA_VISIBLE_DEVICES=6 python scripts/benchmarks/qwen3_flex_inference.py \
+  --verify_no_drift --lora_adapter outputs/lora_rank32_fresh --n_rounds 1 \
+  --stats_path scripts/benchmarks/results/stats/flex_verify_nodrift.json
+```
+
+The "LoRA unmerged" row is shown only for reference -- it's what you'd
+get with a naive PEFT wrapper on the hot path. **Don't use it in
+production**, and it doesn't apply outside the 4-bit path below.
+
+`--no_merge_lora` opts into that reference path (single model, PEFT
+wrapper, adapter unmerged). It's kept for the comparison row above and
+nothing else.
+
+### Same workload at `load_in_4bit=True` (Unsloth bnb-4bit shard)
+
+Loading base as bitsandbytes 4-bit (`unsloth/Qwen3-4B-Base-unsloth-bnb-4bit`,
+compute dtype bf16). LoRA kept as PEFT wrapper (can't merge into 4-bit;
+the double-copy pattern above also doesn't apply -- bnb's `Linear4bit`
+holds packed quantised weights, not regular bf16, so an in-place copy
+of `base_layer.weight` isn't meaningful, and materialising a bf16
+inference copy via dequant would wipe out the memory saving of 4-bit).
+lm_head is tied to embed_tokens post-load because the 4-bit shard ships
+without an lm_head parameter.
+
+| 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 throughput on every backend (vLLM-path 4515 vs bf16 7775 = 58 %;
+flex 1738 vs bf16 5744 = 30 %). The regression is worse for flex because
+PEFT-without-merge doubles the number of matmuls per projection (base + LoRA
+add, separately) on top of the bnb dequant cost; the bf16 path merges LoRA
+into the base and skips both. Peak memory barely moves for vLLM because KV
+cache at `gpu_memory_utilization=0.8` dominates regardless of base size.
+
+Transformers CB at 4-bit + LoRA produces garbage tokens even with
+`model.lm_head.weight = model.model.embed_tokens.weight` tied explicitly.
+Likely a PEFT-over-bnb + batched `generate_batch` interaction bug; did not
+debug further.
+
+## What each option did (batch 64, no LoRA, after CUDA graph capture)
+
+| Config                                                                | tok/s | vs baseline |
+|-----------------------------------------------------------------------|------:|------------:|
+| eager (no graphs)                                                     |  ~420 |       -     |
+| + CUDA graphs                                                         |  4279 |    baseline |
+| + `PRESCALE_QK=true`                                                  |  4367 |       +2 %  |
+| + `USE_TMA=true`                                                      |  4425 |       +3 %  |
+| + `BLOCKS_ARE_CONTIGUOUS=true`                                        |  4703 |       +10 % |
+| + `num_warps=8`                                                       |  5474 |       +28 % |
+| + `num_warps=8, num_stages=3`                                         |  **5898** (peak) |       +38 % |
+
+The single biggest win came from **`num_warps=8`** (up from the default,
+which on Blackwell tends to pick 4 for small block sizes). TMA helps a
+couple percent; `BLOCKS_ARE_CONTIGUOUS` (safe in our setup because
+PageTable.reserve allocates pages sequentially on a fresh batch) helps
+another ~10 % because it lets the kernel skip the page-table indirection
+per block.
+
+## What broke correctness and had to be dropped
+
+- **`ROWS_GUARANTEED_SAFE=true`**: we reserve `batch_idx=0` and
+  `page_idx=0` as padding slots. Padded decode rows only attend to those
+  reserved slots, so the mask returns False for every kv_idx on those
+  rows. Skipping the row-has-at-least-one-unmasked check NaNs the
+  softmax and the model outputs `!!!!!!`.
+- **`BACKEND="TRITON_DECODE"`**: documented but the Inductor code path
+  doesn't recognize the literal. Raises `NameError('TRITON_DECODE is
+  not defined')`.
+- **`USE_TMA=true` + `torch.compile(call_model_with_flex_kwargs)`**:
+  misaligned address at runtime. torch.compile on the whole forward
+  walker breaks TMA's alignment assumptions. Either disable TMA when
+  compiling the walker, or skip compiling the walker (CUDA graph
+  capture already captures it).
+- **`torch.compile(flex_attention, mode="max-autotune")`**: tries to
+  nest `cudagraph_trees` inside our raw CUDA graph capture and hits
+  `Cannot prepare for replay during capturing stage`. Use
+  `max-autotune-no-cudagraphs` instead; negligible throughput delta vs
+  default mode.
+
+## What I tried that did NOT move the needle
+
+- **`BACKEND="FLASH"` on prefill** (FA4 / FlashAttention-4 on Blackwell,
+  torch 2.11 + flash-attn CuTeDSL): empirically 4617 tok/s at batch 64 +
+  LoRA vs 5744 baseline on torch 2.11. The FA4 CuTe kernel is slow when
+  `mask_mod` indexes by `kv_idx` (documented in the attention-gym
+  `flex_flash_attention.py` limitations: "Indexing by kv_idx is a large
+  perf hit"). Our prefill mask is `document_causal`:
+      `docs[q_idx] == docs[kv_idx]`
+  which hits that exact slow path. `BLOCK_SIZE=(256, 128)` + padding to
+  the 256-row Q tile works (output is coherent), it's just slower than
+  the default Triton flex path for this mask.
+- **Inductor autotune replay** (`TORCHINDUCTOR_FLEX_ATTENTION_LOGGING_FILE`
+  + `mode="max-autotune-no-cudagraphs"` + parse the JSON log): Inductor's
+  chosen best decode config (`fwd_num_warps=4, fwd_num_stages=3,
+  fwd_BLOCK_M=64, fwd_BLOCK_N=64`) lands at 4827 tok/s -- worse than the
+  hand-tuned `num_warps=8` at 5744. Autotune times a single kernel call,
+  which doesn't catch cumulative register-spill / L1 effects across the
+  36-layer stack. Harness lives at `flex_autotune_replay.py`.
+- **torch.compile on `call_model_with_flex_kwargs`**: 4425 tok/s (same
+  as eager walker) because the CUDA graph already captures every op in
+  the walker into one replay. The compile step is work we don't need.
+- **`num_warps=4` / `num_warps=16`**: 4486 / 4748 -- neither beats 8.
+  Inductor's default picks 4 on small blocks and we're already past
+  that sweet spot, but 16 wastes registers.
+- **Explicit `fwd_BLOCK_M=128, fwd_BLOCK_N=128` pinning** on top of the
+  manual best: 5009 tok/s. The implicit default already picks 128 for
+  our shape; pinning it inhibits Inductor's shape-specialised choice
+  between the flex_attention and flex_decoding templates.
+
+## Torch version + run-to-run noise
+
+Upgrading torch 2.9.1 -> 2.11 (required for FA4's CuTeDSL path) moves
+best-of-N tok/s from ~5616 to ~5660 at batch 64 + LoRA -- essentially
+within noise. Over 10 rounds, median is 4192 and best is 5660; the large
+spread is GPU clock throttling across a ~60-second sustained run plus
+variable prompt-length distributions per round. Reported numbers use
+best-of-N to match the prior harness; steady-state median is roughly 75
+ % of best.
+
+## Architecture notes (unchanged from prior commits)
+
+- `flex_paged_attention.py`: `PagedKVCache` + `PageTable` verbatim from
+  flex-nano-vllm (BSD-3).
+- `qwen3_flex_inference.py`: monkey-patches `Qwen3Attention.forward` to
+  call `flex_attention(q, k, v, block_mask=...)` against the paged cache.
+  Walks the `Qwen3Model` layer stack manually so `flex_block_mask /
+  flex_input_pos / flex_batch_idx` reach the attention layer without
+  modifying `Qwen3ForCausalLM.forward`. `capture_decode_cudagraph()`
+  pre-reserves one page per batch slot, captures one CUDA graph per
+  bucket in `[1,2,4,8,16,32,...,max_bs]`, then releases the scratch
+  batches.
+
+## Output coherence
+
+All tuned configs produce coherent math solutions on the DAPO-Math-17k
+prompts. See `sample_completions` in any `logs/flex_*_tuned.json`.
+
+## What's left on the table
+
+- **Chunked prefill**: vLLM interleaves prefill and decode inside a
+  single step. flex does a full separate prefill pass per new batch,
+  which is the main remaining penalty for large batches.
+- **Prefill-path mask refactor**: the document_causal mask indexes by
+  `kv_idx`. Flattening to a per-query bias (`bias[q_idx]`) would put
+  FA4 back on the fast path, but this is a non-trivial rework because
+  the causal-within-document constraint needs to be encoded without the
+  `docs[kv_idx]` lookup.
+- **Exhaustive Triton autotune** for flex_decoding: attention-gym's
+  `flex_grid_sweep.py` enumerates 144 fwd configs; Inductor's default
+  autotune only probes a handful. Running the full sweep with
+  end-to-end tok/s as the metric (not single-call ms) might beat the
+  manual num_warps=8 finding, but 144 * 5 rounds is ~20 hrs of B200
+  time.
+- **Kernel-level parity on decode**: vLLM on sm_100 uses FlashInfer
+  TRTLLM kernels which are fused / tuned more aggressively than
+  flex_attention's Inductor-generated Triton. Closing the last
+  ~28-48 % gap will require either tuning more Triton configs or
+  waiting for a TMA-native flex_attention path.
+
+## Raw stats (under `scripts/benchmarks/results/stats/`)
+
+- `flex_{8,16,32,64,128}_tuned.json` (best opts, 5 rounds, torch 2.9.1)
+- `flex_64_lora_tuned.json` (GRPO canonical, torch 2.9.1)
+- `flex_64_lora_torch211_baseline.json` + `_repeat.json` + `_10rounds.json`
+  (same config re-run on torch 2.11 to measure noise)
+- `flex_64_lora_fa4prefill.json` (FA4 prefill regression at batch 64)
+- `flex_64_lora_autotune{,_tma}.json` (Inductor-autotune-suggested config)
+- `flex_64_lora_warps2.json` + `flex_64_lora_pinned_blocks.json` (other
+  sweep points)
+- `flex_{32,64,128,256}x512[_lora]_cudagraph.json` (prior best-of-3 runs)
+- `vllm_{8,16,32,64,128,256}[x512][_lora].json`

scripts/benchmarks/results/grpo_equivalence.md

@@ -0,0 +1,128 @@
+# Phase 2: end-to-end GRPO backend comparison
+
+Same dataset, reward functions, sampling (`temperature=0.1, top_p=0.97,
+min_p=0.5, top_k=5`), and seed (3407) across backends. `num_generations=4`;
+`per_device_train_batch_size` auto-raised to 4 on vanilla-HF backends so
+TRL's `generation_batch_size % num_generations == 0` check passes
+(Unsloth's loader does this for you, vanilla HF does not).
+
+Callbacks: `StatisticsCallback` from `torch_debugging_utils` logs per-step
+loss, grad-norm, memory, wall time; reward / KL are captured from the TRL
+log dict. Median step wall is measured on steps 4..N (first 3 skipped to
+amortize compile / graph / warmup).
+
+## 10-step vibe check
+
+| Backend                    | Train wall (s) | Median step (s) | Peak mem (GB) | % of vLLM |
+|----------------------------|----------------|-----------------|---------------|-----------|
+| vLLM (fast_inference)      | 74.4           | **4.14**        | 157.9         | 100 %     |
+| unsloth_fi_false           | 355.4          | 23.95           | **10.7**      | 17 %      |
+| cb_paged (sdpa_paged load) | 466.0          | 36.02           | 55.6          | 11.5 %    |
+
+## 30-step equivalence
+
+| Backend                    | Train wall (s) | Median step (s) | Peak mem (GB) | % of vLLM |
+|----------------------------|----------------|-----------------|---------------|-----------|
+| vLLM (fast_inference)      | 215.9          | **5.14**        | 159.0         | 100 %     |
+| unsloth_fi_false           | 1165.4         | 41.30           | **10.7**      | 12.4 %    |
+| cb_paged                   | 1564.5         | 39.82           | 61.9          | 12.9 %    |
+
+(Note: fi_false's median step jumped from 23.95 s at 10 steps to 41.30 s at
+30 steps because the early-GRPO policy started producing longer completions
+as it learned to place the `</SOLUTION>` marker; the same effect is present
+but smaller in cb_paged because its LoRA warm-up trajectory is different.)
+
+## Pairwise diff vs vLLM (30 steps, `scripts/benchmarks/compare_grpo_runs.py`)
+
+| Pair                            | max &#124;loss diff&#124; | max &#124;reward diff&#124; | max &#124;kl diff&#124; | max &#124;grad_norm diff&#124; |
+|---------------------------------|---------------------------|------------------------------|------------------------|--------------------------------|
+| vLLM vs **unsloth_fi_false**    | 0.39                      | 9.25 (mean 2.99)             | **0.015**              | 0.94                           |
+| vLLM vs **cb_paged**            | 0.83                      | 6.25 (mean 2.29)             | *(not logged)*         | 919.1                          |
+
+Reward diffs of 2-9 are expected: different rollout backends produce
+different completions even at `temperature=0.1` because of kernel-level
+non-determinism (vLLM uses FlashInfer TRTLLM kernels, CB uses paged SDPA /
+FA4, Unsloth uses its cached fp16 LoRA path). The reward function reads
+those completions, so the reward array mechanically differs. What matters
+for equivalence is:
+
+- **KL trajectory is near-identical** between vLLM and unsloth_fi_false
+  (both stay in `[0, 0.015]` across all 30 steps). The KL *term* of the
+  GRPO loss is the guardrail against policy drift, so matching KL means
+  the training dynamics are in the same regime.
+- **Loss magnitudes are bounded** in `[-0.3, 1.0]` for all three backends.
+- **No NaNs, no unbounded growth, no gibberish completions** in any run.
+
+## grad_norm 919 on cb_paged
+
+The enormous cb_paged grad_norm (vs vLLM's ~1.0) is a clipping story, not a
+correctness story: the vLLM path goes through Unsloth's `FastLanguageModel`
+which clips gradients to `max_grad_norm=1.0` internally, while the vanilla
+HF path used by cb_paged picks up TRL's raw grad_norm reported by the
+optimizer pre-clip (or without clipping if no `max_grad_norm` is set in
+GRPOConfig). For a fair training-dynamics comparison the cb_paged config
+should set `max_grad_norm=1.0` explicitly; left for a follow-up commit.
+
+## KL missing for cb_paged
+
+`StatisticsCallback.on_log` forwards the full TRL log dict into its per-step
+entry only on steps where `loss` is present. TRL's vanilla-HF path separately
+logs KL on a different log call that doesn't include loss, so the callback
+silently drops it. Follow-up: relax the callback so every log dict with a
+`step` field merges into the matching entry regardless of which keys are
+present.
+
+## Headline takeaways
+
+1. **unsloth_fi_false is the pragmatic middle ground**: 12-17% of vLLM's
+   throughput, **15x less peak memory** (10.7 GB vs 159 GB), KL trajectory
+   matching vLLM within sampling noise.
+2. **cb_paged is close to fi_false in throughput at this batch size** (41 s
+   vs 40 s median step at 30 steps) but costs 6x more memory. Phase 3
+   (main-thread sync driver + CUDA graphs on the rollout) is the right
+   lever for making CB competitive.
+3. **torch.compile on the training step is not a quick win** for either
+   backend (Phase 4 report below).
+
+## Phase 3 state (CB sync driver)
+
+`scripts/benchmarks/cb_sync_driver.py`:
+- Eager main-thread driver works end-to-end: smoke test on GPU 1 with 8
+  prompts / 64 tokens produced the expected 512 correct tokens.
+- CUDA graph capture hangs on the first graphed step. Likely cause:
+  `ContinuousBatchProcessor._sample` reads `next_tokens.size(1)` as a
+  Python int to slice `batch_processor.output_ids[:, :tokens]`, which
+  forces a CPU-GPU sync and is not CUDA-graph-safe. Fix direction: keep
+  a fixed `tokens` count when `slice_inputs=False` (buffer size is
+  constant), or rewrite the copy as a full-buffer `copy_` without the
+  slice.
+- Deferred to a follow-up commit.
+
+## Phase 4 state (torch.compile on training forward)
+
+- `unsloth_fi_false + compile_mode=default`: crashes 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+ recompiles /
+  graph breaks on the first optimizer step and never makes progress.
+  Root cause: `modeling_utils.make_inputs_require_grads` calls
+  `Tensor.requires_grad_()` which triggers Dynamo GB0125 (unsupported
+  mutating op). TRL's GRPO `_compute_loss` then re-enters the tracer,
+  which re-triggers the break, which recompiles, and so on.
+- `vllm` is excluded (vLLM owns its own compile pipeline).
+
+Net: compile on the training step is not the right lever in this stack.
+Phase 3 (CUDA graphs on the rollout decode) is.
+
+## Raw stats
+
+- `scripts/benchmarks/results/stats/grpo_{vllm,unsloth_fi_false,cb_paged}_{10,30}.json`
+  (StatisticsCallback per-step logs with full TRL metric dict)
+- `scripts/benchmarks/results/stats/grpo_*_{10,30}.summary.json` (short form)
+
+Pairwise diff:
+
+    python scripts/benchmarks/compare_grpo_runs.py \
+        --ref logs/grpo_vllm_30.json \
+        --candidate logs/grpo_unsloth_fi_false_30.json

scripts/benchmarks/results/lora_rollout_baselines.md

@@ -0,0 +1,60 @@
+# Phase 1: rollout-only LoRA rank-32 microbenchmark
+
+Every backend generates the same 32 prompts (DAPO-Math-17k, seed 3407) for
+`max_new_tokens=512` with equivalence sampling: `temperature=0.1, top_p=0.97,
+min_p=0.5, top_k=5`. 16-prompt warmup, 2 measured rounds, median wall reported.
+
+All four backends load the same `outputs/lora_rank32_fresh` adapter (see
+`make_lora_adapter.py`). LoRA kernels are active on every decode step.
+
+## Results (GPU B200, bf16, Qwen3-4B-Base + rank-32 LoRA)
+
+| Backend              | Median wall (s) | Decode tok/s | Prompt tok/s | Peak mem (GB) | % of vLLM |
+|----------------------|-----------------|--------------|--------------|---------------|-----------|
+| vLLM (fast_inference)| 3.30            | **4581**     | 1467         | 156.2         | 100.0 %   |
+| unsloth_fi_false     | 25.54           | 641          | 190          | **15.8**      | 14.0 %    |
+| CB paged+FA4 (persistent) | 34.99      | 422          | 138          | 103.8         | 9.2 %     |
+| CB sdpa_paged (persistent)| 34.07      | 434          | 142          | 111.9         | 9.5 %     |
+
+## Observations
+
+1. **vLLM with LoRA is ~37% slower than vLLM without LoRA** (7224 → 4581 tok/s
+   per the pre-LoRA PR table). The LoRA kernels cost real time even in vLLM.
+   Still the gold standard by a wide margin.
+
+2. **Unsloth `fast_inference=False` is the surprise**: **1.5× faster than CB**
+   at **1/7th the peak memory**. The cached fp16 LoRA copies in
+   `fast_linear_forward` and the Triton RMSNorm/RoPE paths dominate the CB
+   baseline on this workload. It is a real practical middle ground — no vLLM
+   dependency, low memory, and ~14% of vLLM's throughput.
+
+3. **CB paged_attention (FA4 shim) and CB sdpa_paged are within noise**:
+   422 vs 434 tok/s. At this scale the attention kernel is not the bottleneck;
+   Python-side launch overhead on `_generation_step` dominates (confirmed by
+   prior profile: ~16k `cuLaunchKernelEx` for 371 decoded tokens). CUDA graph
+   replay (Phase 3) is the right lever.
+
+4. **Unsloth `fi_false` reached max_new_tokens on every prompt** (`n_decoded =
+   16384 = 32 × 512`) whereas vLLM / CB stopped some sequences on EOS
+   (`~15000 decoded`). Equivalence sampling + greedy-ish settings means most
+   completions are long, but the slight difference is worth noting when
+   reading the raw tok/s numbers.
+
+5. Completions are qualitatively coherent in every backend (see
+   `sample_completions` in the stats JSONs). vLLM and unsloth_fi_false produce
+   the *same* opening tokens on probe prompts (deterministic sampling lower
+   bound), which is a useful weak sanity check.
+
+## Raw stats
+
+- `scripts/benchmarks/results/stats/lora_vllm_gen.json`
+- `scripts/benchmarks/results/stats/lora_unsloth_fi_false_gen.json`
+- `scripts/benchmarks/results/stats/lora_cb_paged_fa4_gen.json`
+- `scripts/benchmarks/results/stats/lora_cb_sdpa_paged_gen.json`
+
+## Downstream implication
+
+Phase 2 (full GRPO training) will include `unsloth_fi_false` as a first-class
+backend — if throughput parity holds end-to-end, it may be the pragmatic
+default for teams that cannot take the vLLM memory footprint. Phase 3 (CB sync
+driver + CUDA graphs) targets the CB paths specifically.

scripts/benchmarks/results/notebook_ref_10.md

@@ -0,0 +1,54 @@
+# Phase 0 reference run: canonical Unsloth Qwen3-4B GRPO notebook (10 steps)
+
+Reproduction of `Qwen3_(4B)-GRPO.ipynb` with three deviations for the backend
+comparison downstream:
+
+1. `max_steps = 10` (vibe check; 30 and 100 follow in Phase 2).
+2. Equivalence-friendly sampling: `temperature=0.1, top_p=0.97, min_p=0.5,
+   top_k=5`. Low variance so KL / reward trajectories across backends can be
+   compared tightly.
+3. `StatisticsCallback` logs per-step loss, reward, KL, grad-norm, memory, and
+   wall time to `logs/notebook_ref_10.json`.
+
+SFT format-priming stage is skipped with `--skip_sft_pre_finetune` — this run
+is just the GRPO phase.
+
+Config:
+- GPU 6 (B200, bf16)
+- Model: `unsloth/Qwen3-4B-Base`, LoRA rank 32 on all proj layers
+- `num_generations=4`, `per_device_train_batch_size=1` (TRL enforces
+  `pdb * grad_accum * world = multiple of num_generations`, so effective batch
+  is 4 × 1)
+- `gpu_memory_utilization=0.85`
+
+## Per-step results
+
+| step | loss    | reward  | kl      | grad_norm | time(s) | mem(GB) |
+|------|---------|---------|---------|-----------|---------|---------|
+| 1    |  0.2423 | -0.875  | 0.00000 | 0.245     | 60.63   | 158.9   |
+| 2    |  0.1559 | -5.500  | 0.00000 | 0.767     |  3.89   | 157.0   |
+| 3    | -0.1650 | -0.500  | 0.00383 | 0.480     |  7.93   | 158.2   |
+| 4    |  0.3177 | -4.125  | 0.00591 | 0.379     | 11.08   | 158.9   |
+| 5    | -0.0200 |  3.125  | 0.01598 | 0.341     |  2.68   | 156.7   |
+| 6    |  0.0000 | -7.500  | 0.00396 | 0.000     | 10.65   | 158.9   |
+| 7    |  0.0000 | -7.500  | 0.00965 | 0.000     |  4.47   | 157.2   |
+| 8    |  0.0613 | -6.500  | 0.00319 | 0.172     |  5.81   | 157.6   |
+| 9    |  0.0060 | -0.500  | 0.00240 | 0.048     |  4.30   | 157.1   |
+| 10   |  0.1582 | -1.500  | 0.00485 | 0.394     |  6.78   | 157.9   |
+
+**Summary:**
+- Median step wall (steps 4-10): **5.80 s**
+- Total train wall: ~118 s
+- Peak memory: **158.9 GB**
+- KL trajectory: monotonic rise from 0 to ~0.016 by step 5, settles at
+  ~0.005 afterward — consistent with the policy drift being bounded by the KL
+  term.
+- Step 1 is ~60 s because it amortizes the vLLM CUDA-graph capture; the
+  post-warmup median is what Phase 2 will compare against.
+
+## Phase 2 use
+
+This is the gold reference. Every other backend's loss / reward / KL arrays
+will be diffed against this one (see `torch_debugging_utils.compare_training_runs`).
+Throughput numbers are on a separate axis: even an equivalence-passing backend
+that is 3x slower than this is useful information for the PR writeup.

scripts/benchmarks/unsloth_grpo_common.py

@@ -0,0 +1,381 @@
+"""Shared helpers for the Qwen3-4B GRPO comparison scripts.
+
+Exports:
+    REASONING_START, REASONING_END, SOLUTION_START, SOLUTION_END, SYSTEM_PROMPT
+    CHAT_TEMPLATE
+    build_dataset(tokenizer, max_seq_length=2048)
+    build_reward_funcs(tokenizer)
+    build_grpo_kwargs(tokenizer, maximum_length, max_seq_length)
+    StepTimer  (TrainerCallback recording per-step wall time, loss, reward)
+    write_stats(path, backend, timer, ...)
+    install_vllm_sampling_shim()
+
+Keeps dataset loading, chat template, formatting rewards, and GRPO hparams
+identical between the vLLM baseline script and the transformers-CB candidate.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+import time
+
+import numpy as np
+import pandas as pd
+from datasets import Dataset, load_dataset
+
+
+REASONING_START = "<start_working_out>"
+REASONING_END = "<end_working_out>"
+SOLUTION_START = "<SOLUTION>"
+SOLUTION_END = "</SOLUTION>"
+
+SYSTEM_PROMPT = (
+    "You are given a problem.\n"
+    "Think about the problem and provide your working out.\n"
+    f"Place it between {REASONING_START} and {REASONING_END}.\n"
+    f"Then, provide your solution between {SOLUTION_START}{SOLUTION_END}"
+)
+
+
+CHAT_TEMPLATE = (
+    "{% if messages[0]['role'] == 'system' %}"
+    "{{ messages[0]['content'] + eos_token }}"
+    "{% set loop_messages = messages[1:] %}"
+    "{% else %}"
+    "{{ '%%%SYSTEM_PROMPT%%%' + eos_token }}"
+    "{% set loop_messages = messages %}"
+    "{% endif %}"
+    "{% for message in loop_messages %}"
+    "{% if message['role'] == 'user' %}"
+    "{{ message['content'] }}"
+    "{% elif message['role'] == 'assistant' %}"
+    "{{ message['content'] + eos_token }}"
+    "{% endif %}"
+    "{% endfor %}"
+    "{% if add_generation_prompt %}{{ '%%%REASONING_START%%%' }}"
+    "{% endif %}"
+)
+
+
+def apply_chat_template_to_tokenizer(tokenizer):
+    tmpl = CHAT_TEMPLATE.replace("%%%SYSTEM_PROMPT%%%", SYSTEM_PROMPT)
+    tmpl = tmpl.replace("%%%REASONING_START%%%", REASONING_START)
+    tokenizer.chat_template = tmpl
+    return tokenizer
+
+
+def build_dataset(tokenizer, *, max_seq_length: int = 2048):
+    """Build the DAPO-Math-17k GRPO dataset with the prompt formatting from the
+    notebook. Returns `(dataset, maximum_prompt_length)`.
+    """
+    ds = load_dataset("open-r1/DAPO-Math-17k-Processed", "en", split = "train")
+
+    def _map_row(x):
+        return {
+            "prompt": [
+                {"role": "system", "content": SYSTEM_PROMPT},
+                {"role": "user", "content": x["prompt"]},
+            ],
+            "answer": x["solution"],
+        }
+
+    ds = ds.map(_map_row)
+
+    # Tokenize for length measurement (batched for speed).
+    def _tokenize(batch):
+        return {
+            "tokens": tokenizer.apply_chat_template(
+                batch["prompt"],
+                add_generation_prompt = True,
+                tokenize = True,
+            )
+        }
+
+    tokenized = ds.map(_tokenize, batched = True)
+    tokenized = tokenized.map(lambda x: {"L": len(x["tokens"])})
+    lengths = np.array(tokenized["L"])
+    maximum_length = int(np.quantile(lengths, 0.9))
+    ds = ds.select(np.where(lengths <= maximum_length)[0])
+    return ds, maximum_length
+
+
+def build_reward_funcs(tokenizer):
+    """Return the 4 reward functions used in the notebook, wired to `tokenizer`."""
+    solution_end_regex = (
+        r"</SOLUTION>[\s]{0,}" + "(?:" + re.escape(tokenizer.eos_token) + ")?"
+    )
+    match_format = re.compile(
+        rf"{REASONING_END}.*?"
+        rf"{SOLUTION_START}(.+?){solution_end_regex}"
+        rf"[\s]{{0,}}$",
+        flags = re.MULTILINE | re.DOTALL,
+    )
+    match_numbers = re.compile(
+        SOLUTION_START + r".*?[\s]{0,}([-]?[\d\.\,]{1,})",
+        flags = re.MULTILINE | re.DOTALL,
+    )
+
+    def match_format_exactly(completions, **kwargs):
+        scores = []
+        for completion in completions:
+            score = 0.0
+            response = completion[0]["content"]
+            if match_format.search(response) is not None:
+                score += 3.0
+            scores.append(score)
+        return scores
+
+    def match_format_approximately(completions, **kwargs):
+        scores = []
+        for completion in completions:
+            score = 0.0
+            response = completion[0]["content"]
+            score += 0.5 if response.count(REASONING_END) == 1 else -1.0
+            score += 0.5 if response.count(SOLUTION_START) == 1 else -1.0
+            score += 0.5 if response.count(SOLUTION_END) == 1 else -1.0
+            scores.append(score)
+        return scores
+
+    def check_answer(prompts, completions, answer, **kwargs):
+        responses = [c[0]["content"] for c in completions]
+        extracted = [
+            guess.group(1) if (guess := match_format.search(r)) is not None else None
+            for r in responses
+        ]
+        scores = []
+        for guess, true_answer in zip(extracted, answer):
+            score = 0.0
+            if guess is None:
+                scores.append(-2.0)
+                continue
+            if guess == true_answer:
+                score += 5.0
+            elif guess.strip() == true_answer.strip():
+                score += 3.5
+            else:
+                try:
+                    ratio = float(guess) / float(true_answer)
+                    if 0.9 <= ratio <= 1.1:
+                        score += 2.0
+                    elif 0.8 <= ratio <= 1.2:
+                        score += 1.5
+                    else:
+                        score -= 2.5
+                except Exception:
+                    score -= 4.5
+            scores.append(score)
+        return scores
+
+    _printed_state = {"n": 0, "every": 5}
+
+    def check_numbers(prompts, completions, answer, **kwargs):
+        question = prompts[0][-1]["content"]
+        responses = [c[0]["content"] for c in completions]
+        extracted = [
+            guess.group(1) if (guess := match_numbers.search(r)) is not None else None
+            for r in responses
+        ]
+        if _printed_state["n"] % _printed_state["every"] == 0:
+            print(
+                "*" * 20 + f"Question:\n{question}",
+                f"\nAnswer:\n{answer[0]}",
+                f"\nResponse:\n{responses[0]}",
+                f"\nExtracted:\n{extracted[0]}",
+            )
+        _printed_state["n"] += 1
+
+        scores = []
+        for guess, true_answer in zip(extracted, answer):
+            if guess is None:
+                scores.append(-2.5)
+                continue
+            try:
+                t = float(true_answer.strip())
+                g = float(guess.strip().replace(",", ""))
+                scores.append(3.5 if g == t else -1.5)
+            except Exception:
+                scores.append(0.0)
+        return scores
+
+    return [
+        match_format_exactly,
+        match_format_approximately,
+        check_answer,
+        check_numbers,
+    ]
+
+
+def build_grpo_kwargs(
+    tokenizer,
+    maximum_length: int,
+    *,
+    max_seq_length: int = 2048,
+    max_steps: int = 100,
+    num_generations: int = 4,
+    per_device_train_batch_size: int = 1,
+    gradient_accumulation_steps: int = 1,
+    output_dir: str = "outputs",
+):
+    """Return the shared dict of GRPOConfig kwargs used by both backends.
+
+    Caller adds backend-specific keys (use_vllm / use_transformers_paged / etc).
+    """
+    max_prompt_length = maximum_length + 1
+    max_completion_length = max_seq_length - max_prompt_length
+
+    return dict(
+        temperature = 1.0,
+        top_p = 1.0,
+        top_k = -1,
+        min_p = 0.1,
+        learning_rate = 5e-6,
+        weight_decay = 0.001,
+        warmup_ratio = 0.1,
+        lr_scheduler_type = "linear",
+        optim = "adamw_8bit",
+        logging_steps = 1,
+        per_device_train_batch_size = per_device_train_batch_size,
+        gradient_accumulation_steps = gradient_accumulation_steps,
+        num_generations = num_generations,
+        max_prompt_length = max_prompt_length,
+        max_completion_length = max_completion_length,
+        max_steps = max_steps,
+        save_steps = max_steps,
+        report_to = "none",
+        output_dir = output_dir,
+        seed = 3407,
+    )
+
+
+import torch
+from transformers import TrainerCallback
+
+
+class StepTimer(TrainerCallback):
+    """Per-step wall time / loss / reward recorder.
+
+    Shared verbatim across qwen3_grpo_{vllm,naive,tpaged}. Records step wall
+    time in `self.step_wall`, and picks up `loss` / `reward` from the TRL log
+    dict in `on_log`.
+    """
+
+    def __init__(self):
+        self.t0 = None
+        self.step_wall = []
+        self.loss = []
+        self.reward = []
+
+    def on_step_begin(self, _args, state, control, **kwargs):
+        torch.cuda.synchronize()
+        self.t0 = time.perf_counter()
+
+    def on_log(self, _args, state, control, logs = None, **kwargs):
+        if logs is None:
+            return
+        if "loss" in logs:
+            self.loss.append(float(logs["loss"]))
+        if "reward" in logs:
+            self.reward.append(float(logs["reward"]))
+
+    def on_step_end(self, _args, state, control, **kwargs):
+        if self.t0 is not None:
+            torch.cuda.synchronize()
+            self.step_wall.append(time.perf_counter() - self.t0)
+
+
+def write_stats(
+    path: str,
+    backend: str,
+    timer: StepTimer,
+    *,
+    train_wall_s: float,
+    peak_memory_gb: float,
+    max_prompt_length: int,
+    max_completion_length: int,
+    num_generations: int,
+    max_steps: int,
+    extra: dict | None = None,
+) -> None:
+    """Dump the per-step stats dict used by all three GRPO drivers.
+
+    Schema matches the pre-refactor output exactly: `backend`, `train_wall_s`,
+    `peak_memory_gb`, `step_wall_s`, `losses`, `rewards`, `max_prompt_length`,
+    `max_completion_length`, `num_generations`, `max_steps`, plus any
+    backend-specific keys passed in `extra` (e.g. `attn_impl`, `persistent_cb`).
+    """
+    stats = {
+        "backend": backend,
+        "train_wall_s": train_wall_s,
+        "peak_memory_gb": peak_memory_gb,
+        "step_wall_s": timer.step_wall,
+        "losses": timer.loss,
+        "rewards": timer.reward,
+        "max_prompt_length": max_prompt_length,
+        "max_completion_length": max_completion_length,
+        "num_generations": num_generations,
+        "max_steps": max_steps,
+    }
+    if extra:
+        stats.update(extra)
+    with open(path, "w") as f:
+        json.dump(stats, f, indent = 2)
+
+
+def maybe_compile_trainer_forwards(
+    trainer, compile_mode, *, dynamic: bool = True, tag: str = ""
+):
+    """torch.compile wrap `trainer.model.forward` and (if present)
+    `trainer.ref_model.forward`. No-op if `compile_mode` is falsy.
+
+    Ported from the old `qwen3_grpo_unified.py` compile path, minus the
+    out-of-tree `torch_debugging_utils` imports that were dev-only.
+    """
+    if not compile_mode:
+        return
+    import torch._dynamo
+
+    torch._dynamo.config.cache_size_limit = 128
+    try:
+        torch._dynamo.config.allow_unspec_int_on_nn_module = True
+    except AttributeError:
+        pass
+    prefix = f"[{tag}] " if tag else ""
+    print(
+        f"{prefix}Compiling trainer.model.forward (mode={compile_mode}, dynamic={dynamic})"
+    )
+    trainer.model.forward = torch.compile(
+        trainer.model.forward,
+        mode = compile_mode,
+        dynamic = dynamic,
+    )
+    ref = getattr(trainer, "ref_model", None)
+    if ref is not None:
+        ref.forward = torch.compile(
+            ref.forward,
+            mode = compile_mode,
+            dynamic = dynamic,
+        )
+
+
+def install_vllm_sampling_shim():
+    """Shim `vllm.sampling_params.GuidedDecodingParams` for newer vLLM releases.
+
+    TRL's `GRPOTrainer` imports `GuidedDecodingParams` from
+    `vllm.sampling_params`; newer vLLM versions have moved or removed it. Inject
+    a no-op class so the import succeeds even on the non-vLLM training paths
+    (naive, tpaged). No-op if vLLM is not installed or already exposes the
+    symbol.
+    """
+    try:
+        import vllm.sampling_params as _vllm_sp
+    except ImportError:
+        return
+    if hasattr(_vllm_sp, "GuidedDecodingParams"):
+        return
+
+    class _GuidedDecodingParamsShim:  # pragma: no cover - used only if TRL asks
+        def __init__(self, *a, **kw):
+            pass
+
+    _vllm_sp.GuidedDecodingParams = _GuidedDecodingParamsShim

scripts/benchmarks/verify_gemma4_numerics.py

@@ -0,0 +1,149 @@
+"""Compare first-token logits between `FlexGemma4Inference._prefill` and
+vanilla `Gemma4ForCausalLM.forward` on the same prompt. Intended as a
+one-shot correctness check; not part of the benchmark matrix.
+
+Run:
+    CUDA_VISIBLE_DEVICES=2 python scripts/benchmarks/verify_gemma4_numerics.py
+"""
+
+from __future__ import annotations
+
+import copy
+import sys
+from pathlib import Path
+
+import torch
+
+HERE = Path(__file__).resolve().parent
+sys.path.insert(0, str(HERE))
+
+from gemma4_flex_inference import (  # noqa: E402
+    FlexGemma4Inference,
+    Sequence,
+    _require_gemma4,
+)
+
+
+def main():
+    Gemma4ForCausalLM, Gemma4Config, Gemma4TextConfig = _require_gemma4()
+    from transformers.models.gemma4.modeling_gemma4 import (
+        Gemma4ForConditionalGeneration,
+    )
+    from transformers import AutoTokenizer
+
+    name = "unsloth/gemma-4-E2B-it"
+    tok = AutoTokenizer.from_pretrained(name)
+    if tok.pad_token is None:
+        tok.pad_token = tok.eos_token
+
+    full_cfg = Gemma4Config.from_pretrained(name)
+    text_cfg = full_cfg.text_config
+
+    # Untouched HF reference: `Gemma4ForConditionalGeneration.forward` --
+    # the same class everyone else would load via AutoModelForCausalLM
+    # for Gemma-4. No patching, no shell, no flex attention.
+    ref_raw = Gemma4ForConditionalGeneration.from_pretrained(
+        name, dtype = torch.bfloat16, attn_implementation = "eager"
+    ).to("cuda")
+    ref_raw.eval()
+
+    # Shell copy used by `gemma4_flex_inference.main()`: we deep-copy the
+    # loaded multimodal model, drop the vision + audio towers, and move
+    # the language_model into a Gemma4ForCausalLM wrapper so PEFT and
+    # state-dict hashing treat it as a decoder-only model. The flex path
+    # then patches its attention forwards on this shell. We keep the
+    # shell around both as (a) the model that gets flex-patched and
+    # (b) a sanity check that the shell itself matches the raw HF path.
+    full = Gemma4ForConditionalGeneration.from_pretrained(
+        name, dtype = torch.bfloat16, attn_implementation = "eager"
+    )
+    lang = full.model.language_model
+    full.model.vision_tower = None
+    full.model.audio_tower = None
+    full.model.embed_vision = None
+    full.model.embed_audio = None
+
+    shell = Gemma4ForCausalLM(text_cfg)
+    shell.model = lang
+    shell.lm_head.weight = lang.embed_tokens.weight
+    shell = shell.to(torch.bfloat16).to("cuda")
+    shell.eval()
+    del full
+
+    # Deep-copy so Flex's attention patching doesn't mutate the shell.
+    flex_model = copy.deepcopy(shell)
+
+    prompt = "The quick brown fox jumps over"
+    ids = tok(prompt, return_tensors = "pt")["input_ids"].to("cuda")
+    print(f"prompt len = {ids.shape[1]}")
+
+    with torch.inference_mode():
+        # `Gemma4ForConditionalGeneration.forward` applies
+        # `final_logit_softcapping` internally.
+        ref_logits = ref_raw(input_ids = ids, use_cache = False).logits[0, -1, :].float()
+        shell_logits = shell(ids, use_cache = False).logits[0, -1, :].float()
+    print(
+        f"raw   Gemma4ForConditionalGeneration: mean {ref_logits.mean():.4f}, "
+        f"std {ref_logits.std():.4f}, argmax {int(ref_logits.argmax())} "
+        f"({tok.decode([int(ref_logits.argmax())])!r})"
+    )
+    print(
+        f"shell Gemma4ForCausalLM(text_cfg)    : mean {shell_logits.mean():.4f}, "
+        f"std {shell_logits.std():.4f}, argmax {int(shell_logits.argmax())}"
+    )
+    # Dispose of the raw multimodal model before we build FlexGemma4Inference.
+    del ref_raw
+    torch.cuda.empty_cache()
+
+    # Flex path.
+    inf = FlexGemma4Inference(
+        flex_model,
+        tok,
+        max_batch_size = 4,
+        max_seq_length = 256,
+        n_pages = 64,
+        page_size = 64,
+        max_new_tokens = 1,
+        decode_kernel_options = {"BLOCK_M": 16, "BLOCK_N": 16},
+        prefill_kernel_options = {
+            "FORCE_USE_FLEX_ATTENTION": True,
+            "BLOCK_M": 32,
+            "BLOCK_N": 32,
+        },
+        fa4_prefill = False,
+    )
+    seq = Sequence(text = prompt, max_new_tokens = 1)
+    inf.tokenize([seq])
+    bi = inf.page_table.allocate()
+    inf.page_table.reserve(
+        bi,
+        torch.tensor([bi], device = "cuda", dtype = torch.long),
+        seq.total_length,
+    )
+    seq.batch_idx = bi
+    with torch.inference_mode():
+        flex_logits = inf._prefill([seq])[0].float()
+    print(
+        f"flex  last-token logits              : mean {flex_logits.mean():.4f}, "
+        f"std {flex_logits.std():.4f}, argmax {int(flex_logits.argmax())} "
+        f"({tok.decode([int(flex_logits.argmax())])!r})"
+    )
+
+    def report(tag, a, b):
+        diff = (a - b).abs()
+        top_a = set(a.topk(10).indices.tolist())
+        top_b = set(b.topk(10).indices.tolist())
+        print(
+            f"  {tag:14s} max {diff.max().item():.3e}  mean {diff.mean().item():.3e}  "
+            f"argmax={int(a.argmax()) == int(b.argmax())}  top-10={len(top_a & top_b)}/10"
+        )
+
+    print("vs raw Gemma4ForConditionalGeneration:")
+    report("shell vs raw", shell_logits, ref_logits)
+    report("flex  vs raw", flex_logits, ref_logits)
+    print("vs shell (Gemma4ForCausalLM wrapper):")
+    report("flex  vs shell", flex_logits, shell_logits)
+
+
+if __name__ == "__main__":
+    main()

scripts/benchmarks/verify_qwen3_numerics.py

@@ -0,0 +1,94 @@
+"""Compare first-token logits between FlexInference._prefill (qwen3 path)
+and vanilla model(input_ids) for Qwen3 and Llama-3.2.
+
+Run:
+    CUDA_VISIBLE_DEVICES=2 python scripts/benchmarks/verify_qwen3_numerics.py \
+        --model_name unsloth/Qwen3-4B-Base
+    CUDA_VISIBLE_DEVICES=2 python scripts/benchmarks/verify_qwen3_numerics.py \
+        --model_name unsloth/Llama-3.2-3B-Instruct
+"""
+
+from __future__ import annotations
+
+import argparse
+import copy
+import sys
+from pathlib import Path
+
+import torch
+
+HERE = Path(__file__).resolve().parent
+sys.path.insert(0, str(HERE))
+
+from qwen3_flex_inference import FlexInference, Sequence  # noqa: E402
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--model_name", required = True)
+    p.add_argument("--prompt", default = "The quick brown fox jumps over")
+    args = p.parse_args()
+
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    tok = AutoTokenizer.from_pretrained(args.model_name)
+    if tok.pad_token is None:
+        tok.pad_token = tok.eos_token
+
+    base = AutoModelForCausalLM.from_pretrained(
+        args.model_name, dtype = torch.bfloat16, attn_implementation = "eager"
+    ).to("cuda")
+    base.eval()
+
+    flex_model = copy.deepcopy(base)
+
+    ids = tok(args.prompt, return_tensors = "pt")["input_ids"].to("cuda")
+    print(f"prompt len = {ids.shape[1]}")
+
+    with torch.inference_mode():
+        out = base(ids, use_cache = False)
+        ref_logits = out.logits[0, -1, :].float()
+    print(
+        f"vanilla last-token logits: mean {ref_logits.mean().item():.4f}, "
+        f"std {ref_logits.std().item():.4f}, argmax {int(ref_logits.argmax())} "
+        f"({tok.decode([int(ref_logits.argmax())])!r})"
+    )
+
+    inf = FlexInference(
+        flex_model,
+        tok,
+        max_batch_size = 4,
+        max_seq_length = 256,
+        n_pages = 64,
+        page_size = 64,
+        max_new_tokens = 1,
+        fa4_prefill = False,
+    )
+    seq = Sequence(text = args.prompt, max_new_tokens = 1)
+    inf.tokenize([seq])
+    bi = inf.page_table.allocate()
+    inf.page_table.reserve(
+        bi,
+        torch.tensor([bi], device = "cuda", dtype = torch.long),
+        seq.total_length,
+    )
+    seq.batch_idx = bi
+    with torch.inference_mode():
+        flex_logits = inf._prefill([seq])[0].float()
+    print(
+        f"flex    last-token logits: mean {flex_logits.mean().item():.4f}, "
+        f"std {flex_logits.std().item():.4f}, argmax {int(flex_logits.argmax())} "
+        f"({tok.decode([int(flex_logits.argmax())])!r})"
+    )
+
+    diff = (flex_logits - ref_logits).abs()
+    print(f"max abs diff   = {diff.max().item():.4e}")
+    print(f"mean abs diff  = {diff.mean().item():.4e}")
+    print(f"argmax match   = {int(ref_logits.argmax()) == int(flex_logits.argmax())}")
+    top_ref = set(ref_logits.topk(10).indices.tolist())
+    top_flex = set(flex_logits.topk(10).indices.tolist())
+    print(f"top-10 overlap = {len(top_ref & top_flex)} / 10")
+
+
+if __name__ == "__main__":
+    main()

tests/test_fa4_capability_guard.py

@@ -0,0 +1,164 @@
+"""Unit test for the FA4 capability guard in FlexInference.__init__.
+
+Runs on any GPU (and on CPU) because we monkey-patch
+`torch.cuda.get_device_capability` and stub out the page-table / model
+patching that the constructor does after the guard.
+"""
+
+import os
+import sys
+import types
+import warnings
+import unittest
+from unittest import mock
+
+import torch
+
+
+REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+BENCH_DIR = os.path.join(REPO_ROOT, "scripts", "benchmarks")
+if BENCH_DIR not in sys.path:
+    sys.path.insert(0, BENCH_DIR)
+
+# qwen3_flex_inference imports heavy siblings (flex_paged_attention).
+# Stub the PageTable and patch_qwen3_model the constructor calls after the
+# guard so we don't need a real model / CUDA device.
+import qwen3_flex_inference as qfi  # noqa: E402
+
+
+class _FakePageTable:
+    def __init__(self, *a, **kw):
+        pass
+
+    def create_causal_blockmask(self, *a, **kw):
+        return None
+
+
+class _FakeTokenizer:
+    eos_token_id = 0
+
+
+def _make_fake_model(device_str = "cpu"):
+    m = types.SimpleNamespace()
+    m.device = torch.device(device_str)
+    return m
+
+
+def _build(fa4_prefill, cc_major, cc_minor = 0):
+    """Construct a FlexInference with the guard exercised.
+
+    Returns the instance. Patches torch.cuda.get_device_capability,
+    torch.zeros (to avoid CUDA allocation), PageTable, and
+    patch_qwen3_model so __init__ can run to completion without a real
+    model.
+    """
+    fake_model = _make_fake_model("cpu")
+    fake_tok = _FakeTokenizer()
+
+    _real_zeros = torch.zeros
+
+    def _fake_zeros(*a, **kw):
+        kw.pop("device", None)
+        return _real_zeros(*a, **kw)
+
+    with (
+        mock.patch.object(
+            torch.cuda, "get_device_capability", return_value = (cc_major, cc_minor)
+        ),
+        mock.patch.object(qfi, "PageTable", _FakePageTable),
+        mock.patch.object(qfi, "patch_qwen3_model", lambda *a, **kw: None),
+        mock.patch.object(torch, "zeros", _fake_zeros),
+    ):
+        return qfi.FlexInference(
+            model = fake_model,
+            tokenizer = fake_tok,
+            max_batch_size = 2,
+            max_seq_length = 128,
+            n_pages = 4,
+            page_size = 128,
+            max_new_tokens = 16,
+            fa4_prefill = fa4_prefill,
+        )
+
+
+def _fa4_warnings(caught):
+    return [
+        w
+        for w in caught
+        if issubclass(w.category, RuntimeWarning) and "fa4_prefill" in str(w.message)
+    ]
+
+
+class TestFA4CapabilityGuard(unittest.TestCase):
+    # --- explicit opt-in: --fa4_prefill=True ---
+    def test_explicit_on_sub_hopper_disables_and_warns(self):
+        with warnings.catch_warnings(record = True) as caught:
+            warnings.simplefilter("always")
+            fi = _build(fa4_prefill = True, cc_major = 8)
+        self.assertTrue(
+            _fa4_warnings(caught),
+            f"expected RuntimeWarning about fa4_prefill, got {caught!r}",
+        )
+        self.assertIs(fi.fa4_prefill, False)
+        self.assertEqual(fi.prefill_q_block, 128)
+        self.assertNotIn("BACKEND", fi.prefill_kernel_options)
+
+    def _assert_fa4_enabled(self, cc_major, fa4_prefill):
+        with warnings.catch_warnings(record = True) as caught:
+            warnings.simplefilter("always")
+            fi = _build(fa4_prefill = fa4_prefill, cc_major = cc_major)
+        self.assertEqual(
+            _fa4_warnings(caught),
+            [],
+            f"unexpected fa4 RuntimeWarning on sm_{cc_major}0 "
+            f"with fa4_prefill={fa4_prefill}: {caught!r}",
+        )
+        self.assertIs(fi.fa4_prefill, True)
+        self.assertEqual(fi.prefill_q_block, 256)
+        self.assertEqual(fi.prefill_kernel_options.get("BACKEND"), "FLASH")
+
+    def test_explicit_on_hopper_enables(self):
+        self._assert_fa4_enabled(cc_major = 9, fa4_prefill = True)
+
+    def test_explicit_on_blackwell_sm100_enables(self):
+        self._assert_fa4_enabled(cc_major = 10, fa4_prefill = True)
+
+    def test_explicit_on_blackwell_sm120_enables(self):
+        self._assert_fa4_enabled(cc_major = 12, fa4_prefill = True)
+
+    # --- auto-detect: fa4_prefill is None ---
+    def test_auto_on_sub_hopper_disables_silently(self):
+        with warnings.catch_warnings(record = True) as caught:
+            warnings.simplefilter("always")
+            fi = _build(fa4_prefill = None, cc_major = 8)
+        self.assertEqual(
+            _fa4_warnings(caught),
+            [],
+            f"auto-detect must not warn on unsupported GPU: {caught!r}",
+        )
+        self.assertIs(fi.fa4_prefill, False)
+        self.assertEqual(fi.prefill_q_block, 128)
+        self.assertNotIn("BACKEND", fi.prefill_kernel_options)
+
+    def test_auto_on_hopper_enables(self):
+        self._assert_fa4_enabled(cc_major = 9, fa4_prefill = None)
+
+    def test_auto_on_blackwell_sm100_enables(self):
+        self._assert_fa4_enabled(cc_major = 10, fa4_prefill = None)
+
+    def test_auto_on_blackwell_sm120_enables(self):
+        self._assert_fa4_enabled(cc_major = 12, fa4_prefill = None)
+
+    # --- explicit opt-out: --no-fa4_prefill / fa4_prefill=False ---
+    def test_explicit_off_on_blackwell_stays_off(self):
+        with warnings.catch_warnings(record = True) as caught:
+            warnings.simplefilter("always")
+            fi = _build(fa4_prefill = False, cc_major = 10)
+        self.assertEqual(_fa4_warnings(caught), [])
+        self.assertIs(fi.fa4_prefill, False)
+        self.assertEqual(fi.prefill_q_block, 128)
+        self.assertNotIn("BACKEND", fi.prefill_kernel_options)
+
+
+if __name__ == "__main__":
+    unittest.main()