diff --git a/scripts/benchmarks/README.md b/scripts/benchmarks/README.md new file mode 100644 index 0000000000..40bf500761 --- /dev/null +++ b/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. diff --git a/scripts/benchmarks/cb_sync_driver.py b/scripts/benchmarks/cb_sync_driver.py new file mode 100644 index 0000000000..b81a7ae1a2 --- /dev/null +++ b/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) diff --git a/scripts/benchmarks/cb_vs_vllm_generation.py b/scripts/benchmarks/cb_vs_vllm_generation.py new file mode 100644 index 0000000000..155c918ffd --- /dev/null +++ b/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|` 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() diff --git a/scripts/benchmarks/compare_grpo_runs.py b/scripts/benchmarks/compare_grpo_runs.py new file mode 100644 index 0000000000..a6e88e4158 --- /dev/null +++ b/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() diff --git a/scripts/benchmarks/flash_attn_fa4_shim.py b/scripts/benchmarks/flash_attn_fa4_shim.py new file mode 100644 index 0000000000..58b7f8026b --- /dev/null +++ b/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 diff --git a/scripts/benchmarks/flex_autotune_replay.py b/scripts/benchmarks/flex_autotune_replay.py new file mode 100644 index 0000000000..ccec61a53d --- /dev/null +++ b/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', , 32, 8, , , 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() diff --git a/scripts/benchmarks/flex_paged_attention.py b/scripts/benchmarks/flex_paged_attention.py new file mode 100644 index 0000000000..e76716ab46 --- /dev/null +++ b/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 diff --git a/scripts/benchmarks/gemma4_flex_inference.py b/scripts/benchmarks/gemma4_flex_inference.py new file mode 100644 index 0000000000..9e05aa66c5 --- /dev/null +++ b/scripts/benchmarks/gemma4_flex_inference.py @@ -0,0 +1,1148 @@ +"""Gemma-4-E2B-it inference with flex_attention + paged KV cache + CUDA graphs. + +Extends the Qwen3/Llama-3.2 engine in `qwen3_flex_inference.py` to a third +architecture, `unsloth/gemma-4-E2B-it`. Gemma-4 is not a drop-in addition: +its text backbone diverges from Qwen3/Llama in ways that cannot be folded +into a single `hasattr(self, "q_norm")` branch. The divergences, and how +this file handles them: + +1. KV-sharing layers. E2B has 35 layers; the upper 15 lack `k_proj`, + `v_proj`, `k_norm`, `v_norm` entirely and consume the K/V produced + by a "store" layer further up the stack. We link each shared layer's + `_paged_cache` to the store layer's `PagedKVCache` so flex_attention + reads the same pages that the store layer populated a few layers + earlier -- no sidecar, no SDPA fallback, one block mask per regime + works for every layer. +2. Dual attention types. Each layer is either `full_attention` + (`head_dim=512`, `rope_theta=1e6`, `partial_rotary_factor=0.25`) or + `sliding_attention` (`head_dim=256`, `rope_theta=10000`, + `sliding_window=512`). We precompute both (cos, sin) pairs once per + forward and dispatch on `self.layer_type`. +3. Per-layer input embeddings. `embed_tokens_per_layer` produces a + `[B, S, num_layers, 256]` auxiliary table that enters every layer + through a `per_layer_input_gate -> act -> mul -> per_layer_projection + -> post_per_layer_input_norm -> +residual` path after the MLP residual. +4. Four norms per layer. `input_layernorm` / `post_attention_layernorm` + wrap the attention block (double residual); `pre_feedforward_layernorm` + / `post_feedforward_layernorm` wrap the MLP (double residual). A scalar + `layer_scalar` multiplies hidden_states at layer end. +5. Final logit softcap. `logits = tanh(logits / 30.0) * 30.0` applied on + the lm_head output. + +The engine is text-only: `Gemma4ForCausalLM(text_config)` skips the +multimodal `Gemma4ForConditionalGeneration` wrapper and its vision + audio +towers entirely. Shared helpers (`PagedKVCache`, `PageTable`, `Sequence`, +`refresh_lora_merge_from_pristine`, `run_drift_verification`, +`flex_attention_compiled`, `_apply_rotary`, FA4 capability guard) are +imported from `qwen3_flex_inference.py` unchanged. + +Run: + CUDA_VISIBLE_DEVICES=6 python scripts/benchmarks/gemma4_flex_inference.py \ + --n_prompts 64 --max_new_tokens 512 --capture_cudagraph \ + --stats_path logs/flex_gemma4_bf16.json + +Requires `transformers>=5.5.0` (the `gemma4` module). The main workspace +env stays on 4.57.6; this file short-circuits with a clear install hint +if the module is missing. Use `isolated_run.py` with +`--extra_packages "transformers>=5.5.0 peft datasets"` to run on that env. +""" + +from __future__ import annotations + +import argparse +import copy +import json +import os +import sys +import time +import types +from collections import deque +from pathlib import Path +from typing import Optional + +import torch +import torch.nn.functional as F +from torch.nn.attention.flex_attention import BlockMask + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +# Shared helpers from qwen3_flex_inference.py. Import-time cost is paid once; +# we do not re-define any of these locally. +from qwen3_flex_inference import ( # noqa: E402 + DECODE_KERNEL_OPTIONS_DEFAULT, + PREFILL_KERNEL_OPTIONS_DEFAULT, + Sequence, + _apply_rotary, + _hash_state_dict, + _lora_needs_peft_fallback, + flex_attention_compiled, + refresh_lora_merge_from_pristine, + run_drift_verification, +) +from flex_paged_attention import PagedKVCache, PageTable # noqa: E402 +from torch.nn.attention.flex_attention import create_block_mask as _create_block_mask # noqa: E402 + + +# --- sliding-window block mask helpers ------------------------------------ +# +# Gemma-4's `sliding_attention` layers attend only to the last +# `sliding_window` KV positions (`q_idx - kv_idx < W`) in addition to the +# standard causal mask. The shared helpers in `flex_paged_attention.py` +# only expose the pure-causal builders, so we build sliding variants here +# (local to this file so that file is untouched). + + +def _causal_blockmask_with_window( + B: int, L: int, block_size: int, window: int, device: str +): + def causal_windowed(b, h, q_idx, kv_idx): + return (q_idx >= kv_idx) & (q_idx - kv_idx < window) + + return _create_block_mask( + causal_windowed, + B = B, + H = None, + Q_LEN = L, + KV_LEN = L, + BLOCK_SIZE = block_size, + device = device, + ) + + +def _prefill_blockmask_with_window(batch_idx: torch.Tensor, block_size, window: int): + assert batch_idx.ndim == 2 and batch_idx.shape[0] == 1 + L = batch_idx.shape[1] + docs = batch_idx.view(-1) + + def document_causal_windowed(b, h, q_idx, kv_idx): + causal_mask = q_idx >= kv_idx + window_mask = q_idx - kv_idx < window + document_mask = docs[q_idx] == docs[kv_idx] + return causal_mask & window_mask & document_mask + + return _create_block_mask( + document_causal_windowed, + B = 1, + H = None, + Q_LEN = L, + KV_LEN = L, + BLOCK_SIZE = block_size, + ) + + +# --- transformers version guard -------------------------------------------- +# +# Gemma-4 lands in `transformers>=5.5.0`. The main workspace env is on +# 4.57.6, which keeps the Qwen3 + Llama paths in qwen3_flex_inference.py +# working unchanged. We defer the import to call time so `--help` still +# works on 4.57.6. + + +def _require_gemma4(): + try: + from transformers.models.gemma4.modeling_gemma4 import Gemma4ForCausalLM + from transformers.models.gemma4.configuration_gemma4 import ( + Gemma4Config, + Gemma4TextConfig, + ) + except ImportError: + import transformers + + raise SystemExit( + f"Gemma-4 requires transformers>=5.5.0 (`gemma4` module). " + f"Current: transformers=={transformers.__version__}. " + f"Install: uv pip install 'transformers>=5.5.0'" + ) + return Gemma4ForCausalLM, Gemma4Config, Gemma4TextConfig + + +# --- attention forward factory -------------------------------------------- + + +def _apply_rotary_q(q, cos, sin): + """Rotary on Q alone; used on shared layers where K is read + pre-rotated from the store layer's paged cache.""" + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + + def rotate_half(x): + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim = -1) + + return (q * cos) + (rotate_half(q) * sin) + + +def make_flex_gemma4_attention_forward(page_table: PageTable): + """Return a new `forward` method for `Gemma4TextAttention` that routes + through flex_attention against a paged KV cache. + + Two layer kinds: + - non-shared (`self.is_kv_shared_layer == False`): standard q/k/v + projection; writes new K/V into `self._paged_cache` + (which is the layer's own PagedKVCache). + - shared (`self.is_kv_shared_layer == True`): no `k_proj`/ + `v_proj`/`k_norm`/`v_norm`. Reads K/V directly from + the *store* layer's paged cache, which the patching + helper has already linked onto `self._paged_cache`. + No write -- the store layer populated the cache for + the same positions earlier in the walker, so the + shared layer just attends over those pages with the + identical block mask. + + Linking shared layers to the store layer's paged cache keeps one + block-mask + one KV layout across the whole stack and lets + flex_attention handle every layer uniformly (no sidecar, no SDPA + fallback, one CUDA graph capture). + + `position_embeddings` is a dict keyed by `layer_type`; we pick the + right (cos, sin) pair before rotary. + + `self.v_proj` may be None when the config sets `attention_k_eq_v` + (global head dim with shared K=V); that branch reuses k_raw. + """ + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: dict, + attention_mask = None, + past_key_values = None, + cache_position = None, + flex_block_mask: Optional[dict] = None, + flex_input_pos: Optional[torch.Tensor] = None, + flex_batch_idx: Optional[torch.Tensor] = None, + flex_kernel_options: Optional[dict] = None, + **kwargs, + ): + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + cos, sin = position_embeddings[self.layer_type] + + q = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2) + + if getattr(self, "is_kv_shared_layer", False): + # Shared layer reads from the paired store layer's K/V. + # `PagedKVCache.update` returns different shapes for prefill + # vs decode: + # - prefill: the packed k_val/v_val [1, H, L_packed, D] + # - decode : the full paged pool k_cache/v_cache + # [1, H, n_pages*page_size, D] + # The prefill block_mask is sized for L_packed and the decode + # block_mask is sized for the paged pool, so we need to match + # the same shape here. The store layer stashes its + # post-rotary k/v as `_last_k_val` / `_last_v_val` during + # prefill; at decode time we read from its `_paged_cache` + # (the same buffer the shared layer was linked to at patch). + q = _apply_rotary_q(q, cos, sin) + store_attn = self._store_attn + if q.shape[-2] > 1: + k = store_attn._last_k_val + v = store_attn._last_v_val + else: + k = self._paged_cache.k_cache + v = self._paged_cache.v_cache + else: + k_raw = self.k_proj(hidden_states).view(hidden_shape) + k = self.k_norm(k_raw).transpose(1, 2) + # `v_proj` may be None under Gemma-4's K=V global-attention + # option; in that case reuse the raw (un-normed) k projection. + if self.v_proj is not None: + v = self.v_norm( + self.v_proj(hidden_states).view(hidden_shape) + ).transpose(1, 2) + else: + v = k_raw.transpose(1, 2) + q, k = _apply_rotary(q, k, cos, sin) + + # Store layers stash the post-rotary k/v so any shared + # successors can read the same packed prefill tensors. This + # assignment is a pointer rebind, not a copy; CUDA graph + # capture sees a stable attribute reference. Plain + # non-shared layers don't need this. + if getattr(self, "store_full_length_kv", False): + self._last_k_val = k + self._last_v_val = v + + if self._paged_cache is not None and flex_input_pos is not None: + k, v = self._paged_cache.update(flex_input_pos, k, v, flex_batch_idx) + + # flex_block_mask is a dict keyed by layer_type; pick the one + # matching this layer's regime (full_attention vs sliding_attention). + block_mask = flex_block_mask[self.layer_type] + attn_output = flex_attention_compiled( + q, + k, + v, + scale = self.scaling, + block_mask = block_mask, + enable_gqa = True, + kernel_options = flex_kernel_options, + ) + attn_output = attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous() + return self.o_proj(attn_output), None + + return forward + + +def patch_gemma4_attention_forwards(model: torch.nn.Module, page_table: PageTable): + """Attach a PagedKVCache to every non-shared attention layer, link + every shared attention layer to its store layer's cache, and swap in + the flex_attention forward above. + + Three passes: + 1. Allocate a PagedKVCache on each non-shared layer (the cache + shape depends on that layer's head_dim and num_kv_heads, which + vary across Gemma-4 layers). + 2. Walk shared layers and set + `shared._paged_cache = store._paged_cache`, where `store` is + `model.model.layers[shared.kv_shared_layer_index]`. The shared + layer's forward reads `k_cache` / `v_cache` directly; the store + layer's `update()` writes populate the same tensors. + 3. Bind the flex forward. + """ + fwd = make_flex_gemma4_attention_forward(page_table) + for layer in model.model.layers: + attn = layer.self_attn + if getattr(attn, "is_kv_shared_layer", False): + continue + n_kv = attn.k_proj.out_features // attn.head_dim + attn._paged_cache = PagedKVCache( + page_table, + n_heads = n_kv, + head_dim = attn.head_dim, + dtype = model.dtype, + ).to(model.device) + for layer in model.model.layers: + attn = layer.self_attn + if not getattr(attn, "is_kv_shared_layer", False): + continue + store_attn = model.model.layers[attn.kv_shared_layer_index].self_attn + attn._paged_cache = store_attn._paged_cache + attn._store_attn = store_attn + for layer in model.model.layers: + layer.self_attn.forward = types.MethodType(fwd, layer.self_attn) + + +# --- model forward walker -------------------------------------------------- + + +def call_gemma4_model_with_flex_kwargs(model, input_ids, position_ids, flex_kwargs): + """Walk the Gemma-4 text model manually so we can inject flex_* kwargs + into each attention call. Mirrors `call_model_with_flex_kwargs` in + `qwen3_flex_inference.py` but: + + - precomputes both (cos, sin) variants and passes them as a dict + keyed by `layer.self_attn.layer_type`; + - materializes `per_layer_inputs` via the model's own + `get_per_layer_inputs` + `project_per_layer_inputs` helpers; + - runs the double-residual around attn, the double-residual around + MLP, the per-layer-input path, and the `layer_scalar` multiply. + + The final norm is applied here; lm_head + softcap is applied by the + caller (so we can slice to logits_positions before the vocab matmul). + """ + base = model.model + + inputs_embeds = base.embed_tokens(input_ids) + + # Per-layer input table. E2B has hidden_size_per_layer_input = 256 and + # 35 layers, so this is [B, S, 35, 256] -- a local tensor with fixed + # shape across CUDA graph replays (input_ids is pre-allocated upstream). + per_layer_inputs = None + if getattr(base, "hidden_size_per_layer_input", 0): + per_layer_inputs = base.get_per_layer_inputs(input_ids, inputs_embeds) + per_layer_inputs = base.project_per_layer_inputs( + inputs_embeds, per_layer_inputs + ) + + position_embeddings = { + layer_type: base.rotary_emb(inputs_embeds, position_ids, layer_type) + for layer_type in base.unique_layer_types + } + + hidden_states = inputs_embeds + for i, layer in enumerate(base.layers): + residual = hidden_states + hidden_states = layer.input_layernorm(hidden_states) + hidden_states, _ = layer.self_attn( + hidden_states, + position_embeddings = position_embeddings, + **flex_kwargs, + ) + hidden_states = layer.post_attention_layernorm(hidden_states) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = layer.pre_feedforward_layernorm(hidden_states) + hidden_states = layer.mlp(hidden_states) + hidden_states = layer.post_feedforward_layernorm(hidden_states) + hidden_states = residual + hidden_states + + if per_layer_inputs is not None and hasattr(layer, "per_layer_input_gate"): + residual = hidden_states + hidden_states = layer.per_layer_input_gate(hidden_states) + hidden_states = layer.act_fn(hidden_states) + hidden_states = hidden_states * per_layer_inputs[:, :, i, :] + hidden_states = layer.per_layer_projection(hidden_states) + hidden_states = layer.post_per_layer_input_norm(hidden_states) + hidden_states = residual + hidden_states + + if hasattr(layer, "layer_scalar"): + hidden_states = hidden_states * layer.layer_scalar + + hidden_states = base.norm(hidden_states) + return hidden_states + + +# --- inference engine ------------------------------------------------------ + + +class FlexGemma4Inference: + def __init__( + self, + model, + tokenizer, + max_batch_size = 32, + max_seq_length = 2048, + n_pages = 2048, + page_size = 128, + max_new_tokens = 512, + decode_kernel_options = None, + prefill_kernel_options = None, + fa4_prefill = None, + base_model = None, + peft_model = None, + ): + assert max_seq_length % page_size == 0 + self.model = model + self.tokenizer = tokenizer + self.device = model.device + self.eos_token_id = tokenizer.eos_token_id + self.base_model = base_model + self.peft_model = peft_model + self.max_batch_size = max_batch_size + self.max_seq_length = max_seq_length + self.page_size = page_size + self.max_new_tokens = max_new_tokens + + if fa4_prefill is None or fa4_prefill: + major, _ = torch.cuda.get_device_capability(self.device) + supported = major >= 9 + if fa4_prefill and not supported: + import warnings + + warnings.warn( + f"--fa4_prefill needs Hopper (sm_90) or Blackwell " + f"(sm_100 / sm_120); found sm_{major}0. Falling back " + f"to the Triton flex_attention backend.", + RuntimeWarning, + stacklevel = 2, + ) + fa4_prefill = supported + self.fa4_prefill = fa4_prefill + self.prefill_q_block = 256 if fa4_prefill else 128 + self.prefill_kv_block = 128 + + self.decode_kernel_options = ( + decode_kernel_options + if decode_kernel_options is not None + else DECODE_KERNEL_OPTIONS_DEFAULT + ) + base_prefill_opts = ( + prefill_kernel_options + if prefill_kernel_options is not None + else dict(PREFILL_KERNEL_OPTIONS_DEFAULT) + ) + if fa4_prefill: + base_prefill_opts = dict(base_prefill_opts) + base_prefill_opts.pop("FORCE_USE_FLEX_ATTENTION", None) + base_prefill_opts["BACKEND"] = "FLASH" + self.prefill_kernel_options = base_prefill_opts + + self.page_table = PageTable( + n_pages = n_pages, + page_size = page_size, + max_batch_size = max_batch_size, + device = self.device.type, + ) + + patch_gemma4_attention_forwards(model, self.page_table) + + self.input_pos_buffer = torch.zeros( + max_batch_size, dtype = torch.int32, device = self.device + ) + + # Detect the sliding window from any sliding-attention layer. + # We need one block mask per layer type: `full_attention` is pure + # causal; `sliding_attention` is causal AND q_pos - kv_pos < window. + sliding_window = None + for layer in model.model.layers: + if getattr(layer.self_attn, "is_sliding", False): + sliding_window = layer.self_attn.sliding_window + break + + self.sliding_window = sliding_window + self.block_mask_logical_by_type = { + "full_attention": self.page_table.create_causal_blockmask( + B = max_batch_size, L = max_seq_length + ), + } + if sliding_window is not None: + self.block_mask_logical_by_type["sliding_attention"] = ( + _causal_blockmask_with_window( + B = max_batch_size, + L = max_seq_length, + block_size = page_size, + window = sliding_window, + device = self.device.type, + ) + ) + # Legacy alias used by the original page-aware decode slicer. + self.block_mask_logical = self.block_mask_logical_by_type["full_attention"] + + self.cudagraph_captured = False + self.graphs = {} + self.graph_vars = {} + + def tokenize(self, sequences): + for seq in sequences: + ids = self.tokenizer(seq.text, return_tensors = "pt")["input_ids"].squeeze(0) + seq.input_ids = ids + seq.input_length = ids.shape[0] + + def _softcap(self, logits): + sc = getattr(self.model.config, "final_logit_softcapping", None) + if sc is not None and sc > 0: + logits = torch.tanh(logits / sc) * sc + return logits + + def _prefill(self, batch: list) -> torch.Tensor: + input_ids_list = [seq.input_ids.to(self.device) for seq in batch] + input_pos_list = [ + torch.arange(seq.input_length, dtype = torch.long, device = self.device) + for seq in batch + ] + batch_idx_list = [ + torch.full( + (seq.input_length,), + seq.batch_idx, + dtype = torch.long, + device = self.device, + ) + for seq in batch + ] + input_ids = torch.cat(input_ids_list).view(1, -1) + input_pos = torch.cat(input_pos_list).view(1, -1) + batch_idx = torch.cat(batch_idx_list).view(1, -1) + + L = input_ids.shape[1] + q_block = self.prefill_q_block + pad = (q_block - L % q_block) % q_block + if pad > 0: + input_ids = F.pad(input_ids, (0, pad), value = 0) + input_pos = F.pad(input_pos, (0, pad), value = 0) + batch_idx = F.pad(batch_idx, (0, pad), value = 0) + + input_lengths = torch.tensor( + [s.input_length for s in batch], dtype = torch.long, device = self.device + ) + logits_positions = input_lengths.cumsum(dim = 0) - 1 + + prefill_block_size = ( + (self.prefill_q_block, self.prefill_kv_block) + if self.fa4_prefill + else self.prefill_q_block + ) + mask_full = self.page_table.create_prefill_blockmask_no_paging( + batch_idx, BLOCK_SIZE = prefill_block_size + ) + masks_by_type = {"full_attention": mask_full} + if self.sliding_window is not None: + masks_by_type["sliding_attention"] = _prefill_blockmask_with_window( + batch_idx, + block_size = prefill_block_size, + window = self.sliding_window, + ) + + flex_kwargs = dict( + flex_block_mask = masks_by_type, + flex_input_pos = input_pos, + flex_batch_idx = batch_idx, + flex_kernel_options = self.prefill_kernel_options, + ) + hidden = call_gemma4_model_with_flex_kwargs( + self.model, input_ids, input_pos, flex_kwargs + ) + logits = self.model.lm_head(hidden[:, logits_positions, :]).squeeze(0) + return self._softcap(logits) + + def _decode_block_mask(self, batch_idx: torch.Tensor): + """Slice one row of the logical decode mask per sequence, for + both full and sliding regimes. Returns a dict keyed by layer_type + plus the raw `input_pos` tensor (needed for PageTable conversion).""" + input_pos = self.input_pos_buffer[batch_idx] + assert batch_idx.ndim == 1 and input_pos.ndim == 1 + B = batch_idx.shape[0] + + def _slice(block_mask, extra_mask_mod): + input_block_idx = input_pos // block_mask.BLOCK_SIZE[0] + kv_num_blocks = block_mask.kv_num_blocks[ + batch_idx, :, input_block_idx + ].view(B, 1, 1) + kv_indices = block_mask.kv_indices[batch_idx, :, input_block_idx].view( + B, 1, 1, -1 + ) + full_num = full_idx = None + if block_mask.full_kv_num_blocks is not None: + full_num = block_mask.full_kv_num_blocks[ + batch_idx, :, input_block_idx + ].view(B, 1, 1) + full_idx = block_mask.full_kv_indices[ + batch_idx, :, input_block_idx + ].view(B, 1, 1, -1) + + seq_length = (1, block_mask.seq_lengths[1]) + return BlockMask.from_kv_blocks( + kv_num_blocks, + kv_indices, + full_num, + full_idx, + BLOCK_SIZE = block_mask.BLOCK_SIZE, + mask_mod = extra_mask_mod, + seq_lengths = seq_length, + ) + + def causal_offset(off): + def m(b, h, q_idx, kv_idx): + return q_idx + off[b] >= kv_idx + + return m + + def causal_offset_windowed(off, window): + def m(b, h, q_idx, kv_idx): + return (q_idx + off[b] >= kv_idx) & (q_idx + off[b] - kv_idx < window) + + return m + + masks = { + "full_attention": _slice( + self.block_mask_logical_by_type["full_attention"], + causal_offset(input_pos), + ), + } + if self.sliding_window is not None: + masks["sliding_attention"] = _slice( + self.block_mask_logical_by_type["sliding_attention"], + causal_offset_windowed(input_pos, self.sliding_window), + ) + return masks, input_pos + + def _decode_step_eager(self, batch_idx: torch.Tensor, input_ids: torch.Tensor): + B = input_ids.shape[0] + masks, input_pos = self._decode_block_mask(batch_idx) + # Convert each regime's block mask through the page table so the + # logical→physical kv page mapping is correct for every layer. + masks = { + k: self.page_table.convert_logical_block_mask(m, batch_idx) + for k, m in masks.items() + } + position_ids = input_pos.view(B, 1).to(torch.long) + flex_kwargs = dict( + flex_block_mask = masks, + flex_input_pos = input_pos.view(B, 1).to(torch.long), + flex_batch_idx = batch_idx, + flex_kernel_options = self.decode_kernel_options, + ) + hidden = call_gemma4_model_with_flex_kwargs( + self.model, input_ids.view(B, 1), position_ids, flex_kwargs + ) + logits = self.model.lm_head(hidden[:, -1, :]) + return self._softcap(logits) + + def _decode_step( + self, + batch_idx: torch.Tensor, + input_ids: torch.Tensor, + input_pos: torch.Tensor, + ): + self.input_pos_buffer.zero_() + self.input_pos_buffer[batch_idx] = input_pos + if not self.cudagraph_captured: + return self._decode_step_eager(batch_idx, input_ids) + bs = input_ids.size(0) + key = next(x for x in self.graph_bs if x >= bs) + graph = self.graphs[key] + gv = self.graph_vars + for k, v in gv.items(): + if k != "outputs": + v.zero_() + gv["input_ids"][:bs] = input_ids + gv["batch_idx"][:bs] = batch_idx + graph.replay() + return gv["outputs"][:bs] + + def capture_decode_cudagraph(self): + max_bs = self.max_batch_size + reserved_batches = [] + for bi in range(1, max_bs): + try: + allocated = self.page_table.allocate() + self.page_table.reserve( + allocated, + torch.tensor([allocated], device = self.device, dtype = torch.long), + self.page_size, + ) + reserved_batches.append(allocated) + except Exception: + break + + input_ids = torch.zeros(max_bs, dtype = torch.int64, device = self.device) + batch_idx = torch.arange(max_bs, dtype = torch.int64, device = self.device) + outputs = torch.zeros( + (max_bs, self.model.config.vocab_size), + dtype = self.model.dtype, + device = self.device, + ) + self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16)) + pool = None + for bs in reversed(self.graph_bs): + if bs > max_bs: + continue + print(f"[flex-gemma4] capturing CUDA graph for bs={bs}") + torch.cuda.synchronize() + _ = self._decode_step_eager(batch_idx[:bs], input_ids[:bs]) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool): + outputs[:bs] = self._decode_step_eager(batch_idx[:bs], input_ids[:bs]) + if pool is None: + pool = graph.pool() + self.graphs[bs] = graph + torch.cuda.synchronize() + for bi in reserved_batches: + self.page_table.erase(bi) + self.graph_vars = dict( + input_ids = input_ids, batch_idx = batch_idx, outputs = outputs + ) + + def refresh_inference_from_base(self): + if self.base_model is None or self.peft_model is None: + return 0 + return refresh_lora_merge_from_pristine(self.base_model, self.peft_model) + + @torch.inference_mode() + def generate(self, sequences, capture_cudagraph = False): + self.tokenize(sequences) + waiting = deque(sequences) + running = deque() + done = [] + + if capture_cudagraph and not self.cudagraph_captured: + self.capture_decode_cudagraph() + self.cudagraph_captured = True + + while waiting or running: + batch = [] + while waiting and self.page_table.can_reserve(waiting[0].total_length): + seq = waiting.popleft() + bi = self.page_table.allocate() + self.page_table.reserve( + bi, + torch.tensor([bi], device = self.device, dtype = torch.long), + seq.total_length, + ) + seq.batch_idx = bi + batch.append(seq) + if batch: + logits = self._prefill(batch) + next_ids = torch.argmax(logits, dim = -1).tolist() + for i, seq in enumerate(batch): + seq.last_token_id = next_ids[i] + seq.output_ids.append(next_ids[i]) + if ( + seq.last_token_id == self.eos_token_id + or len(seq.output_ids) >= seq.max_new_tokens + ): + seq.finished = True + done.append(seq) + self.page_table.erase(seq.batch_idx) + else: + running.append(seq) + continue + + decode_batch = [] + while running: + seq = running.popleft() + if self.page_table.capacity[seq.batch_idx] >= seq.total_length: + decode_batch.append(seq) + elif self.page_table.can_reserve( + seq.total_length, batch_idx_int = seq.batch_idx + ): + self.page_table.reserve( + seq.batch_idx, + torch.tensor( + [seq.batch_idx], + device = self.device, + dtype = torch.long, + ), + seq.total_length, + ) + decode_batch.append(seq) + else: + running.appendleft(seq) + newest = running.pop() + waiting.appendleft(newest) + self.page_table.erase(newest.batch_idx) + if not decode_batch: + continue + + B = len(decode_batch) + bi_tensor = torch.tensor( + [s.batch_idx for s in decode_batch], + dtype = torch.long, + device = self.device, + ) + last_ids = torch.tensor( + [s.last_token_id for s in decode_batch], + dtype = torch.long, + device = self.device, + ) + cur_pos = torch.tensor( + [s.total_length - 1 for s in decode_batch], + dtype = torch.int32, + device = self.device, + ) + logits = self._decode_step(bi_tensor, last_ids, cur_pos) + next_ids = torch.argmax(logits, dim = -1).tolist() + for i, seq in enumerate(decode_batch): + seq.last_token_id = next_ids[i] + seq.output_ids.append(next_ids[i]) + if ( + seq.last_token_id == self.eos_token_id + or len(seq.output_ids) >= seq.max_new_tokens + ): + seq.finished = True + done.append(seq) + self.page_table.erase(seq.batch_idx) + else: + running.append(seq) + + return done + + +# --- CLI ------------------------------------------------------------------- + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--model_name", default = "unsloth/gemma-4-E2B-it") + p.add_argument("--n_prompts", type = int, default = 64) + p.add_argument("--n_rounds", type = int, default = 2) + p.add_argument("--max_new_tokens", type = int, default = 512) + p.add_argument("--max_batch_size", type = int, default = 64) + p.add_argument("--max_seq_length", type = int, default = 2048) + p.add_argument("--n_pages", type = int, default = 2048) + p.add_argument("--page_size", type = int, default = 128) + p.add_argument("--capture_cudagraph", action = "store_true") + p.add_argument("--lora_adapter", default = None) + p.add_argument("--decode_kernel_options", default = None) + p.add_argument("--prefill_kernel_options", default = None) + p.add_argument( + "--fa4_prefill", + default = None, + action = argparse.BooleanOptionalAction, + help = ( + "Use BLOCK_SIZE=(256,128) + BACKEND=FLASH on prefill. Default " + "auto-enables on Hopper (sm_90) and Blackwell (sm_100, sm_120)." + ), + ) + p.add_argument("--load_in_4bit", action = "store_true") + p.add_argument( + "--no_merge_lora", + action = "store_true", + help = "Keep the LoRA adapter as a PEFT wrapper instead of merging.", + ) + p.add_argument( + "--verify_no_drift", + action = "store_true", + help = "Drift-verify the double-copy LoRA refresh across N cycles.", + ) + p.add_argument("--verify_iterations", type = int, default = 10) + p.add_argument("--model_name_4bit", default = None) + p.add_argument("--stats_path", required = True) + p.add_argument( + "--chat_template", + choices = ["auto", "grpo", "native"], + default = "auto", + help = ( + "Which chat template to use. `auto`: native for Gemma-4. " + "`grpo`: force the GRPO template. `native`: force the " + "tokenizer's built-in template." + ), + ) + args = p.parse_args() + + def _parse_opts(s): + if s is None: + return None + return json.loads(s) + + Gemma4ForCausalLM, Gemma4Config, Gemma4TextConfig = _require_gemma4() + + from transformers import AutoTokenizer + from transformers.models.gemma4.modeling_gemma4 import ( + Gemma4ForConditionalGeneration, + ) + + tok = AutoTokenizer.from_pretrained(args.model_name) + if tok.pad_token is None: + tok.pad_token = tok.eos_token + + base_model = None + peft_model = None + + if args.load_in_4bit: + from transformers import AutoModelForCausalLM + from huggingface_hub import HfApi + + bnb_model_name = args.model_name_4bit or f"{args.model_name}-unsloth-bnb-4bit" + # Probe for the 4-bit shard. If missing, the user asked for a + # quant row we cannot produce; fail loudly rather than silently + # falling back to bf16 (which would mislabel the stats file). + try: + HfApi().model_info(bnb_model_name) + except Exception as e: + raise SystemExit( + f"[flex-gemma4] --load_in_4bit: 4-bit shard {bnb_model_name} " + f"is not available ({e}). Use --model_name_4bit to override " + f"or drop --load_in_4bit for bf16." + ) + print(f"[flex-gemma4] loading 4-bit base: {bnb_model_name}") + # AutoModelForCausalLM resolves to `Gemma4ForConditionalGeneration` + # for Gemma-4 -- mirror the bf16 path and move the + # language_model into a ForCausalLM shell so downstream code can + # reach `model.model.layers` / `model.model.embed_tokens`. + full_model = AutoModelForCausalLM.from_pretrained( + bnb_model_name, + attn_implementation = "eager", + device_map = "cuda:0", + ) + if hasattr(full_model.model, "language_model"): + lang_model = full_model.model.language_model + full_model.model.vision_tower = None + full_model.model.audio_tower = None + full_model.model.embed_vision = None + full_model.model.embed_audio = None + text_cfg = full_model.config.text_config + model = Gemma4ForCausalLM(text_cfg) + model.model = lang_model + model.lm_head.weight = lang_model.embed_tokens.weight + else: + model = full_model + if getattr(model.config, "tie_word_embeddings", False): + model.lm_head.weight = model.model.embed_tokens.weight + model.eval() + del full_model + if args.lora_adapter: + from peft import PeftModel + + peft_wrapper = PeftModel.from_pretrained( + model, + str(Path(args.lora_adapter).resolve()), + is_trainable = False, + ) + model = peft_wrapper.base_model.model + else: + # Text-only bf16. The checkpoint stores text weights under the + # `model.language_model.` prefix (Gemma-4 is natively multimodal). + # We load the full `Gemma4ForConditionalGeneration`, pluck the + # language_model, then drop the vision + audio towers before + # moving to GPU so peak memory reflects text only. + full_cfg = Gemma4Config.from_pretrained(args.model_name) + text_cfg = full_cfg.text_config + + full_model = Gemma4ForConditionalGeneration.from_pretrained( + args.model_name, + dtype = torch.bfloat16, + attn_implementation = "eager", + ) + lang_model = full_model.model.language_model + # Drop the non-text towers. `embed_vision` / `embed_audio` project + # from text hidden size -- harmless when not invoked, but we kill + # them too so deepcopy (below) stays cheap. + full_model.model.vision_tower = None + full_model.model.audio_tower = None + full_model.model.embed_vision = None + full_model.model.embed_audio = None + + # Build a ForCausalLM shell around the language_model so LoRA / + # state-dict hashing treat it like any other HF decoder model. + base_model = Gemma4ForCausalLM(text_cfg) + base_model.model = lang_model + base_model.lm_head.weight = lang_model.embed_tokens.weight + base_model = base_model.to(torch.bfloat16).to("cuda") + base_model.eval() + del full_model + + if not args.lora_adapter: + model = base_model + base_model = None + elif args.no_merge_lora: + from peft import PeftModel + + peft_wrapper = PeftModel.from_pretrained( + base_model, + str(Path(args.lora_adapter).resolve()), + is_trainable = False, + ) + model = peft_wrapper.base_model.model + base_model = None + else: + from peft import PeftModel + + print("[flex-gemma4] deep-copying base model for double-copy rollout") + inference_model = copy.deepcopy(base_model) + inference_model.eval() + peft_model = PeftModel.from_pretrained( + inference_model, + str(Path(args.lora_adapter).resolve()), + is_trainable = False, + ) + model = peft_model.base_model.model + model.eval() + + if args.verify_no_drift: + if args.load_in_4bit: + raise SystemExit( + "--verify_no_drift only applies to the bf16 double-copy path." + ) + if args.no_merge_lora: + raise SystemExit("--verify_no_drift is incompatible with --no_merge_lora.") + if base_model is None or peft_model is None: + raise SystemExit( + "--verify_no_drift requires --lora_adapter against the bf16 path." + ) + print( + f"[flex-gemma4] running drift verification: " + f"{args.verify_iterations} perturb+refresh cycles" + ) + result = run_drift_verification( + base_model, peft_model, n_iters = args.verify_iterations + ) + result = {"mode": "verify_no_drift", "model_name": args.model_name, **result} + 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(result, f, indent = 2) + print(json.dumps(result, indent = 2)) + os._exit(0) + + from unsloth_grpo_common import SYSTEM_PROMPT, apply_chat_template_to_tokenizer + from datasets import load_dataset + + if args.chat_template == "auto": + # Gemma-4 default is the tokenizer native template; only Qwen3 + # in this repo uses GRPO-by-default. + use_grpo = False + elif args.chat_template == "grpo": + use_grpo = True + else: + use_grpo = False + if use_grpo: + apply_chat_template_to_tokenizer(tok) + print("[flex-gemma4] chat_template: GRPO") + else: + print("[flex-gemma4] chat_template: tokenizer native") + 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 + ] + texts = [ + tok.apply_chat_template(m, add_generation_prompt = True, tokenize = False) + for m in messages + ] + + inference = FlexGemma4Inference( + model, + tok, + max_batch_size = args.max_batch_size, + max_seq_length = args.max_seq_length, + n_pages = args.n_pages, + page_size = args.page_size, + max_new_tokens = args.max_new_tokens, + decode_kernel_options = _parse_opts(args.decode_kernel_options), + prefill_kernel_options = _parse_opts(args.prefill_kernel_options), + fa4_prefill = args.fa4_prefill, + base_model = base_model, + peft_model = peft_model, + ) + + if inference.base_model is not None and inference.peft_model is not None: + n = inference.refresh_inference_from_base() + print(f"[flex-gemma4] double-copy rollout: refreshed {n} LoRA-target layers") + + def make_seqs(): + return [Sequence(text = t, max_new_tokens = args.max_new_tokens) for t in texts] + + torch.cuda.reset_peak_memory_stats() + print("[flex-gemma4] warmup (16 prompts)...") + _ = inference.generate(make_seqs()[:16], capture_cudagraph = args.capture_cudagraph) + torch.cuda.synchronize() + + wall_times = [] + total_decoded = 0 + for r in range(args.n_rounds): + torch.cuda.synchronize() + t0 = time.perf_counter() + out = inference.generate(make_seqs()) + torch.cuda.synchronize() + wall_times.append(time.perf_counter() - t0) + total_decoded = sum(len(s.output_ids) for s in out) + print( + f"[flex-gemma4] round {r}: {wall_times[-1]:.2f}s, {total_decoded} " + f"tokens, {total_decoded / wall_times[-1]:.1f} tok/s" + ) + + med = sorted(wall_times)[len(wall_times) // 2] + best = min(wall_times) + peak = torch.cuda.max_memory_allocated() / 1024**3 + sample_completions = [] + for s in out[:3]: + sample_completions.append( + tok.decode(s.output_ids[:80], skip_special_tokens = True) + ) + res = { + "backend": "flex-gemma4", + "model_name": args.model_name, + "capture_cudagraph": args.capture_cudagraph, + "lora_adapter": args.lora_adapter, + "n_prompts": args.n_prompts, + "n_decoded_tokens": total_decoded, + "wall_times_s": wall_times, + "median_wall_s": med, + "best_wall_s": best, + "decode_tps_median": total_decoded / med if med else 0, + "decode_tps_best": total_decoded / best if best else 0, + "max_new_tokens": args.max_new_tokens, + "peak_memory_gb": peak, + "sample_completions": sample_completions, + } + 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(res, f, indent = 2) + print(json.dumps(res, indent = 2)) + os._exit(0) + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmarks/make_lora_adapter.py b/scripts/benchmarks/make_lora_adapter.py new file mode 100644 index 0000000000..ef646e94e8 --- /dev/null +++ b/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() diff --git a/scripts/benchmarks/persistent_cb.py b/scripts/benchmarks/persistent_cb.py new file mode 100644 index 0000000000..8ac9e34b4d --- /dev/null +++ b/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 diff --git a/scripts/benchmarks/qwen3_flex_inference.py b/scripts/benchmarks/qwen3_flex_inference.py new file mode 100644 index 0000000000..e6b8dcf2ea --- /dev/null +++ b/scripts/benchmarks/qwen3_flex_inference.py @@ -0,0 +1,1249 @@ +"""Llama / Qwen3 inference with flex_attention + paged KV cache + CUDA graphs. + +The transformers continuous-batching path tops out at ~10% of vLLM on this +workload because `_generation_step` is Python-heavy (scheduler + paged +attention dispatch per layer + per-request metadata updates). torch.compile +chokes on it (700+ recompile storm, see Phase 4). + +flex-nano-vllm (Chang, 2024) hits 90% of vLLM in 1000 lines of pure PyTorch +by building paged attention on top of `torch.nn.attention.flex_attention`: + +1. The paged KV cache is a single contiguous [1, H, num_pages*page_size, D] + tensor; logical<->physical mapping lives in a PageTable. +2. flex_attention's BlockMask lets us route queries to physical pages via + mask_mod + score_mod callbacks, which compile cleanly. +3. One CUDA graph per batch-size bucket captured during warmup; dispatch to + the nearest bucket on each decode step and pad with batch_idx=0 (reserved + as a no-op slot). + +This file runs the architecture on Qwen3 and Llama-3.2. The attention +forward is monkey-patched to use our PagedKVCache, and the inference loop +runs prefill + decode on the main thread (no background worker, graph +replay works end-to-end). The only arch-specific branch is a per-head QK +RMSNorm that Qwen3 has and Llama does not; everything else (q/k/v/o proj, +head_dim, scaling, rotary_emb, embed_tokens, layers, final norm) is +identical attribute-for-attribute across the two families. + +LoRA: the bf16 path uses a **double-copy rollout pattern** when +`--lora_adapter` is set. A pristine `base_model` lives on GPU alongside a +deep-copy `inference_model` (wrapped by PEFT). Before each rollout -- +or at setup time, here -- the inference copy's LoRA-target base weights +are restored in-place from pristine, then `merge_adapter()` is called +fresh. We never call `unmerge_adapter()`. This avoids the ~1 ULP bf16 +drift per merge/unmerge cycle that would otherwise corrupt the base +model across hundreds of GRPO iterations. `--verify_no_drift` hashes the +base params before and after N cycles and asserts bit-identical. + +Run: + CUDA_VISIBLE_DEVICES=6 python scripts/benchmarks/qwen3_flex_inference.py \ + --n_prompts 32 --max_new_tokens 512 --stats_path logs/qwen3_flex.json + + CUDA_VISIBLE_DEVICES=7 python scripts/benchmarks/qwen3_flex_inference.py \ + --model_name unsloth/Llama-3.2-3B-Instruct --chat_template native \ + --n_prompts 32 --max_new_tokens 512 --stats_path logs/llama32_flex.json + +Add `--capture_cudagraph` to capture per-batch-size decode graphs during +warmup. +""" + +from __future__ import annotations + +import argparse +import copy +import hashlib +import json +import os +import sys +import time +from collections import deque +from dataclasses import dataclass +from pathlib import Path +from typing import Optional + +import torch +import torch.nn.functional as F +from torch.nn.attention.flex_attention import BlockMask, flex_attention + +HERE = Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) + +from flex_paged_attention import PagedKVCache, PageTable # noqa: E402 + +# Compile flex_attention once at import for warm caches. `fullgraph=True` is +# required for the decode CUDA graph capture to be worth anything. +# Allow an environment override to try `mode="max-autotune"` for the kernel +# template search -- pays off on steady-state decode but adds ~minutes of +# warmup time at first import. +_FLEX_COMPILE_MODE = os.environ.get("FLEX_COMPILE_MODE", None) +if _FLEX_COMPILE_MODE: + flex_attention_compiled = torch.compile( + flex_attention, + fullgraph = True, + mode = _FLEX_COMPILE_MODE, + ) +else: + flex_attention_compiled = torch.compile(flex_attention, fullgraph = True) + + +def _apply_rotary(q, k, cos, sin): + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + + def rotate_half(x): + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim = -1) + + q = (q * cos) + (rotate_half(q) * sin) + k = (k * cos) + (rotate_half(k) * sin) + return q, k + + +def make_flex_attention_forward(page_table: PageTable): + """Return a new `forward` method for a decoder-only attention layer + (Qwen3Attention or LlamaAttention) that uses flex_attention against a + paged KV cache. The returned closure captures the shared PageTable; + each layer gets its own PagedKVCache attached to the module as + `self._paged_cache`. + + The only arch-specific branch is Qwen3's per-head QK RMSNorm + (`self.q_norm` / `self.k_norm`), applied after proj+reshape but + before rotary. Llama has no QK-norm so the guard skips. + + Expects the caller to have set on each layer: + self._paged_cache: PagedKVCache + and to pass the following kwargs through the model forward: + flex_block_mask: BlockMask + flex_input_pos: Tensor [B, S] + flex_batch_idx: Tensor [B] (decode) or [1, S] (packed prefill) + flex_kernel_options: dict | None + """ + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask = None, + past_key_values = None, + cache_position = None, + flex_block_mask: Optional[BlockMask] = None, + flex_input_pos: Optional[torch.Tensor] = None, + flex_batch_idx: Optional[torch.Tensor] = None, + flex_kernel_options: Optional[dict] = None, + **kwargs, + ): + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + if hasattr(self, "q_norm"): + # Qwen3: RMSNorm on [B, S, H, D] (per-head), then transpose. + q = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose( + 1, 2 + ) + k = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose( + 1, 2 + ) + else: + # Llama: no QK-norm. + q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) + v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + cos, sin = position_embeddings + q, k = _apply_rotary(q, k, cos, sin) + + # Write to paged KV cache. For prefill, assign_prefill_no_paging + # writes into [1, H, MAX_S, D]; for decode, assign() writes into the + # B decode slots. + if self._paged_cache is not None and flex_input_pos is not None: + k, v = self._paged_cache.update(flex_input_pos, k, v, flex_batch_idx) + + # Flex attention. The block mask routes each query to the correct + # pages; enable_gqa handles num_kv_heads < num_q_heads. + attn_output = flex_attention_compiled( + q, + k, + v, + scale = self.scaling, + block_mask = flex_block_mask, + enable_gqa = True, + kernel_options = flex_kernel_options, + ) + attn_output = attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous() + return self.o_proj(attn_output), None + + return forward + + +def patch_model_attention_forwards(model: torch.nn.Module, page_table: PageTable): + """Attach a `PagedKVCache` to every attention layer of a Qwen3 or Llama + HF decoder model and swap in the flex_attention forward above. + """ + fwd = make_flex_attention_forward(page_table) + for layer in model.model.layers: + attn = layer.self_attn + attn._paged_cache = PagedKVCache( + page_table, + n_heads = model.config.num_key_value_heads, + head_dim = model.config.head_dim, + dtype = model.dtype, + ).to(model.device) + # Bind as method. + import types + + attn.forward = types.MethodType(fwd, attn) + + +# --- model forward helper that passes flex kwargs through ------------------ + + +def call_model_with_flex_kwargs(model, input_ids, position_ids, flex_kwargs): + """`model(**inputs, **flex_kwargs)` would error because the HF ForCausalLM + class doesn't declare the flex_* kwargs. We walk through the model + manually to pass them into the attention layers (which now accept them). + Works identically for Qwen3 and Llama-3.2.""" + base = model.model # Qwen3Model or LlamaModel + inputs_embeds = base.embed_tokens(input_ids) + position_embeddings = base.rotary_emb(inputs_embeds, position_ids) + hidden_states = inputs_embeds + for layer in base.layers: + residual = hidden_states + hidden_states = layer.input_layernorm(hidden_states) + hidden_states, _ = layer.self_attn( + hidden_states, + position_embeddings = position_embeddings, + **flex_kwargs, + ) + hidden_states = residual + hidden_states + residual = hidden_states + hidden_states = layer.post_attention_layernorm(hidden_states) + hidden_states = layer.mlp(hidden_states) + hidden_states = residual + hidden_states + hidden_states = base.norm(hidden_states) + return hidden_states + + +# --- inference engine ------------------------------------------------------ + + +@dataclass +class Sequence: + text: str = "" + input_ids: Optional[torch.Tensor] = None + input_length: int = 0 + output_ids: Optional[list] = None + batch_idx: int = -1 + finished: bool = False + last_token_id: int = -1 + max_new_tokens: int = 512 + + def __post_init__(self): + if self.output_ids is None: + self.output_ids = [] + + @property + def total_length(self) -> int: + return self.input_length + len(self.output_ids) + + +# --- double-copy LoRA rollout helpers ------------------------------------- +# +# PEFT's `merge_adapter` / `unmerge_adapter` pair is asymmetric at bf16: +# merge does `W_bf16 += delta_fp32` (the += upcasts then truncates), while +# unmerge does `W_bf16 -= delta_fp32.to(bf16)` -- the delta is rounded to +# bf16 first, so the round-trip leaves ~1 ULP drift on `base_layer.weight` +# every cycle. Across hundreds of GRPO iterations this corrupts the base +# model; the adapter ends up training against a drifting target. +# +# vLLM avoids this by keeping the base weights pristine and materializing a +# second "base + LoRA" copy for inference. We do the same: keep `base_model` +# (pristine) and `inference_model = deepcopy(base_model)`, wrap the copy +# with PEFT, and before each rollout refresh the LoRA-target base weights +# from pristine in-place and call `merge_adapter()` fresh. Never unmerge -- +# we always re-materialize, so there is no round-trip error to accumulate. + + +def _lora_needs_peft_fallback(module, active_adapters) -> bool: + """True when the module needs PEFT's merge path because the math isn't + plain `W += alpha * B @ A`. rslora is intentionally NOT here: PEFT + folds its scaling (alpha / sqrt(r)) into `module.scaling[adapter]`, so + the fused addmm path handles it transparently via `alpha=`.""" + if getattr(module, "lora_variant", None): + if any(a in module.lora_variant for a in active_adapters): + return True + mag = getattr(module, "lora_magnitude_vector", None) + if mag is not None and len(mag) > 0: + return True + if getattr(module, "fan_in_fan_out", False): + return True + lora_bias = getattr(module, "lora_bias", None) + if lora_bias and any(bool(lora_bias.get(a)) for a in active_adapters): + return True + return False + + +def refresh_lora_merge_from_pristine(base_model, peft_model): + """Fused one-kernel LoRA refresh. For each LoraLayer: + W_inf = W_pristine + sum_active(scaling * (B @ A)) + via `torch.addmm(out=W_inf)`, then set `merged_adapters` directly so + PEFT's forward short-circuits to `base_layer(x)` only. + + In-place addmm writes into the same tensor storage as the merged + weight, so CUDA graphs captured against it stay valid across + refreshes (replay reads the captured address; the new value takes + effect on the next replay without re-capture). + + DoRA / fan_in_fan_out / lora_bias layers fall back to PEFT's + get_delta_weight/merge path via a single `peft_model.merge_adapter()` + call at the end (after restoring their base_layer.weight from + pristine). rslora does not fall back. + + Returns the number of LoraLayer modules refreshed. + """ + from peft.tuners.lora.layer import LoraLayer + + n_refreshed = 0 + needs_fallback = [] + for name, module in peft_model.base_model.model.named_modules(): + if not isinstance(module, LoraLayer): + continue + pristine_w = base_model.get_submodule(name).weight.data + W = module.base_layer.weight.data + active = list(module.active_adapters) + + if _lora_needs_peft_fallback(module, active): + W.copy_(pristine_w) + module.merged_adapters = [] + needs_fallback.append(module) + n_refreshed += 1 + continue + + if not active: + W.copy_(pristine_w) + module.merged_adapters = [] + n_refreshed += 1 + continue + + adapter0 = active[0] + A = module.lora_A[adapter0].weight.data + B = module.lora_B[adapter0].weight.data + torch.addmm( + pristine_w, + B.to(W.dtype), + A.to(W.dtype), + alpha = module.scaling[adapter0], + out = W, + ) + for adapter in active[1:]: + A = module.lora_A[adapter].weight.data + B = module.lora_B[adapter].weight.data + torch.addmm( + W, + B.to(W.dtype), + A.to(W.dtype), + alpha = module.scaling[adapter], + out = W, + ) + module.merged_adapters = list(active) + n_refreshed += 1 + + if needs_fallback: + peft_model.merge_adapter() + + return n_refreshed + + +def _hash_state_dict(model) -> str: + """sha256 over all parameter bytes in name-sorted order. Uses the + bit-level `view(torch.uint8)` reinterpretation so bf16 / int / etc. all + round-trip without any float casting.""" + h = hashlib.sha256() + sd = model.state_dict() + for name in sorted(sd.keys()): + t = sd[name].detach().cpu().contiguous() + h.update(name.encode("utf-8")) + h.update(t.view(torch.uint8).numpy().tobytes()) + return h.hexdigest() + + +def run_drift_verification( + base_model, peft_model, n_iters: int = 10, noise_scale: float = 0.01 +): + """Simulate N GRPO iterations: perturb LoRA weights with random noise, + call `refresh_lora_merge_from_pristine`, repeat. Assert the pristine + `base_model`'s parameters are bit-identical before and after. + + Also checks inference-copy determinism: after restoring the LoRA state + to its initial value, the merged `inference_model` state-dict hash + should match the hash taken right after the first refresh. + """ + from peft.tuners.lora.layer import LoraLayer + + inference_model = peft_model.base_model.model + + # Snapshot initial LoRA A/B weights so we can restore at the end. + initial_lora = {} + for name, module in inference_model.named_modules(): + if not isinstance(module, LoraLayer): + continue + for adapter_name in list(module.lora_A.keys()): + initial_lora[(name, "A", adapter_name)] = module.lora_A[ + adapter_name + ].weight.data.clone() + initial_lora[(name, "B", adapter_name)] = module.lora_B[ + adapter_name + ].weight.data.clone() + + base_hash_before = _hash_state_dict(base_model) + + # Initial refresh: establishes merged-state baseline for the inference copy. + refresh_lora_merge_from_pristine(base_model, peft_model) + inf_hash_initial_merged = _hash_state_dict(inference_model) + + for _ in range(n_iters): + for name, module in inference_model.named_modules(): + if not isinstance(module, LoraLayer): + continue + for adapter_name in list(module.lora_A.keys()): + a = module.lora_A[adapter_name].weight.data + b = module.lora_B[adapter_name].weight.data + a.add_(noise_scale * torch.randn_like(a)) + b.add_(noise_scale * torch.randn_like(b)) + refresh_lora_merge_from_pristine(base_model, peft_model) + + base_hash_after = _hash_state_dict(base_model) + + # Restore initial LoRA weights and re-merge; inference hash must match + # the initial merged-state hash (determinism of the refresh pipeline). + for (name, kind, adapter_name), w in initial_lora.items(): + module = inference_model.get_submodule(name) + tgt = module.lora_A if kind == "A" else module.lora_B + tgt[adapter_name].weight.data.copy_(w) + refresh_lora_merge_from_pristine(base_model, peft_model) + inf_hash_restored = _hash_state_dict(inference_model) + + base_ok = base_hash_before == base_hash_after + inf_ok = inf_hash_initial_merged == inf_hash_restored + + assert base_ok, ( + f"base model drifted across {n_iters} refreshes\n" + f" before: {base_hash_before}\n" + f" after : {base_hash_after}" + ) + assert inf_ok, ( + f"inference model did not revert to deterministic merged-state hash\n" + f" initial : {inf_hash_initial_merged}\n" + f" restored : {inf_hash_restored}" + ) + print(f"[verify] base model bit-identical across {n_iters} refreshes") + print(f"[verify] inference copy deterministic after LoRA restore") + print(f"[verify] sha256 base : {base_hash_before}") + print(f"[verify] sha256 merged : {inf_hash_initial_merged}") + return { + "n_iters": n_iters, + "noise_scale": noise_scale, + "base_hash_before": base_hash_before, + "base_hash_after": base_hash_after, + "base_bit_identical": base_ok, + "inference_hash_initial_merged": inf_hash_initial_merged, + "inference_hash_after_restore": inf_hash_restored, + "inference_deterministic": inf_ok, + } + + +# Default kernel_options per phase. Our defaults stay conservative -- the +# non-default FlexKernelOptions (PRESCALE_QK, ROWS_GUARANTEED_SAFE, USE_TMA) +# are opt-in via CLI because some of them break correctness on our +# paged-attention setup. +# +# Specifically, `ROWS_GUARANTEED_SAFE=True` is unsafe here: we reserve +# batch_idx=0 and page_idx=0 as no-op padding slots. When a decode +# padded batch row maps to only-reserved pages, the block mask returns +# False for every kv_idx, so the row has zero unmasked values. The flag +# tells the kernel to skip the row-has-at-least-one-unmasked check, so +# the softmax NaNs silently -- which manifests as "!!!!!!" token spam. +DECODE_KERNEL_OPTIONS_DEFAULT = None +# Prefill keeps FORCE_USE_FLEX_ATTENTION so we don't auto-dispatch into +# the flex-decoding kernel when the packed q_len gets small. +PREFILL_KERNEL_OPTIONS_DEFAULT = {"FORCE_USE_FLEX_ATTENTION": True} + + +class FlexInference: + def __init__( + self, + model, + tokenizer, + max_batch_size = 32, + max_seq_length = 2048, + n_pages = 2048, + page_size = 128, + max_new_tokens = 512, + decode_kernel_options = None, + prefill_kernel_options = None, + fa4_prefill = None, + base_model = None, + peft_model = None, + ): + assert max_seq_length % page_size == 0 + self.model = model + self.tokenizer = tokenizer + self.device = model.device + self.eos_token_id = tokenizer.eos_token_id + # For double-copy LoRA rollout: `base_model` is the pristine copy + # (never touched); `peft_model` wraps the inference copy (`model` + # above is `peft_model.base_model.model`). Both may be None when + # no LoRA adapter is active, or when the 4-bit naive-wrapper path + # is used. + self.base_model = base_model + self.peft_model = peft_model + self.max_batch_size = max_batch_size + self.max_seq_length = max_seq_length + self.page_size = page_size + self.max_new_tokens = max_new_tokens + # FA4 CuTeDSL kernels ship for Hopper (sm_90) and Blackwell (sm_100, + # sm_120) only. `fa4_prefill=None` means auto-detect: enable where + # supported, silently fall back to the Triton flex_attention backend + # elsewhere. Explicit `fa4_prefill=True` on sub-Hopper still falls + # back, but warns -- the user asked for a kernel that isn't there. + if fa4_prefill is None or fa4_prefill: + major, _ = torch.cuda.get_device_capability(self.device) + supported = major >= 9 + if fa4_prefill and not supported: + import warnings + + warnings.warn( + f"--fa4_prefill needs Hopper (sm_90) or Blackwell " + f"(sm_100 / sm_120); found sm_{major}0. Falling back to " + f"the Triton flex_attention backend.", + RuntimeWarning, + stacklevel = 2, + ) + fa4_prefill = supported + self.fa4_prefill = fa4_prefill + # On SM100 (Blackwell), FA4 via flex_attention requires Q block = 256, + # KV block = 128. See attention-gym `get_flash_block_size`. + self.prefill_q_block = 256 if fa4_prefill else 128 + self.prefill_kv_block = 128 + self.decode_kernel_options = ( + decode_kernel_options + if decode_kernel_options is not None + else DECODE_KERNEL_OPTIONS_DEFAULT + ) + base_prefill_opts = ( + prefill_kernel_options + if prefill_kernel_options is not None + else dict(PREFILL_KERNEL_OPTIONS_DEFAULT) + ) + if fa4_prefill: + # Use the CuTeDSL FA4 kernel on Blackwell. FORCE_USE_FLEX_ATTENTION + # must be off because the FLASH backend is the flex_attention kernel. + base_prefill_opts = dict(base_prefill_opts) + base_prefill_opts.pop("FORCE_USE_FLEX_ATTENTION", None) + base_prefill_opts["BACKEND"] = "FLASH" + self.prefill_kernel_options = base_prefill_opts + + self.page_table = PageTable( + n_pages = n_pages, + page_size = page_size, + max_batch_size = max_batch_size, + device = self.device.type, + ) + patch_model_attention_forwards(model, self.page_table) + + # Pre-allocated decode state. + self.input_pos_buffer = torch.zeros( + max_batch_size, dtype = torch.int32, device = self.device + ) + # Full-length logical causal mask (shared across decode batch). + self.block_mask_logical = self.page_table.create_causal_blockmask( + B = max_batch_size, + L = max_seq_length, + ) + + self.cudagraph_captured = False + self.graphs = {} + self.graph_vars = {} + + def tokenize(self, sequences): + for seq in sequences: + ids = self.tokenizer(seq.text, return_tensors = "pt")["input_ids"].squeeze(0) + seq.input_ids = ids + seq.input_length = ids.shape[0] + + def _prefill(self, batch: list[Sequence]) -> torch.Tensor: + """Packed prefill: concatenate all sequences into [1, L] with a + document_causal mask. Return logits at the last position of each + sequence as [num_seqs, V]. + """ + input_ids_list = [seq.input_ids.to(self.device) for seq in batch] + input_pos_list = [ + torch.arange(seq.input_length, dtype = torch.long, device = self.device) + for seq in batch + ] + batch_idx_list = [ + torch.full( + (seq.input_length,), seq.batch_idx, dtype = torch.long, device = self.device + ) + for seq in batch + ] + input_ids = torch.cat(input_ids_list).view(1, -1) + input_pos = torch.cat(input_pos_list).view(1, -1) + batch_idx = torch.cat(batch_idx_list).view(1, -1) + + # Pad to multiple of Q block size (flex_attention block alignment). + # For FA4 on Blackwell, Q block = 256 -- otherwise 128. + L = input_ids.shape[1] + q_block = self.prefill_q_block + pad = (q_block - L % q_block) % q_block + if pad > 0: + input_ids = F.pad(input_ids, (0, pad), value = 0) + input_pos = F.pad(input_pos, (0, pad), value = 0) + batch_idx = F.pad(batch_idx, (0, pad), value = 0) + + input_lengths = torch.tensor( + [s.input_length for s in batch], dtype = torch.long, device = self.device + ) + logits_positions = input_lengths.cumsum(dim = 0) - 1 # [num_seqs] + + # If FA4 is on, BLOCK_SIZE is a (Q, KV) tuple. Otherwise scalar. + prefill_block_size = ( + (self.prefill_q_block, self.prefill_kv_block) + if self.fa4_prefill + else self.prefill_q_block + ) + mask = self.page_table.create_prefill_blockmask_no_paging( + batch_idx, BLOCK_SIZE = prefill_block_size + ) + + flex_kwargs = dict( + flex_block_mask = mask, + flex_input_pos = input_pos, + flex_batch_idx = batch_idx, + flex_kernel_options = self.prefill_kernel_options, + ) + position_ids = input_pos # Qwen3 uses 0-based; unlike Gemma2 + hidden = call_model_with_flex_kwargs( + self.model, input_ids, position_ids, flex_kwargs + ) + return self.model.lm_head(hidden[:, logits_positions, :]).squeeze(0) + + def _decode_block_mask(self, batch_idx: torch.Tensor): + """Slice a single-row BlockMask for every seq in the decode batch, + then translate logical→physical pages.""" + block_mask = self.block_mask_logical + input_pos = self.input_pos_buffer[batch_idx] + assert batch_idx.ndim == 1 and input_pos.ndim == 1 + B = batch_idx.shape[0] + input_block_idx = input_pos // block_mask.BLOCK_SIZE[0] + kv_num_blocks = block_mask.kv_num_blocks[batch_idx, :, input_block_idx].view( + B, 1, 1 + ) + kv_indices = block_mask.kv_indices[batch_idx, :, input_block_idx].view( + B, 1, 1, -1 + ) + full_num = full_idx = None + if block_mask.full_kv_num_blocks is not None: + full_num = block_mask.full_kv_num_blocks[ + batch_idx, :, input_block_idx + ].view(B, 1, 1) + full_idx = block_mask.full_kv_indices[batch_idx, :, input_block_idx].view( + B, 1, 1, -1 + ) + + def causal_offset(off): + def offset(b, h, q_idx, kv_idx): + return q_idx + off[b] >= kv_idx + + return offset + + seq_length = (1, block_mask.seq_lengths[1]) + mask = BlockMask.from_kv_blocks( + kv_num_blocks, + kv_indices, + full_num, + full_idx, + BLOCK_SIZE = block_mask.BLOCK_SIZE, + mask_mod = causal_offset(input_pos), + seq_lengths = seq_length, + ) + return mask, input_pos + + def _decode_step_eager(self, batch_idx: torch.Tensor, input_ids: torch.Tensor): + B = input_ids.shape[0] + mask, input_pos = self._decode_block_mask(batch_idx) + mask = self.page_table.convert_logical_block_mask(mask, batch_idx) + position_ids = (input_pos).view(B, 1).to(torch.long) + flex_kwargs = dict( + flex_block_mask = mask, + flex_input_pos = input_pos.view(B, 1).to(torch.long), + flex_batch_idx = batch_idx, + flex_kernel_options = self.decode_kernel_options, + ) + hidden = call_model_with_flex_kwargs( + self.model, input_ids.view(B, 1), position_ids, flex_kwargs + ) + return self.model.lm_head(hidden[:, -1, :]) # [B, V] + + def _decode_step( + self, batch_idx: torch.Tensor, input_ids: torch.Tensor, input_pos: torch.Tensor + ): + self.input_pos_buffer.zero_() + self.input_pos_buffer[batch_idx] = input_pos + if not self.cudagraph_captured: + return self._decode_step_eager(batch_idx, input_ids) + bs = input_ids.size(0) + key = next(x for x in self.graph_bs if x >= bs) + graph = self.graphs[key] + gv = self.graph_vars + # batch_idx=0 is the reserved no-op slot. Zero out the unused part + # of each capture-shape buffer so padded entries don't write into + # real KV pages. + for k, v in gv.items(): + if k != "outputs": + v.zero_() + gv["input_ids"][:bs] = input_ids + gv["batch_idx"][:bs] = batch_idx + graph.replay() + return gv["outputs"][:bs] + + def capture_decode_cudagraph(self): + """Capture one CUDA graph per batch-size bucket. + + Pre-reserves a page for every batch_idx slot so the KV cache writes + during capture hit valid physical addresses. After capture we erase + the batches -- the graph replay reads/writes the same physical + pages regardless of whether the logical batch currently owns them, + because batch_idx 0 is reserved as a padding slot. + """ + max_bs = self.max_batch_size + # Reserve a dummy page for every slot we're going to use during + # capture. Without this, assign() does k_cache[:, :, -1, :] = ... + # and we get an illegal memory access. + reserved_batches = [] + for bi in range(1, max_bs): + try: + allocated = self.page_table.allocate() + self.page_table.reserve( + allocated, + torch.tensor([allocated], device = self.device, dtype = torch.long), + self.page_size, # just one page + ) + reserved_batches.append(allocated) + except Exception: + break + + input_ids = torch.zeros(max_bs, dtype = torch.int64, device = self.device) + batch_idx = torch.arange(max_bs, dtype = torch.int64, device = self.device) + outputs = torch.zeros( + (max_bs, self.model.config.vocab_size), + dtype = self.model.dtype, + device = self.device, + ) + self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16)) + pool = None + for bs in reversed(self.graph_bs): + if bs > max_bs: + continue + print(f"[flex] capturing CUDA graph for bs={bs}") + torch.cuda.synchronize() + _ = self._decode_step_eager(batch_idx[:bs], input_ids[:bs]) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool): + outputs[:bs] = self._decode_step_eager(batch_idx[:bs], input_ids[:bs]) + if pool is None: + pool = graph.pool() + self.graphs[bs] = graph + torch.cuda.synchronize() + # Release the scratch batches; real requests will re-allocate them. + for bi in reserved_batches: + self.page_table.erase(bi) + self.graph_vars = dict( + input_ids = input_ids, batch_idx = batch_idx, outputs = outputs + ) + + def refresh_inference_from_base(self): + """Re-materialize the inference copy's merged LoRA weights from the + pristine `base_model`. Call this once at setup (before CUDA graph + capture) and, in a real GRPO loop, once after every training step + that updates the LoRA adapter. Never call `unmerge_adapter()` -- + we always re-merge from pristine, so no drift accumulates. + + No-op when the double-copy pair wasn't configured (e.g. no LoRA, + or 4-bit naive PEFT-wrapper path). + """ + if self.base_model is None or self.peft_model is None: + return 0 + return refresh_lora_merge_from_pristine(self.base_model, self.peft_model) + + @torch.inference_mode() + def generate(self, sequences: list[Sequence], capture_cudagraph = False): + self.tokenize(sequences) + waiting = deque(sequences) + running = deque() + done = [] + + if capture_cudagraph and not self.cudagraph_captured: + self.capture_decode_cudagraph() + self.cudagraph_captured = True + + while waiting or running: + # 1. Try to schedule new requests into running. + batch = [] + while waiting and self.page_table.can_reserve(waiting[0].total_length): + seq = waiting.popleft() + bi = self.page_table.allocate() + self.page_table.reserve( + bi, + torch.tensor([bi], device = self.device, dtype = torch.long), + seq.total_length, + ) + seq.batch_idx = bi + batch.append(seq) + if batch: + logits = self._prefill(batch) + next_ids = torch.argmax(logits, dim = -1).tolist() + for i, seq in enumerate(batch): + seq.last_token_id = next_ids[i] + seq.output_ids.append(next_ids[i]) + if ( + seq.last_token_id == self.eos_token_id + or len(seq.output_ids) >= seq.max_new_tokens + ): + seq.finished = True + done.append(seq) + self.page_table.erase(seq.batch_idx) + else: + running.append(seq) + continue + + # 2. Reserve pages for running seqs that need more capacity. + decode_batch = [] + while running: + seq = running.popleft() + if self.page_table.capacity[seq.batch_idx] >= seq.total_length: + decode_batch.append(seq) + elif self.page_table.can_reserve( + seq.total_length, batch_idx_int = seq.batch_idx + ): + self.page_table.reserve( + seq.batch_idx, + torch.tensor( + [seq.batch_idx], device = self.device, dtype = torch.long + ), + seq.total_length, + ) + decode_batch.append(seq) + else: + running.appendleft(seq) + newest = running.pop() + waiting.appendleft(newest) + self.page_table.erase(newest.batch_idx) + if not decode_batch: + continue + + B = len(decode_batch) + bi_tensor = torch.tensor( + [s.batch_idx for s in decode_batch], + dtype = torch.long, + device = self.device, + ) + last_ids = torch.tensor( + [s.last_token_id for s in decode_batch], + dtype = torch.long, + device = self.device, + ) + cur_pos = torch.tensor( + [s.total_length - 1 for s in decode_batch], + dtype = torch.int32, + device = self.device, + ) + logits = self._decode_step(bi_tensor, last_ids, cur_pos) + next_ids = torch.argmax(logits, dim = -1).tolist() + for i, seq in enumerate(decode_batch): + seq.last_token_id = next_ids[i] + seq.output_ids.append(next_ids[i]) + if ( + seq.last_token_id == self.eos_token_id + or len(seq.output_ids) >= seq.max_new_tokens + ): + seq.finished = True + done.append(seq) + self.page_table.erase(seq.batch_idx) + else: + running.append(seq) + + return done + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--model_name", default = "unsloth/Qwen3-4B-Base") + 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("--max_batch_size", type = int, default = 64) + p.add_argument("--max_seq_length", type = int, default = 2048) + p.add_argument("--n_pages", type = int, default = 2048) + p.add_argument("--page_size", type = int, default = 128) + p.add_argument("--capture_cudagraph", action = "store_true") + p.add_argument("--lora_adapter", default = None) + # Kernel tuning (optional JSON-valued CLI args so we can sweep quickly): + p.add_argument( + "--decode_kernel_options", + default = None, + help = "JSON for FlexKernelOptions applied in decode, " + 'e.g. \'{"PRESCALE_QK":true,"USE_TMA":true}\'.', + ) + p.add_argument( + "--prefill_kernel_options", default = None, help = "Same but for prefill." + ) + # If set, torch.compile the full attention-stack closure in addition to + # (or instead of) compiling just flex_attention. `reduce-overhead` is + # the interesting mode; it nests with our CUDA graph capture. + p.add_argument( + "--compile_model_forward", + default = None, + choices = [None, "default", "reduce-overhead", "max-autotune-no-cudagraphs"], + ) + p.add_argument( + "--fa4_prefill", + default = None, + action = argparse.BooleanOptionalAction, + help = ( + "Use BLOCK_SIZE=(256,128) + BACKEND=FLASH on prefill to unlock the " + "CuTeDSL FA4 kernel. Default auto-enables on Hopper (sm_90) and " + "Blackwell (sm_100, sm_120); use --no-fa4_prefill to force off." + ), + ) + p.add_argument( + "--load_in_4bit", + action = "store_true", + help = ( + "Load the base model as bitsandbytes 4-bit. When set with " + "--lora_adapter, the LoRA is kept as a PEFT wrapper (no merge) " + "because merging into 4-bit weights is not supported." + ), + ) + p.add_argument( + "--no_merge_lora", + action = "store_true", + help = ( + "Reference path: keep the LoRA adapter as a PEFT wrapper " + "instead of merging it. Runs three matmuls per projection; " + "slow. Useful for the unmerged row in the writeup's comparison " + "table. The default is now the double-copy pattern, which is " + "both merge-speed and drift-free." + ), + ) + p.add_argument( + "--verify_no_drift", + action = "store_true", + help = ( + "Drift-verification mode. Hash the pristine base model params, " + "run N perturb+refresh cycles (simulating N GRPO iterations) " + "on a copy, re-hash, and assert bit-identical. Requires a " + "--lora_adapter; skips rollout generation." + ), + ) + p.add_argument( + "--verify_iterations", + type = int, + default = 10, + help = "Number of perturb+refresh cycles for --verify_no_drift.", + ) + p.add_argument( + "--model_name_4bit", + default = None, + help = ( + "Override the 4-bit shard name. Defaults to " + "`{model_name}-unsloth-bnb-4bit`." + ), + ) + p.add_argument("--stats_path", required = True) + p.add_argument( + "--chat_template", + choices = ["auto", "grpo", "native"], + default = "auto", + help = ( + "Which chat template to use for building prompts. " + "`auto`: GRPO template for Qwen3, tokenizer's native template " + "otherwise. `grpo`: force the GRPO template (matches prior " + "Qwen3 baselines). `native`: force the tokenizer's built-in " + "template (required for Llama-3.2-Instruct)." + ), + ) + args = p.parse_args() + + def _parse_opts(s): + if s is None: + return None + return json.loads(s) + + 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_model = None + peft_model = None + + if args.load_in_4bit: + # Load the pre-quantized Unsloth 4-bit shard. Compute dtype comes + # from the packaged config (bf16 for these shards). + # + # 4-bit keeps the naive PEFT-wrapper path: bnb's `Linear4bit` holds + # packed quantised weights, not regular bf16, so the double-copy + # refresh (in-place copy of `base_layer.weight`) doesn't apply. + # Materializing a full bf16 inference copy via dequant would wipe + # out the memory saving of 4-bit. + bnb_model_name = args.model_name_4bit or f"{args.model_name}-unsloth-bnb-4bit" + print(f"[flex] loading 4-bit base: {bnb_model_name}") + model = AutoModelForCausalLM.from_pretrained( + bnb_model_name, + attn_implementation = "eager", + device_map = "cuda:0", + ) + # See note in cb_vs_vllm_generation.py: tie lm_head to embed_tokens + # for bnb-4bit shards of tied-embedding models. + if getattr(model.config, "tie_word_embeddings", False): + model.lm_head.weight = model.model.embed_tokens.weight + model.eval() + + if args.lora_adapter: + from peft import PeftModel + + peft_wrapper = PeftModel.from_pretrained( + model, + str(Path(args.lora_adapter).resolve()), + is_trainable = False, + ) + # LoRA stays as a wrapper around Params4bit; three matmuls per + # projection. This is the slow reference row in the writeup. + model = peft_wrapper.base_model.model + else: + # bf16 path -- double-copy LoRA rollout. + # + # `base_model` stays pristine; we deep-copy it to `inference_model`, + # wrap the copy with PEFT, and re-materialize the merged LoRA on + # the copy whenever the LoRA weights change. Memory cost: +~8 GB + # for Qwen3-4B bf16 (two copies on GPU) -- well within budget vs + # vLLM's 156 GB. + base_model = AutoModelForCausalLM.from_pretrained( + args.model_name, + dtype = torch.bfloat16, + attn_implementation = "eager", + ).to("cuda") + base_model.eval() + + if not args.lora_adapter: + # No adapter -- use base_model directly, no inference copy. + model = base_model + base_model = None + elif args.no_merge_lora: + # Reference path: PEFT wrapper on the only model copy, + # adapter unmerged. Three matmuls per projection. Kept for + # the comparison row in the writeup. + from peft import PeftModel + + peft_wrapper = PeftModel.from_pretrained( + base_model, + str(Path(args.lora_adapter).resolve()), + is_trainable = False, + ) + model = peft_wrapper.base_model.model + base_model = None + else: + # Double-copy rollout path. + from peft import PeftModel + + print("[flex] deep-copying base model for double-copy LoRA rollout") + inference_model = copy.deepcopy(base_model) + inference_model.eval() + + peft_model = PeftModel.from_pretrained( + inference_model, + str(Path(args.lora_adapter).resolve()), + is_trainable = False, + ) + model = peft_model.base_model.model + model.eval() + # Do NOT call merge_adapter here -- FlexInference.refresh_ + # inference_from_base() below handles the initial merge so the + # same code path runs at setup and on every GRPO refresh. + + # Drift-verification mode: skip rollout generation, just hash-check. + if args.verify_no_drift: + if args.load_in_4bit: + raise SystemExit( + "--verify_no_drift only applies to the bf16 double-copy path " + "(4-bit keeps the naive PEFT-wrapper path, no merge refresh)." + ) + if args.no_merge_lora: + raise SystemExit( + "--verify_no_drift is incompatible with --no_merge_lora " + "(nothing is merged; nothing to drift)." + ) + if base_model is None or peft_model is None: + raise SystemExit( + "--verify_no_drift requires --lora_adapter so there is a " + "LoRA to merge/refresh against the pristine base." + ) + print( + f"[flex] running drift verification: {args.verify_iterations} " + f"perturb+refresh cycles" + ) + result = run_drift_verification( + base_model, peft_model, n_iters = args.verify_iterations + ) + result = {"mode": "verify_no_drift", **result} + 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(result, f, indent = 2) + print(json.dumps(result, indent = 2)) + os._exit(0) + + from unsloth_grpo_common import ( + SYSTEM_PROMPT, + apply_chat_template_to_tokenizer, + ) + from datasets import load_dataset + + # Pick which chat template builds the prompts. Qwen3 baselines in + # this repo were recorded against the GRPO template; Llama-3.2-Instruct + # only produces coherent completions with its shipped Instruct + # template. + if args.chat_template == "auto": + use_grpo = type(model).__name__.startswith("Qwen3") + elif args.chat_template == "grpo": + use_grpo = True + else: # "native" + use_grpo = False + if use_grpo: + apply_chat_template_to_tokenizer(tok) + print("[flex] chat_template: GRPO") + else: + print("[flex] chat_template: tokenizer native") + 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 + ] + texts = [ + tok.apply_chat_template(m, add_generation_prompt = True, tokenize = False) + for m in messages + ] + + # Make sure the base HF model the attention layers belong to isn't + # wrapped by PeftModel anymore (we merged); `.model` should be + # Qwen3ForCausalLM or LlamaForCausalLM. + inference = FlexInference( + model, + tok, + max_batch_size = args.max_batch_size, + max_seq_length = args.max_seq_length, + n_pages = args.n_pages, + page_size = args.page_size, + max_new_tokens = args.max_new_tokens, + decode_kernel_options = _parse_opts(args.decode_kernel_options), + prefill_kernel_options = _parse_opts(args.prefill_kernel_options), + fa4_prefill = args.fa4_prefill, + base_model = base_model, + peft_model = peft_model, + ) + + # Initial merge from pristine. Done via `refresh_inference_from_base` + # (not raw `merge_adapter`) so the exact same code path runs at setup + # and at every GRPO refresh -- the CUDA graph capture below sees the + # merged weights already in place. In a real GRPO loop, call + # `inference.refresh_inference_from_base()` after every training step + # that updates the LoRA adapter. We skip per-round refresh in this + # benchmark because the LoRA weights don't change between rounds. + if inference.base_model is not None and inference.peft_model is not None: + n = inference.refresh_inference_from_base() + print(f"[flex] double-copy rollout: refreshed {n} LoRA-target layers") + + # Optionally compile the manual forward walker. This fuses the layer-stack + # ops around flex_attention. Under CUDA graph capture, the compiled + # function gets captured into the same graph. + if args.compile_model_forward: + torch._dynamo.config.cache_size_limit = 256 + print( + f"[flex] torch.compile(call_model_with_flex_kwargs, " + f"mode={args.compile_model_forward!r})" + ) + import sys as _sys + + _this = _sys.modules[__name__] + _this.call_model_with_flex_kwargs = torch.compile( + call_model_with_flex_kwargs, + mode = args.compile_model_forward, + dynamic = True, + fullgraph = False, + ) + + def make_seqs(): + return [Sequence(text = t, max_new_tokens = args.max_new_tokens) for t in texts] + + # Warmup. + torch.cuda.reset_peak_memory_stats() + print("[flex] warmup (16 prompts)...") + _ = inference.generate(make_seqs()[:16], capture_cudagraph = args.capture_cudagraph) + torch.cuda.synchronize() + + wall_times = [] + total_decoded = 0 + for r in range(args.n_rounds): + torch.cuda.synchronize() + t0 = time.perf_counter() + out = inference.generate(make_seqs()) + torch.cuda.synchronize() + wall_times.append(time.perf_counter() - t0) + total_decoded = sum(len(s.output_ids) for s in out) + print( + f"[flex] 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] + best = min(wall_times) + peak = torch.cuda.max_memory_allocated() / 1024**3 + # Sample a couple of completions so we can eyeball coherence. + sample_completions = [] + for s in out[:3]: + sample_completions.append( + tok.decode(s.output_ids[:80], skip_special_tokens = True) + ) + res = { + "backend": "flex", + "model_name": args.model_name, + "capture_cudagraph": args.capture_cudagraph, + "lora_adapter": args.lora_adapter, + "n_prompts": args.n_prompts, + "n_decoded_tokens": total_decoded, + "wall_times_s": wall_times, + "median_wall_s": med, + "best_wall_s": best, + "decode_tps_median": total_decoded / med if med else 0, + "decode_tps_best": total_decoded / best if best else 0, + "max_new_tokens": args.max_new_tokens, + "peak_memory_gb": peak, + "sample_completions": sample_completions, + } + 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(res, f, indent = 2) + print(json.dumps(res, indent = 2)) + os._exit(0) + + +if __name__ == "__main__": + main() diff --git a/scripts/benchmarks/qwen3_grpo_naive.py b/scripts/benchmarks/qwen3_grpo_naive.py new file mode 100644 index 0000000000..3d8a5dcb9e --- /dev/null +++ b/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() diff --git a/scripts/benchmarks/qwen3_grpo_tpaged.py b/scripts/benchmarks/qwen3_grpo_tpaged.py new file mode 100644 index 0000000000..c2b8839704 --- /dev/null +++ b/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|` 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() diff --git a/scripts/benchmarks/qwen3_grpo_vllm.py b/scripts/benchmarks/qwen3_grpo_vllm.py new file mode 100644 index 0000000000..d78b131c46 --- /dev/null +++ b/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() diff --git a/scripts/benchmarks/results/flex_vs_vllm.md b/scripts/benchmarks/results/flex_vs_vllm.md new file mode 100644 index 0000000000..1a6cfb13f1 --- /dev/null +++ b/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` diff --git a/scripts/benchmarks/results/grpo_equivalence.md b/scripts/benchmarks/results/grpo_equivalence.md new file mode 100644 index 0000000000..6c8288cece --- /dev/null +++ b/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 `` 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 |loss diff| | max |reward diff| | max |kl diff| | max |grad_norm diff| | +|---------------------------------|---------------------------|------------------------------|------------------------|--------------------------------| +| 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 diff --git a/scripts/benchmarks/results/lora_rollout_baselines.md b/scripts/benchmarks/results/lora_rollout_baselines.md new file mode 100644 index 0000000000..28f3a54798 --- /dev/null +++ b/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. diff --git a/scripts/benchmarks/results/notebook_ref_10.md b/scripts/benchmarks/results/notebook_ref_10.md new file mode 100644 index 0000000000..c04d44fcfa --- /dev/null +++ b/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. diff --git a/scripts/benchmarks/unsloth_grpo_common.py b/scripts/benchmarks/unsloth_grpo_common.py new file mode 100644 index 0000000000..b01c3971d0 --- /dev/null +++ b/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 = "" +REASONING_END = "" +SOLUTION_START = "" +SOLUTION_END = "" + +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"[\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 diff --git a/scripts/benchmarks/verify_gemma4_numerics.py b/scripts/benchmarks/verify_gemma4_numerics.py new file mode 100644 index 0000000000..a6b1bc3cbf --- /dev/null +++ b/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() diff --git a/scripts/benchmarks/verify_qwen3_numerics.py b/scripts/benchmarks/verify_qwen3_numerics.py new file mode 100644 index 0000000000..56b91f2398 --- /dev/null +++ b/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() diff --git a/tests/test_fa4_capability_guard.py b/tests/test_fa4_capability_guard.py new file mode 100644 index 0000000000..70aea8934a --- /dev/null +++ b/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()