"""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 adapts that architecture to Qwen3-4B. 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). 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 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_qwen3_attention_forward(page_table: PageTable): """Return a new `forward` method for `Qwen3Attention` 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`. 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) 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) 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_qwen3_model(model: torch.nn.Module, page_table: PageTable): """Attach a `PagedKVCache` to every `Qwen3Attention` layer and swap in the flex_attention forward above. """ fwd = make_flex_qwen3_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 Qwen3ForCausalLM doesn't declare the flex_* kwargs. We walk through the model manually to pass them into the attention layers (which now accept them).""" base = model.model # Qwen3Model 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_qwen3_model(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) 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 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 ] texts = [ tok.apply_chat_template(m, add_generation_prompt = True, tokenize = False) for m in messages ] # Make sure the base HF model that Qwen3Attention belongs to isn't wrapped # by PeftModel anymore (we merged); `.model` should be Qwen3ForCausalLM. 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": "qwen3_flex", "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()