qwen3_5 flex: fix bs=max_batch_size serialization regression

PageTable reserves batch_idx=0 as a no-op slot, so passing
``max_batch_size=N`` to ``PageTable.__init__`` yields only N-1
user-allocatable slots. When a user requests bs=N concurrent sequences,
the Nth sequence gets stranded in the ``waiting`` queue until the first
N-1 sequences finish, then runs serially. For bs=8 × 64 tokens this
split the decode into 63 iterations at B=7 followed by 63 at B=1 —
throughput collapsed from a predicted ~800 tok/s to 143 tok/s.

Fix: bump the internal PageTable capacity by 1 so the user's
max_batch_size maps to that many concurrent slots. All batch-indexed
buffers (input_pos_buffer, block_mask_logical, _linear_conv_states,
_linear_recurrent_states, _linear_caches) size accordingly. Slot 0
remains internally reserved; user-facing ``self.max_batch_size``
unchanged.

Throughput (Qwen3.5-4B, 128 new tokens, post-warmup):

| bs | pre-fix tok/s | post-fix tok/s | change |
|---:|---:|---:|---:|
|  1 | 191.2 | 191.5 | same |
|  2 | 338.1 | 339.4 | same |
|  4 | 548.6 | 542.0 | same |
|  8 | 143.5 | **804.7** | **5.6x** |

MoE Qwen3.6-35B-A3B bs=8 × 64: 38.2 → **305.2 tok/s** (8x).

Parity at bs=4 × 16 tokens: was 3/4 exact (P1 matched 13/16), now 4/4
exact — the serialization also exposed an extra bf16-drift path.

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_gemma4_moe_merge_parity.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: GNU Affero General Public License v3.0
+# Copyright 2023-present the Unsloth team. All rights reserved.
+
+"""Numerical parity check for ``refresh_moe_lora_merge_from_pristine``
+at Gemma 4 26B-A4B expert shapes.
+
+Reuses ``merge_under_test`` and ``reference_merge`` from
+``tests/flex_moe_merge_parity.py`` (both are standalone helpers — they
+import nothing from the flex engine, only exercise the bitwise kernel
+logic lifted from ``unsloth/inference/flex_moe.py:682-800``).
+
+Shapes covered (Gemma 4 26B-A4B):
+- gate_up_proj: (E=128, 2I=1408, H=2816)  standard ``(E, out, in)``
+- down_proj:    (E=128, H=2816, I=704)    standard ``(E, out, in)``
+Rank-16 LoRA, single + dual adapter, bf16 + fp32.
+
+Gemma 4 MoE uses the same ``F.linear``-oriented expert layout as Qwen3,
+so only the ``transposed=False`` branch of
+``refresh_moe_lora_merge_from_pristine`` is exercised here. (The
+transposed branch is already covered for gpt-oss in
+``flex_moe_merge_parity.py``.)
+
+Usage::
+    CUDA_VISIBLE_DEVICES=2 python -u tests/flex_gemma4_moe_merge_parity.py
+"""
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import torch
+
+_REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+from tests.flex_moe_merge_parity import (  # noqa: E402
+    merge_under_test,
+    reference_merge,
+    test_correctness,
+)
+
+
+def main():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    torch.manual_seed(3407)
+    print(f"[merge-parity-gemma4] device={device} dtype=bf16 + fp32")
+    print(f"[merge-parity-gemma4] torch={torch.__version__}")
+    print()
+
+    print("== Correctness (fp32 golden + bf16 realistic Gemma 4 shapes) ==")
+    # gate_up_proj: (E=128, 2I=1408, H=2816) standard
+    # down_proj:    (E=128, H=2816,  I=704 ) standard
+    cases = [
+        # (E, in_dim, out_dim, R, dtype, transposed, n_adapters)
+        ( 8,  64, 128, 4,  torch.float32, False, 1),
+        ( 8,  64, 128, 4,  torch.float32, False, 2),
+        (16, 128, 256, 8,  torch.float32, False, 2),
+        # bf16 at Gemma 4 26B-A4B MoE shapes.
+        (128, 2816, 1408, 16, torch.bfloat16, False, 1),  # gate_up_proj
+        (128,  704, 2816, 16, torch.bfloat16, False, 1),  # down_proj
+        (128, 2816, 1408, 16, torch.bfloat16, False, 2),  # two adapters
+        (128, 2816, 1408, 64, torch.bfloat16, False, 1),  # higher rank
+    ]
+    all_ok = True
+    for E, in_dim, out_dim, R, dtype, tr, na in cases:
+        ok = test_correctness(E, in_dim, out_dim, R, dtype, device,
+                              transposed=tr, n_adapters=na)
+        all_ok = all_ok and ok
+    print(f"\n  overall: {'PASS' if all_ok else 'FAIL'}")
+
+
+if __name__ == "__main__":
+    main()

tests/flex_gemma4_parity.py

@@ -0,0 +1,191 @@
+# SPDX-License-Identifier: GNU Affero General Public License v3.0
+# Copyright 2023-present the Unsloth team. All rights reserved.
+
+"""Token parity for FlexGemma4Inference / FlexGemma4MoEInference vs pure HF.
+
+Usage::
+    CUDA_VISIBLE_DEVICES=2 UNSLOTH_FAST_INFERENCE=1 python -u \\
+        tests/flex_gemma4_parity.py --backend flex --model unsloth/gemma-4-31B-it
+
+    CUDA_VISIBLE_DEVICES=3 python -u tests/flex_gemma4_parity.py \\
+        --backend hf --model unsloth/gemma-4-31B-it
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+
+_REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+
+CHAT_PROMPTS = [
+    "In one sentence, what is Paris?",
+    "What is 23 + 19? Answer in one word.",
+    "Continue: The quick brown fox jumps over",
+]
+
+
+def _run_flex(args, dtype, *, capture: bool, lora_path=None):
+    import torch
+    os.environ["UNSLOTH_FAST_INFERENCE"] = "1"
+    import unsloth  # noqa
+    from unsloth import FastLanguageModel
+
+    if not capture:
+        # Force the eager decode path across arches.
+        try:
+            from unsloth.inference.flex_gemma4 import FlexGemma4Inference
+            FlexGemma4Inference.capture_decode_cudagraph = lambda self: None
+        except Exception:
+            pass
+        try:
+            from unsloth.inference.flex_gemma4_moe import FlexGemma4MoEInference
+            FlexGemma4MoEInference.capture_decode_cudagraph = lambda self: None
+        except Exception:
+            pass
+
+    model, tok = FastLanguageModel.from_pretrained(
+        model_name=args.model,
+        max_seq_length=args.max_seq_length,
+        dtype=dtype,
+        load_in_4bit=False,
+        fast_inference=True,
+        max_batch_size=4,
+        gpu_memory_utilization=0.6,
+    )
+
+    if lora_path is not None:
+        model.load_adapter(lora_path, adapter_name="default")
+        print(f"[parity-flex] LoRA attached from {lora_path}")
+
+    prompts = [
+        tok.apply_chat_template(
+            [{"role": "user", "content": p}],
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+        for p in CHAT_PROMPTS[: args.num_prompts]
+    ]
+
+    class _SP:
+        max_tokens = args.max_new_tokens
+        temperature = 0.0
+
+    _ = model.fast_generate(prompts, sampling_params=_SP(), use_tqdm=False)
+    outs = model.fast_generate(prompts, sampling_params=_SP(), use_tqdm=False)
+    token_ids = [list(o.outputs[0].token_ids) for o in outs]
+    texts = [o.outputs[0].text for o in outs]
+    return token_ids, texts, tok
+
+
+def _run_hf(args, dtype, *, lora_path=None):
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    tok = AutoTokenizer.from_pretrained(args.model)
+
+    try:
+        model = AutoModelForCausalLM.from_pretrained(
+            args.model, dtype=dtype, device_map="cuda",
+            attn_implementation="eager",
+        )
+    except Exception:
+        # Multimodal Gemma 4 (ConditionalGeneration) — load the top-level class.
+        from transformers import AutoModelForImageTextToText
+        model = AutoModelForImageTextToText.from_pretrained(
+            args.model, dtype=dtype, device_map="cuda",
+            attn_implementation="eager",
+        )
+    model.eval()
+
+    if lora_path is not None:
+        from peft import PeftModel
+        model = PeftModel.from_pretrained(model, lora_path)
+        model.eval()
+        print(f"[parity-hf] LoRA attached from {lora_path}")
+
+    if tok.pad_token_id is None:
+        tok.pad_token_id = tok.eos_token_id
+    tok.padding_side = "left"
+
+    prompts = [
+        tok.apply_chat_template(
+            [{"role": "user", "content": p}],
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+        for p in CHAT_PROMPTS[: args.num_prompts]
+    ]
+    inputs = tok(prompts, return_tensors="pt", padding=True).to("cuda")
+    out = model.generate(
+        **inputs,
+        max_new_tokens=args.max_new_tokens,
+        do_sample=False,
+        temperature=1.0,
+        pad_token_id=tok.pad_token_id,
+    )
+    prompt_len = inputs["input_ids"].shape[1]
+    eos = tok.eos_token_id
+    pad = tok.pad_token_id
+    token_ids = []
+    texts = []
+    for row in out:
+        ids = row[prompt_len:].tolist()
+        while ids and ids[-1] in (eos, pad):
+            ids.pop()
+        token_ids.append(ids)
+        texts.append(tok.decode(ids, skip_special_tokens=False))
+    return token_ids, texts, tok
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--model", default="unsloth/gemma-4-31B-it")
+    p.add_argument("--backend", choices=["flex", "flex_eager", "hf"], required=True)
+    p.add_argument("--max_new_tokens", type=int, default=64)
+    p.add_argument("--num_prompts", type=int, default=3)
+    p.add_argument("--lora_path", default=None)
+    p.add_argument("--max_seq_length", type=int, default=1024)
+    p.add_argument("--out_dir", default="async_task_outputs/gemma4_moe_bench")
+    args = p.parse_args()
+
+    import torch
+    dtype = torch.bfloat16
+
+    if args.backend == "flex":
+        token_ids, texts, _ = _run_flex(args, dtype, capture=True, lora_path=args.lora_path)
+    elif args.backend == "flex_eager":
+        token_ids, texts, _ = _run_flex(args, dtype, capture=False, lora_path=args.lora_path)
+    else:
+        token_ids, texts, _ = _run_hf(args, dtype, lora_path=args.lora_path)
+
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    model_tag = args.model.replace("/", "_")
+    suffix = "_lora" if args.lora_path else ""
+    out_path = out_dir / f"parity_{model_tag}_{args.backend}{suffix}.json"
+    with open(out_path, "w") as f:
+        json.dump(
+            {
+                "backend": args.backend,
+                "model": args.model,
+                "prompts": CHAT_PROMPTS[: args.num_prompts],
+                "token_ids": token_ids,
+                "texts": texts,
+            },
+            f,
+            indent=2,
+        )
+    print(f"[parity-{args.backend}] wrote {out_path}")
+    for i, (pp, tt) in enumerate(zip(CHAT_PROMPTS[: args.num_prompts], texts)):
+        print(f"[parity-{args.backend}] prompt {i}: {pp!r}")
+        print(f"[parity-{args.backend}] completion {i}: {tt!r}")
+
+
+if __name__ == "__main__":
+    main()

tests/flex_gpt_oss_parity.py

@@ -0,0 +1,176 @@
+# SPDX-License-Identifier: GNU Affero General Public License v3.0
+# Copyright 2023-present the Unsloth team. All rights reserved.
+
+"""Token parity for FlexGptOssInference vs pure HF generate on gpt-oss-20b.
+
+Usage::
+    # Flex side (cudagraph off in Phase 1).
+    CUDA_VISIBLE_DEVICES=2 UNSLOTH_FAST_INFERENCE=1 python -u \\
+        tests/flex_gpt_oss_parity.py --backend flex
+
+    # HF reference.
+    CUDA_VISIBLE_DEVICES=3 python -u tests/flex_gpt_oss_parity.py \\
+        --backend hf
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+from pathlib import Path
+
+_REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+
+CHAT_PROMPTS = [
+    "In one sentence, what is Paris?",
+    "What is 23 + 19? Answer in one word.",
+    "Continue: The quick brown fox jumps over",
+]
+
+
+def _run_flex(args, dtype, *, capture: bool, lora_path=None):
+    import torch
+    os.environ["UNSLOTH_FAST_INFERENCE"] = "1"
+    import unsloth  # noqa
+    from unsloth import FastLanguageModel
+
+    if not capture:
+        # Force the eager decode path by stubbing capture_decode_cudagraph.
+        from unsloth.inference.flex_gpt_oss import FlexGptOssInference
+        FlexGptOssInference.capture_decode_cudagraph = lambda self: None
+
+    model, tok = FastLanguageModel.from_pretrained(
+        model_name=args.model,
+        max_seq_length=args.max_seq_length,
+        dtype=dtype,
+        load_in_4bit=False,
+        fast_inference=True,
+        max_batch_size=4,
+        gpu_memory_utilization=0.6,
+    )
+
+    if lora_path is not None:
+        model.load_adapter(lora_path, adapter_name="default")
+        print(f"[parity-flex] LoRA attached from {lora_path}")
+
+    prompts = [
+        tok.apply_chat_template(
+            [{"role": "user", "content": p}],
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+        for p in CHAT_PROMPTS
+    ]
+
+    class _SP:
+        max_tokens = args.max_new_tokens
+        temperature = 0.0
+
+    # First call primes; second is the measurement.
+    _ = model.fast_generate(prompts, sampling_params=_SP(), use_tqdm=False)
+    outs = model.fast_generate(prompts, sampling_params=_SP(), use_tqdm=False)
+    token_ids = [list(o.outputs[0].token_ids) for o in outs]
+    texts = [o.outputs[0].text for o in outs]
+    return token_ids, texts, tok
+
+
+def _run_hf(args, dtype, *, lora_path=None):
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    tok = AutoTokenizer.from_pretrained(args.model)
+    model = AutoModelForCausalLM.from_pretrained(
+        args.model, dtype=dtype, device_map="cuda",
+        attn_implementation="eager",
+    )
+    model.eval()
+
+    if lora_path is not None:
+        from peft import PeftModel
+        model = PeftModel.from_pretrained(model, lora_path)
+        model.eval()
+        print(f"[parity-hf] LoRA attached from {lora_path}")
+
+    if tok.pad_token_id is None:
+        tok.pad_token_id = tok.eos_token_id
+    tok.padding_side = "left"
+
+    prompts = [
+        tok.apply_chat_template(
+            [{"role": "user", "content": p}],
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+        for p in CHAT_PROMPTS
+    ]
+    inputs = tok(prompts, return_tensors="pt", padding=True).to("cuda")
+    out = model.generate(
+        **inputs,
+        max_new_tokens=args.max_new_tokens,
+        do_sample=False,
+        temperature=1.0,
+        pad_token_id=tok.pad_token_id,
+    )
+    prompt_len = inputs["input_ids"].shape[1]
+    eos = tok.eos_token_id
+    pad = tok.pad_token_id
+    token_ids = []
+    texts = []
+    for row in out:
+        ids = row[prompt_len:].tolist()
+        while ids and ids[-1] in (eos, pad):
+            ids.pop()
+        token_ids.append(ids)
+        texts.append(tok.decode(ids, skip_special_tokens=False))
+    return token_ids, texts, tok
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--model", default="unsloth/gpt-oss-20b-BF16")
+    p.add_argument("--backend", choices=["flex", "flex_eager", "hf"], required=True)
+    p.add_argument("--max_new_tokens", type=int, default=64)
+    p.add_argument("--lora_path", default=None,
+                   help="optional LoRA adapter dir to attach before decode")
+    p.add_argument("--max_seq_length", type=int, default=1024)
+    p.add_argument("--out_dir", default="async_task_outputs/qwen3_moe_grpo_bench_v2")
+    args = p.parse_args()
+
+    import torch
+    dtype = torch.bfloat16
+
+    if args.backend == "flex":
+        token_ids, texts, _ = _run_flex(args, dtype, capture=True, lora_path=args.lora_path)
+    elif args.backend == "flex_eager":
+        token_ids, texts, _ = _run_flex(args, dtype, capture=False, lora_path=args.lora_path)
+    else:
+        token_ids, texts, _ = _run_hf(args, dtype, lora_path=args.lora_path)
+
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    suffix = "_lora" if args.lora_path else ""
+    out_path = out_dir / f"parity_gptoss_{args.backend}{suffix}.json"
+    with open(out_path, "w") as f:
+        json.dump(
+            {
+                "backend": args.backend,
+                "model": args.model,
+                "prompts": CHAT_PROMPTS,
+                "token_ids": token_ids,
+                "texts": texts,
+            },
+            f,
+            indent=2,
+        )
+    print(f"[parity-{args.backend}] wrote {out_path}")
+    for i, (pp, tt) in enumerate(zip(CHAT_PROMPTS, texts)):
+        print(f"[parity-{args.backend}] prompt {i}: {pp!r}")
+        print(f"[parity-{args.backend}] completion {i}: {tt!r}")
+
+
+if __name__ == "__main__":
+    main()

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,199 @@
+# SPDX-License-Identifier: GNU Affero General Public License v3.0
+# Copyright 2023-present the Unsloth team. All rights reserved.
+
+"""Decode throughput bench: FlexMoEInference vs HF generate on Qwen3 MoE.
+
+Apples-to-apples decode on the same prompt set + new-token budget,
+same LoRA rank, same precision. Mirrors PR #5123's
+``tests/flex_fastlm_bench.py`` shape.
+
+Usage:
+    CUDA_VISIBLE_DEVICES=0 UNSLOTH_FAST_INFERENCE=1 python -u \
+        tests/flex_moe_bench.py --backend flex --load_in_4bit
+
+    CUDA_VISIBLE_DEVICES=0 python -u \
+        tests/flex_moe_bench.py --backend hf   --load_in_4bit
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+import time
+from pathlib import Path
+
+_REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--backend", choices = ["flex", "hf", "hf_naive"], default = "flex")
+    p.add_argument("--model", default = "unsloth/Qwen3-30B-A3B-Instruct-2507")
+    p.add_argument("--dtype", choices = ["bf16", "fp16"], default = "bf16")
+    p.add_argument("--load_in_4bit", action = "store_true")
+    p.add_argument("--n_prompts", type = int, default = 8)
+    p.add_argument("--max_new_tokens", type = int, default = 64)
+    p.add_argument("--max_seq_length", type = int, default = 1024)
+    p.add_argument("--warmup_rounds", type = int, default = 1)
+    p.add_argument("--timed_rounds", type = int, default = 2)
+    p.add_argument("--out_dir", default = "async_task_outputs/qwen3_moe_grpo_bench")
+    p.add_argument("--chat_template", action = "store_true",
+                   help = "wrap each prompt with tokenizer.apply_chat_template")
+    p.add_argument("--user_prompt",
+                   default = "Continue this sentence: The quick brown fox jumps over fence {i}, then")
+    args = p.parse_args()
+
+    import torch
+
+    if args.backend == "flex":
+        os.environ["UNSLOTH_FAST_INFERENCE"] = "1"
+    os.environ.setdefault("UNSLOTH_MOE_BACKEND", "grouped_mm")
+
+    dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
+    torch.cuda.reset_peak_memory_stats()
+    t_load0 = time.perf_counter()
+
+    if args.backend == "hf_naive":
+        # Pure transformers path: NO ``import unsloth`` so none of
+        # Unsloth's Qwen3 MoE attention / MLP patches run. This is the
+        # fair naive reference to compare flex fast-inference against.
+        from transformers import (
+            AutoModelForCausalLM,
+            AutoTokenizer,
+            BitsAndBytesConfig,
+        )
+        quant_cfg = None
+        if args.load_in_4bit:
+            quant_cfg = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_compute_dtype=dtype,
+                bnb_4bit_use_double_quant=True,
+                bnb_4bit_quant_type="nf4",
+            )
+        tokenizer = AutoTokenizer.from_pretrained(args.model)
+        if tokenizer.pad_token_id is None or tokenizer.pad_token == "<|PAD_TOKEN|>":
+            tokenizer.pad_token = "<|vision_pad|>"
+        tokenizer.padding_side = "left"
+        attn_impl = os.environ.get("HF_ATTN_IMPL", "eager")
+        print(f"[bench] HF attn_implementation={attn_impl}")
+        model = AutoModelForCausalLM.from_pretrained(
+            args.model,
+            dtype=dtype,
+            quantization_config=quant_cfg,
+            device_map="cuda",
+            attn_implementation=attn_impl,
+        )
+        model.eval()
+        print(f"[bench] pure transformers (no unsloth patches)")
+    else:
+        import unsloth
+        print(f"[bench] unsloth={unsloth.__file__}")
+        from unsloth import FastLanguageModel
+        model, tokenizer = FastLanguageModel.from_pretrained(
+            model_name = args.model,
+            max_seq_length = args.max_seq_length,
+            dtype = dtype,
+            load_in_4bit = args.load_in_4bit,
+            fast_inference = args.backend == "flex",
+        )
+
+    t_load = time.perf_counter() - t_load0
+    peak_load = torch.cuda.max_memory_reserved() / 1024**3
+    print(f"[bench] loaded in {t_load:.1f}s peak {peak_load:.1f} GB")
+
+    if args.chat_template:
+        prompts = []
+        for i in range(args.n_prompts):
+            msg = [{"role": "user", "content": args.user_prompt.format(i=i)}]
+            prompts.append(tokenizer.apply_chat_template(
+                msg, tokenize = False, add_generation_prompt = True,
+            ))
+        print(f"[bench] chat_template enabled. prompt[0] (first 200 chars):\n"
+              f"  {prompts[0][:200]!r}")
+    else:
+        prompts = [f"The quick brown fox jumps over fence {i}, then"
+                   for i in range(args.n_prompts)]
+
+    if args.backend == "flex":
+        class _SP:
+            max_tokens = args.max_new_tokens
+            temperature = 0.0
+
+        # Warmup
+        for _ in range(args.warmup_rounds):
+            _ = model.fast_generate(prompts, sampling_params = _SP(), use_tqdm = False)
+
+        # Timed
+        wall_times = []
+        tok_counts = []
+        for _ in range(args.timed_rounds):
+            t0 = time.perf_counter()
+            outs = model.fast_generate(prompts, sampling_params = _SP(), use_tqdm = False)
+            wall_times.append(time.perf_counter() - t0)
+            tok_counts.append(
+                sum(len(o.outputs[0].token_ids) for o in outs)
+            )
+    else:
+        # HF generate (shared for "hf" unsloth-patched and "hf_naive" pure).
+        inputs = tokenizer(prompts, return_tensors = "pt", padding = True).to("cuda")
+        gen_kwargs = dict(
+            max_new_tokens = args.max_new_tokens,
+            do_sample = False,
+            temperature = 1.0,
+            pad_token_id = tokenizer.pad_token_id or tokenizer.eos_token_id,
+        )
+        # Warmup
+        for _ in range(args.warmup_rounds):
+            _ = model.generate(**inputs, **gen_kwargs)
+            torch.cuda.synchronize()
+
+        wall_times = []
+        tok_counts = []
+        for _ in range(args.timed_rounds):
+            torch.cuda.synchronize()
+            t0 = time.perf_counter()
+            out = model.generate(**inputs, **gen_kwargs)
+            torch.cuda.synchronize()
+            wall_times.append(time.perf_counter() - t0)
+            n_new = (out.shape[1] - inputs["input_ids"].shape[1]) * out.shape[0]
+            tok_counts.append(n_new)
+
+    peak_gen = torch.cuda.max_memory_reserved() / 1024**3
+    median_wall = sorted(wall_times)[len(wall_times) // 2]
+    median_tok = tok_counts[len(wall_times) // 2]
+    tok_per_s = median_tok / median_wall if median_wall > 0 else 0.0
+
+    print(f"[bench] wall: {wall_times}")
+    print(f"[bench] tok counts: {tok_counts}")
+    print(f"[bench] median wall: {median_wall:.2f}s  median tok/s: {tok_per_s:.1f}")
+    print(f"[bench] peak VRAM after gen: {peak_gen:.1f} GB")
+
+    precision = "4bit" if args.load_in_4bit else args.dtype
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents = True, exist_ok = True)
+    summary = {
+        "phase": "bench_decode",
+        "backend": args.backend,
+        "model": args.model,
+        "precision": precision,
+        "n_prompts": args.n_prompts,
+        "max_new_tokens": args.max_new_tokens,
+        "wall_times_s": wall_times,
+        "tok_counts": tok_counts,
+        "median_wall_s": round(median_wall, 3),
+        "median_tok_s": round(tok_per_s, 1),
+        "peak_vram_load_gb": round(peak_load, 2),
+        "peak_vram_after_gen_gb": round(peak_gen, 2),
+        "t_load_s": round(t_load, 1),
+    }
+    with open(out_dir / f"bench_decode_{args.backend}_{precision}.json", "w") as f:
+        json.dump(summary, f, indent = 2)
+    print(f"[bench] wrote {out_dir / f'bench_decode_{args.backend}_{precision}.json'}")
+
+
+if __name__ == "__main__":
+    main()

tests/flex_moe_merge_parity.py

@@ -0,0 +1,290 @@
+# SPDX-License-Identifier: GNU Affero General Public License v3.0
+# Copyright 2023-present the Unsloth team. All rights reserved.
+
+"""Numerical + performance check for ``refresh_moe_lora_merge_from_pristine``.
+
+Does what we never actually did on this branch: verify the stacked MoE
+LoRA merge kernel produces the same W_inf as the textbook reference
+``W_ref[e] = W_pristine[e] + sum_a scaling_a * B_a[e] @ A_a[e]`` for both
+standard (E, 2I, H) and transposed (E, H, 2I) expert layouts, with 1 and
+2 active adapters, then benchmarks the batched ``baddbmm`` path against
+the dense-layer-style ``addmm`` loop at Qwen3-30B-A3B MoE shapes.
+
+Usage::
+    CUDA_VISIBLE_DEVICES=2 python -u tests/flex_moe_merge_parity.py
+"""
+from __future__ import annotations
+
+import sys
+import time
+from pathlib import Path
+
+import torch
+
+_REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+
+# ---------------------------------------------------------------------------
+# Reference: the exact formula from the docstring, no tricks.
+# ---------------------------------------------------------------------------
+
+def reference_merge(W_pristine, adapters, *, transposed):
+    """adapters: list of (A_stacked [E*R, in], B_stacked [out, E*R], scaling)."""
+    E = W_pristine.shape[0]
+    out = W_pristine.clone()
+    for A_stack, B_stack, scaling in adapters:
+        R = A_stack.shape[0] // E
+        in_dim = A_stack.shape[1]
+        out_dim = B_stack.shape[0]
+        for e in range(E):
+            A_e = A_stack[e * R : (e + 1) * R]           # (R, in_dim)
+            B_e = B_stack[:, e * R : (e + 1) * R]         # (out_dim, R)
+            delta = (B_e @ A_e).to(W_pristine.dtype)      # (out_dim, in_dim)
+            if transposed:
+                out[e] += scaling * delta.t()
+            else:
+                out[e] += scaling * delta
+    return out
+
+
+# ---------------------------------------------------------------------------
+# Under-test: a standalone copy of the merge body from flex_moe.py so we
+# can drive it without loading a 30B model. The logic is a character-for-
+# character lift of lines 747-794 in unsloth/inference/flex_moe.py.
+# ---------------------------------------------------------------------------
+
+def merge_under_test(W_inf, W_pristine, adapters):
+    """adapters: list of (A_stacked, B_stacked, scaling)."""
+    E = W_inf.shape[0]
+    # Orientation detection (mirrors flex_moe.py:752-763).
+    A_w0, B_w0, _ = adapters[0]
+    in_dim = A_w0.shape[1]
+    out_dim = B_w0.shape[0]
+    d0, d1 = W_inf.shape[1], W_inf.shape[2]
+    if d0 == out_dim and d1 == in_dim:
+        is_standard = True
+    elif d0 == in_dim and d1 == out_dim:
+        is_standard = False
+    else:
+        raise RuntimeError("orientation")
+
+    W_inf.copy_(W_pristine)
+
+    for A_w, B_w, scaling in adapters:
+        R = A_w.shape[0] // E
+        A_3d = A_w.view(E, R, in_dim)
+        B_3d = B_w.view(out_dim, E, R).permute(1, 0, 2).contiguous()
+        if is_standard:
+            torch.baddbmm(
+                W_inf,
+                B_3d.to(W_inf.dtype),
+                A_3d.to(W_inf.dtype),
+                alpha=float(scaling),
+                beta=1.0,
+                out=W_inf,
+            )
+        else:
+            torch.baddbmm(
+                W_inf,
+                A_3d.transpose(-2, -1).contiguous().to(W_inf.dtype),
+                B_3d.transpose(-2, -1).contiguous().to(W_inf.dtype),
+                alpha=float(scaling),
+                beta=1.0,
+                out=W_inf,
+            )
+    return W_inf
+
+
+# ---------------------------------------------------------------------------
+# Correctness
+# ---------------------------------------------------------------------------
+
+def _make_adapter(E, R, in_dim, out_dim, dtype, device, seed):
+    g = torch.Generator(device=device).manual_seed(seed)
+    A = torch.randn(E * R, in_dim, generator=g, device=device, dtype=dtype) * 0.02
+    B = torch.randn(out_dim, E * R, generator=g, device=device, dtype=dtype) * 0.02
+    return A, B
+
+
+def test_correctness(E, in_dim, out_dim, R, dtype, device, *, transposed, n_adapters):
+    scalings = [2.0, 0.5][:n_adapters]
+    adapters = [
+        (*_make_adapter(E, R, in_dim, out_dim, dtype, device, seed=11 + i), s)
+        for i, s in enumerate(scalings)
+    ]
+
+    if transposed:
+        W_pristine = torch.randn(E, in_dim, out_dim, device=device, dtype=dtype) * 0.02
+    else:
+        W_pristine = torch.randn(E, out_dim, in_dim, device=device, dtype=dtype) * 0.02
+    W_inf = torch.empty_like(W_pristine)
+
+    merge_under_test(W_inf, W_pristine, adapters)
+    W_ref = reference_merge(W_pristine, adapters, transposed=transposed)
+
+    tol = dict(atol=3e-3, rtol=3e-3) if dtype == torch.bfloat16 else dict(atol=1e-5, rtol=1e-5)
+    close = torch.allclose(W_inf, W_ref, **tol)
+    max_err = (W_inf - W_ref).abs().max().item()
+    scale = W_ref.abs().max().item()
+    rel = max_err / max(scale, 1e-8)
+    label = f"E={E} in={in_dim} out={out_dim} R={R} {dtype} transposed={transposed} n_adapters={n_adapters}"
+    verdict = "OK" if close else "FAIL"
+    print(f"  [{verdict}] {label}  max_abs={max_err:.3e}  rel={rel:.3e}")
+    return close
+
+
+# ---------------------------------------------------------------------------
+# Performance: batched baddbmm vs per-expert addmm loop (same total flops).
+# ---------------------------------------------------------------------------
+
+def _sync():
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+
+
+def bench(fn, *args, iters=50, warmup=10):
+    for _ in range(warmup):
+        fn(*args)
+    _sync()
+    t0 = time.perf_counter()
+    for _ in range(iters):
+        fn(*args)
+    _sync()
+    return (time.perf_counter() - t0) / iters
+
+
+def perf_vs_dense_addmm(out_dim, in_dim, R, dtype, device):
+    """Dense-layer LoRA refresh: a single torch.addmm call like the one
+    in flex_qwen3_llama.py::refresh_lora_merge_from_pristine. Used as a
+    per-matrix baseline to compare against the MoE baddbmm kernel cost."""
+    W_pristine = torch.randn(out_dim, in_dim, device=device, dtype=dtype) * 0.02
+    W = torch.empty_like(W_pristine)
+    A = torch.randn(R, in_dim, device=device, dtype=dtype) * 0.02
+    B = torch.randn(out_dim, R, device=device, dtype=dtype) * 0.02
+    scaling = 2.0
+
+    def run_addmm():
+        torch.addmm(W_pristine, B, A, alpha=scaling, out=W)
+
+    t = bench(run_addmm, iters=200, warmup=20)
+    print(f"  dense addmm  out={out_dim:>5} in={in_dim:>5} R={R:>2} "
+          f"{t * 1e3:7.4f}ms")
+    return t
+
+
+def perf_compare(E, in_dim, out_dim, R, dtype, device):
+    A_w, B_w = _make_adapter(E, R, in_dim, out_dim, dtype, device, seed=0)
+    W_pristine = torch.randn(E, out_dim, in_dim, device=device, dtype=dtype) * 0.02
+    W_inf_a = torch.empty_like(W_pristine)
+    W_inf_b = torch.empty_like(W_pristine)
+
+    A_3d = A_w.view(E, R, in_dim)
+    B_3d = B_w.view(out_dim, E, R).permute(1, 0, 2).contiguous()
+    scaling = 2.0
+
+    def run_baddbmm():
+        W_inf_a.copy_(W_pristine)
+        torch.baddbmm(W_inf_a, B_3d, A_3d, alpha=scaling, beta=1.0, out=W_inf_a)
+
+    def run_addmm_loop():
+        W_inf_b.copy_(W_pristine)
+        for e in range(E):
+            torch.addmm(
+                W_inf_b[e],
+                B_3d[e],
+                A_3d[e],
+                alpha=scaling,
+                out=W_inf_b[e],
+            )
+
+    t_bad = bench(run_baddbmm)
+    t_loop = bench(run_addmm_loop)
+
+    # Sanity: they produce the same result.
+    run_baddbmm()
+    run_addmm_loop()
+    max_err = (W_inf_a - W_inf_b).abs().max().item()
+
+    print(f"  E={E:>3} out={out_dim:>5} in={in_dim:>5} R={R:>2} "
+          f"baddbmm={t_bad * 1e3:7.3f}ms  addmm_loop={t_loop * 1e3:7.3f}ms  "
+          f"speedup={t_loop / t_bad:5.2f}x  max_abs_diff={max_err:.1e}")
+
+
+def main():
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    torch.manual_seed(3407)
+    print(f"[merge-parity] device={device} dtype=bf16 + fp32")
+    print(f"[merge-parity] torch={torch.__version__}")
+    print()
+
+    print("== Correctness (fp32 golden + bf16 realistic) ==")
+    # Qwen3-30B-A3B shapes.
+    #   gate_up_proj: (E=128, 2I=1536, H=2048) standard
+    #   down_proj:    (E=128, H=2048, I=768)  standard
+    # Transposed variant: (E, H, 2I) — exercised via the `transposed=True` test.
+    cases = [
+        # (E, in_dim, out_dim, R, dtype, transposed, n_adapters)
+        ( 8, 128, 192, 4, torch.float32, False, 1),
+        ( 8, 128, 192, 4, torch.float32, True,  1),
+        ( 8, 128, 192, 4, torch.float32, False, 2),
+        (16, 256, 512, 8, torch.float32, True,  2),
+        # bf16 at realistic Qwen3 MoE shapes. Mild tol to allow 3e-3 bf16 noise.
+        (128, 2048, 1536, 16, torch.bfloat16, False, 1),  # gate_up_proj-like
+        (128, 768,  2048, 16, torch.bfloat16, False, 1),  # down_proj-like
+        (128, 2048, 1536, 16, torch.bfloat16, True,  1),  # transposed gate_up
+        (128, 2048, 1536, 16, torch.bfloat16, False, 2),  # two adapters
+    ]
+    all_ok = True
+    for E, in_dim, out_dim, R, dtype, tr, na in cases:
+        ok = test_correctness(E, in_dim, out_dim, R, dtype, device,
+                              transposed=tr, n_adapters=na)
+        all_ok = all_ok and ok
+    print(f"\n  overall: {'PASS' if all_ok else 'FAIL'}")
+
+    print()
+    print("== Performance: batched baddbmm vs per-expert addmm loop (bf16) ==")
+    print("   (same arithmetic, different dispatch pattern — baddbmm = 1 kernel,")
+    print("    addmm loop = E kernels)")
+    moe_cases = [
+        (128, 2048, 1536, 16),   # gate_up_proj
+        (128, 768,  2048, 16),   # down_proj
+        (128, 2048, 1536, 64),   # larger rank
+    ]
+    for E, in_dim, out_dim, R in moe_cases:
+        perf_compare(E, in_dim, out_dim, R, torch.bfloat16, device)
+
+    print()
+    print("== Dense baseline: torch.addmm on a single expert-sized matrix ==")
+    print("   (this is what flex_qwen3_llama.py:354-370 does for dense layers)")
+    dense_times = {}
+    for E, in_dim, out_dim, R in moe_cases:
+        t = perf_vs_dense_addmm(out_dim, in_dim, R, torch.bfloat16, device)
+        dense_times[(E, in_dim, out_dim, R)] = t
+
+    print()
+    print("== Per-expert cost comparison ==")
+    print("   (MoE baddbmm cost / E) vs single dense addmm for the same per-expert matrix")
+    for E, in_dim, out_dim, R in moe_cases:
+        # Measure baddbmm cost again for this exact shape to get the numerator.
+        A_w, B_w = _make_adapter(E, R, in_dim, out_dim, torch.bfloat16, device, seed=0)
+        W_pristine = torch.randn(E, out_dim, in_dim, device=device, dtype=torch.bfloat16) * 0.02
+        W_inf = torch.empty_like(W_pristine)
+        A_3d = A_w.view(E, R, in_dim)
+        B_3d = B_w.view(out_dim, E, R).permute(1, 0, 2).contiguous()
+
+        def run_baddbmm():
+            W_inf.copy_(W_pristine)
+            torch.baddbmm(W_inf, B_3d, A_3d, alpha=2.0, beta=1.0, out=W_inf)
+        t_moe = bench(run_baddbmm, iters=50, warmup=10)
+        t_dense = dense_times[(E, in_dim, out_dim, R)]
+        per_expert = t_moe / E
+        print(f"  E={E:>3} out={out_dim:>5} in={in_dim:>5} R={R:>2}  "
+              f"moe_per_expert={per_expert * 1e6:7.2f}us  "
+              f"dense_addmm={t_dense * 1e6:7.2f}us  "
+              f"ratio(moe/dense)={per_expert / t_dense:5.2f}x")
+
+
+if __name__ == "__main__":
+    main()

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_moe_vllm_bench.py

@@ -0,0 +1,217 @@
+# SPDX-License-Identifier: GNU Affero General Public License v3.0
+# Copyright 2023-present the Unsloth team. All rights reserved.
+
+"""vLLM direct throughput comparison for Qwen3 MoE.
+
+Uses ``vllm.LLM`` directly (no unsloth) to bench the same workload
+as ``tests/flex_moe_bench.py`` so the flex+compile_walker numbers
+have an apples-to-apples vLLM baseline.
+
+Usage:
+    CUDA_VISIBLE_DEVICES=2 python -u \\
+        tests/flex_moe_vllm_bench.py --load_in_4bit
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+# DeepGEMM isn't installed in this env; vLLM's FP8 warmup crashes the
+# engine-core subprocess without it. Disable before importing vllm.
+os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
+
+_REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+
+def _gpu_mem_used_gb() -> float:
+    try:
+        out = subprocess.run(
+            ["nvidia-smi", "--query-gpu=memory.used",
+             "--format=csv,noheader,nounits", "-i", "0"],
+            capture_output=True, text=True, timeout=5,
+        )
+        return int(out.stdout.strip().splitlines()[0]) / 1024
+    except Exception:
+        return 0.0
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--model", default="unsloth/Qwen3-30B-A3B-Instruct-2507")
+    p.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16")
+    p.add_argument("--load_in_4bit", action="store_true",
+                   help="use the -bnb-4bit checkpoint variant")
+    p.add_argument("--n_prompts", type=int, default=8)
+    p.add_argument("--max_new_tokens", type=int, default=64)
+    p.add_argument("--max_model_len", type=int, default=1024)
+    p.add_argument("--warmup_rounds", type=int, default=1)
+    p.add_argument("--timed_rounds", type=int, default=2)
+    p.add_argument("--out_dir", default="async_task_outputs/qwen3_moe_grpo_bench_v2")
+    p.add_argument("--gpu_memory_utilization", type=float, default=0.85)
+    p.add_argument("--enable_lora", action="store_true",
+                   help="enable vLLM LoRA serving; requires --lora_path")
+    p.add_argument("--lora_path", default=None,
+                   help="path to a LoRA adapter dir to load per-request")
+    p.add_argument("--max_lora_rank", type=int, default=16)
+    p.add_argument("--enforce_eager", action="store_true",
+                   help="disable vLLM CUDA graph capture (diagnostic)")
+    p.add_argument("--tag", default=None,
+                   help="extra label for the output json filename")
+    p.add_argument("--chat_template", action="store_true",
+                   help="wrap each prompt with the tokenizer's chat "
+                        "template (needed for chat-tuned models like "
+                        "gpt-oss that produce gibberish on raw strings)")
+    p.add_argument("--user_prompt",
+                   default="Continue this sentence: The quick brown fox jumps over fence {i}, then",
+                   help="user-message template when --chat_template is set "
+                        "(``{i}`` gets replaced with the prompt index)")
+    args = p.parse_args()
+
+    import torch
+    from vllm import LLM, SamplingParams
+    import vllm
+    print(f"[vllm] version: {vllm.__version__}")
+
+    dtype_str = "bfloat16" if args.dtype == "bf16" else "float16"
+
+    # For 4bit, tell vLLM to apply bnb quantization on-the-fly to the
+    # bf16 checkpoint (there's no pre-quantized unsloth 4bit variant for
+    # this model on HF). For FP8, point at the FP8 variant directly.
+    model_id = args.model
+    quantization = None
+    if args.load_in_4bit:
+        quantization = "bitsandbytes"
+        print(f"[vllm] 4bit: model_id={model_id} quantization=bitsandbytes (on-the-fly)")
+    elif os.environ.get("USE_FP8") == "1":
+        model_id = args.model + "-FP8"
+        quantization = "fp8"
+        print(f"[vllm] fp8: model_id={model_id}")
+
+    llm_kwargs = dict(
+        model=model_id,
+        dtype=dtype_str,
+        quantization=quantization,
+        gpu_memory_utilization=args.gpu_memory_utilization,
+        max_model_len=args.max_model_len,
+        max_num_seqs=max(args.n_prompts, 8),
+        enforce_eager=args.enforce_eager,
+        trust_remote_code=False,
+    )
+    if args.enable_lora:
+        llm_kwargs.update(
+            enable_lora=True,
+            max_lora_rank=args.max_lora_rank,
+            max_loras=1,
+        )
+
+    torch.cuda.reset_peak_memory_stats()
+    t0 = time.perf_counter()
+    llm = LLM(**llm_kwargs)
+    t_load = time.perf_counter() - t0
+    # vLLM uses its own allocator; torch's reserved bytes miss it.
+    # Query nvidia-smi instead (visible GPU only — i.e. whatever
+    # CUDA_VISIBLE_DEVICES exposes as device 0).
+    peak_load = _gpu_mem_used_gb()
+    print(f"[vllm] loaded in {t_load:.1f}s peak {peak_load:.1f} GB")
+
+    # Use vLLM's native ``llm.chat`` when chat-template is requested —
+    # it threads harmony-aware templating + tokenization through vLLM's
+    # chat_utils, which is what gpt-oss needs (hand-rolled
+    # apply_chat_template + raw ``llm.generate`` was producing gibberish
+    # on gpt-oss-20b-BF16 no-LoRA).
+    if args.chat_template:
+        chat_messages = [
+            [{"role": "user", "content": args.user_prompt.format(i=i)}]
+            for i in range(args.n_prompts)
+        ]
+        print(f"[vllm] chat mode: llm.chat() with {args.n_prompts} single-turn "
+              f"messages. first: {chat_messages[0][0]['content']!r}")
+        prompts = None
+    else:
+        prompts = [f"The quick brown fox jumps over fence {i}, then"
+                   for i in range(args.n_prompts)]
+        chat_messages = None
+
+    sp = SamplingParams(max_tokens=args.max_new_tokens, temperature=0.0)
+
+    gen_kwargs = {}
+    if args.enable_lora and args.lora_path is not None:
+        from vllm.lora.request import LoRARequest
+        gen_kwargs["lora_request"] = LoRARequest(
+            lora_name="flex_lora",
+            lora_int_id=1,
+            lora_path=args.lora_path,
+        )
+        print(f"[vllm] LoRA enabled: path={args.lora_path} rank<={args.max_lora_rank}")
+
+    def _run_once():
+        if chat_messages is not None:
+            return llm.chat(chat_messages, sampling_params=sp,
+                            use_tqdm=False, **gen_kwargs)
+        return llm.generate(prompts, sampling_params=sp,
+                            use_tqdm=False, **gen_kwargs)
+
+    # Warmup.
+    for _ in range(args.warmup_rounds):
+        _ = _run_once()
+
+    # Timed.
+    wall = []
+    tok_counts = []
+    for _ in range(args.timed_rounds):
+        torch.cuda.synchronize()
+        t0 = time.perf_counter()
+        outs = _run_once()
+        torch.cuda.synchronize()
+        wall.append(time.perf_counter() - t0)
+        tok_counts.append(sum(len(o.outputs[0].token_ids) for o in outs))
+
+    peak_gen = _gpu_mem_used_gb()
+    med_wall = sorted(wall)[len(wall) // 2]
+    med_tok = tok_counts[len(wall) // 2]
+    tps = med_tok / med_wall if med_wall > 0 else 0.0
+
+    sample_text = outs[0].outputs[0].text if outs else ""
+
+    print(f"[vllm] wall: {wall}")
+    print(f"[vllm] tok counts: {tok_counts}")
+    print(f"[vllm] median wall: {med_wall:.3f}s  median tok/s: {tps:.1f}")
+    print(f"[vllm] peak VRAM after gen: {peak_gen:.1f} GB")
+    print(f"[vllm] sample completion[0]: {sample_text[:200]!r}")
+
+    precision = "4bit" if args.load_in_4bit else args.dtype
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    summary = {
+        "phase": "bench_decode",
+        "backend": "vllm",
+        "vllm_version": vllm.__version__,
+        "model": model_id,
+        "precision": precision,
+        "n_prompts": args.n_prompts,
+        "max_new_tokens": args.max_new_tokens,
+        "wall_times_s": wall,
+        "tok_counts": tok_counts,
+        "median_wall_s": round(med_wall, 3),
+        "median_tok_s": round(tps, 1),
+        "peak_vram_load_gb": round(peak_load, 2),
+        "peak_vram_after_gen_gb": round(peak_gen, 2),
+        "t_load_s": round(t_load, 1),
+        "sample_completion": sample_text[:500],
+    }
+    suffix = f"_{args.tag}" if args.tag else ""
+    out_path = out_dir / f"bench_decode_vllm_{precision}{suffix}.json"
+    with open(out_path, "w") as f:
+        json.dump(summary, f, indent=2)
+    print(f"[vllm] wrote {out_path}")
+
+
+if __name__ == "__main__":
+    main()

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

tests/gemma4_fast_inference_parity.py

@@ -0,0 +1,141 @@
+"""Greedy parity: FastLanguageModel(fast_inference=True) vs HF naive on Gemma 4.
+
+Covers all four Gemma 4 variants (E2B / E4B dense, 31B dense, 26B-A4B MoE) via
+``--model``. Runs HF FIRST, frees the GPU, then runs vLLM — loading vLLM first
+in the same process leaves global state (allocator, compile cache, patched
+functions) that subtly perturbs a later plain HF run, producing false
+divergences. HF-first ordering gives bitwise matches.
+
+Requires vLLM nightly (>= 2026-04-17 for `vllm#39291` Gemma 4 LoRA) plus the
+`unsloth-zoo#603` vLLM Gemma 4 patches.
+
+Example:
+    CUDA_VISIBLE_DEVICES=0 python -u tests/gemma4_fast_inference_parity.py \\
+        --model unsloth/gemma-4-26b-a4b-it --max_new_tokens 32
+"""
+import argparse
+import gc
+import json
+import os
+
+os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
+os.environ.setdefault("UNSLOTH_MOE_BACKEND", "grouped_mm")
+
+# Import unsloth at module scope so its patches apply to the HF run as well —
+# users always import unsloth before touching transformers in practice.
+import unsloth  # noqa: F401, E402
+import torch  # noqa: E402
+
+
+def _render(tok, raw_prompts):
+    return [
+        tok.apply_chat_template(
+            [{"role": "user", "content": p}],
+            tokenize=False, add_generation_prompt=True,
+        )
+        for p in raw_prompts
+    ]
+
+
+def run_hf_first(model_name, raw_prompts, max_new_tokens):
+    from transformers import AutoModelForImageTextToText, AutoProcessor
+    proc = AutoProcessor.from_pretrained(model_name)
+    tok = proc.tokenizer if hasattr(proc, "tokenizer") else proc
+    prompts = _render(tok, raw_prompts)
+
+    model = AutoModelForImageTextToText.from_pretrained(
+        model_name, dtype=torch.bfloat16,
+        attn_implementation="sdpa", device_map="cuda:0",
+    )
+    model.eval()
+    ids, texts = [], []
+    for p in prompts:
+        enc = tok(p, return_tensors="pt").to("cuda:0")
+        with torch.inference_mode():
+            gen = model.generate(
+                **enc, max_new_tokens=max_new_tokens,
+                do_sample=False, temperature=None, top_p=None,
+            )
+        new = gen[0, enc.input_ids.shape[1]:]
+        ids.append(new.tolist())
+        texts.append(tok.decode(new, skip_special_tokens=True))
+
+    del model, proc
+    gc.collect()
+    torch.cuda.empty_cache()
+    return ids, texts, prompts
+
+
+def run_vllm_second(model_name, prompts, max_new_tokens):
+    from unsloth import FastLanguageModel
+    from vllm import SamplingParams
+
+    model, _ = FastLanguageModel.from_pretrained(
+        model_name=model_name, max_seq_length=1024, dtype=torch.bfloat16,
+        load_in_4bit=False, fast_inference=True, max_batch_size=8,
+        gpu_memory_utilization=0.6, max_lora_rank=16,
+    )
+    sp = SamplingParams(max_tokens=max_new_tokens, temperature=0.0)
+    outs = model.fast_generate(prompts, sampling_params=sp, use_tqdm=False)
+    ids = [list(o.outputs[0].token_ids) for o in outs]
+    texts = [o.outputs[0].text for o in outs]
+    return ids, texts
+
+
+DEFAULT_PROMPTS = [
+    "In one sentence, what is Paris?",
+    "What is 23 + 19? Answer in one word.",
+    "Continue this phrase: The quick brown fox jumps over",
+]
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", default="unsloth/gemma-4-26b-a4b-it")
+    ap.add_argument("--max_new_tokens", type=int, default=32)
+    ap.add_argument("--json_out", default=None)
+    args = ap.parse_args()
+
+    print(f"=== HF first ({args.model}) ===")
+    hf_ids, hf_texts, prompts = run_hf_first(
+        args.model, DEFAULT_PROMPTS, args.max_new_tokens,
+    )
+    for i, t in enumerate(hf_texts):
+        print(f"[hf]   P{i}: {t[:120]!r}")
+
+    print("\n=== vLLM second ===")
+    fast_ids, fast_texts = run_vllm_second(
+        args.model, prompts, args.max_new_tokens,
+    )
+    for i, t in enumerate(fast_texts):
+        print(f"[fast] P{i}: {t[:120]!r}")
+
+    print("\n=== Match ===")
+    total, matched, rows = 0, 0, []
+    for i, (a, b) in enumerate(zip(fast_ids, hf_ids)):
+        n = min(len(a), len(b))
+        m = sum(1 for x, y in zip(a[:n], b[:n]) if x == y)
+        total += n
+        matched += m
+        print(f"[P{i}] {m}/{n}")
+        rows.append({
+            "prompt_id": i, "match": m, "total": n,
+            "fast": fast_texts[i], "hf": hf_texts[i],
+        })
+    print(f"\nTotal {matched}/{total}")
+
+    if args.json_out:
+        os.makedirs(os.path.dirname(args.json_out) or ".", exist_ok=True)
+        with open(args.json_out, "w") as f:
+            json.dump({
+                "model": args.model, "rows": rows,
+                "matched": matched, "total": total,
+            }, f, indent=2)
+
+    # Exit non-zero on any divergence so CI can gate on bitwise parity.
+    if matched != total:
+        raise SystemExit(f"Divergence: {matched}/{total}")
+
+
+if __name__ == "__main__":
+    main()

tests/gemma4_flex_bench.py

@@ -0,0 +1,132 @@
+"""Throughput bench for Unsloth flex fast-inference on Gemma 4.
+
+Runs ``FastLanguageModel(fast_inference=True)`` through the flex engine
+(``UNSLOTH_FAST_INFERENCE=1``) and measures tok/s at a list of batch sizes.
+Captures CUDA graphs at each bucket; the bench reports the post-warmup pass.
+
+Example:
+    CUDA_VISIBLE_DEVICES=0 UNSLOTH_FAST_INFERENCE=1 \\
+        UNSLOTH_MOE_BACKEND=grouped_mm \\
+        python -u tests/gemma4_flex_bench.py \\
+        --model unsloth/gemma-4-26b-a4b-it --batch_sizes 1 4 8 16 \\
+        --max_new_tokens 64 --json_out async_task_outputs/flex_bench_26b.json
+
+Pair with tests/gemma4_fast_inference_parity.py for HF-naive parity and
+with tests/gemma4_fast_bench.py (vLLM nightly path) for cross-engine
+throughput comparison.
+"""
+import argparse
+import json
+import os
+import time
+
+os.environ.setdefault("UNSLOTH_FAST_INFERENCE", "1")
+os.environ.setdefault("UNSLOTH_MOE_BACKEND", "grouped_mm")
+os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
+
+import unsloth  # noqa: F401, E402
+import torch  # noqa: E402
+
+
+SHORT_PROMPT = "In one sentence, what is Paris?"
+LONG_POOL = [
+    "In one sentence, what is Paris?",
+    "What is 23 + 19? Answer in one word.",
+    "Continue this phrase: The quick brown fox jumps over",
+    "Name three primary colors.",
+    "Who wrote the play Hamlet?",
+    "What is the capital of Japan?",
+    "What does the acronym 'NASA' stand for?",
+    "Explain the water cycle in one sentence.",
+    "Give a two-word summary of the French Revolution.",
+    "What is the tallest mountain on Earth?",
+    "Define 'entropy' in one sentence.",
+    "Who painted the Mona Lisa?",
+    "What is the boiling point of water in Celsius?",
+    "Name one prime number larger than 10.",
+    "What color do you get when you mix blue and yellow?",
+    "Give one example of an amphibian.",
+]
+
+
+def _render(tok, prompts):
+    return [
+        tok.apply_chat_template(
+            [{"role": "user", "content": p}],
+            tokenize=False, add_generation_prompt=True,
+        )
+        for p in prompts
+    ]
+
+
+class _SP:
+    def __init__(self, n):
+        self.max_tokens = n
+        self.temperature = 0.0
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", required=True)
+    ap.add_argument("--batch_sizes", nargs="+", type=int, default=[1, 4, 8, 16])
+    ap.add_argument("--max_new_tokens", type=int, default=64)
+    ap.add_argument("--json_out", required=True)
+    ap.add_argument("--max_seq_length", type=int, default=1024)
+    args = ap.parse_args()
+
+    from unsloth import FastLanguageModel
+
+    t_load = time.perf_counter()
+    model, tok_raw = FastLanguageModel.from_pretrained(
+        model_name=args.model, max_seq_length=args.max_seq_length,
+        dtype=torch.bfloat16, load_in_4bit=False, fast_inference=True,
+        max_batch_size=max(args.batch_sizes),
+        gpu_memory_utilization=0.6,
+    )
+    t_load = time.perf_counter() - t_load
+    tok = tok_raw.tokenizer if hasattr(tok_raw, "tokenizer") else tok_raw
+    print(f"[flex] load {t_load:.1f}s")
+
+    sp = _SP(args.max_new_tokens)
+    results = []
+
+    for bs in args.batch_sizes:
+        prompts = _render(tok, LONG_POOL[:bs])
+        # Warmup: primes the cudagraph bucket and walker compile.
+        _ = model.fast_generate(prompts, sampling_params=sp, use_tqdm=False)
+        t0 = time.perf_counter()
+        outs = model.fast_generate(prompts, sampling_params=sp, use_tqdm=False)
+        dt = time.perf_counter() - t0
+        total = sum(len(o.outputs[0].token_ids) for o in outs)
+        tok_s = total / dt if dt > 0 else 0.0
+        results.append({
+            "batch_size": bs, "gen_s": round(dt, 3),
+            "total_tokens": total, "tok_s": round(tok_s, 1),
+        })
+        print(f"[flex bs={bs}] {total} tok in {dt:.2f}s -> {tok_s:.1f} tok/s")
+
+    # Parity sanity: run SHORT_PROMPT at bs=1, save ids.
+    p = _render(tok, [SHORT_PROMPT])
+    outs = model.fast_generate(p, sampling_params=sp, use_tqdm=False)
+    parity_ids = list(outs[0].outputs[0].token_ids)
+    parity_text = outs[0].outputs[0].text
+
+    peak_gb = torch.cuda.max_memory_allocated() / 1024**3
+
+    summary = {
+        "model": args.model,
+        "max_new_tokens": args.max_new_tokens,
+        "load_s": round(t_load, 2),
+        "peak_vram_gb": round(peak_gb, 2),
+        "batches": results,
+        "parity_ids": parity_ids,
+        "parity_text": parity_text,
+    }
+    os.makedirs(os.path.dirname(args.json_out) or ".", exist_ok=True)
+    with open(args.json_out, "w") as f:
+        json.dump(summary, f, indent=2)
+    print(f"[flex] wrote {args.json_out}")
+
+
+if __name__ == "__main__":
+    main()

tests/qwen3_5_flex_parity_bs4.py

@@ -0,0 +1,124 @@
+"""B=4 parity test so the captured graph bucket actually replays.
+
+Compares flex captured-graph decode vs flex eager-decode (capture
+disabled via ``UNSLOTH_FLEX_QWEN3_5_NO_CAPTURE=1``). Exact token-for-token
+match is the success criterion — both paths use identical math, only
+the kernel-launch wrapper differs.
+
+    CUDA_VISIBLE_DEVICES=2 UNSLOTH_FAST_INFERENCE=1 PYTHONUNBUFFERED=1 \\
+        python -u tests/qwen3_5_flex_parity_bs4.py \\
+        --model Qwen/Qwen3.5-4B --max_new_tokens 16 \\
+        --json_out async_task_outputs/qwen3_5/flex_parity_bs4.json
+"""
+import argparse
+import json
+import os
+import time
+
+
+CHAT_PROMPTS = [
+    "In one sentence, what is Paris?",
+    "What is 23 + 19? Answer in one word.",
+    "Continue: The quick brown fox jumps over",
+    "Name one primary color.",
+]
+
+
+def _run(model, tok, args, capture: bool):
+    class SP:
+        def __init__(self, n): self.max_tokens = n; self.temperature = 0.0
+
+    sp = SP(args.max_new_tokens)
+    rendered = [
+        tok.apply_chat_template(
+            [{"role": "user", "content": p}],
+            tokenize=False, add_generation_prompt=True,
+        )
+        for p in CHAT_PROMPTS
+    ]
+    os.environ["UNSLOTH_FLEX_QWEN3_5_NO_CAPTURE"] = "0" if capture else "1"
+    t0 = time.perf_counter()
+    outs = model.fast_generate(rendered, sampling_params=sp, use_tqdm=False)
+    dt = time.perf_counter() - t0
+    results = []
+    for p, o in zip(CHAT_PROMPTS, outs):
+        results.append({
+            "prompt": p,
+            "text": o.outputs[0].text,
+            "ids": list(o.outputs[0].token_ids),
+        })
+    return dt, results
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", required=True)
+    ap.add_argument("--max_new_tokens", type=int, default=16)
+    ap.add_argument("--max_seq_length", type=int, default=1024)
+    ap.add_argument("--max_batch_size", type=int, default=4)
+    ap.add_argument("--gpu_memory_utilization", type=float, default=0.6)
+    ap.add_argument("--json_out", required=True)
+    args = ap.parse_args()
+
+    import torch
+    os.environ["UNSLOTH_FAST_INFERENCE"] = "1"
+    import unsloth  # noqa
+    from unsloth import FastLanguageModel
+
+    model, tok_raw = FastLanguageModel.from_pretrained(
+        model_name=args.model,
+        max_seq_length=args.max_seq_length,
+        dtype=torch.bfloat16,
+        load_in_4bit=False,
+        fast_inference=True,
+        max_batch_size=args.max_batch_size,
+        gpu_memory_utilization=args.gpu_memory_utilization,
+    )
+    tok = tok_raw.tokenizer if hasattr(tok_raw, "tokenizer") else tok_raw
+
+    # Eager first (so the graph capture doesn't pollute state).
+    dt_eager, eager = _run(model, tok, args, capture=False)
+    print(f"[eager] gen_s={dt_eager:.2f}")
+    for r in eager:
+        print(f"[eager] {r['prompt'][:40]!r} -> {r['text'][:80]!r}")
+
+    dt_cap, cap = _run(model, tok, args, capture=True)
+    print(f"[cap]   gen_s={dt_cap:.2f}")
+    for r in cap:
+        print(f"[cap]   {r['prompt'][:40]!r} -> {r['text'][:80]!r}")
+
+    # Token-by-token parity.
+    mismatches = []
+    for i, (e, c) in enumerate(zip(eager, cap)):
+        matches = 0
+        for a, b in zip(e["ids"], c["ids"]):
+            if a == b:
+                matches += 1
+            else:
+                break
+        mismatches.append({
+            "prompt_idx": i,
+            "prompt": e["prompt"],
+            "matches": matches,
+            "total": min(len(e["ids"]), len(c["ids"])),
+            "eager_text": e["text"],
+            "cap_text": c["text"],
+        })
+        tag = "OK " if matches == min(len(e["ids"]), len(c["ids"])) else "MM "
+        print(f"[{tag}] P{i} {matches}/{min(len(e['ids']), len(c['ids']))}")
+
+    summary = {
+        "model": args.model,
+        "max_new_tokens": args.max_new_tokens,
+        "eager_gen_s": round(dt_eager, 3),
+        "capture_gen_s": round(dt_cap, 3),
+        "mismatches": mismatches,
+    }
+    os.makedirs(os.path.dirname(args.json_out) or ".", exist_ok=True)
+    with open(args.json_out, "w") as f:
+        json.dump(summary, f, indent=2)
+    print(f"[parity-bs4] wrote {args.json_out}")
+
+
+if __name__ == "__main__":
+    main()

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,45 @@
+# SPDX-License-Identifier: GNU Affero General Public License v3.0
+# Copyright 2023-present the Unsloth team. All rights reserved.
+
+"""Flex-attention inference engines.
+
+``UNSLOTH_FAST_INFERENCE=1`` routes ``FastLanguageModel.from_pretrained``
+through :func:`load_flex`, which wraps the selected HF model with a
+:class:`FlexEngine` that presents the vLLM ``LLM`` surface used by
+Unsloth / TRL GRPO (``.generate``, ``.chat``, ``.sleep``, ``.wake_up``,
+``.llm_engine``, plus ``save_lora`` / ``load_lora`` via the module
+shim).
+
+Four architectures are supported today: Qwen3 (dense), Qwen3-MoE,
+Llama-3, Gemma-4-E2B-it. Anything else raises
+:class:`NotImplementedError`; unset the env var or use vLLM instead."""
+
+from .flex_engine import (
+    FlexEngine,
+    build_flex_engine,
+    install_flex_sentinel,
+    load_flex,
+)
+from .flex_moe import FlexMoEInference
+from .flex_gemma4_moe import FlexGemma4MoEInference
+from .vllm_shim import (
+    CompletionOutput,
+    LoRARequest,
+    RequestOutput,
+    load_lora,
+    save_lora,
+)
+
+__all__ = [
+    "FlexEngine",
+    "FlexMoEInference",
+    "FlexGemma4MoEInference",
+    "load_flex",
+    "build_flex_engine",
+    "install_flex_sentinel",
+    "LoRARequest",
+    "RequestOutput",
+    "CompletionOutput",
+    "save_lora",
+    "load_lora",
+]

unsloth/inference/flex_gemma4_moe.py

@@ -0,0 +1,937 @@
+# SPDX-License-Identifier: GNU Affero General Public License v3.0
+# Copyright 2023-present the Unsloth team. All rights reserved.
+
+"""Gemma 4 MoE inference with flex_attention + paged KV cache.
+
+Sibling of ``flex_moe.py`` (Qwen3-MoE) and ``flex_gpt_oss.py`` (gpt-oss).
+Scope: ``unsloth/gemma-4-26B-A4B-it`` (30 layers, 128 experts top-k 8,
+hidden 2816, moe_intermediate 704, dense intermediate 2112, ~3.8B active).
+
+Arch-specific pieces on top of the shared engine:
+
+1. **Dual dense MLP + MoE per layer.** Unlike Qwen3/gpt-oss which replace
+   the dense MLP with an expert block, Gemma 4 MoE runs both in parallel
+   per layer and sums their normed outputs before the residual add:
+
+       residual = h
+       h = pre_ffw_norm(h)
+       mlp_out = layer.mlp(h)                        # dense SwiGLU
+       h1 = post_ffw_norm_1(mlp_out)
+       h2 = pre_ffw_norm_2(residual.reshape(-1, H))  # experts input
+       h2 = layer.experts(h2, top_k_idx, top_k_w)
+       h2 = post_ffw_norm_2(h2).reshape(residual)
+       h = post_ffw_norm(h1 + h2)
+       h = residual + h
+       h *= layer.layer_scalar
+
+2. **Two-tier RoPE.**
+   - Sliding layers (25/30): ``rope_type="default"``, theta=10K,
+     full ``head_dim=256``.
+   - Full-attn layers (5/30): ``rope_type="proportional"``, theta=1M,
+     ``global_head_dim=512``, ``partial_rotary_factor=0.25`` (rotate
+     only leading 25% of head dim; pass through remainder).
+   Walker computes both (cos, sin) tuples per generate entry; attention
+   forward picks the right one based on ``self_attn.layer_type``.
+
+3. **Per-head Q/K/V RMSNorm pre-rotary.** ``q_norm`` and ``k_norm`` apply
+   RMSNorm on the last (head_dim) axis before RoPE. ``v_norm`` applies
+   with ``with_scale=False`` (RMSNorm that just divides, no gain).
+
+4. **K=V alternative on full-attn layers.** ``attention_k_eq_v=True`` +
+   full-attn layer ⇒ ``v_proj is None`` and ``value_states`` is the raw
+   ``k_proj(hidden)`` output (before ``k_norm`` and before RoPE),
+   followed only by ``v_norm``. Sliding layers use the normal q/k/v path.
+
+5. **Expert grouped_mm.** Reuses ``Gemma4TextExperts.forward`` through
+   ``unsloth_zoo.temporary_patches.gemma4_moe.patch_gemma4_moe`` — the
+   ``per_expert_scale`` is pre-folded into routing weights, so the
+   generic ``forward_native_grouped_mm`` (``moe_utils.py``) handles the
+   standard ``(E, 2I, H)``/``(E, H, I)`` layout with ``act_fn =
+   gelu_pytorch_tanh`` via the default ``elif hasattr(self, 'act_fn')``
+   fallback — zero changes needed in moe_utils.
+
+6. **Embedding scale + final ``layer_scalar``.** Embed output gets
+   multiplied by ``sqrt(hidden_size)``. Each decoder layer's output is
+   multiplied by ``self.layer_scalar`` (a ``torch.ones(1)`` buffer —
+   numerically a no-op today, but must not be dropped so the walker
+   matches HF bitwise).
+
+Out of scope (this file errors out if the config requests them):
+- ``num_kv_shared_layers > 0`` (E2B/E4B KV-share variants)
+- ``hidden_size_per_layer_input > 0`` (E2B/E4B per-layer input gate)
+- Mixed sliding-window sizes across layers
+- bnb-4bit stacked experts (no ``Gemma4TextExpertsBnb4bit`` ships yet)
+"""
+
+from __future__ import annotations
+
+import math
+import os
+import types
+from collections import deque
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.attention.flex_attention import BlockMask, create_block_mask
+
+try:
+    from .flex_qwen3_llama import (
+        DECODE_KERNEL_OPTIONS_DEFAULT,
+        PREFILL_KERNEL_OPTIONS_DEFAULT,
+        Sequence,
+        flex_attention_compiled,
+        refresh_lora_merge_from_pristine,
+    )
+    from .flex_moe import refresh_moe_lora_merge_from_pristine
+    from .flex_paged_attention import PagedKVCache, PageTable
+except ImportError:  # script-mode fallback
+    import sys
+    from pathlib import Path
+    sys.path.insert(0, str(Path(__file__).resolve().parent))
+    from flex_qwen3_llama import (  # noqa: E402
+        DECODE_KERNEL_OPTIONS_DEFAULT,
+        PREFILL_KERNEL_OPTIONS_DEFAULT,
+        Sequence,
+        flex_attention_compiled,
+        refresh_lora_merge_from_pristine,
+    )
+    from flex_moe import refresh_moe_lora_merge_from_pristine  # noqa: E402
+    from flex_paged_attention import PagedKVCache, PageTable  # noqa: E402
+
+
+# ---------------------------------------------------------------------------
+# Rotary helpers
+# ---------------------------------------------------------------------------
+
+
+def _rotate_half(x):
+    x1, x2 = torch.chunk(x, 2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def _apply_rotary_full(x, cos, sin):
+    """Standard Llama-style RoPE on the full head dim.
+
+    ``cos``/``sin`` shape ``(B, S, D)`` → unsqueeze at head-dim to
+    ``(B, 1, S, D)`` so it broadcasts over ``(B, H, S, D)`` q/k.
+
+    For Gemma 4 full-attention layers, ``rope_type="proportional"`` emits
+    an ``inv_freq`` with zeros in the tail ``(1 - partial_rotary_factor)``
+    fraction of positions. Those zero entries make the corresponding
+    cos=1 / sin=0, so ``rotate_half`` passes the tail dims through
+    unchanged — no separate partial-RoPE helper is needed.
+    """
+    cos = cos.unsqueeze(1)
+    sin = sin.unsqueeze(1)
+    return (x * cos) + (_rotate_half(x) * sin)
+
+
+# ---------------------------------------------------------------------------
+# Attention forward: paged KV + flex_attention + per-layer sliding + k=v.
+# ---------------------------------------------------------------------------
+
+
+def make_gemma4_moe_attention_forward(page_table: PageTable):
+    """Return a ``forward`` method for ``Gemma4TextAttention``.
+
+    Differences from ``flex_gpt_oss``:
+    - No sinks; call ``flex_attention_compiled`` without ``return_lse``.
+    - ``q_norm`` / ``k_norm`` / ``v_norm`` before rotary / KV write.
+    - ``v_proj is None`` (k=v full-attn) ⇒ reuse raw ``k_proj`` as v.
+    - Partial RoPE on full-attn layers via ``_partial_rotary_dim``.
+    """
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] = None,
+        position_embeddings_sliding: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+        position_embeddings_full: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+        attention_mask=None,
+        past_key_values=None,
+        cache_position=None,
+        flex_block_mask: Optional[BlockMask] = None,
+        flex_block_mask_sliding: Optional[BlockMask] = None,
+        flex_input_pos: Optional[torch.Tensor] = None,
+        flex_batch_idx: Optional[torch.Tensor] = None,
+        flex_kernel_options: Optional[dict] = None,
+        **kwargs,
+    ):
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        # pick cos/sin for this layer
+        if self.is_sliding:
+            cos, sin = position_embeddings_sliding
+        else:
+            cos, sin = position_embeddings_full
+
+        # q projection
+        q = self.q_proj(hidden_states).view(hidden_shape)
+        q = self.q_norm(q)
+        # shape now (B, S, Hq, D)
+
+        # k/v projections
+        k_raw = self.k_proj(hidden_states).view(hidden_shape)
+
+        if self.v_proj is not None:
+            v_raw = self.v_proj(hidden_states).view(hidden_shape)
+        else:
+            # k=v alternative: value uses the RAW k projection
+            # (before k_norm, before rotary). v_norm then applies with
+            # with_scale=False.
+            v_raw = k_raw
+
+        k = self.k_norm(k_raw)
+        v = self.v_norm(v_raw)
+
+        # Transpose for flex_attention: (B, S, H, D) -> (B, H, S, D).
+        # Rotary applied AFTER the transpose so cos/sin broadcast on dim 1.
+        q = q.transpose(1, 2)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+
+        q = _apply_rotary_full(q, cos, sin)
+        k = _apply_rotary_full(k, cos, sin)
+
+        # Paged-KV write.
+        if self._paged_cache is not None and flex_input_pos is not None:
+            cache_dtype = self._paged_cache.k_cache.dtype
+            if k.dtype != cache_dtype:
+                k = k.to(cache_dtype)
+            if v.dtype != cache_dtype:
+                v = v.to(cache_dtype)
+            k, v = self._paged_cache.update(flex_input_pos, k, v, flex_batch_idx)
+
+        # Per-layer block mask dispatch.
+        if self.is_sliding and flex_block_mask_sliding is not None:
+            block_mask = flex_block_mask_sliding
+        else:
+            block_mask = flex_block_mask
+
+        attn_output = flex_attention_compiled(
+            q,
+            k,
+            v,
+            scale=self.scaling,
+            block_mask=block_mask,
+            enable_gqa=True,
+            kernel_options=flex_kernel_options,
+        )
+
+        attn_output = attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous()
+        return self.o_proj(attn_output), None
+
+    return forward
+
+
+def patch_gemma4_moe_attention_forwards(model: torch.nn.Module, page_table: PageTable):
+    """Attach a ``PagedKVCache`` + replace ``forward`` on every
+    ``Gemma4TextAttention`` layer AND replace ``Gemma4TextExperts.forward``
+    with the grouped_mm MoE backend.
+
+    Gemma 4 has per-layer-type head_dim and kv_heads, so each layer's
+    cache is sized from its own attrs.
+
+    The stock ``Gemma4TextExperts.forward`` (modeling_gemma4.py:1263) is a
+    Python loop over active experts — slow for decode and not
+    CUDA-graph-capturable. We rebind to the generic
+    ``forward_native_grouped_mm`` (moe_utils.py:771) which handles the
+    standard ``(E, 2I, H)`` / ``(E, H, I)`` layout with the ``act_fn``
+    fallback activation path (``gelu_pytorch_tanh`` via ``ACT2FN``).
+
+    The router's ``per_expert_scale`` is already folded into
+    ``top_k_weights`` by the modeling-native ``Gemma4TextRouter.forward``
+    (modeling_gemma4.py:1309), so no separate folding is needed here.
+    (The unsloth-zoo ``gemma4_moe.patch_gemma4_moe`` targets a legacy
+    ``Gemma4TextMoEBlock`` class that no longer ships in transformers
+    5.5+; it no-ops on current transformers.)
+    """
+    fwd = make_gemma4_moe_attention_forward(page_table)
+    text_cfg = getattr(model.config, "text_config", model.config)
+
+    if getattr(text_cfg, "num_kv_shared_layers", 0):
+        raise NotImplementedError(
+            "Gemma 4 MoE with KV-shared layers not supported yet "
+            f"(num_kv_shared_layers={text_cfg.num_kv_shared_layers})."
+        )
+    if getattr(text_cfg, "hidden_size_per_layer_input", 0):
+        raise NotImplementedError(
+            "Gemma 4 MoE with per-layer input gate not supported yet "
+            f"(hidden_size_per_layer_input={text_cfg.hidden_size_per_layer_input})."
+        )
+
+    # Pick grouped_mm MoE backend once; bind per-layer below.
+    try:
+        from unsloth_zoo.temporary_patches.moe_utils import get_forward_moe_backend
+        _moe_forward = get_forward_moe_backend()
+    except Exception:  # pragma: no cover - defensive
+        _moe_forward = None
+
+    for layer in model.model.layers:
+        attn = layer.self_attn
+        num_q_heads = attn.q_proj.out_features // attn.head_dim
+        num_kv_heads = max(1, num_q_heads // attn.num_key_value_groups)
+        attn._paged_cache = PagedKVCache(
+            page_table,
+            n_heads=num_kv_heads,
+            head_dim=attn.head_dim,
+            dtype=model.dtype,
+        ).to(model.device)
+        attn.forward = types.MethodType(fwd, attn)
+
+        if _moe_forward is not None and getattr(layer, "enable_moe_block", False):
+            experts = layer.experts
+            # The grouped_mm backend inspects ``self.act_fn`` — already
+            # set by ``Gemma4TextExperts.__init__`` to
+            # ``ACT2FN[config.hidden_activation]``. No other setup needed.
+            experts.forward = types.MethodType(_moe_forward, experts)
+
+
+# ---------------------------------------------------------------------------
+# Walker: dual dense MLP + MoE per decoder layer + layer_scalar.
+# ---------------------------------------------------------------------------
+
+
+def _compute_rotary_per_layer_type(base, inputs_embeds, position_ids):
+    """Call ``Gemma4TextRotaryEmbedding`` once per known ``layer_type``.
+
+    Returns ``(cos_sliding, sin_sliding), (cos_full, sin_full)``. cos/sin
+    for the full-attention type are ALREADY partial-sized in the HF impl
+    if ``rope_type="proportional"`` + ``partial_rotary_factor<1.0`` — the
+    init fn slices the head dim to ``partial_rotary_factor * head_dim``
+    before emitting ``inv_freq``. So we just forward through.
+    """
+    rot = base.rotary_emb
+    layer_types = set(getattr(rot, "layer_types", {"sliding_attention", "full_attention"}))
+    pos_sliding = pos_full = None
+    if "sliding_attention" in layer_types:
+        pos_sliding = rot(inputs_embeds, position_ids, layer_type="sliding_attention")
+    if "full_attention" in layer_types:
+        pos_full = rot(inputs_embeds, position_ids, layer_type="full_attention")
+    # Fallback: if only one type exists, use it for both slots so the
+    # attention forward never reads None.
+    if pos_sliding is None:
+        pos_sliding = pos_full
+    if pos_full is None:
+        pos_full = pos_sliding
+    return pos_sliding, pos_full
+
+
+def call_gemma4_moe_model_with_flex_kwargs(model, input_ids, position_ids, flex_kwargs):
+    """Walk a ``Gemma4TextModel`` manually, injecting flex kwargs into each
+    attention call. Per-layer dual MLP + MoE with ``layer_scalar``.
+
+    ``model.model.embed_tokens`` is a ``Gemma4TextScaledWordEmbedding`` that
+    already applies ``* sqrt(hidden_size)`` in its forward, so we just
+    call it (no manual scale).
+    """
+    base = model.model
+    inputs_embeds = base.embed_tokens(input_ids)
+    position_embeddings_sliding, position_embeddings_full = _compute_rotary_per_layer_type(
+        base, inputs_embeds, position_ids
+    )
+
+    hidden_states = inputs_embeds
+    compute_dtype = inputs_embeds.dtype
+
+    for layer in base.layers:
+        residual = hidden_states
+        hidden_states = layer.input_layernorm(hidden_states).to(compute_dtype)
+
+        attn_out, _ = layer.self_attn(
+            hidden_states,
+            position_embeddings_sliding=position_embeddings_sliding,
+            position_embeddings_full=position_embeddings_full,
+            **flex_kwargs,
+        )
+        attn_out = layer.post_attention_layernorm(attn_out).to(compute_dtype)
+        hidden_states = residual + attn_out
+
+        residual = hidden_states
+        pre_ffw = layer.pre_feedforward_layernorm(hidden_states).to(compute_dtype)
+        mlp_out = layer.mlp(pre_ffw)
+
+        if getattr(layer, "enable_moe_block", False):
+            h1 = layer.post_feedforward_layernorm_1(mlp_out).to(compute_dtype)
+
+            flat = residual.reshape(-1, residual.shape[-1])
+            _, top_k_w, top_k_idx = layer.router(flat)
+            h2 = layer.pre_feedforward_layernorm_2(flat).to(compute_dtype)
+            h2 = layer.experts(h2, top_k_idx, top_k_w)
+            h2 = h2.reshape(residual.shape)
+            h2 = layer.post_feedforward_layernorm_2(h2).to(compute_dtype)
+
+            mlp_out = h1 + h2
+
+        mlp_out = layer.post_feedforward_layernorm(mlp_out).to(compute_dtype)
+        hidden_states = residual + mlp_out
+        hidden_states = hidden_states * layer.layer_scalar
+
+    hidden_states = base.norm(hidden_states)
+    return hidden_states
+
+
+# ---------------------------------------------------------------------------
+# Sliding-window block-mask builders. Paged-KV aware.
+# ---------------------------------------------------------------------------
+
+
+def _create_sliding_causal_blockmask(page_table: PageTable, B: int, L: int, W: int):
+    def sliding_causal(b, h, q_idx, kv_idx):
+        return (q_idx >= kv_idx) & (q_idx - kv_idx < W)
+
+    return create_block_mask(
+        sliding_causal,
+        B=B,
+        H=None,
+        Q_LEN=L,
+        KV_LEN=L,
+        BLOCK_SIZE=page_table.page_size,
+        device=page_table.device,
+    )
+
+
+def _create_prefill_sliding_blockmask(
+    page_table: PageTable, batch_idx: torch.Tensor, W: int, BLOCK_SIZE: int = 128
+):
+    assert batch_idx.ndim == 2 and batch_idx.shape[0] == 1
+    L = batch_idx.shape[1]
+    docs = batch_idx.view(-1)
+
+    def document_causal_sliding(b, h, q_idx, kv_idx):
+        causal = q_idx >= kv_idx
+        window = (q_idx - kv_idx) < W
+        document = docs[q_idx] == docs[kv_idx]
+        return causal & window & document
+
+    return create_block_mask(
+        document_causal_sliding,
+        B=1,
+        H=None,
+        Q_LEN=L,
+        KV_LEN=L,
+        BLOCK_SIZE=BLOCK_SIZE,
+    )
+
+
+# ---------------------------------------------------------------------------
+# FlexGemma4MoEInference
+# ---------------------------------------------------------------------------
+
+
+class FlexGemma4MoEInference:
+    """Gemma 4 MoE inference engine. API-compatible with ``FlexMoEInference``
+    and ``FlexGptOssInference``.
+
+    Phase 1: eager decode (capture disabled). Phase 3 enables capture.
+    """
+
+    def __init__(
+        self,
+        model,
+        tokenizer,
+        max_batch_size=32,
+        max_seq_length=2048,
+        n_pages=2048,
+        page_size=128,
+        max_new_tokens=512,
+        decode_kernel_options=None,
+        prefill_kernel_options=None,
+        fa4_prefill=None,
+        base_model=None,
+        peft_model=None,
+        cumem_allocator=None,
+        compile_walker=None,
+    ):
+        assert max_seq_length % page_size == 0
+        assert hasattr(model, "model") and hasattr(model.model, "layers"), (
+            "FlexGemma4MoEInference expects a HF CausalLM shape (.model.layers)."
+        )
+        for i, layer in enumerate(model.model.layers):
+            if not getattr(layer, "enable_moe_block", False):
+                # Dense-only Gemma 4 layers mixed with MoE is not expected
+                # for 26B-A4B; if it ever appears we still don't crash —
+                # the walker's ``if enable_moe_block`` guard handles it.
+                continue
+            assert hasattr(layer, "router") and hasattr(layer, "experts"), (
+                f"Layer {i} claims enable_moe_block=True but is missing router/experts."
+            )
+
+        self.model = model
+        self.tokenizer = tokenizer
+        self.device = model.device
+        self.eos_token_id = tokenizer.eos_token_id
+        self.base_model = base_model
+        self.peft_model = peft_model
+        self.max_batch_size = max_batch_size
+        self.max_seq_length = max_seq_length
+        self.page_size = page_size
+        self.max_new_tokens = max_new_tokens
+
+        # Detect per-layer sliding window. Gemma 4 uses one sliding size
+        # across sliding_attention layers.
+        sliding_windows = set()
+        for layer in model.model.layers:
+            sw = getattr(layer.self_attn, "sliding_window", None)
+            if sw is not None:
+                sliding_windows.add(int(sw))
+        if len(sliding_windows) > 1:
+            raise NotImplementedError(
+                f"Mixed sliding-window sizes not supported: {sliding_windows}"
+            )
+        self.sliding_window = next(iter(sliding_windows), None)
+
+        # Detect bnb-4bit stacked experts (no such class ships yet for
+        # Gemma 4; if it appears, flag and force eager).
+        self._has_bnb_experts = any(
+            getattr(layer, "enable_moe_block", False)
+            and type(layer.experts).__name__.endswith("Bnb4bit")
+            for layer in model.model.layers
+        )
+
+        # FA4 kernel branch.
+        if fa4_prefill is None or fa4_prefill:
+            major, _ = torch.cuda.get_device_capability(self.device)
+            supported = major >= 9
+            if fa4_prefill and not supported:
+                import warnings
+                warnings.warn(
+                    f"--fa4_prefill needs sm_90+; found sm_{major}0. Falling "
+                    f"back to the Triton flex_attention backend.",
+                    RuntimeWarning,
+                    stacklevel=2,
+                )
+            fa4_prefill = supported
+        self.fa4_prefill = fa4_prefill
+        self.prefill_q_block = 256 if fa4_prefill else 128
+        self.prefill_kv_block = 128
+        self.decode_kernel_options = (
+            decode_kernel_options
+            if decode_kernel_options is not None
+            else DECODE_KERNEL_OPTIONS_DEFAULT
+        )
+        base_prefill_opts = (
+            prefill_kernel_options
+            if prefill_kernel_options is not None
+            else dict(PREFILL_KERNEL_OPTIONS_DEFAULT)
+        )
+        if fa4_prefill:
+            base_prefill_opts = dict(base_prefill_opts)
+            base_prefill_opts.pop("FORCE_USE_FLEX_ATTENTION", None)
+            base_prefill_opts["BACKEND"] = "FLASH"
+        self.prefill_kernel_options = base_prefill_opts
+
+        from .sleep_mode import kv_cache_pool as _kv_cache_pool
+        with _kv_cache_pool(cumem_allocator):
+            self.page_table = PageTable(
+                n_pages=n_pages,
+                page_size=page_size,
+                max_batch_size=max_batch_size,
+                device=self.device.type,
+            )
+            patch_gemma4_moe_attention_forwards(model, self.page_table)
+
+        self.input_pos_buffer = torch.zeros(
+            max_batch_size, dtype=torch.int32, device=self.device
+        )
+        self.block_mask_logical = self.page_table.create_causal_blockmask(
+            B=max_batch_size,
+            L=max_seq_length,
+        )
+        if self.sliding_window is not None:
+            self.block_mask_logical_sliding = _create_sliding_causal_blockmask(
+                self.page_table,
+                B=max_batch_size,
+                L=max_seq_length,
+                W=self.sliding_window,
+            )
+        else:
+            self.block_mask_logical_sliding = None
+
+        self.cudagraph_captured = False
+        self.graphs = {}
+        self.graph_vars = {}
+        self.graph_bs = None
+
+        # Optional torch.compile walker.
+        if compile_walker is None:
+            compile_walker = os.environ.get("UNSLOTH_FLEX_COMPILE_WALKER", "") == "1"
+        self._moe_walker = call_gemma4_moe_model_with_flex_kwargs
+        if compile_walker:
+            try:
+                self._moe_walker = torch.compile(
+                    call_gemma4_moe_model_with_flex_kwargs,
+                    fullgraph=False,
+                    dynamic=False,
+                )
+                print(
+                    "[flex-gemma4moe] wrapped call_gemma4_moe_model_with_flex_kwargs "
+                    "with torch.compile(fullgraph=False, dynamic=False)"
+                )
+            except Exception as e:
+                print(f"[flex-gemma4moe] torch.compile wrap failed: {e}")
+                self._moe_walker = call_gemma4_moe_model_with_flex_kwargs
+
+        if self._has_bnb_experts:
+            print(
+                "[flex-gemma4moe] bnb-4bit experts detected; CUDA graph capture "
+                "disabled (decode stays eager)."
+            )
+
+    # --- tokenize / prefill / decode --------------------------------------
+
+    def tokenize(self, sequences):
+        for seq in sequences:
+            if seq.input_ids is not None and seq.input_length > 0:
+                continue
+            ids = self.tokenizer(seq.text, return_tensors="pt")["input_ids"].squeeze(0)
+            seq.input_ids = ids
+            seq.input_length = ids.shape[0]
+
+    def _prefill(self, batch: list[Sequence]) -> torch.Tensor:
+        input_ids_list = [seq.input_ids.to(self.device) for seq in batch]
+        input_pos_list = [
+            torch.arange(seq.input_length, dtype=torch.long, device=self.device)
+            for seq in batch
+        ]
+        batch_idx_list = [
+            torch.full(
+                (seq.input_length,), seq.batch_idx, dtype=torch.long, device=self.device
+            )
+            for seq in batch
+        ]
+        input_ids = torch.cat(input_ids_list).view(1, -1)
+        input_pos = torch.cat(input_pos_list).view(1, -1)
+        batch_idx = torch.cat(batch_idx_list).view(1, -1)
+
+        L = input_ids.shape[1]
+        q_block = self.prefill_q_block
+        pad = (q_block - L % q_block) % q_block
+        if pad > 0:
+            input_ids = F.pad(input_ids, (0, pad), value=0)
+            input_pos = F.pad(input_pos, (0, pad), value=0)
+            batch_idx = F.pad(batch_idx, (0, pad), value=0)
+
+        input_lengths = torch.tensor(
+            [s.input_length for s in batch], dtype=torch.long, device=self.device
+        )
+        logits_positions = input_lengths.cumsum(dim=0) - 1
+
+        prefill_block_size = (
+            (self.prefill_q_block, self.prefill_kv_block)
+            if self.fa4_prefill
+            else self.prefill_q_block
+        )
+        mask = self.page_table.create_prefill_blockmask_no_paging(
+            batch_idx, BLOCK_SIZE=prefill_block_size
+        )
+        mask_sliding = None
+        if self.sliding_window is not None:
+            mask_sliding = _create_prefill_sliding_blockmask(
+                self.page_table,
+                batch_idx,
+                W=self.sliding_window,
+                BLOCK_SIZE=prefill_block_size,
+            )
+
+        flex_kwargs = dict(
+            flex_block_mask=mask,
+            flex_block_mask_sliding=mask_sliding,
+            flex_input_pos=input_pos,
+            flex_batch_idx=batch_idx,
+            flex_kernel_options=self.prefill_kernel_options,
+        )
+        position_ids = input_pos
+        hidden = self._moe_walker(
+            self.model, input_ids, position_ids, flex_kwargs
+        )
+        return self.model.lm_head(hidden[:, logits_positions, :]).squeeze(0)
+
+    def _decode_block_mask(self, batch_idx: torch.Tensor, *, sliding: bool):
+        block_mask = (
+            self.block_mask_logical_sliding if sliding else self.block_mask_logical
+        )
+        input_pos = self.input_pos_buffer[batch_idx]
+        assert batch_idx.ndim == 1 and input_pos.ndim == 1
+        B = batch_idx.shape[0]
+        input_block_idx = input_pos // block_mask.BLOCK_SIZE[0]
+        kv_num_blocks = block_mask.kv_num_blocks[batch_idx, :, input_block_idx].view(
+            B, 1, 1
+        )
+        kv_indices = block_mask.kv_indices[batch_idx, :, input_block_idx].view(
+            B, 1, 1, -1
+        )
+        full_num = full_idx = None
+        if block_mask.full_kv_num_blocks is not None:
+            full_num = block_mask.full_kv_num_blocks[
+                batch_idx, :, input_block_idx
+            ].view(B, 1, 1)
+            full_idx = block_mask.full_kv_indices[
+                batch_idx, :, input_block_idx
+            ].view(B, 1, 1, -1)
+
+        if sliding:
+            W = self.sliding_window
+
+            def mask_fn(off):
+                def m(b, h, q_idx, kv_idx):
+                    pos = q_idx + off[b]
+                    return (pos >= kv_idx) & (pos - kv_idx < W)
+                return m
+        else:
+            def mask_fn(off):
+                def m(b, h, q_idx, kv_idx):
+                    return q_idx + off[b] >= kv_idx
+                return m
+
+        seq_length = (1, block_mask.seq_lengths[1])
+        mask = BlockMask.from_kv_blocks(
+            kv_num_blocks,
+            kv_indices,
+            full_num,
+            full_idx,
+            BLOCK_SIZE=block_mask.BLOCK_SIZE,
+            mask_mod=mask_fn(input_pos),
+            seq_lengths=seq_length,
+        )
+        return mask, input_pos
+
+    def _decode_step_eager(self, batch_idx: torch.Tensor, input_ids: torch.Tensor):
+        B = input_ids.shape[0]
+        mask_full, input_pos = self._decode_block_mask(batch_idx, sliding=False)
+        mask_full = self.page_table.convert_logical_block_mask(mask_full, batch_idx)
+        mask_sliding = None
+        if self.sliding_window is not None:
+            ms, _ = self._decode_block_mask(batch_idx, sliding=True)
+            mask_sliding = self.page_table.convert_logical_block_mask(ms, batch_idx)
+
+        position_ids = input_pos.view(B, 1).to(torch.long)
+        flex_kwargs = dict(
+            flex_block_mask=mask_full,
+            flex_block_mask_sliding=mask_sliding,
+            flex_input_pos=input_pos.view(B, 1).to(torch.long),
+            flex_batch_idx=batch_idx,
+            flex_kernel_options=self.decode_kernel_options,
+        )
+        hidden = self._moe_walker(
+            self.model, input_ids.view(B, 1), position_ids, flex_kwargs
+        )
+        return self.model.lm_head(hidden[:, -1, :])
+
+    def _decode_step(
+        self, batch_idx: torch.Tensor, input_ids: torch.Tensor, input_pos: torch.Tensor
+    ):
+        self.input_pos_buffer.zero_()
+        self.input_pos_buffer[batch_idx] = input_pos
+        if not self.cudagraph_captured or self.graph_bs is None:
+            return self._decode_step_eager(batch_idx, input_ids)
+        bs = input_ids.size(0)
+        key = next(x for x in self.graph_bs if x >= bs)
+        graph = self.graphs[key]
+        gv = self.graph_vars
+        for k, v in gv.items():
+            if k != "outputs":
+                v.zero_()
+        gv["input_ids"][:bs] = input_ids
+        gv["batch_idx"][:bs] = batch_idx
+        graph.replay()
+        return gv["outputs"][:bs]
+
+    def capture_decode_cudagraph(self):
+        """Capture one CUDA graph per bs bucket.
+
+        Phase 3 wires this in. bnb-4bit stays eager either way (no
+        stacked bnb experts class ships for Gemma 4 today; guard is
+        defensive).
+        """
+        if self._has_bnb_experts:
+            print(
+                "[flex-gemma4moe] bnb-4bit experts: skipping cudagraph capture; "
+                "decode stays eager."
+            )
+            return
+        try:
+            from unsloth_zoo.temporary_patches.moe_utils import select_moe_backend
+            backend = select_moe_backend()
+        except Exception:
+            backend = None
+        if backend != "grouped_mm":
+            print(
+                f"[flex-gemma4moe] MoE CUDA graph capture requires the "
+                f"'grouped_mm' backend (got {backend!r}); skipping capture."
+            )
+            return
+
+        max_bs = self.max_batch_size
+        reserved_batches = []
+        for bi in range(1, max_bs):
+            try:
+                allocated = self.page_table.allocate()
+                self.page_table.reserve(
+                    allocated,
+                    torch.tensor([allocated], device=self.device, dtype=torch.long),
+                    self.page_size,
+                )
+                reserved_batches.append(allocated)
+            except Exception:
+                break
+
+        input_ids = torch.zeros(max_bs, dtype=torch.int64, device=self.device)
+        batch_idx = torch.arange(max_bs, dtype=torch.int64, device=self.device)
+        outputs = torch.zeros(
+            (max_bs, self.model.config.vocab_size),
+            dtype=self.model.dtype,
+            device=self.device,
+        )
+
+        _env_bs = os.environ.get("UNSLOTH_FLEX_GRAPH_BS")
+        if _env_bs:
+            try:
+                self.graph_bs = [int(x) for x in _env_bs.split(",") if x.strip()]
+            except ValueError:
+                print(
+                    f"[flex-gemma4moe] invalid UNSLOTH_FLEX_GRAPH_BS={_env_bs!r}; "
+                    f"using default bucket ladder"
+                )
+                self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16))
+        else:
+            self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16))
+
+        pool = None
+        for bs in reversed(self.graph_bs):
+            if bs > max_bs:
+                continue
+            print(f"[flex-gemma4moe] capturing CUDA graph for bs={bs}")
+            torch.cuda.synchronize()
+            _ = self._decode_step_eager(batch_idx[:bs], input_ids[:bs])
+            torch.cuda.synchronize()
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph, pool):
+                outputs[:bs] = self._decode_step_eager(batch_idx[:bs], input_ids[:bs])
+            if pool is None:
+                pool = graph.pool()
+            self.graphs[bs] = graph
+            torch.cuda.synchronize()
+        for bi in reserved_batches:
+            self.page_table.erase(bi)
+        self.graph_vars = dict(
+            input_ids=input_ids, batch_idx=batch_idx, outputs=outputs
+        )
+
+    def refresh_inference_from_base(self):
+        """Refresh merged LoRA weights from pristine. Dense refresh on
+        q/k/v/o/router/dense-MLP; MoE refresh on ``gate_up_proj`` /
+        ``down_proj`` in standard ``(E, 2I, H)``/``(E, H, I)``
+        orientation — reuses ``refresh_moe_lora_merge_from_pristine``
+        verbatim."""
+        if self.base_model is None or self.peft_model is None:
+            return 0
+        n = refresh_lora_merge_from_pristine(self.base_model, self.peft_model)
+        try:
+            n += refresh_moe_lora_merge_from_pristine(
+                self.base_model, self.peft_model
+            )
+        except Exception:
+            pass
+        return n
+
+    @torch.inference_mode()
+    def generate(self, sequences: list[Sequence], capture_cudagraph=False):
+        self.tokenize(sequences)
+        waiting = deque(sequences)
+        running = deque()
+        done = []
+
+        if capture_cudagraph and not self.cudagraph_captured:
+            self.capture_decode_cudagraph()
+            if self.graphs:
+                self.cudagraph_captured = True
+
+        while waiting or running:
+            batch = []
+            while waiting and self.page_table.can_reserve(waiting[0].total_length):
+                seq = waiting.popleft()
+                bi = self.page_table.allocate()
+                self.page_table.reserve(
+                    bi,
+                    torch.tensor([bi], device=self.device, dtype=torch.long),
+                    seq.total_length,
+                )
+                seq.batch_idx = bi
+                batch.append(seq)
+            if batch:
+                logits = self._prefill(batch)
+                next_ids = torch.argmax(logits, dim=-1).tolist()
+                for i, seq in enumerate(batch):
+                    seq.last_token_id = next_ids[i]
+                    seq.output_ids.append(next_ids[i])
+                    if (
+                        seq.last_token_id == self.eos_token_id
+                        or len(seq.output_ids) >= seq.max_new_tokens
+                    ):
+                        seq.finished = True
+                        done.append(seq)
+                        self.page_table.erase(seq.batch_idx)
+                    else:
+                        running.append(seq)
+                continue
+
+            decode_batch = []
+            while running:
+                seq = running.popleft()
+                if self.page_table.capacity[seq.batch_idx] >= seq.total_length:
+                    decode_batch.append(seq)
+                elif self.page_table.can_reserve(
+                    seq.total_length, batch_idx_int=seq.batch_idx
+                ):
+                    self.page_table.reserve(
+                        seq.batch_idx,
+                        torch.tensor(
+                            [seq.batch_idx], device=self.device, dtype=torch.long
+                        ),
+                        seq.total_length,
+                    )
+                    decode_batch.append(seq)
+                else:
+                    running.appendleft(seq)
+                    newest = running.pop()
+                    waiting.appendleft(newest)
+                    self.page_table.erase(newest.batch_idx)
+            if not decode_batch:
+                continue
+
+            B = len(decode_batch)
+            bi_tensor = torch.tensor(
+                [s.batch_idx for s in decode_batch],
+                dtype=torch.long,
+                device=self.device,
+            )
+            last_ids = torch.tensor(
+                [s.last_token_id for s in decode_batch],
+                dtype=torch.long,
+                device=self.device,
+            )
+            cur_pos = torch.tensor(
+                [s.total_length - 1 for s in decode_batch],
+                dtype=torch.int32,
+                device=self.device,
+            )
+            logits = self._decode_step(bi_tensor, last_ids, cur_pos)
+            next_ids = torch.argmax(logits, dim=-1).tolist()
+            for i, seq in enumerate(decode_batch):
+                seq.last_token_id = next_ids[i]
+                seq.output_ids.append(next_ids[i])
+                if (
+                    seq.last_token_id == self.eos_token_id
+                    or len(seq.output_ids) >= seq.max_new_tokens
+                ):
+                    seq.finished = True
+                    done.append(seq)
+                    self.page_table.erase(seq.batch_idx)
+                else:
+                    running.append(seq)
+
+        return done

unsloth/inference/flex_gpt_oss.py

@@ -0,0 +1,794 @@
+# SPDX-License-Identifier: GNU Affero General Public License v3.0
+# Copyright 2023-present the Unsloth team. All rights reserved.
+
+"""gpt-oss inference with flex_attention + paged KV cache + attention sinks.
+
+Sibling of ``flex_moe.py`` (Qwen3-MoE). Three arch-specific pieces:
+
+1. **Attention sinks.** Each layer has a learned per-head ``sinks``
+   parameter that biases a virtual sink token in the softmax. We
+   implement it the same way ``unsloth_zoo.flex_attention.attention_sink``
+   does: call flex_attention with ``return_lse=True``, then scale the
+   output by ``sigmoid(lse - sinks[h])``.
+
+2. **Per-layer sliding window.** gpt-oss layer types alternate between
+   full attention and sliding-128. The walker passes two ``BlockMask``
+   objects (``flex_block_mask`` + ``flex_block_mask_sliding``) and each
+   attention forward picks one based on ``self.sliding_window``.
+
+3. **MoE expert stack.** Uses unsloth_zoo's ``GptOssExperts`` forward
+   (``forward_native_grouped_mm`` with the ``"GptOssExperts"`` branch
+   — interleaved gate/up split + ``gate * sigmoid(gate * 1.702)`` — is
+   already there at moe_utils.py:918-972, reused verbatim). The MoE
+   LoRA merge reuses ``refresh_moe_lora_merge_from_pristine`` from
+   ``flex_moe.py`` — the transposed-orientation branch covers
+   gpt-oss's ``(E, H, 2I)`` gate_up_proj layout.
+
+bnb-4bit: ``GptOssExpertsBnb4bit`` uses an ``nn.ModuleList`` per-expert
+loop that can't be CUDA-graph-captured. Detect it in ``__init__`` and
+set ``capture_cudagraph=False`` — decode still benefits from paged KV
+and flex_attention, just not the graph replay.
+"""
+
+from __future__ import annotations
+
+import os
+import types
+from collections import deque
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch.nn.attention.flex_attention import BlockMask, create_block_mask
+
+try:
+    from .flex_qwen3_llama import (
+        DECODE_KERNEL_OPTIONS_DEFAULT,
+        PREFILL_KERNEL_OPTIONS_DEFAULT,
+        Sequence,
+        flex_attention_compiled,
+        refresh_lora_merge_from_pristine,
+    )
+    from .flex_moe import refresh_moe_lora_merge_from_pristine
+    from .flex_paged_attention import PagedKVCache, PageTable
+except ImportError:  # script-mode fallback
+    import sys
+    from pathlib import Path
+    sys.path.insert(0, str(Path(__file__).resolve().parent))
+    from flex_qwen3_llama import (  # noqa: E402
+        DECODE_KERNEL_OPTIONS_DEFAULT,
+        PREFILL_KERNEL_OPTIONS_DEFAULT,
+        Sequence,
+        flex_attention_compiled,
+        refresh_lora_merge_from_pristine,
+    )
+    from flex_moe import refresh_moe_lora_merge_from_pristine  # noqa: E402
+    from flex_paged_attention import PagedKVCache, PageTable  # noqa: E402
+
+
+# gpt-oss rotary is NOT the Llama ``(q * cos) + (rotate_half(q) * sin)``
+# half-dim sin/cos rotation. It does a first/second-half split with
+# ``cos`` / ``sin`` sized to ``head_dim / 2``. Import the reference here
+# and unsqueeze head-dim once per call — same behaviour as
+# ``apply_rotary_pos_emb`` in ``transformers.models.gpt_oss``.
+def _gpt_oss_apply_rotary(q, k, cos, sin):
+    cos = cos.unsqueeze(1)  # [B, 1, S, D/2]
+    sin = sin.unsqueeze(1)
+
+    def rotate(x):
+        first, second = torch.chunk(x, 2, dim=-1)
+        first_ = first * cos - second * sin
+        second_ = second * cos + first * sin
+        return torch.cat((first_, second_), dim=-1)
+
+    return rotate(q), rotate(k)
+
+
+# ---------------------------------------------------------------------------
+# Attention forward: paged KV + flex_attention + sinks + per-layer sliding.
+# ---------------------------------------------------------------------------
+
+
+def make_gptoss_attention_forward(page_table: PageTable):
+    """Return a ``forward`` method for ``GptOssAttention``.
+
+    Differences from the dense Qwen3/Llama flex forward:
+    - No ``q_norm`` / ``k_norm`` (gpt-oss has neither).
+    - Uses ``flex_attention(..., return_lse=True)`` and scales output by
+      ``sigmoid(lse - sinks[h])`` for the sink token.
+    - Picks ``flex_block_mask_sliding`` when the layer has
+      ``self.sliding_window is not None``.
+    """
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask=None,
+        past_key_values=None,
+        cache_position=None,
+        flex_block_mask: Optional[BlockMask] = None,
+        flex_block_mask_sliding: Optional[BlockMask] = None,
+        flex_input_pos: Optional[torch.Tensor] = None,
+        flex_batch_idx: Optional[torch.Tensor] = None,
+        flex_kernel_options: Optional[dict] = None,
+        **kwargs,
+    ):
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        q = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        k = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        q, k = _gpt_oss_apply_rotary(q, k, cos, sin)
+
+        # Paged-KV write. Match the pre-allocated cache dtype; attention
+        # linears may compute in fp32 under autocast.
+        if self._paged_cache is not None and flex_input_pos is not None:
+            cache_dtype = self._paged_cache.k_cache.dtype
+            if k.dtype != cache_dtype:
+                k = k.to(cache_dtype)
+            if v.dtype != cache_dtype:
+                v = v.to(cache_dtype)
+            k, v = self._paged_cache.update(flex_input_pos, k, v, flex_batch_idx)
+
+        # Per-layer block mask dispatch. sliding_window is None on full
+        # attention layers, an int on sliding layers.
+        if self.sliding_window is not None and flex_block_mask_sliding is not None:
+            block_mask = flex_block_mask_sliding
+        else:
+            block_mask = flex_block_mask
+
+        attn_output, logsumexp = flex_attention_compiled(
+            q,
+            k,
+            v,
+            scale=self.scaling,
+            block_mask=block_mask,
+            enable_gqa=True,
+            kernel_options=flex_kernel_options,
+            return_lse=True,
+        )
+
+        # Attention sink. Equivalent to concatenating a sink column
+        # ``sinks[h]`` to the attention logits before softmax. With
+        # ``return_lse=True`` we have ``lse = log sum_k exp(QK[q, k] *
+        # scale)``; the sink-aware softmax scales the output by
+        # ``exp(lse) / (exp(lse) + exp(sinks[h])) = sigmoid(lse -
+        # sinks[h])``. Mirrors ``flex_attention_add_sinks`` in
+        # ``unsloth_zoo/flex_attention/attention_sink.py``.
+        logsumexp = logsumexp - self.sinks.view(1, -1, 1)
+        sink_scale = torch.sigmoid(logsumexp)
+        attn_output = attn_output * sink_scale.unsqueeze(-1).to(attn_output.dtype)
+
+        attn_output = attn_output.transpose(1, 2).reshape(*input_shape, -1).contiguous()
+        return self.o_proj(attn_output), None
+
+    return forward
+
+
+def patch_gptoss_attention_forwards(model: torch.nn.Module, page_table: PageTable):
+    """Attach a ``PagedKVCache`` + replace ``forward`` on every
+    ``GptOssAttention`` layer."""
+    fwd = make_gptoss_attention_forward(page_table)
+    for layer in model.model.layers:
+        attn = layer.self_attn
+        attn._paged_cache = PagedKVCache(
+            page_table,
+            n_heads=model.config.num_key_value_heads,
+            head_dim=model.config.head_dim,
+            dtype=model.dtype,
+        ).to(model.device)
+        attn.forward = types.MethodType(fwd, attn)
+
+
+# ---------------------------------------------------------------------------
+# Walker: pass flex kwargs through each layer, unpack MLP tuple return.
+# ---------------------------------------------------------------------------
+
+
+def call_gpt_oss_model_with_flex_kwargs(model, input_ids, position_ids, flex_kwargs):
+    """Walk a ``GptOssModel`` manually, injecting flex kwargs into each
+    attention call. ``GptOssMLP.forward`` returns
+    ``(hidden_states, router_scores)``; we discard router scores at
+    inference (no load-balance loss)."""
+    base = model.model
+    inputs_embeds = base.embed_tokens(input_ids)
+    position_embeddings = base.rotary_emb(inputs_embeds, position_ids)
+    _cos, _sin = position_embeddings
+    if _cos.dim() == 2:
+        _cos = _cos[position_ids]
+        _sin = _sin[position_ids]
+        position_embeddings = (_cos, _sin)
+    hidden_states = inputs_embeds
+    compute_dtype = inputs_embeds.dtype
+    for layer in base.layers:
+        residual = hidden_states
+        hidden_states = layer.input_layernorm(hidden_states).to(compute_dtype)
+        hidden_states, _ = layer.self_attn(
+            hidden_states,
+            position_embeddings=position_embeddings,
+            **flex_kwargs,
+        )
+        hidden_states = residual + hidden_states.to(compute_dtype)
+        residual = hidden_states
+        hidden_states = layer.post_attention_layernorm(hidden_states).to(compute_dtype)
+        mlp_out = layer.mlp(hidden_states)
+        if isinstance(mlp_out, tuple):
+            hidden_states = mlp_out[0]
+        else:
+            hidden_states = mlp_out
+        hidden_states = residual + hidden_states.to(compute_dtype)
+    hidden_states = base.norm(hidden_states)
+    return hidden_states
+
+
+# ---------------------------------------------------------------------------
+# Sliding-window block-mask builders. Paged-KV aware.
+# ---------------------------------------------------------------------------
+
+
+def _create_sliding_causal_blockmask(page_table: PageTable, B: int, L: int, W: int):
+    """Full-window causal block mask: ``q >= kv`` and ``q - kv < W``.
+    Built against the logical block layout; ``convert_logical_block_mask``
+    wraps it to paged indices later for decode."""
+
+    def sliding_causal(b, h, q_idx, kv_idx):
+        return (q_idx >= kv_idx) & (q_idx - kv_idx < W)
+
+    return create_block_mask(
+        sliding_causal,
+        B=B,
+        H=None,
+        Q_LEN=L,
+        KV_LEN=L,
+        BLOCK_SIZE=page_table.page_size,
+        device=page_table.device,
+    )
+
+
+def _create_prefill_sliding_blockmask(
+    page_table: PageTable, batch_idx: torch.Tensor, W: int, BLOCK_SIZE: int = 128
+):
+    """Document-causal + sliding: ``q >= kv``, ``q - kv < W``, and
+    ``docs[q] == docs[kv]``. Mirrors
+    ``create_prefill_blockmask_no_paging`` with the window term added."""
+    assert batch_idx.ndim == 2 and batch_idx.shape[0] == 1
+    L = batch_idx.shape[1]
+    docs = batch_idx.view(-1)
+
+    def document_causal_sliding(b, h, q_idx, kv_idx):
+        causal = q_idx >= kv_idx
+        window = (q_idx - kv_idx) < W
+        document = docs[q_idx] == docs[kv_idx]
+        return causal & window & document
+
+    return create_block_mask(
+        document_causal_sliding,
+        B=1,
+        H=None,
+        Q_LEN=L,
+        KV_LEN=L,
+        BLOCK_SIZE=BLOCK_SIZE,
+    )
+
+
+# ---------------------------------------------------------------------------
+# FlexGptOssInference
+# ---------------------------------------------------------------------------
+
+
+class FlexGptOssInference:
+    """gpt-oss inference engine. API-compatible with ``FlexMoEInference``.
+
+    Phase 1: eager decode only (no CUDA graph capture). Phase 3 adds the
+    capture path for the bf16 variant; bnb-4bit stays eager because
+    ``GptOssExpertsBnb4bit`` uses an nn.ModuleList per-expert loop that
+    isn't graph-capturable.
+    """
+
+    def __init__(
+        self,
+        model,
+        tokenizer,
+        max_batch_size=32,
+        max_seq_length=2048,
+        n_pages=2048,
+        page_size=128,
+        max_new_tokens=512,
+        decode_kernel_options=None,
+        prefill_kernel_options=None,
+        fa4_prefill=None,
+        base_model=None,
+        peft_model=None,
+        cumem_allocator=None,
+        compile_walker=None,
+    ):
+        assert max_seq_length % page_size == 0
+        assert hasattr(model, "model") and hasattr(model.model, "layers"), (
+            "FlexGptOssInference expects a HF CausalLM shape (.model.layers)."
+        )
+        for i, layer in enumerate(model.model.layers):
+            assert hasattr(layer.self_attn, "sinks"), (
+                f"Layer {i}.self_attn has no sinks — not a gpt-oss attention?"
+            )
+
+        self.model = model
+        self.tokenizer = tokenizer
+        self.device = model.device
+        self.eos_token_id = tokenizer.eos_token_id
+        self.base_model = base_model
+        self.peft_model = peft_model
+        self.max_batch_size = max_batch_size
+        self.max_seq_length = max_seq_length
+        self.page_size = page_size
+        self.max_new_tokens = max_new_tokens
+
+        # Detect per-layer sliding window. gpt-oss alternates full /
+        # sliding-128; record whichever window is used so block-mask
+        # builders pick it up. If no layer has a sliding window, skip
+        # the sliding mask construction entirely.
+        sliding_windows = {
+            int(layer.self_attn.sliding_window)
+            for layer in model.model.layers
+            if layer.self_attn.sliding_window is not None
+        }
+        if len(sliding_windows) > 1:
+            raise NotImplementedError(
+                f"Mixed sliding-window sizes not supported: {sliding_windows}"
+            )
+        self.sliding_window = next(iter(sliding_windows), None)
+
+        # Detect bnb-4bit experts. The bnb variant uses an nn.ModuleList
+        # per-expert loop — disable CUDA graph capture if present.
+        self._has_bnb_experts = any(
+            type(layer.mlp.experts).__name__ == "GptOssExpertsBnb4bit"
+            for layer in model.model.layers
+        )
+
+        # FA4 kernel branch (same as dense/MoE).
+        if fa4_prefill is None or fa4_prefill:
+            major, _ = torch.cuda.get_device_capability(self.device)
+            supported = major >= 9
+            if fa4_prefill and not supported:
+                import warnings
+                warnings.warn(
+                    f"--fa4_prefill needs sm_90+; found sm_{major}0. Falling "
+                    f"back to the Triton flex_attention backend.",
+                    RuntimeWarning,
+                    stacklevel=2,
+                )
+            fa4_prefill = supported
+        self.fa4_prefill = fa4_prefill
+        self.prefill_q_block = 256 if fa4_prefill else 128
+        self.prefill_kv_block = 128
+        self.decode_kernel_options = (
+            decode_kernel_options
+            if decode_kernel_options is not None
+            else DECODE_KERNEL_OPTIONS_DEFAULT
+        )
+        base_prefill_opts = (
+            prefill_kernel_options
+            if prefill_kernel_options is not None
+            else dict(PREFILL_KERNEL_OPTIONS_DEFAULT)
+        )
+        if fa4_prefill:
+            base_prefill_opts = dict(base_prefill_opts)
+            base_prefill_opts.pop("FORCE_USE_FLEX_ATTENTION", None)
+            base_prefill_opts["BACKEND"] = "FLASH"
+        self.prefill_kernel_options = base_prefill_opts
+
+        from .sleep_mode import kv_cache_pool as _kv_cache_pool
+        with _kv_cache_pool(cumem_allocator):
+            self.page_table = PageTable(
+                n_pages=n_pages,
+                page_size=page_size,
+                max_batch_size=max_batch_size,
+                device=self.device.type,
+            )
+            patch_gptoss_attention_forwards(model, self.page_table)
+
+        self.input_pos_buffer = torch.zeros(
+            max_batch_size, dtype=torch.int32, device=self.device
+        )
+        self.block_mask_logical = self.page_table.create_causal_blockmask(
+            B=max_batch_size,
+            L=max_seq_length,
+        )
+        # Sliding-window logical block mask (only if any layer uses it).
+        if self.sliding_window is not None:
+            self.block_mask_logical_sliding = _create_sliding_causal_blockmask(
+                self.page_table,
+                B=max_batch_size,
+                L=max_seq_length,
+                W=self.sliding_window,
+            )
+        else:
+            self.block_mask_logical_sliding = None
+
+        # Cudagraph capture disabled for Phase 1. Phase 3 wires this up
+        # for bf16; bnb-4bit stays False either way.
+        self.cudagraph_captured = False
+        self.graphs = {}
+        self.graph_vars = {}
+
+        # Optional torch.compile walker (Phase 5 turns this on).
+        if compile_walker is None:
+            compile_walker = os.environ.get("UNSLOTH_FLEX_COMPILE_WALKER", "") == "1"
+        self._moe_walker = call_gpt_oss_model_with_flex_kwargs
+        if compile_walker:
+            try:
+                self._moe_walker = torch.compile(
+                    call_gpt_oss_model_with_flex_kwargs,
+                    fullgraph=False,
+                    dynamic=False,
+                )
+                print(
+                    "[flex-gptoss] wrapped call_gpt_oss_model_with_flex_kwargs "
+                    "with torch.compile(fullgraph=False, dynamic=False)"
+                )
+            except Exception as e:
+                print(f"[flex-gptoss] torch.compile wrap failed: {e}")
+                self._moe_walker = call_gpt_oss_model_with_flex_kwargs
+
+        if self._has_bnb_experts:
+            print(
+                "[flex-gptoss] bnb-4bit experts detected; CUDA graph capture "
+                "disabled (decode stays eager)."
+            )
+
+    # --- tokenize / prefill / decode --------------------------------------
+
+    def tokenize(self, sequences):
+        for seq in sequences:
+            if seq.input_ids is not None and seq.input_length > 0:
+                continue
+            ids = self.tokenizer(seq.text, return_tensors="pt")["input_ids"].squeeze(0)
+            seq.input_ids = ids
+            seq.input_length = ids.shape[0]
+
+    def _prefill(self, batch: list[Sequence]) -> torch.Tensor:
+        input_ids_list = [seq.input_ids.to(self.device) for seq in batch]
+        input_pos_list = [
+            torch.arange(seq.input_length, dtype=torch.long, device=self.device)
+            for seq in batch
+        ]
+        batch_idx_list = [
+            torch.full(
+                (seq.input_length,), seq.batch_idx, dtype=torch.long, device=self.device
+            )
+            for seq in batch
+        ]
+        input_ids = torch.cat(input_ids_list).view(1, -1)
+        input_pos = torch.cat(input_pos_list).view(1, -1)
+        batch_idx = torch.cat(batch_idx_list).view(1, -1)
+
+        L = input_ids.shape[1]
+        q_block = self.prefill_q_block
+        pad = (q_block - L % q_block) % q_block
+        if pad > 0:
+            input_ids = F.pad(input_ids, (0, pad), value=0)
+            input_pos = F.pad(input_pos, (0, pad), value=0)
+            batch_idx = F.pad(batch_idx, (0, pad), value=0)
+
+        input_lengths = torch.tensor(
+            [s.input_length for s in batch], dtype=torch.long, device=self.device
+        )
+        logits_positions = input_lengths.cumsum(dim=0) - 1
+
+        prefill_block_size = (
+            (self.prefill_q_block, self.prefill_kv_block)
+            if self.fa4_prefill
+            else self.prefill_q_block
+        )
+        mask = self.page_table.create_prefill_blockmask_no_paging(
+            batch_idx, BLOCK_SIZE=prefill_block_size
+        )
+        mask_sliding = None
+        if self.sliding_window is not None:
+            mask_sliding = _create_prefill_sliding_blockmask(
+                self.page_table,
+                batch_idx,
+                W=self.sliding_window,
+                BLOCK_SIZE=prefill_block_size,
+            )
+
+        flex_kwargs = dict(
+            flex_block_mask=mask,
+            flex_block_mask_sliding=mask_sliding,
+            flex_input_pos=input_pos,
+            flex_batch_idx=batch_idx,
+            flex_kernel_options=self.prefill_kernel_options,
+        )
+        position_ids = input_pos
+        hidden = self._moe_walker(
+            self.model, input_ids, position_ids, flex_kwargs
+        )
+        return self.model.lm_head(hidden[:, logits_positions, :]).squeeze(0)
+
+    def _decode_block_mask(self, batch_idx: torch.Tensor, *, sliding: bool):
+        block_mask = (
+            self.block_mask_logical_sliding if sliding else self.block_mask_logical
+        )
+        input_pos = self.input_pos_buffer[batch_idx]
+        assert batch_idx.ndim == 1 and input_pos.ndim == 1
+        B = batch_idx.shape[0]
+        input_block_idx = input_pos // block_mask.BLOCK_SIZE[0]
+        kv_num_blocks = block_mask.kv_num_blocks[batch_idx, :, input_block_idx].view(
+            B, 1, 1
+        )
+        kv_indices = block_mask.kv_indices[batch_idx, :, input_block_idx].view(
+            B, 1, 1, -1
+        )
+        full_num = full_idx = None
+        if block_mask.full_kv_num_blocks is not None:
+            full_num = block_mask.full_kv_num_blocks[
+                batch_idx, :, input_block_idx
+            ].view(B, 1, 1)
+            full_idx = block_mask.full_kv_indices[
+                batch_idx, :, input_block_idx
+            ].view(B, 1, 1, -1)
+
+        if sliding:
+            W = self.sliding_window
+
+            def mask_fn(off):
+                def m(b, h, q_idx, kv_idx):
+                    pos = q_idx + off[b]
+                    return (pos >= kv_idx) & (pos - kv_idx < W)
+                return m
+        else:
+            def mask_fn(off):
+                def m(b, h, q_idx, kv_idx):
+                    return q_idx + off[b] >= kv_idx
+                return m
+
+        seq_length = (1, block_mask.seq_lengths[1])
+        mask = BlockMask.from_kv_blocks(
+            kv_num_blocks,
+            kv_indices,
+            full_num,
+            full_idx,
+            BLOCK_SIZE=block_mask.BLOCK_SIZE,
+            mask_mod=mask_fn(input_pos),
+            seq_lengths=seq_length,
+        )
+        return mask, input_pos
+
+    def _decode_step_eager(self, batch_idx: torch.Tensor, input_ids: torch.Tensor):
+        B = input_ids.shape[0]
+        mask_full, input_pos = self._decode_block_mask(batch_idx, sliding=False)
+        mask_full = self.page_table.convert_logical_block_mask(mask_full, batch_idx)
+        mask_sliding = None
+        if self.sliding_window is not None:
+            ms, _ = self._decode_block_mask(batch_idx, sliding=True)
+            mask_sliding = self.page_table.convert_logical_block_mask(ms, batch_idx)
+
+        position_ids = input_pos.view(B, 1).to(torch.long)
+        flex_kwargs = dict(
+            flex_block_mask=mask_full,
+            flex_block_mask_sliding=mask_sliding,
+            flex_input_pos=input_pos.view(B, 1).to(torch.long),
+            flex_batch_idx=batch_idx,
+            flex_kernel_options=self.decode_kernel_options,
+        )
+        hidden = self._moe_walker(
+            self.model, input_ids.view(B, 1), position_ids, flex_kwargs
+        )
+        return self.model.lm_head(hidden[:, -1, :])
+
+    def _decode_step(
+        self, batch_idx: torch.Tensor, input_ids: torch.Tensor, input_pos: torch.Tensor
+    ):
+        self.input_pos_buffer.zero_()
+        self.input_pos_buffer[batch_idx] = input_pos
+        if not self.cudagraph_captured:
+            return self._decode_step_eager(batch_idx, input_ids)
+        bs = input_ids.size(0)
+        key = next(x for x in self.graph_bs if x >= bs)
+        graph = self.graphs[key]
+        gv = self.graph_vars
+        for k, v in gv.items():
+            if k != "outputs":
+                v.zero_()
+        gv["input_ids"][:bs] = input_ids
+        gv["batch_idx"][:bs] = batch_idx
+        graph.replay()
+        return gv["outputs"][:bs]
+
+    def capture_decode_cudagraph(self):
+        """Capture one CUDA graph per bs bucket. bf16 experts only;
+        bnb-4bit (nn.ModuleList per-expert loop) can't be captured and
+        ``__init__`` already set ``capture_cudagraph`` to False for that
+        variant — this method is only reached when the engine-level
+        flag is still True."""
+        if self._has_bnb_experts:
+            print(
+                "[flex-gptoss] bnb-4bit experts: skipping cudagraph capture; "
+                "decode stays eager."
+            )
+            return
+        try:
+            from unsloth_zoo.temporary_patches.moe_utils import select_moe_backend
+            backend = select_moe_backend()
+        except Exception:
+            backend = None
+        if backend != "grouped_mm":
+            print(
+                f"[flex-gptoss] MoE CUDA graph capture requires the "
+                f"'grouped_mm' backend (got {backend!r}); skipping capture."
+            )
+            return
+
+        max_bs = self.max_batch_size
+        reserved_batches = []
+        for bi in range(1, max_bs):
+            try:
+                allocated = self.page_table.allocate()
+                self.page_table.reserve(
+                    allocated,
+                    torch.tensor([allocated], device=self.device, dtype=torch.long),
+                    self.page_size,
+                )
+                reserved_batches.append(allocated)
+            except Exception:
+                break
+
+        input_ids = torch.zeros(max_bs, dtype=torch.int64, device=self.device)
+        batch_idx = torch.arange(max_bs, dtype=torch.int64, device=self.device)
+        outputs = torch.zeros(
+            (max_bs, self.model.config.vocab_size),
+            dtype=self.model.dtype,
+            device=self.device,
+        )
+
+        _env_bs = os.environ.get("UNSLOTH_FLEX_GRAPH_BS")
+        if _env_bs:
+            try:
+                self.graph_bs = [int(x) for x in _env_bs.split(",") if x.strip()]
+            except ValueError:
+                print(
+                    f"[flex-gptoss] invalid UNSLOTH_FLEX_GRAPH_BS={_env_bs!r}; "
+                    f"using default bucket ladder"
+                )
+                self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16))
+        else:
+            self.graph_bs = [1, 2, 4, 8] + list(range(16, max_bs + 1, 16))
+
+        pool = None
+        for bs in reversed(self.graph_bs):
+            if bs > max_bs:
+                continue
+            print(f"[flex-gptoss] capturing CUDA graph for bs={bs}")
+            torch.cuda.synchronize()
+            _ = self._decode_step_eager(batch_idx[:bs], input_ids[:bs])
+            torch.cuda.synchronize()
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph, pool):
+                outputs[:bs] = self._decode_step_eager(batch_idx[:bs], input_ids[:bs])
+            if pool is None:
+                pool = graph.pool()
+            self.graphs[bs] = graph
+            torch.cuda.synchronize()
+        for bi in reserved_batches:
+            self.page_table.erase(bi)
+        self.graph_vars = dict(
+            input_ids=input_ids, batch_idx=batch_idx, outputs=outputs
+        )
+
+    def refresh_inference_from_base(self):
+        """Refresh merged LoRA weights from pristine. Reuses the dense
+        refresh for q/k/v/o/router and the MoE refresh
+        (``refresh_moe_lora_merge_from_pristine``) for ``gate_up_proj``
+        / ``down_proj`` in the ``(E, H, 2I)`` transposed orientation."""
+        if self.base_model is None or self.peft_model is None:
+            return 0
+        n = refresh_lora_merge_from_pristine(self.base_model, self.peft_model)
+        try:
+            n += refresh_moe_lora_merge_from_pristine(
+                self.base_model, self.peft_model
+            )
+        except Exception:
+            pass
+        return n
+
+    @torch.inference_mode()
+    def generate(self, sequences: list[Sequence], capture_cudagraph=False):
+        """Decode loop. Phase 1: always eager."""
+        self.tokenize(sequences)
+        waiting = deque(sequences)
+        running = deque()
+        done = []
+
+        if capture_cudagraph and not self.cudagraph_captured:
+            self.capture_decode_cudagraph()
+            if self.graphs:
+                self.cudagraph_captured = True
+
+        while waiting or running:
+            batch = []
+            while waiting and self.page_table.can_reserve(waiting[0].total_length):
+                seq = waiting.popleft()
+                bi = self.page_table.allocate()
+                self.page_table.reserve(
+                    bi,
+                    torch.tensor([bi], device=self.device, dtype=torch.long),
+                    seq.total_length,
+                )
+                seq.batch_idx = bi
+                batch.append(seq)
+            if batch:
+                logits = self._prefill(batch)
+                next_ids = torch.argmax(logits, dim=-1).tolist()
+                for i, seq in enumerate(batch):
+                    seq.last_token_id = next_ids[i]
+                    seq.output_ids.append(next_ids[i])
+                    if (
+                        seq.last_token_id == self.eos_token_id
+                        or len(seq.output_ids) >= seq.max_new_tokens
+                    ):
+                        seq.finished = True
+                        done.append(seq)
+                        self.page_table.erase(seq.batch_idx)
+                    else:
+                        running.append(seq)
+                continue
+
+            decode_batch = []
+            while running:
+                seq = running.popleft()
+                if self.page_table.capacity[seq.batch_idx] >= seq.total_length:
+                    decode_batch.append(seq)
+                elif self.page_table.can_reserve(
+                    seq.total_length, batch_idx_int=seq.batch_idx
+                ):
+                    self.page_table.reserve(
+                        seq.batch_idx,
+                        torch.tensor(
+                            [seq.batch_idx], device=self.device, dtype=torch.long
+                        ),
+                        seq.total_length,
+                    )
+                    decode_batch.append(seq)
+                else:
+                    running.appendleft(seq)
+                    newest = running.pop()
+                    waiting.appendleft(newest)
+                    self.page_table.erase(newest.batch_idx)
+            if not decode_batch:
+                continue
+
+            B = len(decode_batch)
+            bi_tensor = torch.tensor(
+                [s.batch_idx for s in decode_batch],
+                dtype=torch.long,
+                device=self.device,
+            )
+            last_ids = torch.tensor(
+                [s.last_token_id for s in decode_batch],
+                dtype=torch.long,
+                device=self.device,
+            )
+            cur_pos = torch.tensor(
+                [s.total_length - 1 for s in decode_batch],
+                dtype=torch.int32,
+                device=self.device,
+            )
+            logits = self._decode_step(bi_tensor, last_ids, cur_pos)
+            next_ids = torch.argmax(logits, dim=-1).tolist()
+            for i, seq in enumerate(decode_batch):
+                seq.last_token_id = next_ids[i]
+                seq.output_ids.append(next_ids[i])
+                if (
+                    seq.last_token_id == self.eos_token_id
+                    or len(seq.output_ids) >= seq.max_new_tokens
+                ):
+                    seq.finished = True
+                    done.append(seq)
+                    self.page_table.erase(seq.batch_idx)
+                else:
+                    running.append(seq)
+
+        return done

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

@@ -245,6 +245,9 @@ VLLM_SUPPORTED_VLM = [
     "mistral3",
     "qwen3_vl",
     "qwen3_vl_moe",
+    "gemma4",
+    "qwen3_5",
+    "qwen3_5_moe",
 ]
 VLLM_NON_LORA_VLM = [
     "mllama",
@@ -606,7 +609,16 @@ class FastBaseModel:
         vllm_enable_lora = True
 
         if is_vlm and fast_inference:
-            if not any(arch in VLLM_SUPPORTED_VLM for arch in model_types):
+            # The UNSLOTH_FAST_INFERENCE=1 flex backend ships with text-only
+            # support for Gemma-4-E2B-it today. Allow gemma4 through the
+            # vision-model path when the flex backend is selected; the
+            # vLLM-only compat list below still gates the default path.
+            _use_flex = os.environ.get("UNSLOTH_FAST_INFERENCE", "0") == "1"
+            _flex_allowed = {"gemma4", "qwen3_5", "qwen3_5_moe"}
+            if not (
+                any(arch in VLLM_SUPPORTED_VLM for arch in model_types)
+                or (_use_flex and any(arch in _flex_allowed for arch in model_types))
+            ):
                 raise RuntimeError(
                     f"Unsloth: Fast inference is only supported for Language models and Qwen2.5-VL, Gemma3 among vision models. "
                     f"Found architectures: {', '.join(model_types)}!"
@@ -912,9 +924,15 @@ class FastBaseModel:
         verify_fp8_support_if_applicable(model_config)
 
         raise_handler = RaiseUninitialized()
-        if not fast_inference:
+        _use_flex_fast_inference = (
+            fast_inference and os.environ.get("UNSLOTH_FAST_INFERENCE", "0") == "1"
+        )
+        if (not fast_inference) or _use_flex_fast_inference:
+            # Shared by the standard HF load and the flex-inference load.
             # Prevent load_in_fp8 from being forwarded into HF internal model loading
             load_in_fp8 = kwargs.pop("load_in_fp8", None)
+            # ``max_batch_size`` is a FlexEngine kwarg, not an HF one.
+            _flex_max_batch_size = kwargs.pop("max_batch_size", 32)
             # Transformers 5.x @strict config classes reject unexpected kwargs.
             # Move config-level attributes onto the config object directly.
             _num_labels = kwargs.pop("num_labels", None)
@@ -944,6 +962,16 @@ class FastBaseModel:
                 fast_inference = fast_inference,
             )
             if hasattr(model, "generate"):
+                if _use_flex_fast_inference:
+                    # Defer engine construction until tokenizer/processor is
+                    # loaded; see the wiring block further down.
+                    model._unsloth_needs_flex_engine = dict(
+                        dtype = dtype,
+                        max_seq_length = max_seq_length,
+                        max_lora_rank = max_lora_rank,
+                        max_batch_size = _flex_max_batch_size,
+                        gpu_memory_utilization = gpu_memory_utilization,
+                    )
                 model.fast_generate = make_fast_generate_wrapper(model.generate)
                 model.fast_generate_batches = error_out_no_vllm
             if offload_embedding:
@@ -1028,6 +1056,22 @@ class FastBaseModel:
                 is_vision_model = is_vlm,
                 fp8_mode = fp8_mode,
             )
+            # why: vLLM's -O3 compile backend (compilation_config=3) hits
+            # _decompose_size_nodes "Tried to erase Node size_N but it still
+            # had N users" on Gemma 4 multimodal models (both dense E2B and
+            # MoE 26B-A4B). Until the upstream FX splitter is fixed, fall
+            # back to compilation_config=0 (no piecewise compile) for gemma4
+            # so fast_inference still loads. Runtime is slower than -O3 but
+            # correctness is preserved.
+            #
+            # The same FX splitter bug reproduces on Qwen3.5 / Qwen3.6
+            # (qwen3_5 dense, qwen3_5_moe). The hybrid GatedDeltaNet +
+            # full-attention layers defeat piecewise compile's node erasure
+            # heuristics. Force eager for these archs too.
+            _eager_arches = {"gemma4", "qwen3_5", "qwen3_5_moe"}
+            if any(arch in _eager_arches for arch in (model_types or [])):
+                load_vllm_kwargs.setdefault("compilation_config", 0)
+                load_vllm_kwargs.setdefault("enforce_eager", True)
             for allowed_arg in allowed_args:
                 if allowed_arg not in load_vllm_kwargs and allowed_arg in kwargs:
                     load_vllm_kwargs[allowed_arg] = kwargs[allowed_arg]
@@ -1223,6 +1267,18 @@ class FastBaseModel:
                 raise _patch_err
         model = post_patch_loss_function(model)
 
+        # UNSLOTH_FAST_INFERENCE=1 path: install the lazy ``vllm_engine``
+        # sentinel now that the tokenizer/processor is available. The real
+        # ``FlexEngine(...)`` construction is deferred to
+        # :func:`build_flex_engine` so the batch-size dimension can be
+        # sized from the GRPO rollout shape. Keeps the training model's
+        # forward intact — the engine carries its own deep-copy.
+        if hasattr(model, "_unsloth_needs_flex_engine"):
+            from unsloth.inference.flex_engine import install_flex_sentinel
+
+            _tok_for_flex = getattr(tokenizer, "tokenizer", tokenizer)
+            install_flex_sentinel(model, _tok_for_flex)
+
         # Log Unsloth version for future fastpaths for inference
         if hasattr(model, "config"):
             model.config.update({"unsloth_version": __version__})
@@ -1512,6 +1568,10 @@ class FastBaseModel:
         patch_saving_functions(model, vision = True)
         patch_peft_fast_inference(model)
 
+        _engine = getattr(model, "vllm_engine", None)
+        if _engine is not None and hasattr(_engine, "bind_peft_model"):
+            _engine.bind_peft_model(model)
+
         # Add for_inference and for_training
         model.for_training = functools.partial(FastBaseModel.for_training, model)
         model.for_inference = functools.partial(FastBaseModel.for_inference, model)