diff --git a/tests/flex_fastlm_bench.py b/tests/flex_fastlm_bench.py new file mode 100644 index 0000000000..1f1ea7fa3f --- /dev/null +++ b/tests/flex_fastlm_bench.py @@ -0,0 +1,78 @@ +"""Batched steady-state throughput bench through ``FastLanguageModel`` + +``UNSLOTH_FAST_INFERENCE=1``. First ``generate`` call primes CUDA graphs; +subsequent calls report steady state. Compare against April CLI-only +numbers for the same workload.""" + +import os, sys, time +from pathlib import Path +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +import argparse +import torch + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--model", default="unsloth/Qwen3-4B-Base") + p.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16") + p.add_argument("--n_prompts", type=int, default=8) + p.add_argument("--max_new_tokens", type=int, default=64) + p.add_argument("--max_batch_size", type=int, default=16) + p.add_argument("--max_seq_length", type=int, default=1024) + p.add_argument("--n_rounds", type=int, default=3) + args = p.parse_args() + + os.environ.setdefault("UNSLOTH_FAST_INFERENCE", "1") + import unsloth + from unsloth import FastLanguageModel + + dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16 + + model, tok = FastLanguageModel.from_pretrained( + model_name=args.model, + max_seq_length=args.max_seq_length, + dtype=dtype, + load_in_4bit=False, + fast_inference=True, + max_batch_size=args.max_batch_size, + ) + print(f"[bench] model={args.model} dtype={args.dtype}") + prompts = [f"In one sentence, a fact about {t} is" for t in [ + "the moon", "gravity", "the ocean", "the sun", "honey", + "rain", "trees", "mountains" + ][:args.n_prompts]] + + class _SP: + max_tokens = args.max_new_tokens + temperature = 0.0 + + # Warmup round — captures CUDA graphs. + print("[bench] warmup (CUDA graph capture)...") + t0 = time.perf_counter() + _ = model.fast_generate(prompts, sampling_params=_SP(), use_tqdm=False) + print(f"[bench] warmup wall: {time.perf_counter() - t0:.2f}s") + + walls = [] + tok_counts = [] + for r in range(args.n_rounds): + torch.cuda.synchronize() + t1 = time.perf_counter() + outs = model.fast_generate(prompts, sampling_params=_SP(), use_tqdm=False) + torch.cuda.synchronize() + dt = time.perf_counter() - t1 + n_tok = sum(len(o.outputs[0].token_ids) for o in outs) + walls.append(dt) + tok_counts.append(n_tok) + print(f"[bench] round {r}: {n_tok} toks in {dt:.2f}s -> {n_tok/dt:.1f} tok/s") + + if walls: + wall_med = sorted(walls)[len(walls) // 2] + tok_med = tok_counts[len(walls) // 2] + print(f"[bench] median: {tok_med} toks in {wall_med:.2f}s " + f"=> {tok_med / wall_med:.1f} tok/s") + + +if __name__ == "__main__": + main() diff --git a/tests/flex_fastlm_smoke.py b/tests/flex_fastlm_smoke.py new file mode 100644 index 0000000000..05b620552c --- /dev/null +++ b/tests/flex_fastlm_smoke.py @@ -0,0 +1,93 @@ +"""Smoke-test the ``UNSLOTH_FAST_INFERENCE=1`` path through +``FastLanguageModel.from_pretrained``. + +Invoked as: + CUDA_VISIBLE_DEVICES=2 UNSLOTH_FAST_INFERENCE=1 python tests/flex_fastlm_smoke.py \ + --model unsloth/Qwen3-4B-Base --dtype bf16 --no-lora + +Prints tokens/s + the first generated string. +""" + +from __future__ import annotations + +import argparse +import os +import sys +import time +from pathlib import Path + +# Make the local fork importable. +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--model", default="unsloth/Qwen3-4B-Base") + p.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16") + p.add_argument("--load_in_4bit", action="store_true") + p.add_argument("--with_lora", action="store_true") + p.add_argument("--max_new_tokens", type=int, default=32) + p.add_argument("--max_seq_length", type=int, default=1024) + p.add_argument("--prompt", default="The quick brown fox jumps over") + args = p.parse_args() + + import torch + + os.environ.setdefault("UNSLOTH_FAST_INFERENCE", "1") + print(f"[smoke] UNSLOTH_FAST_INFERENCE={os.environ.get('UNSLOTH_FAST_INFERENCE')}") + + import unsloth + print(f"[smoke] unsloth={unsloth.__file__}") + from unsloth import FastLanguageModel + + dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16 + + t0 = time.perf_counter() + model, tokenizer = FastLanguageModel.from_pretrained( + model_name=args.model, + max_seq_length=args.max_seq_length, + dtype=dtype, + load_in_4bit=args.load_in_4bit, + fast_inference=True, + ) + t_load = time.perf_counter() - t0 + print(f"[smoke] loaded model in {t_load:.1f}s; dtype={model.dtype}") + print(f"[smoke] hasattr(model, 'vllm_engine'): {hasattr(model, 'vllm_engine')}") + print(f"[smoke] vllm_engine type: {type(model.vllm_engine).__name__}") + + if args.with_lora: + model = FastLanguageModel.get_peft_model( + model, + r=16, + target_modules=["q_proj", "k_proj", "v_proj", "o_proj"], + lora_alpha=16, + lora_dropout=0.0, + bias="none", + use_gradient_checkpointing="unsloth", + random_state=3407, + ) + print(f"[smoke] PEFT model type: {type(model).__name__}") + print(f"[smoke] model.vllm_engine bound to PEFT: " + f"{hasattr(model, 'vllm_engine')}") + + from unsloth.inference.vllm_shim import LoRARequest + prompts = [args.prompt] + # Minimal SamplingParams stand-in + class _SP: + max_tokens = args.max_new_tokens + temperature = 0.0 + t1 = time.perf_counter() + outputs = model.fast_generate(prompts, sampling_params=_SP(), use_tqdm=False) + dt = time.perf_counter() - t1 + out = outputs[0] + n_tok = len(out.outputs[0].token_ids) + print(f"[smoke] generated {n_tok} tokens in {dt:.2f}s " + f"({n_tok / dt:.1f} tok/s)") + print(f"[smoke] prompt: {args.prompt!r}") + print(f"[smoke] completion: {out.outputs[0].text!r}") + + +if __name__ == "__main__": + main() diff --git a/unsloth/inference/__init__.py b/unsloth/inference/__init__.py new file mode 100644 index 0000000000..cc0deb9efa --- /dev/null +++ b/unsloth/inference/__init__.py @@ -0,0 +1,31 @@ +"""Flex-attention inference engines. + +``UNSLOTH_FAST_INFERENCE=1`` routes ``FastLanguageModel.from_pretrained`` +through :func:`load_flex`, which wraps the selected HF model with a +:class:`FlexEngine` that presents the vLLM ``LLM`` surface used by +Unsloth / TRL GRPO (``.generate``, ``.chat``, ``.sleep``, ``.wake_up``, +``.llm_engine``, plus ``save_lora`` / ``load_lora`` via the module +shim). + +Three architectures are supported today: Qwen3, Llama-3, Gemma-4-E2B-it. +Anything else raises :class:`NotImplementedError` — there is no silent +fallback; unset the env var or use vLLM instead.""" + +from .flex_engine import FlexEngine, load_flex +from .vllm_shim import ( + CompletionOutput, + LoRARequest, + RequestOutput, + load_lora, + save_lora, +) + +__all__ = [ + "FlexEngine", + "load_flex", + "LoRARequest", + "RequestOutput", + "CompletionOutput", + "save_lora", + "load_lora", +] diff --git a/unsloth/inference/flex_engine.py b/unsloth/inference/flex_engine.py new file mode 100644 index 0000000000..b17eb39d9c --- /dev/null +++ b/unsloth/inference/flex_engine.py @@ -0,0 +1,720 @@ +"""FlexEngine: vLLM-compatible LLM surface for the flex inference backends. + +When ``UNSLOTH_FAST_INFERENCE=1`` is set, :func:`load_flex` wraps the HF model +with a :class:`FlexEngine`. TRL's GRPO trainer and Unsloth's own callers treat +it as an ``LLM``: + + engine.generate(prompts, sampling_params=..., lora_request=..., use_tqdm=False) + engine.chat(messages, sampling_params=..., lora_request=..., use_tqdm=False) + engine.sleep(level=2) + engine.wake_up(tags=["kv_cache"]) + engine.llm_engine # minimal stub, see below + +Three architectures are supported: Qwen3, Llama-3, Gemma-4-E2B-it. Anything +else raises :class:`NotImplementedError`. + +Colocation model (no change from today): +- ``self.hf_model`` is the HF model instance returned by ``from_pretrained`` + (or its PEFT-wrapped descendant after :meth:`bind_peft_model`). +- The underlying :class:`FlexInference` / :class:`FlexGemma4Inference` implements + the usual double-copy rollout pattern (pristine ``base_model`` + inference + copy wrapped by PEFT), and ``refresh_lora_merge_from_pristine`` re-merges on + every rollout. +- vLLM's sleep mode is not implemented for the flex backend. ``.sleep()`` / + ``.wake_up()`` are no-op stubs so code paths that gate on + ``UNSLOTH_VLLM_STANDBY`` stay valid. +""" + +from __future__ import annotations + +import copy +import functools +import importlib.util +import os +import types +import warnings +from typing import Any, Optional + +import torch + +from .flex_paged_attention import PagedKVCache, PageTable # noqa: F401 +from .flex_qwen3_llama import ( + DECODE_KERNEL_OPTIONS_DEFAULT, + PREFILL_KERNEL_OPTIONS_DEFAULT, + FlexInference, + Sequence, + refresh_lora_merge_from_pristine, +) +from .flex_gemma4 import FlexGemma4Inference +from .vllm_shim import CompletionOutput, LoRARequest, RequestOutput + + +# --------------------------------------------------------------------------- +# Hardware / kernel auto-tune +# --------------------------------------------------------------------------- + + +def _flash_attn_4_importable() -> bool: + """FA4's CuTeDSL backend lives in the ``flash_attn_interface`` ships with + the ``flash-attn`` 4.x wheel. We don't import it here (it can be slow); + just check whether the module is resolvable.""" + return importlib.util.find_spec("flash_attn_interface") is not None + + +def _fa4_ok_for_head_dim(head_dim: int, device_cap: tuple[int, int]) -> bool: + """See the plan for the hardware tier table. Returns whether FA4 is + safe to use for an attention layer with the given ``head_dim`` on a GPU + with the given ``(major, minor)`` capability.""" + major, _ = device_cap + if major < 9: + return False # Ampere and older: Triton only + if not _flash_attn_4_importable(): + return False + if major == 9: + return 8 <= head_dim <= 256 # Hopper + return 8 <= head_dim <= 128 # Blackwell (sm_100 / sm_120) + + +def _triton_block_defaults( + head_dim_max: int, device_cap: tuple[int, int] +) -> tuple[int, int, int, int]: + """Returns ``(prefill_BM, prefill_BN, decode_BM, decode_BN)``. + + Blackwell (sm_100 / sm_120) shared-memory budget is tight; at + ``head_dim >= 256`` the default 128x128 prefill block overflows and the + Triton autotuner raises ``OutOfMemoryError: out of resource``. See the + ``--prefill_kernel_options`` probes in ``scripts/benchmarks/*.py``.""" + major, _ = device_cap + if major >= 10 and head_dim_max >= 256: + return (32, 32, 16, 16) + if major >= 10 and head_dim_max >= 128: + return (64, 64, 32, 32) + return (128, 128, 64, 64) + + +def _collect_head_dims(hf_model) -> list[int]: + """Collect the per-layer ``head_dim`` for every attention sub-module in + the model. Gemma-4 text stack mixes 256 / 512 between full / sliding + layers; Qwen3 + Llama-3 use a single head_dim.""" + head_dims = set() + for module in hf_model.modules(): + hd = getattr(module, "head_dim", None) + if hd is not None and isinstance(hd, int) and hd > 0: + head_dims.add(hd) + if not head_dims: + cfg = getattr(hf_model, "config", None) + if cfg is not None: + hd = getattr(cfg, "head_dim", None) + if hd is None: + num_heads = getattr(cfg, "num_attention_heads", None) + hidden = getattr(cfg, "hidden_size", None) + if num_heads and hidden: + hd = hidden // num_heads + if hd: + head_dims.add(int(hd)) + return sorted(head_dims) + + +def _auto_kernel_options( + hf_model, + device, + prefill_kernel_options: Optional[dict] = None, + decode_kernel_options: Optional[dict] = None, + fa4_prefill: Optional[bool] = None, +) -> tuple[Optional[bool], dict, dict]: + """Derive safe FA4 + Triton block defaults from the GPU + model head_dim + band. Explicit kwargs always win. + + FA4 is OFF by default. The CuTeDSL FLASH backend currently crashes on + short B200 prompts (Llama-3.2-3B, 7-token input hits + ``handle_block_sparse_empty_tile_correction_sm100`` + ``'NoneType' object is not subscriptable``). Users who want FA4 can + pass ``fa4_prefill=True`` explicitly to the engine after confirming it + works on their workload.""" + cap = torch.cuda.get_device_capability(device) if torch.cuda.is_available() else (0, 0) + head_dims = _collect_head_dims(hf_model) or [128] + head_dim_max = max(head_dims) + + if fa4_prefill is None: + fa4_prefill = False # conservative default — see docstring above + + bm_p, bn_p, bm_d, bn_d = _triton_block_defaults(head_dim_max, cap) + + if prefill_kernel_options is None: + if fa4_prefill: + # FA4 path: let the FLASH backend pick its own block shapes. + prefill_kernel_options = None + else: + prefill_kernel_options = { + "FORCE_USE_FLEX_ATTENTION": True, + "BLOCK_M": bm_p, + "BLOCK_N": bn_p, + } + if decode_kernel_options is None: + decode_kernel_options = {"BLOCK_M": bm_d, "BLOCK_N": bn_d} + + return fa4_prefill, prefill_kernel_options, decode_kernel_options + + +# --------------------------------------------------------------------------- +# Arch detection +# --------------------------------------------------------------------------- + + +def _detect_arch(hf_model) -> str: + """Return one of ``"gemma4"``, ``"qwen3"``, ``"llama3"`` or raises.""" + # Look at the inner base model's class; PEFT wrappers delegate to + # ``.base_model.model``. + target = hf_model + for attr in ("base_model", "model"): + inner = getattr(target, attr, None) + if inner is not None and inner is not target: + target = inner + name = type(target).__name__ + # Take from the class hierarchy too so we don't miss Gemma4ForCausalLM + # nested under a PEFT wrapper's ``base_model.model``. + names = [c.__name__ for c in type(hf_model).__mro__] + candidates = set(names + [name]) + lowered = " ".join(n.lower() for n in candidates) + if "gemma4" in lowered or "gemma_4" in lowered or "gemma-4" in lowered: + return "gemma4" + if "qwen3" in lowered: + return "qwen3" + if "llama" in lowered: + return "llama3" + raise NotImplementedError( + "UNSLOTH_FAST_INFERENCE=1 only supports Qwen3, Llama-3, Gemma-4 " + f"today; got {type(hf_model).__name__}. Unset the env var or use vLLM." + ) + + +# --------------------------------------------------------------------------- +# Gemma-4 text-only shell extraction +# --------------------------------------------------------------------------- + + +def _extract_gemma4_text_shell(full_model): + """Return a ``Gemma4ForCausalLM(text_cfg)`` shell wrapping the text + backbone of ``full_model``. Mirrors the CLI path in + ``scripts/benchmarks/gemma4_flex_inference.py:954-968``: strip the + vision + audio towers, point ``shell.model`` at + ``full_model.model.language_model``, tie ``lm_head.weight`` to the + embeddings so the forward pass matches plain HF.""" + from transformers.models.gemma4.modeling_gemma4 import ( + Gemma4ForCausalLM, + ) + + lang = getattr(full_model.model, "language_model", None) + if lang is None: + return full_model # already a text-only CausalLM + full_model.model.vision_tower = None + full_model.model.audio_tower = None + full_model.model.embed_vision = None + full_model.model.embed_audio = None + text_cfg = full_model.config.text_config + shell = Gemma4ForCausalLM(text_cfg) + shell.model = lang + shell.lm_head.weight = lang.embed_tokens.weight + shell = shell.to(next(lang.parameters()).device) + shell.eval() + return shell + + +# --------------------------------------------------------------------------- +# FlexEngine +# --------------------------------------------------------------------------- + + +class _LLMEngineStub: + """Minimal stand-in for ``vllm.LLM.llm_engine``. + + ``unsloth/models/rl.py:104`` reaches into + ``trainer.model.vllm_engine`` to pull the LLM out; `GRPOTrainer`'s + patched ``_move_model_to_vllm`` path calls ``driver_worker.model_runner + .model.load_weights`` but Unsloth's RL patch rewrites those calls to + ``pass`` (rl.py:1795-1810), so a nested attribute chain that simply + exists is enough.""" + + def __init__(self): + self.vllm_config = types.SimpleNamespace(lora_config = types.SimpleNamespace()) + self.model_executor = types.SimpleNamespace( + driver_worker = types.SimpleNamespace( + model_runner = types.SimpleNamespace( + model = types.SimpleNamespace(load_weights = lambda *a, **kw: None), + ) + ) + ) + + +class FlexEngine: + """vLLM-compatible wrapper around :class:`FlexInference` / + :class:`FlexGemma4Inference`. + + Args: + hf_model: The HF model returned by ``from_pretrained``. It will be + swapped to the PEFT-wrapped instance after ``bind_peft_model``. + tokenizer: The companion tokenizer. + dtype: ``torch.bfloat16`` or ``torch.float16``. Used for the autocast + context wrapping forward/prefill/decode. + max_seq_length: Upper bound on ``prompt + completion`` length. + max_lora_rank: Unused by flex today (documented for parity with vLLM). + max_batch_size: Concurrent sequence bucket size for the paged KV. + page_size: Paged-KV page size (tokens per page). + gpu_memory_utilization: Controls ``n_pages``. 0.5 means half the free + VRAM is dedicated to the KV cache. + capture_cudagraph: Capture CUDA graphs on the first decode step. + """ + + def __init__( + self, + hf_model, + tokenizer, + *, + dtype: torch.dtype = torch.bfloat16, + max_seq_length: int = 2048, + max_lora_rank: int = 64, # accepted for API parity; no-op + max_batch_size: int = 32, + page_size: int = 128, + gpu_memory_utilization: float = 0.5, + max_new_tokens: int = 512, + prefill_kernel_options: Optional[dict] = None, + decode_kernel_options: Optional[dict] = None, + fa4_prefill: Optional[bool] = None, + capture_cudagraph: bool = True, + base_model = None, + peft_model = None, + inference_model = None, + ): + assert dtype in (torch.bfloat16, torch.float16), ( + f"FlexEngine requires bf16 or fp16 dtype; got {dtype}." + ) + self.hf_model = hf_model + self.tokenizer = tokenizer + self.compute_dtype = dtype + self.max_seq_length = max_seq_length + self.max_batch_size = max_batch_size + self.page_size = page_size + self.max_new_tokens = max_new_tokens + self.capture_cudagraph = capture_cudagraph + self._cudagraph_primed = False + self._current_lora_int_id: Optional[int] = None + + self.device = hf_model.device + + # Colocate pattern (mirrors vLLM's "colocate" mode): the engine + # runs on its own deep-copy of the HF model so that flex attention + # patching + KV-cache attachment does not mutate the training + # model's forward path. The user's ``model`` stays untouched for + # gradient computation; rollouts go through this copy. + # + # If ``inference_model`` is provided by the caller, use it directly + # (the loader uses this to hand in a copy captured BEFORE Unsloth's + # post-patching; passing the patched model here would break + # rotary_emb / QKV dispatch). + # + # The pristine-base copy (second deep-copy) is LoRA-only — it is + # the refresh source for ``refresh_lora_merge_from_pristine``. We + # defer materialising it until :meth:`bind_peft_model`, so the + # no-LoRA path stays at 2x the base model's VRAM instead of 3x. + if inference_model is None: + inference_model = copy.deepcopy(hf_model) + inference_model.eval() + self._pristine_base = base_model # None until bind_peft_model runs + self._inference_model = inference_model + self._inference_peft = peft_model # filled in by bind_peft_model + + # Autocast wrapping (applied by the impl via a context manager). + fa4_prefill, prefill_kernel_options, decode_kernel_options = ( + _auto_kernel_options( + inference_model, + self.device, + prefill_kernel_options = prefill_kernel_options, + decode_kernel_options = decode_kernel_options, + fa4_prefill = fa4_prefill, + ) + ) + + # Size n_pages from available VRAM so big prompts don't OOM. + n_pages = self._compute_n_pages(gpu_memory_utilization, max_batch_size, page_size) + + arch = _detect_arch(inference_model) + self.arch = arch + if arch == "gemma4": + inference_model = _extract_gemma4_text_shell(inference_model) + self._inference_model = inference_model + Impl = FlexGemma4Inference if arch == "gemma4" else FlexInference + self._impl = Impl( + inference_model, + tokenizer, + max_batch_size = max_batch_size, + max_seq_length = max_seq_length, + n_pages = n_pages, + page_size = page_size, + max_new_tokens = max_new_tokens, + decode_kernel_options = decode_kernel_options, + prefill_kernel_options = prefill_kernel_options, + fa4_prefill = fa4_prefill, + base_model = self._pristine_base, + peft_model = peft_model, + ) + self._llm_engine_stub = _LLMEngineStub() + + # ----- configuration helpers ----- + + def _compute_n_pages( + self, gpu_mem_util: float, max_batch: int, page_size: int + ) -> int: + """Size the paged-KV allocation. + + ``n_pages`` = ceil(``max_batch * max_seq_length / page_size``) with a + small headroom factor; any more and we waste VRAM on ghost pages + that the scheduler never fills. ``gpu_memory_utilization`` scales the + headroom.""" + min_pages = max( + 1, max_batch * ((self.max_seq_length + page_size - 1) // page_size) + ) + factor = 1.0 + max(0.1, min(1.0, gpu_mem_util)) + return max(min_pages, int(min_pages * factor)) + + # ----- generate ----- + + def generate( + self, + prompts = None, + sampling_params = None, + lora_request = None, + use_tqdm: bool = False, + **kwargs, + ): + """Drop-in replacement for ``vllm.LLM.generate``. + + Returns a list of :class:`RequestOutput` mirroring vLLM's shape. + """ + if prompts is None and "prompt_token_ids" in kwargs: + prompts = kwargs.pop("prompt_token_ids") + prompts = self._normalize_prompts(prompts) + + max_new_tokens, _ = self._extract_sampling(sampling_params) + + # Apply LoRA: if the request carries tensors, refresh the merged copy. + if lora_request is not None: + self._apply_lora_request(lora_request) + + seqs = self._build_sequences(prompts, max_new_tokens) + + with torch.amp.autocast("cuda", dtype = self.compute_dtype): + done = self._impl.generate( + seqs, + capture_cudagraph = self.capture_cudagraph and not self._cudagraph_primed, + ) + self._cudagraph_primed = self.capture_cudagraph + + # Reorder by the input prompt index (the impl returns done seqs in + # completion order, not input order). + done_by_input = sorted(done, key = lambda s: getattr(s, "_input_idx", 0)) + outs = [] + for idx, seq in enumerate(done_by_input): + text = self.tokenizer.decode(seq.output_ids, skip_special_tokens = False) + prompt_text = getattr(seq, "text", "") or self.tokenizer.decode( + seq.input_ids.tolist() if seq.input_ids is not None else [], + skip_special_tokens = False, + ) + co = CompletionOutput( + index = 0, + text = text, + token_ids = list(seq.output_ids), + finish_reason = ( + "stop" if seq.last_token_id == self._impl.eos_token_id else "length" + ), + ) + ro = RequestOutput( + request_id = str(idx), + prompt = prompt_text, + prompt_token_ids = ( + seq.input_ids.tolist() if seq.input_ids is not None else [] + ), + outputs = [co], + ) + outs.append(ro) + return outs + + def chat( + self, + messages, + sampling_params = None, + lora_request = None, + use_tqdm: bool = False, + **kwargs, + ): + """Apply the tokenizer's chat template and defer to :meth:`generate`.""" + if messages is None: + return [] + # ``messages`` from vLLM is either a list of message-lists or a single + # message-list (one conversation). + if isinstance(messages, list) and messages and isinstance(messages[0], dict): + # Single conversation. + convos = [messages] + else: + convos = list(messages) + prompts = [ + self.tokenizer.apply_chat_template( + m, tokenize = False, add_generation_prompt = True + ) + for m in convos + ] + return self.generate( + prompts, + sampling_params = sampling_params, + lora_request = lora_request, + use_tqdm = use_tqdm, + **kwargs, + ) + + # ----- LoRA refresh ----- + + def _apply_lora_request(self, lora_request): + """Copy the training LoRA tensors onto the inference-side PEFT + wrapper and refresh the merged inference weights from the pristine + base. + + Expects ``lora_request.lora_tensors`` to be a ``state_dict`` slice + filtered for ``.lora_A.`` / ``.lora_B.`` keys (which is what + :func:`~unsloth.inference.vllm_shim.load_lora` produces).""" + if lora_request is None: + return + base_model = getattr(self._impl, "base_model", None) + peft_model = getattr(self._impl, "peft_model", None) + if base_model is None or peft_model is None: + # 4-bit / no double-copy: LoRA tensors already live on the PEFT + # wrappers (bnb-4bit packed weights can't be in-place refreshed; + # PEFT's three-matmul wrapper does the math at runtime). + return + + tensors = getattr(lora_request, "lora_tensors", None) + if tensors: + target_sd = peft_model.state_dict() + renamed = {} + for k, v in tensors.items(): + # ``load_lora`` strips ``.default``; PEFT's own state_dict + # keeps it. Try both forms. + if k in target_sd: + renamed[k] = v + continue + candidate = k.replace(".lora_A.", ".lora_A.default.").replace( + ".lora_B.", ".lora_B.default." + ) + if candidate in target_sd: + renamed[candidate] = v + if renamed: + missing, unexpected = peft_model.load_state_dict( + renamed, strict = False + ) + # ``missing`` will be every non-LoRA param; that's fine. + refresh_lora_merge_from_pristine(base_model, peft_model) + self._current_lora_int_id = getattr(lora_request, "lora_int_id", None) + + # ----- prompt normalization ----- + + @staticmethod + def _normalize_prompts(prompts): + if prompts is None: + return [] + if isinstance(prompts, str): + return [prompts] + if isinstance(prompts, list): + if not prompts: + return [] + # Already a list. Leave dict / str / list[int] elements as-is. + return prompts + return [prompts] + + def _build_sequences(self, prompts, max_new_tokens: int) -> list: + seqs = [] + for idx, p in enumerate(prompts): + if isinstance(p, str): + seq = Sequence(text = p, max_new_tokens = max_new_tokens) + elif isinstance(p, dict): + # vLLM TokensPrompt / TextPrompt dict + if "prompt_token_ids" in p: + ids = torch.tensor(p["prompt_token_ids"], dtype = torch.long) + seq = Sequence( + text = p.get("prompt", ""), + input_ids = ids, + input_length = int(ids.shape[0]), + max_new_tokens = max_new_tokens, + ) + elif "prompt" in p: + seq = Sequence( + text = p["prompt"], max_new_tokens = max_new_tokens, + ) + else: + raise ValueError( + f"Unsupported prompt dict (no 'prompt' or " + f"'prompt_token_ids'): {list(p)}" + ) + elif isinstance(p, (list, tuple)) and p and isinstance(p[0], int): + ids = torch.tensor(list(p), dtype = torch.long) + seq = Sequence( + text = "", + input_ids = ids, + input_length = int(ids.shape[0]), + max_new_tokens = max_new_tokens, + ) + else: + raise ValueError( + "FlexEngine.generate accepts str, list[int], or a " + "TokensPrompt/TextPrompt dict; got " + f"{type(p).__name__} at index {idx}." + ) + seq._input_idx = idx # preserve input order across the scheduler + seqs.append(seq) + return seqs + + def _extract_sampling(self, sampling_params) -> tuple[int, dict]: + """Pull what we actually honour out of a ``SamplingParams``. + + Today the flex path does argmax sampling only. We read ``max_tokens`` + (mapped to ``max_new_tokens``) and ignore ``temperature`` / ``top_p`` + / ``top_k`` with a warning the first time we see something non-greedy. + Sufficient for GRPO's on-policy rollout, which already accepts that + the generation is deterministic per-prompt under a fixed seed.""" + if sampling_params is None: + return self.max_new_tokens, {} + max_tokens = getattr(sampling_params, "max_tokens", None) + if max_tokens is None: + max_tokens = self.max_new_tokens + temp = getattr(sampling_params, "temperature", 0.0) or 0.0 + if temp and temp > 0 and not getattr(self, "_warned_sampling", False): + warnings.warn( + "FlexEngine (UNSLOTH_FAST_INFERENCE=1) does argmax sampling " + "only; sampling_params.temperature / top_p / top_k are " + "ignored. Unset the env var or use vLLM for stochastic " + "sampling.", + RuntimeWarning, + stacklevel = 3, + ) + self._warned_sampling = True + return int(max_tokens), {} + + # ----- sleep-mode stubs ----- + # + # vLLM's sleep mode offloads engine weights to CPU between rollouts so + # training can use the freed VRAM. The flex backend does not implement + # that yet; these stubs exist so code paths that assume the API are safe. + + def sleep(self, level: int = 2): + """No-op. Real implementation in a follow-up PR (move PagedKVCache + + shared buffers + the inference-copy HF shell to CPU pinned memory + and swap back on wake_up).""" + return None + + def wake_up(self, tags: Optional[list] = None): + """No-op companion to :meth:`sleep`.""" + return None + + @property + def llm_engine(self): + """Nested-attribute stub used by Unsloth's RL patch — see docstring + on :class:`_LLMEngineStub`.""" + return self._llm_engine_stub + + # ----- PEFT wiring ----- + + def bind_peft_model(self, training_peft_model): + """Hook called at the end of ``FastLanguageModel.get_peft_model``. + + * Updates ``self.hf_model`` so ``save_lora`` / ``load_lora`` — which + call ``model.state_dict()`` — see the training LoRA weights. + * Materialises the pristine-base copy (second deep-copy of the + inference model's un-patched weights) used as the refresh source + for ``refresh_lora_merge_from_pristine``. This is lazy so the + no-LoRA path stays at 2x the base model's VRAM. + * Mirrors the training PEFT config onto the inference copy by + wrapping ``self._inference_model`` with a matching ``PeftModel``. + The inference-side adapter starts at zero; LoRA tensors are then + loaded from the request on each ``generate`` call, and + ``refresh_lora_merge_from_pristine`` fuses them into the merged + base weights before the kernel is launched. + * The deep-copy pattern keeps the training model's forward pass + untouched (its attention is not flex-patched).""" + self.hf_model = training_peft_model + + # Materialise the pristine source the first time we see a LoRA. + if self._pristine_base is None: + # The inference model has already been flex-patched; its + # linear weights (what LoRA merges into) are still pristine, + # so we can clone it and just not call flex attention on the + # pristine copy. + self._pristine_base = copy.deepcopy(self._inference_model) + self._pristine_base.eval() + self._impl.base_model = self._pristine_base + + if self._inference_peft is None: + try: + from peft import get_peft_model as _get_peft_model + + peft_cfg = training_peft_model.peft_config["default"] + # Wrap the already-patched inference copy with a fresh LoRA + # adapter of the same shape. LoraLayer insertion is + # attention-forward-agnostic; it wraps Linear modules. + self._inference_peft = _get_peft_model( + self._inference_model, peft_cfg + ) + self._inference_peft.eval() + except Exception as e: + warnings.warn( + f"FlexEngine.bind_peft_model: could not build an " + f"inference-side PEFT wrapper ({e}). Falling back to " + "single-copy mode — LoRA refresh will rely on the " + "training model's state_dict directly.", + RuntimeWarning, + stacklevel = 2, + ) + self._inference_peft = training_peft_model + + self._impl.peft_model = self._inference_peft + + +# --------------------------------------------------------------------------- +# load_flex — mirrors unsloth_zoo.vllm_utils.load_vllm's call shape so the +# dispatcher in unsloth/models/llama.py + vision.py can route cleanly. +# --------------------------------------------------------------------------- + + +def load_flex( + hf_model, + tokenizer, + *, + dtype: torch.dtype = torch.bfloat16, + max_seq_length: int = 2048, + max_lora_rank: int = 64, + max_batch_size: int = 32, + gpu_memory_utilization: float = 0.5, + page_size: int = 128, + capture_cudagraph: bool = True, + base_model = None, + peft_model = None, + **_unused_vllm_kwargs, +) -> FlexEngine: + """Construct a :class:`FlexEngine` around an already-loaded HF model. + + ``_unused_vllm_kwargs`` swallows vLLM-only kwargs passed through by + `unsloth/models/llama.py` (``use_bitsandbytes``, ``enable_lora``, + ``disable_log_stats``, ``fp8_mode``, ``float8_kv_cache``, ...) so the + dispatcher doesn't have to know which backend is selected.""" + return FlexEngine( + hf_model, + tokenizer, + dtype = dtype, + max_seq_length = max_seq_length, + max_lora_rank = max_lora_rank, + max_batch_size = max_batch_size, + gpu_memory_utilization = gpu_memory_utilization, + page_size = page_size, + capture_cudagraph = capture_cudagraph, + base_model = base_model, + peft_model = peft_model, + ) + + +__all__ = ["FlexEngine", "load_flex"] diff --git a/unsloth/inference/flex_gemma4.py b/unsloth/inference/flex_gemma4.py new file mode 100644 index 0000000000..42664c1469 --- /dev/null +++ b/unsloth/inference/flex_gemma4.py @@ -0,0 +1,1173 @@ +"""Gemma-4-E2B-it inference with flex_attention + paged KV cache + CUDA graphs. + +Extends the Qwen3/Llama-3.2 engine in `qwen3_flex_inference.py` to a third +architecture, `unsloth/gemma-4-E2B-it`. Gemma-4 is not a drop-in addition: +its text backbone diverges from Qwen3/Llama in ways that cannot be folded +into a single `hasattr(self, "q_norm")` branch. The divergences, and how +this file handles them: + +1. KV-sharing layers. E2B has 35 layers; the upper 15 lack `k_proj`, + `v_proj`, `k_norm`, `v_norm` entirely and consume the K/V produced + by a "store" layer further up the stack. We link each shared layer's + `_paged_cache` to the store layer's `PagedKVCache` so flex_attention + reads the same pages that the store layer populated a few layers + earlier -- no sidecar, no SDPA fallback, one block mask per regime + works for every layer. +2. Dual attention types. Each layer is either `full_attention` + (`head_dim=512`, `rope_theta=1e6`, `partial_rotary_factor=0.25`) or + `sliding_attention` (`head_dim=256`, `rope_theta=10000`, + `sliding_window=512`). We precompute both (cos, sin) pairs once per + forward and dispatch on `self.layer_type`. +3. Per-layer input embeddings. `embed_tokens_per_layer` produces a + `[B, S, num_layers, 256]` auxiliary table that enters every layer + through a `per_layer_input_gate -> act -> mul -> per_layer_projection + -> post_per_layer_input_norm -> +residual` path after the MLP residual. +4. Four norms per layer. `input_layernorm` / `post_attention_layernorm` + wrap the attention block (double residual); `pre_feedforward_layernorm` + / `post_feedforward_layernorm` wrap the MLP (double residual). A scalar + `layer_scalar` multiplies hidden_states at layer end. +5. Final logit softcap. `logits = tanh(logits / 30.0) * 30.0` applied on + the lm_head output. + +The engine is text-only: `Gemma4ForCausalLM(text_config)` skips the +multimodal `Gemma4ForConditionalGeneration` wrapper and its vision + audio +towers entirely. Shared helpers (`PagedKVCache`, `PageTable`, `Sequence`, +`refresh_lora_merge_from_pristine`, `run_drift_verification`, +`flex_attention_compiled`, `_apply_rotary`, FA4 capability guard) are +imported from `qwen3_flex_inference.py` unchanged. + +Run: + CUDA_VISIBLE_DEVICES=6 python scripts/benchmarks/gemma4_flex_inference.py \ + --n_prompts 64 --max_new_tokens 512 --capture_cudagraph \ + --stats_path logs/flex_gemma4_bf16.json + +Requires `transformers>=5.5.0` (the `gemma4` module). The main workspace +env stays on 4.57.6; this file short-circuits with a clear install hint +if the module is missing. Use `isolated_run.py` with +`--extra_packages "transformers>=5.5.0 peft datasets"` to run on that env. +""" + +from __future__ import annotations + +import argparse +import copy +import json +import os +import sys +import time +import types +from collections import deque +from pathlib import Path +from typing import Optional + +import torch +import torch.nn.functional as F +from torch.nn.attention.flex_attention import BlockMask + +HERE = Path(__file__).resolve().parent + +try: + from .flex_qwen3_llama import ( + DECODE_KERNEL_OPTIONS_DEFAULT, + PREFILL_KERNEL_OPTIONS_DEFAULT, + Sequence, + _apply_rotary, + _hash_state_dict, + _lora_needs_peft_fallback, + flex_attention_compiled, + refresh_lora_merge_from_pristine, + run_drift_verification, + ) + from .flex_paged_attention import PagedKVCache, PageTable +except ImportError: # script-mode fallback (scripts/benchmarks CLI shim) + sys.path.insert(0, str(HERE)) + from qwen3_flex_inference import ( # noqa: E402 + DECODE_KERNEL_OPTIONS_DEFAULT, + PREFILL_KERNEL_OPTIONS_DEFAULT, + Sequence, + _apply_rotary, + _hash_state_dict, + _lora_needs_peft_fallback, + flex_attention_compiled, + refresh_lora_merge_from_pristine, + run_drift_verification, + ) + from flex_paged_attention import PagedKVCache, PageTable # noqa: E402 +from torch.nn.attention.flex_attention import create_block_mask as _create_block_mask # noqa: E402 + + +# --- sliding-window block mask helpers ------------------------------------ +# +# Gemma-4's `sliding_attention` layers attend only to the last +# `sliding_window` KV positions (`q_idx - kv_idx < W`) in addition to the +# standard causal mask. The shared helpers in `flex_paged_attention.py` +# only expose the pure-causal builders, so we build sliding variants here +# (local to this file so that file is untouched). + + +def _causal_blockmask_with_window(B: int, L: int, block_size: int, window: int, device: str): + def causal_windowed(b, h, q_idx, kv_idx): + return (q_idx >= kv_idx) & (q_idx - kv_idx < window) + + return _create_block_mask( + causal_windowed, + B = B, + H = None, + Q_LEN = L, + KV_LEN = L, + BLOCK_SIZE = block_size, + device = device, + ) + + +def _prefill_blockmask_with_window(batch_idx: torch.Tensor, block_size, window: int): + assert batch_idx.ndim == 2 and batch_idx.shape[0] == 1 + L = batch_idx.shape[1] + docs = batch_idx.view(-1) + + def document_causal_windowed(b, h, q_idx, kv_idx): + causal_mask = q_idx >= kv_idx + window_mask = q_idx - kv_idx < window + document_mask = docs[q_idx] == docs[kv_idx] + return causal_mask & window_mask & document_mask + + return _create_block_mask( + document_causal_windowed, + B = 1, + H = None, + Q_LEN = L, + KV_LEN = L, + BLOCK_SIZE = block_size, + ) + + +# --- transformers version guard -------------------------------------------- +# +# Gemma-4 lands in `transformers>=5.5.0`. The main workspace env is on +# 4.57.6, which keeps the Qwen3 + Llama paths in qwen3_flex_inference.py +# working unchanged. We defer the import to call time so `--help` still +# works on 4.57.6. + + +def _require_gemma4(): + try: + from transformers.models.gemma4.modeling_gemma4 import Gemma4ForCausalLM + from transformers.models.gemma4.configuration_gemma4 import ( + Gemma4Config, + Gemma4TextConfig, + ) + except ImportError: + import transformers + + raise SystemExit( + f"Gemma-4 requires transformers>=5.5.0 (`gemma4` module). " + f"Current: transformers=={transformers.__version__}. " + f"Install: uv pip install 'transformers>=5.5.0'" + ) + return Gemma4ForCausalLM, Gemma4Config, Gemma4TextConfig + + +# --- attention forward factory -------------------------------------------- + + +def _apply_rotary_q(q, cos, sin): + """Rotary on Q alone; used on shared layers where K is read + pre-rotated from the store layer's paged cache.""" + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + + def rotate_half(x): + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim = -1) + + return (q * cos) + (rotate_half(q) * sin) + + +def make_flex_gemma4_attention_forward(page_table: PageTable): + """Return a new `forward` method for `Gemma4TextAttention` that routes + through flex_attention against a paged KV cache. + + Two layer kinds: + - non-shared (`self.is_kv_shared_layer == False`): standard q/k/v + projection; writes new K/V into `self._paged_cache` + (which is the layer's own PagedKVCache). + - shared (`self.is_kv_shared_layer == True`): no `k_proj`/ + `v_proj`/`k_norm`/`v_norm`. Reads K/V directly from + the *store* layer's paged cache, which the patching + helper has already linked onto `self._paged_cache`. + No write -- the store layer populated the cache for + the same positions earlier in the walker, so the + shared layer just attends over those pages with the + identical block mask. + + Linking shared layers to the store layer's paged cache keeps one + block-mask + one KV layout across the whole stack and lets + flex_attention handle every layer uniformly (no sidecar, no SDPA + fallback, one CUDA graph capture). + + `position_embeddings` is a dict keyed by `layer_type`; we pick the + right (cos, sin) pair before rotary. + + `self.v_proj` may be None when the config sets `attention_k_eq_v` + (global head dim with shared K=V); that branch reuses k_raw. + """ + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: dict, + attention_mask = None, + past_key_values = None, + cache_position = None, + flex_block_mask: Optional[dict] = None, + flex_input_pos: Optional[torch.Tensor] = None, + flex_batch_idx: Optional[torch.Tensor] = None, + flex_kernel_options: Optional[dict] = None, + **kwargs, + ): + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + cos, sin = position_embeddings[self.layer_type] + + q = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2) + + if getattr(self, "is_kv_shared_layer", False): + # Shared layer reads from the paired store layer's K/V. + # `PagedKVCache.update` returns different shapes for prefill + # vs decode: + # - prefill: the packed k_val/v_val [1, H, L_packed, D] + # - decode : the full paged pool k_cache/v_cache + # [1, H, n_pages*page_size, D] + # The prefill block_mask is sized for L_packed and the decode + # block_mask is sized for the paged pool, so we need to match + # the same shape here. The store layer stashes its + # post-rotary k/v as `_last_k_val` / `_last_v_val` during + # prefill; at decode time we read from its `_paged_cache` + # (the same buffer the shared layer was linked to at patch). + q = _apply_rotary_q(q, cos, sin) + store_attn = self._store_attn + if q.shape[-2] > 1: + k = store_attn._last_k_val + v = store_attn._last_v_val + else: + k = self._paged_cache.k_cache + v = self._paged_cache.v_cache + else: + k_raw = self.k_proj(hidden_states).view(hidden_shape) + k = self.k_norm(k_raw).transpose(1, 2) + # `v_proj` may be None under Gemma-4's K=V global-attention + # option; in that case reuse the raw (un-normed) k projection. + if self.v_proj is not None: + v = self.v_norm( + self.v_proj(hidden_states).view(hidden_shape) + ).transpose(1, 2) + else: + v = k_raw.transpose(1, 2) + q, k = _apply_rotary(q, k, cos, sin) + + # Store layers stash the post-rotary k/v so any shared + # successors can read the same packed prefill tensors. This + # assignment is a pointer rebind, not a copy; CUDA graph + # capture sees a stable attribute reference. Plain + # non-shared layers don't need this. + if getattr(self, "store_full_length_kv", False): + self._last_k_val = k + self._last_v_val = v + + if self._paged_cache is not None and flex_input_pos is not None: + k, v = self._paged_cache.update(flex_input_pos, k, v, flex_batch_idx) + + # flex_block_mask is a dict keyed by layer_type; pick the one + # matching this layer's regime (full_attention vs sliding_attention). + block_mask = flex_block_mask[self.layer_type] + attn_output = flex_attention_compiled( + q, + k, + v, + scale = self.scaling, + block_mask = block_mask, + enable_gqa = True, + kernel_options = flex_kernel_options, + ) + attn_output = attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous() + return self.o_proj(attn_output), None + + return forward + + +def patch_gemma4_attention_forwards( + model: torch.nn.Module, page_table: PageTable +): + """Attach a PagedKVCache to every non-shared attention layer, link + every shared attention layer to its store layer's cache, and swap in + the flex_attention forward above. + + Three passes: + 1. Allocate a PagedKVCache on each non-shared layer (the cache + shape depends on that layer's head_dim and num_kv_heads, which + vary across Gemma-4 layers). + 2. Walk shared layers and set + `shared._paged_cache = store._paged_cache`, where `store` is + `model.model.layers[shared.kv_shared_layer_index]`. The shared + layer's forward reads `k_cache` / `v_cache` directly; the store + layer's `update()` writes populate the same tensors. + 3. Bind the flex forward. + """ + fwd = make_flex_gemma4_attention_forward(page_table) + for layer in model.model.layers: + attn = layer.self_attn + if getattr(attn, "is_kv_shared_layer", False): + continue + n_kv = attn.k_proj.out_features // attn.head_dim + attn._paged_cache = PagedKVCache( + page_table, + n_heads = n_kv, + head_dim = attn.head_dim, + dtype = model.dtype, + ).to(model.device) + for layer in model.model.layers: + attn = layer.self_attn + if not getattr(attn, "is_kv_shared_layer", False): + continue + store_attn = model.model.layers[attn.kv_shared_layer_index].self_attn + attn._paged_cache = store_attn._paged_cache + attn._store_attn = store_attn + for layer in model.model.layers: + layer.self_attn.forward = types.MethodType(fwd, layer.self_attn) + + +# --- model forward walker -------------------------------------------------- + + +def call_gemma4_model_with_flex_kwargs( + model, input_ids, position_ids, flex_kwargs +): + """Walk the Gemma-4 text model manually so we can inject flex_* kwargs + into each attention call. Mirrors `call_model_with_flex_kwargs` in + `qwen3_flex_inference.py` but: + + - precomputes both (cos, sin) variants and passes them as a dict + keyed by `layer.self_attn.layer_type`; + - materializes `per_layer_inputs` via the model's own + `get_per_layer_inputs` + `project_per_layer_inputs` helpers; + - runs the double-residual around attn, the double-residual around + MLP, the per-layer-input path, and the `layer_scalar` multiply. + + The final norm is applied here; lm_head + softcap is applied by the + caller (so we can slice to logits_positions before the vocab matmul). + """ + base = model.model + + inputs_embeds = base.embed_tokens(input_ids) + + # Per-layer input table. E2B has hidden_size_per_layer_input = 256 and + # 35 layers, so this is [B, S, 35, 256] -- a local tensor with fixed + # shape across CUDA graph replays (input_ids is pre-allocated upstream). + per_layer_inputs = None + if getattr(base, "hidden_size_per_layer_input", 0): + per_layer_inputs = base.get_per_layer_inputs(input_ids, inputs_embeds) + per_layer_inputs = base.project_per_layer_inputs( + inputs_embeds, per_layer_inputs + ) + + position_embeddings = { + layer_type: base.rotary_emb(inputs_embeds, position_ids, layer_type) + for layer_type in base.unique_layer_types + } + + hidden_states = inputs_embeds + for i, layer in enumerate(base.layers): + residual = hidden_states + hidden_states = layer.input_layernorm(hidden_states) + hidden_states, _ = layer.self_attn( + hidden_states, + position_embeddings = position_embeddings, + **flex_kwargs, + ) + hidden_states = layer.post_attention_layernorm(hidden_states) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = layer.pre_feedforward_layernorm(hidden_states) + hidden_states = layer.mlp(hidden_states) + hidden_states = layer.post_feedforward_layernorm(hidden_states) + hidden_states = residual + hidden_states + + if per_layer_inputs is not None and hasattr(layer, "per_layer_input_gate"): + residual = hidden_states + hidden_states = layer.per_layer_input_gate(hidden_states) + hidden_states = layer.act_fn(hidden_states) + hidden_states = hidden_states * per_layer_inputs[:, :, i, :] + hidden_states = layer.per_layer_projection(hidden_states) + hidden_states = layer.post_per_layer_input_norm(hidden_states) + hidden_states = residual + hidden_states + + if hasattr(layer, "layer_scalar"): + hidden_states = hidden_states * layer.layer_scalar + + hidden_states = base.norm(hidden_states) + return hidden_states + + +# --- inference engine ------------------------------------------------------ + + +class FlexGemma4Inference: + def __init__( + self, + model, + tokenizer, + max_batch_size = 32, + max_seq_length = 2048, + n_pages = 2048, + page_size = 128, + max_new_tokens = 512, + decode_kernel_options = None, + prefill_kernel_options = None, + fa4_prefill = None, + base_model = None, + peft_model = None, + ): + assert max_seq_length % page_size == 0 + self.model = model + self.tokenizer = tokenizer + self.device = model.device + self.eos_token_id = tokenizer.eos_token_id + self.base_model = base_model + self.peft_model = peft_model + self.max_batch_size = max_batch_size + self.max_seq_length = max_seq_length + self.page_size = page_size + self.max_new_tokens = max_new_tokens + + if fa4_prefill is None or fa4_prefill: + major, _ = torch.cuda.get_device_capability(self.device) + supported = major >= 9 + if fa4_prefill and not supported: + import warnings + + warnings.warn( + f"--fa4_prefill needs Hopper (sm_90) or Blackwell " + f"(sm_100 / sm_120); found sm_{major}0. Falling back " + f"to the Triton flex_attention backend.", + RuntimeWarning, + stacklevel = 2, + ) + fa4_prefill = supported + self.fa4_prefill = fa4_prefill + self.prefill_q_block = 256 if fa4_prefill else 128 + self.prefill_kv_block = 128 + + self.decode_kernel_options = ( + decode_kernel_options + if decode_kernel_options is not None + else DECODE_KERNEL_OPTIONS_DEFAULT + ) + base_prefill_opts = ( + prefill_kernel_options + if prefill_kernel_options is not None + else dict(PREFILL_KERNEL_OPTIONS_DEFAULT) + ) + if fa4_prefill: + base_prefill_opts = dict(base_prefill_opts) + base_prefill_opts.pop("FORCE_USE_FLEX_ATTENTION", None) + base_prefill_opts["BACKEND"] = "FLASH" + self.prefill_kernel_options = base_prefill_opts + + self.page_table = PageTable( + n_pages = n_pages, + page_size = page_size, + max_batch_size = max_batch_size, + device = self.device.type, + ) + + patch_gemma4_attention_forwards(model, self.page_table) + + self.input_pos_buffer = torch.zeros( + max_batch_size, dtype = torch.int32, device = self.device + ) + + # Detect the sliding window from any sliding-attention layer. + # We need one block mask per layer type: `full_attention` is pure + # causal; `sliding_attention` is causal AND q_pos - kv_pos < window. + sliding_window = None + for layer in model.model.layers: + if getattr(layer.self_attn, "is_sliding", False): + sliding_window = layer.self_attn.sliding_window + break + + self.sliding_window = sliding_window + self.block_mask_logical_by_type = { + "full_attention": self.page_table.create_causal_blockmask( + B = max_batch_size, L = max_seq_length + ), + } + if sliding_window is not None: + self.block_mask_logical_by_type["sliding_attention"] = ( + _causal_blockmask_with_window( + B = max_batch_size, + L = max_seq_length, + block_size = page_size, + window = sliding_window, + device = self.device.type, + ) + ) + # Legacy alias used by the original page-aware decode slicer. + self.block_mask_logical = self.block_mask_logical_by_type["full_attention"] + + self.cudagraph_captured = False + self.graphs = {} + self.graph_vars = {} + + def tokenize(self, sequences): + for seq in sequences: + if seq.input_ids is not None and seq.input_length > 0: + # Pre-tokenized input (see FlexEngine). Skip. + continue + ids = self.tokenizer(seq.text, return_tensors = "pt")["input_ids"].squeeze(0) + seq.input_ids = ids + seq.input_length = ids.shape[0] + + def _softcap(self, logits): + sc = getattr(self.model.config, "final_logit_softcapping", None) + if sc is not None and sc > 0: + logits = torch.tanh(logits / sc) * sc + return logits + + def _prefill(self, batch: list) -> torch.Tensor: + input_ids_list = [seq.input_ids.to(self.device) for seq in batch] + input_pos_list = [ + torch.arange(seq.input_length, dtype = torch.long, device = self.device) + for seq in batch + ] + batch_idx_list = [ + torch.full( + (seq.input_length,), + seq.batch_idx, + dtype = torch.long, + device = self.device, + ) + for seq in batch + ] + input_ids = torch.cat(input_ids_list).view(1, -1) + input_pos = torch.cat(input_pos_list).view(1, -1) + batch_idx = torch.cat(batch_idx_list).view(1, -1) + + L = input_ids.shape[1] + q_block = self.prefill_q_block + pad = (q_block - L % q_block) % q_block + if pad > 0: + input_ids = F.pad(input_ids, (0, pad), value = 0) + input_pos = F.pad(input_pos, (0, pad), value = 0) + batch_idx = F.pad(batch_idx, (0, pad), value = 0) + + input_lengths = torch.tensor( + [s.input_length for s in batch], dtype = torch.long, device = self.device + ) + logits_positions = input_lengths.cumsum(dim = 0) - 1 + + prefill_block_size = ( + (self.prefill_q_block, self.prefill_kv_block) + if self.fa4_prefill + else self.prefill_q_block + ) + mask_full = self.page_table.create_prefill_blockmask_no_paging( + batch_idx, BLOCK_SIZE = prefill_block_size + ) + masks_by_type = {"full_attention": mask_full} + if self.sliding_window is not None: + masks_by_type["sliding_attention"] = _prefill_blockmask_with_window( + batch_idx, + block_size = prefill_block_size, + window = self.sliding_window, + ) + + flex_kwargs = dict( + flex_block_mask = masks_by_type, + flex_input_pos = input_pos, + flex_batch_idx = batch_idx, + flex_kernel_options = self.prefill_kernel_options, + ) + hidden = call_gemma4_model_with_flex_kwargs( + self.model, input_ids, input_pos, flex_kwargs + ) + logits = self.model.lm_head(hidden[:, logits_positions, :]).squeeze(0) + return self._softcap(logits) + + def _decode_block_mask(self, batch_idx: torch.Tensor): + """Slice one row of the logical decode mask per sequence, for + both full and sliding regimes. Returns a dict keyed by layer_type + plus the raw `input_pos` tensor (needed for PageTable conversion).""" + input_pos = self.input_pos_buffer[batch_idx] + assert batch_idx.ndim == 1 and input_pos.ndim == 1 + B = batch_idx.shape[0] + + def _slice(block_mask, extra_mask_mod): + input_block_idx = input_pos // block_mask.BLOCK_SIZE[0] + kv_num_blocks = block_mask.kv_num_blocks[ + batch_idx, :, input_block_idx + ].view(B, 1, 1) + kv_indices = block_mask.kv_indices[ + batch_idx, :, input_block_idx + ].view(B, 1, 1, -1) + full_num = full_idx = None + if block_mask.full_kv_num_blocks is not None: + full_num = block_mask.full_kv_num_blocks[ + batch_idx, :, input_block_idx + ].view(B, 1, 1) + full_idx = block_mask.full_kv_indices[ + batch_idx, :, input_block_idx + ].view(B, 1, 1, -1) + + seq_length = (1, block_mask.seq_lengths[1]) + return BlockMask.from_kv_blocks( + kv_num_blocks, + kv_indices, + full_num, + full_idx, + BLOCK_SIZE = block_mask.BLOCK_SIZE, + mask_mod = extra_mask_mod, + seq_lengths = seq_length, + ) + + def causal_offset(off): + def m(b, h, q_idx, kv_idx): + return q_idx + off[b] >= kv_idx + + return m + + def causal_offset_windowed(off, window): + def m(b, h, q_idx, kv_idx): + return (q_idx + off[b] >= kv_idx) & ( + q_idx + off[b] - kv_idx < window + ) + + return m + + masks = { + "full_attention": _slice( + self.block_mask_logical_by_type["full_attention"], + causal_offset(input_pos), + ), + } + if self.sliding_window is not None: + masks["sliding_attention"] = _slice( + self.block_mask_logical_by_type["sliding_attention"], + causal_offset_windowed(input_pos, self.sliding_window), + ) + return masks, input_pos + + def _decode_step_eager(self, batch_idx: torch.Tensor, input_ids: torch.Tensor): + B = input_ids.shape[0] + masks, input_pos = self._decode_block_mask(batch_idx) + # Convert each regime's block mask through the page table so the + # logical→physical kv page mapping is correct for every layer. + masks = { + k: self.page_table.convert_logical_block_mask(m, batch_idx) + for k, m in masks.items() + } + position_ids = input_pos.view(B, 1).to(torch.long) + flex_kwargs = dict( + flex_block_mask = masks, + flex_input_pos = input_pos.view(B, 1).to(torch.long), + flex_batch_idx = batch_idx, + flex_kernel_options = self.decode_kernel_options, + ) + hidden = call_gemma4_model_with_flex_kwargs( + self.model, input_ids.view(B, 1), position_ids, flex_kwargs + ) + logits = self.model.lm_head(hidden[:, -1, :]) + return self._softcap(logits) + + def _decode_step( + self, + batch_idx: torch.Tensor, + input_ids: torch.Tensor, + input_pos: torch.Tensor, + ): + self.input_pos_buffer.zero_() + self.input_pos_buffer[batch_idx] = input_pos + if not self.cudagraph_captured: + return self._decode_step_eager(batch_idx, input_ids) + bs = input_ids.size(0) + key = next(x for x in self.graph_bs if x >= bs) + graph = self.graphs[key] + gv = self.graph_vars + for k, v in gv.items(): + if k != "outputs": + v.zero_() + gv["input_ids"][:bs] = input_ids + gv["batch_idx"][:bs] = batch_idx + graph.replay() + return gv["outputs"][:bs] + + def capture_decode_cudagraph(self): + max_bs = self.max_batch_size + reserved_batches = [] + for bi in range(1, max_bs): + try: + allocated = self.page_table.allocate() + self.page_table.reserve( + allocated, + torch.tensor( + [allocated], device = self.device, dtype = torch.long + ), + self.page_size, + ) + reserved_batches.append(allocated) + except Exception: + break + + input_ids = torch.zeros(max_bs, dtype = torch.int64, device = self.device) + batch_idx = torch.arange(max_bs, dtype = torch.int64, device = self.device) + outputs = torch.zeros( + (max_bs, self.model.config.vocab_size), + dtype = self.model.dtype, + device = self.device, + ) + self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16)) + pool = None + for bs in reversed(self.graph_bs): + if bs > max_bs: + continue + print(f"[flex-gemma4] capturing CUDA graph for bs={bs}") + torch.cuda.synchronize() + _ = self._decode_step_eager(batch_idx[:bs], input_ids[:bs]) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool): + outputs[:bs] = self._decode_step_eager(batch_idx[:bs], input_ids[:bs]) + if pool is None: + pool = graph.pool() + self.graphs[bs] = graph + torch.cuda.synchronize() + for bi in reserved_batches: + self.page_table.erase(bi) + self.graph_vars = dict( + input_ids = input_ids, batch_idx = batch_idx, outputs = outputs + ) + + def refresh_inference_from_base(self): + if self.base_model is None or self.peft_model is None: + return 0 + return refresh_lora_merge_from_pristine(self.base_model, self.peft_model) + + @torch.inference_mode() + def generate(self, sequences, capture_cudagraph = False): + self.tokenize(sequences) + waiting = deque(sequences) + running = deque() + done = [] + + if capture_cudagraph and not self.cudagraph_captured: + self.capture_decode_cudagraph() + self.cudagraph_captured = True + + while waiting or running: + batch = [] + while waiting and self.page_table.can_reserve(waiting[0].total_length): + seq = waiting.popleft() + bi = self.page_table.allocate() + self.page_table.reserve( + bi, + torch.tensor([bi], device = self.device, dtype = torch.long), + seq.total_length, + ) + seq.batch_idx = bi + batch.append(seq) + if batch: + logits = self._prefill(batch) + next_ids = torch.argmax(logits, dim = -1).tolist() + for i, seq in enumerate(batch): + seq.last_token_id = next_ids[i] + seq.output_ids.append(next_ids[i]) + if ( + seq.last_token_id == self.eos_token_id + or len(seq.output_ids) >= seq.max_new_tokens + ): + seq.finished = True + done.append(seq) + self.page_table.erase(seq.batch_idx) + else: + running.append(seq) + continue + + decode_batch = [] + while running: + seq = running.popleft() + if self.page_table.capacity[seq.batch_idx] >= seq.total_length: + decode_batch.append(seq) + elif self.page_table.can_reserve( + seq.total_length, batch_idx_int = seq.batch_idx + ): + self.page_table.reserve( + seq.batch_idx, + torch.tensor( + [seq.batch_idx], + device = self.device, + dtype = torch.long, + ), + seq.total_length, + ) + decode_batch.append(seq) + else: + running.appendleft(seq) + newest = running.pop() + waiting.appendleft(newest) + self.page_table.erase(newest.batch_idx) + if not decode_batch: + continue + + B = len(decode_batch) + bi_tensor = torch.tensor( + [s.batch_idx for s in decode_batch], + dtype = torch.long, + device = self.device, + ) + last_ids = torch.tensor( + [s.last_token_id for s in decode_batch], + dtype = torch.long, + device = self.device, + ) + cur_pos = torch.tensor( + [s.total_length - 1 for s in decode_batch], + dtype = torch.int32, + device = self.device, + ) + logits = self._decode_step(bi_tensor, last_ids, cur_pos) + next_ids = torch.argmax(logits, dim = -1).tolist() + for i, seq in enumerate(decode_batch): + seq.last_token_id = next_ids[i] + seq.output_ids.append(next_ids[i]) + if ( + seq.last_token_id == self.eos_token_id + or len(seq.output_ids) >= seq.max_new_tokens + ): + seq.finished = True + done.append(seq) + self.page_table.erase(seq.batch_idx) + else: + running.append(seq) + + return done + + +# --- CLI ------------------------------------------------------------------- + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--model_name", default = "unsloth/gemma-4-E2B-it") + p.add_argument("--n_prompts", type = int, default = 64) + p.add_argument("--n_rounds", type = int, default = 2) + p.add_argument("--max_new_tokens", type = int, default = 512) + p.add_argument("--max_batch_size", type = int, default = 64) + p.add_argument("--max_seq_length", type = int, default = 2048) + p.add_argument("--n_pages", type = int, default = 2048) + p.add_argument("--page_size", type = int, default = 128) + p.add_argument("--capture_cudagraph", action = "store_true") + p.add_argument("--lora_adapter", default = None) + p.add_argument("--decode_kernel_options", default = None) + p.add_argument("--prefill_kernel_options", default = None) + p.add_argument( + "--fa4_prefill", + default = None, + action = argparse.BooleanOptionalAction, + help = ( + "Use BLOCK_SIZE=(256,128) + BACKEND=FLASH on prefill. Default " + "auto-enables on Hopper (sm_90) and Blackwell (sm_100, sm_120)." + ), + ) + p.add_argument("--load_in_4bit", action = "store_true") + p.add_argument( + "--no_merge_lora", + action = "store_true", + help = "Keep the LoRA adapter as a PEFT wrapper instead of merging.", + ) + p.add_argument( + "--verify_no_drift", + action = "store_true", + help = "Drift-verify the double-copy LoRA refresh across N cycles.", + ) + p.add_argument("--verify_iterations", type = int, default = 10) + p.add_argument("--model_name_4bit", default = None) + p.add_argument("--stats_path", required = True) + p.add_argument( + "--chat_template", + choices = ["auto", "grpo", "native"], + default = "auto", + help = ( + "Which chat template to use. `auto`: native for Gemma-4. " + "`grpo`: force the GRPO template. `native`: force the " + "tokenizer's built-in template." + ), + ) + args = p.parse_args() + + def _parse_opts(s): + if s is None: + return None + return json.loads(s) + + Gemma4ForCausalLM, Gemma4Config, Gemma4TextConfig = _require_gemma4() + + from transformers import AutoTokenizer + from transformers.models.gemma4.modeling_gemma4 import ( + Gemma4ForConditionalGeneration, + ) + + tok = AutoTokenizer.from_pretrained(args.model_name) + if tok.pad_token is None: + tok.pad_token = tok.eos_token + + base_model = None + peft_model = None + + if args.load_in_4bit: + from transformers import AutoModelForCausalLM + from huggingface_hub import HfApi + + bnb_model_name = args.model_name_4bit or f"{args.model_name}-unsloth-bnb-4bit" + # Probe for the 4-bit shard. If missing, the user asked for a + # quant row we cannot produce; fail loudly rather than silently + # falling back to bf16 (which would mislabel the stats file). + try: + HfApi().model_info(bnb_model_name) + except Exception as e: + raise SystemExit( + f"[flex-gemma4] --load_in_4bit: 4-bit shard {bnb_model_name} " + f"is not available ({e}). Use --model_name_4bit to override " + f"or drop --load_in_4bit for bf16." + ) + print(f"[flex-gemma4] loading 4-bit base: {bnb_model_name}") + # AutoModelForCausalLM resolves to `Gemma4ForConditionalGeneration` + # for Gemma-4 -- mirror the bf16 path and move the + # language_model into a ForCausalLM shell so downstream code can + # reach `model.model.layers` / `model.model.embed_tokens`. + full_model = AutoModelForCausalLM.from_pretrained( + bnb_model_name, + attn_implementation = "eager", + device_map = "cuda:0", + ) + if hasattr(full_model.model, "language_model"): + lang_model = full_model.model.language_model + full_model.model.vision_tower = None + full_model.model.audio_tower = None + full_model.model.embed_vision = None + full_model.model.embed_audio = None + text_cfg = full_model.config.text_config + model = Gemma4ForCausalLM(text_cfg) + model.model = lang_model + model.lm_head.weight = lang_model.embed_tokens.weight + else: + model = full_model + if getattr(model.config, "tie_word_embeddings", False): + model.lm_head.weight = model.model.embed_tokens.weight + model.eval() + del full_model + if args.lora_adapter: + from peft import PeftModel + + peft_wrapper = PeftModel.from_pretrained( + model, + str(Path(args.lora_adapter).resolve()), + is_trainable = False, + ) + model = peft_wrapper.base_model.model + else: + # Text-only bf16. The checkpoint stores text weights under the + # `model.language_model.` prefix (Gemma-4 is natively multimodal). + # We load the full `Gemma4ForConditionalGeneration`, pluck the + # language_model, then drop the vision + audio towers before + # moving to GPU so peak memory reflects text only. + full_cfg = Gemma4Config.from_pretrained(args.model_name) + text_cfg = full_cfg.text_config + + full_model = Gemma4ForConditionalGeneration.from_pretrained( + args.model_name, + dtype = torch.bfloat16, + attn_implementation = "eager", + ) + lang_model = full_model.model.language_model + # Drop the non-text towers. `embed_vision` / `embed_audio` project + # from text hidden size -- harmless when not invoked, but we kill + # them too so deepcopy (below) stays cheap. + full_model.model.vision_tower = None + full_model.model.audio_tower = None + full_model.model.embed_vision = None + full_model.model.embed_audio = None + + # Build a ForCausalLM shell around the language_model so LoRA / + # state-dict hashing treat it like any other HF decoder model. + base_model = Gemma4ForCausalLM(text_cfg) + base_model.model = lang_model + base_model.lm_head.weight = lang_model.embed_tokens.weight + base_model = base_model.to(torch.bfloat16).to("cuda") + base_model.eval() + del full_model + + if not args.lora_adapter: + model = base_model + base_model = None + elif args.no_merge_lora: + from peft import PeftModel + + peft_wrapper = PeftModel.from_pretrained( + base_model, + str(Path(args.lora_adapter).resolve()), + is_trainable = False, + ) + model = peft_wrapper.base_model.model + base_model = None + else: + from peft import PeftModel + + print("[flex-gemma4] deep-copying base model for double-copy rollout") + inference_model = copy.deepcopy(base_model) + inference_model.eval() + peft_model = PeftModel.from_pretrained( + inference_model, + str(Path(args.lora_adapter).resolve()), + is_trainable = False, + ) + model = peft_model.base_model.model + model.eval() + + if args.verify_no_drift: + if args.load_in_4bit: + raise SystemExit( + "--verify_no_drift only applies to the bf16 double-copy path." + ) + if args.no_merge_lora: + raise SystemExit( + "--verify_no_drift is incompatible with --no_merge_lora." + ) + if base_model is None or peft_model is None: + raise SystemExit( + "--verify_no_drift requires --lora_adapter against the bf16 path." + ) + print( + f"[flex-gemma4] running drift verification: " + f"{args.verify_iterations} perturb+refresh cycles" + ) + result = run_drift_verification( + base_model, peft_model, n_iters = args.verify_iterations + ) + result = {"mode": "verify_no_drift", "model_name": args.model_name, **result} + os.makedirs( + os.path.dirname(os.path.abspath(args.stats_path)) or ".", + exist_ok = True, + ) + with open(args.stats_path, "w") as f: + json.dump(result, f, indent = 2) + print(json.dumps(result, indent = 2)) + os._exit(0) + + from unsloth_grpo_common import SYSTEM_PROMPT, apply_chat_template_to_tokenizer + from datasets import load_dataset + + if args.chat_template == "auto": + # Gemma-4 default is the tokenizer native template; only Qwen3 + # in this repo uses GRPO-by-default. + use_grpo = False + elif args.chat_template == "grpo": + use_grpo = True + else: + use_grpo = False + if use_grpo: + apply_chat_template_to_tokenizer(tok) + print("[flex-gemma4] chat_template: GRPO") + else: + print("[flex-gemma4] chat_template: tokenizer native") + ds = load_dataset("open-r1/DAPO-Math-17k-Processed", "en", split = "train") + ds = ds.shuffle(seed = 3407).select(range(args.n_prompts)) + messages = [ + [ + {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "user", "content": x["prompt"]}, + ] + for x in ds + ] + texts = [ + tok.apply_chat_template(m, add_generation_prompt = True, tokenize = False) + for m in messages + ] + + inference = FlexGemma4Inference( + model, + tok, + max_batch_size = args.max_batch_size, + max_seq_length = args.max_seq_length, + n_pages = args.n_pages, + page_size = args.page_size, + max_new_tokens = args.max_new_tokens, + decode_kernel_options = _parse_opts(args.decode_kernel_options), + prefill_kernel_options = _parse_opts(args.prefill_kernel_options), + fa4_prefill = args.fa4_prefill, + base_model = base_model, + peft_model = peft_model, + ) + + if inference.base_model is not None and inference.peft_model is not None: + n = inference.refresh_inference_from_base() + print(f"[flex-gemma4] double-copy rollout: refreshed {n} LoRA-target layers") + + def make_seqs(): + return [Sequence(text = t, max_new_tokens = args.max_new_tokens) for t in texts] + + torch.cuda.reset_peak_memory_stats() + print("[flex-gemma4] warmup (16 prompts)...") + _ = inference.generate(make_seqs()[:16], capture_cudagraph = args.capture_cudagraph) + torch.cuda.synchronize() + + wall_times = [] + total_decoded = 0 + for r in range(args.n_rounds): + torch.cuda.synchronize() + t0 = time.perf_counter() + out = inference.generate(make_seqs()) + torch.cuda.synchronize() + wall_times.append(time.perf_counter() - t0) + total_decoded = sum(len(s.output_ids) for s in out) + print( + f"[flex-gemma4] round {r}: {wall_times[-1]:.2f}s, {total_decoded} " + f"tokens, {total_decoded / wall_times[-1]:.1f} tok/s" + ) + + med = sorted(wall_times)[len(wall_times) // 2] + best = min(wall_times) + peak = torch.cuda.max_memory_allocated() / 1024**3 + sample_completions = [] + for s in out[:3]: + sample_completions.append( + tok.decode(s.output_ids[:80], skip_special_tokens = True) + ) + res = { + "backend": "flex-gemma4", + "model_name": args.model_name, + "capture_cudagraph": args.capture_cudagraph, + "lora_adapter": args.lora_adapter, + "n_prompts": args.n_prompts, + "n_decoded_tokens": total_decoded, + "wall_times_s": wall_times, + "median_wall_s": med, + "best_wall_s": best, + "decode_tps_median": total_decoded / med if med else 0, + "decode_tps_best": total_decoded / best if best else 0, + "max_new_tokens": args.max_new_tokens, + "peak_memory_gb": peak, + "sample_completions": sample_completions, + } + os.makedirs( + os.path.dirname(os.path.abspath(args.stats_path)) or ".", exist_ok = True + ) + with open(args.stats_path, "w") as f: + json.dump(res, f, indent = 2) + print(json.dumps(res, indent = 2)) + os._exit(0) + + +if __name__ == "__main__": + main() diff --git a/unsloth/inference/flex_paged_attention.py b/unsloth/inference/flex_paged_attention.py new file mode 100644 index 0000000000..e76716ab46 --- /dev/null +++ b/unsloth/inference/flex_paged_attention.py @@ -0,0 +1,502 @@ +# Adapted from attention-gym +# Original source: https://github.com/pytorch-labs/attention-gym +# License: BSD 3-Clause (see THIRD_PARTY_LICENSES.md) +# Copyright (c) 2023, Driss Guessous + +# the original implementation has some bugs and has some feature that lives outside of the PageTable class + +from typing import Optional +import torch +from torch import Tensor +from torch.nn.attention.flex_attention import ( + _identity, + _mask_mod_signature, + _score_mod_signature, + BlockMask, + noop_mask, + create_block_mask, +) + +create_block_mask = torch.compile(create_block_mask, dynamic = True) + + +def _cdiv(x: int | float | torch.Tensor, multiple: int | float | torch.Tensor): + return (x + multiple - 1) // multiple + + +class PagedKVCache(torch.nn.Module): + def __init__(self, page_table, n_heads, head_dim, dtype): + super().__init__() + cache_shape = (1, n_heads, page_table.n_pages * page_table.page_size, head_dim) + self.register_buffer("k_cache", torch.zeros(cache_shape, dtype = dtype)) + self.register_buffer("v_cache", torch.zeros(cache_shape, dtype = dtype)) + + self.page_table = page_table + + def update(self, input_pos, k_val, v_val, batch_idx = None): + assert ( + batch_idx is not None + ), "batch_idx is required for paged kv cache, are you using non-paged attention?" + + if batch_idx.ndim == 1: + # batch_idx should be [B] (decode) + return self.page_table.assign( + batch_idx, input_pos, k_val, v_val, self.k_cache, self.v_cache + ) + else: + assert batch_idx.ndim == 2, "batch_idx must be 1D or 2D" + # batch_idx should be [1, L] (batch prefill) + return self.page_table.assign_prefill_no_paging( + batch_idx, input_pos, k_val, v_val, self.k_cache, self.v_cache + ) + + +class PageTable: + """ + PageTable is a modified version of PagedAttention from attention-gym. + + PageTable improves it by: + - maintaining a cpu copy of the page table, to avoid device-to-host transfers + - support batch prefill + - fix the bug in the original code in mask_mod and score_mod by mapping physical batch index to logical batch index + - subsuming the free_batch_idx into the page table, so we don't need to maintain it separately + """ + + def __init__( + self, + n_pages: int, + page_size: int, + max_batch_size: int, + device: str = "cuda", + ): + self.n_pages = n_pages + self.page_size = page_size + self.max_batch_size = max_batch_size + self.device = device + + # page table: [logical_batch_idx, logical_block_idx] -> physical_page_idx + self.page_table = -torch.ones( + (max_batch_size, self.n_pages), dtype = torch.int64, device = device + ) + self.page_table[0, :] = ( + 0 # page 0 is reserved for simpler code in assign_prefill_no_paging + ) + self.page_table_cpu = [[] for _ in range(max_batch_size)] + + self.capacity = [ + 0 for _ in range(max_batch_size) + ] # capacity: batch_idx -> number of pages allocated * page size + self.free_pages = list( + reversed(range(1, n_pages)) + ) # page 0 is reserved for simpler code in assign_prefill_no_paging + self.free_batch_idx = list( + reversed(range(1, max_batch_size)) + ) # batch_idx 0 is reserved for no-op + + # [logical_batch_idx, physical_page_idx] -> logical_page_idx + self.physical_to_logical = -torch.ones( + (max_batch_size, n_pages), dtype = torch.int64, device = device + ) + + def can_reserve(self, size: int, batch_idx_int: int | None = None) -> bool: + """check if we can reserve new pages for an existing request or a new request, without gpu operations""" + if batch_idx_int is None: + # check if we can schedule a new request + return ( + self.pages_available * self.page_size >= size + and len(self.free_batch_idx) > 0 + ) + else: + # check if we can reserve new pages for an existing request + return self.reserve(batch_idx_int, None, size, dry_run = True) + + def allocate(self) -> int: + """allocate a new batch""" + batch_idx = self.free_batch_idx.pop() + + self.capacity[batch_idx] = 0 + self.physical_to_logical[batch_idx, :] = -1 + self.page_table[batch_idx, :] = -1 + return batch_idx + + @property + def pages_available(self) -> int: + return len(self.free_pages) + + def reserve( + self, + batch_idx_int: int, + batch_idx: torch.Tensor, + seq_len: int, + dry_run: bool = False, + ) -> bool: + """ + Requests the capacity of a given batch to be at least enough to + hold `seq_len` elements. + + Args: + batch_idx_int (int): batch index to be reserved; + batch_idx (Tensor): batch index to be reserved; shape :math:`(1)`. + seq_len (Tensor): minimum capacity for the given batch; shape :math:`(1)`. + + Returns: + bool: True if the reservation was successful, False if the reservation was not successful (no space, and in this case, no update is done) + """ + + if seq_len <= self.capacity[batch_idx_int]: + return True + + num_pages_to_allocate = _cdiv( + seq_len - self.capacity[batch_idx_int], self.page_size + ) + + can_allocate = num_pages_to_allocate <= self.pages_available + if dry_run: + return can_allocate + + if not can_allocate: + raise RuntimeError( + f"Cannot reserve {num_pages_to_allocate} pages for a sequence of length {seq_len} " + f"in batch {batch_idx_int}. Only {self.pages_available} pages available. " + f"Current capacity is {self.capacity[batch_idx_int]} tokens." + ) + + start_page_idx = self.capacity[batch_idx_int] // self.page_size + end_page_idx = start_page_idx + num_pages_to_allocate + + # find empty physical pages + allocated_pages_list = self.free_pages[-num_pages_to_allocate:] + allocated_pages = torch.tensor(allocated_pages_list, device = self.device) + # update page table + self.page_table[batch_idx, start_page_idx:end_page_idx] = allocated_pages + + # update metadata + self.physical_to_logical[batch_idx, allocated_pages] = torch.arange( + start_page_idx, + end_page_idx, + device = self.device, + ) + # update cpu side metadata + self.page_table_cpu[batch_idx_int] += allocated_pages_list + self.free_pages = self.free_pages[:-num_pages_to_allocate] + self.capacity[batch_idx_int] += num_pages_to_allocate * self.page_size + return True + + def erase(self, batch_idx: int) -> None: + """ + Removes a single batch from paged attention. + + Args: + batch_idx (int): batch index to be removed; + """ + # NOTE: the GPU side data will only be reset/overwritten when we allocate it for a new batch + self.free_batch_idx.append(batch_idx) + allocated_pages_cpu = self.page_table_cpu[batch_idx] + self.free_pages.extend(reversed(allocated_pages_cpu)) + self.page_table_cpu[batch_idx] = [] + + def assign( + self, + batch_idx: torch.Tensor, + input_pos: torch.Tensor, + k_val: torch.Tensor, + v_val: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + ) -> None: + """ + Assigns new contents `val` to the storage `cache` at the location + `batch_idx` and `input_pos`. + + Args: + batch_idx (Tensor): batch index; shape :math:`(B)`. + input_pos (Tensor): input positions to be assigned for the given batch; shape :math:`(B, S)`. + val (Tensor): value to be assigned; shape :math:`(B, H, S, D)` + cache (Tensor): the cache to store the values; shape:`(1, H, MAX_S, D)` + """ + if k_val.requires_grad: + raise RuntimeError("val must not require gradient") + + B, H, S, K_D = k_val.shape + _, H_cache, MAX_S, D_cache = k_cache.shape + assert H_cache == H, "number of heads must match" + assert MAX_S >= S, "cache must have enough space" + assert D_cache == K_D, "hidden dim must match" + assert input_pos.shape == (B, S), "input_pos must have the same shape as val" + assert batch_idx.shape == (B,), "batch_idx must have one dimension only" + + V_D = v_val.shape[3] + if B != batch_idx.shape[0]: + raise RuntimeError( + f"Expect val and batch_idx have the same batch size but got B={B} and B={batch_idx.shape[0]}." + ) + if H != k_cache.shape[1]: + raise RuntimeError( + f"Expect val and cache has the same number of heads but got H={H} and H={k_cache.shape[1]}." + ) + if S != input_pos.shape[1]: + raise RuntimeError( + f"Expect val and input_pos has the same length but got S={S} and S={input_pos.shape[0]}." + ) + if K_D != k_cache.shape[3]: + raise RuntimeError( + f"Expect k_val and k_cache has the same hidden dim but got D={K_D} and D={k_cache.shape[3]}." + ) + if V_D != v_cache.shape[3]: + raise RuntimeError( + f"Expect v_val and v_cache has the same hidden dim but got D={V_D} and D={v_cache.shape[3]}." + ) + + # find address + logical_block_idx = input_pos // self.page_size # [B, S] + logical_block_offset = input_pos % self.page_size # [B, S] + + # NOTE: this code path is only used for decoding. For batch prefill, use assign_prefill_no_paging() instead + physical_block_idx = torch.gather( + self.page_table[batch_idx], 1, logical_block_idx.to(torch.int64) + ).to(torch.int32) # [B, S] + + addr = (physical_block_idx * self.page_size + logical_block_offset).view( + -1 + ) # [B*S] + + k_val = k_val.permute(1, 0, 2, 3).contiguous().view(1, H, B * S, K_D) + v_val = v_val.permute(1, 0, 2, 3).contiguous().view(1, H, B * S, V_D) + + k_cache[:, :, addr, :] = k_val + v_cache[:, :, addr, :] = v_val + + return k_cache, v_cache + + def convert_logical_block_mask( + self, + block_mask: BlockMask, + batch_idx: Optional[torch.Tensor] = None, + ) -> BlockMask: + """ + Converts a logical block mask by mapping its logical kv indices to the corresponding + physical kv indices. + + Args: + block_mask (BlockMask): logical block mask; + kv_indices shape :math:`(B, H, ROWS, MAX_BLOCKS_IN_COL)`. + batch_idx (Tensor): batch index corresponding to the block_mask + batch dimension. This provides flexibility to convert a + block mask with smaller batch size than the page table; + shape :math:`(B)`. + """ + B, H, ROWS, MAX_BLOCKS_IN_COL = block_mask.kv_indices.shape + + if block_mask.BLOCK_SIZE[1] != self.page_size: + raise RuntimeError( + f"Expect block_mask has the same column block size as page_sizebut got size={block_mask.BLOCK_SIZE[1]} and size={self.page_size}" + ) + + device = block_mask.kv_num_blocks.device + + if batch_idx is None: + batch_idx = torch.arange(B, device = device) + + assert batch_idx.ndim == 1, "batch_idx must be a 1D tensor" + assert ( + batch_idx.shape[0] == B + ), "batch_idx must have the same shape as block_mask" + assert ( + B <= self.max_batch_size + ), "batch_idx must be less than or equal to max_batch_size" + + page_table = self.page_table[batch_idx] + + def transform(num_blocks, indices): + """ + transform the block mask from [B, H, num_q_blocks, num_logical_kv_blocks] + to [B, H, num_q_blocks, num_physical_kv_blocks] + + kv_num_blocks: [B, H, num_q_blocks] -> unchanged + kv_indices: [B, H, num_q_blocks, num_logical_kv_blocks] -> [B, H, num_q_blocks, num_physical_kv_blocks] + """ + if num_blocks is None: + return None, None + new_kv_num_blocks = num_blocks.clone() + new_kv_indices = torch.zeros( + (B, H, ROWS, self.n_pages), dtype = torch.int32, device = device + ) + new_kv_indices[:, :, :, :MAX_BLOCKS_IN_COL] = ( + torch.gather(page_table, 1, indices.view(B, -1).to(torch.int64)) + .view(block_mask.kv_indices.shape) + .to(torch.int32) + ) + return new_kv_num_blocks, new_kv_indices + + new_kv_num_blocks, new_kv_indices = transform( + block_mask.kv_num_blocks, block_mask.kv_indices + ) + new_full_kv_num_blocks, new_full_kv_indices = transform( + block_mask.full_kv_num_blocks, block_mask.full_kv_indices + ) + + new_mask_mod = self.get_mask_mod(block_mask.mask_mod, batch_idx) + + seq_lengths = (block_mask.seq_lengths[0], self.n_pages * self.page_size) + return BlockMask.from_kv_blocks( + new_kv_num_blocks, + new_kv_indices, + new_full_kv_num_blocks, + new_full_kv_indices, + block_mask.BLOCK_SIZE, + new_mask_mod, + seq_lengths = seq_lengths, + ) + + def get_logical_kv_idx( + self, + physical_batch_idx: torch.Tensor, + physical_kv_idx: torch.Tensor, + batch_idx: torch.Tensor, + ): + logical_batch_idx = batch_idx[physical_batch_idx] + physical_kv_block = physical_kv_idx // self.page_size + physical_kv_offset = physical_kv_idx % self.page_size + logical_block_idx = self.physical_to_logical[ + logical_batch_idx, physical_kv_block + ] + logical_kv_idx = logical_block_idx * self.page_size + physical_kv_offset + is_valid = logical_block_idx >= 0 + safe_logical_kv_idx = logical_kv_idx.clamp(min = 0) + return is_valid, safe_logical_kv_idx + + def get_mask_mod( + self, mask_mod: Optional[_mask_mod_signature], batch_idx: torch.Tensor + ) -> _mask_mod_signature: + """ + Converts a mask_mod based on mapping from the physical block index to the logical + block index. + + Args: + mask_mod (_mask_mod_signature): mask_mod based on the logical block index. + """ + if mask_mod is None: + mask_mod = noop_mask + + def new_mask_mod( + b: torch.Tensor, + h: torch.Tensor, + q_idx: torch.Tensor, + physical_kv_idx: torch.Tensor, + ): + is_valid, safe_logical_kv_idx = self.get_logical_kv_idx( + b, physical_kv_idx, batch_idx + ) + return torch.where( + is_valid, mask_mod(b, h, q_idx, safe_logical_kv_idx), False + ) + + return new_mask_mod + + # NOTE: not used in the current codebase + def get_score_mod( + self, score_mod: Optional[_score_mod_signature], batch_idx: torch.Tensor + ) -> _score_mod_signature: + """ + Converts a score_mod based on mapping from the physical block index to the logical + block index. + + Args: + score_mod (_score_mod_signature): score_mod based on the logical block index. + """ + if score_mod is None: + score_mod = _identity + + def new_score_mod( + score: torch.Tensor, + b: torch.Tensor, + h: torch.Tensor, + q_idx: torch.Tensor, + physical_kv_idx: torch.Tensor, + ): + is_valid, safe_logical_kv_idx = self.get_logical_kv_idx( + b, physical_kv_idx, batch_idx + ) + return torch.where( + is_valid, + score_mod(score, b, h, q_idx, safe_logical_kv_idx), + float("-inf"), + ) + + return new_score_mod + + def create_causal_blockmask(self, B, L): + """A minimal, unoptimized causal block mask creation function""" + + def causal(b, h, q_idx, kv_idx): + return q_idx >= kv_idx + + return create_block_mask( + causal, + B = B, + H = None, + Q_LEN = L, + KV_LEN = L, + BLOCK_SIZE = self.page_size, + device = self.device, + ) + + def create_prefill_blockmask_no_paging( + self, batch_idx: Tensor, BLOCK_SIZE: int = 128 + ): + """ + there's no prefix sharing implemented, batch_idx is the document id, batch_idx is not guaranteed to be sorted + """ + assert batch_idx.ndim == 2, "batch_idx must be a 2D tensor" + assert batch_idx.shape[0] == 1, "batch_idx must have batch size 1" + L = batch_idx.shape[1] + docs = batch_idx.view(-1) + + def document_causal(b, h, q_idx, kv_idx): + causal_mask = q_idx >= kv_idx + document_mask = docs[q_idx] == docs[kv_idx] + return causal_mask & document_mask + + return create_block_mask( + document_causal, B = 1, H = None, Q_LEN = L, KV_LEN = L, BLOCK_SIZE = BLOCK_SIZE + ) + + # we assign prefill to the cache, similar to assign(), except we don't return the k_cache, v_cache, we only return the k_val, v_val + def assign_prefill_no_paging( + self, + batch_idx: torch.Tensor, + input_pos: torch.Tensor, + k_val: torch.Tensor, + v_val: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + ) -> None: + """ + assigns kv and returns the original kv + + batch_idx: [1, L] + input_pos: [1, L] + k_val: [1, H, L, D] + v_val: [1, H, L, D] + k_cache: [1, H, MAX_S, D] + v_cache: [1, H, MAX_S, D] + """ + + assert batch_idx.ndim == 2, "batch_idx must be a 2D tensor" + assert input_pos.ndim == 2, "input_pos must be a 2D tensor" + assert k_val.ndim == 4, "k_val must be a 4D tensor" + assert v_val.ndim == 4, "v_val must be a 4D tensor" + assert k_cache.ndim == 4, "k_cache must be a 4D tensor" + assert v_cache.ndim == 4, "v_cache must be a 4D tensor" + assert batch_idx.shape[0] == 1, "batch_idx must have batch size 1" + + input_pos_block_idx = input_pos // self.page_size + input_pos_offset_in_block = input_pos % self.page_size + physical_kv_idx = ( + self.page_table[batch_idx, input_pos_block_idx] * self.page_size + + input_pos_offset_in_block + ) + k_cache[:, :, physical_kv_idx.view(-1), :] = k_val + v_cache[:, :, physical_kv_idx.view(-1), :] = v_val + + return k_val, v_val diff --git a/unsloth/inference/flex_qwen3_llama.py b/unsloth/inference/flex_qwen3_llama.py new file mode 100644 index 0000000000..83bf3b1e73 --- /dev/null +++ b/unsloth/inference/flex_qwen3_llama.py @@ -0,0 +1,1266 @@ +"""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 + +try: + from .flex_paged_attention import PagedKVCache, PageTable +except ImportError: # script-mode fallback (scripts/benchmarks CLI shim) + 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) + # Unsloth's ``LlamaRotaryEmbedding`` drop-in (swapped in for + # ``Qwen3RotaryEmbedding`` at import time, see + # ``unsloth/models/qwen3.py:445``) returns the full cached cos/sin as + # ``[max_seq, D]`` 2D tensors, expecting the caller to slice. Stock HF + # returns ``[B, S, D]`` already-sliced. Detect which one we got. + _cos, _sin = position_embeddings + if _cos.dim() == 2: + _cos = _cos[position_ids] + _sin = _sin[position_ids] + position_embeddings = (_cos, _sin) + 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: + if seq.input_ids is not None and seq.input_length > 0: + # Pre-tokenized input (e.g. from FlexEngine which accepts + # vLLM-style list[int] / TokensPrompt prompts). Skip. + continue + ids = self.tokenizer(seq.text, return_tensors = "pt")["input_ids"].squeeze(0) + seq.input_ids = ids + seq.input_length = ids.shape[0] + + def _prefill(self, batch: list[Sequence]) -> torch.Tensor: + """Packed prefill: concatenate all sequences into [1, L] with a + document_causal mask. Return logits at the last position of each + sequence as [num_seqs, V]. + """ + input_ids_list = [seq.input_ids.to(self.device) for seq in batch] + input_pos_list = [ + torch.arange(seq.input_length, dtype = torch.long, device = self.device) + for seq in batch + ] + batch_idx_list = [ + torch.full( + (seq.input_length,), seq.batch_idx, dtype = torch.long, device = self.device + ) + for seq in batch + ] + input_ids = torch.cat(input_ids_list).view(1, -1) + input_pos = torch.cat(input_pos_list).view(1, -1) + batch_idx = torch.cat(batch_idx_list).view(1, -1) + + # Pad to multiple of Q block size (flex_attention block alignment). + # For FA4 on Blackwell, Q block = 256 -- otherwise 128. + L = input_ids.shape[1] + q_block = self.prefill_q_block + pad = (q_block - L % q_block) % q_block + if pad > 0: + input_ids = F.pad(input_ids, (0, pad), value = 0) + input_pos = F.pad(input_pos, (0, pad), value = 0) + batch_idx = F.pad(batch_idx, (0, pad), value = 0) + + input_lengths = torch.tensor( + [s.input_length for s in batch], dtype = torch.long, device = self.device + ) + logits_positions = input_lengths.cumsum(dim = 0) - 1 # [num_seqs] + + # If FA4 is on, BLOCK_SIZE is a (Q, KV) tuple. Otherwise scalar. + prefill_block_size = ( + (self.prefill_q_block, self.prefill_kv_block) + if self.fa4_prefill + else self.prefill_q_block + ) + mask = self.page_table.create_prefill_blockmask_no_paging( + batch_idx, BLOCK_SIZE = prefill_block_size + ) + + flex_kwargs = dict( + flex_block_mask = mask, + flex_input_pos = input_pos, + flex_batch_idx = batch_idx, + flex_kernel_options = self.prefill_kernel_options, + ) + position_ids = input_pos # Qwen3 uses 0-based; unlike Gemma2 + hidden = call_model_with_flex_kwargs( + self.model, input_ids, position_ids, flex_kwargs + ) + return self.model.lm_head(hidden[:, logits_positions, :]).squeeze(0) + + def _decode_block_mask(self, batch_idx: torch.Tensor): + """Slice a single-row BlockMask for every seq in the decode batch, + then translate logical→physical pages.""" + block_mask = self.block_mask_logical + input_pos = self.input_pos_buffer[batch_idx] + assert batch_idx.ndim == 1 and input_pos.ndim == 1 + B = batch_idx.shape[0] + input_block_idx = input_pos // block_mask.BLOCK_SIZE[0] + kv_num_blocks = block_mask.kv_num_blocks[batch_idx, :, input_block_idx].view( + B, 1, 1 + ) + kv_indices = block_mask.kv_indices[batch_idx, :, input_block_idx].view( + B, 1, 1, -1 + ) + full_num = full_idx = None + if block_mask.full_kv_num_blocks is not None: + full_num = block_mask.full_kv_num_blocks[ + batch_idx, :, input_block_idx + ].view(B, 1, 1) + full_idx = block_mask.full_kv_indices[batch_idx, :, input_block_idx].view( + B, 1, 1, -1 + ) + + def causal_offset(off): + def offset(b, h, q_idx, kv_idx): + return q_idx + off[b] >= kv_idx + + return offset + + seq_length = (1, block_mask.seq_lengths[1]) + mask = BlockMask.from_kv_blocks( + kv_num_blocks, + kv_indices, + full_num, + full_idx, + BLOCK_SIZE = block_mask.BLOCK_SIZE, + mask_mod = causal_offset(input_pos), + seq_lengths = seq_length, + ) + return mask, input_pos + + def _decode_step_eager(self, batch_idx: torch.Tensor, input_ids: torch.Tensor): + B = input_ids.shape[0] + mask, input_pos = self._decode_block_mask(batch_idx) + mask = self.page_table.convert_logical_block_mask(mask, batch_idx) + position_ids = (input_pos).view(B, 1).to(torch.long) + flex_kwargs = dict( + flex_block_mask = mask, + flex_input_pos = input_pos.view(B, 1).to(torch.long), + flex_batch_idx = batch_idx, + flex_kernel_options = self.decode_kernel_options, + ) + hidden = call_model_with_flex_kwargs( + self.model, input_ids.view(B, 1), position_ids, flex_kwargs + ) + return self.model.lm_head(hidden[:, -1, :]) # [B, V] + + def _decode_step( + self, batch_idx: torch.Tensor, input_ids: torch.Tensor, input_pos: torch.Tensor + ): + self.input_pos_buffer.zero_() + self.input_pos_buffer[batch_idx] = input_pos + if not self.cudagraph_captured: + return self._decode_step_eager(batch_idx, input_ids) + bs = input_ids.size(0) + key = next(x for x in self.graph_bs if x >= bs) + graph = self.graphs[key] + gv = self.graph_vars + # batch_idx=0 is the reserved no-op slot. Zero out the unused part + # of each capture-shape buffer so padded entries don't write into + # real KV pages. + for k, v in gv.items(): + if k != "outputs": + v.zero_() + gv["input_ids"][:bs] = input_ids + gv["batch_idx"][:bs] = batch_idx + graph.replay() + return gv["outputs"][:bs] + + def capture_decode_cudagraph(self): + """Capture one CUDA graph per batch-size bucket. + + Pre-reserves a page for every batch_idx slot so the KV cache writes + during capture hit valid physical addresses. After capture we erase + the batches -- the graph replay reads/writes the same physical + pages regardless of whether the logical batch currently owns them, + because batch_idx 0 is reserved as a padding slot. + """ + max_bs = self.max_batch_size + # Reserve a dummy page for every slot we're going to use during + # capture. Without this, assign() does k_cache[:, :, -1, :] = ... + # and we get an illegal memory access. + reserved_batches = [] + for bi in range(1, max_bs): + try: + allocated = self.page_table.allocate() + self.page_table.reserve( + allocated, + torch.tensor([allocated], device = self.device, dtype = torch.long), + self.page_size, # just one page + ) + reserved_batches.append(allocated) + except Exception: + break + + input_ids = torch.zeros(max_bs, dtype = torch.int64, device = self.device) + batch_idx = torch.arange(max_bs, dtype = torch.int64, device = self.device) + outputs = torch.zeros( + (max_bs, self.model.config.vocab_size), + dtype = self.model.dtype, + device = self.device, + ) + self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16)) + pool = None + for bs in reversed(self.graph_bs): + if bs > max_bs: + continue + print(f"[flex] capturing CUDA graph for bs={bs}") + torch.cuda.synchronize() + _ = self._decode_step_eager(batch_idx[:bs], input_ids[:bs]) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool): + outputs[:bs] = self._decode_step_eager(batch_idx[:bs], input_ids[:bs]) + if pool is None: + pool = graph.pool() + self.graphs[bs] = graph + torch.cuda.synchronize() + # Release the scratch batches; real requests will re-allocate them. + for bi in reserved_batches: + self.page_table.erase(bi) + self.graph_vars = dict( + input_ids = input_ids, batch_idx = batch_idx, outputs = outputs + ) + + def refresh_inference_from_base(self): + """Re-materialize the inference copy's merged LoRA weights from the + pristine `base_model`. Call this once at setup (before CUDA graph + capture) and, in a real GRPO loop, once after every training step + that updates the LoRA adapter. Never call `unmerge_adapter()` -- + we always re-merge from pristine, so no drift accumulates. + + No-op when the double-copy pair wasn't configured (e.g. no LoRA, + or 4-bit naive PEFT-wrapper path). + """ + if self.base_model is None or self.peft_model is None: + return 0 + return refresh_lora_merge_from_pristine(self.base_model, self.peft_model) + + @torch.inference_mode() + def generate(self, sequences: list[Sequence], capture_cudagraph = False): + self.tokenize(sequences) + waiting = deque(sequences) + running = deque() + done = [] + + if capture_cudagraph and not self.cudagraph_captured: + self.capture_decode_cudagraph() + self.cudagraph_captured = True + + while waiting or running: + # 1. Try to schedule new requests into running. + batch = [] + while waiting and self.page_table.can_reserve(waiting[0].total_length): + seq = waiting.popleft() + bi = self.page_table.allocate() + self.page_table.reserve( + bi, + torch.tensor([bi], device = self.device, dtype = torch.long), + seq.total_length, + ) + seq.batch_idx = bi + batch.append(seq) + if batch: + logits = self._prefill(batch) + next_ids = torch.argmax(logits, dim = -1).tolist() + for i, seq in enumerate(batch): + seq.last_token_id = next_ids[i] + seq.output_ids.append(next_ids[i]) + if ( + seq.last_token_id == self.eos_token_id + or len(seq.output_ids) >= seq.max_new_tokens + ): + seq.finished = True + done.append(seq) + self.page_table.erase(seq.batch_idx) + else: + running.append(seq) + continue + + # 2. Reserve pages for running seqs that need more capacity. + decode_batch = [] + while running: + seq = running.popleft() + if self.page_table.capacity[seq.batch_idx] >= seq.total_length: + decode_batch.append(seq) + elif self.page_table.can_reserve( + seq.total_length, batch_idx_int = seq.batch_idx + ): + self.page_table.reserve( + seq.batch_idx, + torch.tensor( + [seq.batch_idx], device = self.device, dtype = torch.long + ), + seq.total_length, + ) + decode_batch.append(seq) + else: + running.appendleft(seq) + newest = running.pop() + waiting.appendleft(newest) + self.page_table.erase(newest.batch_idx) + if not decode_batch: + continue + + B = len(decode_batch) + bi_tensor = torch.tensor( + [s.batch_idx for s in decode_batch], + dtype = torch.long, + device = self.device, + ) + last_ids = torch.tensor( + [s.last_token_id for s in decode_batch], + dtype = torch.long, + device = self.device, + ) + cur_pos = torch.tensor( + [s.total_length - 1 for s in decode_batch], + dtype = torch.int32, + device = self.device, + ) + logits = self._decode_step(bi_tensor, last_ids, cur_pos) + next_ids = torch.argmax(logits, dim = -1).tolist() + for i, seq in enumerate(decode_batch): + seq.last_token_id = next_ids[i] + seq.output_ids.append(next_ids[i]) + if ( + seq.last_token_id == self.eos_token_id + or len(seq.output_ids) >= seq.max_new_tokens + ): + seq.finished = True + done.append(seq) + self.page_table.erase(seq.batch_idx) + else: + running.append(seq) + + return done + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--model_name", default = "unsloth/Qwen3-4B-Base") + p.add_argument("--n_prompts", type = int, default = 32) + p.add_argument("--n_rounds", type = int, default = 2) + p.add_argument("--max_new_tokens", type = int, default = 512) + p.add_argument("--max_batch_size", type = int, default = 64) + p.add_argument("--max_seq_length", type = int, default = 2048) + p.add_argument("--n_pages", type = int, default = 2048) + p.add_argument("--page_size", type = int, default = 128) + p.add_argument("--capture_cudagraph", action = "store_true") + p.add_argument("--lora_adapter", default = None) + # Kernel tuning (optional JSON-valued CLI args so we can sweep quickly): + p.add_argument( + "--decode_kernel_options", + default = None, + help = "JSON for FlexKernelOptions applied in decode, " + 'e.g. \'{"PRESCALE_QK":true,"USE_TMA":true}\'.', + ) + p.add_argument( + "--prefill_kernel_options", default = None, help = "Same but for prefill." + ) + # If set, torch.compile the full attention-stack closure in addition to + # (or instead of) compiling just flex_attention. `reduce-overhead` is + # the interesting mode; it nests with our CUDA graph capture. + p.add_argument( + "--compile_model_forward", + default = None, + choices = [None, "default", "reduce-overhead", "max-autotune-no-cudagraphs"], + ) + p.add_argument( + "--fa4_prefill", + default = None, + action = argparse.BooleanOptionalAction, + help = ( + "Use BLOCK_SIZE=(256,128) + BACKEND=FLASH on prefill to unlock the " + "CuTeDSL FA4 kernel. Default auto-enables on Hopper (sm_90) and " + "Blackwell (sm_100, sm_120); use --no-fa4_prefill to force off." + ), + ) + p.add_argument( + "--load_in_4bit", + action = "store_true", + help = ( + "Load the base model as bitsandbytes 4-bit. When set with " + "--lora_adapter, the LoRA is kept as a PEFT wrapper (no merge) " + "because merging into 4-bit weights is not supported." + ), + ) + p.add_argument( + "--no_merge_lora", + action = "store_true", + help = ( + "Reference path: keep the LoRA adapter as a PEFT wrapper " + "instead of merging it. Runs three matmuls per projection; " + "slow. Useful for the unmerged row in the writeup's comparison " + "table. The default is now the double-copy pattern, which is " + "both merge-speed and drift-free." + ), + ) + p.add_argument( + "--verify_no_drift", + action = "store_true", + help = ( + "Drift-verification mode. Hash the pristine base model params, " + "run N perturb+refresh cycles (simulating N GRPO iterations) " + "on a copy, re-hash, and assert bit-identical. Requires a " + "--lora_adapter; skips rollout generation." + ), + ) + p.add_argument( + "--verify_iterations", + type = int, + default = 10, + help = "Number of perturb+refresh cycles for --verify_no_drift.", + ) + p.add_argument( + "--model_name_4bit", + default = None, + help = ( + "Override the 4-bit shard name. Defaults to " + "`{model_name}-unsloth-bnb-4bit`." + ), + ) + p.add_argument("--stats_path", required = True) + p.add_argument( + "--chat_template", + choices = ["auto", "grpo", "native"], + default = "auto", + help = ( + "Which chat template to use for building prompts. " + "`auto`: GRPO template for Qwen3, tokenizer's native template " + "otherwise. `grpo`: force the GRPO template (matches prior " + "Qwen3 baselines). `native`: force the tokenizer's built-in " + "template (required for Llama-3.2-Instruct)." + ), + ) + args = p.parse_args() + + def _parse_opts(s): + if s is None: + return None + return json.loads(s) + + from transformers import AutoModelForCausalLM, AutoTokenizer + + tok = AutoTokenizer.from_pretrained(args.model_name) + if tok.pad_token is None: + tok.pad_token = tok.eos_token + + base_model = None + peft_model = None + + if args.load_in_4bit: + # Load the pre-quantized Unsloth 4-bit shard. Compute dtype comes + # from the packaged config (bf16 for these shards). + # + # 4-bit keeps the naive PEFT-wrapper path: bnb's `Linear4bit` holds + # packed quantised weights, not regular bf16, so the double-copy + # refresh (in-place copy of `base_layer.weight`) doesn't apply. + # Materializing a full bf16 inference copy via dequant would wipe + # out the memory saving of 4-bit. + bnb_model_name = args.model_name_4bit or f"{args.model_name}-unsloth-bnb-4bit" + print(f"[flex] loading 4-bit base: {bnb_model_name}") + model = AutoModelForCausalLM.from_pretrained( + bnb_model_name, + attn_implementation = "eager", + device_map = "cuda:0", + ) + # See note in cb_vs_vllm_generation.py: tie lm_head to embed_tokens + # for bnb-4bit shards of tied-embedding models. + if getattr(model.config, "tie_word_embeddings", False): + model.lm_head.weight = model.model.embed_tokens.weight + model.eval() + + if args.lora_adapter: + from peft import PeftModel + + peft_wrapper = PeftModel.from_pretrained( + model, + str(Path(args.lora_adapter).resolve()), + is_trainable = False, + ) + # LoRA stays as a wrapper around Params4bit; three matmuls per + # projection. This is the slow reference row in the writeup. + model = peft_wrapper.base_model.model + else: + # bf16 path -- double-copy LoRA rollout. + # + # `base_model` stays pristine; we deep-copy it to `inference_model`, + # wrap the copy with PEFT, and re-materialize the merged LoRA on + # the copy whenever the LoRA weights change. Memory cost: +~8 GB + # for Qwen3-4B bf16 (two copies on GPU) -- well within budget vs + # vLLM's 156 GB. + base_model = AutoModelForCausalLM.from_pretrained( + args.model_name, + dtype = torch.bfloat16, + attn_implementation = "eager", + ).to("cuda") + base_model.eval() + + if not args.lora_adapter: + # No adapter -- use base_model directly, no inference copy. + model = base_model + base_model = None + elif args.no_merge_lora: + # Reference path: PEFT wrapper on the only model copy, + # adapter unmerged. Three matmuls per projection. Kept for + # the comparison row in the writeup. + from peft import PeftModel + + peft_wrapper = PeftModel.from_pretrained( + base_model, + str(Path(args.lora_adapter).resolve()), + is_trainable = False, + ) + model = peft_wrapper.base_model.model + base_model = None + else: + # Double-copy rollout path. + from peft import PeftModel + + print("[flex] deep-copying base model for double-copy LoRA rollout") + inference_model = copy.deepcopy(base_model) + inference_model.eval() + + peft_model = PeftModel.from_pretrained( + inference_model, + str(Path(args.lora_adapter).resolve()), + is_trainable = False, + ) + model = peft_model.base_model.model + model.eval() + # Do NOT call merge_adapter here -- FlexInference.refresh_ + # inference_from_base() below handles the initial merge so the + # same code path runs at setup and on every GRPO refresh. + + # Drift-verification mode: skip rollout generation, just hash-check. + if args.verify_no_drift: + if args.load_in_4bit: + raise SystemExit( + "--verify_no_drift only applies to the bf16 double-copy path " + "(4-bit keeps the naive PEFT-wrapper path, no merge refresh)." + ) + if args.no_merge_lora: + raise SystemExit( + "--verify_no_drift is incompatible with --no_merge_lora " + "(nothing is merged; nothing to drift)." + ) + if base_model is None or peft_model is None: + raise SystemExit( + "--verify_no_drift requires --lora_adapter so there is a " + "LoRA to merge/refresh against the pristine base." + ) + print( + f"[flex] running drift verification: {args.verify_iterations} " + f"perturb+refresh cycles" + ) + result = run_drift_verification( + base_model, peft_model, n_iters = args.verify_iterations + ) + result = {"mode": "verify_no_drift", **result} + os.makedirs( + os.path.dirname(os.path.abspath(args.stats_path)) or ".", + exist_ok = True, + ) + with open(args.stats_path, "w") as f: + json.dump(result, f, indent = 2) + print(json.dumps(result, indent = 2)) + os._exit(0) + + from unsloth_grpo_common import ( + SYSTEM_PROMPT, + apply_chat_template_to_tokenizer, + ) + from datasets import load_dataset + + # Pick which chat template builds the prompts. Qwen3 baselines in + # this repo were recorded against the GRPO template; Llama-3.2-Instruct + # only produces coherent completions with its shipped Instruct + # template. + if args.chat_template == "auto": + use_grpo = type(model).__name__.startswith("Qwen3") + elif args.chat_template == "grpo": + use_grpo = True + else: # "native" + use_grpo = False + if use_grpo: + apply_chat_template_to_tokenizer(tok) + print("[flex] chat_template: GRPO") + else: + print("[flex] chat_template: tokenizer native") + ds = load_dataset("open-r1/DAPO-Math-17k-Processed", "en", split = "train") + ds = ds.shuffle(seed = 3407).select(range(args.n_prompts)) + messages = [ + [ + {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "user", "content": x["prompt"]}, + ] + for x in ds + ] + texts = [ + tok.apply_chat_template(m, add_generation_prompt = True, tokenize = False) + for m in messages + ] + + # Make sure the base HF model the attention layers belong to isn't + # wrapped by PeftModel anymore (we merged); `.model` should be + # Qwen3ForCausalLM or LlamaForCausalLM. + inference = FlexInference( + model, + tok, + max_batch_size = args.max_batch_size, + max_seq_length = args.max_seq_length, + n_pages = args.n_pages, + page_size = args.page_size, + max_new_tokens = args.max_new_tokens, + decode_kernel_options = _parse_opts(args.decode_kernel_options), + prefill_kernel_options = _parse_opts(args.prefill_kernel_options), + fa4_prefill = args.fa4_prefill, + base_model = base_model, + peft_model = peft_model, + ) + + # Initial merge from pristine. Done via `refresh_inference_from_base` + # (not raw `merge_adapter`) so the exact same code path runs at setup + # and at every GRPO refresh -- the CUDA graph capture below sees the + # merged weights already in place. In a real GRPO loop, call + # `inference.refresh_inference_from_base()` after every training step + # that updates the LoRA adapter. We skip per-round refresh in this + # benchmark because the LoRA weights don't change between rounds. + if inference.base_model is not None and inference.peft_model is not None: + n = inference.refresh_inference_from_base() + print(f"[flex] double-copy rollout: refreshed {n} LoRA-target layers") + + # Optionally compile the manual forward walker. This fuses the layer-stack + # ops around flex_attention. Under CUDA graph capture, the compiled + # function gets captured into the same graph. + if args.compile_model_forward: + torch._dynamo.config.cache_size_limit = 256 + print( + f"[flex] torch.compile(call_model_with_flex_kwargs, " + f"mode={args.compile_model_forward!r})" + ) + import sys as _sys + + _this = _sys.modules[__name__] + _this.call_model_with_flex_kwargs = torch.compile( + call_model_with_flex_kwargs, + mode = args.compile_model_forward, + dynamic = True, + fullgraph = False, + ) + + def make_seqs(): + return [Sequence(text = t, max_new_tokens = args.max_new_tokens) for t in texts] + + # Warmup. + torch.cuda.reset_peak_memory_stats() + print("[flex] warmup (16 prompts)...") + _ = inference.generate(make_seqs()[:16], capture_cudagraph = args.capture_cudagraph) + torch.cuda.synchronize() + + wall_times = [] + total_decoded = 0 + for r in range(args.n_rounds): + torch.cuda.synchronize() + t0 = time.perf_counter() + out = inference.generate(make_seqs()) + torch.cuda.synchronize() + wall_times.append(time.perf_counter() - t0) + total_decoded = sum(len(s.output_ids) for s in out) + print( + f"[flex] round {r}: {wall_times[-1]:.2f}s, {total_decoded} tokens, " + f"{total_decoded / wall_times[-1]:.1f} tok/s" + ) + + med = sorted(wall_times)[len(wall_times) // 2] + best = min(wall_times) + peak = torch.cuda.max_memory_allocated() / 1024**3 + # Sample a couple of completions so we can eyeball coherence. + sample_completions = [] + for s in out[:3]: + sample_completions.append( + tok.decode(s.output_ids[:80], skip_special_tokens = True) + ) + res = { + "backend": "flex", + "model_name": args.model_name, + "capture_cudagraph": args.capture_cudagraph, + "lora_adapter": args.lora_adapter, + "n_prompts": args.n_prompts, + "n_decoded_tokens": total_decoded, + "wall_times_s": wall_times, + "median_wall_s": med, + "best_wall_s": best, + "decode_tps_median": total_decoded / med if med else 0, + "decode_tps_best": total_decoded / best if best else 0, + "max_new_tokens": args.max_new_tokens, + "peak_memory_gb": peak, + "sample_completions": sample_completions, + } + os.makedirs(os.path.dirname(os.path.abspath(args.stats_path)) or ".", exist_ok = True) + with open(args.stats_path, "w") as f: + json.dump(res, f, indent = 2) + print(json.dumps(res, indent = 2)) + os._exit(0) + + +if __name__ == "__main__": + main() diff --git a/unsloth/inference/vllm_shim.py b/unsloth/inference/vllm_shim.py new file mode 100644 index 0000000000..16b7025d3b --- /dev/null +++ b/unsloth/inference/vllm_shim.py @@ -0,0 +1,171 @@ +"""vLLM-API surface for the flex inference backend. + +`FlexEngine.generate` / `.chat` return :class:`RequestOutput` objects with the +same attribute shape (`prompt_token_ids`, `outputs[i].token_ids`, +`outputs[i].text`, `outputs[i].logprobs`) that vLLM's `LLM.generate` does, so +TRL's GRPO trainer — which reads `output.prompt_token_ids` and iterates +`output.outputs[i].token_ids` (trl/trainer/grpo_trainer.py:1274-1279) — does +not care which backend produced the result. + +`save_lora` / `load_lora` mirror `unsloth_zoo.vllm_utils.save_lora` / +`load_lora` (vllm_utils.py:2389-2628). The only backend-visible difference is +the `LoRARequest` class: we emit our shim, not vllm's, so the GRPO patch in +`unsloth/models/rl.py:1880` passes our object straight to `FlexEngine.generate`. +""" + +from __future__ import annotations + +import functools +import os +from dataclasses import dataclass, field +from typing import Any, Optional + + +# --------------------------------------------------------------------------- +# vLLM result objects +# --------------------------------------------------------------------------- + + +@dataclass +class CompletionOutput: + """vLLM `CompletionOutput` stand-in. + + TRL reads ``.token_ids`` (list of int) and ``.logprobs`` (optional list of + dict[int, Logprob]); we populate both.""" + + index: int = 0 + text: str = "" + token_ids: list = field(default_factory = list) + cumulative_logprob: Optional[float] = None + logprobs: Optional[list] = None + finish_reason: Optional[str] = "stop" + stop_reason: Optional[str] = None + + +@dataclass +class RequestOutput: + """vLLM `RequestOutput` stand-in. + + TRL reads ``.prompt_token_ids`` (list of int) and iterates + ``.outputs`` (list[CompletionOutput]).""" + + request_id: str = "" + prompt: str = "" + prompt_token_ids: list = field(default_factory = list) + outputs: list = field(default_factory = list) + finished: bool = True + + +# --------------------------------------------------------------------------- +# LoRARequest stand-in +# --------------------------------------------------------------------------- +# +# The vLLM + unsloth_zoo combo imports ``vllm.lora.request.LoRARequest`` at +# call time. When the flex backend is selected we never hit that import; we +# pass our own dataclass that carries the same three attributes the caller +# reads: ``lora_name``, ``lora_int_id``, and ``lora_tensors`` + +# ``lora_config`` (for the in-memory LoRA fast path). + + +@dataclass +class LoRARequest: + lora_name: str = "" + lora_int_id: int = 0 + lora_path: Optional[str] = None + lora_tensors: Optional[dict] = None # dict[str, torch.Tensor] + lora_config: Any = None + + +# --------------------------------------------------------------------------- +# save_lora / load_lora — drop-in for unsloth_zoo.vllm_utils equivalents +# --------------------------------------------------------------------------- + +_LORA_REQUEST_ID: Optional[int] = None + + +def save_lora(model, save_directory, *args, **kwargs): + """Dump the PEFT LoRA tensors (``.lora_A.`` / ``.lora_B.``) into a + PEFT-compatible directory. Mirrors + ``unsloth_zoo.vllm_utils.save_lora`` (vllm_utils.py:2389-2397) byte-for-byte + so existing callers (e.g. TRL's GRPOTrainer patch) are untouched.""" + state_dict = model.state_dict() + dtype = model.get_input_embeddings().weight.dtype + state_dict = { + k: v.to(dtype) + for k, v in state_dict.items() + if ".lora_A." in k or ".lora_B." in k + } + kwargs["state_dict"] = state_dict + model.save_pretrained(save_directory = save_directory, *args, **kwargs) + + +def _get_peft_config(save_directory): + """Late-imported to keep the `peft` dep optional at module-load time.""" + from peft import PeftConfig + + return PeftConfig.from_pretrained(save_directory) + + +def load_lora( + model, + save_directory, + load_tensors: bool = False, + lora_request_id: Optional[int] = None, +): + """Build a :class:`LoRARequest` the flex backend can consume. + + Mirrors ``unsloth_zoo.vllm_utils.load_lora`` (vllm_utils.py:2574-2628): + increments a module-level counter so each request gets a fresh + ``lora_int_id``, writes the PEFT adapter config to ``save_directory`` on + first call (or when ``load_tensors=True``) and captures the current + state-dict LoRA tensors so the engine can merge them without an extra + disk round-trip.""" + global _LORA_REQUEST_ID + if _LORA_REQUEST_ID is None: + _LORA_REQUEST_ID = 1 + if lora_request_id is None: + lora_request_id = _LORA_REQUEST_ID + if not os.path.exists(save_directory) or lora_request_id == 1: + if load_tensors: + model.peft_config["default"].save_pretrained(save_directory) + elif not os.path.exists(save_directory): + raise OSError( + f"Unsloth: LoRA filepath = {save_directory} does not exist!" + ) + + if load_tensors: + peft_config = _get_peft_config(save_directory) + state_dict = model.state_dict() + state_dict = { + k.replace(".default", ""): v + for k, v in state_dict.items() + if ".lora_A." in k or ".lora_B." in k + } + req = LoRARequest( + lora_name = str(lora_request_id), + lora_int_id = lora_request_id, + lora_tensors = state_dict, + lora_config = peft_config, + ) + else: + req = LoRARequest( + lora_name = str(lora_request_id), + lora_int_id = lora_request_id, + lora_path = save_directory, + ) + _LORA_REQUEST_ID += 1 + return req + + +# Partial-applied variants, mirroring ``patch_peft_fast_inference``: +# model.save_lora = functools.partial(save_lora, model) +# model.load_lora = functools.partial(load_lora, model) +# are attached by the caller (unsloth/models/_utils.py:2707-2708). + +__all__ = [ + "CompletionOutput", + "RequestOutput", + "LoRARequest", + "save_lora", + "load_lora", +] diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 34450549c5..18e1902979 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -2672,6 +2672,10 @@ def validate_loftq_config(loftq_config, lora_dropout, bias, init_lora_weights, m def fast_inference_setup(model_name, model_config): fast_inference = True + if os.environ.get("UNSLOTH_FAST_INFERENCE", "0") == "1": + # Flex inference backend (Qwen3 / Llama-3 / Gemma-4-E2B-it). Skip + # the vLLM setup entirely — the flex engine does not use vllm. + return fast_inference, model_name if not is_vLLM_available(): logger.warning_once( "Unsloth: vLLM is not installed! Will use Unsloth inference!" @@ -2701,8 +2705,15 @@ def patch_peft_fast_inference(model): model.fast_generate = model.model.fast_generate model.fast_generate_batches = model.model.fast_generate_batches - # Also saving and loading LoRA - from unsloth_zoo.vllm_utils import save_lora, load_lora + # Pick the right save_lora / load_lora implementation. The flex + # backend (UNSLOTH_FAST_INFERENCE=1) never imports vllm; reading + # from unsloth_zoo.vllm_utils would try to import + # ``vllm.lora.request.LoRARequest`` inside load_lora and fail. + _is_flex = type(vllm_engine).__name__ == "FlexEngine" + if _is_flex: + from unsloth.inference.vllm_shim import save_lora, load_lora + else: + from unsloth_zoo.vllm_utils import save_lora, load_lora model.save_lora = functools.partial(save_lora, model) model.load_lora = functools.partial(load_lora, model) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index d39f2588ef..55edbc59e1 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2488,7 +2488,18 @@ class FastLlamaModel: offload_embedding = False, fast_inference = fast_inference, ) - elif not fast_inference: + elif not fast_inference or ( + os.environ.get("UNSLOTH_FAST_INFERENCE", "0") == "1" + ): + # Two callers share this branch: + # * the standard HF load (``fast_inference=False``) + # * the flex-inference load (``fast_inference=True`` + + # ``UNSLOTH_FAST_INFERENCE=1``) -- we load the HF model the + # same way and then wrap it with a ``FlexEngine`` below. + # ``max_batch_size`` is a FlexEngine-specific kwarg, not an HF + # one; stash it before ``AutoModelForCausalLM.from_pretrained`` + # rejects it as an unexpected argument. + _flex_max_batch_size = kwargs.pop("max_batch_size", 32) model = AutoModelForCausalLM.from_pretrained( model_name, device_map = device_map, @@ -2508,10 +2519,32 @@ class FastLlamaModel: load_in_4bit = load_in_4bit, load_in_8bit = kwargs.get("load_in_8bit", False), offload_embedding = False, - fast_inference = False, + fast_inference = fast_inference, ) - model.fast_generate = make_fast_generate_wrapper(model.generate) - model.fast_generate_batches = None + if fast_inference and os.environ.get("UNSLOTH_FAST_INFERENCE", "0") == "1": + # Snapshot the HF model BEFORE Unsloth post-patching so the + # flex inference copy uses the original transformers + # rotary / QKV layout. A later block (after the tokenizer is + # loaded) constructs the FlexEngine around this copy. + import copy as _copy + + model._unsloth_flex_inference_copy = _copy.deepcopy(model) + model._unsloth_flex_inference_copy.eval() + model._unsloth_needs_flex_engine = dict( + dtype = dtype, + max_seq_length = max_seq_length, + max_lora_rank = max_lora_rank, + max_batch_size = _flex_max_batch_size, + gpu_memory_utilization = gpu_memory_utilization, + ) + # Provide a placeholder so downstream code that checks + # hasattr(model, "fast_generate") doesn't explode; we swap + # it out once the engine is live. + model.fast_generate = make_fast_generate_wrapper(model.generate) + model.fast_generate_batches = None + else: + model.fast_generate = make_fast_generate_wrapper(model.generate) + model.fast_generate_batches = None else: from unsloth_zoo.vllm_utils import ( load_vllm, @@ -2584,6 +2617,32 @@ class FastLlamaModel: model, tokenizer, correct_dtype = dtype ) + # UNSLOTH_FAST_INFERENCE=1 path: build the FlexEngine now that the + # tokenizer is available. The HF model itself is NOT flex-patched -- + # the engine owns its own pre-patched deep-copy (captured above + # before Unsloth's post_patch), so `model.forward` (used by the + # training loop) stays intact. + _flex_args = getattr(model, "_unsloth_needs_flex_engine", None) + if _flex_args is not None: + del model._unsloth_needs_flex_engine + _inference_copy = model._unsloth_flex_inference_copy + del model._unsloth_flex_inference_copy + from unsloth.inference.flex_engine import FlexEngine + + flex_engine = FlexEngine( + hf_model = model, + tokenizer = tokenizer, + inference_model = _inference_copy, + base_model = None, + peft_model = None, + **_flex_args, + ) + model.vllm_engine = flex_engine + model.fast_generate = flex_engine.generate + model.fast_generate_batches = functools.partial( + flex_engine.generate, use_tqdm = False + ) + # Patch up QKV / O and MLP for idx, layer in enumerate(model.model.layers): layer.self_attn.apply_qkv = original_apply_qkv @@ -3290,6 +3349,13 @@ class FastLlamaModel: patch_peft_fast_inference(model) + # Hand the training-side PEFT wrapper to the flex engine so that + # state_dict() reads the current LoRA weights and the inference + # copy can mirror them. + _engine = getattr(model, "vllm_engine", None) + if _engine is not None and hasattr(_engine, "bind_peft_model"): + _engine.bind_peft_model(model) + # Add for_inference and for_training model.for_training = functools.partial(FastLlamaModel.for_training, model) model.for_inference = functools.partial(FastLlamaModel.for_inference, model) diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index fc91178d88..ebd3ae8f9e 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -353,7 +353,15 @@ class FastLanguageModel(FastLlamaModel): or dtype == torch.float32 ) - if fast_inference: + _use_flex_fast_inference = ( + os.environ.get("UNSLOTH_FAST_INFERENCE", "0") == "1" + ) + if fast_inference and _use_flex_fast_inference: + # Flex backend path: skip the vLLM import gate entirely. The + # actual engine is attached further down in + # ``FastLlamaModel.from_pretrained``. + pass + elif fast_inference: if importlib.util.find_spec("vllm") is None: raise ImportError( "Unsloth: Please install vLLM before enabling `fast_inference`!\n" @@ -979,7 +987,14 @@ class FastModel(FastBaseModel): ) load_in_4bit = False - if fast_inference: + _use_flex_fast_inference = ( + os.environ.get("UNSLOTH_FAST_INFERENCE", "0") == "1" + ) + if fast_inference and _use_flex_fast_inference: + # Flex backend path: skip the vLLM import gate. The engine is + # attached further down in the ``FastBaseModel`` loader. + pass + elif fast_inference: if importlib.util.find_spec("vllm") is None: raise ImportError( "Unsloth: Please install vLLM before enabling `fast_inference`!\n" diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index df371e00c8..0e44cca711 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -606,7 +606,16 @@ class FastBaseModel: vllm_enable_lora = True if is_vlm and fast_inference: - if not any(arch in VLLM_SUPPORTED_VLM for arch in model_types): + # The UNSLOTH_FAST_INFERENCE=1 flex backend ships with text-only + # support for Gemma-4-E2B-it today. Allow gemma4 through the + # vision-model path when the flex backend is selected; the + # vLLM-only compat list below still gates the default path. + _use_flex = os.environ.get("UNSLOTH_FAST_INFERENCE", "0") == "1" + _flex_allowed = {"gemma4"} + if not ( + any(arch in VLLM_SUPPORTED_VLM for arch in model_types) + or (_use_flex and any(arch in _flex_allowed for arch in model_types)) + ): raise RuntimeError( f"Unsloth: Fast inference is only supported for Language models and Qwen2.5-VL, Gemma3 among vision models. " f"Found architectures: {', '.join(model_types)}!" @@ -912,9 +921,15 @@ class FastBaseModel: verify_fp8_support_if_applicable(model_config) raise_handler = RaiseUninitialized() - if not fast_inference: + _use_flex_fast_inference = ( + fast_inference and os.environ.get("UNSLOTH_FAST_INFERENCE", "0") == "1" + ) + if (not fast_inference) or _use_flex_fast_inference: + # Shared by the standard HF load and the flex-inference load. # Prevent load_in_fp8 from being forwarded into HF internal model loading load_in_fp8 = kwargs.pop("load_in_fp8", None) + # ``max_batch_size`` is a FlexEngine kwarg, not an HF one. + _flex_max_batch_size = kwargs.pop("max_batch_size", 32) # Transformers 5.x @strict config classes reject unexpected kwargs. # Move config-level attributes onto the config object directly. _num_labels = kwargs.pop("num_labels", None) @@ -944,6 +959,16 @@ class FastBaseModel: fast_inference = fast_inference, ) if hasattr(model, "generate"): + if _use_flex_fast_inference: + # Defer engine construction until tokenizer/processor is + # loaded; see the wiring block further down. + model._unsloth_needs_flex_engine = dict( + dtype = dtype, + max_seq_length = max_seq_length, + max_lora_rank = max_lora_rank, + max_batch_size = _flex_max_batch_size, + gpu_memory_utilization = gpu_memory_utilization, + ) model.fast_generate = make_fast_generate_wrapper(model.generate) model.fast_generate_batches = error_out_no_vllm if offload_embedding: @@ -1223,6 +1248,22 @@ class FastBaseModel: raise _patch_err model = post_patch_loss_function(model) + # UNSLOTH_FAST_INFERENCE=1 path: build the FlexEngine here so that + # the tokenizer/processor is available. Keeps the training model's + # forward intact — the engine carries its own deep-copy. + _flex_args = getattr(model, "_unsloth_needs_flex_engine", None) + if _flex_args is not None: + del model._unsloth_needs_flex_engine + from unsloth.inference.flex_engine import load_flex + + _tok_for_flex = getattr(tokenizer, "tokenizer", tokenizer) + flex_engine = load_flex(model, _tok_for_flex, **_flex_args) + model.vllm_engine = flex_engine + model.fast_generate = flex_engine.generate + model.fast_generate_batches = functools.partial( + flex_engine.generate, use_tqdm = False + ) + # Log Unsloth version for future fastpaths for inference if hasattr(model, "config"): model.config.update({"unsloth_version": __version__}) @@ -1512,6 +1553,10 @@ class FastBaseModel: patch_saving_functions(model, vision = True) patch_peft_fast_inference(model) + _engine = getattr(model, "vllm_engine", None) + if _engine is not None and hasattr(_engine, "bind_peft_model"): + _engine.bind_peft_model(model) + # Add for_inference and for_training model.for_training = functools.partial(FastBaseModel.for_training, model) model.for_inference = functools.partial(FastBaseModel.for_inference, model)