diff --git a/tests/flex_fastlm_bench.py b/tests/flex_fastlm_bench.py new file mode 100644 index 0000000000..96311d5be0 --- /dev/null +++ b/tests/flex_fastlm_bench.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""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..4c8ce8d1ea --- /dev/null +++ b/tests/flex_fastlm_smoke.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""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/tests/flex_gemma4_moe_merge_parity.py b/tests/flex_gemma4_moe_merge_parity.py new file mode 100644 index 0000000000..5e3d67e7a8 --- /dev/null +++ b/tests/flex_gemma4_moe_merge_parity.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""Numerical parity check for ``refresh_moe_lora_merge_from_pristine`` +at Gemma 4 26B-A4B expert shapes. + +Reuses ``merge_under_test`` and ``reference_merge`` from +``tests/flex_moe_merge_parity.py`` (both are standalone helpers — they +import nothing from the flex engine, only exercise the bitwise kernel +logic lifted from ``unsloth/inference/flex_moe.py:682-800``). + +Shapes covered (Gemma 4 26B-A4B): +- gate_up_proj: (E=128, 2I=1408, H=2816) standard ``(E, out, in)`` +- down_proj: (E=128, H=2816, I=704) standard ``(E, out, in)`` +Rank-16 LoRA, single + dual adapter, bf16 + fp32. + +Gemma 4 MoE uses the same ``F.linear``-oriented expert layout as Qwen3, +so only the ``transposed=False`` branch of +``refresh_moe_lora_merge_from_pristine`` is exercised here. (The +transposed branch is already covered for gpt-oss in +``flex_moe_merge_parity.py``.) + +Usage:: + CUDA_VISIBLE_DEVICES=2 python -u tests/flex_gemma4_moe_merge_parity.py +""" +from __future__ import annotations + +import sys +from pathlib import Path + +import torch + +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + +from tests.flex_moe_merge_parity import ( # noqa: E402 + merge_under_test, + reference_merge, + test_correctness, +) + + +def main(): + device = "cuda" if torch.cuda.is_available() else "cpu" + torch.manual_seed(3407) + print(f"[merge-parity-gemma4] device={device} dtype=bf16 + fp32") + print(f"[merge-parity-gemma4] torch={torch.__version__}") + print() + + print("== Correctness (fp32 golden + bf16 realistic Gemma 4 shapes) ==") + # gate_up_proj: (E=128, 2I=1408, H=2816) standard + # down_proj: (E=128, H=2816, I=704 ) standard + cases = [ + # (E, in_dim, out_dim, R, dtype, transposed, n_adapters) + ( 8, 64, 128, 4, torch.float32, False, 1), + ( 8, 64, 128, 4, torch.float32, False, 2), + (16, 128, 256, 8, torch.float32, False, 2), + # bf16 at Gemma 4 26B-A4B MoE shapes. + (128, 2816, 1408, 16, torch.bfloat16, False, 1), # gate_up_proj + (128, 704, 2816, 16, torch.bfloat16, False, 1), # down_proj + (128, 2816, 1408, 16, torch.bfloat16, False, 2), # two adapters + (128, 2816, 1408, 64, torch.bfloat16, False, 1), # higher rank + ] + all_ok = True + for E, in_dim, out_dim, R, dtype, tr, na in cases: + ok = test_correctness(E, in_dim, out_dim, R, dtype, device, + transposed=tr, n_adapters=na) + all_ok = all_ok and ok + print(f"\n overall: {'PASS' if all_ok else 'FAIL'}") + + +if __name__ == "__main__": + main() diff --git a/tests/flex_gemma4_parity.py b/tests/flex_gemma4_parity.py new file mode 100644 index 0000000000..c892ad85cd --- /dev/null +++ b/tests/flex_gemma4_parity.py @@ -0,0 +1,191 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""Token parity for FlexGemma4Inference / FlexGemma4MoEInference vs pure HF. + +Usage:: + CUDA_VISIBLE_DEVICES=2 UNSLOTH_FAST_INFERENCE=1 python -u \\ + tests/flex_gemma4_parity.py --backend flex --model unsloth/gemma-4-31B-it + + CUDA_VISIBLE_DEVICES=3 python -u tests/flex_gemma4_parity.py \\ + --backend hf --model unsloth/gemma-4-31B-it +""" +from __future__ import annotations + +import argparse +import json +import os +import sys +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)) + + +CHAT_PROMPTS = [ + "In one sentence, what is Paris?", + "What is 23 + 19? Answer in one word.", + "Continue: The quick brown fox jumps over", +] + + +def _run_flex(args, dtype, *, capture: bool, lora_path=None): + import torch + os.environ["UNSLOTH_FAST_INFERENCE"] = "1" + import unsloth # noqa + from unsloth import FastLanguageModel + + if not capture: + # Force the eager decode path across arches. + try: + from unsloth.inference.flex_gemma4 import FlexGemma4Inference + FlexGemma4Inference.capture_decode_cudagraph = lambda self: None + except Exception: + pass + try: + from unsloth.inference.flex_gemma4_moe import FlexGemma4MoEInference + FlexGemma4MoEInference.capture_decode_cudagraph = lambda self: None + except Exception: + pass + + 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=4, + gpu_memory_utilization=0.6, + ) + + if lora_path is not None: + model.load_adapter(lora_path, adapter_name="default") + print(f"[parity-flex] LoRA attached from {lora_path}") + + prompts = [ + tok.apply_chat_template( + [{"role": "user", "content": p}], + tokenize=False, + add_generation_prompt=True, + ) + for p in CHAT_PROMPTS[: args.num_prompts] + ] + + class _SP: + max_tokens = args.max_new_tokens + temperature = 0.0 + + _ = model.fast_generate(prompts, sampling_params=_SP(), use_tqdm=False) + outs = model.fast_generate(prompts, sampling_params=_SP(), use_tqdm=False) + token_ids = [list(o.outputs[0].token_ids) for o in outs] + texts = [o.outputs[0].text for o in outs] + return token_ids, texts, tok + + +def _run_hf(args, dtype, *, lora_path=None): + import torch + from transformers import AutoModelForCausalLM, AutoTokenizer + + tok = AutoTokenizer.from_pretrained(args.model) + + try: + model = AutoModelForCausalLM.from_pretrained( + args.model, dtype=dtype, device_map="cuda", + attn_implementation="eager", + ) + except Exception: + # Multimodal Gemma 4 (ConditionalGeneration) — load the top-level class. + from transformers import AutoModelForImageTextToText + model = AutoModelForImageTextToText.from_pretrained( + args.model, dtype=dtype, device_map="cuda", + attn_implementation="eager", + ) + model.eval() + + if lora_path is not None: + from peft import PeftModel + model = PeftModel.from_pretrained(model, lora_path) + model.eval() + print(f"[parity-hf] LoRA attached from {lora_path}") + + if tok.pad_token_id is None: + tok.pad_token_id = tok.eos_token_id + tok.padding_side = "left" + + prompts = [ + tok.apply_chat_template( + [{"role": "user", "content": p}], + tokenize=False, + add_generation_prompt=True, + ) + for p in CHAT_PROMPTS[: args.num_prompts] + ] + inputs = tok(prompts, return_tensors="pt", padding=True).to("cuda") + out = model.generate( + **inputs, + max_new_tokens=args.max_new_tokens, + do_sample=False, + temperature=1.0, + pad_token_id=tok.pad_token_id, + ) + prompt_len = inputs["input_ids"].shape[1] + eos = tok.eos_token_id + pad = tok.pad_token_id + token_ids = [] + texts = [] + for row in out: + ids = row[prompt_len:].tolist() + while ids and ids[-1] in (eos, pad): + ids.pop() + token_ids.append(ids) + texts.append(tok.decode(ids, skip_special_tokens=False)) + return token_ids, texts, tok + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--model", default="unsloth/gemma-4-31B-it") + p.add_argument("--backend", choices=["flex", "flex_eager", "hf"], required=True) + p.add_argument("--max_new_tokens", type=int, default=64) + p.add_argument("--num_prompts", type=int, default=3) + p.add_argument("--lora_path", default=None) + p.add_argument("--max_seq_length", type=int, default=1024) + p.add_argument("--out_dir", default="async_task_outputs/gemma4_moe_bench") + args = p.parse_args() + + import torch + dtype = torch.bfloat16 + + if args.backend == "flex": + token_ids, texts, _ = _run_flex(args, dtype, capture=True, lora_path=args.lora_path) + elif args.backend == "flex_eager": + token_ids, texts, _ = _run_flex(args, dtype, capture=False, lora_path=args.lora_path) + else: + token_ids, texts, _ = _run_hf(args, dtype, lora_path=args.lora_path) + + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + model_tag = args.model.replace("/", "_") + suffix = "_lora" if args.lora_path else "" + out_path = out_dir / f"parity_{model_tag}_{args.backend}{suffix}.json" + with open(out_path, "w") as f: + json.dump( + { + "backend": args.backend, + "model": args.model, + "prompts": CHAT_PROMPTS[: args.num_prompts], + "token_ids": token_ids, + "texts": texts, + }, + f, + indent=2, + ) + print(f"[parity-{args.backend}] wrote {out_path}") + for i, (pp, tt) in enumerate(zip(CHAT_PROMPTS[: args.num_prompts], texts)): + print(f"[parity-{args.backend}] prompt {i}: {pp!r}") + print(f"[parity-{args.backend}] completion {i}: {tt!r}") + + +if __name__ == "__main__": + main() diff --git a/tests/flex_gpt_oss_parity.py b/tests/flex_gpt_oss_parity.py new file mode 100644 index 0000000000..9b165cc030 --- /dev/null +++ b/tests/flex_gpt_oss_parity.py @@ -0,0 +1,176 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""Token parity for FlexGptOssInference vs pure HF generate on gpt-oss-20b. + +Usage:: + # Flex side (cudagraph off in Phase 1). + CUDA_VISIBLE_DEVICES=2 UNSLOTH_FAST_INFERENCE=1 python -u \\ + tests/flex_gpt_oss_parity.py --backend flex + + # HF reference. + CUDA_VISIBLE_DEVICES=3 python -u tests/flex_gpt_oss_parity.py \\ + --backend hf +""" +from __future__ import annotations + +import argparse +import json +import os +import sys +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)) + + +CHAT_PROMPTS = [ + "In one sentence, what is Paris?", + "What is 23 + 19? Answer in one word.", + "Continue: The quick brown fox jumps over", +] + + +def _run_flex(args, dtype, *, capture: bool, lora_path=None): + import torch + os.environ["UNSLOTH_FAST_INFERENCE"] = "1" + import unsloth # noqa + from unsloth import FastLanguageModel + + if not capture: + # Force the eager decode path by stubbing capture_decode_cudagraph. + from unsloth.inference.flex_gpt_oss import FlexGptOssInference + FlexGptOssInference.capture_decode_cudagraph = lambda self: None + + 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=4, + gpu_memory_utilization=0.6, + ) + + if lora_path is not None: + model.load_adapter(lora_path, adapter_name="default") + print(f"[parity-flex] LoRA attached from {lora_path}") + + prompts = [ + tok.apply_chat_template( + [{"role": "user", "content": p}], + tokenize=False, + add_generation_prompt=True, + ) + for p in CHAT_PROMPTS + ] + + class _SP: + max_tokens = args.max_new_tokens + temperature = 0.0 + + # First call primes; second is the measurement. + _ = model.fast_generate(prompts, sampling_params=_SP(), use_tqdm=False) + outs = model.fast_generate(prompts, sampling_params=_SP(), use_tqdm=False) + token_ids = [list(o.outputs[0].token_ids) for o in outs] + texts = [o.outputs[0].text for o in outs] + return token_ids, texts, tok + + +def _run_hf(args, dtype, *, lora_path=None): + import torch + from transformers import AutoModelForCausalLM, AutoTokenizer + + tok = AutoTokenizer.from_pretrained(args.model) + model = AutoModelForCausalLM.from_pretrained( + args.model, dtype=dtype, device_map="cuda", + attn_implementation="eager", + ) + model.eval() + + if lora_path is not None: + from peft import PeftModel + model = PeftModel.from_pretrained(model, lora_path) + model.eval() + print(f"[parity-hf] LoRA attached from {lora_path}") + + if tok.pad_token_id is None: + tok.pad_token_id = tok.eos_token_id + tok.padding_side = "left" + + prompts = [ + tok.apply_chat_template( + [{"role": "user", "content": p}], + tokenize=False, + add_generation_prompt=True, + ) + for p in CHAT_PROMPTS + ] + inputs = tok(prompts, return_tensors="pt", padding=True).to("cuda") + out = model.generate( + **inputs, + max_new_tokens=args.max_new_tokens, + do_sample=False, + temperature=1.0, + pad_token_id=tok.pad_token_id, + ) + prompt_len = inputs["input_ids"].shape[1] + eos = tok.eos_token_id + pad = tok.pad_token_id + token_ids = [] + texts = [] + for row in out: + ids = row[prompt_len:].tolist() + while ids and ids[-1] in (eos, pad): + ids.pop() + token_ids.append(ids) + texts.append(tok.decode(ids, skip_special_tokens=False)) + return token_ids, texts, tok + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--model", default="unsloth/gpt-oss-20b-BF16") + p.add_argument("--backend", choices=["flex", "flex_eager", "hf"], required=True) + p.add_argument("--max_new_tokens", type=int, default=64) + p.add_argument("--lora_path", default=None, + help="optional LoRA adapter dir to attach before decode") + p.add_argument("--max_seq_length", type=int, default=1024) + p.add_argument("--out_dir", default="async_task_outputs/qwen3_moe_grpo_bench_v2") + args = p.parse_args() + + import torch + dtype = torch.bfloat16 + + if args.backend == "flex": + token_ids, texts, _ = _run_flex(args, dtype, capture=True, lora_path=args.lora_path) + elif args.backend == "flex_eager": + token_ids, texts, _ = _run_flex(args, dtype, capture=False, lora_path=args.lora_path) + else: + token_ids, texts, _ = _run_hf(args, dtype, lora_path=args.lora_path) + + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + suffix = "_lora" if args.lora_path else "" + out_path = out_dir / f"parity_gptoss_{args.backend}{suffix}.json" + with open(out_path, "w") as f: + json.dump( + { + "backend": args.backend, + "model": args.model, + "prompts": CHAT_PROMPTS, + "token_ids": token_ids, + "texts": texts, + }, + f, + indent=2, + ) + print(f"[parity-{args.backend}] wrote {out_path}") + for i, (pp, tt) in enumerate(zip(CHAT_PROMPTS, texts)): + print(f"[parity-{args.backend}] prompt {i}: {pp!r}") + print(f"[parity-{args.backend}] completion {i}: {tt!r}") + + +if __name__ == "__main__": + main() diff --git a/tests/flex_lazy_batch_smoke.py b/tests/flex_lazy_batch_smoke.py new file mode 100644 index 0000000000..87b3fcf917 --- /dev/null +++ b/tests/flex_lazy_batch_smoke.py @@ -0,0 +1,227 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""Smoke tests for the deferred FlexEngine batch-size sizing. + +Unit-level: covers the four cases described in the implementation plan +by monkey-patching :class:`FlexEngine` with a cheap stand-in so the +tests run on any box (no CUDA / no model download). The dispatch logic +lives entirely in :func:`build_flex_engine`, +:func:`install_flex_sentinel`, and :func:`_build_flex_from_args`, which +are the units under test. + +Run as: + python tests/flex_lazy_batch_smoke.py +""" + +from __future__ import annotations + +import sys +import types +import warnings +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)) + + +class _StubFlexEngine: + """Records the ``max_batch_size`` construction arg and nothing else. + + ``_cudagraph_primed`` flips True after the first ``generate`` so the + post-warmup refuse path can be exercised without touching CUDA. + """ + + instances: list = [] + + def __init__( + self, + hf_model, + tokenizer, + *, + dtype = None, + max_seq_length: int = 2048, + max_lora_rank: int = 64, + max_batch_size: int = 32, + page_size: int = 128, + gpu_memory_utilization: float = 0.5, + max_new_tokens: int = 512, + prefill_kernel_options = None, + decode_kernel_options = None, + fa4_prefill = None, + capture_cudagraph: bool = True, + base_model = None, + peft_model = None, + inference_model = None, + ): + self.hf_model = hf_model + self.tokenizer = tokenizer + self.max_batch_size = max_batch_size + self.max_seq_length = max_seq_length + self.compute_dtype = dtype + self._cudagraph_primed = False + self.generate_calls = 0 + _StubFlexEngine.instances.append(self) + + def generate(self, prompts = None, *args, **kwargs): + self.generate_calls += 1 + self._cudagraph_primed = True + return [("stub", prompts)] + + +def _make_stub_model(): + """An object that quacks like an HF model for ``install_flex_sentinel``.""" + + model = types.SimpleNamespace() + model._unsloth_needs_flex_engine = dict( + dtype = "bf16", + max_seq_length = 2048, + max_lora_rank = 64, + max_batch_size = 32, + gpu_memory_utilization = 0.5, + ) + model._unsloth_flex_inference_copy = object() # never dereferenced + return model + + +def _install_stub(): + """Patch FlexEngine with :class:`_StubFlexEngine` for the duration of the + test process. Imports happen lazily inside ``build_flex_engine``, so we + patch the module attribute before those calls fire.""" + + import unsloth.inference.flex_engine as fe + + _StubFlexEngine.instances.clear() + fe.FlexEngine = _StubFlexEngine + + +def _case1_default_floor(): + """No trainer, no kwargs: fast_generate builds at floor=32.""" + + from unsloth.inference.flex_engine import install_flex_sentinel + + _install_stub() + model = _make_stub_model() + install_flex_sentinel(model, tokenizer = object()) + + assert hasattr(model, "vllm_engine"), "sentinel not installed" + assert not hasattr( + model, "_flex_engine_instance" + ), "engine should NOT exist before first use" + + out = model.fast_generate(["hello"]) + assert out == [("stub", ["hello"])] + + engine = model._flex_engine_instance + assert engine.max_batch_size == 32, engine.max_batch_size + # Sentinel was replaced with the real engine after build. + assert model.vllm_engine is engine + print(" [1/4] default path: floor=32 build on first fast_generate OK") + + +def _case2_grpo_bump(): + """User kwarg=16 + GRPO target=64 → engine built at 64, warning logged.""" + + from unsloth.inference.flex_engine import ( + _build_flex_from_args, + install_flex_sentinel, + ) + + _install_stub() + model = _make_stub_model() + model._unsloth_needs_flex_engine["max_batch_size"] = 16 # user floor + install_flex_sentinel(model, tokenizer = object()) + + args = types.SimpleNamespace( + per_device_train_batch_size = 2, + steps_per_generation = 4, + num_generations = 8, + gradient_accumulation_steps = 1, + ) + with warnings.catch_warnings(record = True) as caught: + warnings.simplefilter("always") + _build_flex_from_args(model, args) + + engine = model._flex_engine_instance + assert engine.max_batch_size == 64, engine.max_batch_size + assert any("16 -> 64" in str(w.message) for w in caught), [ + str(w.message) for w in caught + ] + print(" [2/4] GRPO bump: 16 -> 64 with warning OK") + + +def _case3_user_floor_wins(): + """User kwarg=128 + GRPO target=8 → engine stays at 128, no warning.""" + + from unsloth.inference.flex_engine import ( + _build_flex_from_args, + install_flex_sentinel, + ) + + _install_stub() + model = _make_stub_model() + model._unsloth_needs_flex_engine["max_batch_size"] = 128 + install_flex_sentinel(model, tokenizer = object()) + + args = types.SimpleNamespace( + per_device_train_batch_size = 1, + steps_per_generation = 2, + num_generations = 4, + gradient_accumulation_steps = 1, + ) + with warnings.catch_warnings(record = True) as caught: + warnings.simplefilter("always") + _build_flex_from_args(model, args) + + engine = model._flex_engine_instance + assert engine.max_batch_size == 128, engine.max_batch_size + assert not any("FlexEngine" in str(w.message) for w in caught), [ + str(w.message) for w in caught + ] + print(" [3/4] user floor wins: engine.max_batch_size=128 OK") + + +def _case4_post_warmup_refused(): + """fast_generate primes the engine; later GRPO target=64 must raise.""" + + from unsloth.inference.flex_engine import ( + _build_flex_from_args, + install_flex_sentinel, + ) + + _install_stub() + model = _make_stub_model() + install_flex_sentinel(model, tokenizer = object()) + + model.fast_generate(["hi"]) # builds at floor=32, sets _cudagraph_primed + assert model._flex_engine_instance.max_batch_size == 32 + + args = types.SimpleNamespace( + per_device_train_batch_size = 2, + steps_per_generation = 4, + num_generations = 8, + gradient_accumulation_steps = 1, + ) + try: + _build_flex_from_args(model, args) + except RuntimeError as exc: + msg = str(exc) + assert "32" in msg and "64" in msg, msg + assert "max_batch_size=64" in msg, msg + print(" [4/4] post-warmup rebuild refused with actionable msg OK") + return + raise AssertionError("expected RuntimeError when growing a built engine") + + +def main(): + print("flex_lazy_batch_smoke:") + _case1_default_floor() + _case2_grpo_bump() + _case3_user_floor_wins() + _case4_post_warmup_refused() + print("ALL CASES PASSED") + + +if __name__ == "__main__": + main() diff --git a/tests/flex_lazy_live_smoke.py b/tests/flex_lazy_live_smoke.py new file mode 100644 index 0000000000..093aceb64c --- /dev/null +++ b/tests/flex_lazy_live_smoke.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""Live smoke-test for deferred FlexEngine construction. + +Loads a small flex-supported model with ``UNSLOTH_FAST_INFERENCE=1`` and +checks: + + 1. After ``from_pretrained``, ``model.vllm_engine`` is the lazy + sentinel -- no ``_flex_engine_instance`` yet. + 2. ``build_flex_engine(model)`` constructs the engine at the stashed + ``max_batch_size`` floor (default 32) and wires + ``model.vllm_engine`` / ``fast_generate`` onto it. + 3. ``build_flex_engine(model, max_batch_size=X)`` with ``X`` larger + than the built size raises :class:`RuntimeError` with an actionable + hint pointing the user back to ``max_batch_size=`` in + ``from_pretrained``. + +This intentionally does NOT call ``engine.generate`` -- that path is +covered by the existing ``tests/flex_fastlm_smoke.py`` and would +duplicate its warm-up cost. The goal here is to confirm the lazy +dispatch behavior end-to-end. + +Run as: + CUDA_VISIBLE_DEVICES=0 UNSLOTH_FAST_INFERENCE=1 \ + python tests/flex_lazy_live_smoke.py +""" + +from __future__ import annotations + +import os +import sys +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)) + + +def main(): + assert ( + os.environ.get("UNSLOTH_FAST_INFERENCE", "0") == "1" + ), "export UNSLOTH_FAST_INFERENCE=1 before running this smoke" + import unsloth # noqa: F401 (must import before transformers) + from unsloth import FastLanguageModel + from unsloth.inference.flex_engine import ( + FlexEngine, + _LazyFlexEngineSentinel, + build_flex_engine, + ) + + model_name = os.environ.get("FLEX_LAZY_SMOKE_MODEL", "unsloth/Qwen3-0.6B-Base") + print(f"loading {model_name} ...") + model, tokenizer = FastLanguageModel.from_pretrained( + model_name = model_name, + max_seq_length = 1024, + fast_inference = True, + load_in_4bit = False, + ) + + assert hasattr(model, "vllm_engine"), "vllm_engine attr missing" + sentinel = model.vllm_engine + assert isinstance( + sentinel, _LazyFlexEngineSentinel + ), f"expected sentinel, got {type(sentinel)}" + assert not hasattr( + model, "_flex_engine_instance" + ), "engine should NOT be built before first use" + print(" [1/3] sentinel installed, no engine yet") + + engine = build_flex_engine(model) + assert isinstance(engine, FlexEngine), type(engine) + assert engine.max_batch_size == 32, engine.max_batch_size + assert model._flex_engine_instance is engine + assert model.vllm_engine is engine + print( + f" [2/3] build_flex_engine built engine at max_batch_size={engine.max_batch_size}" + ) + + try: + build_flex_engine(model, max_batch_size = 64) + except RuntimeError as exc: + msg = str(exc) + assert "32" in msg and "64" in msg, msg + assert "max_batch_size=64" in msg, msg + print(" [3/3] post-build resize refused with actionable msg") + else: + raise AssertionError("expected RuntimeError when growing a built engine") + + print("LIVE SMOKE PASSED") + + +if __name__ == "__main__": + main() diff --git a/tests/flex_moe_bench.py b/tests/flex_moe_bench.py new file mode 100644 index 0000000000..dc1aba2071 --- /dev/null +++ b/tests/flex_moe_bench.py @@ -0,0 +1,199 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""Decode throughput bench: FlexMoEInference vs HF generate on Qwen3 MoE. + +Apples-to-apples decode on the same prompt set + new-token budget, +same LoRA rank, same precision. Mirrors PR #5123's +``tests/flex_fastlm_bench.py`` shape. + +Usage: + CUDA_VISIBLE_DEVICES=0 UNSLOTH_FAST_INFERENCE=1 python -u \ + tests/flex_moe_bench.py --backend flex --load_in_4bit + + CUDA_VISIBLE_DEVICES=0 python -u \ + tests/flex_moe_bench.py --backend hf --load_in_4bit +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import 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)) + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--backend", choices = ["flex", "hf", "hf_naive"], default = "flex") + p.add_argument("--model", default = "unsloth/Qwen3-30B-A3B-Instruct-2507") + p.add_argument("--dtype", choices = ["bf16", "fp16"], default = "bf16") + p.add_argument("--load_in_4bit", action = "store_true") + p.add_argument("--n_prompts", type = int, default = 8) + p.add_argument("--max_new_tokens", type = int, default = 64) + p.add_argument("--max_seq_length", type = int, default = 1024) + p.add_argument("--warmup_rounds", type = int, default = 1) + p.add_argument("--timed_rounds", type = int, default = 2) + p.add_argument("--out_dir", default = "async_task_outputs/qwen3_moe_grpo_bench") + p.add_argument("--chat_template", action = "store_true", + help = "wrap each prompt with tokenizer.apply_chat_template") + p.add_argument("--user_prompt", + default = "Continue this sentence: The quick brown fox jumps over fence {i}, then") + args = p.parse_args() + + import torch + + if args.backend == "flex": + os.environ["UNSLOTH_FAST_INFERENCE"] = "1" + os.environ.setdefault("UNSLOTH_MOE_BACKEND", "grouped_mm") + + dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16 + torch.cuda.reset_peak_memory_stats() + t_load0 = time.perf_counter() + + if args.backend == "hf_naive": + # Pure transformers path: NO ``import unsloth`` so none of + # Unsloth's Qwen3 MoE attention / MLP patches run. This is the + # fair naive reference to compare flex fast-inference against. + from transformers import ( + AutoModelForCausalLM, + AutoTokenizer, + BitsAndBytesConfig, + ) + quant_cfg = None + if args.load_in_4bit: + quant_cfg = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_compute_dtype=dtype, + bnb_4bit_use_double_quant=True, + bnb_4bit_quant_type="nf4", + ) + tokenizer = AutoTokenizer.from_pretrained(args.model) + if tokenizer.pad_token_id is None or tokenizer.pad_token == "<|PAD_TOKEN|>": + tokenizer.pad_token = "<|vision_pad|>" + tokenizer.padding_side = "left" + attn_impl = os.environ.get("HF_ATTN_IMPL", "eager") + print(f"[bench] HF attn_implementation={attn_impl}") + model = AutoModelForCausalLM.from_pretrained( + args.model, + dtype=dtype, + quantization_config=quant_cfg, + device_map="cuda", + attn_implementation=attn_impl, + ) + model.eval() + print(f"[bench] pure transformers (no unsloth patches)") + else: + import unsloth + print(f"[bench] unsloth={unsloth.__file__}") + from unsloth import FastLanguageModel + 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 = args.backend == "flex", + ) + + t_load = time.perf_counter() - t_load0 + peak_load = torch.cuda.max_memory_reserved() / 1024**3 + print(f"[bench] loaded in {t_load:.1f}s peak {peak_load:.1f} GB") + + if args.chat_template: + prompts = [] + for i in range(args.n_prompts): + msg = [{"role": "user", "content": args.user_prompt.format(i=i)}] + prompts.append(tokenizer.apply_chat_template( + msg, tokenize = False, add_generation_prompt = True, + )) + print(f"[bench] chat_template enabled. prompt[0] (first 200 chars):\n" + f" {prompts[0][:200]!r}") + else: + prompts = [f"The quick brown fox jumps over fence {i}, then" + for i in range(args.n_prompts)] + + if args.backend == "flex": + class _SP: + max_tokens = args.max_new_tokens + temperature = 0.0 + + # Warmup + for _ in range(args.warmup_rounds): + _ = model.fast_generate(prompts, sampling_params = _SP(), use_tqdm = False) + + # Timed + wall_times = [] + tok_counts = [] + for _ in range(args.timed_rounds): + t0 = time.perf_counter() + outs = model.fast_generate(prompts, sampling_params = _SP(), use_tqdm = False) + wall_times.append(time.perf_counter() - t0) + tok_counts.append( + sum(len(o.outputs[0].token_ids) for o in outs) + ) + else: + # HF generate (shared for "hf" unsloth-patched and "hf_naive" pure). + inputs = tokenizer(prompts, return_tensors = "pt", padding = True).to("cuda") + gen_kwargs = dict( + max_new_tokens = args.max_new_tokens, + do_sample = False, + temperature = 1.0, + pad_token_id = tokenizer.pad_token_id or tokenizer.eos_token_id, + ) + # Warmup + for _ in range(args.warmup_rounds): + _ = model.generate(**inputs, **gen_kwargs) + torch.cuda.synchronize() + + wall_times = [] + tok_counts = [] + for _ in range(args.timed_rounds): + torch.cuda.synchronize() + t0 = time.perf_counter() + out = model.generate(**inputs, **gen_kwargs) + torch.cuda.synchronize() + wall_times.append(time.perf_counter() - t0) + n_new = (out.shape[1] - inputs["input_ids"].shape[1]) * out.shape[0] + tok_counts.append(n_new) + + peak_gen = torch.cuda.max_memory_reserved() / 1024**3 + median_wall = sorted(wall_times)[len(wall_times) // 2] + median_tok = tok_counts[len(wall_times) // 2] + tok_per_s = median_tok / median_wall if median_wall > 0 else 0.0 + + print(f"[bench] wall: {wall_times}") + print(f"[bench] tok counts: {tok_counts}") + print(f"[bench] median wall: {median_wall:.2f}s median tok/s: {tok_per_s:.1f}") + print(f"[bench] peak VRAM after gen: {peak_gen:.1f} GB") + + precision = "4bit" if args.load_in_4bit else args.dtype + out_dir = Path(args.out_dir) + out_dir.mkdir(parents = True, exist_ok = True) + summary = { + "phase": "bench_decode", + "backend": args.backend, + "model": args.model, + "precision": precision, + "n_prompts": args.n_prompts, + "max_new_tokens": args.max_new_tokens, + "wall_times_s": wall_times, + "tok_counts": tok_counts, + "median_wall_s": round(median_wall, 3), + "median_tok_s": round(tok_per_s, 1), + "peak_vram_load_gb": round(peak_load, 2), + "peak_vram_after_gen_gb": round(peak_gen, 2), + "t_load_s": round(t_load, 1), + } + with open(out_dir / f"bench_decode_{args.backend}_{precision}.json", "w") as f: + json.dump(summary, f, indent = 2) + print(f"[bench] wrote {out_dir / f'bench_decode_{args.backend}_{precision}.json'}") + + +if __name__ == "__main__": + main() diff --git a/tests/flex_moe_merge_parity.py b/tests/flex_moe_merge_parity.py new file mode 100644 index 0000000000..844aad6e81 --- /dev/null +++ b/tests/flex_moe_merge_parity.py @@ -0,0 +1,290 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""Numerical + performance check for ``refresh_moe_lora_merge_from_pristine``. + +Does what we never actually did on this branch: verify the stacked MoE +LoRA merge kernel produces the same W_inf as the textbook reference +``W_ref[e] = W_pristine[e] + sum_a scaling_a * B_a[e] @ A_a[e]`` for both +standard (E, 2I, H) and transposed (E, H, 2I) expert layouts, with 1 and +2 active adapters, then benchmarks the batched ``baddbmm`` path against +the dense-layer-style ``addmm`` loop at Qwen3-30B-A3B MoE shapes. + +Usage:: + CUDA_VISIBLE_DEVICES=2 python -u tests/flex_moe_merge_parity.py +""" +from __future__ import annotations + +import sys +import time +from pathlib import Path + +import torch + +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + + +# --------------------------------------------------------------------------- +# Reference: the exact formula from the docstring, no tricks. +# --------------------------------------------------------------------------- + +def reference_merge(W_pristine, adapters, *, transposed): + """adapters: list of (A_stacked [E*R, in], B_stacked [out, E*R], scaling).""" + E = W_pristine.shape[0] + out = W_pristine.clone() + for A_stack, B_stack, scaling in adapters: + R = A_stack.shape[0] // E + in_dim = A_stack.shape[1] + out_dim = B_stack.shape[0] + for e in range(E): + A_e = A_stack[e * R : (e + 1) * R] # (R, in_dim) + B_e = B_stack[:, e * R : (e + 1) * R] # (out_dim, R) + delta = (B_e @ A_e).to(W_pristine.dtype) # (out_dim, in_dim) + if transposed: + out[e] += scaling * delta.t() + else: + out[e] += scaling * delta + return out + + +# --------------------------------------------------------------------------- +# Under-test: a standalone copy of the merge body from flex_moe.py so we +# can drive it without loading a 30B model. The logic is a character-for- +# character lift of lines 747-794 in unsloth/inference/flex_moe.py. +# --------------------------------------------------------------------------- + +def merge_under_test(W_inf, W_pristine, adapters): + """adapters: list of (A_stacked, B_stacked, scaling).""" + E = W_inf.shape[0] + # Orientation detection (mirrors flex_moe.py:752-763). + A_w0, B_w0, _ = adapters[0] + in_dim = A_w0.shape[1] + out_dim = B_w0.shape[0] + d0, d1 = W_inf.shape[1], W_inf.shape[2] + if d0 == out_dim and d1 == in_dim: + is_standard = True + elif d0 == in_dim and d1 == out_dim: + is_standard = False + else: + raise RuntimeError("orientation") + + W_inf.copy_(W_pristine) + + for A_w, B_w, scaling in adapters: + R = A_w.shape[0] // E + A_3d = A_w.view(E, R, in_dim) + B_3d = B_w.view(out_dim, E, R).permute(1, 0, 2).contiguous() + if is_standard: + torch.baddbmm( + W_inf, + B_3d.to(W_inf.dtype), + A_3d.to(W_inf.dtype), + alpha=float(scaling), + beta=1.0, + out=W_inf, + ) + else: + torch.baddbmm( + W_inf, + A_3d.transpose(-2, -1).contiguous().to(W_inf.dtype), + B_3d.transpose(-2, -1).contiguous().to(W_inf.dtype), + alpha=float(scaling), + beta=1.0, + out=W_inf, + ) + return W_inf + + +# --------------------------------------------------------------------------- +# Correctness +# --------------------------------------------------------------------------- + +def _make_adapter(E, R, in_dim, out_dim, dtype, device, seed): + g = torch.Generator(device=device).manual_seed(seed) + A = torch.randn(E * R, in_dim, generator=g, device=device, dtype=dtype) * 0.02 + B = torch.randn(out_dim, E * R, generator=g, device=device, dtype=dtype) * 0.02 + return A, B + + +def test_correctness(E, in_dim, out_dim, R, dtype, device, *, transposed, n_adapters): + scalings = [2.0, 0.5][:n_adapters] + adapters = [ + (*_make_adapter(E, R, in_dim, out_dim, dtype, device, seed=11 + i), s) + for i, s in enumerate(scalings) + ] + + if transposed: + W_pristine = torch.randn(E, in_dim, out_dim, device=device, dtype=dtype) * 0.02 + else: + W_pristine = torch.randn(E, out_dim, in_dim, device=device, dtype=dtype) * 0.02 + W_inf = torch.empty_like(W_pristine) + + merge_under_test(W_inf, W_pristine, adapters) + W_ref = reference_merge(W_pristine, adapters, transposed=transposed) + + tol = dict(atol=3e-3, rtol=3e-3) if dtype == torch.bfloat16 else dict(atol=1e-5, rtol=1e-5) + close = torch.allclose(W_inf, W_ref, **tol) + max_err = (W_inf - W_ref).abs().max().item() + scale = W_ref.abs().max().item() + rel = max_err / max(scale, 1e-8) + label = f"E={E} in={in_dim} out={out_dim} R={R} {dtype} transposed={transposed} n_adapters={n_adapters}" + verdict = "OK" if close else "FAIL" + print(f" [{verdict}] {label} max_abs={max_err:.3e} rel={rel:.3e}") + return close + + +# --------------------------------------------------------------------------- +# Performance: batched baddbmm vs per-expert addmm loop (same total flops). +# --------------------------------------------------------------------------- + +def _sync(): + if torch.cuda.is_available(): + torch.cuda.synchronize() + + +def bench(fn, *args, iters=50, warmup=10): + for _ in range(warmup): + fn(*args) + _sync() + t0 = time.perf_counter() + for _ in range(iters): + fn(*args) + _sync() + return (time.perf_counter() - t0) / iters + + +def perf_vs_dense_addmm(out_dim, in_dim, R, dtype, device): + """Dense-layer LoRA refresh: a single torch.addmm call like the one + in flex_qwen3_llama.py::refresh_lora_merge_from_pristine. Used as a + per-matrix baseline to compare against the MoE baddbmm kernel cost.""" + W_pristine = torch.randn(out_dim, in_dim, device=device, dtype=dtype) * 0.02 + W = torch.empty_like(W_pristine) + A = torch.randn(R, in_dim, device=device, dtype=dtype) * 0.02 + B = torch.randn(out_dim, R, device=device, dtype=dtype) * 0.02 + scaling = 2.0 + + def run_addmm(): + torch.addmm(W_pristine, B, A, alpha=scaling, out=W) + + t = bench(run_addmm, iters=200, warmup=20) + print(f" dense addmm out={out_dim:>5} in={in_dim:>5} R={R:>2} " + f"{t * 1e3:7.4f}ms") + return t + + +def perf_compare(E, in_dim, out_dim, R, dtype, device): + A_w, B_w = _make_adapter(E, R, in_dim, out_dim, dtype, device, seed=0) + W_pristine = torch.randn(E, out_dim, in_dim, device=device, dtype=dtype) * 0.02 + W_inf_a = torch.empty_like(W_pristine) + W_inf_b = torch.empty_like(W_pristine) + + A_3d = A_w.view(E, R, in_dim) + B_3d = B_w.view(out_dim, E, R).permute(1, 0, 2).contiguous() + scaling = 2.0 + + def run_baddbmm(): + W_inf_a.copy_(W_pristine) + torch.baddbmm(W_inf_a, B_3d, A_3d, alpha=scaling, beta=1.0, out=W_inf_a) + + def run_addmm_loop(): + W_inf_b.copy_(W_pristine) + for e in range(E): + torch.addmm( + W_inf_b[e], + B_3d[e], + A_3d[e], + alpha=scaling, + out=W_inf_b[e], + ) + + t_bad = bench(run_baddbmm) + t_loop = bench(run_addmm_loop) + + # Sanity: they produce the same result. + run_baddbmm() + run_addmm_loop() + max_err = (W_inf_a - W_inf_b).abs().max().item() + + print(f" E={E:>3} out={out_dim:>5} in={in_dim:>5} R={R:>2} " + f"baddbmm={t_bad * 1e3:7.3f}ms addmm_loop={t_loop * 1e3:7.3f}ms " + f"speedup={t_loop / t_bad:5.2f}x max_abs_diff={max_err:.1e}") + + +def main(): + device = "cuda" if torch.cuda.is_available() else "cpu" + torch.manual_seed(3407) + print(f"[merge-parity] device={device} dtype=bf16 + fp32") + print(f"[merge-parity] torch={torch.__version__}") + print() + + print("== Correctness (fp32 golden + bf16 realistic) ==") + # Qwen3-30B-A3B shapes. + # gate_up_proj: (E=128, 2I=1536, H=2048) standard + # down_proj: (E=128, H=2048, I=768) standard + # Transposed variant: (E, H, 2I) — exercised via the `transposed=True` test. + cases = [ + # (E, in_dim, out_dim, R, dtype, transposed, n_adapters) + ( 8, 128, 192, 4, torch.float32, False, 1), + ( 8, 128, 192, 4, torch.float32, True, 1), + ( 8, 128, 192, 4, torch.float32, False, 2), + (16, 256, 512, 8, torch.float32, True, 2), + # bf16 at realistic Qwen3 MoE shapes. Mild tol to allow 3e-3 bf16 noise. + (128, 2048, 1536, 16, torch.bfloat16, False, 1), # gate_up_proj-like + (128, 768, 2048, 16, torch.bfloat16, False, 1), # down_proj-like + (128, 2048, 1536, 16, torch.bfloat16, True, 1), # transposed gate_up + (128, 2048, 1536, 16, torch.bfloat16, False, 2), # two adapters + ] + all_ok = True + for E, in_dim, out_dim, R, dtype, tr, na in cases: + ok = test_correctness(E, in_dim, out_dim, R, dtype, device, + transposed=tr, n_adapters=na) + all_ok = all_ok and ok + print(f"\n overall: {'PASS' if all_ok else 'FAIL'}") + + print() + print("== Performance: batched baddbmm vs per-expert addmm loop (bf16) ==") + print(" (same arithmetic, different dispatch pattern — baddbmm = 1 kernel,") + print(" addmm loop = E kernels)") + moe_cases = [ + (128, 2048, 1536, 16), # gate_up_proj + (128, 768, 2048, 16), # down_proj + (128, 2048, 1536, 64), # larger rank + ] + for E, in_dim, out_dim, R in moe_cases: + perf_compare(E, in_dim, out_dim, R, torch.bfloat16, device) + + print() + print("== Dense baseline: torch.addmm on a single expert-sized matrix ==") + print(" (this is what flex_qwen3_llama.py:354-370 does for dense layers)") + dense_times = {} + for E, in_dim, out_dim, R in moe_cases: + t = perf_vs_dense_addmm(out_dim, in_dim, R, torch.bfloat16, device) + dense_times[(E, in_dim, out_dim, R)] = t + + print() + print("== Per-expert cost comparison ==") + print(" (MoE baddbmm cost / E) vs single dense addmm for the same per-expert matrix") + for E, in_dim, out_dim, R in moe_cases: + # Measure baddbmm cost again for this exact shape to get the numerator. + A_w, B_w = _make_adapter(E, R, in_dim, out_dim, torch.bfloat16, device, seed=0) + W_pristine = torch.randn(E, out_dim, in_dim, device=device, dtype=torch.bfloat16) * 0.02 + W_inf = torch.empty_like(W_pristine) + A_3d = A_w.view(E, R, in_dim) + B_3d = B_w.view(out_dim, E, R).permute(1, 0, 2).contiguous() + + def run_baddbmm(): + W_inf.copy_(W_pristine) + torch.baddbmm(W_inf, B_3d, A_3d, alpha=2.0, beta=1.0, out=W_inf) + t_moe = bench(run_baddbmm, iters=50, warmup=10) + t_dense = dense_times[(E, in_dim, out_dim, R)] + per_expert = t_moe / E + print(f" E={E:>3} out={out_dim:>5} in={in_dim:>5} R={R:>2} " + f"moe_per_expert={per_expert * 1e6:7.2f}us " + f"dense_addmm={t_dense * 1e6:7.2f}us " + f"ratio(moe/dense)={per_expert / t_dense:5.2f}x") + + +if __name__ == "__main__": + main() diff --git a/tests/flex_moe_micro_bench.py b/tests/flex_moe_micro_bench.py new file mode 100644 index 0000000000..d2bf5a67ca --- /dev/null +++ b/tests/flex_moe_micro_bench.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""Tight, single-load decode throughput probe for Qwen3 MoE. + +Loads the model once, captures CUDA graphs, then sweeps batch sizes. +Avoids the 30s cold-load tax of the full bench so optimization +iterations can run in under a minute per config. + +Usage: + CUDA_VISIBLE_DEVICES=5 UNSLOTH_FAST_INFERENCE=1 \\ + UNSLOTH_MOE_BACKEND=grouped_mm python -u \\ + tests/flex_moe_micro_bench.py --load_in_4bit --bs 1,4,8,16,32 +""" +from __future__ import annotations + +import argparse +import json +import os +import sys +import 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)) + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--model", default="unsloth/Qwen3-30B-A3B-Instruct-2507") + p.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16") + p.add_argument("--load_in_4bit", action="store_true") + p.add_argument("--bs", default="1,4,8,16,32", + help="comma-separated batch sizes to sweep") + p.add_argument("--max_new_tokens", type=int, default=128) + p.add_argument("--max_seq_length", type=int, default=1024) + p.add_argument("--max_batch_size", type=int, default=32) + p.add_argument("--warmup_rounds", type=int, default=1) + p.add_argument("--timed_rounds", type=int, default=2) + p.add_argument("--tag", default="baseline", + help="label for this config in the output JSON") + p.add_argument("--compile_mode", choices=["off", "walker", "walker_fullgraph"], + default="off", + help="wrap call_moe_model_with_flex_kwargs in torch.compile") + p.add_argument("--compile_opts", choices=["stock", "unsloth_O3", "inference_freeze", "coord_descent"], + default="stock", + help="which inductor / dynamo options profile to apply before compile") + p.add_argument("--explain", action="store_true", + help="run torch._dynamo.explain on the walker first to list breaks") + p.add_argument("--out_dir", default="async_task_outputs/qwen3_moe_grpo_bench_v2") + args = p.parse_args() + + bs_list = [int(x) for x in args.bs.split(",") if x.strip()] + os.environ["UNSLOTH_FAST_INFERENCE"] = "1" + os.environ.setdefault("UNSLOTH_MOE_BACKEND", "grouped_mm") + import torch + + import unsloth # noqa: F401 + from unsloth import FastLanguageModel + from unsloth.inference import flex_moe as _flex_moe_mod + + # Apply inductor / dynamo config BEFORE wrapping with torch.compile. + if args.compile_opts == "unsloth_O3": + # Aggressive autotune + coord descent + aggressive_fusion. Matches + # unsloth_zoo.patching_utils.patch_torch_compile(O3=True). + import torch._inductor.config as _ic + import torch._dynamo.config as _dc + _ic.max_autotune = True + _ic.max_autotune_pointwise = True + _ic.coordinate_descent_tuning = True + _ic.aggressive_fusion = True + _ic.cuda.use_fast_math = True + _dc.cache_size_limit = 1024 + _dc.recompile_limit = 1024 + _dc.capture_scalar_outputs = True + _dc.capture_dynamic_output_shape_ops = True + print("[micro] inductor/dynamo options: unsloth_O3") + elif args.compile_opts == "coord_descent": + # Just ``coordinate_descent_tuning = True`` — fast compile, small + # fusion upside. + import torch._inductor.config as _ic + _ic.coordinate_descent_tuning = True + print("[micro] inductor options: coord_descent only") + elif args.compile_opts == "inference_freeze": + # Inference-friendly: constant-fold weights via freezing=True. + # Only safe when the model weights won't be updated after compile + # (true here — we capture graphs post-load and never refresh + # during bench). + import torch._inductor.config as _ic + import torch._dynamo.config as _dc + _ic.freezing = True + _ic.max_autotune = True + _ic.coordinate_descent_tuning = True + _ic.cuda.use_fast_math = True + _dc.cache_size_limit = 1024 + _dc.capture_scalar_outputs = True + print("[micro] inductor/dynamo options: inference_freeze") + + # Apply torch.compile to the decode walker BEFORE the engine is built + # / graphs are captured, so the compiled kernels get recorded into the + # CUDA graph. + if args.compile_mode != "off": + fullgraph = args.compile_mode == "walker_fullgraph" + orig_walker = _flex_moe_mod.call_moe_model_with_flex_kwargs + compile_kwargs = dict(fullgraph=fullgraph, dynamic=False) + tmode = os.environ.get("FLEX_COMPILE_MODE", "") + if tmode: + compile_kwargs["mode"] = tmode + compiled = torch.compile(orig_walker, **compile_kwargs) + _flex_moe_mod.call_moe_model_with_flex_kwargs = compiled + print(f"[micro] wrapped call_moe_model_with_flex_kwargs with " + f"torch.compile(fullgraph={fullgraph}, mode={tmode or 'default'})") + + dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16 + torch.cuda.reset_peak_memory_stats() + + 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, + max_batch_size=args.max_batch_size, + ) + print(f"[micro] loaded in {time.perf_counter() - t0:.1f}s") + + class _SP: + max_tokens = args.max_new_tokens + temperature = 0.0 + + results = [] + for bs in bs_list: + prompts = [f"The quick brown fox jumps over fence {i}, then" + for i in range(bs)] + + # Warmup (first call captures the graphs for all buckets). + for _ in range(args.warmup_rounds): + _ = model.fast_generate(prompts, sampling_params=_SP(), use_tqdm=False) + + # Timed. + wall = [] + n_tok = [] + for _ in range(args.timed_rounds): + torch.cuda.synchronize() + t0 = time.perf_counter() + outs = model.fast_generate(prompts, sampling_params=_SP(), use_tqdm=False) + torch.cuda.synchronize() + wall.append(time.perf_counter() - t0) + n_tok.append(sum(len(o.outputs[0].token_ids) for o in outs)) + + med_wall = sorted(wall)[len(wall) // 2] + med_tok = n_tok[len(wall) // 2] + tps = med_tok / med_wall if med_wall > 0 else 0.0 + # Sanity: print the first completion so we can eyeball for + # gibberish. A compile bug or bad capture shows up here first. + sample_text = outs[0].outputs[0].text if outs else "" + sample_preview = sample_text.replace("\n", "\\n")[:120] + print(f"[micro] bs={bs:>3} tok={med_tok:>5} " + f"wall={med_wall:.3f}s tok/s={tps:.1f}") + print(f"[micro] bs={bs:>3} completion[0]: {sample_preview!r}") + results.append({ + "bs": bs, + "max_new_tokens": args.max_new_tokens, + "median_wall_s": round(med_wall, 3), + "median_tok": med_tok, + "tok_per_s": round(tps, 1), + "wall_times_s": wall, + "sample_completion": sample_text[:400], + }) + + peak = torch.cuda.max_memory_reserved() / 1024**3 + precision = "4bit" if args.load_in_4bit else args.dtype + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + out_path = out_dir / f"micro_bench_{args.tag}_{precision}.json" + with open(out_path, "w") as f: + json.dump({ + "tag": args.tag, + "precision": precision, + "peak_vram_gb": round(peak, 2), + "results": results, + }, f, indent=2) + print(f"[micro] peak VRAM: {peak:.1f} GB") + print(f"[micro] wrote {out_path}") + + +if __name__ == "__main__": + main() diff --git a/tests/flex_moe_parity.py b/tests/flex_moe_parity.py new file mode 100644 index 0000000000..6cf84d1830 --- /dev/null +++ b/tests/flex_moe_parity.py @@ -0,0 +1,184 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""Token-level parity: FlexMoEInference (CUDA-graph capture) vs HF generate. + +Same prompt set, temperature=0, max_new_tokens fixed. Reports per-prompt +token-id match rate + first divergence index. Serves as the correctness +check for the v2 grouped_mm + CUDA-graph-capture changes. + +Usage: + CUDA_VISIBLE_DEVICES=5 UNSLOTH_FAST_INFERENCE=1 \ + UNSLOTH_MOE_BACKEND=grouped_mm python -u \ + tests/flex_moe_parity.py --load_in_4bit +""" +from __future__ import annotations + +import argparse +import json +import os +import sys +import 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)) + + +def _run_flex(prompts, args, dtype, *, capture: bool): + import torch + os.environ["UNSLOTH_FAST_INFERENCE"] = "1" + os.environ.setdefault("UNSLOTH_MOE_BACKEND", "grouped_mm") + import unsloth # noqa: F401 + from unsloth import FastLanguageModel + + if not capture: + # Monkey-patch ``capture_decode_cudagraph`` to a no-op BEFORE the + # engine is built so ``generate`` takes the eager branch. The + # engine's ``self.graphs`` stays empty, ``cudagraph_captured`` + # stays False, and every step goes through ``_decode_step_eager``. + from unsloth.inference.flex_moe import FlexMoEInference + FlexMoEInference.capture_decode_cudagraph = lambda self: None + + # Opt-in torch.compile wrap of the decode walker for parity check + # — enabled via env var to avoid cluttering the CLI further. + if os.environ.get("FLEX_MOE_COMPILE_WALKER") == "1": + import torch as _torch + from unsloth.inference import flex_moe as _flex_moe_mod + _orig = _flex_moe_mod.call_moe_model_with_flex_kwargs + _flex_moe_mod.call_moe_model_with_flex_kwargs = _torch.compile( + _orig, fullgraph=False, dynamic=False + ) + print("[parity] torch.compile(call_moe_model_with_flex_kwargs) enabled") + + 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, + ) + + class _SP: + max_tokens = args.max_new_tokens + temperature = 0.0 + + # First call warms / captures; second call is the measurement. + _ = model.fast_generate(prompts, sampling_params=_SP(), use_tqdm=False) + outputs = model.fast_generate(prompts, sampling_params=_SP(), use_tqdm=False) + token_ids = [list(o.outputs[0].token_ids) for o in outputs] + texts = [o.outputs[0].text for o in outputs] + return token_ids, texts, tokenizer + + +def _run_hf(prompts, args, dtype): + import torch + # Pure Hugging Face: NO ``import unsloth`` — we want the unpatched + # reference forward to compare flex against. Quantization via + # transformers' ``BitsAndBytesConfig`` matches what unsloth loads + # under the hood for ``load_in_4bit=True``. + from transformers import ( + AutoModelForCausalLM, + AutoTokenizer, + BitsAndBytesConfig, + ) + # Translate the unsloth-flavoured model id (unsloth/Qwen3-30B-A3B-Instruct-2507) + # to the 4bit variant if load_in_4bit was requested (FastLanguageModel + # does this implicitly; do it explicitly here for the naive path). + model_id = args.model + quant_cfg = None + if args.load_in_4bit: + quant_cfg = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_compute_dtype=dtype, + bnb_4bit_use_double_quant=True, + bnb_4bit_quant_type="nf4", + ) + tokenizer = AutoTokenizer.from_pretrained(model_id) + model = AutoModelForCausalLM.from_pretrained( + model_id, + dtype=dtype, + quantization_config=quant_cfg, + device_map="cuda", + attn_implementation="eager", + ) + model.eval() + # Qwen3-30B-A3B-Instruct-2507 uses <|vision_pad|> as its pad token. + # Unsloth's loader may swap it to a sentinel; reset to the HF default + # so batched left-padded generation matches the authoritative config. + if tokenizer.pad_token_id is None or tokenizer.pad_token == "<|PAD_TOKEN|>": + tokenizer.pad_token = "<|vision_pad|>" + tokenizer.padding_side = "left" + gen_kwargs = dict( + max_new_tokens=args.max_new_tokens, + do_sample=False, + temperature=1.0, + pad_token_id=tokenizer.pad_token_id, + ) + inputs = tokenizer(prompts, return_tensors="pt", padding=True).to("cuda") + out = model.generate(**inputs, **gen_kwargs) + prompt_len = inputs["input_ids"].shape[1] + eos = tokenizer.eos_token_id + pad = tokenizer.pad_token_id + token_ids = [] + texts = [] + for row in out: + ids = row[prompt_len:].tolist() + while ids and ids[-1] in (eos, pad): + ids.pop() + token_ids.append(ids) + texts.append(tokenizer.decode(ids, skip_special_tokens=True)) + return token_ids, texts, tokenizer + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--model", default="unsloth/Qwen3-30B-A3B-Instruct-2507") + p.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16") + p.add_argument("--load_in_4bit", 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("--backend", choices=["flex", "flex_eager", "hf"], required=True) + p.add_argument("--out_dir", default="async_task_outputs/qwen3_moe_grpo_bench_v2") + args = p.parse_args() + + import torch + prompts = [ + "The quick brown fox jumps over", + "Q: What is 23 + 19?\nA:", + "Paris is the capital of", + ] + dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16 + + if args.backend == "flex": + token_ids, texts, tok = _run_flex(prompts, args, dtype, capture=True) + elif args.backend == "flex_eager": + token_ids, texts, tok = _run_flex(prompts, args, dtype, capture=False) + else: + token_ids, texts, tok = _run_hf(prompts, args, dtype) + + precision = "4bit" if args.load_in_4bit else args.dtype + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + out_path = out_dir / f"parity_{args.backend}_{precision}.json" + with open(out_path, "w") as f: + json.dump( + { + "backend": args.backend, + "precision": precision, + "prompts": prompts, + "token_ids": token_ids, + "texts": texts, + }, + f, + indent=2, + ) + print(f"[parity-{args.backend}] wrote {out_path}") + for i, (p_, t_) in enumerate(zip(prompts, texts)): + print(f"[parity-{args.backend}] prompt {i}: {p_!r}") + print(f"[parity-{args.backend}] completion {i}: {t_!r}") + + +if __name__ == "__main__": + main() diff --git a/tests/flex_moe_smoke.py b/tests/flex_moe_smoke.py new file mode 100644 index 0000000000..beab04cc33 --- /dev/null +++ b/tests/flex_moe_smoke.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""Smoke-test ``UNSLOTH_FAST_INFERENCE=1`` on a Qwen3 MoE model. + +Mirrors ``tests/flex_fastlm_smoke.py`` but targets the new +``FlexMoEInference`` path added for ``Qwen3MoeForCausalLM``. + +Invoked as: + CUDA_VISIBLE_DEVICES=0 UNSLOTH_FAST_INFERENCE=1 python -u \ + tests/flex_moe_smoke.py \ + --model unsloth/Qwen3-30B-A3B-Instruct-2507 \ + --load_in_4bit + +Writes a small JSON summary to +``async_task_outputs/qwen3_moe_grpo_bench/smoke_A_{precision}.json``. +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import 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)) + + +def main(): + p = argparse.ArgumentParser() + p.add_argument( + "--model", default = "unsloth/Qwen3-30B-A3B-Instruct-2507" + ) + 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") + p.add_argument("--out_dir", default = "async_task_outputs/qwen3_moe_grpo_bench") + args = p.parse_args() + + import torch + + os.environ.setdefault("UNSLOTH_FAST_INFERENCE", "1") + os.environ.setdefault("UNSLOTH_MOE_BACKEND", "grouped_mm") + print(f"[smoke] UNSLOTH_FAST_INFERENCE={os.environ.get('UNSLOTH_FAST_INFERENCE')}") + print(f"[smoke] UNSLOTH_MOE_BACKEND={os.environ.get('UNSLOTH_MOE_BACKEND')}") + + import unsloth + + print(f"[smoke] unsloth={unsloth.__file__}") + from unsloth import FastLanguageModel + + dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16 + + torch.cuda.reset_peak_memory_stats() + 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 + peak_after_load = torch.cuda.max_memory_reserved() / 1024**3 + print(f"[smoke] loaded model in {t_load:.1f}s; peak VRAM after load: {peak_after_load:.2f} GB") + print(f"[smoke] hasattr(model, 'vllm_engine'): {hasattr(model, 'vllm_engine')}") + print(f"[smoke] vllm_engine type: {type(model.vllm_engine).__name__}") + arch = getattr(model.vllm_engine, "arch", "?") + impl = type(model.vllm_engine._impl).__name__ + print(f"[smoke] FlexEngine.arch={arch} impl={impl}") + + if args.with_lora: + model = FastLanguageModel.get_peft_model( + model, + r = 16, + target_modules = [ + "q_proj", "k_proj", "v_proj", "o_proj", + "gate_proj", "up_proj", "down_proj", "gate_up_proj", + ], + lora_alpha = 32, + lora_dropout = 0.0, + bias = "none", + use_gradient_checkpointing = "unsloth", + random_state = 3407, + ) + print(f"[smoke] PEFT model type: {type(model).__name__}") + + prompts = [args.prompt] + + class _SP: + max_tokens = args.max_new_tokens + temperature = 0.0 + + # First call includes prefill + any lazy engine bring-up; measure separately. + t_first0 = time.perf_counter() + outputs = model.fast_generate(prompts, sampling_params = _SP(), use_tqdm = False) + t_first = time.perf_counter() - t_first0 + out = outputs[0] + n_tok = len(out.outputs[0].token_ids) + + # Warm steady-state: run again and measure. + t_warm0 = time.perf_counter() + outputs2 = model.fast_generate(prompts, sampling_params = _SP(), use_tqdm = False) + t_warm = time.perf_counter() - t_warm0 + n_tok_warm = len(outputs2[0].outputs[0].token_ids) + + peak_after_gen = torch.cuda.max_memory_reserved() / 1024**3 + print( + f"[smoke] first call: generated {n_tok} tokens in {t_first:.2f}s " + f"({n_tok / t_first:.1f} tok/s)" + ) + print( + f"[smoke] warm call: generated {n_tok_warm} tokens in {t_warm:.2f}s " + f"({n_tok_warm / t_warm:.1f} tok/s)" + ) + print(f"[smoke] peak VRAM after gen: {peak_after_gen:.2f} GB") + print(f"[smoke] prompt: {args.prompt!r}") + print(f"[smoke] completion: {out.outputs[0].text!r}") + + precision = "4bit" if args.load_in_4bit else args.dtype + out_dir = Path(args.out_dir) + out_dir.mkdir(parents = True, exist_ok = True) + summary = { + "phase": "smoke_A", + "model": args.model, + "precision": precision, + "dtype": str(dtype), + "max_seq_length": args.max_seq_length, + "max_new_tokens": args.max_new_tokens, + "with_lora": args.with_lora, + "t_load_s": round(t_load, 2), + "peak_vram_after_load_gb": round(peak_after_load, 2), + "peak_vram_after_gen_gb": round(peak_after_gen, 2), + "first_call_s": round(t_first, 2), + "first_call_tok_s": round(n_tok / t_first, 1), + "warm_call_s": round(t_warm, 2), + "warm_call_tok_s": round(n_tok_warm / t_warm, 1), + "arch": arch, + "impl": impl, + "prompt": args.prompt, + "completion": out.outputs[0].text, + } + with open(out_dir / f"smoke_A_{precision}.json", "w") as f: + json.dump(summary, f, indent = 2) + print(f"[smoke] wrote {out_dir / f'smoke_A_{precision}.json'}") + + +if __name__ == "__main__": + main() diff --git a/tests/flex_moe_vllm_bench.py b/tests/flex_moe_vllm_bench.py new file mode 100644 index 0000000000..1e2844a331 --- /dev/null +++ b/tests/flex_moe_vllm_bench.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""vLLM direct throughput comparison for Qwen3 MoE. + +Uses ``vllm.LLM`` directly (no unsloth) to bench the same workload +as ``tests/flex_moe_bench.py`` so the flex+compile_walker numbers +have an apples-to-apples vLLM baseline. + +Usage: + CUDA_VISIBLE_DEVICES=2 python -u \\ + tests/flex_moe_vllm_bench.py --load_in_4bit +""" +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +import time +from pathlib import Path + +# DeepGEMM isn't installed in this env; vLLM's FP8 warmup crashes the +# engine-core subprocess without it. Disable before importing vllm. +os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0") + +_REPO_ROOT = Path(__file__).resolve().parents[1] +if str(_REPO_ROOT) not in sys.path: + sys.path.insert(0, str(_REPO_ROOT)) + + +def _gpu_mem_used_gb() -> float: + try: + out = subprocess.run( + ["nvidia-smi", "--query-gpu=memory.used", + "--format=csv,noheader,nounits", "-i", "0"], + capture_output=True, text=True, timeout=5, + ) + return int(out.stdout.strip().splitlines()[0]) / 1024 + except Exception: + return 0.0 + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--model", default="unsloth/Qwen3-30B-A3B-Instruct-2507") + p.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16") + p.add_argument("--load_in_4bit", action="store_true", + help="use the -bnb-4bit checkpoint variant") + p.add_argument("--n_prompts", type=int, default=8) + p.add_argument("--max_new_tokens", type=int, default=64) + p.add_argument("--max_model_len", type=int, default=1024) + p.add_argument("--warmup_rounds", type=int, default=1) + p.add_argument("--timed_rounds", type=int, default=2) + p.add_argument("--out_dir", default="async_task_outputs/qwen3_moe_grpo_bench_v2") + p.add_argument("--gpu_memory_utilization", type=float, default=0.85) + p.add_argument("--enable_lora", action="store_true", + help="enable vLLM LoRA serving; requires --lora_path") + p.add_argument("--lora_path", default=None, + help="path to a LoRA adapter dir to load per-request") + p.add_argument("--max_lora_rank", type=int, default=16) + p.add_argument("--enforce_eager", action="store_true", + help="disable vLLM CUDA graph capture (diagnostic)") + p.add_argument("--tag", default=None, + help="extra label for the output json filename") + p.add_argument("--chat_template", action="store_true", + help="wrap each prompt with the tokenizer's chat " + "template (needed for chat-tuned models like " + "gpt-oss that produce gibberish on raw strings)") + p.add_argument("--user_prompt", + default="Continue this sentence: The quick brown fox jumps over fence {i}, then", + help="user-message template when --chat_template is set " + "(``{i}`` gets replaced with the prompt index)") + args = p.parse_args() + + import torch + from vllm import LLM, SamplingParams + import vllm + print(f"[vllm] version: {vllm.__version__}") + + dtype_str = "bfloat16" if args.dtype == "bf16" else "float16" + + # For 4bit, tell vLLM to apply bnb quantization on-the-fly to the + # bf16 checkpoint (there's no pre-quantized unsloth 4bit variant for + # this model on HF). For FP8, point at the FP8 variant directly. + model_id = args.model + quantization = None + if args.load_in_4bit: + quantization = "bitsandbytes" + print(f"[vllm] 4bit: model_id={model_id} quantization=bitsandbytes (on-the-fly)") + elif os.environ.get("USE_FP8") == "1": + model_id = args.model + "-FP8" + quantization = "fp8" + print(f"[vllm] fp8: model_id={model_id}") + + llm_kwargs = dict( + model=model_id, + dtype=dtype_str, + quantization=quantization, + gpu_memory_utilization=args.gpu_memory_utilization, + max_model_len=args.max_model_len, + max_num_seqs=max(args.n_prompts, 8), + enforce_eager=args.enforce_eager, + trust_remote_code=False, + ) + if args.enable_lora: + llm_kwargs.update( + enable_lora=True, + max_lora_rank=args.max_lora_rank, + max_loras=1, + ) + + torch.cuda.reset_peak_memory_stats() + t0 = time.perf_counter() + llm = LLM(**llm_kwargs) + t_load = time.perf_counter() - t0 + # vLLM uses its own allocator; torch's reserved bytes miss it. + # Query nvidia-smi instead (visible GPU only — i.e. whatever + # CUDA_VISIBLE_DEVICES exposes as device 0). + peak_load = _gpu_mem_used_gb() + print(f"[vllm] loaded in {t_load:.1f}s peak {peak_load:.1f} GB") + + # Use vLLM's native ``llm.chat`` when chat-template is requested — + # it threads harmony-aware templating + tokenization through vLLM's + # chat_utils, which is what gpt-oss needs (hand-rolled + # apply_chat_template + raw ``llm.generate`` was producing gibberish + # on gpt-oss-20b-BF16 no-LoRA). + if args.chat_template: + chat_messages = [ + [{"role": "user", "content": args.user_prompt.format(i=i)}] + for i in range(args.n_prompts) + ] + print(f"[vllm] chat mode: llm.chat() with {args.n_prompts} single-turn " + f"messages. first: {chat_messages[0][0]['content']!r}") + prompts = None + else: + prompts = [f"The quick brown fox jumps over fence {i}, then" + for i in range(args.n_prompts)] + chat_messages = None + + sp = SamplingParams(max_tokens=args.max_new_tokens, temperature=0.0) + + gen_kwargs = {} + if args.enable_lora and args.lora_path is not None: + from vllm.lora.request import LoRARequest + gen_kwargs["lora_request"] = LoRARequest( + lora_name="flex_lora", + lora_int_id=1, + lora_path=args.lora_path, + ) + print(f"[vllm] LoRA enabled: path={args.lora_path} rank<={args.max_lora_rank}") + + def _run_once(): + if chat_messages is not None: + return llm.chat(chat_messages, sampling_params=sp, + use_tqdm=False, **gen_kwargs) + return llm.generate(prompts, sampling_params=sp, + use_tqdm=False, **gen_kwargs) + + # Warmup. + for _ in range(args.warmup_rounds): + _ = _run_once() + + # Timed. + wall = [] + tok_counts = [] + for _ in range(args.timed_rounds): + torch.cuda.synchronize() + t0 = time.perf_counter() + outs = _run_once() + torch.cuda.synchronize() + wall.append(time.perf_counter() - t0) + tok_counts.append(sum(len(o.outputs[0].token_ids) for o in outs)) + + peak_gen = _gpu_mem_used_gb() + med_wall = sorted(wall)[len(wall) // 2] + med_tok = tok_counts[len(wall) // 2] + tps = med_tok / med_wall if med_wall > 0 else 0.0 + + sample_text = outs[0].outputs[0].text if outs else "" + + print(f"[vllm] wall: {wall}") + print(f"[vllm] tok counts: {tok_counts}") + print(f"[vllm] median wall: {med_wall:.3f}s median tok/s: {tps:.1f}") + print(f"[vllm] peak VRAM after gen: {peak_gen:.1f} GB") + print(f"[vllm] sample completion[0]: {sample_text[:200]!r}") + + precision = "4bit" if args.load_in_4bit else args.dtype + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + summary = { + "phase": "bench_decode", + "backend": "vllm", + "vllm_version": vllm.__version__, + "model": model_id, + "precision": precision, + "n_prompts": args.n_prompts, + "max_new_tokens": args.max_new_tokens, + "wall_times_s": wall, + "tok_counts": tok_counts, + "median_wall_s": round(med_wall, 3), + "median_tok_s": round(tps, 1), + "peak_vram_load_gb": round(peak_load, 2), + "peak_vram_after_gen_gb": round(peak_gen, 2), + "t_load_s": round(t_load, 1), + "sample_completion": sample_text[:500], + } + suffix = f"_{args.tag}" if args.tag else "" + out_path = out_dir / f"bench_decode_vllm_{precision}{suffix}.json" + with open(out_path, "w") as f: + json.dump(summary, f, indent=2) + print(f"[vllm] wrote {out_path}") + + +if __name__ == "__main__": + main() diff --git a/tests/flex_sleep_mode_smoke.py b/tests/flex_sleep_mode_smoke.py new file mode 100644 index 0000000000..d206eae773 --- /dev/null +++ b/tests/flex_sleep_mode_smoke.py @@ -0,0 +1,297 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""Smoke-test :meth:`FlexEngine.sleep` and :meth:`FlexEngine.wake_up`. + +Invoked as: + CUDA_VISIBLE_DEVICES=3 \ + UNSLOTH_FAST_INFERENCE=1 UNSLOTH_VLLM_STANDBY=1 \ + python tests/flex_sleep_mode_smoke.py --model unsloth/Qwen3-4B-Base + +What it checks: + 1. ``model.vllm_engine._sleep_mode_enabled`` is True when vLLM is + importable. + 2. Captured CUDA graphs survive a sleep/wake round-trip: a second + ``fast_generate`` call after ``sleep`` -> ``wake_up`` does not + re-capture and still produces the same token ids. + 3. ``torch.cuda.memory_allocated()`` drops on sleep and returns close + to the pre-sleep value on wake. With ``--no-standby``, memory + should be identical across the three probes (no-op path). + 4. Hardens the regression guard by running the same workflow with + ``UNSLOTH_VLLM_STANDBY`` unset (``--no-standby``); ``sleep`` / + ``wake_up`` must be exact no-ops. + +This is a smoke test, not the full verification matrix from the plan +(that lives under ``scripts/benchmarks``). +""" + +from __future__ import annotations + +import argparse +import gc +import os +import sys +import 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)) + + +def _gb(n: int) -> float: + return round(n / 1e9, 3) + + +def _probe(label: str) -> dict: + import torch + + gc.collect() + torch.cuda.synchronize() + # ``memory_allocated`` / ``memory_reserved`` only track torch's + # caching allocator and ignore cuMem-backed pools, so they do NOT + # drop on sleep even though cuMem has unmapped the pages. + # ``mem_get_info()`` asks the CUDA runtime directly, so it sees + # cuMem unmaps and is the right probe for sleep / wake verification. + free_bytes, total_bytes = torch.cuda.mem_get_info() + return { + "label": label, + "allocated_gb": _gb(torch.cuda.memory_allocated()), + "reserved_gb": _gb(torch.cuda.memory_reserved()), + "cuda_free_gb": _gb(free_bytes), + "cuda_used_gb": _gb(total_bytes - free_bytes), + } + + +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("--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") + p.add_argument( + "--no-standby", + action = "store_true", + help = "Force UNSLOTH_VLLM_STANDBY=0 to validate the no-op regression path.", + ) + p.add_argument( + "--cycles", + type = int, + default = 1, + help = "Number of sleep / wake / generate cycles after warmup. " + ">1 exercises the repeated-cycle regression (run #6).", + ) + args = p.parse_args() + + os.environ.setdefault("UNSLOTH_FAST_INFERENCE", "1") + if args.no_standby: + os.environ["UNSLOTH_VLLM_STANDBY"] = "0" + else: + os.environ.setdefault("UNSLOTH_VLLM_STANDBY", "1") + standby = os.environ.get("UNSLOTH_VLLM_STANDBY", "0") == "1" + print( + f"[sleep-smoke] UNSLOTH_FAST_INFERENCE=" + f"{os.environ.get('UNSLOTH_FAST_INFERENCE')} " + f"UNSLOTH_VLLM_STANDBY={os.environ.get('UNSLOTH_VLLM_STANDBY')}" + ) + + import torch + + import unsloth + from unsloth import FastLanguageModel + + print(f"[sleep-smoke] unsloth={unsloth.__file__}") + + 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, + ) + print( + f"[sleep-smoke] loaded {args.model} in " + f"{time.perf_counter() - t0:.1f}s; dtype={model.dtype}" + ) + + engine = model.vllm_engine + print(f"[sleep-smoke] engine type: {type(engine).__name__}") + sleep_enabled = getattr(engine, "_sleep_mode_enabled", None) + print(f"[sleep-smoke] engine._sleep_mode_enabled: {sleep_enabled}") + if standby: + if sleep_enabled is not True: + # Most likely vLLM is not importable in this environment. + print( + "[sleep-smoke] WARNING: UNSLOTH_VLLM_STANDBY=1 was set " + "but engine._sleep_mode_enabled is False (vLLM missing?)" + ) + else: + assert sleep_enabled is False, ( + f"Expected sleep mode to be disabled with UNSLOTH_VLLM_STANDBY=0, " + f"got {sleep_enabled}" + ) + + ll_cfg = engine.llm_engine.vllm_config.model_config + print( + f"[sleep-smoke] llm_engine.vllm_config.model_config.enable_sleep_mode: " + f"{getattr(ll_cfg, 'enable_sleep_mode', None)}" + ) + assert getattr(ll_cfg, "enable_sleep_mode", None) == bool(sleep_enabled), ( + "_LLMEngineStub.model_config.enable_sleep_mode must mirror " + "engine._sleep_mode_enabled" + ) + + from unsloth.inference.vllm_shim import LoRARequest # noqa: F401 + + prompts = [args.prompt] + + # ----- warmup ----- + t0 = time.perf_counter() + out1 = engine.generate( + prompts, + sampling_params = type( + "SP", + (), + {"max_tokens": args.max_new_tokens, "temperature": 0.0}, + )(), + ) + print( + f"[sleep-smoke] warmup generate: {time.perf_counter() - t0:.2f}s; " + f"tok_ids[:10]={out1[0].outputs[0].token_ids[:10]}" + ) + pre_tokens = list(out1[0].outputs[0].token_ids) + + probe_pre = _probe("pre-sleep") + print(f"[sleep-smoke] {probe_pre}") + + # Diagnostic: checksum the inference-model weights so we can detect + # if cuMem's sleep/wake round-trip corrupts any parameter. + def _checksum_params(mod, limit = 16): + import torch as _t + + out = [] + for i, (name, p) in enumerate(mod.named_parameters()): + if i >= limit: + break + t = p.detach() + out.append((name, list(t.shape), float(t.float().abs().sum().item()))) + return out + + pre_sums = _checksum_params(engine._inference_model) + print("[sleep-smoke] pre-sleep first-16 param |sum|:") + for n, s, v in pre_sums: + print(f" {n} {s} {v:.4f}") + + # Repeated sleep / wake / generate cycle test (plan matrix run #6). + # Each cycle validates that the engine does not drift: tokens remain + # bitwise identical, memory returns to baseline, weights round-trip + # cleanly. A bug that only surfaces on the second or third cycle + # (stale Python state, double-wake, leaked handles) fails here. + for cycle in range(args.cycles): + if args.cycles > 1: + print(f"[sleep-smoke] --- cycle {cycle + 1}/{args.cycles} ---") + + # ----- sleep ----- + t0 = time.perf_counter() + engine.sleep(level = 1) + t_sleep = time.perf_counter() - t0 + probe_post = _probe(f"post-sleep[{cycle + 1}]") + print(f"[sleep-smoke] sleep(level=1) took {t_sleep:.3f}s; {probe_post}") + + if sleep_enabled: + drop = probe_pre["cuda_used_gb"] - probe_post["cuda_used_gb"] + print( + f"[sleep-smoke] process-level VRAM drop on sleep: " + f"{drop:+.3f} GB (cuMem-managed; not visible in " + f"torch.memory_allocated)" + ) + if not getattr(engine, "_single_copy_mode", False): + # 16-bit path: both weights + kv_cache pools are dropped. + assert drop >= 1.0, ( + f"Expected multi-GB drop in process-level VRAM on " + f"sleep(level=1); got {drop:+.3f} GB" + ) + else: + # 4-bit single-copy: only KV cache drops; weights stay. + assert drop > 0.0, ( + f"Expected KV-cache drop in process-level VRAM on " + f"sleep(level=1); got {drop:+.3f} GB" + ) + else: + # With sleep mode off, the sleep() call must not free VRAM. + # Process-level jitter is allowed (shared GPU); torch-owned + # allocations must be untouched. + assert probe_post["allocated_gb"] == probe_pre["allocated_gb"], ( + "With sleep mode disabled, torch.memory_allocated must " + "be unchanged by sleep()" + ) + + # ----- wake ----- + t0 = time.perf_counter() + engine.wake_up() + t_wake = time.perf_counter() - t0 + probe_wake = _probe(f"post-wake[{cycle + 1}]") + print(f"[sleep-smoke] wake_up() took {t_wake:.3f}s; {probe_wake}") + + post_sums = _checksum_params(engine._inference_model) + diffs = [] + for (n1, s1, v1), (n2, s2, v2) in zip(pre_sums, post_sums): + delta = abs(v1 - v2) + if delta > 0.0: + diffs.append((n1, v1, v2, delta)) + print( + f"[sleep-smoke] post-wake weight diff: " + f"{len(diffs)}/{len(pre_sums)} params changed " + f"(bitwise-exact restore expected)" + ) + for n, v1, v2, d in diffs[:8]: + print(f" diff {n}: pre={v1:.4f} post={v2:.4f} delta={d:.4f}") + assert len(diffs) == 0, ( + f"Weight corruption on sleep / wake (cycle {cycle + 1}): " + f"{len(diffs)}/{len(pre_sums)} first-layer params changed" + ) + + # ----- verify we can still generate ----- + t0 = time.perf_counter() + out2 = engine.generate( + prompts, + sampling_params = type( + "SP", + (), + {"max_tokens": args.max_new_tokens, "temperature": 0.0}, + )(), + ) + t_regen = time.perf_counter() - t0 + post_tokens = list(out2[0].outputs[0].token_ids) + match = pre_tokens == post_tokens + print( + f"[sleep-smoke] post-wake generate: {t_regen:.2f}s; " + f"tok_ids[:10]={post_tokens[:10]}; matches_pre={match}" + ) + assert match, ( + f"Pre-sleep / post-wake token ids must match exactly " + f"(cycle {cycle + 1}).\n" + f"Pre: {pre_tokens}\nPost: {post_tokens}" + ) + + if sleep_enabled: + delta = probe_wake["cuda_used_gb"] - probe_pre["cuda_used_gb"] + print( + f"[sleep-smoke] post-wake vs pre-sleep process-level " + f"VRAM delta: {delta:+.3f} GB (tolerance: +/- 1.5 GB)" + ) + assert abs(delta) < 1.5, ( + f"Post-wake VRAM diverged from pre-sleep " + f"(cycle {cycle + 1}): delta={delta:+.3f} GB" + ) + + print(f"[sleep-smoke] PASS ({args.cycles} cycle(s))") + + +if __name__ == "__main__": + main() diff --git a/tests/gemma4_fast_inference_parity.py b/tests/gemma4_fast_inference_parity.py new file mode 100644 index 0000000000..03e3b9913c --- /dev/null +++ b/tests/gemma4_fast_inference_parity.py @@ -0,0 +1,141 @@ +"""Greedy parity: FastLanguageModel(fast_inference=True) vs HF naive on Gemma 4. + +Covers all four Gemma 4 variants (E2B / E4B dense, 31B dense, 26B-A4B MoE) via +``--model``. Runs HF FIRST, frees the GPU, then runs vLLM — loading vLLM first +in the same process leaves global state (allocator, compile cache, patched +functions) that subtly perturbs a later plain HF run, producing false +divergences. HF-first ordering gives bitwise matches. + +Requires vLLM nightly (>= 2026-04-17 for `vllm#39291` Gemma 4 LoRA) plus the +`unsloth-zoo#603` vLLM Gemma 4 patches. + +Example: + CUDA_VISIBLE_DEVICES=0 python -u tests/gemma4_fast_inference_parity.py \\ + --model unsloth/gemma-4-26b-a4b-it --max_new_tokens 32 +""" +import argparse +import gc +import json +import os + +os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0") +os.environ.setdefault("UNSLOTH_MOE_BACKEND", "grouped_mm") + +# Import unsloth at module scope so its patches apply to the HF run as well — +# users always import unsloth before touching transformers in practice. +import unsloth # noqa: F401, E402 +import torch # noqa: E402 + + +def _render(tok, raw_prompts): + return [ + tok.apply_chat_template( + [{"role": "user", "content": p}], + tokenize=False, add_generation_prompt=True, + ) + for p in raw_prompts + ] + + +def run_hf_first(model_name, raw_prompts, max_new_tokens): + from transformers import AutoModelForImageTextToText, AutoProcessor + proc = AutoProcessor.from_pretrained(model_name) + tok = proc.tokenizer if hasattr(proc, "tokenizer") else proc + prompts = _render(tok, raw_prompts) + + model = AutoModelForImageTextToText.from_pretrained( + model_name, dtype=torch.bfloat16, + attn_implementation="sdpa", device_map="cuda:0", + ) + model.eval() + ids, texts = [], [] + for p in prompts: + enc = tok(p, return_tensors="pt").to("cuda:0") + with torch.inference_mode(): + gen = model.generate( + **enc, max_new_tokens=max_new_tokens, + do_sample=False, temperature=None, top_p=None, + ) + new = gen[0, enc.input_ids.shape[1]:] + ids.append(new.tolist()) + texts.append(tok.decode(new, skip_special_tokens=True)) + + del model, proc + gc.collect() + torch.cuda.empty_cache() + return ids, texts, prompts + + +def run_vllm_second(model_name, prompts, max_new_tokens): + from unsloth import FastLanguageModel + from vllm import SamplingParams + + model, _ = FastLanguageModel.from_pretrained( + model_name=model_name, max_seq_length=1024, dtype=torch.bfloat16, + load_in_4bit=False, fast_inference=True, max_batch_size=8, + gpu_memory_utilization=0.6, max_lora_rank=16, + ) + sp = SamplingParams(max_tokens=max_new_tokens, temperature=0.0) + outs = model.fast_generate(prompts, sampling_params=sp, use_tqdm=False) + ids = [list(o.outputs[0].token_ids) for o in outs] + texts = [o.outputs[0].text for o in outs] + return ids, texts + + +DEFAULT_PROMPTS = [ + "In one sentence, what is Paris?", + "What is 23 + 19? Answer in one word.", + "Continue this phrase: The quick brown fox jumps over", +] + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model", default="unsloth/gemma-4-26b-a4b-it") + ap.add_argument("--max_new_tokens", type=int, default=32) + ap.add_argument("--json_out", default=None) + args = ap.parse_args() + + print(f"=== HF first ({args.model}) ===") + hf_ids, hf_texts, prompts = run_hf_first( + args.model, DEFAULT_PROMPTS, args.max_new_tokens, + ) + for i, t in enumerate(hf_texts): + print(f"[hf] P{i}: {t[:120]!r}") + + print("\n=== vLLM second ===") + fast_ids, fast_texts = run_vllm_second( + args.model, prompts, args.max_new_tokens, + ) + for i, t in enumerate(fast_texts): + print(f"[fast] P{i}: {t[:120]!r}") + + print("\n=== Match ===") + total, matched, rows = 0, 0, [] + for i, (a, b) in enumerate(zip(fast_ids, hf_ids)): + n = min(len(a), len(b)) + m = sum(1 for x, y in zip(a[:n], b[:n]) if x == y) + total += n + matched += m + print(f"[P{i}] {m}/{n}") + rows.append({ + "prompt_id": i, "match": m, "total": n, + "fast": fast_texts[i], "hf": hf_texts[i], + }) + print(f"\nTotal {matched}/{total}") + + if args.json_out: + os.makedirs(os.path.dirname(args.json_out) or ".", exist_ok=True) + with open(args.json_out, "w") as f: + json.dump({ + "model": args.model, "rows": rows, + "matched": matched, "total": total, + }, f, indent=2) + + # Exit non-zero on any divergence so CI can gate on bitwise parity. + if matched != total: + raise SystemExit(f"Divergence: {matched}/{total}") + + +if __name__ == "__main__": + main() diff --git a/tests/gemma4_flex_bench.py b/tests/gemma4_flex_bench.py new file mode 100644 index 0000000000..f896005263 --- /dev/null +++ b/tests/gemma4_flex_bench.py @@ -0,0 +1,132 @@ +"""Throughput bench for Unsloth flex fast-inference on Gemma 4. + +Runs ``FastLanguageModel(fast_inference=True)`` through the flex engine +(``UNSLOTH_FAST_INFERENCE=1``) and measures tok/s at a list of batch sizes. +Captures CUDA graphs at each bucket; the bench reports the post-warmup pass. + +Example: + CUDA_VISIBLE_DEVICES=0 UNSLOTH_FAST_INFERENCE=1 \\ + UNSLOTH_MOE_BACKEND=grouped_mm \\ + python -u tests/gemma4_flex_bench.py \\ + --model unsloth/gemma-4-26b-a4b-it --batch_sizes 1 4 8 16 \\ + --max_new_tokens 64 --json_out async_task_outputs/flex_bench_26b.json + +Pair with tests/gemma4_fast_inference_parity.py for HF-naive parity and +with tests/gemma4_fast_bench.py (vLLM nightly path) for cross-engine +throughput comparison. +""" +import argparse +import json +import os +import time + +os.environ.setdefault("UNSLOTH_FAST_INFERENCE", "1") +os.environ.setdefault("UNSLOTH_MOE_BACKEND", "grouped_mm") +os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0") + +import unsloth # noqa: F401, E402 +import torch # noqa: E402 + + +SHORT_PROMPT = "In one sentence, what is Paris?" +LONG_POOL = [ + "In one sentence, what is Paris?", + "What is 23 + 19? Answer in one word.", + "Continue this phrase: The quick brown fox jumps over", + "Name three primary colors.", + "Who wrote the play Hamlet?", + "What is the capital of Japan?", + "What does the acronym 'NASA' stand for?", + "Explain the water cycle in one sentence.", + "Give a two-word summary of the French Revolution.", + "What is the tallest mountain on Earth?", + "Define 'entropy' in one sentence.", + "Who painted the Mona Lisa?", + "What is the boiling point of water in Celsius?", + "Name one prime number larger than 10.", + "What color do you get when you mix blue and yellow?", + "Give one example of an amphibian.", +] + + +def _render(tok, prompts): + return [ + tok.apply_chat_template( + [{"role": "user", "content": p}], + tokenize=False, add_generation_prompt=True, + ) + for p in prompts + ] + + +class _SP: + def __init__(self, n): + self.max_tokens = n + self.temperature = 0.0 + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model", required=True) + ap.add_argument("--batch_sizes", nargs="+", type=int, default=[1, 4, 8, 16]) + ap.add_argument("--max_new_tokens", type=int, default=64) + ap.add_argument("--json_out", required=True) + ap.add_argument("--max_seq_length", type=int, default=1024) + args = ap.parse_args() + + from unsloth import FastLanguageModel + + t_load = time.perf_counter() + model, tok_raw = FastLanguageModel.from_pretrained( + model_name=args.model, max_seq_length=args.max_seq_length, + dtype=torch.bfloat16, load_in_4bit=False, fast_inference=True, + max_batch_size=max(args.batch_sizes), + gpu_memory_utilization=0.6, + ) + t_load = time.perf_counter() - t_load + tok = tok_raw.tokenizer if hasattr(tok_raw, "tokenizer") else tok_raw + print(f"[flex] load {t_load:.1f}s") + + sp = _SP(args.max_new_tokens) + results = [] + + for bs in args.batch_sizes: + prompts = _render(tok, LONG_POOL[:bs]) + # Warmup: primes the cudagraph bucket and walker compile. + _ = model.fast_generate(prompts, sampling_params=sp, use_tqdm=False) + t0 = time.perf_counter() + outs = model.fast_generate(prompts, sampling_params=sp, use_tqdm=False) + dt = time.perf_counter() - t0 + total = sum(len(o.outputs[0].token_ids) for o in outs) + tok_s = total / dt if dt > 0 else 0.0 + results.append({ + "batch_size": bs, "gen_s": round(dt, 3), + "total_tokens": total, "tok_s": round(tok_s, 1), + }) + print(f"[flex bs={bs}] {total} tok in {dt:.2f}s -> {tok_s:.1f} tok/s") + + # Parity sanity: run SHORT_PROMPT at bs=1, save ids. + p = _render(tok, [SHORT_PROMPT]) + outs = model.fast_generate(p, sampling_params=sp, use_tqdm=False) + parity_ids = list(outs[0].outputs[0].token_ids) + parity_text = outs[0].outputs[0].text + + peak_gb = torch.cuda.max_memory_allocated() / 1024**3 + + summary = { + "model": args.model, + "max_new_tokens": args.max_new_tokens, + "load_s": round(t_load, 2), + "peak_vram_gb": round(peak_gb, 2), + "batches": results, + "parity_ids": parity_ids, + "parity_text": parity_text, + } + os.makedirs(os.path.dirname(args.json_out) or ".", exist_ok=True) + with open(args.json_out, "w") as f: + json.dump(summary, f, indent=2) + print(f"[flex] wrote {args.json_out}") + + +if __name__ == "__main__": + main() diff --git a/tests/qwen3_5_flex_parity_bs4.py b/tests/qwen3_5_flex_parity_bs4.py new file mode 100644 index 0000000000..75a208fb75 --- /dev/null +++ b/tests/qwen3_5_flex_parity_bs4.py @@ -0,0 +1,124 @@ +"""B=4 parity test so the captured graph bucket actually replays. + +Compares flex captured-graph decode vs flex eager-decode (capture +disabled via ``UNSLOTH_FLEX_QWEN3_5_NO_CAPTURE=1``). Exact token-for-token +match is the success criterion — both paths use identical math, only +the kernel-launch wrapper differs. + + CUDA_VISIBLE_DEVICES=2 UNSLOTH_FAST_INFERENCE=1 PYTHONUNBUFFERED=1 \\ + python -u tests/qwen3_5_flex_parity_bs4.py \\ + --model Qwen/Qwen3.5-4B --max_new_tokens 16 \\ + --json_out async_task_outputs/qwen3_5/flex_parity_bs4.json +""" +import argparse +import json +import os +import time + + +CHAT_PROMPTS = [ + "In one sentence, what is Paris?", + "What is 23 + 19? Answer in one word.", + "Continue: The quick brown fox jumps over", + "Name one primary color.", +] + + +def _run(model, tok, args, capture: bool): + class SP: + def __init__(self, n): self.max_tokens = n; self.temperature = 0.0 + + sp = SP(args.max_new_tokens) + rendered = [ + tok.apply_chat_template( + [{"role": "user", "content": p}], + tokenize=False, add_generation_prompt=True, + ) + for p in CHAT_PROMPTS + ] + os.environ["UNSLOTH_FLEX_QWEN3_5_NO_CAPTURE"] = "0" if capture else "1" + t0 = time.perf_counter() + outs = model.fast_generate(rendered, sampling_params=sp, use_tqdm=False) + dt = time.perf_counter() - t0 + results = [] + for p, o in zip(CHAT_PROMPTS, outs): + results.append({ + "prompt": p, + "text": o.outputs[0].text, + "ids": list(o.outputs[0].token_ids), + }) + return dt, results + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--model", required=True) + ap.add_argument("--max_new_tokens", type=int, default=16) + ap.add_argument("--max_seq_length", type=int, default=1024) + ap.add_argument("--max_batch_size", type=int, default=4) + ap.add_argument("--gpu_memory_utilization", type=float, default=0.6) + ap.add_argument("--json_out", required=True) + args = ap.parse_args() + + import torch + os.environ["UNSLOTH_FAST_INFERENCE"] = "1" + import unsloth # noqa + from unsloth import FastLanguageModel + + model, tok_raw = FastLanguageModel.from_pretrained( + model_name=args.model, + max_seq_length=args.max_seq_length, + dtype=torch.bfloat16, + load_in_4bit=False, + fast_inference=True, + max_batch_size=args.max_batch_size, + gpu_memory_utilization=args.gpu_memory_utilization, + ) + tok = tok_raw.tokenizer if hasattr(tok_raw, "tokenizer") else tok_raw + + # Eager first (so the graph capture doesn't pollute state). + dt_eager, eager = _run(model, tok, args, capture=False) + print(f"[eager] gen_s={dt_eager:.2f}") + for r in eager: + print(f"[eager] {r['prompt'][:40]!r} -> {r['text'][:80]!r}") + + dt_cap, cap = _run(model, tok, args, capture=True) + print(f"[cap] gen_s={dt_cap:.2f}") + for r in cap: + print(f"[cap] {r['prompt'][:40]!r} -> {r['text'][:80]!r}") + + # Token-by-token parity. + mismatches = [] + for i, (e, c) in enumerate(zip(eager, cap)): + matches = 0 + for a, b in zip(e["ids"], c["ids"]): + if a == b: + matches += 1 + else: + break + mismatches.append({ + "prompt_idx": i, + "prompt": e["prompt"], + "matches": matches, + "total": min(len(e["ids"]), len(c["ids"])), + "eager_text": e["text"], + "cap_text": c["text"], + }) + tag = "OK " if matches == min(len(e["ids"]), len(c["ids"])) else "MM " + print(f"[{tag}] P{i} {matches}/{min(len(e['ids']), len(c['ids']))}") + + summary = { + "model": args.model, + "max_new_tokens": args.max_new_tokens, + "eager_gen_s": round(dt_eager, 3), + "capture_gen_s": round(dt_cap, 3), + "mismatches": mismatches, + } + os.makedirs(os.path.dirname(args.json_out) or ".", exist_ok=True) + with open(args.json_out, "w") as f: + json.dump(summary, f, indent=2) + print(f"[parity-bs4] wrote {args.json_out}") + + +if __name__ == "__main__": + main() diff --git a/unsloth/__init__.py b/unsloth/__init__.py index 73a9ce8fef..52114bb544 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -137,6 +137,7 @@ from .import_fixes import ( fix_vllm_aimv2_issue, check_vllm_torch_sm100_compatibility, fix_vllm_guided_decoding_params, + fix_trl_vllm_ascend, fix_vllm_pdl_blackwell, fix_triton_compiled_kernel_missing_attrs, patch_trunc_normal_precision_issue, @@ -159,6 +160,7 @@ fix_vllm_aimv2_issue() # Check vLLM + torch < 2.9.0 + SM100 compatibility BEFORE importing vLLM check_vllm_torch_sm100_compatibility() fix_vllm_guided_decoding_params() +fix_trl_vllm_ascend() fix_vllm_pdl_blackwell() fix_triton_compiled_kernel_missing_attrs() patch_trunc_normal_precision_issue() @@ -179,6 +181,7 @@ del fix_xformers_performance_issue del fix_vllm_aimv2_issue del check_vllm_torch_sm100_compatibility del fix_vllm_guided_decoding_params +del fix_trl_vllm_ascend del fix_vllm_pdl_blackwell del fix_triton_compiled_kernel_missing_attrs del patch_trunc_normal_precision_issue diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index ca44a0ce7e..3dc8b2d184 100644 --- a/unsloth/import_fixes.py +++ b/unsloth/import_fixes.py @@ -489,6 +489,33 @@ def fix_vllm_guided_decoding_params(): ) +def fix_trl_vllm_ascend(): + # transformers >= 4.48's `_is_package_available(name)` returns a + # tuple (bool, version_or_None). TRL caches that tuple in + # module-level `_*_available` flags and the matching + # `is_*_available()` accessors return the tuple directly. A + # non-empty tuple is always truthy, so `if is_X_available():` + # fires even when X is absent, triggering an unconditional + # `import X` that fails. The surfaced case is `vllm_ascend` + # (blocks `from trl import GRPOConfig, GRPOTrainer` outside + # Huawei Ascend hosts); `llm_blender`, `deepspeed`, `joblib` + # share the same shape. Coerce every tuple-cached flag in + # trl.import_utils to bool; the existing accessors that just + # return the cached value then naturally yield a bool. + if importlib.util.find_spec("trl") is None: + return + try: + import trl.import_utils as tiu + except Exception: + return + for attr in list(vars(tiu)): + if not (attr.startswith("_") and attr.endswith("_available")): + continue + cached = getattr(tiu, attr) + if isinstance(cached, tuple): + setattr(tiu, attr, bool(cached and cached[0])) + + def ignore_logger_messages(): # Ignore Environment variable `HF_TOKEN` is set try: diff --git a/unsloth/inference/__init__.py b/unsloth/inference/__init__.py new file mode 100644 index 0000000000..67063f7aef --- /dev/null +++ b/unsloth/inference/__init__.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""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). + +Four architectures are supported today: Qwen3 (dense), Qwen3-MoE, +Llama-3, Gemma-4-E2B-it. Anything else raises +:class:`NotImplementedError`; unset the env var or use vLLM instead.""" + +from .flex_engine import ( + FlexEngine, + build_flex_engine, + install_flex_sentinel, + load_flex, +) +from .flex_moe import FlexMoEInference +from .flex_gemma4_moe import FlexGemma4MoEInference +from .vllm_shim import ( + CompletionOutput, + LoRARequest, + RequestOutput, + load_lora, + save_lora, +) + +__all__ = [ + "FlexEngine", + "FlexMoEInference", + "FlexGemma4MoEInference", + "load_flex", + "build_flex_engine", + "install_flex_sentinel", + "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..43c7cb0303 --- /dev/null +++ b/unsloth/inference/flex_engine.py @@ -0,0 +1,1118 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""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 gc +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 .flex_moe import FlexMoEInference +from .sleep_mode import ( + _get_cumem_allocator, + kv_cache_pool, + sleep_mode_enabled, + weight_pool, +) +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"``, ``"gemma4_moe"``, ``"qwen3_moe"``, ``"qwen3"``, ``"gpt_oss"``, ``"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: + # Gemma 4 26B-A4B carries ``Gemma4TextConfig.num_experts > 1`` — the + # text class name is the same for dense/MoE variants, so we distinguish + # on the config. Dense E2B/E4B/31B fall through to ``"gemma4"``. + cfg = getattr(hf_model, "config", None) + text_cfg = getattr(cfg, "text_config", None) if cfg is not None else None + # Some PEFT / shell wrappers surface the text config directly. + if text_cfg is None: + text_cfg = cfg + num_experts = getattr(text_cfg, "num_experts", 0) if text_cfg is not None else 0 + if num_experts and num_experts > 1: + return "gemma4_moe" + return "gemma4" + # Check MoE before dense; ``Qwen3MoeForCausalLM`` contains both + # ``"qwen3moe"`` and ``"qwen3"`` substrings. + # Same for Qwen3.5 / Qwen3.6 — check qwen3_5 variants BEFORE + # plain qwen3 so ``Qwen3_5MoeForConditionalGeneration`` doesn't get + # dispatched to the Qwen3 backend. + if "qwen3_5moe" in lowered or "qwen3_5_moe" in lowered or "qwen35moe" in lowered: + return "qwen3_5_moe" + if "qwen3_5" in lowered or "qwen35" in lowered: + return "qwen3_5" + if "qwen3moe" in lowered or "qwen3_moe" in lowered: + return "qwen3_moe" + if "qwen3" in lowered: + return "qwen3" + if "gptoss" in lowered or "gpt_oss" in lowered: + return "gpt_oss" + if "llama" in lowered: + return "llama3" + raise NotImplementedError( + "UNSLOTH_FAST_INFERENCE=1 only supports Qwen3, Qwen3-MoE, " + "Qwen3.5 / Qwen3.6 (dense + MoE), gpt-oss, Llama-3, Gemma-4 " + f"(dense + MoE) 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, sleep_enabled: bool = False): + self.vllm_config = types.SimpleNamespace( + lora_config = types.SimpleNamespace(), + model_config = types.SimpleNamespace( + enable_sleep_mode = bool(sleep_enabled), + ), + ) + 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 + + # Sleep-mode setup. When ``UNSLOTH_VLLM_STANDBY=1`` is set AND + # vLLM is importable, we route the engine's heavy allocations + # (the inference deep-copies + per-layer PagedKVCache buffers) + # through cuMem-backed pools so ``FlexEngine.sleep`` can offload + # weights to pinned CPU memory without destroying captured CUDA + # graphs. The allocator assigns stable GPU virtual addresses, so + # unmapping on sleep and re-mapping on wake preserves pointer + # validity. ``expandable_segments:True`` on + # ``PYTORCH_CUDA_ALLOC_CONF`` is incompatible with cuMem; if the + # user has it set, ``_get_cumem_allocator`` returns None and + # sleep stays a no-op. + self._sleep_mode_enabled = sleep_mode_enabled() + self._cumem_allocator = ( + _get_cumem_allocator() if self._sleep_mode_enabled else None + ) + + # 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: + with weight_pool(self._cumem_allocator): + 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 + # ``_inference_model is hf_model`` is the 4-bit / single-copy + # fallback path (see ``bind_peft_model``); flex skips the second + # deep-copy there because bnb-4bit packed weights can't be + # in-place refreshed. In that mode there is no CPU-backup step + # to do on ``sleep`` — only the KV cache gets dropped. + self._single_copy_mode = inference_model is hf_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 in ("gemma4", "gemma4_moe"): + inference_model = _extract_gemma4_text_shell(inference_model) + self._inference_model = inference_model + if arch == "gemma4": + Impl = FlexGemma4Inference + elif arch == "gemma4_moe": + from .flex_gemma4_moe import FlexGemma4MoEInference + Impl = FlexGemma4MoEInference + elif arch == "qwen3_moe": + Impl = FlexMoEInference + # CUDA graph capture is supported on the ``grouped_mm`` MoE + # backend only. ``FlexMoEInference.capture_decode_cudagraph`` + # re-checks the active backend at capture time and skips + # capture on any other backend, leaving ``capture_cudagraph`` + # alone here. + elif arch == "gpt_oss": + from .flex_gpt_oss import FlexGptOssInference + Impl = FlexGptOssInference + elif arch in ("qwen3_5", "qwen3_5_moe"): + from .flex_qwen3_5 import FlexQwen3_5Inference + Impl = FlexQwen3_5Inference + else: + Impl = FlexInference + # Pass the cuMem allocator through so the impl can wrap ONLY + # the paged-KV allocations (``PageTable`` + per-layer + # ``PagedKVCache``) in the ``kv_cache`` pool. Everything else + # the impl creates (``input_pos_buffer``, ``block_mask_logical``, + # captured CUDA-graph scratch / graph_vars) stays in torch's + # default allocator — those buffers are tiny and, critically, + # the captured CUDA graphs reference block_mask indices by + # address, so they must survive sleep / wake unchanged. + 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, + cumem_allocator = self._cumem_allocator, + ) + self._llm_engine_stub = _LLMEngineStub( + sleep_enabled = self._sleep_mode_enabled, + ) + + # ----- 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 ----- + # + # vLLM's sleep mode offloads engine weights to CPU between rollouts so + # training can use the freed VRAM. The flex backend implements level 1 + # (weights offloaded to pinned CPU, KV cache dropped and re-zeroed on + # wake) via :class:`vllm.device_allocator.cumem.CuMemAllocator`. + # Captured CUDA graphs survive the round-trip because cuMem keeps the + # GPU virtual addresses stable across sleep / wake. + # + # Sleep activates only when ``UNSLOTH_VLLM_STANDBY=1`` is set AND + # vLLM is importable (evaluated at ``__init__`` time). Otherwise + # ``sleep`` / ``wake_up`` are no-ops so code that unconditionally + # calls the API (TRL's GRPO trainer) stays correct. + + def sleep(self, level: int = 1): + """Offload inference weights to pinned CPU memory (level 1). + + ``level=2`` is not implemented on the flex backend (it would + require rebuilding the inference deep-copy from the training + model on wake); requesting it emits a warning and falls back to + level 1. + + In the 4-bit single-copy fallback path the inference model + shares storage with the training model, so only the KV cache is + dropped on sleep; the weights stay resident. + """ + if not self._sleep_mode_enabled or self._cumem_allocator is None: + return None + if level not in (1, 2): + raise ValueError(f"FlexEngine.sleep: level must be 1 or 2, got {level}") + if level == 2: + warnings.warn( + "FlexEngine.sleep(level=2) is not implemented on the " + "flex backend; falling back to level=1 (CPU-pinned " + "weight offload).", + RuntimeWarning, + stacklevel = 2, + ) + if self._single_copy_mode: + # Weights are shared with the training model (4-bit path); + # only the kv_cache pool is ours to drop. + self._cumem_allocator.sleep(offload_tags = ()) + else: + self._cumem_allocator.sleep(offload_tags = ("weights",)) + gc.collect() + # NOTE: we deliberately do NOT call torch.cuda.empty_cache() here. + # Captured CUDA graphs may retain scratch / workspace tensors in + # torch's default caching allocator at fixed addresses; emptying + # the cache between sleep and wake can invalidate those + # addresses and cause the next graph replay to read freed + # memory. The cuMem pools have already released their physical + # pages; there is no additional VRAM to reclaim via empty_cache. + return None + + def wake_up(self, tags: Optional[list] = None): + """Re-map cuMem handles and restore offloaded weights. + + ``tags=None`` wakes everything; TRL's GRPO trainer calls + ``wake_up(tags=["kv_cache"])`` and then ``wake_up(tags=["weights"])`` + on consecutive steps to stagger the VRAM reclaim. + """ + if not self._sleep_mode_enabled or self._cumem_allocator is None: + return None + if tags is not None and not isinstance(tags, list): + tags = list(tags) + self._cumem_allocator.wake_up(tags = tags) + 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: + if self.arch in ("qwen3_moe", "gpt_oss", "gemma4_moe", "qwen3_5_moe"): + # For MoE architectures (Qwen3-MoE, gpt-oss, gemma4-MoE), + # LoRA lives on the stacked-expert ParamWrapper and never + # merges in-place into the expert tensors during training + # (see unsloth_zoo.temporary_patches.moe_utils + # _patched_param_wrapper_forward). That means the + # training model's expert weights ARE the pristine + # source — no third 20-60 GB deep-copy is needed. Point + # the LoRA-refresh helper at the training model's base + # directly. This keeps the model at 2x residency instead + # of 3x. + try: + pristine = training_peft_model.get_base_model() + except AttributeError: + pristine = training_peft_model + self._pristine_base = pristine + self._impl.base_model = pristine + else: + # Dense path: the inference model has already been + # flex-patched; its linear weights (what LoRA merges + # into) are still pristine, so we clone it and just do + # not call flex attention on the pristine copy. + with weight_pool(self._cumem_allocator): + 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. + with weight_pool(self._cumem_allocator): + 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, + ) + + +# --------------------------------------------------------------------------- +# Lazy engine construction +# +# FlexEngine's ``max_batch_size`` drives fixed-shape GPU page tables, the +# ``input_pos_buffer``, the ``block_mask_logical`` build, and the CUDA-graph +# bucket list. These are allocated inside ``FlexEngine.__init__`` and there is +# no post-init resize path. Picking the wrong value at ``from_pretrained`` +# time forces users to either overshoot (wasted KV pages + slower graph +# capture) or undershoot (``PageTable.can_reserve`` stalls the rollout). +# +# ``build_flex_engine`` defers the construction until the real rollout batch +# size is known: GRPOTrainer's ``__init__`` patch in ``unsloth/models/rl.py`` +# calls :func:`_build_flex_from_args` to pass +# ``per_device_train_batch_size * steps_per_generation * num_generations`` +# through; the plain ``model.fast_generate`` path falls back to the +# ``max_batch_size`` kwarg originally passed to ``from_pretrained``. +# --------------------------------------------------------------------------- + + +class _LazyFlexEngineSentinel: + """Placeholder for ``model.vllm_engine`` before the FlexEngine is built. + + Forwards attribute access to the real engine, triggering construction + (with the stashed ``max_batch_size`` floor) on first access. This keeps + ``hasattr(model, "vllm_engine")`` True between ``from_pretrained`` and + the first build, which matters for ``rl.py``'s ``args.use_vllm`` setter. + """ + + __slots__ = ("_model",) + + def __init__(self, model): + object.__setattr__(self, "_model", model) + + def _resolve(self): + engine = getattr(self._model, "_flex_engine_instance", None) + if engine is None: + engine = build_flex_engine(self._model) + return engine + + def __getattr__(self, name): + if name == "_model": + raise AttributeError(name) + # Keep dunder lookups out of the resolve path so ``copy.deepcopy`` + # (which probes ``__deepcopy__``, ``__reduce_ex__``, ...) doesn't + # recursively trigger an engine build. This matters when FlexEngine + # itself calls ``copy.deepcopy(hf_model)`` to build its inference + # copy — the training model still holds the sentinel, and without + # this guard the copy.deepcopy attribute probe infinite-loops. + if name.startswith("__") and name.endswith("__"): + raise AttributeError(name) + return getattr(self._resolve(), name) + + def __deepcopy__(self, memo): + # The sentinel is bound to the training ``model`` object; a copy + # that points at a different model would be meaningless. Treat it + # as a singleton for copy purposes. + return self + + def __copy__(self): + return self + + def __bool__(self): + return True + + def __repr__(self): + engine = getattr(self._model, "_flex_engine_instance", None) + if engine is None: + return "" + return repr(engine) + + +def install_flex_sentinel(model, tokenizer): + """Wire the lazy ``vllm_engine`` / ``fast_generate`` placeholders. + + Called from ``unsloth/models/llama.py`` and ``unsloth/models/vision.py`` + in place of the eager ``FlexEngine(...)`` construction. The real build + happens inside :func:`build_flex_engine`, triggered either by the + GRPOTrainer patch (``_build_flex_from_args``) or by the first + ``model.fast_generate`` call. + """ + model._unsloth_flex_tokenizer = tokenizer + model.vllm_engine = _LazyFlexEngineSentinel(model) + + def _lazy_fast_generate(prompts = None, *gen_args, **gen_kwargs): + engine = getattr(model, "_flex_engine_instance", None) + if engine is None: + engine = build_flex_engine(model) + return engine.generate(prompts, *gen_args, **gen_kwargs) + + def _lazy_fast_generate_batches(prompts = None, *gen_args, **gen_kwargs): + engine = getattr(model, "_flex_engine_instance", None) + if engine is None: + engine = build_flex_engine(model) + gen_kwargs.setdefault("use_tqdm", False) + return engine.generate(prompts, *gen_args, **gen_kwargs) + + model.fast_generate = _lazy_fast_generate + model.fast_generate_batches = _lazy_fast_generate_batches + + +def _construct_and_attach(model, max_batch_size: int): + """Construct the FlexEngine with the given batch size and wire it up.""" + pending = model._unsloth_needs_flex_engine + tokenizer = getattr(model, "_unsloth_flex_tokenizer", None) + inference_copy = getattr(model, "_unsloth_flex_inference_copy", None) + + engine_kwargs = dict(pending) + engine_kwargs["max_batch_size"] = int(max_batch_size) + + engine = FlexEngine( + hf_model = model, + tokenizer = tokenizer, + inference_model = inference_copy, + base_model = None, + peft_model = None, + **engine_kwargs, + ) + model._flex_engine_instance = engine + + # Drop the sentinel (plain attribute; setting replaces it). + model.vllm_engine = engine + model.fast_generate = engine.generate + model.fast_generate_batches = functools.partial(engine.generate, use_tqdm = False) + + # Consume the one-shot stashes: the deep-copy is owned by the engine now, + # and the needs-dict's role as "build spec" is over. + for _attr in ( + "_unsloth_flex_inference_copy", + "_unsloth_flex_tokenizer", + "_unsloth_needs_flex_engine", + ): + if hasattr(model, _attr): + try: + delattr(model, _attr) + except AttributeError: + pass + return engine + + +def build_flex_engine(model, max_batch_size: Optional[int] = None): + """Construct or return the FlexEngine attached to ``model``. + + Called lazily: by the RL patch (:func:`_build_flex_from_args`) once + ``GRPOTrainer.args`` is resolved, or by ``model.fast_generate`` on + first use. + + ``max_batch_size`` resolution (first build): + - ``floor = model._unsloth_needs_flex_engine['max_batch_size']`` + - ``effective = max(floor, max_batch_size or 0)`` + - A warning is emitted when ``effective > floor`` so users see the + GRPO-driven bump. + + Once the engine is built, it is the sole source of truth for the + batch-size dimension. Subsequent calls are idempotent when the + requested size fits; requesting a larger size raises + :class:`RuntimeError` because the engine's fixed-shape GPU buffers + and captured CUDA graphs cannot be grown in place. + """ + existing = getattr(model, "_flex_engine_instance", None) + pending = getattr(model, "_unsloth_needs_flex_engine", None) + + # Non-flex model (plain HF or plain vLLM). No-op so ``rl.py``'s patch + # stays unconditional. + if existing is None and pending is None: + return None + + requested = int(max_batch_size) if max_batch_size else 0 + + if existing is not None: + if requested <= existing.max_batch_size: + return existing + raise RuntimeError( + f"Unsloth: FlexEngine was built at max_batch_size=" + f"{existing.max_batch_size}; cannot grow to {requested} after " + f"construction (fixed-shape GPU page tables + CUDA graphs). " + f"Pass max_batch_size={requested} to " + f"FastLanguageModel.from_pretrained before the first " + f"fast_generate / GRPOTrainer call." + ) + + floor = int(pending["max_batch_size"]) + target = max(floor, requested) + if target > floor: + warnings.warn( + f"Unsloth: increasing FlexEngine max_batch_size {floor} -> " + f"{target} to fit the GRPO rollout batch " + f"(per_device_train_batch_size * steps_per_generation * " + f"num_generations). Pass max_batch_size={target} to " + f"FastLanguageModel.from_pretrained to silence this warning.", + stacklevel = 2, + ) + return _construct_and_attach(model, target) + + +def _build_flex_from_args(model, args): + """Helper used by the ``rl.py`` GRPOTrainer-init patch. + + Reads the rollout batch size from the TRL args and triggers the + FlexEngine build. No-op when ``model`` wasn't loaded through the + flex-inference path (plain vLLM / plain HF). + """ + if not hasattr(model, "_unsloth_needs_flex_engine") and not hasattr( + model, "_flex_engine_instance" + ): + return None + pdbs = int(getattr(args, "per_device_train_batch_size", 1) or 1) + spg = int( + getattr(args, "steps_per_generation", None) + or getattr(args, "gradient_accumulation_steps", 1) + or 1 + ) + ngen = int(getattr(args, "num_generations", 1) or 1) + # Written as ``max(A, B)`` for reviewer clarity. Reduces to the second + # term whenever ``num_generations >= 1`` (always). + grpo_target = max(pdbs * spg, pdbs * spg * ngen) + return build_flex_engine(model, max_batch_size = grpo_target) + + +__all__ = [ + "FlexEngine", + "load_flex", + "build_flex_engine", + "install_flex_sentinel", + "_build_flex_from_args", +] diff --git a/unsloth/inference/flex_gemma4.py b/unsloth/inference/flex_gemma4.py new file mode 100644 index 0000000000..c75a1f0f4a --- /dev/null +++ b/unsloth/inference/flex_gemma4.py @@ -0,0 +1,1172 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""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, + cumem_allocator = 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 + + # See ``FlexInference.__init__`` for why only the paged-KV + # allocations go through the cuMem ``kv_cache`` pool. + from .sleep_mode import kv_cache_pool as _kv_cache_pool + + with _kv_cache_pool(cumem_allocator): + 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_gemma4_moe.py b/unsloth/inference/flex_gemma4_moe.py new file mode 100644 index 0000000000..afea7b4bf6 --- /dev/null +++ b/unsloth/inference/flex_gemma4_moe.py @@ -0,0 +1,937 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""Gemma 4 MoE inference with flex_attention + paged KV cache. + +Sibling of ``flex_moe.py`` (Qwen3-MoE) and ``flex_gpt_oss.py`` (gpt-oss). +Scope: ``unsloth/gemma-4-26B-A4B-it`` (30 layers, 128 experts top-k 8, +hidden 2816, moe_intermediate 704, dense intermediate 2112, ~3.8B active). + +Arch-specific pieces on top of the shared engine: + +1. **Dual dense MLP + MoE per layer.** Unlike Qwen3/gpt-oss which replace + the dense MLP with an expert block, Gemma 4 MoE runs both in parallel + per layer and sums their normed outputs before the residual add: + + residual = h + h = pre_ffw_norm(h) + mlp_out = layer.mlp(h) # dense SwiGLU + h1 = post_ffw_norm_1(mlp_out) + h2 = pre_ffw_norm_2(residual.reshape(-1, H)) # experts input + h2 = layer.experts(h2, top_k_idx, top_k_w) + h2 = post_ffw_norm_2(h2).reshape(residual) + h = post_ffw_norm(h1 + h2) + h = residual + h + h *= layer.layer_scalar + +2. **Two-tier RoPE.** + - Sliding layers (25/30): ``rope_type="default"``, theta=10K, + full ``head_dim=256``. + - Full-attn layers (5/30): ``rope_type="proportional"``, theta=1M, + ``global_head_dim=512``, ``partial_rotary_factor=0.25`` (rotate + only leading 25% of head dim; pass through remainder). + Walker computes both (cos, sin) tuples per generate entry; attention + forward picks the right one based on ``self_attn.layer_type``. + +3. **Per-head Q/K/V RMSNorm pre-rotary.** ``q_norm`` and ``k_norm`` apply + RMSNorm on the last (head_dim) axis before RoPE. ``v_norm`` applies + with ``with_scale=False`` (RMSNorm that just divides, no gain). + +4. **K=V alternative on full-attn layers.** ``attention_k_eq_v=True`` + + full-attn layer ⇒ ``v_proj is None`` and ``value_states`` is the raw + ``k_proj(hidden)`` output (before ``k_norm`` and before RoPE), + followed only by ``v_norm``. Sliding layers use the normal q/k/v path. + +5. **Expert grouped_mm.** Reuses ``Gemma4TextExperts.forward`` through + ``unsloth_zoo.temporary_patches.gemma4_moe.patch_gemma4_moe`` — the + ``per_expert_scale`` is pre-folded into routing weights, so the + generic ``forward_native_grouped_mm`` (``moe_utils.py``) handles the + standard ``(E, 2I, H)``/``(E, H, I)`` layout with ``act_fn = + gelu_pytorch_tanh`` via the default ``elif hasattr(self, 'act_fn')`` + fallback — zero changes needed in moe_utils. + +6. **Embedding scale + final ``layer_scalar``.** Embed output gets + multiplied by ``sqrt(hidden_size)``. Each decoder layer's output is + multiplied by ``self.layer_scalar`` (a ``torch.ones(1)`` buffer — + numerically a no-op today, but must not be dropped so the walker + matches HF bitwise). + +Out of scope (this file errors out if the config requests them): +- ``num_kv_shared_layers > 0`` (E2B/E4B KV-share variants) +- ``hidden_size_per_layer_input > 0`` (E2B/E4B per-layer input gate) +- Mixed sliding-window sizes across layers +- bnb-4bit stacked experts (no ``Gemma4TextExpertsBnb4bit`` ships yet) +""" + +from __future__ import annotations + +import math +import os +import types +from collections import deque +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn.attention.flex_attention import BlockMask, create_block_mask + +try: + from .flex_qwen3_llama import ( + DECODE_KERNEL_OPTIONS_DEFAULT, + PREFILL_KERNEL_OPTIONS_DEFAULT, + Sequence, + flex_attention_compiled, + refresh_lora_merge_from_pristine, + ) + from .flex_moe import refresh_moe_lora_merge_from_pristine + from .flex_paged_attention import PagedKVCache, PageTable +except ImportError: # script-mode fallback + import sys + from pathlib import Path + sys.path.insert(0, str(Path(__file__).resolve().parent)) + from flex_qwen3_llama import ( # noqa: E402 + DECODE_KERNEL_OPTIONS_DEFAULT, + PREFILL_KERNEL_OPTIONS_DEFAULT, + Sequence, + flex_attention_compiled, + refresh_lora_merge_from_pristine, + ) + from flex_moe import refresh_moe_lora_merge_from_pristine # noqa: E402 + from flex_paged_attention import PagedKVCache, PageTable # noqa: E402 + + +# --------------------------------------------------------------------------- +# Rotary helpers +# --------------------------------------------------------------------------- + + +def _rotate_half(x): + x1, x2 = torch.chunk(x, 2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + + +def _apply_rotary_full(x, cos, sin): + """Standard Llama-style RoPE on the full head dim. + + ``cos``/``sin`` shape ``(B, S, D)`` → unsqueeze at head-dim to + ``(B, 1, S, D)`` so it broadcasts over ``(B, H, S, D)`` q/k. + + For Gemma 4 full-attention layers, ``rope_type="proportional"`` emits + an ``inv_freq`` with zeros in the tail ``(1 - partial_rotary_factor)`` + fraction of positions. Those zero entries make the corresponding + cos=1 / sin=0, so ``rotate_half`` passes the tail dims through + unchanged — no separate partial-RoPE helper is needed. + """ + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + return (x * cos) + (_rotate_half(x) * sin) + + +# --------------------------------------------------------------------------- +# Attention forward: paged KV + flex_attention + per-layer sliding + k=v. +# --------------------------------------------------------------------------- + + +def make_gemma4_moe_attention_forward(page_table: PageTable): + """Return a ``forward`` method for ``Gemma4TextAttention``. + + Differences from ``flex_gpt_oss``: + - No sinks; call ``flex_attention_compiled`` without ``return_lse``. + - ``q_norm`` / ``k_norm`` / ``v_norm`` before rotary / KV write. + - ``v_proj is None`` (k=v full-attn) ⇒ reuse raw ``k_proj`` as v. + - Partial RoPE on full-attn layers via ``_partial_rotary_dim``. + """ + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor] = None, + position_embeddings_sliding: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + position_embeddings_full: Optional[tuple[torch.Tensor, torch.Tensor]] = None, + attention_mask=None, + past_key_values=None, + cache_position=None, + flex_block_mask: Optional[BlockMask] = None, + flex_block_mask_sliding: 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) + + # pick cos/sin for this layer + if self.is_sliding: + cos, sin = position_embeddings_sliding + else: + cos, sin = position_embeddings_full + + # q projection + q = self.q_proj(hidden_states).view(hidden_shape) + q = self.q_norm(q) + # shape now (B, S, Hq, D) + + # k/v projections + k_raw = self.k_proj(hidden_states).view(hidden_shape) + + if self.v_proj is not None: + v_raw = self.v_proj(hidden_states).view(hidden_shape) + else: + # k=v alternative: value uses the RAW k projection + # (before k_norm, before rotary). v_norm then applies with + # with_scale=False. + v_raw = k_raw + + k = self.k_norm(k_raw) + v = self.v_norm(v_raw) + + # Transpose for flex_attention: (B, S, H, D) -> (B, H, S, D). + # Rotary applied AFTER the transpose so cos/sin broadcast on dim 1. + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + + q = _apply_rotary_full(q, cos, sin) + k = _apply_rotary_full(k, cos, sin) + + # Paged-KV write. + if self._paged_cache is not None and flex_input_pos is not None: + cache_dtype = self._paged_cache.k_cache.dtype + if k.dtype != cache_dtype: + k = k.to(cache_dtype) + if v.dtype != cache_dtype: + v = v.to(cache_dtype) + k, v = self._paged_cache.update(flex_input_pos, k, v, flex_batch_idx) + + # Per-layer block mask dispatch. + if self.is_sliding and flex_block_mask_sliding is not None: + block_mask = flex_block_mask_sliding + else: + block_mask = flex_block_mask + + 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_moe_attention_forwards(model: torch.nn.Module, page_table: PageTable): + """Attach a ``PagedKVCache`` + replace ``forward`` on every + ``Gemma4TextAttention`` layer AND replace ``Gemma4TextExperts.forward`` + with the grouped_mm MoE backend. + + Gemma 4 has per-layer-type head_dim and kv_heads, so each layer's + cache is sized from its own attrs. + + The stock ``Gemma4TextExperts.forward`` (modeling_gemma4.py:1263) is a + Python loop over active experts — slow for decode and not + CUDA-graph-capturable. We rebind to the generic + ``forward_native_grouped_mm`` (moe_utils.py:771) which handles the + standard ``(E, 2I, H)`` / ``(E, H, I)`` layout with the ``act_fn`` + fallback activation path (``gelu_pytorch_tanh`` via ``ACT2FN``). + + The router's ``per_expert_scale`` is already folded into + ``top_k_weights`` by the modeling-native ``Gemma4TextRouter.forward`` + (modeling_gemma4.py:1309), so no separate folding is needed here. + (The unsloth-zoo ``gemma4_moe.patch_gemma4_moe`` targets a legacy + ``Gemma4TextMoEBlock`` class that no longer ships in transformers + 5.5+; it no-ops on current transformers.) + """ + fwd = make_gemma4_moe_attention_forward(page_table) + text_cfg = getattr(model.config, "text_config", model.config) + + if getattr(text_cfg, "num_kv_shared_layers", 0): + raise NotImplementedError( + "Gemma 4 MoE with KV-shared layers not supported yet " + f"(num_kv_shared_layers={text_cfg.num_kv_shared_layers})." + ) + if getattr(text_cfg, "hidden_size_per_layer_input", 0): + raise NotImplementedError( + "Gemma 4 MoE with per-layer input gate not supported yet " + f"(hidden_size_per_layer_input={text_cfg.hidden_size_per_layer_input})." + ) + + # Pick grouped_mm MoE backend once; bind per-layer below. + try: + from unsloth_zoo.temporary_patches.moe_utils import get_forward_moe_backend + _moe_forward = get_forward_moe_backend() + except Exception: # pragma: no cover - defensive + _moe_forward = None + + for layer in model.model.layers: + attn = layer.self_attn + num_q_heads = attn.q_proj.out_features // attn.head_dim + num_kv_heads = max(1, num_q_heads // attn.num_key_value_groups) + attn._paged_cache = PagedKVCache( + page_table, + n_heads=num_kv_heads, + head_dim=attn.head_dim, + dtype=model.dtype, + ).to(model.device) + attn.forward = types.MethodType(fwd, attn) + + if _moe_forward is not None and getattr(layer, "enable_moe_block", False): + experts = layer.experts + # The grouped_mm backend inspects ``self.act_fn`` — already + # set by ``Gemma4TextExperts.__init__`` to + # ``ACT2FN[config.hidden_activation]``. No other setup needed. + experts.forward = types.MethodType(_moe_forward, experts) + + +# --------------------------------------------------------------------------- +# Walker: dual dense MLP + MoE per decoder layer + layer_scalar. +# --------------------------------------------------------------------------- + + +def _compute_rotary_per_layer_type(base, inputs_embeds, position_ids): + """Call ``Gemma4TextRotaryEmbedding`` once per known ``layer_type``. + + Returns ``(cos_sliding, sin_sliding), (cos_full, sin_full)``. cos/sin + for the full-attention type are ALREADY partial-sized in the HF impl + if ``rope_type="proportional"`` + ``partial_rotary_factor<1.0`` — the + init fn slices the head dim to ``partial_rotary_factor * head_dim`` + before emitting ``inv_freq``. So we just forward through. + """ + rot = base.rotary_emb + layer_types = set(getattr(rot, "layer_types", {"sliding_attention", "full_attention"})) + pos_sliding = pos_full = None + if "sliding_attention" in layer_types: + pos_sliding = rot(inputs_embeds, position_ids, layer_type="sliding_attention") + if "full_attention" in layer_types: + pos_full = rot(inputs_embeds, position_ids, layer_type="full_attention") + # Fallback: if only one type exists, use it for both slots so the + # attention forward never reads None. + if pos_sliding is None: + pos_sliding = pos_full + if pos_full is None: + pos_full = pos_sliding + return pos_sliding, pos_full + + +def call_gemma4_moe_model_with_flex_kwargs(model, input_ids, position_ids, flex_kwargs): + """Walk a ``Gemma4TextModel`` manually, injecting flex kwargs into each + attention call. Per-layer dual MLP + MoE with ``layer_scalar``. + + ``model.model.embed_tokens`` is a ``Gemma4TextScaledWordEmbedding`` that + already applies ``* sqrt(hidden_size)`` in its forward, so we just + call it (no manual scale). + """ + base = model.model + inputs_embeds = base.embed_tokens(input_ids) + position_embeddings_sliding, position_embeddings_full = _compute_rotary_per_layer_type( + base, inputs_embeds, position_ids + ) + + hidden_states = inputs_embeds + compute_dtype = inputs_embeds.dtype + + for layer in base.layers: + residual = hidden_states + hidden_states = layer.input_layernorm(hidden_states).to(compute_dtype) + + attn_out, _ = layer.self_attn( + hidden_states, + position_embeddings_sliding=position_embeddings_sliding, + position_embeddings_full=position_embeddings_full, + **flex_kwargs, + ) + attn_out = layer.post_attention_layernorm(attn_out).to(compute_dtype) + hidden_states = residual + attn_out + + residual = hidden_states + pre_ffw = layer.pre_feedforward_layernorm(hidden_states).to(compute_dtype) + mlp_out = layer.mlp(pre_ffw) + + if getattr(layer, "enable_moe_block", False): + h1 = layer.post_feedforward_layernorm_1(mlp_out).to(compute_dtype) + + flat = residual.reshape(-1, residual.shape[-1]) + _, top_k_w, top_k_idx = layer.router(flat) + h2 = layer.pre_feedforward_layernorm_2(flat).to(compute_dtype) + h2 = layer.experts(h2, top_k_idx, top_k_w) + h2 = h2.reshape(residual.shape) + h2 = layer.post_feedforward_layernorm_2(h2).to(compute_dtype) + + mlp_out = h1 + h2 + + mlp_out = layer.post_feedforward_layernorm(mlp_out).to(compute_dtype) + hidden_states = residual + mlp_out + hidden_states = hidden_states * layer.layer_scalar + + hidden_states = base.norm(hidden_states) + return hidden_states + + +# --------------------------------------------------------------------------- +# Sliding-window block-mask builders. Paged-KV aware. +# --------------------------------------------------------------------------- + + +def _create_sliding_causal_blockmask(page_table: PageTable, B: int, L: int, W: int): + def sliding_causal(b, h, q_idx, kv_idx): + return (q_idx >= kv_idx) & (q_idx - kv_idx < W) + + return create_block_mask( + sliding_causal, + B=B, + H=None, + Q_LEN=L, + KV_LEN=L, + BLOCK_SIZE=page_table.page_size, + device=page_table.device, + ) + + +def _create_prefill_sliding_blockmask( + page_table: PageTable, batch_idx: torch.Tensor, W: int, BLOCK_SIZE: int = 128 +): + assert batch_idx.ndim == 2 and batch_idx.shape[0] == 1 + L = batch_idx.shape[1] + docs = batch_idx.view(-1) + + def document_causal_sliding(b, h, q_idx, kv_idx): + causal = q_idx >= kv_idx + window = (q_idx - kv_idx) < W + document = docs[q_idx] == docs[kv_idx] + return causal & window & document + + return create_block_mask( + document_causal_sliding, + B=1, + H=None, + Q_LEN=L, + KV_LEN=L, + BLOCK_SIZE=BLOCK_SIZE, + ) + + +# --------------------------------------------------------------------------- +# FlexGemma4MoEInference +# --------------------------------------------------------------------------- + + +class FlexGemma4MoEInference: + """Gemma 4 MoE inference engine. API-compatible with ``FlexMoEInference`` + and ``FlexGptOssInference``. + + Phase 1: eager decode (capture disabled). Phase 3 enables capture. + """ + + 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, + cumem_allocator=None, + compile_walker=None, + ): + assert max_seq_length % page_size == 0 + assert hasattr(model, "model") and hasattr(model.model, "layers"), ( + "FlexGemma4MoEInference expects a HF CausalLM shape (.model.layers)." + ) + for i, layer in enumerate(model.model.layers): + if not getattr(layer, "enable_moe_block", False): + # Dense-only Gemma 4 layers mixed with MoE is not expected + # for 26B-A4B; if it ever appears we still don't crash — + # the walker's ``if enable_moe_block`` guard handles it. + continue + assert hasattr(layer, "router") and hasattr(layer, "experts"), ( + f"Layer {i} claims enable_moe_block=True but is missing router/experts." + ) + + 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 + + # Detect per-layer sliding window. Gemma 4 uses one sliding size + # across sliding_attention layers. + sliding_windows = set() + for layer in model.model.layers: + sw = getattr(layer.self_attn, "sliding_window", None) + if sw is not None: + sliding_windows.add(int(sw)) + if len(sliding_windows) > 1: + raise NotImplementedError( + f"Mixed sliding-window sizes not supported: {sliding_windows}" + ) + self.sliding_window = next(iter(sliding_windows), None) + + # Detect bnb-4bit stacked experts (no such class ships yet for + # Gemma 4; if it appears, flag and force eager). + self._has_bnb_experts = any( + getattr(layer, "enable_moe_block", False) + and type(layer.experts).__name__.endswith("Bnb4bit") + for layer in model.model.layers + ) + + # FA4 kernel branch. + 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 sm_90+; found sm_{major}0. Falling " + f"back 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 + + from .sleep_mode import kv_cache_pool as _kv_cache_pool + with _kv_cache_pool(cumem_allocator): + self.page_table = PageTable( + n_pages=n_pages, + page_size=page_size, + max_batch_size=max_batch_size, + device=self.device.type, + ) + patch_gemma4_moe_attention_forwards(model, self.page_table) + + self.input_pos_buffer = torch.zeros( + max_batch_size, dtype=torch.int32, device=self.device + ) + self.block_mask_logical = self.page_table.create_causal_blockmask( + B=max_batch_size, + L=max_seq_length, + ) + if self.sliding_window is not None: + self.block_mask_logical_sliding = _create_sliding_causal_blockmask( + self.page_table, + B=max_batch_size, + L=max_seq_length, + W=self.sliding_window, + ) + else: + self.block_mask_logical_sliding = None + + self.cudagraph_captured = False + self.graphs = {} + self.graph_vars = {} + self.graph_bs = None + + # Optional torch.compile walker. + if compile_walker is None: + compile_walker = os.environ.get("UNSLOTH_FLEX_COMPILE_WALKER", "") == "1" + self._moe_walker = call_gemma4_moe_model_with_flex_kwargs + if compile_walker: + try: + self._moe_walker = torch.compile( + call_gemma4_moe_model_with_flex_kwargs, + fullgraph=False, + dynamic=False, + ) + print( + "[flex-gemma4moe] wrapped call_gemma4_moe_model_with_flex_kwargs " + "with torch.compile(fullgraph=False, dynamic=False)" + ) + except Exception as e: + print(f"[flex-gemma4moe] torch.compile wrap failed: {e}") + self._moe_walker = call_gemma4_moe_model_with_flex_kwargs + + if self._has_bnb_experts: + print( + "[flex-gemma4moe] bnb-4bit experts detected; CUDA graph capture " + "disabled (decode stays eager)." + ) + + # --- tokenize / prefill / decode -------------------------------------- + + def tokenize(self, sequences): + for seq in sequences: + if seq.input_ids is not None and seq.input_length > 0: + 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: + 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 = self.page_table.create_prefill_blockmask_no_paging( + batch_idx, BLOCK_SIZE=prefill_block_size + ) + mask_sliding = None + if self.sliding_window is not None: + mask_sliding = _create_prefill_sliding_blockmask( + self.page_table, + batch_idx, + W=self.sliding_window, + BLOCK_SIZE=prefill_block_size, + ) + + flex_kwargs = dict( + flex_block_mask=mask, + flex_block_mask_sliding=mask_sliding, + flex_input_pos=input_pos, + flex_batch_idx=batch_idx, + flex_kernel_options=self.prefill_kernel_options, + ) + position_ids = input_pos + hidden = self._moe_walker( + 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, *, sliding: bool): + block_mask = ( + self.block_mask_logical_sliding if sliding else 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) + + if sliding: + W = self.sliding_window + + def mask_fn(off): + def m(b, h, q_idx, kv_idx): + pos = q_idx + off[b] + return (pos >= kv_idx) & (pos - kv_idx < W) + return m + else: + def mask_fn(off): + def m(b, h, q_idx, kv_idx): + return q_idx + off[b] >= kv_idx + return m + + 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=mask_fn(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_full, input_pos = self._decode_block_mask(batch_idx, sliding=False) + mask_full = self.page_table.convert_logical_block_mask(mask_full, batch_idx) + mask_sliding = None + if self.sliding_window is not None: + ms, _ = self._decode_block_mask(batch_idx, sliding=True) + mask_sliding = self.page_table.convert_logical_block_mask(ms, batch_idx) + + position_ids = input_pos.view(B, 1).to(torch.long) + flex_kwargs = dict( + flex_block_mask=mask_full, + flex_block_mask_sliding=mask_sliding, + flex_input_pos=input_pos.view(B, 1).to(torch.long), + flex_batch_idx=batch_idx, + flex_kernel_options=self.decode_kernel_options, + ) + hidden = self._moe_walker( + self.model, input_ids.view(B, 1), position_ids, flex_kwargs + ) + return self.model.lm_head(hidden[:, -1, :]) + + 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 or self.graph_bs is None: + 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): + """Capture one CUDA graph per bs bucket. + + Phase 3 wires this in. bnb-4bit stays eager either way (no + stacked bnb experts class ships for Gemma 4 today; guard is + defensive). + """ + if self._has_bnb_experts: + print( + "[flex-gemma4moe] bnb-4bit experts: skipping cudagraph capture; " + "decode stays eager." + ) + return + try: + from unsloth_zoo.temporary_patches.moe_utils import select_moe_backend + backend = select_moe_backend() + except Exception: + backend = None + if backend != "grouped_mm": + print( + f"[flex-gemma4moe] MoE CUDA graph capture requires the " + f"'grouped_mm' backend (got {backend!r}); skipping capture." + ) + return + + 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, + ) + + _env_bs = os.environ.get("UNSLOTH_FLEX_GRAPH_BS") + if _env_bs: + try: + self.graph_bs = [int(x) for x in _env_bs.split(",") if x.strip()] + except ValueError: + print( + f"[flex-gemma4moe] invalid UNSLOTH_FLEX_GRAPH_BS={_env_bs!r}; " + f"using default bucket ladder" + ) + self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16)) + else: + 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-gemma4moe] 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): + """Refresh merged LoRA weights from pristine. Dense refresh on + q/k/v/o/router/dense-MLP; MoE refresh on ``gate_up_proj`` / + ``down_proj`` in standard ``(E, 2I, H)``/``(E, H, I)`` + orientation — reuses ``refresh_moe_lora_merge_from_pristine`` + verbatim.""" + if self.base_model is None or self.peft_model is None: + return 0 + n = refresh_lora_merge_from_pristine(self.base_model, self.peft_model) + try: + n += refresh_moe_lora_merge_from_pristine( + self.base_model, self.peft_model + ) + except Exception: + pass + return n + + @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() + if self.graphs: + 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 diff --git a/unsloth/inference/flex_gpt_oss.py b/unsloth/inference/flex_gpt_oss.py new file mode 100644 index 0000000000..375d2273bb --- /dev/null +++ b/unsloth/inference/flex_gpt_oss.py @@ -0,0 +1,794 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""gpt-oss inference with flex_attention + paged KV cache + attention sinks. + +Sibling of ``flex_moe.py`` (Qwen3-MoE). Three arch-specific pieces: + +1. **Attention sinks.** Each layer has a learned per-head ``sinks`` + parameter that biases a virtual sink token in the softmax. We + implement it the same way ``unsloth_zoo.flex_attention.attention_sink`` + does: call flex_attention with ``return_lse=True``, then scale the + output by ``sigmoid(lse - sinks[h])``. + +2. **Per-layer sliding window.** gpt-oss layer types alternate between + full attention and sliding-128. The walker passes two ``BlockMask`` + objects (``flex_block_mask`` + ``flex_block_mask_sliding``) and each + attention forward picks one based on ``self.sliding_window``. + +3. **MoE expert stack.** Uses unsloth_zoo's ``GptOssExperts`` forward + (``forward_native_grouped_mm`` with the ``"GptOssExperts"`` branch + — interleaved gate/up split + ``gate * sigmoid(gate * 1.702)`` — is + already there at moe_utils.py:918-972, reused verbatim). The MoE + LoRA merge reuses ``refresh_moe_lora_merge_from_pristine`` from + ``flex_moe.py`` — the transposed-orientation branch covers + gpt-oss's ``(E, H, 2I)`` gate_up_proj layout. + +bnb-4bit: ``GptOssExpertsBnb4bit`` uses an ``nn.ModuleList`` per-expert +loop that can't be CUDA-graph-captured. Detect it in ``__init__`` and +set ``capture_cudagraph=False`` — decode still benefits from paged KV +and flex_attention, just not the graph replay. +""" + +from __future__ import annotations + +import os +import types +from collections import deque +from typing import Optional + +import torch +import torch.nn.functional as F +from torch.nn.attention.flex_attention import BlockMask, create_block_mask + +try: + from .flex_qwen3_llama import ( + DECODE_KERNEL_OPTIONS_DEFAULT, + PREFILL_KERNEL_OPTIONS_DEFAULT, + Sequence, + flex_attention_compiled, + refresh_lora_merge_from_pristine, + ) + from .flex_moe import refresh_moe_lora_merge_from_pristine + from .flex_paged_attention import PagedKVCache, PageTable +except ImportError: # script-mode fallback + import sys + from pathlib import Path + sys.path.insert(0, str(Path(__file__).resolve().parent)) + from flex_qwen3_llama import ( # noqa: E402 + DECODE_KERNEL_OPTIONS_DEFAULT, + PREFILL_KERNEL_OPTIONS_DEFAULT, + Sequence, + flex_attention_compiled, + refresh_lora_merge_from_pristine, + ) + from flex_moe import refresh_moe_lora_merge_from_pristine # noqa: E402 + from flex_paged_attention import PagedKVCache, PageTable # noqa: E402 + + +# gpt-oss rotary is NOT the Llama ``(q * cos) + (rotate_half(q) * sin)`` +# half-dim sin/cos rotation. It does a first/second-half split with +# ``cos`` / ``sin`` sized to ``head_dim / 2``. Import the reference here +# and unsqueeze head-dim once per call — same behaviour as +# ``apply_rotary_pos_emb`` in ``transformers.models.gpt_oss``. +def _gpt_oss_apply_rotary(q, k, cos, sin): + cos = cos.unsqueeze(1) # [B, 1, S, D/2] + sin = sin.unsqueeze(1) + + def rotate(x): + first, second = torch.chunk(x, 2, dim=-1) + first_ = first * cos - second * sin + second_ = second * cos + first * sin + return torch.cat((first_, second_), dim=-1) + + return rotate(q), rotate(k) + + +# --------------------------------------------------------------------------- +# Attention forward: paged KV + flex_attention + sinks + per-layer sliding. +# --------------------------------------------------------------------------- + + +def make_gptoss_attention_forward(page_table: PageTable): + """Return a ``forward`` method for ``GptOssAttention``. + + Differences from the dense Qwen3/Llama flex forward: + - No ``q_norm`` / ``k_norm`` (gpt-oss has neither). + - Uses ``flex_attention(..., return_lse=True)`` and scales output by + ``sigmoid(lse - sinks[h])`` for the sink token. + - Picks ``flex_block_mask_sliding`` when the layer has + ``self.sliding_window is not 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_block_mask_sliding: Optional[BlockMask] = None, + flex_input_pos: Optional[torch.Tensor] = None, + flex_batch_idx: Optional[torch.Tensor] = None, + flex_kernel_options: Optional[dict] = None, + **kwargs, + ): + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + q = self.q_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 = _gpt_oss_apply_rotary(q, k, cos, sin) + + # Paged-KV write. Match the pre-allocated cache dtype; attention + # linears may compute in fp32 under autocast. + if self._paged_cache is not None and flex_input_pos is not None: + cache_dtype = self._paged_cache.k_cache.dtype + if k.dtype != cache_dtype: + k = k.to(cache_dtype) + if v.dtype != cache_dtype: + v = v.to(cache_dtype) + k, v = self._paged_cache.update(flex_input_pos, k, v, flex_batch_idx) + + # Per-layer block mask dispatch. sliding_window is None on full + # attention layers, an int on sliding layers. + if self.sliding_window is not None and flex_block_mask_sliding is not None: + block_mask = flex_block_mask_sliding + else: + block_mask = flex_block_mask + + attn_output, logsumexp = flex_attention_compiled( + q, + k, + v, + scale=self.scaling, + block_mask=block_mask, + enable_gqa=True, + kernel_options=flex_kernel_options, + return_lse=True, + ) + + # Attention sink. Equivalent to concatenating a sink column + # ``sinks[h]`` to the attention logits before softmax. With + # ``return_lse=True`` we have ``lse = log sum_k exp(QK[q, k] * + # scale)``; the sink-aware softmax scales the output by + # ``exp(lse) / (exp(lse) + exp(sinks[h])) = sigmoid(lse - + # sinks[h])``. Mirrors ``flex_attention_add_sinks`` in + # ``unsloth_zoo/flex_attention/attention_sink.py``. + logsumexp = logsumexp - self.sinks.view(1, -1, 1) + sink_scale = torch.sigmoid(logsumexp) + attn_output = attn_output * sink_scale.unsqueeze(-1).to(attn_output.dtype) + + attn_output = attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous() + return self.o_proj(attn_output), None + + return forward + + +def patch_gptoss_attention_forwards(model: torch.nn.Module, page_table: PageTable): + """Attach a ``PagedKVCache`` + replace ``forward`` on every + ``GptOssAttention`` layer.""" + fwd = make_gptoss_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) + attn.forward = types.MethodType(fwd, attn) + + +# --------------------------------------------------------------------------- +# Walker: pass flex kwargs through each layer, unpack MLP tuple return. +# --------------------------------------------------------------------------- + + +def call_gpt_oss_model_with_flex_kwargs(model, input_ids, position_ids, flex_kwargs): + """Walk a ``GptOssModel`` manually, injecting flex kwargs into each + attention call. ``GptOssMLP.forward`` returns + ``(hidden_states, router_scores)``; we discard router scores at + inference (no load-balance loss).""" + base = model.model + inputs_embeds = base.embed_tokens(input_ids) + position_embeddings = base.rotary_emb(inputs_embeds, position_ids) + _cos, _sin = position_embeddings + if _cos.dim() == 2: + _cos = _cos[position_ids] + _sin = _sin[position_ids] + position_embeddings = (_cos, _sin) + hidden_states = inputs_embeds + compute_dtype = inputs_embeds.dtype + for layer in base.layers: + residual = hidden_states + hidden_states = layer.input_layernorm(hidden_states).to(compute_dtype) + hidden_states, _ = layer.self_attn( + hidden_states, + position_embeddings=position_embeddings, + **flex_kwargs, + ) + hidden_states = residual + hidden_states.to(compute_dtype) + residual = hidden_states + hidden_states = layer.post_attention_layernorm(hidden_states).to(compute_dtype) + mlp_out = layer.mlp(hidden_states) + if isinstance(mlp_out, tuple): + hidden_states = mlp_out[0] + else: + hidden_states = mlp_out + hidden_states = residual + hidden_states.to(compute_dtype) + hidden_states = base.norm(hidden_states) + return hidden_states + + +# --------------------------------------------------------------------------- +# Sliding-window block-mask builders. Paged-KV aware. +# --------------------------------------------------------------------------- + + +def _create_sliding_causal_blockmask(page_table: PageTable, B: int, L: int, W: int): + """Full-window causal block mask: ``q >= kv`` and ``q - kv < W``. + Built against the logical block layout; ``convert_logical_block_mask`` + wraps it to paged indices later for decode.""" + + def sliding_causal(b, h, q_idx, kv_idx): + return (q_idx >= kv_idx) & (q_idx - kv_idx < W) + + return create_block_mask( + sliding_causal, + B=B, + H=None, + Q_LEN=L, + KV_LEN=L, + BLOCK_SIZE=page_table.page_size, + device=page_table.device, + ) + + +def _create_prefill_sliding_blockmask( + page_table: PageTable, batch_idx: torch.Tensor, W: int, BLOCK_SIZE: int = 128 +): + """Document-causal + sliding: ``q >= kv``, ``q - kv < W``, and + ``docs[q] == docs[kv]``. Mirrors + ``create_prefill_blockmask_no_paging`` with the window term added.""" + assert batch_idx.ndim == 2 and batch_idx.shape[0] == 1 + L = batch_idx.shape[1] + docs = batch_idx.view(-1) + + def document_causal_sliding(b, h, q_idx, kv_idx): + causal = q_idx >= kv_idx + window = (q_idx - kv_idx) < W + document = docs[q_idx] == docs[kv_idx] + return causal & window & document + + return create_block_mask( + document_causal_sliding, + B=1, + H=None, + Q_LEN=L, + KV_LEN=L, + BLOCK_SIZE=BLOCK_SIZE, + ) + + +# --------------------------------------------------------------------------- +# FlexGptOssInference +# --------------------------------------------------------------------------- + + +class FlexGptOssInference: + """gpt-oss inference engine. API-compatible with ``FlexMoEInference``. + + Phase 1: eager decode only (no CUDA graph capture). Phase 3 adds the + capture path for the bf16 variant; bnb-4bit stays eager because + ``GptOssExpertsBnb4bit`` uses an nn.ModuleList per-expert loop that + isn't graph-capturable. + """ + + 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, + cumem_allocator=None, + compile_walker=None, + ): + assert max_seq_length % page_size == 0 + assert hasattr(model, "model") and hasattr(model.model, "layers"), ( + "FlexGptOssInference expects a HF CausalLM shape (.model.layers)." + ) + for i, layer in enumerate(model.model.layers): + assert hasattr(layer.self_attn, "sinks"), ( + f"Layer {i}.self_attn has no sinks — not a gpt-oss attention?" + ) + + 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 + + # Detect per-layer sliding window. gpt-oss alternates full / + # sliding-128; record whichever window is used so block-mask + # builders pick it up. If no layer has a sliding window, skip + # the sliding mask construction entirely. + sliding_windows = { + int(layer.self_attn.sliding_window) + for layer in model.model.layers + if layer.self_attn.sliding_window is not None + } + if len(sliding_windows) > 1: + raise NotImplementedError( + f"Mixed sliding-window sizes not supported: {sliding_windows}" + ) + self.sliding_window = next(iter(sliding_windows), None) + + # Detect bnb-4bit experts. The bnb variant uses an nn.ModuleList + # per-expert loop — disable CUDA graph capture if present. + self._has_bnb_experts = any( + type(layer.mlp.experts).__name__ == "GptOssExpertsBnb4bit" + for layer in model.model.layers + ) + + # FA4 kernel branch (same as dense/MoE). + 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 sm_90+; found sm_{major}0. Falling " + f"back 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 + + from .sleep_mode import kv_cache_pool as _kv_cache_pool + with _kv_cache_pool(cumem_allocator): + self.page_table = PageTable( + n_pages=n_pages, + page_size=page_size, + max_batch_size=max_batch_size, + device=self.device.type, + ) + patch_gptoss_attention_forwards(model, self.page_table) + + self.input_pos_buffer = torch.zeros( + max_batch_size, dtype=torch.int32, device=self.device + ) + self.block_mask_logical = self.page_table.create_causal_blockmask( + B=max_batch_size, + L=max_seq_length, + ) + # Sliding-window logical block mask (only if any layer uses it). + if self.sliding_window is not None: + self.block_mask_logical_sliding = _create_sliding_causal_blockmask( + self.page_table, + B=max_batch_size, + L=max_seq_length, + W=self.sliding_window, + ) + else: + self.block_mask_logical_sliding = None + + # Cudagraph capture disabled for Phase 1. Phase 3 wires this up + # for bf16; bnb-4bit stays False either way. + self.cudagraph_captured = False + self.graphs = {} + self.graph_vars = {} + + # Optional torch.compile walker (Phase 5 turns this on). + if compile_walker is None: + compile_walker = os.environ.get("UNSLOTH_FLEX_COMPILE_WALKER", "") == "1" + self._moe_walker = call_gpt_oss_model_with_flex_kwargs + if compile_walker: + try: + self._moe_walker = torch.compile( + call_gpt_oss_model_with_flex_kwargs, + fullgraph=False, + dynamic=False, + ) + print( + "[flex-gptoss] wrapped call_gpt_oss_model_with_flex_kwargs " + "with torch.compile(fullgraph=False, dynamic=False)" + ) + except Exception as e: + print(f"[flex-gptoss] torch.compile wrap failed: {e}") + self._moe_walker = call_gpt_oss_model_with_flex_kwargs + + if self._has_bnb_experts: + print( + "[flex-gptoss] bnb-4bit experts detected; CUDA graph capture " + "disabled (decode stays eager)." + ) + + # --- tokenize / prefill / decode -------------------------------------- + + def tokenize(self, sequences): + for seq in sequences: + if seq.input_ids is not None and seq.input_length > 0: + 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: + 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 = self.page_table.create_prefill_blockmask_no_paging( + batch_idx, BLOCK_SIZE=prefill_block_size + ) + mask_sliding = None + if self.sliding_window is not None: + mask_sliding = _create_prefill_sliding_blockmask( + self.page_table, + batch_idx, + W=self.sliding_window, + BLOCK_SIZE=prefill_block_size, + ) + + flex_kwargs = dict( + flex_block_mask=mask, + flex_block_mask_sliding=mask_sliding, + flex_input_pos=input_pos, + flex_batch_idx=batch_idx, + flex_kernel_options=self.prefill_kernel_options, + ) + position_ids = input_pos + hidden = self._moe_walker( + 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, *, sliding: bool): + block_mask = ( + self.block_mask_logical_sliding if sliding else 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) + + if sliding: + W = self.sliding_window + + def mask_fn(off): + def m(b, h, q_idx, kv_idx): + pos = q_idx + off[b] + return (pos >= kv_idx) & (pos - kv_idx < W) + return m + else: + def mask_fn(off): + def m(b, h, q_idx, kv_idx): + return q_idx + off[b] >= kv_idx + return m + + 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=mask_fn(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_full, input_pos = self._decode_block_mask(batch_idx, sliding=False) + mask_full = self.page_table.convert_logical_block_mask(mask_full, batch_idx) + mask_sliding = None + if self.sliding_window is not None: + ms, _ = self._decode_block_mask(batch_idx, sliding=True) + mask_sliding = self.page_table.convert_logical_block_mask(ms, batch_idx) + + position_ids = input_pos.view(B, 1).to(torch.long) + flex_kwargs = dict( + flex_block_mask=mask_full, + flex_block_mask_sliding=mask_sliding, + flex_input_pos=input_pos.view(B, 1).to(torch.long), + flex_batch_idx=batch_idx, + flex_kernel_options=self.decode_kernel_options, + ) + hidden = self._moe_walker( + self.model, input_ids.view(B, 1), position_ids, flex_kwargs + ) + return self.model.lm_head(hidden[:, -1, :]) + + 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): + """Capture one CUDA graph per bs bucket. bf16 experts only; + bnb-4bit (nn.ModuleList per-expert loop) can't be captured and + ``__init__`` already set ``capture_cudagraph`` to False for that + variant — this method is only reached when the engine-level + flag is still True.""" + if self._has_bnb_experts: + print( + "[flex-gptoss] bnb-4bit experts: skipping cudagraph capture; " + "decode stays eager." + ) + return + try: + from unsloth_zoo.temporary_patches.moe_utils import select_moe_backend + backend = select_moe_backend() + except Exception: + backend = None + if backend != "grouped_mm": + print( + f"[flex-gptoss] MoE CUDA graph capture requires the " + f"'grouped_mm' backend (got {backend!r}); skipping capture." + ) + return + + 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, + ) + + _env_bs = os.environ.get("UNSLOTH_FLEX_GRAPH_BS") + if _env_bs: + try: + self.graph_bs = [int(x) for x in _env_bs.split(",") if x.strip()] + except ValueError: + print( + f"[flex-gptoss] invalid UNSLOTH_FLEX_GRAPH_BS={_env_bs!r}; " + f"using default bucket ladder" + ) + self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16)) + else: + 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-gptoss] 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): + """Refresh merged LoRA weights from pristine. Reuses the dense + refresh for q/k/v/o/router and the MoE refresh + (``refresh_moe_lora_merge_from_pristine``) for ``gate_up_proj`` + / ``down_proj`` in the ``(E, H, 2I)`` transposed orientation.""" + if self.base_model is None or self.peft_model is None: + return 0 + n = refresh_lora_merge_from_pristine(self.base_model, self.peft_model) + try: + n += refresh_moe_lora_merge_from_pristine( + self.base_model, self.peft_model + ) + except Exception: + pass + return n + + @torch.inference_mode() + def generate(self, sequences: list[Sequence], capture_cudagraph=False): + """Decode loop. Phase 1: always eager.""" + self.tokenize(sequences) + waiting = deque(sequences) + running = deque() + done = [] + + if capture_cudagraph and not self.cudagraph_captured: + self.capture_decode_cudagraph() + if self.graphs: + 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 diff --git a/unsloth/inference/flex_moe.py b/unsloth/inference/flex_moe.py new file mode 100644 index 0000000000..af61e12e4b --- /dev/null +++ b/unsloth/inference/flex_moe.py @@ -0,0 +1,800 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""Qwen3-MoE inference with flex_attention + paged KV cache. + +Sibling of ``flex_qwen3_llama.py`` (dense Qwen3 / Llama-3). Handles +``Qwen3MoeForCausalLM`` where each decoder layer's ``mlp`` is a +``Qwen3MoeSparseMoeBlock``. Two differences from the dense path: + +1. The walker (``call_moe_model_with_flex_kwargs``) unpacks whatever the + MoE MLP returns. Stock HF 5.x returns a plain tensor; Unsloth's + patched ``Qwen3MoeSparseMoeBlock_fast_forward`` returns + ``(hidden_states, router_logits)``. The ``isinstance(_, tuple)`` + guard handles both without coupling this file to either forward. + +2. Decode runs eager (no CUDA-graph capture). The MoE expert routing + uses ``torch.where`` + a Python for-loop over experts, which is + data-dependent-shape and not graph-capturable. Prefill still uses + flex_attention compiled. A future cut can swap in a padded-fixed- + shape dispatch via ``UNSLOTH_MOE_STATIC_DISPATCH=1``. + +Everything else — paged-KV cache, attention forward, prefill block-mask, +LoRA double-copy refresh — is shared verbatim with the dense path. +""" + +from __future__ import annotations + +import os +import types +from collections import deque +from typing import Optional + +import torch +import torch.nn.functional as F +from torch.nn.attention.flex_attention import BlockMask + +try: + from .flex_qwen3_llama import ( + DECODE_KERNEL_OPTIONS_DEFAULT, + PREFILL_KERNEL_OPTIONS_DEFAULT, + Sequence, + patch_model_attention_forwards, + refresh_lora_merge_from_pristine, + ) + from .flex_paged_attention import PagedKVCache, PageTable +except ImportError: # script-mode fallback + import sys + from pathlib import Path + sys.path.insert(0, str(Path(__file__).resolve().parent)) + from flex_qwen3_llama import ( # noqa: E402 + DECODE_KERNEL_OPTIONS_DEFAULT, + PREFILL_KERNEL_OPTIONS_DEFAULT, + Sequence, + patch_model_attention_forwards, + refresh_lora_merge_from_pristine, + ) + from flex_paged_attention import PagedKVCache, PageTable # noqa: E402 + + +def call_moe_model_with_flex_kwargs(model, input_ids, position_ids, flex_kwargs): + """Walk a Qwen3-MoE model manually, injecting flex_* kwargs into each + attention call. Mirrors ``call_model_with_flex_kwargs`` from + ``flex_qwen3_llama.py`` but handles the MoE MLP return shape. + + For Qwen3-MoE, ``layer.mlp`` is a ``Qwen3MoeSparseMoeBlock``. Its + forward signature varies by patch: + + - stock HF 5.x: returns a single tensor (``final_hidden_states``). + - Unsloth's ``Qwen3MoeSparseMoeBlock_fast_forward``: returns + ``(final_X, router_logits)``. + - unsloth_zoo's ``sparse_moe_block_forward``: returns a single + tensor. + + We call ``layer.mlp(...)`` and unpack whatever comes back. At + inference we discard ``router_logits`` — no load-balance loss. + """ + base = model.model # Qwen3MoeModel + inputs_embeds = base.embed_tokens(input_ids) + position_embeddings = base.rotary_emb(inputs_embeds, position_ids) + # Unsloth's ``LlamaRotaryEmbedding`` drop-in 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. + _cos, _sin = position_embeddings + if _cos.dim() == 2: + _cos = _cos[position_ids] + _sin = _sin[position_ids] + position_embeddings = (_cos, _sin) + hidden_states = inputs_embeds + # RMSNorm + bnb-4bit Linear compute can promote activations to fp32 + # along the Qwen3 MoE path even under autocast. Lock activations to + # the embed dtype so paged-KV writes (which index_put_ into a + # pre-allocated bf16 cache) see a matching dtype. + compute_dtype = inputs_embeds.dtype + for layer in base.layers: + # Attention block — identical to dense Qwen3 / Llama. + residual = hidden_states + hidden_states = layer.input_layernorm(hidden_states).to(compute_dtype) + hidden_states, _ = layer.self_attn( + hidden_states, + position_embeddings = position_embeddings, + **flex_kwargs, + ) + hidden_states = residual + hidden_states.to(compute_dtype) + # MoE MLP. + residual = hidden_states + hidden_states = layer.post_attention_layernorm(hidden_states).to(compute_dtype) + mlp_out = layer.mlp(hidden_states) + if isinstance(mlp_out, tuple): + hidden_states = mlp_out[0] + else: + hidden_states = mlp_out + hidden_states = residual + hidden_states.to(compute_dtype) + hidden_states = base.norm(hidden_states) + return hidden_states + + +class FlexMoEInference: + """MoE inference engine. API-compatible with ``FlexInference`` so + ``FlexEngine`` dispatch is a one-line change. + + Differences: + - uses ``call_moe_model_with_flex_kwargs`` (tuple-aware walker). + - ``cudagraph_captured`` is permanently False; ``generate`` always + runs the eager decode path. ``capture_decode_cudagraph`` raises + ``NotImplementedError`` so a stray ``capture_cudagraph=True`` + fails loudly rather than silently producing wrong output. + """ + + 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, + cumem_allocator = None, + compile_walker = None, + ): + # FastQwen3MoeModel.pre_patch (unsloth/models/qwen3_moe.py) installs + # a legacy Qwen3MoeSparseMoeBlock_fast_forward that expects + # ``self.gate_proj``; transformers 5.x Qwen3MoE uses + # ``self.gate`` / ``self.experts`` instead, so that forward is dead + # code on this env. Unsloth-zoo's ``patch_qwen3_moe`` re-patches it + # to the correct ``sparse_moe_block_forward``, but Unsloth's + # pre_patch can run later and silently clobber it (patch_function + # bails via can_safely_patch on a second pass). Force-restore the + # stock HF forward here so the flex walker sees a working MLP. + try: + import transformers.models.qwen3_moe.modeling_qwen3_moe as _hf_mod + _BlockCls = _hf_mod.Qwen3MoeSparseMoeBlock + cur_forward = getattr(_BlockCls, "forward", None) + cur_name = getattr(cur_forward, "__name__", "") + if "fast_forward" in cur_name or cur_name == "Qwen3MoeSparseMoeBlock_fast_forward": + # Prefer unsloth_zoo's patched version if present; + # fall back to the stock HF forward otherwise. + unique = getattr(_BlockCls, "_original_forward_Qwen3MoeSparseMoeBlock", None) or getattr(_BlockCls, "_Qwen3MoeSparseMoeBlock_original_forward", None) + if unique is not None: + _BlockCls.forward = unique + else: + # Re-run unsloth_zoo patch to install sparse_moe_block_forward. + from unsloth_zoo.temporary_patches.qwen3_moe import patch_qwen3_moe + patch_qwen3_moe() + # If patch_function still skipped due to can_safely_patch, + # fall back to stock HF as a last resort. + cur_forward_after = getattr(_BlockCls, "forward", None) + cur_name_after = getattr(cur_forward_after, "__name__", "") + if "fast_forward" in cur_name_after: + # Lazy-load pristine forward by reloading the module. + import importlib + _fresh_mod = importlib.reload(_hf_mod) + _BlockCls.forward = _fresh_mod.Qwen3MoeSparseMoeBlock.forward + except Exception: + pass + assert max_seq_length % page_size == 0 + # Startup sanity checks. If any of these fail the architecture + # isn't a Qwen3-MoE variant we know how to drive. + assert hasattr(model, "model") and hasattr(model.model, "layers"), ( + "FlexMoEInference expects a HF CausalLM shape (.model.layers)." + ) + for i, layer in enumerate(model.model.layers): + assert hasattr(layer, "post_attention_layernorm"), ( + f"Layer {i} has no post_attention_layernorm." + ) + assert hasattr(layer, "mlp") and callable( + getattr(layer.mlp, "forward", None) + ), f"Layer {i}.mlp has no callable forward." + + 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 + + # Kernel-options / FA4 branch — copied from FlexInference. + 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 + 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 + + # Route paged-KV allocations through the cuMem pool when sleep + # mode is active. Block-mask / input_pos scratch stays in the + # default allocator. + from .sleep_mode import kv_cache_pool as _kv_cache_pool + with _kv_cache_pool(cumem_allocator): + 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) + + self.input_pos_buffer = torch.zeros( + max_batch_size, dtype = torch.int32, device = self.device + ) + 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 = {} + + # Optional: wrap ``call_moe_model_with_flex_kwargs`` with + # ``torch.compile(fullgraph=False, dynamic=False)`` BEFORE CUDA + # graph capture. On Qwen3-30B-A3B 4bit this gives ~2x decode + # throughput (378 → 753 tok/s at bs=8, 1735 → 3383 tok/s at + # bs=48) because Inductor fuses the layernorm + residual + + # router pointwise ops and the compiled kernels get recorded + # into the captured graph. Opt-in for now: either pass + # ``compile_walker=True`` explicitly or set + # ``UNSLOTH_FLEX_COMPILE_WALKER=1``. + if compile_walker is None: + compile_walker = os.environ.get("UNSLOTH_FLEX_COMPILE_WALKER", "") == "1" + self._moe_walker = call_moe_model_with_flex_kwargs + if compile_walker: + try: + self._moe_walker = torch.compile( + call_moe_model_with_flex_kwargs, + fullgraph = False, + dynamic = False, + ) + print( + "[flex-moe] wrapped call_moe_model_with_flex_kwargs " + "with torch.compile(fullgraph=False, dynamic=False)" + ) + except Exception as e: + print(f"[flex-moe] torch.compile wrap failed, falling back: {e}") + self._moe_walker = call_moe_model_with_flex_kwargs + + # --- tokenize / prefill / decode --------------------------------------- + # Near-verbatim from FlexInference. Only difference is the walker. + + def tokenize(self, sequences): + for seq in sequences: + if seq.input_ids is not None and seq.input_length > 0: + 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: + 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 = 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 + hidden = self._moe_walker( + 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): + 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 = self._moe_walker( + self.model, input_ids.view(B, 1), position_ids, flex_kwargs + ) + return self.model.lm_head(hidden[:, -1, :]) + + 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 for MoE decode. + + Supported on the ``grouped_mm`` MoE backend only. On that backend + the decode path is fixed-shape: + ``bincount(minlength=num_experts) → cumsum → argsort → + torch._grouped_mm × 2 → index_add_``. Python control flow in + ``sparse_moe_block_forward`` runs once at capture time; only the + recorded CUDA kernels replay. + + For any other backend (``unsloth_triton``, ``native_torch``) this + method logs a warning and returns without enabling replay, so + ``generate(capture_cudagraph=True)`` silently falls back to eager + decode instead of failing inside the captured graph. + + Pre-reserves a page for every ``batch_idx`` slot so the paged-KV + ``index_put_`` during capture hits valid physical addresses. The + reservations are erased after capture — replay reads / writes the + same physical pages regardless of the logical batch state, + because ``batch_idx = 0`` is reserved as a padding slot. + """ + try: + from unsloth_zoo.temporary_patches.moe_utils import select_moe_backend + backend = select_moe_backend() + except Exception: + backend = None + if backend != "grouped_mm": + print( + f"[flex] MoE CUDA graph capture requires the 'grouped_mm' " + f"backend (got {backend!r}); skipping capture, decode stays " + f"eager." + ) + return + + 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, + ) + # Bucket ladder for CUDA graph capture. Default mirrors the dense + # FlexInference pattern. ``UNSLOTH_FLEX_GRAPH_BS=1,8,32,64`` etc + # lets you override (useful for tight memory or for bigger batch + # bench). Values > max_bs are silently skipped below. + _env_bs = os.environ.get("UNSLOTH_FLEX_GRAPH_BS") + if _env_bs: + try: + self.graph_bs = [int(x) for x in _env_bs.split(",") if x.strip()] + except ValueError: + print( + f"[flex-moe] invalid UNSLOTH_FLEX_GRAPH_BS={_env_bs!r}; " + f"using default bucket ladder" + ) + self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16)) + else: + 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-moe] 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): + """Re-materialize the inference copy's merged LoRA weights from + the pristine base. No-op when no adapter is configured. + + For Qwen3-MoE, dense LoRA targets (q/k/v/o and potentially the + router ``gate``) are handled by the dense refresh. Stacked + expert LoRA targets (``gate_up_proj`` / ``down_proj``) are + handled by the MoE refresh, which writes via ``torch.baddbmm`` + into the same stacked-tensor storage so captured replay + addresses stay valid. + """ + if self.base_model is None or self.peft_model is None: + return 0 + n = refresh_lora_merge_from_pristine(self.base_model, self.peft_model) + try: + n += refresh_moe_lora_merge_from_pristine( + self.base_model, self.peft_model + ) + except Exception: + # MoE LoRA merge is best-effort for now: ZOO's MoE PEFT wrapper + # varies by transformers version. If the wrapper shape isn't + # recognised we fall back to the dense refresh only (which + # already handled any LoraLayer-wrapped modules). + pass + return n + + @torch.inference_mode() + def generate(self, sequences: list[Sequence], capture_cudagraph = False): + """Decode loop. Captures one CUDA graph per bucket on first call + when ``capture_cudagraph=True`` and the ``grouped_mm`` MoE + backend is active; otherwise falls back to eager decode.""" + self.tokenize(sequences) + waiting = deque(sequences) + running = deque() + done = [] + + if capture_cudagraph and not self.cudagraph_captured: + self.capture_decode_cudagraph() + # ``capture_decode_cudagraph`` leaves ``cudagraph_captured`` + # alone when it skips (non-grouped_mm backend), so only flip + # the flag when at least one bucket was actually captured. + if self.graphs: + 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 + + +# =========================================================================== +# MoE LoRA refresh — phase-4 companion to ``refresh_lora_merge_from_pristine``. +# =========================================================================== + + +def _get_moe_wrapper_tensor(wrapper): + """Return the underlying 3D expert tensor for a PEFT-wrapped MoE + parameter. Tries the common attribute paths in order.""" + if hasattr(wrapper, "get_base_layer"): + base = wrapper.get_base_layer() + if hasattr(base, "data"): + return base.data + return base + if isinstance(wrapper, torch.Tensor): + return wrapper.data + if hasattr(wrapper, "data"): + return wrapper.data + return None + + +def _pristine_moe_tensor(pristine_module, param_name): + p = getattr(pristine_module, param_name, None) + if p is None: + return None + if isinstance(p, torch.Tensor): + return p.data if hasattr(p, "data") else p + if hasattr(p, "data"): + return p.data + return p + + +def refresh_moe_lora_merge_from_pristine(base_model, peft_model): + """Batched in-place LoRA merge for Qwen3-MoE stacked expert tensors. + + For each PEFT ParamWrapper on a ``Qwen3MoeExperts.gate_up_proj`` / + ``down_proj``, compute:: + + W_inf[e] = W_pristine[e] + sum_active(scaling * B[e] @ A[e]) + + via ``torch.baddbmm`` into the same storage, mirroring the dense + ``refresh_lora_merge_from_pristine`` semantics (in-place write so + captured CUDA-graph replay reads the refreshed values). + + Handles both standard (``E, 2I, H``) and transposed (``E, H, 2I``) + stacked orientations via a runtime shape check against the flat + ``lora_A``/``lora_B`` shapes. + + Returns the count of expert tensors refreshed. No-op when no + ParamWrapper-style MoE LoRA is present (e.g. dense-only LoRA, or + a transformers version that hasn't introduced stacked experts). + """ + if base_model is None or peft_model is None: + return 0 + + inference_model = peft_model.base_model.model + n_refreshed = 0 + + for name, module in inference_model.named_modules(): + if not (hasattr(module, "gate_up_proj") and hasattr(module, "down_proj")): + continue + if not hasattr(module, "num_experts"): + continue + E = int(module.num_experts) + try: + pristine = base_model.get_submodule(name) + except AttributeError: + continue + + for param_name in ("gate_up_proj", "down_proj"): + wrapper = getattr(module, param_name, None) + pristine_data = _pristine_moe_tensor(pristine, param_name) + if wrapper is None or pristine_data is None: + continue + has_lora = hasattr(wrapper, "lora_A") and hasattr(wrapper, "lora_B") + if not has_lora: + # No PEFT wrapping — keep the plain parameter in sync + # with pristine (covers the no-LoRA case where the + # inference copy otherwise diverges via training). + W_inf = _get_moe_wrapper_tensor(wrapper) + if W_inf is not None and W_inf.shape == pristine_data.shape: + W_inf.copy_(pristine_data) + n_refreshed += 1 + continue + + W_inf = _get_moe_wrapper_tensor(wrapper) + if W_inf is None or W_inf.dim() != 3: + continue + + adapter_names = list(wrapper.lora_A.keys()) + if not adapter_names: + W_inf.copy_(pristine_data) + if hasattr(wrapper, "merged_adapters"): + wrapper.merged_adapters = [] + n_refreshed += 1 + continue + + # Determine orientation from lora shapes vs W_inf shape. + lora_A_w0 = wrapper.lora_A[adapter_names[0]].weight.data + lora_B_w0 = wrapper.lora_B[adapter_names[0]].weight.data + in_dim = lora_A_w0.shape[1] + out_dim = lora_B_w0.shape[0] + d0, d1 = W_inf.shape[1], W_inf.shape[2] + if d0 == out_dim and d1 == in_dim: + is_standard = True + elif d0 == in_dim and d1 == out_dim: + is_standard = False + else: + raise RuntimeError( + f"[refresh_moe_lora_merge_from_pristine] cannot " + f"determine orientation for {name}.{param_name}: " + f"W_inf.shape={tuple(W_inf.shape)}, " + f"in_dim={in_dim}, out_dim={out_dim}" + ) + + # Reset to pristine, then accumulate per-adapter. + W_inf.copy_(pristine_data) + + for adapter_name in adapter_names: + scaling = wrapper.scaling[adapter_name] + A_w = wrapper.lora_A[adapter_name].weight.data + B_w = wrapper.lora_B[adapter_name].weight.data + R = A_w.shape[0] // E + # A_w: (E*R, in_dim) -> A_3d: (E, R, in_dim) + A_3d = A_w.view(E, R, in_dim) + # B_w: (out_dim, E*R) -> (out_dim, E, R) -> (E, out_dim, R) + B_3d = B_w.view(out_dim, E, R).permute(1, 0, 2).contiguous() + if is_standard: + torch.baddbmm( + W_inf, + B_3d.to(W_inf.dtype), + A_3d.to(W_inf.dtype), + alpha = float(scaling), + beta = 1.0, + out = W_inf, + ) + else: + torch.baddbmm( + W_inf, + A_3d.transpose(-2, -1).contiguous().to(W_inf.dtype), + B_3d.transpose(-2, -1).contiguous().to(W_inf.dtype), + alpha = float(scaling), + beta = 1.0, + out = W_inf, + ) + + if hasattr(wrapper, "merged_adapters"): + wrapper.merged_adapters = list(adapter_names) + n_refreshed += 1 + + return n_refreshed diff --git a/unsloth/inference/flex_paged_attention.py b/unsloth/inference/flex_paged_attention.py new file mode 100644 index 0000000000..f5b3aabc89 --- /dev/null +++ b/unsloth/inference/flex_paged_attention.py @@ -0,0 +1,504 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. +# +# Adapted from attention-gym (https://github.com/pytorch-labs/attention-gym) +# Copyright (c) 2023, Driss Guessous, licensed under BSD 3-Clause +# (see THIRD_PARTY_LICENSES.md). + +# 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_5.py b/unsloth/inference/flex_qwen3_5.py new file mode 100644 index 0000000000..a44070905c --- /dev/null +++ b/unsloth/inference/flex_qwen3_5.py @@ -0,0 +1,1083 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""Flex-attention inference backend for Qwen3.5 / Qwen3.6 dense and MoE. + +Qwen3.5 / Qwen3.6 text models use a hybrid layer stack: every 4th layer is a +standard full-attention block (partial rotary_factor=0.25, attn_output_gate, +per-head Q/K RMSNorm), the remaining 3 of 4 are Gated DeltaNet linear +attention blocks backed by ``flash-linear-attention``'s Triton kernels +(``fused_recurrent_gated_delta_rule`` for decode, +``chunk_gated_delta_rule`` for prefill). + +Flex Attention can only express softmax attention, so the 75% linear_attn +layers can't use it directly. This engine therefore dispatches per layer: + +- ``full_attention`` layers: swap ``Qwen3_5Attention.forward`` with a + flex_attention-based forward bound to a shared ``PageTable`` (same + machinery as :mod:`flex_qwen3_llama`). +- ``linear_attention`` layers: keep HF's ``Qwen3_5GatedDeltaNet.forward`` + intact. It calls FLA directly when the fast path is installed and falls + back to the pure-torch reference kernels otherwise. + +State management for linear_attn layers is delegated to HF's +``DynamicCache`` populated with ``LinearAttentionAndFullAttentionLayer`` +layers. That class owns both a conv_state (size ``[B, conv_dim, 4]``) +and a recurrent_state (``[B, num_v_heads, head_k_dim, head_v_dim]`` +float32) per layer. Full-attn layers use the same layer type but the +``DynamicLayer`` half is bypassed because the flex forward writes +directly into the PageTable-backed paged KV cache. + +This is a minimal first-iteration backend. See +tests/qwen3_5_flex_parity.py for the smoke. Follow-ups: +- CUDA graph capture for the decode step. +- torch.compile on the walker. +- Qwen3.5-MoE: expert dispatch uses the existing ``forward_moe_backend`` + from moe_utils; the layer-level dispatch in this file is already + arch-agnostic to the expert kernel. +""" +from __future__ import annotations + +import os +import types +from collections import deque +from dataclasses import dataclass +from typing import Optional + +import torch +import torch.nn.functional as F +from torch.nn.attention.flex_attention import BlockMask, flex_attention + +from .flex_paged_attention import PagedKVCache, PageTable + + +# Same pattern as flex_qwen3_llama: compile once at import for warm caches. +_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 _rotate_half(x: torch.Tensor) -> torch.Tensor: + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def _apply_partial_rotary(q, k, cos, sin): + """Qwen3.5's partial_rotary_factor=0.25 form. + + ``cos`` / ``sin`` arrive shaped ``[B, S, rotary_dim]`` where + ``rotary_dim = head_dim * 0.25 = 64`` for head_dim=256. Split q / k + along ``head_dim`` into a rotary prefix and an untouched tail, + apply rotary to the prefix, concat back. Mirrors + ``modeling_qwen3_5.apply_rotary_pos_emb``. + """ + cos = cos.unsqueeze(1) # [B, 1, S, rotary_dim] + sin = sin.unsqueeze(1) + rotary_dim = cos.shape[-1] + q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] + k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] + q_rot = (q_rot * cos) + (_rotate_half(q_rot) * sin) + k_rot = (k_rot * cos) + (_rotate_half(k_rot) * sin) + return torch.cat([q_rot, q_pass], dim=-1), torch.cat([k_rot, k_pass], dim=-1) + + +def make_qwen3_5_flex_attention_forward(page_table: PageTable): + """Swap-in for ``Qwen3_5Attention.forward`` that uses flex_attention + + paged KV cache. Handles: + - attn_output_gate: ``q_proj`` outputs ``2 * hidden_shape`` and we + chunk into (query, gate); the attn_output is multiplied by + ``sigmoid(gate)`` before ``o_proj``. + - per-head Q/K RMSNorm before rotary. + - partial rotary (factor=0.25) via :func:`_apply_partial_rotary`. + """ + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings, + attention_mask=None, + past_key_values=None, + flex_block_mask: Optional[BlockMask] = None, + flex_input_pos: Optional[torch.Tensor] = None, + flex_batch_idx: Optional[torch.Tensor] = None, + flex_kernel_options: Optional[dict] = None, + **kwargs, + ): + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + # q_proj output: [..., n_heads, 2 * head_dim]; split into q / gate. + q_and_gate = self.q_proj(hidden_states).view(*input_shape, -1, self.head_dim * 2) + query_states, gate = torch.chunk(q_and_gate, 2, dim=-1) + gate = gate.reshape(*input_shape, -1) # [..., n_heads * head_dim] + + # q_norm / k_norm apply per-head RMSNorm on the head_dim axis. + query_states = self.q_norm(query_states.view(hidden_shape)).transpose(1, 2) + key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = _apply_partial_rotary( + query_states, key_states, 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 + # B decode slots. + if self._paged_cache is not None and flex_input_pos is not None: + cache_dtype = self._paged_cache.k_cache.dtype + if key_states.dtype != cache_dtype: + key_states = key_states.to(cache_dtype) + if value_states.dtype != cache_dtype: + value_states = value_states.to(cache_dtype) + key_states, value_states = self._paged_cache.update( + flex_input_pos, key_states, value_states, flex_batch_idx, + ) + + attn_output = flex_attention_compiled( + query_states, + key_states, + value_states, + 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() + attn_output = attn_output * torch.sigmoid(gate) + attn_output = self.o_proj(attn_output) + return attn_output, None + + return forward + + +def _patch_qwen3_5_full_attn_forwards(text_model, page_table: PageTable): + """Walk ``text_model.layers`` and, for each layer whose ``layer_type == + 'full_attention'``, attach a ``PagedKVCache`` and swap in the flex + attention forward. Linear_attn layers are left alone (HF forward + dispatches to FLA directly). + + ``text_model`` is ``Qwen3_5TextModel`` or ``Qwen3_5MoeTextModel`` — + both expose ``.layers`` / ``.config``. The outer multimodal + wrapper's ``.model.language_model`` is what the engine resolves and + passes in. + """ + cfg = text_model.config + fwd = make_qwen3_5_flex_attention_forward(page_table) + for layer_idx, layer in enumerate(text_model.layers): + if cfg.layer_types[layer_idx] != "full_attention": + continue + attn = layer.self_attn + attn._paged_cache = PagedKVCache( + page_table, + n_heads=cfg.num_key_value_heads, + head_dim=cfg.head_dim, + dtype=text_model.dtype, + ).to(text_model.dtype).to(next(text_model.parameters()).device) + attn.forward = types.MethodType(fwd, attn) + + +def _resolve_text_model(model: torch.nn.Module): + """Qwen3.5 / Qwen3.6 ship the text backbone two levels down: + ``Qwen3_5ForConditionalGeneration.model.language_model`` is the + Qwen3_5TextModel (or Qwen3_5MoeTextModel) that owns ``.layers``, + ``.embed_tokens``, ``.norm``, ``.rotary_emb``. + """ + base = getattr(model, "model", model) + lang = getattr(base, "language_model", None) + if lang is not None and hasattr(lang, "layers"): + return lang + if hasattr(base, "layers"): + return base + # PEFT wrapper. + inner = getattr(model, "base_model", None) + if inner is not None: + return _resolve_text_model(getattr(inner, "model", inner)) + raise RuntimeError( + "Cannot locate the Qwen3.5 text backbone (expected " + "``.model.language_model.layers``)" + ) + + +def _torch_causal_conv1d_update( + hidden_states: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor] = None, +): + """CUDA-graph-safe causal conv1d update step for Qwen3.5 DeltaNet. + + Mirrors ``transformers.models.qwen3_5.torch_causal_conv1d_update``: + rolls ``conv_state`` forward by one step in place, convolves the + (state || hidden_states) window with the depthwise conv1d weight, then + applies SiLU. Used in place of the ``causal_conv1d_update`` kernel which + writes via a CUDA memcpy that cudagraph-capture refuses on some inputs. + """ + _, hidden_size, seq_len = hidden_states.shape + state_len = conv_state.shape[-1] + hs_ext = torch.cat([conv_state, hidden_states], dim=-1).to(weight.dtype) + conv_state.copy_(hs_ext[:, :, -state_len:]) + out = F.conv1d( + hs_ext, weight.unsqueeze(1), bias, padding=0, groups=hidden_size, + ) + out = F.silu(out[:, :, -seq_len:]) + return out.to(hidden_states.dtype) + + +def _deltanet_decode_step_static( + layer, + hidden_states: torch.Tensor, + conv_state_gathered: torch.Tensor, + recurrent_state_gathered: torch.Tensor, +): + """Single-token decode for ``Qwen3_5GatedDeltaNet`` with caller-owned + state tensors (no ``DynamicCache``). ``conv_state_gathered`` is updated + in place by :func:`_torch_causal_conv1d_update`; the recurrent state + new value is freshly allocated and returned for the caller to scatter. + + Shapes: + ``hidden_states`` [B, 1, H] + ``conv_state_gathered`` [B, conv_dim, K] (bf16) + ``recurrent_state_gathered`` [B, HV, Hk, Hv] (fp32) + Returns: output [B, 1, H], updated_conv_state, new_recurrent_state. + """ + B, S, _ = hidden_states.shape + mixed_qkv = layer.in_proj_qkv(hidden_states).transpose(1, 2) # [B, conv_dim, 1] + z = layer.in_proj_z(hidden_states).reshape(B, S, -1, layer.head_v_dim) + b = layer.in_proj_b(hidden_states) + a = layer.in_proj_a(hidden_states) + + mixed_qkv = _torch_causal_conv1d_update( + mixed_qkv, + conv_state_gathered, + layer.conv1d.weight.squeeze(1), + layer.conv1d.bias, + ) + + mixed_qkv = mixed_qkv.transpose(1, 2) # [B, 1, conv_dim] + query, key, value = torch.split( + mixed_qkv, + [layer.key_dim, layer.key_dim, layer.value_dim], + dim=-1, + ) + query = query.reshape(B, S, -1, layer.head_k_dim) + key = key.reshape(B, S, -1, layer.head_k_dim) + value = value.reshape(B, S, -1, layer.head_v_dim) + + beta = b.sigmoid() + g = -layer.A_log.float().exp() * F.softplus(a.float() + layer.dt_bias) + r = layer.num_v_heads // layer.num_k_heads + if r > 1: + query = query.repeat_interleave(r, dim=2) + key = key.repeat_interleave(r, dim=2) + + core_attn_out, last_recurrent_state = layer.recurrent_gated_delta_rule( + query, key, value, g=g, beta=beta, + initial_state=recurrent_state_gathered, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + ) + + core_attn_out = core_attn_out.reshape(-1, layer.head_v_dim) + z = z.reshape(-1, layer.head_v_dim) + core_attn_out = layer.norm(core_attn_out, z) + core_attn_out = core_attn_out.reshape(B, S, -1) + output = layer.out_proj(core_attn_out) + return output, conv_state_gathered, last_recurrent_state + + +def _call_qwen3_5_decode_static( + text_model, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + full_attn_flex_kwargs: dict, + batch_idx: torch.Tensor, + linear_conv_states: dict, + linear_recurrent_states: dict, + *, + lm_head_fn, +): + """Batched decode walker using static per-layer state buffers. + + Linear layers ``index_select`` → deltanet-decode → ``index_copy_`` back. + Full-attention layers use the pre-patched flex+paged KV forward. + """ + cfg = text_model.config + inputs_embeds = text_model.embed_tokens(input_ids) + position_embeddings = text_model.rotary_emb(inputs_embeds, position_ids) + + hidden_states = inputs_embeds + layer_types = cfg.layer_types + for layer_idx, layer in enumerate(text_model.layers): + residual = hidden_states + hidden_states = layer.input_layernorm(hidden_states) + + if layer_types[layer_idx] == "linear_attention": + conv_buf = linear_conv_states[layer_idx] + rec_buf = linear_recurrent_states[layer_idx] + conv_gath = conv_buf.index_select(0, batch_idx) + rec_gath = rec_buf.index_select(0, batch_idx) + out, conv_new, rec_new = _deltanet_decode_step_static( + layer.linear_attn, hidden_states, conv_gath, rec_gath, + ) + conv_buf.index_copy_(0, batch_idx, conv_new) + rec_buf.index_copy_(0, batch_idx, rec_new) + hidden_states = out + else: + hidden_states, _ = layer.self_attn( + hidden_states, + position_embeddings=position_embeddings, + **full_attn_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 = text_model.norm(hidden_states) + return lm_head_fn(hidden_states) + + +def _call_qwen3_5_with_flex_kwargs( + text_model, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + full_attn_flex_kwargs: dict, + linear_cache, + linear_attn_mask, + *, + lm_head_fn, +): + """Per-layer walker. + + - ``full_attention`` layers: called with ``**full_attn_flex_kwargs`` + (block_mask / input_pos / batch_idx / kernel_options). The patched + attention forward consumes these; ``past_key_values`` is ignored. + - ``linear_attention`` layers: HF's ``Qwen3_5GatedDeltaNet.forward`` + needs ``cache_params=`` and the linear_attn_mask + (left-padding aware, ``None`` during cached decode). + - Residual / MLP / layernorms are the stock HF modules so we run + them inline without swapping. + """ + cfg = text_model.config + inputs_embeds = text_model.embed_tokens(input_ids) + + # Qwen3_5TextRotaryEmbedding internally expands 2D position_ids to + # 3D via [None, ...].expand(3, B, S). For text-only, all 3 copies + # are identical so the result collapses to standard rope on the + # first ``partial_rotary_factor`` fraction of head_dim. + position_embeddings = text_model.rotary_emb(inputs_embeds, position_ids) + + hidden_states = inputs_embeds + layer_types = cfg.layer_types + + for layer_idx, layer in enumerate(text_model.layers): + residual = hidden_states + hidden_states = layer.input_layernorm(hidden_states) + + if layer_types[layer_idx] == "linear_attention": + hidden_states = layer.linear_attn( + hidden_states=hidden_states, + cache_params=linear_cache, + attention_mask=linear_attn_mask, + ) + else: + hidden_states, _ = layer.self_attn( + hidden_states, + position_embeddings=position_embeddings, + **full_attn_flex_kwargs, + ) + + hidden_states = residual + hidden_states + + # MLP path. For dense Qwen3_5 the layer has ``layer.mlp``; for + # Qwen3_5Moe it has the same name, but the submodule is a + # Qwen3_5MoeSparseMoeBlock that owns gate + experts + shared + # expert. Both implement ``forward(hidden_states) -> hidden`` + # so we don't dispatch here. + residual = hidden_states + hidden_states = layer.post_attention_layernorm(hidden_states) + hidden_states = layer.mlp(hidden_states) + hidden_states = residual + hidden_states + + hidden_states = text_model.norm(hidden_states) + return lm_head_fn(hidden_states) + + +@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 + _linear_cache = None # per-seq DynamicCache for linear_attn state + + 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) + + +DECODE_KERNEL_OPTIONS_DEFAULT = None +PREFILL_KERNEL_OPTIONS_DEFAULT = {"FORCE_USE_FLEX_ATTENTION": True} + + +class FlexQwen3_5Inference: + """Flex-attention backend for Qwen3.5 / Qwen3.6 dense and MoE.""" + + arch = "qwen3_5" + + def __init__( + self, + model, + tokenizer, + max_batch_size: int = 8, + max_seq_length: int = 2048, + n_pages: int = 2048, + page_size: int = 128, + max_new_tokens: int = 512, + decode_kernel_options=None, + prefill_kernel_options=None, + fa4_prefill=None, + base_model=None, + peft_model=None, + cumem_allocator=None, + ): + assert max_seq_length % page_size == 0 + self.model = model + self.tokenizer = tokenizer + self.device = next(model.parameters()).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 + # PageTable reserves batch_idx=0 as a no-op slot, so the first + # ``max_batch_size`` user-allocatable slots are 1..max_batch_size. + # Bump the page table's capacity by 1 so the user actually gets + # ``max_batch_size`` concurrent sequences (otherwise bs=N requests + # serialise into bs=N-1 + bs=1, collapsing throughput — see + # commit on this file for the bs=8 regression diagnosis). + page_table_cap = max_batch_size + 1 + self._page_table_cap = page_table_cap + self.max_seq_length = max_seq_length + self.page_size = page_size + self.max_new_tokens = max_new_tokens + + self.fa4_prefill = bool(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 + + # Resolve the text backbone so we walk .layers / .rotary_emb + # directly, bypassing the outer multimodal wrapper's complex + # forward. lm_head still lives on the outer wrapper. + self.text_model = _resolve_text_model(model) + if hasattr(model, "lm_head"): + self._lm_head = model.lm_head + else: + # Multimodal wrapper routes lm_head through .model. + self._lm_head = model.model.lm_head + + # Page table for the full-attention layers' paged KV cache. + from .sleep_mode import kv_cache_pool as _kv_cache_pool + with _kv_cache_pool(cumem_allocator): + self.page_table = PageTable( + n_pages=n_pages, + page_size=page_size, + max_batch_size=page_table_cap, + device=self.device.type, + ) + _patch_qwen3_5_full_attn_forwards(self.text_model, self.page_table) + + self.input_pos_buffer = torch.zeros( + page_table_cap, dtype=torch.int32, device=self.device, + ) + self.block_mask_logical = self.page_table.create_causal_blockmask( + B=page_table_cap, L=max_seq_length, + ) + + # Pre-allocate per-slot LinearAttention caches. Each sequence's + # decode loop reuses the same DynamicCache object for its slot; + # we reset it via .reset() at prefill so the conv/recurrent + # states start clean. + self._linear_caches = [ + self._build_linear_cache() for _ in range(page_table_cap) + ] + + # Static state buffers for CUDA-graph-captured batched decode. + # We keep one conv_state tensor and one recurrent_state tensor per + # linear_attention layer, sized ``[max_batch_size, ...]``. Prefill + # still writes into per-seq ``DynamicCache`` objects; after + # prefill, :meth:`_sync_static_from_dynamic_cache` copies the + # prefilled state into the right slot of these buffers, and the + # batched decode step gather/scatters through them. + cfg = self.text_model.config + self._linear_layer_indices = [ + i for i, t in enumerate(cfg.layer_types) if t == "linear_attention" + ] + self._linear_conv_states = {} + self._linear_recurrent_states = {} + for layer_idx in self._linear_layer_indices: + la = self.text_model.layers[layer_idx].linear_attn + self._linear_conv_states[layer_idx] = torch.zeros( + page_table_cap, la.conv_dim, la.conv_kernel_size, + dtype=la.conv1d.weight.dtype, device=self.device, + ) + self._linear_recurrent_states[layer_idx] = torch.zeros( + page_table_cap, la.num_v_heads, la.head_k_dim, la.head_v_dim, + dtype=torch.float32, device=self.device, + ) + try: + torch._dynamo.mark_static_address( + self._linear_conv_states[layer_idx], + ) + torch._dynamo.mark_static_address( + self._linear_recurrent_states[layer_idx], + ) + except Exception: + pass + + # CUDA-graph state. Populated by :meth:`capture_decode_cudagraph` + # on first decode when ``capture_cudagraph=True`` is passed to + # :meth:`generate`. + self.graphs = {} + self.graph_vars = None + self._captured = False + + # ----- cache helpers ------------------------------------------------- + def _build_linear_cache(self): + """Allocate a ``DynamicCache`` seeded from ``model.config`` so + ``config.layer_types`` dictates the per-layer cache class: + ``"linear_attention"`` → ``LinearAttentionLayer`` + (owns conv/recurrent state) and ``"full_attention"`` → + ``DynamicLayer`` (dense KV). The full-attn DynamicLayer half is + untouched in this engine (the flex forward writes directly into + the paged KV cache) — only the linear layers actually see use. + """ + from transformers.cache_utils import DynamicCache + cache = DynamicCache(config=self.model.config) + return cache + + # ----- prefill ------------------------------------------------------- + def tokenize(self, sequences): + for seq in sequences: + if seq.input_ids is not None and seq.input_length > 0: + continue + ids = self.tokenizer(seq.text, return_tensors="pt")["input_ids"].squeeze(0) + seq.input_ids = ids + seq.input_length = ids.shape[0] + + @torch.inference_mode() + def _prefill_single(self, seq: Sequence) -> int: + """One-sequence prefill. Unlike :class:`FlexInference`, we can't + trivially pack multiple linear_attn sequences into a [1, L_total] + run because the conv1d and recurrent state carry independent + context per sequence. First iteration runs one prefill per seq. + Full-attn layers still use document-causal paging internally so + the paged KV cache is populated correctly.""" + device = self.device + input_ids = seq.input_ids.to(device).view(1, -1) + L = seq.input_length + batch_idx = torch.full((1, L), seq.batch_idx, dtype=torch.long, device=device) + input_pos = torch.arange(L, dtype=torch.long, device=device).view(1, -1) + + # Pad packed prefill to the flex_attention Q-block boundary. + 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=seq.batch_idx) + + 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, + ) + full_attn_kwargs = dict( + flex_block_mask=mask, + flex_input_pos=input_pos, + flex_batch_idx=batch_idx, + flex_kernel_options=self.prefill_kernel_options, + ) + linear_cache = seq._linear_cache + # HF uses torch.all(attention_mask == 1) to decide to drop the + # mask. For a fresh prefill pass None, meaning the DeltaNet + # layer computes without a padding mask (our single packed + # sequence has no padding). + logits = _call_qwen3_5_with_flex_kwargs( + self.text_model, + input_ids, + input_pos, + full_attn_kwargs, + linear_cache, + linear_attn_mask=None, + lm_head_fn=self._lm_head, + ) + # Logits at the last real token (before pad) — position L-1. + return int(torch.argmax(logits[0, L - 1, :]).item()) + + @torch.inference_mode() + def _decode_step_eager( + self, + batch_idx: torch.Tensor, + input_ids: torch.Tensor, + linear_caches: list, + ): + """Single decode step across a batch of sequences. We run the + DeltaNet layers with each sequence's own cache sequentially — + the FLA recurrent kernel is per-sequence. The full_attn layers + are batched (paged KV + flex_attention). + + Returns logits [B, V]. + """ + 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) + + full_attn_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, + ) + + # Path B: per-sequence decode. The FLA recurrent kernel expects + # a single ``initial_state`` tensor per call; the state lives in + # each sequence's DynamicCache. Batched decode across different + # cache tensors would require fused batched FLA kernels which + # aren't wired up yet. Serialize the DeltaNet halves of the + # step, batch the full-attn halves. + # + # Simplest working implementation: decode one sequence at a time + # end-to-end. The per-seq paged block mask already slices to + # that row, so the full-attn math still runs correctly. + logits_list = [] + for i in range(B): + single_bi = batch_idx[i : i + 1] + single_ids = input_ids[i : i + 1].view(1, 1) + single_mask, single_pos = self._decode_block_mask(single_bi) + single_mask = self.page_table.convert_logical_block_mask( + single_mask, single_bi, + ) + single_kwargs = dict( + flex_block_mask=single_mask, + flex_input_pos=single_pos.view(1, 1).to(torch.long), + flex_batch_idx=single_bi, + flex_kernel_options=self.decode_kernel_options, + ) + logits = _call_qwen3_5_with_flex_kwargs( + self.text_model, + single_ids, + single_pos.view(1, 1).to(torch.long), + single_kwargs, + linear_caches[i], + linear_attn_mask=None, + lm_head_fn=self._lm_head, + ) + logits_list.append(logits[0, -1, :]) + return torch.stack(logits_list, dim=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. Copy of the same helper + in ``FlexInference`` — dedup if this pattern grows a third + consumer.""" + block_mask = self.block_mask_logical + input_pos = self.input_pos_buffer[batch_idx] + 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 _reset_linear_cache(self, seq: Sequence): + """Fresh linear-attn cache for a newly-scheduled sequence.""" + seq._linear_cache = self._build_linear_cache() + + # ----- static-state decode ------------------------------------------ + def _sync_static_from_dynamic_cache(self, seq: Sequence): + """Copy prefilled conv/recurrent state from ``seq._linear_cache`` + into the engine's static per-layer buffers at ``seq.batch_idx``. + + The HF prefill path writes into each seq's ``DynamicCache``; to + hand off to the batched/captured decode we copy those tensors into + the fixed slot in the static buffer.""" + slot = seq.batch_idx + cache = seq._linear_cache + for layer_idx in self._linear_layer_indices: + layer_cache = cache.layers[layer_idx] + conv = getattr(layer_cache, "conv_states", None) + rec = getattr(layer_cache, "recurrent_states", None) + if conv is not None and conv.numel() > 0: + self._linear_conv_states[layer_idx][slot].copy_(conv[0]) + if rec is not None and rec.numel() > 0: + self._linear_recurrent_states[layer_idx][slot].copy_(rec[0]) + + def _reset_static_state(self, batch_idx: int): + """Zero the static state for a slot (e.g. when a seq finishes).""" + for layer_idx in self._linear_layer_indices: + self._linear_conv_states[layer_idx][batch_idx].zero_() + self._linear_recurrent_states[layer_idx][batch_idx].zero_() + + @torch.inference_mode() + def _decode_step_batched_static( + self, batch_idx: torch.Tensor, input_ids: torch.Tensor, + ) -> torch.Tensor: + """Eager batched decode using the engine's static state buffers. + + Walks every transformer layer once on the batch: full-attn layers + go through the paged KV + flex attention forward; linear layers + gather/scatter through static buffers. Returns logits ``[B, V]`` + at the last position. + """ + 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) + ids = input_ids.view(B, 1) + pos = input_pos.view(B, 1).to(torch.long) + full_attn_kwargs = dict( + flex_block_mask=mask, + flex_input_pos=pos, + flex_batch_idx=batch_idx, + flex_kernel_options=self.decode_kernel_options, + ) + logits = _call_qwen3_5_decode_static( + self.text_model, + ids, + pos, + full_attn_kwargs, + batch_idx, + self._linear_conv_states, + self._linear_recurrent_states, + lm_head_fn=self._lm_head, + ) + return logits[:, -1, :] + + def capture_decode_cudagraph(self): + """Capture decode-step CUDA graphs across a bucket ladder. + + Reserves dummy page-table slots so the paged KV cache machinery + has somewhere to write during the capture warmups, then captures + one graph per bucket size sharing a single memory pool. Buckets + are ``[1, 2, 4, 8, 16, 32, ...]`` capped at ``max_batch_size``. + + After capture the engine zeros the static conv/recurrent states + and releases the dummy page reservations so generation starts + clean. + """ + max_bs = self.max_batch_size + # Temporarily reserve page-table slots 0..N so the decode block + # mask builder has state to index into. This only matters at + # capture time — we release these after. + reserved = [] + for _ in range(max_bs): + try: + bi = self.page_table.allocate() + self.page_table.reserve( + bi, + torch.tensor([bi], device=self.device, dtype=torch.long), + self.page_size, + ) + reserved.append(bi) + except Exception: + break + if not reserved: + raise RuntimeError( + "capture_decode_cudagraph: could not allocate any " + "page-table slots for capture.", + ) + + # Static input + output buffers. + input_ids_buf = torch.zeros( + max_bs, dtype=torch.int64, device=self.device, + ) + # Use reserved slot ids as the default indirection; replay time + # overrides with ``.copy_``. + bi_init = reserved + list(range(len(reserved), max_bs)) + batch_idx_buf = torch.tensor( + bi_init[:max_bs], dtype=torch.int64, device=self.device, + ) + vocab_size = getattr( + self.text_model.config, "vocab_size", None, + ) or getattr(self.model.config, "vocab_size", None) + if vocab_size is None: + # Multimodal wrapper: vocab_size lives on text_config. + vocab_size = self.model.config.text_config.vocab_size + outputs_buf = torch.zeros( + (max_bs, vocab_size), + dtype=self.model.dtype, + device=self.device, + ) + try: + torch._dynamo.mark_static_address(input_ids_buf) + torch._dynamo.mark_static_address(batch_idx_buf) + torch._dynamo.mark_static_address(outputs_buf) + except Exception: + pass + + # Bucket ladder. Start small, expand to max_bs. + ladder = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16)) + ladder = sorted(set(bs for bs in ladder if bs <= max_bs)) + self.graph_bs = ladder + + pool = None + for bs in reversed(ladder): + torch.cuda.synchronize() + # Warmup eager run seeds the flex_attention compiled cache + # for this bs and populates any autograd-off kernels. + _ = self._decode_step_batched_static( + batch_idx_buf[:bs], input_ids_buf[:bs], + ) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + if pool is None: + with torch.cuda.graph(graph): + out = self._decode_step_batched_static( + batch_idx_buf[:bs], input_ids_buf[:bs], + ) + outputs_buf[:bs].copy_(out) + pool = graph.pool() + else: + with torch.cuda.graph(graph, pool=pool): + out = self._decode_step_batched_static( + batch_idx_buf[:bs], input_ids_buf[:bs], + ) + outputs_buf[:bs].copy_(out) + self.graphs[bs] = graph + torch.cuda.synchronize() + + # Release capture reservations + wipe polluted static state. + for bi in reserved: + self.page_table.erase(bi) + for layer_idx in self._linear_layer_indices: + self._linear_conv_states[layer_idx].zero_() + self._linear_recurrent_states[layer_idx].zero_() + + self.graph_vars = dict( + input_ids=input_ids_buf, + batch_idx=batch_idx_buf, + outputs=outputs_buf, + ) + self._captured = True + + def _pick_bucket(self, B: int) -> Optional[int]: + """Exact-match bucket lookup. Padding the batch with duplicate + ``batch_idx`` entries breaks the ``index_copy_`` scatter into the + static state buffers (duplicate indices make the last write win, + so the primary slot's conv/recurrent state gets overwritten by + the padding slot's advance). Requiring exact match keeps the + captured replay correct and falls back to the eager batched + decode for sizes without a dedicated graph.""" + if not self._captured: + return None + if B in self.graphs: + return B + return None + + @torch.inference_mode() + def _decode_step( + self, batch_idx: torch.Tensor, input_ids: torch.Tensor, + ) -> torch.Tensor: + """Decode dispatch: captured graph replay when the active batch + size exactly matches a captured bucket, otherwise eager batched + decode. See :meth:`_pick_bucket` for why we require an exact + match (no padding).""" + B = batch_idx.shape[0] + bucket = self._pick_bucket(B) + if bucket is None: + return self._decode_step_batched_static(batch_idx, input_ids) + gv = self.graph_vars + gv["input_ids"][:bucket].copy_(input_ids) + gv["batch_idx"][:bucket].copy_(batch_idx) + self.graphs[bucket].replay() + return gv["outputs"][:B].clone() + + # ----- generate loop ------------------------------------------------- + @torch.inference_mode() + def generate(self, sequences: list, capture_cudagraph: bool = False): + """Main entry. + + Prefill is single-seq through HF's DeltaNet forward (writes into + per-seq ``DynamicCache``). After each prefill we sync that state + into the engine's static buffers so decode can run batched. + Decode is batched across all active sequences. If + ``capture_cudagraph=True`` and the bucket ladder hasn't been + captured yet, we capture on the first decode; subsequent decodes + replay the matching graph. + """ + # ``UNSLOTH_FLEX_QWEN3_5_NO_CAPTURE=1`` forces the eager batched + # decode path for debugging (bypasses graph capture entirely). + if os.environ.get("UNSLOTH_FLEX_QWEN3_5_NO_CAPTURE", "0") == "1": + capture_cudagraph = False + if capture_cudagraph and not self._captured: + self.capture_decode_cudagraph() + self.tokenize(sequences) + waiting = deque(sequences) + running = deque() + done = [] + + while waiting or running: + # Schedule waiting -> running via prefill. + if 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 + self._reset_linear_cache(seq) + self._reset_static_state(bi) + next_id = self._prefill_single(seq) + # Hand prefilled conv/recurrent state over to the static + # buffers so the batched decode step can read it. + self._sync_static_from_dynamic_cache(seq) + seq.last_token_id = next_id + seq.output_ids.append(next_id) + 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 + + if not running: + # No running seqs yet and prefill budget exhausted; wait. + # In practice the page table can always reserve at least + # one pending prompt thanks to the headroom factor. + break + + # Decode step for everything running. + decode_batch = [] + while running: + seq = running.popleft() + needed = seq.total_length + if self.page_table.capacity[seq.batch_idx] >= needed: + decode_batch.append(seq) + elif self.page_table.can_reserve( + needed, batch_idx_int=seq.batch_idx, + ): + self.page_table.reserve( + seq.batch_idx, + torch.tensor( + [seq.batch_idx], + device=self.device, + dtype=torch.long, + ), + needed, + ) + decode_batch.append(seq) + else: + # Evict: push newest back to waiting. + 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, + ) + self.input_pos_buffer.zero_() + self.input_pos_buffer[bi_tensor] = cur_pos + + logits = self._decode_step(bi_tensor, last_ids) + 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 + + +__all__ = ["FlexQwen3_5Inference", "Sequence"] diff --git a/unsloth/inference/flex_qwen3_llama.py b/unsloth/inference/flex_qwen3_llama.py new file mode 100644 index 0000000000..a13caf449a --- /dev/null +++ b/unsloth/inference/flex_qwen3_llama.py @@ -0,0 +1,1287 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""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. + # Match the pre-allocated KV cache dtype; bnb-4bit Linear compute + # can produce fp32 k/v even under autocast, and the paged-cache + # index_put_ refuses mixed dtypes. + if self._paged_cache is not None and flex_input_pos is not None: + cache_dtype = self._paged_cache.k_cache.dtype + if k.dtype != cache_dtype: + k = k.to(cache_dtype) + if v.dtype != cache_dtype: + v = v.to(cache_dtype) + 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, + cumem_allocator = 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 + + # Route the paged-KV allocations through the cuMem "kv_cache" + # pool when sleep mode is active, so FlexEngine.sleep can drop + # them and wake_up can re-map fresh zeroed pages at the same + # virtual addresses. The block_mask / input_pos scratch below + # stays in the default allocator so captured CUDA graphs that + # reference them stay valid across sleep / wake. + from .sleep_mode import kv_cache_pool as _kv_cache_pool + + with _kv_cache_pool(cumem_allocator): + 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/sleep_mode.py b/unsloth/inference/sleep_mode.py new file mode 100644 index 0000000000..92b0f43ba0 --- /dev/null +++ b/unsloth/inference/sleep_mode.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""Sleep-mode helpers for the flex inference backend. + +When ``UNSLOTH_VLLM_STANDBY=1`` is set and vLLM's ``CuMemAllocator`` is +importable, :class:`~unsloth.inference.flex_engine.FlexEngine` routes its +heavy GPU allocations (the inference deep-copies, the PEFT wrapper, and +per-layer :class:`~unsloth.inference.flex_paged_attention.PagedKVCache` +buffers) through cuMem-backed pools. ``sleep(level=1)`` then offloads the +``weights`` pool to pinned CPU memory and discards the ``kv_cache`` pool +outright; ``wake_up`` re-maps the handles at the same virtual addresses +so previously captured CUDA graphs and compiled artifacts stay valid. + +If vLLM is not installed, :func:`_get_cumem_allocator` returns ``None`` +and the helpers fall back to :func:`contextlib.nullcontext`, which keeps +the flex engine working without sleep support. +""" + +from __future__ import annotations + +import contextlib +import os +import warnings +from typing import Any, Optional + + +_WARNED_NO_VLLM = False + + +def _get_cumem_allocator() -> Optional[Any]: + """Return ``vllm.device_allocator.cumem.CuMemAllocator.get_instance()`` + or ``None`` if vLLM is not importable. + + Emits a single warning on the first failed import so users who opt + into sleep mode with ``UNSLOTH_VLLM_STANDBY=1`` see a clear message + about the soft dependency on vLLM.""" + global _WARNED_NO_VLLM + try: + from vllm.device_allocator.cumem import CuMemAllocator + except Exception as e: + if not _WARNED_NO_VLLM: + warnings.warn( + "FlexEngine sleep mode requires vLLM's CuMemAllocator " + f"(import failed: {e}). Sleep / wake_up will be no-ops. " + "Install vLLM to enable level-1 sleep mode on the flex " + "backend.", + RuntimeWarning, + stacklevel = 2, + ) + _WARNED_NO_VLLM = True + return None + return CuMemAllocator.get_instance() + + +def sleep_mode_enabled() -> bool: + """``True`` iff ``UNSLOTH_VLLM_STANDBY=1`` is set AND vLLM is + available. Evaluated at :class:`FlexEngine.__init__` time so the + choice of allocator is stable for the engine's lifetime.""" + if os.environ.get("UNSLOTH_VLLM_STANDBY", "0") != "1": + return False + return _get_cumem_allocator() is not None + + +def _pool(allocator: Optional[Any], tag: str): + """Return a ``use_memory_pool(tag=...)`` context manager if the + allocator is available, else ``nullcontext`` so callers can wrap + allocation sites unconditionally.""" + if allocator is None: + return contextlib.nullcontext() + return allocator.use_memory_pool(tag = tag) + + +def weight_pool(allocator: Optional[Any]): + """Context manager for ``tag="weights"`` allocations (offloaded to + pinned CPU on sleep, restored via ``cudaMemcpy`` on wake).""" + return _pool(allocator, "weights") + + +def kv_cache_pool(allocator: Optional[Any]): + """Context manager for ``tag="kv_cache"`` allocations (discarded on + sleep, re-mapped with zeros on wake).""" + return _pool(allocator, "kv_cache") + + +def describe_sleep_state(engine) -> dict: + """Return a small dict summarising the flex engine's sleep-mode + state plus a torch CUDA memory snapshot. Useful from bench scripts + and tests.""" + import torch + + allocator = getattr(engine, "_cumem_allocator", None) + state: dict = { + "sleep_mode_enabled": bool(getattr(engine, "_sleep_mode_enabled", False)), + "cumem_allocator": type(allocator).__name__ if allocator is not None else None, + } + if torch.cuda.is_available(): + state["allocated_gb"] = round(torch.cuda.memory_allocated() / 1e9, 3) + state["reserved_gb"] = round(torch.cuda.memory_reserved() / 1e9, 3) + if allocator is not None and hasattr(allocator, "get_current_usage"): + try: + state["cumem_current_usage"] = allocator.get_current_usage() + except Exception: + pass + return state + + +__all__ = [ + "_get_cumem_allocator", + "sleep_mode_enabled", + "weight_pool", + "kv_cache_pool", + "describe_sleep_state", +] diff --git a/unsloth/inference/vllm_shim.py b/unsloth/inference/vllm_shim.py new file mode 100644 index 0000000000..13a3960453 --- /dev/null +++ b/unsloth/inference/vllm_shim.py @@ -0,0 +1,172 @@ +# SPDX-License-Identifier: GNU Affero General Public License v3.0 +# Copyright 2023-present the Unsloth team. All rights reserved. + +"""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..42bb9dc67b 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -1397,12 +1397,26 @@ def _LlamaModel_fast_forward_inference( XX2 = XX2, variance = variance, ) - X = mlp_fast_forward_inference( - decoder_layer.mlp, - X, - temp_gate = temp_gates[device_index], - temp_up = temp_ups[device_index], - ) + # MoE blocks (Qwen3MoeSparseMoeBlock, etc.) do not have the + # dense gate_proj / up_proj / down_proj attributes that + # mlp_fast_forward_inference requires. Delegate to the + # class-level forward (patched by unsloth_zoo for MoE) and + # unpack the (hidden_states, router_logits) tuple. + _mlp_mod = decoder_layer.mlp + if not ( + hasattr(_mlp_mod, "gate_proj") + and hasattr(_mlp_mod, "up_proj") + and hasattr(_mlp_mod, "down_proj") + ): + _mlp_out = _mlp_mod(X) + X = _mlp_out[0] if isinstance(_mlp_out, tuple) else _mlp_out + else: + X = mlp_fast_forward_inference( + _mlp_mod, + X, + temp_gate = temp_gates[device_index], + temp_up = temp_ups[device_index], + ) X += residual next_decoder_cache.append(present_key_value) @@ -2488,7 +2502,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 +2533,48 @@ 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. + # + # If ``UNSLOTH_VLLM_STANDBY=1`` is set AND vLLM is + # importable, wrap the deep-copy in the cuMem + # ``weights`` pool so ``FlexEngine.sleep(level=1)`` can + # offload it to pinned CPU memory without invalidating + # the engine's captured CUDA graphs. This must happen at + # the deep-copy site, not later in FlexEngine.__init__, + # because the engine receives the copy as + # ``inference_model=`` and never re-allocates it. + import copy as _copy + from unsloth.inference.sleep_mode import ( + _get_cumem_allocator as _flex_get_cumem, + sleep_mode_enabled as _flex_sleep_enabled, + weight_pool as _flex_weight_pool, + ) + + _flex_allocator = _flex_get_cumem() if _flex_sleep_enabled() else None + with _flex_weight_pool(_flex_allocator): + 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 +2647,20 @@ class FastLlamaModel: model, tokenizer, correct_dtype = dtype ) + # UNSLOTH_FAST_INFERENCE=1 path: install the lazy ``vllm_engine`` + # sentinel now that the tokenizer is available. The real + # ``FlexEngine(...)`` construction is deferred to + # :func:`build_flex_engine` so the batch-size dimension can be + # sized from the GRPO rollout shape instead of frozen here at a + # guessed default. 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. + if hasattr(model, "_unsloth_needs_flex_engine"): + from unsloth.inference.flex_engine import install_flex_sentinel + + install_flex_sentinel(model, tokenizer) + # Patch up QKV / O and MLP for idx, layer in enumerate(model.model.layers): layer.self_attn.apply_qkv = original_apply_qkv @@ -3290,6 +3367,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) @@ -3338,7 +3422,7 @@ class FastLlamaModel: apply_lora_mlp = apply_lora_mlp_swiglu elif model_type == "falcon_h1": apply_lora_mlp = apply_lora_mlp_swiglu - elif model_type == "qwen3moe": + elif model_type == "qwen3_moe": apply_lora_mlp = apply_lora_mlp_swiglu else: raise NotImplementedError(f"Unsloth: {model_type} is not yet implemented!") @@ -3412,6 +3496,16 @@ class FastLlamaModel: # MLP patching mlp_module = layer.mlp + # Qwen3 MoE uses Qwen3MoeSparseMoeBlock which holds + # stacked expert tensors on .experts; the dense + # gate/up/down fusion does not apply. MoE LoRA is + # wired through unsloth_zoo/moe_utils instead. + if not ( + hasattr(mlp_module, "gate_proj") + and hasattr(mlp_module, "up_proj") + and hasattr(mlp_module, "down_proj") + ): + continue gate_proj = mlp_module.gate_proj up_proj = mlp_module.up_proj down_proj = mlp_module.down_proj diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index fc91178d88..5174fbcc94 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -353,7 +353,13 @@ 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" @@ -632,7 +638,7 @@ class FastLanguageModel(FastLlamaModel): dispatch_model = FastGemma2Model elif model_type == "qwen2": dispatch_model = FastQwen2Model - elif model_type == "qwen3": # or model_type == "qwen3_moe": + elif model_type == "qwen3" or model_type == "qwen3_moe": if not SUPPORTS_QWEN3 or not SUPPORTS_QWEN3_MOE: raise ImportError( f"Unsloth: Your transformers version of {transformers_version} does not support Qwen3.\n" @@ -979,7 +985,12 @@ 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/qwen3_moe.py b/unsloth/models/qwen3_moe.py index e1f8c71b6b..52506449ae 100644 --- a/unsloth/models/qwen3_moe.py +++ b/unsloth/models/qwen3_moe.py @@ -137,9 +137,15 @@ def Qwen3MoeDecoderLayer_fast_forward( hidden_states = fast_rms_layernorm_inference( self.post_attention_layernorm, hidden_states ) - hidden_states, router_logits = Qwen3MoeSparseMoeBlock_fast_forward( - self.mlp, hidden_states - ) + # Use the class-level forward (patched by unsloth_zoo to + # sparse_moe_block_forward for transformers 5.x) instead of + # directly calling the legacy fast_forward, which breaks on + # stacked-expert MoE blocks that lack self.gate_proj. + mlp_out = self.mlp(hidden_states) + if isinstance(mlp_out, tuple): + hidden_states, router_logits = mlp_out[0], mlp_out[1] + else: + hidden_states, router_logits = mlp_out, None hidden_states += residual else: residual = hidden_states @@ -160,7 +166,17 @@ def Qwen3MoeDecoderLayer_fast_forward( # MoE Router MLP residual = hidden_states hidden_states = fast_rms_layernorm(self.post_attention_layernorm, hidden_states) - hidden_states, router_logits = self.mlp(hidden_states) + # unsloth_zoo's sparse_moe_block_forward returns a plain tensor + # for transformers 5.x stacked experts; the legacy patched + # forward returned a (hidden_states, router_logits) tuple. + # Handle both. + mlp_out = self.mlp(hidden_states) + if isinstance(mlp_out, tuple): + hidden_states = mlp_out[0] + router_logits = mlp_out[1] if len(mlp_out) > 1 else None + else: + hidden_states = mlp_out + router_logits = None hidden_states = residual + hidden_states outputs = (hidden_states,) @@ -188,7 +204,14 @@ class FastQwen3MoeModel(FastQwen3Model): Qwen3MoeAttention.forward = Qwen3Attention_fast_forward # Qwen3SdpaAttention .forward = Qwen3Attention_fast_forward # Qwen3FlashAttention2 .forward = Qwen3Attention_fast_forward - Qwen3MoeSparseMoeBlock.forward = Qwen3MoeSparseMoeBlock_fast_forward + # Qwen3MoeSparseMoeBlock.forward is patched by unsloth_zoo's + # patch_qwen3_moe (temporary_patches) to a transformers-5.x-aware + # sparse_moe_block_forward that correctly handles + # self.gate / self.experts. The legacy + # Qwen3MoeSparseMoeBlock_fast_forward below assumed a flat + # self.gate_proj attribute which no longer exists on stacked + # transformers 5.x experts. Skip the legacy override. + # Qwen3MoeSparseMoeBlock.forward = Qwen3MoeSparseMoeBlock_fast_forward Qwen3MoeMLP.forward = ( fast_swiglu_inference # This is analogous to Dense models' MLP ) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index ac9b35a822..f3be2dbdeb 100755 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -1752,9 +1752,22 @@ def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, import 'GuidedDecodingParams(backend="outlines", regex=args.vllm_guided_decoding_regex) ' 'if getattr(args, "vllm_guided_decoding_regex", None) is not None else None,', ) - # Replace with our vLLM engine + # Replace with our vLLM engine. + # + # ``_build_flex_from_args`` sizes the FlexEngine's fixed-shape + # page tables / CUDA-graph buckets from the GRPO rollout batch + # (per_device_train_batch_size * steps_per_generation * + # num_generations) BEFORE the first ``model.vllm_engine`` + # access triggers the lazy build. No-op when the model wasn't + # loaded through the flex-inference path (plain vLLM / plain + # HF), so this injection is safe for every backend. The + # import is inlined on the same line so ``create_new_function`` + # doesn't need a cross-module import entry. sampling_params = ( " " * 12 + + "from unsloth.inference.flex_engine import " + + "_build_flex_from_args as __unsloth_build_flex_from_args; " + + "__unsloth_build_flex_from_args(model, args); " + "self.llm = model.vllm_engine; self._last_loaded_step = 0; " + sampling_params ) # Add spaces @@ -1790,9 +1803,18 @@ def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, import ) if trl_version >= Version("0.18.0"): - # Replace LLM init with already existing vLLM engine for colocate mode + # Replace LLM init with already existing vLLM engine for colocate mode. + # Prepend the FlexEngine build call so the engine is sized from the + # GRPO rollout batch before ``model.vllm_engine`` is dereferenced. + # Inline the import so ``create_new_function`` doesn't need a + # cross-module entry; the helper is a no-op on non-flex models. vllm_llm_init_pattern = r"self\.llm\s*=\s*LLM\(.*?\)*\)\s*?\n(?!,)" - vllm_llm_replacement = "self.llm = model.vllm_engine\n" + vllm_llm_replacement = ( + "from unsloth.inference.flex_engine import " + "_build_flex_from_args as __unsloth_build_flex_from_args; " + "__unsloth_build_flex_from_args(model, args); " + "self.llm = model.vllm_engine\n" + ) new_vllm_part = re.sub( vllm_llm_init_pattern, vllm_llm_replacement, diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index df371e00c8..a5fc68d3f1 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -245,6 +245,9 @@ VLLM_SUPPORTED_VLM = [ "mistral3", "qwen3_vl", "qwen3_vl_moe", + "gemma4", + "qwen3_5", + "qwen3_5_moe", ] VLLM_NON_LORA_VLM = [ "mllama", @@ -606,7 +609,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", "qwen3_5", "qwen3_5_moe"} + 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 +924,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 +962,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: @@ -1028,6 +1056,22 @@ class FastBaseModel: is_vision_model = is_vlm, fp8_mode = fp8_mode, ) + # why: vLLM's -O3 compile backend (compilation_config=3) hits + # _decompose_size_nodes "Tried to erase Node size_N but it still + # had N users" on Gemma 4 multimodal models (both dense E2B and + # MoE 26B-A4B). Until the upstream FX splitter is fixed, fall + # back to compilation_config=0 (no piecewise compile) for gemma4 + # so fast_inference still loads. Runtime is slower than -O3 but + # correctness is preserved. + # + # The same FX splitter bug reproduces on Qwen3.5 / Qwen3.6 + # (qwen3_5 dense, qwen3_5_moe). The hybrid GatedDeltaNet + + # full-attention layers defeat piecewise compile's node erasure + # heuristics. Force eager for these archs too. + _eager_arches = {"gemma4", "qwen3_5", "qwen3_5_moe"} + if any(arch in _eager_arches for arch in (model_types or [])): + load_vllm_kwargs.setdefault("compilation_config", 0) + load_vllm_kwargs.setdefault("enforce_eager", True) for allowed_arg in allowed_args: if allowed_arg not in load_vllm_kwargs and allowed_arg in kwargs: load_vllm_kwargs[allowed_arg] = kwargs[allowed_arg] @@ -1223,6 +1267,18 @@ class FastBaseModel: raise _patch_err model = post_patch_loss_function(model) + # UNSLOTH_FAST_INFERENCE=1 path: install the lazy ``vllm_engine`` + # sentinel now that the tokenizer/processor is available. The real + # ``FlexEngine(...)`` construction is deferred to + # :func:`build_flex_engine` so the batch-size dimension can be + # sized from the GRPO rollout shape. Keeps the training model's + # forward intact — the engine carries its own deep-copy. + if hasattr(model, "_unsloth_needs_flex_engine"): + from unsloth.inference.flex_engine import install_flex_sentinel + + _tok_for_flex = getattr(tokenizer, "tokenizer", tokenizer) + install_flex_sentinel(model, _tok_for_flex) + # Log Unsloth version for future fastpaths for inference if hasattr(model, "config"): model.config.update({"unsloth_version": __version__}) @@ -1512,6 +1568,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)