tests/flex_moe_micro_bench: coord_descent compile_opts profile

Adds ``--compile_opts coord_descent`` that sets
``torch._inductor.config.coordinate_descent_tuning = True`` alone.
Tested; it regresses 4bit at bs=48 from 3383 → 2147 tok/s so
``coord_descent`` alone is not worth shipping. Keeping the option
in the bisection bench so the regression stays reproducible.

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()

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()

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()

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()

tests/flex_moe_bench.py

@@ -0,0 +1,185 @@
+# 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")
+    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")
+
+    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()

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()

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()

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()

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()

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

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:

unsloth/inference/__init__.py

@@ -0,0 +1,43 @@
+# 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 .vllm_shim import (
+    CompletionOutput,
+    LoRARequest,
+    RequestOutput,
+    load_lora,
+    save_lora,
+)
+
+__all__ = [
+    "FlexEngine",
+    "FlexMoEInference",
+    "load_flex",
+    "build_flex_engine",
+    "install_flex_sentinel",
+    "LoRARequest",
+    "RequestOutput",
+    "CompletionOutput",
+    "save_lora",
+    "load_lora",
+]

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

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

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",
+]

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",
+]

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)

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

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"

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
         )

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,

unsloth/models/vision.py

@@ -606,7 +606,16 @@ class FastBaseModel:
         vllm_enable_lora = True
 
         if is_vlm and fast_inference:
-            if not any(arch in VLLM_SUPPORTED_VLM for arch in model_types):
+            # The UNSLOTH_FAST_INFERENCE=1 flex backend ships with text-only
+            # support for Gemma-4-E2B-it today. Allow gemma4 through the
+            # vision-model path when the flex backend is selected; the
+            # vLLM-only compat list below still gates the default path.
+            _use_flex = os.environ.get("UNSLOTH_FAST_INFERENCE", "0") == "1"
+            _flex_allowed = {"gemma4"}
+            if not (
+                any(arch in VLLM_SUPPORTED_VLM for arch in model_types)
+                or (_use_flex and any(arch in _flex_allowed for arch in model_types))
+            ):
                 raise RuntimeError(
                     f"Unsloth: Fast inference is only supported for Language models and Qwen2.5-VL, Gemma3 among vision models. "
                     f"Found architectures: {', '.join(model_types)}!"
@@ -912,9 +921,15 @@ class FastBaseModel:
         verify_fp8_support_if_applicable(model_config)
 
         raise_handler = RaiseUninitialized()
-        if not fast_inference:
+        _use_flex_fast_inference = (
+            fast_inference and os.environ.get("UNSLOTH_FAST_INFERENCE", "0") == "1"
+        )
+        if (not fast_inference) or _use_flex_fast_inference:
+            # Shared by the standard HF load and the flex-inference load.
             # Prevent load_in_fp8 from being forwarded into HF internal model loading
             load_in_fp8 = kwargs.pop("load_in_fp8", None)
+            # ``max_batch_size`` is a FlexEngine kwarg, not an HF one.
+            _flex_max_batch_size = kwargs.pop("max_batch_size", 32)
             # Transformers 5.x @strict config classes reject unexpected kwargs.
             # Move config-level attributes onto the config object directly.
             _num_labels = kwargs.pop("num_labels", None)
@@ -944,6 +959,16 @@ class FastBaseModel:
                 fast_inference = fast_inference,
             )
             if hasattr(model, "generate"):
+                if _use_flex_fast_inference:
+                    # Defer engine construction until tokenizer/processor is
+                    # loaded; see the wiring block further down.
+                    model._unsloth_needs_flex_engine = dict(
+                        dtype = dtype,
+                        max_seq_length = max_seq_length,
+                        max_lora_rank = max_lora_rank,
+                        max_batch_size = _flex_max_batch_size,
+                        gpu_memory_utilization = gpu_memory_utilization,
+                    )
                 model.fast_generate = make_fast_generate_wrapper(model.generate)
                 model.fast_generate_batches = error_out_no_vllm
             if offload_embedding:
@@ -1223,6 +1248,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 +1549,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)