314ab6ae86
Replace the copy+merge_adapter pair in refresh_lora_merge_from_pristine with one torch.addmm(pristine, B, A, alpha=scaling, out=W_inf) per LoraLayer, then set merged_adapters directly so PEFT's forward short-circuits to base_layer(x). Previously each refresh did two passes per weight: a bf16 copy from pristine, then PEFT merge_adapter which materialises a full [out, in] fp32 delta via get_delta_weight and in-place adds it back. The fused path skips the transient delta allocation and runs one cuBLAS GEMM instead. cuBLAS accumulates the bf16 matmul in fp32 internally, so the numerical result stays within 1 bf16 ULP of PEFT's path (verified on the rank-32 Qwen3-4B adapter: max abs diff 1.22e-04). DoRA, fan_in_fan_out, and lora_bias=True layers fall back to PEFT's get_delta_weight/merge path via a single trailing merge_adapter call after restoring their base_layer.weight from pristine. rslora is not a fallback -- PEFT folds alpha/sqrt(r) into module.scaling[adapter], so the fused addmm picks it up transparently via alpha=. Drift verification still passes: base bit-identical across 10 perturb+refresh cycles, inference state deterministic after LoRA restore.
1176 lines
46 KiB
Python
1176 lines
46 KiB
Python
"""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 = False,
|
|
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
|
|
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",
|
|
action = "store_true",
|
|
help = (
|
|
"Use BLOCK_SIZE=(256,128) + BACKEND=FLASH on prefill to unlock the "
|
|
"CuTeDSL FA4 kernel on Blackwell (SM100)."
|
|
),
|
|
)
|
|
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()
|