1249 lines
49 KiB
Python
1249 lines
49 KiB
Python
"""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()
|