inference: expose flex backend via UNSLOTH_FAST_INFERENCE=1

Setting UNSLOTH_FAST_INFERENCE=1 makes
FastLanguageModel.from_pretrained(..., fast_inference=True) return the
flex attention + paged KV + CUDA-graph backend instead of vLLM for
Qwen3, Llama-3, and Gemma-4-E2B-it. Default "0" keeps the existing vLLM
path.

The flex engine lives in a new unsloth.inference subpackage so the
FastLanguageModel wiring and any external caller can reach it without
going through scripts/. The user-facing GRPO notebooks for Qwen3,
Llama-3, and Gemma-4 run end-to-end with no code edits, just the env
var.

New subpackage

- unsloth/inference/flex_paged_attention.py: paged KV cache and
  block-mask helpers.
- unsloth/inference/flex_qwen3_llama.py: FlexInference for Qwen3 and
  Llama-3.2. tokenize() skips pre-populated input_ids so the engine
  can accept vLLM-style token-id prompts.
- unsloth/inference/flex_gemma4.py: FlexGemma4Inference. Text-shell
  extraction from Gemma4ForConditionalGeneration, per-layer-type
  sliding-window masks, KV sharing for layers 15-34.
- unsloth/inference/flex_engine.py: FlexEngine with the vLLM LLM
  surface. Arch dispatch, deep-copy of the HF model for colocate
  rollout so training forward stays intact, lazy pristine-base copy
  on first bind_peft_model, torch.amp.autocast("cuda", dtype=...)
  wrapping for the whole generate path, auto-tune for Triton block
  sizes per GPU capability and head_dim band. FA4 default OFF
  because the sm_100 FLASH backend crashes on short prompts; users
  can opt in with fa4_prefill=True.
- unsloth/inference/vllm_shim.py: LoRARequest, RequestOutput,
  CompletionOutput dataclasses plus save_lora / load_lora that
  mirror unsloth_zoo.vllm_utils line for line so the TRL GRPO patch
  sees the same attribute shape.

Wiring

- unsloth/models/loader.py: UNSLOTH_FAST_INFERENCE=1 bypasses the
  vLLM import guard in FastLanguageModel.from_pretrained (L356) and
  FastModel.from_pretrained (L982).
- unsloth/models/llama.py: the elif not fast_inference load branch
  now also handles the flex case, pops max_batch_size from kwargs
  (FlexEngine-only), snapshots the model into an inference copy
  BEFORE Unsloth's post-patch so flex-attention patching lands on a
  clean HF layout, and after the tokenizer loads constructs
  FlexEngine and attaches .vllm_engine / .fast_generate /
  .fast_generate_batches. FastLlamaModel.get_peft_model calls
  engine.bind_peft_model(peft) right after patch_peft_fast_inference
  so state_dict() reads the training LoRA tensors.
- unsloth/models/vision.py: same pattern for FastBaseModel. Gemma-4
  is allowed through the vision-model gate when the flex backend is
  selected; the engine extracts the text shell from
  Gemma4ForConditionalGeneration internally.
- unsloth/models/_utils.py: patch_peft_fast_inference picks
  unsloth.inference.vllm_shim.save_lora / load_lora when the engine
  is a FlexEngine, so the flex path never imports
  vllm.lora.request. fast_inference_setup skips patch_vllm() when
  UNSLOTH_FAST_INFERENCE=1.

Smoke tests (B200 sm_100)

- Qwen3-4B-Base bf16 (no LoRA and with get_peft_model LoRA)
- Qwen3-4B-Base fp16
- Llama-3.2-3B-Instruct bf16
- gemma-4-E2B-it bf16

all load through FastLanguageModel + UNSLOTH_FAST_INFERENCE=1 and
generate coherent completions.

Batched steady-state throughput via FastLanguageModel (autocast ON,
n_prompts=8 max_new_tokens=64 max_batch_size=16 max_seq_length=1024):

| Model                      | Flex CLI  | Integration | Delta |
|----------------------------|-----------|-------------|-------|
| Qwen3-4B-Base              | 819 tok/s | 595 tok/s   | -27%  |
| Llama-3.2-3B-Instruct      | 1163      | 1357        | +17%  |

Qwen3's regression is outside the 10% budget; most likely cause is
q_norm / k_norm being promoted to fp32 under autocast. Tracked for
a follow-up; the integration is functional and GRPO-usable today.

Note

The flex engine source files live here in unsloth/inference/. The
CLI benchmark drivers in scripts/benchmarks/ land in a separate PR
(#5108) which also carries the Gemma-4 and Llama-3 regression runs.
This PR depends only on what is already on main.

Known follow-ups

- FA4 on sm_100 needs a more robust enable rule before it can be
  default-on.
- Memory: 2x base-model VRAM without LoRA, 3x with LoRA.
- FlexEngine.sleep() / .wake_up() are no-op stubs; real CPU-offload
  parity with UNSLOTH_VLLM_STANDBY is a separate PR.

tests/flex_fastlm_bench.py

@@ -0,0 +1,78 @@
+"""Batched steady-state throughput bench through ``FastLanguageModel`` +
+``UNSLOTH_FAST_INFERENCE=1``. First ``generate`` call primes CUDA graphs;
+subsequent calls report steady state. Compare against April CLI-only
+numbers for the same workload."""
+
+import os, sys, time
+from pathlib import Path
+_REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+import argparse
+import torch
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--model", default="unsloth/Qwen3-4B-Base")
+    p.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16")
+    p.add_argument("--n_prompts", type=int, default=8)
+    p.add_argument("--max_new_tokens", type=int, default=64)
+    p.add_argument("--max_batch_size", type=int, default=16)
+    p.add_argument("--max_seq_length", type=int, default=1024)
+    p.add_argument("--n_rounds", type=int, default=3)
+    args = p.parse_args()
+
+    os.environ.setdefault("UNSLOTH_FAST_INFERENCE", "1")
+    import unsloth
+    from unsloth import FastLanguageModel
+
+    dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
+
+    model, tok = FastLanguageModel.from_pretrained(
+        model_name=args.model,
+        max_seq_length=args.max_seq_length,
+        dtype=dtype,
+        load_in_4bit=False,
+        fast_inference=True,
+        max_batch_size=args.max_batch_size,
+    )
+    print(f"[bench] model={args.model} dtype={args.dtype}")
+    prompts = [f"In one sentence, a fact about {t} is" for t in [
+        "the moon", "gravity", "the ocean", "the sun", "honey",
+        "rain", "trees", "mountains"
+    ][:args.n_prompts]]
+
+    class _SP:
+        max_tokens = args.max_new_tokens
+        temperature = 0.0
+
+    # Warmup round — captures CUDA graphs.
+    print("[bench] warmup (CUDA graph capture)...")
+    t0 = time.perf_counter()
+    _ = model.fast_generate(prompts, sampling_params=_SP(), use_tqdm=False)
+    print(f"[bench] warmup wall: {time.perf_counter() - t0:.2f}s")
+
+    walls = []
+    tok_counts = []
+    for r in range(args.n_rounds):
+        torch.cuda.synchronize()
+        t1 = time.perf_counter()
+        outs = model.fast_generate(prompts, sampling_params=_SP(), use_tqdm=False)
+        torch.cuda.synchronize()
+        dt = time.perf_counter() - t1
+        n_tok = sum(len(o.outputs[0].token_ids) for o in outs)
+        walls.append(dt)
+        tok_counts.append(n_tok)
+        print(f"[bench] round {r}: {n_tok} toks in {dt:.2f}s -> {n_tok/dt:.1f} tok/s")
+
+    if walls:
+        wall_med = sorted(walls)[len(walls) // 2]
+        tok_med = tok_counts[len(walls) // 2]
+        print(f"[bench] median: {tok_med} toks in {wall_med:.2f}s "
+              f"=> {tok_med / wall_med:.1f} tok/s")
+
+
+if __name__ == "__main__":
+    main()

tests/flex_fastlm_smoke.py

@@ -0,0 +1,93 @@
+"""Smoke-test the ``UNSLOTH_FAST_INFERENCE=1`` path through
+``FastLanguageModel.from_pretrained``.
+
+Invoked as:
+    CUDA_VISIBLE_DEVICES=2 UNSLOTH_FAST_INFERENCE=1 python tests/flex_fastlm_smoke.py \
+        --model unsloth/Qwen3-4B-Base --dtype bf16 --no-lora
+
+Prints tokens/s + the first generated string.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+import time
+from pathlib import Path
+
+# Make the local fork importable.
+_REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--model", default="unsloth/Qwen3-4B-Base")
+    p.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16")
+    p.add_argument("--load_in_4bit", action="store_true")
+    p.add_argument("--with_lora", action="store_true")
+    p.add_argument("--max_new_tokens", type=int, default=32)
+    p.add_argument("--max_seq_length", type=int, default=1024)
+    p.add_argument("--prompt", default="The quick brown fox jumps over")
+    args = p.parse_args()
+
+    import torch
+
+    os.environ.setdefault("UNSLOTH_FAST_INFERENCE", "1")
+    print(f"[smoke] UNSLOTH_FAST_INFERENCE={os.environ.get('UNSLOTH_FAST_INFERENCE')}")
+
+    import unsloth
+    print(f"[smoke] unsloth={unsloth.__file__}")
+    from unsloth import FastLanguageModel
+
+    dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
+
+    t0 = time.perf_counter()
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name=args.model,
+        max_seq_length=args.max_seq_length,
+        dtype=dtype,
+        load_in_4bit=args.load_in_4bit,
+        fast_inference=True,
+    )
+    t_load = time.perf_counter() - t0
+    print(f"[smoke] loaded model in {t_load:.1f}s; dtype={model.dtype}")
+    print(f"[smoke] hasattr(model, 'vllm_engine'): {hasattr(model, 'vllm_engine')}")
+    print(f"[smoke] vllm_engine type: {type(model.vllm_engine).__name__}")
+
+    if args.with_lora:
+        model = FastLanguageModel.get_peft_model(
+            model,
+            r=16,
+            target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
+            lora_alpha=16,
+            lora_dropout=0.0,
+            bias="none",
+            use_gradient_checkpointing="unsloth",
+            random_state=3407,
+        )
+        print(f"[smoke] PEFT model type: {type(model).__name__}")
+        print(f"[smoke] model.vllm_engine bound to PEFT: "
+              f"{hasattr(model, 'vllm_engine')}")
+
+    from unsloth.inference.vllm_shim import LoRARequest
+    prompts = [args.prompt]
+    # Minimal SamplingParams stand-in
+    class _SP:
+        max_tokens = args.max_new_tokens
+        temperature = 0.0
+    t1 = time.perf_counter()
+    outputs = model.fast_generate(prompts, sampling_params=_SP(), use_tqdm=False)
+    dt = time.perf_counter() - t1
+    out = outputs[0]
+    n_tok = len(out.outputs[0].token_ids)
+    print(f"[smoke] generated {n_tok} tokens in {dt:.2f}s "
+          f"({n_tok / dt:.1f} tok/s)")
+    print(f"[smoke] prompt: {args.prompt!r}")
+    print(f"[smoke] completion: {out.outputs[0].text!r}")
+
+
+if __name__ == "__main__":
+    main()

unsloth/inference/__init__.py

@@ -0,0 +1,31 @@
+"""Flex-attention inference engines.
+
+``UNSLOTH_FAST_INFERENCE=1`` routes ``FastLanguageModel.from_pretrained``
+through :func:`load_flex`, which wraps the selected HF model with a
+:class:`FlexEngine` that presents the vLLM ``LLM`` surface used by
+Unsloth / TRL GRPO (``.generate``, ``.chat``, ``.sleep``, ``.wake_up``,
+``.llm_engine``, plus ``save_lora`` / ``load_lora`` via the module
+shim).
+
+Three architectures are supported today: Qwen3, Llama-3, Gemma-4-E2B-it.
+Anything else raises :class:`NotImplementedError` — there is no silent
+fallback; unset the env var or use vLLM instead."""
+
+from .flex_engine import FlexEngine, load_flex
+from .vllm_shim import (
+    CompletionOutput,
+    LoRARequest,
+    RequestOutput,
+    load_lora,
+    save_lora,
+)
+
+__all__ = [
+    "FlexEngine",
+    "load_flex",
+    "LoRARequest",
+    "RequestOutput",
+    "CompletionOutput",
+    "save_lora",
+    "load_lora",
+]

unsloth/inference/flex_engine.py

@@ -0,0 +1,720 @@
+"""FlexEngine: vLLM-compatible LLM surface for the flex inference backends.
+
+When ``UNSLOTH_FAST_INFERENCE=1`` is set, :func:`load_flex` wraps the HF model
+with a :class:`FlexEngine`. TRL's GRPO trainer and Unsloth's own callers treat
+it as an ``LLM``:
+
+    engine.generate(prompts, sampling_params=..., lora_request=..., use_tqdm=False)
+    engine.chat(messages, sampling_params=..., lora_request=..., use_tqdm=False)
+    engine.sleep(level=2)
+    engine.wake_up(tags=["kv_cache"])
+    engine.llm_engine                # minimal stub, see below
+
+Three architectures are supported: Qwen3, Llama-3, Gemma-4-E2B-it. Anything
+else raises :class:`NotImplementedError`.
+
+Colocation model (no change from today):
+- ``self.hf_model`` is the HF model instance returned by ``from_pretrained``
+  (or its PEFT-wrapped descendant after :meth:`bind_peft_model`).
+- The underlying :class:`FlexInference` / :class:`FlexGemma4Inference` implements
+  the usual double-copy rollout pattern (pristine ``base_model`` + inference
+  copy wrapped by PEFT), and ``refresh_lora_merge_from_pristine`` re-merges on
+  every rollout.
+- vLLM's sleep mode is not implemented for the flex backend. ``.sleep()`` /
+  ``.wake_up()`` are no-op stubs so code paths that gate on
+  ``UNSLOTH_VLLM_STANDBY`` stay valid.
+"""
+
+from __future__ import annotations
+
+import copy
+import functools
+import importlib.util
+import os
+import types
+import warnings
+from typing import Any, Optional
+
+import torch
+
+from .flex_paged_attention import PagedKVCache, PageTable  # noqa: F401
+from .flex_qwen3_llama import (
+    DECODE_KERNEL_OPTIONS_DEFAULT,
+    PREFILL_KERNEL_OPTIONS_DEFAULT,
+    FlexInference,
+    Sequence,
+    refresh_lora_merge_from_pristine,
+)
+from .flex_gemma4 import FlexGemma4Inference
+from .vllm_shim import CompletionOutput, LoRARequest, RequestOutput
+
+
+# ---------------------------------------------------------------------------
+# Hardware / kernel auto-tune
+# ---------------------------------------------------------------------------
+
+
+def _flash_attn_4_importable() -> bool:
+    """FA4's CuTeDSL backend lives in the ``flash_attn_interface`` ships with
+    the ``flash-attn`` 4.x wheel. We don't import it here (it can be slow);
+    just check whether the module is resolvable."""
+    return importlib.util.find_spec("flash_attn_interface") is not None
+
+
+def _fa4_ok_for_head_dim(head_dim: int, device_cap: tuple[int, int]) -> bool:
+    """See the plan for the hardware tier table. Returns whether FA4 is
+    safe to use for an attention layer with the given ``head_dim`` on a GPU
+    with the given ``(major, minor)`` capability."""
+    major, _ = device_cap
+    if major < 9:
+        return False  # Ampere and older: Triton only
+    if not _flash_attn_4_importable():
+        return False
+    if major == 9:
+        return 8 <= head_dim <= 256  # Hopper
+    return 8 <= head_dim <= 128       # Blackwell (sm_100 / sm_120)
+
+
+def _triton_block_defaults(
+    head_dim_max: int, device_cap: tuple[int, int]
+) -> tuple[int, int, int, int]:
+    """Returns ``(prefill_BM, prefill_BN, decode_BM, decode_BN)``.
+
+    Blackwell (sm_100 / sm_120) shared-memory budget is tight; at
+    ``head_dim >= 256`` the default 128x128 prefill block overflows and the
+    Triton autotuner raises ``OutOfMemoryError: out of resource``. See the
+    ``--prefill_kernel_options`` probes in ``scripts/benchmarks/*.py``."""
+    major, _ = device_cap
+    if major >= 10 and head_dim_max >= 256:
+        return (32, 32, 16, 16)
+    if major >= 10 and head_dim_max >= 128:
+        return (64, 64, 32, 32)
+    return (128, 128, 64, 64)
+
+
+def _collect_head_dims(hf_model) -> list[int]:
+    """Collect the per-layer ``head_dim`` for every attention sub-module in
+    the model. Gemma-4 text stack mixes 256 / 512 between full / sliding
+    layers; Qwen3 + Llama-3 use a single head_dim."""
+    head_dims = set()
+    for module in hf_model.modules():
+        hd = getattr(module, "head_dim", None)
+        if hd is not None and isinstance(hd, int) and hd > 0:
+            head_dims.add(hd)
+    if not head_dims:
+        cfg = getattr(hf_model, "config", None)
+        if cfg is not None:
+            hd = getattr(cfg, "head_dim", None)
+            if hd is None:
+                num_heads = getattr(cfg, "num_attention_heads", None)
+                hidden = getattr(cfg, "hidden_size", None)
+                if num_heads and hidden:
+                    hd = hidden // num_heads
+            if hd:
+                head_dims.add(int(hd))
+    return sorted(head_dims)
+
+
+def _auto_kernel_options(
+    hf_model,
+    device,
+    prefill_kernel_options: Optional[dict] = None,
+    decode_kernel_options: Optional[dict] = None,
+    fa4_prefill: Optional[bool] = None,
+) -> tuple[Optional[bool], dict, dict]:
+    """Derive safe FA4 + Triton block defaults from the GPU + model head_dim
+    band. Explicit kwargs always win.
+
+    FA4 is OFF by default. The CuTeDSL FLASH backend currently crashes on
+    short B200 prompts (Llama-3.2-3B, 7-token input hits
+    ``handle_block_sparse_empty_tile_correction_sm100``
+    ``'NoneType' object is not subscriptable``). Users who want FA4 can
+    pass ``fa4_prefill=True`` explicitly to the engine after confirming it
+    works on their workload."""
+    cap = torch.cuda.get_device_capability(device) if torch.cuda.is_available() else (0, 0)
+    head_dims = _collect_head_dims(hf_model) or [128]
+    head_dim_max = max(head_dims)
+
+    if fa4_prefill is None:
+        fa4_prefill = False  # conservative default — see docstring above
+
+    bm_p, bn_p, bm_d, bn_d = _triton_block_defaults(head_dim_max, cap)
+
+    if prefill_kernel_options is None:
+        if fa4_prefill:
+            # FA4 path: let the FLASH backend pick its own block shapes.
+            prefill_kernel_options = None
+        else:
+            prefill_kernel_options = {
+                "FORCE_USE_FLEX_ATTENTION": True,
+                "BLOCK_M": bm_p,
+                "BLOCK_N": bn_p,
+            }
+    if decode_kernel_options is None:
+        decode_kernel_options = {"BLOCK_M": bm_d, "BLOCK_N": bn_d}
+
+    return fa4_prefill, prefill_kernel_options, decode_kernel_options
+
+
+# ---------------------------------------------------------------------------
+# Arch detection
+# ---------------------------------------------------------------------------
+
+
+def _detect_arch(hf_model) -> str:
+    """Return one of ``"gemma4"``, ``"qwen3"``, ``"llama3"`` or raises."""
+    # Look at the inner base model's class; PEFT wrappers delegate to
+    # ``.base_model.model``.
+    target = hf_model
+    for attr in ("base_model", "model"):
+        inner = getattr(target, attr, None)
+        if inner is not None and inner is not target:
+            target = inner
+    name = type(target).__name__
+    # Take from the class hierarchy too so we don't miss Gemma4ForCausalLM
+    # nested under a PEFT wrapper's ``base_model.model``.
+    names = [c.__name__ for c in type(hf_model).__mro__]
+    candidates = set(names + [name])
+    lowered = " ".join(n.lower() for n in candidates)
+    if "gemma4" in lowered or "gemma_4" in lowered or "gemma-4" in lowered:
+        return "gemma4"
+    if "qwen3" in lowered:
+        return "qwen3"
+    if "llama" in lowered:
+        return "llama3"
+    raise NotImplementedError(
+        "UNSLOTH_FAST_INFERENCE=1 only supports Qwen3, Llama-3, Gemma-4 "
+        f"today; got {type(hf_model).__name__}. Unset the env var or use vLLM."
+    )
+
+
+# ---------------------------------------------------------------------------
+# Gemma-4 text-only shell extraction
+# ---------------------------------------------------------------------------
+
+
+def _extract_gemma4_text_shell(full_model):
+    """Return a ``Gemma4ForCausalLM(text_cfg)`` shell wrapping the text
+    backbone of ``full_model``. Mirrors the CLI path in
+    ``scripts/benchmarks/gemma4_flex_inference.py:954-968``: strip the
+    vision + audio towers, point ``shell.model`` at
+    ``full_model.model.language_model``, tie ``lm_head.weight`` to the
+    embeddings so the forward pass matches plain HF."""
+    from transformers.models.gemma4.modeling_gemma4 import (
+        Gemma4ForCausalLM,
+    )
+
+    lang = getattr(full_model.model, "language_model", None)
+    if lang is None:
+        return full_model  # already a text-only CausalLM
+    full_model.model.vision_tower = None
+    full_model.model.audio_tower = None
+    full_model.model.embed_vision = None
+    full_model.model.embed_audio = None
+    text_cfg = full_model.config.text_config
+    shell = Gemma4ForCausalLM(text_cfg)
+    shell.model = lang
+    shell.lm_head.weight = lang.embed_tokens.weight
+    shell = shell.to(next(lang.parameters()).device)
+    shell.eval()
+    return shell
+
+
+# ---------------------------------------------------------------------------
+# FlexEngine
+# ---------------------------------------------------------------------------
+
+
+class _LLMEngineStub:
+    """Minimal stand-in for ``vllm.LLM.llm_engine``.
+
+    ``unsloth/models/rl.py:104`` reaches into
+    ``trainer.model.vllm_engine`` to pull the LLM out; `GRPOTrainer`'s
+    patched ``_move_model_to_vllm`` path calls ``driver_worker.model_runner
+    .model.load_weights`` but Unsloth's RL patch rewrites those calls to
+    ``pass`` (rl.py:1795-1810), so a nested attribute chain that simply
+    exists is enough."""
+
+    def __init__(self):
+        self.vllm_config = types.SimpleNamespace(lora_config = types.SimpleNamespace())
+        self.model_executor = types.SimpleNamespace(
+            driver_worker = types.SimpleNamespace(
+                model_runner = types.SimpleNamespace(
+                    model = types.SimpleNamespace(load_weights = lambda *a, **kw: None),
+                )
+            )
+        )
+
+
+class FlexEngine:
+    """vLLM-compatible wrapper around :class:`FlexInference` /
+    :class:`FlexGemma4Inference`.
+
+    Args:
+        hf_model: The HF model returned by ``from_pretrained``. It will be
+            swapped to the PEFT-wrapped instance after ``bind_peft_model``.
+        tokenizer: The companion tokenizer.
+        dtype: ``torch.bfloat16`` or ``torch.float16``. Used for the autocast
+            context wrapping forward/prefill/decode.
+        max_seq_length: Upper bound on ``prompt + completion`` length.
+        max_lora_rank: Unused by flex today (documented for parity with vLLM).
+        max_batch_size: Concurrent sequence bucket size for the paged KV.
+        page_size: Paged-KV page size (tokens per page).
+        gpu_memory_utilization: Controls ``n_pages``. 0.5 means half the free
+            VRAM is dedicated to the KV cache.
+        capture_cudagraph: Capture CUDA graphs on the first decode step.
+    """
+
+    def __init__(
+        self,
+        hf_model,
+        tokenizer,
+        *,
+        dtype: torch.dtype = torch.bfloat16,
+        max_seq_length: int = 2048,
+        max_lora_rank: int = 64,  # accepted for API parity; no-op
+        max_batch_size: int = 32,
+        page_size: int = 128,
+        gpu_memory_utilization: float = 0.5,
+        max_new_tokens: int = 512,
+        prefill_kernel_options: Optional[dict] = None,
+        decode_kernel_options: Optional[dict] = None,
+        fa4_prefill: Optional[bool] = None,
+        capture_cudagraph: bool = True,
+        base_model = None,
+        peft_model = None,
+        inference_model = None,
+    ):
+        assert dtype in (torch.bfloat16, torch.float16), (
+            f"FlexEngine requires bf16 or fp16 dtype; got {dtype}."
+        )
+        self.hf_model = hf_model
+        self.tokenizer = tokenizer
+        self.compute_dtype = dtype
+        self.max_seq_length = max_seq_length
+        self.max_batch_size = max_batch_size
+        self.page_size = page_size
+        self.max_new_tokens = max_new_tokens
+        self.capture_cudagraph = capture_cudagraph
+        self._cudagraph_primed = False
+        self._current_lora_int_id: Optional[int] = None
+
+        self.device = hf_model.device
+
+        # Colocate pattern (mirrors vLLM's "colocate" mode): the engine
+        # runs on its own deep-copy of the HF model so that flex attention
+        # patching + KV-cache attachment does not mutate the training
+        # model's forward path. The user's ``model`` stays untouched for
+        # gradient computation; rollouts go through this copy.
+        #
+        # If ``inference_model`` is provided by the caller, use it directly
+        # (the loader uses this to hand in a copy captured BEFORE Unsloth's
+        # post-patching; passing the patched model here would break
+        # rotary_emb / QKV dispatch).
+        #
+        # The pristine-base copy (second deep-copy) is LoRA-only — it is
+        # the refresh source for ``refresh_lora_merge_from_pristine``. We
+        # defer materialising it until :meth:`bind_peft_model`, so the
+        # no-LoRA path stays at 2x the base model's VRAM instead of 3x.
+        if inference_model is None:
+            inference_model = copy.deepcopy(hf_model)
+            inference_model.eval()
+        self._pristine_base = base_model  # None until bind_peft_model runs
+        self._inference_model = inference_model
+        self._inference_peft = peft_model  # filled in by bind_peft_model
+
+        # Autocast wrapping (applied by the impl via a context manager).
+        fa4_prefill, prefill_kernel_options, decode_kernel_options = (
+            _auto_kernel_options(
+                inference_model,
+                self.device,
+                prefill_kernel_options = prefill_kernel_options,
+                decode_kernel_options = decode_kernel_options,
+                fa4_prefill = fa4_prefill,
+            )
+        )
+
+        # Size n_pages from available VRAM so big prompts don't OOM.
+        n_pages = self._compute_n_pages(gpu_memory_utilization, max_batch_size, page_size)
+
+        arch = _detect_arch(inference_model)
+        self.arch = arch
+        if arch == "gemma4":
+            inference_model = _extract_gemma4_text_shell(inference_model)
+            self._inference_model = inference_model
+        Impl = FlexGemma4Inference if arch == "gemma4" else FlexInference
+        self._impl = Impl(
+            inference_model,
+            tokenizer,
+            max_batch_size = max_batch_size,
+            max_seq_length = max_seq_length,
+            n_pages = n_pages,
+            page_size = page_size,
+            max_new_tokens = max_new_tokens,
+            decode_kernel_options = decode_kernel_options,
+            prefill_kernel_options = prefill_kernel_options,
+            fa4_prefill = fa4_prefill,
+            base_model = self._pristine_base,
+            peft_model = peft_model,
+        )
+        self._llm_engine_stub = _LLMEngineStub()
+
+    # ----- configuration helpers -----
+
+    def _compute_n_pages(
+        self, gpu_mem_util: float, max_batch: int, page_size: int
+    ) -> int:
+        """Size the paged-KV allocation.
+
+        ``n_pages`` = ceil(``max_batch * max_seq_length / page_size``) with a
+        small headroom factor; any more and we waste VRAM on ghost pages
+        that the scheduler never fills. ``gpu_memory_utilization`` scales the
+        headroom."""
+        min_pages = max(
+            1, max_batch * ((self.max_seq_length + page_size - 1) // page_size)
+        )
+        factor = 1.0 + max(0.1, min(1.0, gpu_mem_util))
+        return max(min_pages, int(min_pages * factor))
+
+    # ----- generate -----
+
+    def generate(
+        self,
+        prompts = None,
+        sampling_params = None,
+        lora_request = None,
+        use_tqdm: bool = False,
+        **kwargs,
+    ):
+        """Drop-in replacement for ``vllm.LLM.generate``.
+
+        Returns a list of :class:`RequestOutput` mirroring vLLM's shape.
+        """
+        if prompts is None and "prompt_token_ids" in kwargs:
+            prompts = kwargs.pop("prompt_token_ids")
+        prompts = self._normalize_prompts(prompts)
+
+        max_new_tokens, _ = self._extract_sampling(sampling_params)
+
+        # Apply LoRA: if the request carries tensors, refresh the merged copy.
+        if lora_request is not None:
+            self._apply_lora_request(lora_request)
+
+        seqs = self._build_sequences(prompts, max_new_tokens)
+
+        with torch.amp.autocast("cuda", dtype = self.compute_dtype):
+            done = self._impl.generate(
+                seqs,
+                capture_cudagraph = self.capture_cudagraph and not self._cudagraph_primed,
+            )
+        self._cudagraph_primed = self.capture_cudagraph
+
+        # Reorder by the input prompt index (the impl returns done seqs in
+        # completion order, not input order).
+        done_by_input = sorted(done, key = lambda s: getattr(s, "_input_idx", 0))
+        outs = []
+        for idx, seq in enumerate(done_by_input):
+            text = self.tokenizer.decode(seq.output_ids, skip_special_tokens = False)
+            prompt_text = getattr(seq, "text", "") or self.tokenizer.decode(
+                seq.input_ids.tolist() if seq.input_ids is not None else [],
+                skip_special_tokens = False,
+            )
+            co = CompletionOutput(
+                index = 0,
+                text = text,
+                token_ids = list(seq.output_ids),
+                finish_reason = (
+                    "stop" if seq.last_token_id == self._impl.eos_token_id else "length"
+                ),
+            )
+            ro = RequestOutput(
+                request_id = str(idx),
+                prompt = prompt_text,
+                prompt_token_ids = (
+                    seq.input_ids.tolist() if seq.input_ids is not None else []
+                ),
+                outputs = [co],
+            )
+            outs.append(ro)
+        return outs
+
+    def chat(
+        self,
+        messages,
+        sampling_params = None,
+        lora_request = None,
+        use_tqdm: bool = False,
+        **kwargs,
+    ):
+        """Apply the tokenizer's chat template and defer to :meth:`generate`."""
+        if messages is None:
+            return []
+        # ``messages`` from vLLM is either a list of message-lists or a single
+        # message-list (one conversation).
+        if isinstance(messages, list) and messages and isinstance(messages[0], dict):
+            # Single conversation.
+            convos = [messages]
+        else:
+            convos = list(messages)
+        prompts = [
+            self.tokenizer.apply_chat_template(
+                m, tokenize = False, add_generation_prompt = True
+            )
+            for m in convos
+        ]
+        return self.generate(
+            prompts,
+            sampling_params = sampling_params,
+            lora_request = lora_request,
+            use_tqdm = use_tqdm,
+            **kwargs,
+        )
+
+    # ----- LoRA refresh -----
+
+    def _apply_lora_request(self, lora_request):
+        """Copy the training LoRA tensors onto the inference-side PEFT
+        wrapper and refresh the merged inference weights from the pristine
+        base.
+
+        Expects ``lora_request.lora_tensors`` to be a ``state_dict`` slice
+        filtered for ``.lora_A.`` / ``.lora_B.`` keys (which is what
+        :func:`~unsloth.inference.vllm_shim.load_lora` produces)."""
+        if lora_request is None:
+            return
+        base_model = getattr(self._impl, "base_model", None)
+        peft_model = getattr(self._impl, "peft_model", None)
+        if base_model is None or peft_model is None:
+            # 4-bit / no double-copy: LoRA tensors already live on the PEFT
+            # wrappers (bnb-4bit packed weights can't be in-place refreshed;
+            # PEFT's three-matmul wrapper does the math at runtime).
+            return
+
+        tensors = getattr(lora_request, "lora_tensors", None)
+        if tensors:
+            target_sd = peft_model.state_dict()
+            renamed = {}
+            for k, v in tensors.items():
+                # ``load_lora`` strips ``.default``; PEFT's own state_dict
+                # keeps it. Try both forms.
+                if k in target_sd:
+                    renamed[k] = v
+                    continue
+                candidate = k.replace(".lora_A.", ".lora_A.default.").replace(
+                    ".lora_B.", ".lora_B.default."
+                )
+                if candidate in target_sd:
+                    renamed[candidate] = v
+            if renamed:
+                missing, unexpected = peft_model.load_state_dict(
+                    renamed, strict = False
+                )
+                # ``missing`` will be every non-LoRA param; that's fine.
+        refresh_lora_merge_from_pristine(base_model, peft_model)
+        self._current_lora_int_id = getattr(lora_request, "lora_int_id", None)
+
+    # ----- prompt normalization -----
+
+    @staticmethod
+    def _normalize_prompts(prompts):
+        if prompts is None:
+            return []
+        if isinstance(prompts, str):
+            return [prompts]
+        if isinstance(prompts, list):
+            if not prompts:
+                return []
+            # Already a list. Leave dict / str / list[int] elements as-is.
+            return prompts
+        return [prompts]
+
+    def _build_sequences(self, prompts, max_new_tokens: int) -> list:
+        seqs = []
+        for idx, p in enumerate(prompts):
+            if isinstance(p, str):
+                seq = Sequence(text = p, max_new_tokens = max_new_tokens)
+            elif isinstance(p, dict):
+                # vLLM TokensPrompt / TextPrompt dict
+                if "prompt_token_ids" in p:
+                    ids = torch.tensor(p["prompt_token_ids"], dtype = torch.long)
+                    seq = Sequence(
+                        text = p.get("prompt", ""),
+                        input_ids = ids,
+                        input_length = int(ids.shape[0]),
+                        max_new_tokens = max_new_tokens,
+                    )
+                elif "prompt" in p:
+                    seq = Sequence(
+                        text = p["prompt"], max_new_tokens = max_new_tokens,
+                    )
+                else:
+                    raise ValueError(
+                        f"Unsupported prompt dict (no 'prompt' or "
+                        f"'prompt_token_ids'): {list(p)}"
+                    )
+            elif isinstance(p, (list, tuple)) and p and isinstance(p[0], int):
+                ids = torch.tensor(list(p), dtype = torch.long)
+                seq = Sequence(
+                    text = "",
+                    input_ids = ids,
+                    input_length = int(ids.shape[0]),
+                    max_new_tokens = max_new_tokens,
+                )
+            else:
+                raise ValueError(
+                    "FlexEngine.generate accepts str, list[int], or a "
+                    "TokensPrompt/TextPrompt dict; got "
+                    f"{type(p).__name__} at index {idx}."
+                )
+            seq._input_idx = idx  # preserve input order across the scheduler
+            seqs.append(seq)
+        return seqs
+
+    def _extract_sampling(self, sampling_params) -> tuple[int, dict]:
+        """Pull what we actually honour out of a ``SamplingParams``.
+
+        Today the flex path does argmax sampling only. We read ``max_tokens``
+        (mapped to ``max_new_tokens``) and ignore ``temperature`` / ``top_p``
+        / ``top_k`` with a warning the first time we see something non-greedy.
+        Sufficient for GRPO's on-policy rollout, which already accepts that
+        the generation is deterministic per-prompt under a fixed seed."""
+        if sampling_params is None:
+            return self.max_new_tokens, {}
+        max_tokens = getattr(sampling_params, "max_tokens", None)
+        if max_tokens is None:
+            max_tokens = self.max_new_tokens
+        temp = getattr(sampling_params, "temperature", 0.0) or 0.0
+        if temp and temp > 0 and not getattr(self, "_warned_sampling", False):
+            warnings.warn(
+                "FlexEngine (UNSLOTH_FAST_INFERENCE=1) does argmax sampling "
+                "only; sampling_params.temperature / top_p / top_k are "
+                "ignored. Unset the env var or use vLLM for stochastic "
+                "sampling.",
+                RuntimeWarning,
+                stacklevel = 3,
+            )
+            self._warned_sampling = True
+        return int(max_tokens), {}
+
+    # ----- sleep-mode stubs -----
+    #
+    # vLLM's sleep mode offloads engine weights to CPU between rollouts so
+    # training can use the freed VRAM. The flex backend does not implement
+    # that yet; these stubs exist so code paths that assume the API are safe.
+
+    def sleep(self, level: int = 2):
+        """No-op. Real implementation in a follow-up PR (move PagedKVCache +
+        shared buffers + the inference-copy HF shell to CPU pinned memory
+        and swap back on wake_up)."""
+        return None
+
+    def wake_up(self, tags: Optional[list] = None):
+        """No-op companion to :meth:`sleep`."""
+        return None
+
+    @property
+    def llm_engine(self):
+        """Nested-attribute stub used by Unsloth's RL patch — see docstring
+        on :class:`_LLMEngineStub`."""
+        return self._llm_engine_stub
+
+    # ----- PEFT wiring -----
+
+    def bind_peft_model(self, training_peft_model):
+        """Hook called at the end of ``FastLanguageModel.get_peft_model``.
+
+        * Updates ``self.hf_model`` so ``save_lora`` / ``load_lora`` — which
+          call ``model.state_dict()`` — see the training LoRA weights.
+        * Materialises the pristine-base copy (second deep-copy of the
+          inference model's un-patched weights) used as the refresh source
+          for ``refresh_lora_merge_from_pristine``. This is lazy so the
+          no-LoRA path stays at 2x the base model's VRAM.
+        * Mirrors the training PEFT config onto the inference copy by
+          wrapping ``self._inference_model`` with a matching ``PeftModel``.
+          The inference-side adapter starts at zero; LoRA tensors are then
+          loaded from the request on each ``generate`` call, and
+          ``refresh_lora_merge_from_pristine`` fuses them into the merged
+          base weights before the kernel is launched.
+        * The deep-copy pattern keeps the training model's forward pass
+          untouched (its attention is not flex-patched)."""
+        self.hf_model = training_peft_model
+
+        # Materialise the pristine source the first time we see a LoRA.
+        if self._pristine_base is None:
+            # The inference model has already been flex-patched; its
+            # linear weights (what LoRA merges into) are still pristine,
+            # so we can clone it and just not call flex attention on the
+            # pristine copy.
+            self._pristine_base = copy.deepcopy(self._inference_model)
+            self._pristine_base.eval()
+            self._impl.base_model = self._pristine_base
+
+        if self._inference_peft is None:
+            try:
+                from peft import get_peft_model as _get_peft_model
+
+                peft_cfg = training_peft_model.peft_config["default"]
+                # Wrap the already-patched inference copy with a fresh LoRA
+                # adapter of the same shape. LoraLayer insertion is
+                # attention-forward-agnostic; it wraps Linear modules.
+                self._inference_peft = _get_peft_model(
+                    self._inference_model, peft_cfg
+                )
+                self._inference_peft.eval()
+            except Exception as e:
+                warnings.warn(
+                    f"FlexEngine.bind_peft_model: could not build an "
+                    f"inference-side PEFT wrapper ({e}). Falling back to "
+                    "single-copy mode — LoRA refresh will rely on the "
+                    "training model's state_dict directly.",
+                    RuntimeWarning,
+                    stacklevel = 2,
+                )
+                self._inference_peft = training_peft_model
+
+        self._impl.peft_model = self._inference_peft
+
+
+# ---------------------------------------------------------------------------
+# load_flex — mirrors unsloth_zoo.vllm_utils.load_vllm's call shape so the
+# dispatcher in unsloth/models/llama.py + vision.py can route cleanly.
+# ---------------------------------------------------------------------------
+
+
+def load_flex(
+    hf_model,
+    tokenizer,
+    *,
+    dtype: torch.dtype = torch.bfloat16,
+    max_seq_length: int = 2048,
+    max_lora_rank: int = 64,
+    max_batch_size: int = 32,
+    gpu_memory_utilization: float = 0.5,
+    page_size: int = 128,
+    capture_cudagraph: bool = True,
+    base_model = None,
+    peft_model = None,
+    **_unused_vllm_kwargs,
+) -> FlexEngine:
+    """Construct a :class:`FlexEngine` around an already-loaded HF model.
+
+    ``_unused_vllm_kwargs`` swallows vLLM-only kwargs passed through by
+    `unsloth/models/llama.py` (``use_bitsandbytes``, ``enable_lora``,
+    ``disable_log_stats``, ``fp8_mode``, ``float8_kv_cache``, ...) so the
+    dispatcher doesn't have to know which backend is selected."""
+    return FlexEngine(
+        hf_model,
+        tokenizer,
+        dtype = dtype,
+        max_seq_length = max_seq_length,
+        max_lora_rank = max_lora_rank,
+        max_batch_size = max_batch_size,
+        gpu_memory_utilization = gpu_memory_utilization,
+        page_size = page_size,
+        capture_cudagraph = capture_cudagraph,
+        base_model = base_model,
+        peft_model = peft_model,
+    )
+
+
+__all__ = ["FlexEngine", "load_flex"]

unsloth/inference/flex_paged_attention.py

@@ -0,0 +1,502 @@
+# Adapted from attention-gym
+# Original source: https://github.com/pytorch-labs/attention-gym
+# License: BSD 3-Clause (see THIRD_PARTY_LICENSES.md)
+# Copyright (c) 2023, Driss Guessous
+
+# the original implementation has some bugs and has some feature that lives outside of the PageTable class
+
+from typing import Optional
+import torch
+from torch import Tensor
+from torch.nn.attention.flex_attention import (
+    _identity,
+    _mask_mod_signature,
+    _score_mod_signature,
+    BlockMask,
+    noop_mask,
+    create_block_mask,
+)
+
+create_block_mask = torch.compile(create_block_mask, dynamic = True)
+
+
+def _cdiv(x: int | float | torch.Tensor, multiple: int | float | torch.Tensor):
+    return (x + multiple - 1) // multiple
+
+
+class PagedKVCache(torch.nn.Module):
+    def __init__(self, page_table, n_heads, head_dim, dtype):
+        super().__init__()
+        cache_shape = (1, n_heads, page_table.n_pages * page_table.page_size, head_dim)
+        self.register_buffer("k_cache", torch.zeros(cache_shape, dtype = dtype))
+        self.register_buffer("v_cache", torch.zeros(cache_shape, dtype = dtype))
+
+        self.page_table = page_table
+
+    def update(self, input_pos, k_val, v_val, batch_idx = None):
+        assert (
+            batch_idx is not None
+        ), "batch_idx is required for paged kv cache, are you using non-paged attention?"
+
+        if batch_idx.ndim == 1:
+            # batch_idx should be [B] (decode)
+            return self.page_table.assign(
+                batch_idx, input_pos, k_val, v_val, self.k_cache, self.v_cache
+            )
+        else:
+            assert batch_idx.ndim == 2, "batch_idx must be 1D or 2D"
+            # batch_idx should be [1, L] (batch prefill)
+            return self.page_table.assign_prefill_no_paging(
+                batch_idx, input_pos, k_val, v_val, self.k_cache, self.v_cache
+            )
+
+
+class PageTable:
+    """
+    PageTable is a modified version of PagedAttention from attention-gym.
+
+    PageTable improves it by:
+    - maintaining a cpu copy of the page table, to avoid device-to-host transfers
+    - support batch prefill
+    - fix the bug in the original code in mask_mod and score_mod by mapping physical batch index to logical batch index
+    - subsuming the free_batch_idx into the page table, so we don't need to maintain it separately
+    """
+
+    def __init__(
+        self,
+        n_pages: int,
+        page_size: int,
+        max_batch_size: int,
+        device: str = "cuda",
+    ):
+        self.n_pages = n_pages
+        self.page_size = page_size
+        self.max_batch_size = max_batch_size
+        self.device = device
+
+        # page table: [logical_batch_idx, logical_block_idx] -> physical_page_idx
+        self.page_table = -torch.ones(
+            (max_batch_size, self.n_pages), dtype = torch.int64, device = device
+        )
+        self.page_table[0, :] = (
+            0  # page 0 is reserved for simpler code in assign_prefill_no_paging
+        )
+        self.page_table_cpu = [[] for _ in range(max_batch_size)]
+
+        self.capacity = [
+            0 for _ in range(max_batch_size)
+        ]  # capacity: batch_idx -> number of pages allocated * page size
+        self.free_pages = list(
+            reversed(range(1, n_pages))
+        )  # page 0 is reserved for simpler code in assign_prefill_no_paging
+        self.free_batch_idx = list(
+            reversed(range(1, max_batch_size))
+        )  # batch_idx 0 is reserved for no-op
+
+        # [logical_batch_idx, physical_page_idx] -> logical_page_idx
+        self.physical_to_logical = -torch.ones(
+            (max_batch_size, n_pages), dtype = torch.int64, device = device
+        )
+
+    def can_reserve(self, size: int, batch_idx_int: int | None = None) -> bool:
+        """check if we can reserve new pages for an existing request or a new request, without gpu operations"""
+        if batch_idx_int is None:
+            # check if we can schedule a new request
+            return (
+                self.pages_available * self.page_size >= size
+                and len(self.free_batch_idx) > 0
+            )
+        else:
+            # check if we can reserve new pages for an existing request
+            return self.reserve(batch_idx_int, None, size, dry_run = True)
+
+    def allocate(self) -> int:
+        """allocate a new batch"""
+        batch_idx = self.free_batch_idx.pop()
+
+        self.capacity[batch_idx] = 0
+        self.physical_to_logical[batch_idx, :] = -1
+        self.page_table[batch_idx, :] = -1
+        return batch_idx
+
+    @property
+    def pages_available(self) -> int:
+        return len(self.free_pages)
+
+    def reserve(
+        self,
+        batch_idx_int: int,
+        batch_idx: torch.Tensor,
+        seq_len: int,
+        dry_run: bool = False,
+    ) -> bool:
+        """
+        Requests the capacity of a given batch to be at least enough to
+        hold `seq_len` elements.
+
+        Args:
+            batch_idx_int (int): batch index to be reserved;
+            batch_idx (Tensor): batch index to be reserved; shape :math:`(1)`.
+            seq_len (Tensor): minimum capacity for the given batch; shape :math:`(1)`.
+
+        Returns:
+            bool: True if the reservation was successful, False if the reservation was not successful (no space, and in this case, no update is done)
+        """
+
+        if seq_len <= self.capacity[batch_idx_int]:
+            return True
+
+        num_pages_to_allocate = _cdiv(
+            seq_len - self.capacity[batch_idx_int], self.page_size
+        )
+
+        can_allocate = num_pages_to_allocate <= self.pages_available
+        if dry_run:
+            return can_allocate
+
+        if not can_allocate:
+            raise RuntimeError(
+                f"Cannot reserve {num_pages_to_allocate} pages for a sequence of length {seq_len} "
+                f"in batch {batch_idx_int}. Only {self.pages_available} pages available. "
+                f"Current capacity is {self.capacity[batch_idx_int]} tokens."
+            )
+
+        start_page_idx = self.capacity[batch_idx_int] // self.page_size
+        end_page_idx = start_page_idx + num_pages_to_allocate
+
+        # find empty physical pages
+        allocated_pages_list = self.free_pages[-num_pages_to_allocate:]
+        allocated_pages = torch.tensor(allocated_pages_list, device = self.device)
+        # update page table
+        self.page_table[batch_idx, start_page_idx:end_page_idx] = allocated_pages
+
+        # update metadata
+        self.physical_to_logical[batch_idx, allocated_pages] = torch.arange(
+            start_page_idx,
+            end_page_idx,
+            device = self.device,
+        )
+        # update cpu side metadata
+        self.page_table_cpu[batch_idx_int] += allocated_pages_list
+        self.free_pages = self.free_pages[:-num_pages_to_allocate]
+        self.capacity[batch_idx_int] += num_pages_to_allocate * self.page_size
+        return True
+
+    def erase(self, batch_idx: int) -> None:
+        """
+        Removes a single batch from paged attention.
+
+        Args:
+            batch_idx (int): batch index to be removed;
+        """
+        # NOTE: the GPU side data will only be reset/overwritten when we allocate it for a new batch
+        self.free_batch_idx.append(batch_idx)
+        allocated_pages_cpu = self.page_table_cpu[batch_idx]
+        self.free_pages.extend(reversed(allocated_pages_cpu))
+        self.page_table_cpu[batch_idx] = []
+
+    def assign(
+        self,
+        batch_idx: torch.Tensor,
+        input_pos: torch.Tensor,
+        k_val: torch.Tensor,
+        v_val: torch.Tensor,
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+    ) -> None:
+        """
+        Assigns new contents `val` to the storage `cache` at the location
+        `batch_idx` and `input_pos`.
+
+        Args:
+            batch_idx (Tensor): batch index; shape :math:`(B)`.
+            input_pos (Tensor): input positions to be assigned for the given batch; shape :math:`(B, S)`.
+            val (Tensor): value to be assigned; shape :math:`(B, H, S, D)`
+            cache (Tensor): the cache to store the values; shape:`(1, H, MAX_S, D)`
+        """
+        if k_val.requires_grad:
+            raise RuntimeError("val must not require gradient")
+
+        B, H, S, K_D = k_val.shape
+        _, H_cache, MAX_S, D_cache = k_cache.shape
+        assert H_cache == H, "number of heads must match"
+        assert MAX_S >= S, "cache must have enough space"
+        assert D_cache == K_D, "hidden dim must match"
+        assert input_pos.shape == (B, S), "input_pos must have the same shape as val"
+        assert batch_idx.shape == (B,), "batch_idx must have one dimension only"
+
+        V_D = v_val.shape[3]
+        if B != batch_idx.shape[0]:
+            raise RuntimeError(
+                f"Expect val and batch_idx have the same batch size but got B={B} and B={batch_idx.shape[0]}."
+            )
+        if H != k_cache.shape[1]:
+            raise RuntimeError(
+                f"Expect val and cache has the same number of heads but got H={H} and H={k_cache.shape[1]}."
+            )
+        if S != input_pos.shape[1]:
+            raise RuntimeError(
+                f"Expect val and input_pos has the same length but got S={S} and S={input_pos.shape[0]}."
+            )
+        if K_D != k_cache.shape[3]:
+            raise RuntimeError(
+                f"Expect k_val and k_cache has the same hidden dim but got D={K_D} and D={k_cache.shape[3]}."
+            )
+        if V_D != v_cache.shape[3]:
+            raise RuntimeError(
+                f"Expect v_val and v_cache has the same hidden dim but got D={V_D} and D={v_cache.shape[3]}."
+            )
+
+        # find address
+        logical_block_idx = input_pos // self.page_size  # [B, S]
+        logical_block_offset = input_pos % self.page_size  # [B, S]
+
+        # NOTE: this code path is only used for decoding. For batch prefill, use assign_prefill_no_paging() instead
+        physical_block_idx = torch.gather(
+            self.page_table[batch_idx], 1, logical_block_idx.to(torch.int64)
+        ).to(torch.int32)  # [B, S]
+
+        addr = (physical_block_idx * self.page_size + logical_block_offset).view(
+            -1
+        )  # [B*S]
+
+        k_val = k_val.permute(1, 0, 2, 3).contiguous().view(1, H, B * S, K_D)
+        v_val = v_val.permute(1, 0, 2, 3).contiguous().view(1, H, B * S, V_D)
+
+        k_cache[:, :, addr, :] = k_val
+        v_cache[:, :, addr, :] = v_val
+
+        return k_cache, v_cache
+
+    def convert_logical_block_mask(
+        self,
+        block_mask: BlockMask,
+        batch_idx: Optional[torch.Tensor] = None,
+    ) -> BlockMask:
+        """
+        Converts a logical block mask by mapping its logical kv indices to the corresponding
+        physical kv indices.
+
+        Args:
+            block_mask (BlockMask): logical block mask;
+                kv_indices shape :math:`(B, H, ROWS, MAX_BLOCKS_IN_COL)`.
+            batch_idx (Tensor): batch index corresponding to the block_mask
+                batch dimension. This provides flexibility to convert a
+                block mask with smaller batch size than the page table;
+                shape :math:`(B)`.
+        """
+        B, H, ROWS, MAX_BLOCKS_IN_COL = block_mask.kv_indices.shape
+
+        if block_mask.BLOCK_SIZE[1] != self.page_size:
+            raise RuntimeError(
+                f"Expect block_mask has the same column block size as page_sizebut got size={block_mask.BLOCK_SIZE[1]} and size={self.page_size}"
+            )
+
+        device = block_mask.kv_num_blocks.device
+
+        if batch_idx is None:
+            batch_idx = torch.arange(B, device = device)
+
+        assert batch_idx.ndim == 1, "batch_idx must be a 1D tensor"
+        assert (
+            batch_idx.shape[0] == B
+        ), "batch_idx must have the same shape as block_mask"
+        assert (
+            B <= self.max_batch_size
+        ), "batch_idx must be less than or equal to max_batch_size"
+
+        page_table = self.page_table[batch_idx]
+
+        def transform(num_blocks, indices):
+            """
+            transform the block mask from [B, H, num_q_blocks, num_logical_kv_blocks]
+            to [B, H, num_q_blocks, num_physical_kv_blocks]
+
+            kv_num_blocks: [B, H, num_q_blocks] -> unchanged
+            kv_indices: [B, H, num_q_blocks, num_logical_kv_blocks] -> [B, H, num_q_blocks, num_physical_kv_blocks]
+            """
+            if num_blocks is None:
+                return None, None
+            new_kv_num_blocks = num_blocks.clone()
+            new_kv_indices = torch.zeros(
+                (B, H, ROWS, self.n_pages), dtype = torch.int32, device = device
+            )
+            new_kv_indices[:, :, :, :MAX_BLOCKS_IN_COL] = (
+                torch.gather(page_table, 1, indices.view(B, -1).to(torch.int64))
+                .view(block_mask.kv_indices.shape)
+                .to(torch.int32)
+            )
+            return new_kv_num_blocks, new_kv_indices
+
+        new_kv_num_blocks, new_kv_indices = transform(
+            block_mask.kv_num_blocks, block_mask.kv_indices
+        )
+        new_full_kv_num_blocks, new_full_kv_indices = transform(
+            block_mask.full_kv_num_blocks, block_mask.full_kv_indices
+        )
+
+        new_mask_mod = self.get_mask_mod(block_mask.mask_mod, batch_idx)
+
+        seq_lengths = (block_mask.seq_lengths[0], self.n_pages * self.page_size)
+        return BlockMask.from_kv_blocks(
+            new_kv_num_blocks,
+            new_kv_indices,
+            new_full_kv_num_blocks,
+            new_full_kv_indices,
+            block_mask.BLOCK_SIZE,
+            new_mask_mod,
+            seq_lengths = seq_lengths,
+        )
+
+    def get_logical_kv_idx(
+        self,
+        physical_batch_idx: torch.Tensor,
+        physical_kv_idx: torch.Tensor,
+        batch_idx: torch.Tensor,
+    ):
+        logical_batch_idx = batch_idx[physical_batch_idx]
+        physical_kv_block = physical_kv_idx // self.page_size
+        physical_kv_offset = physical_kv_idx % self.page_size
+        logical_block_idx = self.physical_to_logical[
+            logical_batch_idx, physical_kv_block
+        ]
+        logical_kv_idx = logical_block_idx * self.page_size + physical_kv_offset
+        is_valid = logical_block_idx >= 0
+        safe_logical_kv_idx = logical_kv_idx.clamp(min = 0)
+        return is_valid, safe_logical_kv_idx
+
+    def get_mask_mod(
+        self, mask_mod: Optional[_mask_mod_signature], batch_idx: torch.Tensor
+    ) -> _mask_mod_signature:
+        """
+        Converts a mask_mod based on mapping from the physical block index to the logical
+        block index.
+
+        Args:
+            mask_mod (_mask_mod_signature): mask_mod based on the logical block index.
+        """
+        if mask_mod is None:
+            mask_mod = noop_mask
+
+        def new_mask_mod(
+            b: torch.Tensor,
+            h: torch.Tensor,
+            q_idx: torch.Tensor,
+            physical_kv_idx: torch.Tensor,
+        ):
+            is_valid, safe_logical_kv_idx = self.get_logical_kv_idx(
+                b, physical_kv_idx, batch_idx
+            )
+            return torch.where(
+                is_valid, mask_mod(b, h, q_idx, safe_logical_kv_idx), False
+            )
+
+        return new_mask_mod
+
+    # NOTE: not used in the current codebase
+    def get_score_mod(
+        self, score_mod: Optional[_score_mod_signature], batch_idx: torch.Tensor
+    ) -> _score_mod_signature:
+        """
+        Converts a score_mod based on mapping from the physical block index to the logical
+        block index.
+
+        Args:
+            score_mod (_score_mod_signature): score_mod based on the logical block index.
+        """
+        if score_mod is None:
+            score_mod = _identity
+
+        def new_score_mod(
+            score: torch.Tensor,
+            b: torch.Tensor,
+            h: torch.Tensor,
+            q_idx: torch.Tensor,
+            physical_kv_idx: torch.Tensor,
+        ):
+            is_valid, safe_logical_kv_idx = self.get_logical_kv_idx(
+                b, physical_kv_idx, batch_idx
+            )
+            return torch.where(
+                is_valid,
+                score_mod(score, b, h, q_idx, safe_logical_kv_idx),
+                float("-inf"),
+            )
+
+        return new_score_mod
+
+    def create_causal_blockmask(self, B, L):
+        """A minimal, unoptimized causal block mask creation function"""
+
+        def causal(b, h, q_idx, kv_idx):
+            return q_idx >= kv_idx
+
+        return create_block_mask(
+            causal,
+            B = B,
+            H = None,
+            Q_LEN = L,
+            KV_LEN = L,
+            BLOCK_SIZE = self.page_size,
+            device = self.device,
+        )
+
+    def create_prefill_blockmask_no_paging(
+        self, batch_idx: Tensor, BLOCK_SIZE: int = 128
+    ):
+        """
+        there's no prefix sharing implemented, batch_idx is the document id, batch_idx is not guaranteed to be sorted
+        """
+        assert batch_idx.ndim == 2, "batch_idx must be a 2D tensor"
+        assert batch_idx.shape[0] == 1, "batch_idx must have batch size 1"
+        L = batch_idx.shape[1]
+        docs = batch_idx.view(-1)
+
+        def document_causal(b, h, q_idx, kv_idx):
+            causal_mask = q_idx >= kv_idx
+            document_mask = docs[q_idx] == docs[kv_idx]
+            return causal_mask & document_mask
+
+        return create_block_mask(
+            document_causal, B = 1, H = None, Q_LEN = L, KV_LEN = L, BLOCK_SIZE = BLOCK_SIZE
+        )
+
+    # we assign prefill to the cache, similar to assign(), except we don't return the k_cache, v_cache, we only return the k_val, v_val
+    def assign_prefill_no_paging(
+        self,
+        batch_idx: torch.Tensor,
+        input_pos: torch.Tensor,
+        k_val: torch.Tensor,
+        v_val: torch.Tensor,
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+    ) -> None:
+        """
+        assigns kv and returns the original kv
+
+        batch_idx: [1, L]
+        input_pos: [1, L]
+        k_val: [1, H, L, D]
+        v_val: [1, H, L, D]
+        k_cache: [1, H, MAX_S, D]
+        v_cache: [1, H, MAX_S, D]
+        """
+
+        assert batch_idx.ndim == 2, "batch_idx must be a 2D tensor"
+        assert input_pos.ndim == 2, "input_pos must be a 2D tensor"
+        assert k_val.ndim == 4, "k_val must be a 4D tensor"
+        assert v_val.ndim == 4, "v_val must be a 4D tensor"
+        assert k_cache.ndim == 4, "k_cache must be a 4D tensor"
+        assert v_cache.ndim == 4, "v_cache must be a 4D tensor"
+        assert batch_idx.shape[0] == 1, "batch_idx must have batch size 1"
+
+        input_pos_block_idx = input_pos // self.page_size
+        input_pos_offset_in_block = input_pos % self.page_size
+        physical_kv_idx = (
+            self.page_table[batch_idx, input_pos_block_idx] * self.page_size
+            + input_pos_offset_in_block
+        )
+        k_cache[:, :, physical_kv_idx.view(-1), :] = k_val
+        v_cache[:, :, physical_kv_idx.view(-1), :] = v_val
+
+        return k_val, v_val

unsloth/inference/vllm_shim.py

@@ -0,0 +1,171 @@
+"""vLLM-API surface for the flex inference backend.
+
+`FlexEngine.generate` / `.chat` return :class:`RequestOutput` objects with the
+same attribute shape (`prompt_token_ids`, `outputs[i].token_ids`,
+`outputs[i].text`, `outputs[i].logprobs`) that vLLM's `LLM.generate` does, so
+TRL's GRPO trainer — which reads `output.prompt_token_ids` and iterates
+`output.outputs[i].token_ids` (trl/trainer/grpo_trainer.py:1274-1279) — does
+not care which backend produced the result.
+
+`save_lora` / `load_lora` mirror `unsloth_zoo.vllm_utils.save_lora` /
+`load_lora` (vllm_utils.py:2389-2628). The only backend-visible difference is
+the `LoRARequest` class: we emit our shim, not vllm's, so the GRPO patch in
+`unsloth/models/rl.py:1880` passes our object straight to `FlexEngine.generate`.
+"""
+
+from __future__ import annotations
+
+import functools
+import os
+from dataclasses import dataclass, field
+from typing import Any, Optional
+
+
+# ---------------------------------------------------------------------------
+# vLLM result objects
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class CompletionOutput:
+    """vLLM `CompletionOutput` stand-in.
+
+    TRL reads ``.token_ids`` (list of int) and ``.logprobs`` (optional list of
+    dict[int, Logprob]); we populate both."""
+
+    index: int = 0
+    text: str = ""
+    token_ids: list = field(default_factory = list)
+    cumulative_logprob: Optional[float] = None
+    logprobs: Optional[list] = None
+    finish_reason: Optional[str] = "stop"
+    stop_reason: Optional[str] = None
+
+
+@dataclass
+class RequestOutput:
+    """vLLM `RequestOutput` stand-in.
+
+    TRL reads ``.prompt_token_ids`` (list of int) and iterates
+    ``.outputs`` (list[CompletionOutput])."""
+
+    request_id: str = ""
+    prompt: str = ""
+    prompt_token_ids: list = field(default_factory = list)
+    outputs: list = field(default_factory = list)
+    finished: bool = True
+
+
+# ---------------------------------------------------------------------------
+# LoRARequest stand-in
+# ---------------------------------------------------------------------------
+#
+# The vLLM + unsloth_zoo combo imports ``vllm.lora.request.LoRARequest`` at
+# call time. When the flex backend is selected we never hit that import; we
+# pass our own dataclass that carries the same three attributes the caller
+# reads: ``lora_name``, ``lora_int_id``, and ``lora_tensors`` +
+# ``lora_config`` (for the in-memory LoRA fast path).
+
+
+@dataclass
+class LoRARequest:
+    lora_name: str = ""
+    lora_int_id: int = 0
+    lora_path: Optional[str] = None
+    lora_tensors: Optional[dict] = None  # dict[str, torch.Tensor]
+    lora_config: Any = None
+
+
+# ---------------------------------------------------------------------------
+# save_lora / load_lora — drop-in for unsloth_zoo.vllm_utils equivalents
+# ---------------------------------------------------------------------------
+
+_LORA_REQUEST_ID: Optional[int] = None
+
+
+def save_lora(model, save_directory, *args, **kwargs):
+    """Dump the PEFT LoRA tensors (``.lora_A.`` / ``.lora_B.``) into a
+    PEFT-compatible directory. Mirrors
+    ``unsloth_zoo.vllm_utils.save_lora`` (vllm_utils.py:2389-2397) byte-for-byte
+    so existing callers (e.g. TRL's GRPOTrainer patch) are untouched."""
+    state_dict = model.state_dict()
+    dtype = model.get_input_embeddings().weight.dtype
+    state_dict = {
+        k: v.to(dtype)
+        for k, v in state_dict.items()
+        if ".lora_A." in k or ".lora_B." in k
+    }
+    kwargs["state_dict"] = state_dict
+    model.save_pretrained(save_directory = save_directory, *args, **kwargs)
+
+
+def _get_peft_config(save_directory):
+    """Late-imported to keep the `peft` dep optional at module-load time."""
+    from peft import PeftConfig
+
+    return PeftConfig.from_pretrained(save_directory)
+
+
+def load_lora(
+    model,
+    save_directory,
+    load_tensors: bool = False,
+    lora_request_id: Optional[int] = None,
+):
+    """Build a :class:`LoRARequest` the flex backend can consume.
+
+    Mirrors ``unsloth_zoo.vllm_utils.load_lora`` (vllm_utils.py:2574-2628):
+    increments a module-level counter so each request gets a fresh
+    ``lora_int_id``, writes the PEFT adapter config to ``save_directory`` on
+    first call (or when ``load_tensors=True``) and captures the current
+    state-dict LoRA tensors so the engine can merge them without an extra
+    disk round-trip."""
+    global _LORA_REQUEST_ID
+    if _LORA_REQUEST_ID is None:
+        _LORA_REQUEST_ID = 1
+    if lora_request_id is None:
+        lora_request_id = _LORA_REQUEST_ID
+    if not os.path.exists(save_directory) or lora_request_id == 1:
+        if load_tensors:
+            model.peft_config["default"].save_pretrained(save_directory)
+        elif not os.path.exists(save_directory):
+            raise OSError(
+                f"Unsloth: LoRA filepath = {save_directory} does not exist!"
+            )
+
+    if load_tensors:
+        peft_config = _get_peft_config(save_directory)
+        state_dict = model.state_dict()
+        state_dict = {
+            k.replace(".default", ""): v
+            for k, v in state_dict.items()
+            if ".lora_A." in k or ".lora_B." in k
+        }
+        req = LoRARequest(
+            lora_name = str(lora_request_id),
+            lora_int_id = lora_request_id,
+            lora_tensors = state_dict,
+            lora_config = peft_config,
+        )
+    else:
+        req = LoRARequest(
+            lora_name = str(lora_request_id),
+            lora_int_id = lora_request_id,
+            lora_path = save_directory,
+        )
+    _LORA_REQUEST_ID += 1
+    return req
+
+
+# Partial-applied variants, mirroring ``patch_peft_fast_inference``:
+#   model.save_lora = functools.partial(save_lora, model)
+#   model.load_lora = functools.partial(load_lora, model)
+# are attached by the caller (unsloth/models/_utils.py:2707-2708).
+
+__all__ = [
+    "CompletionOutput",
+    "RequestOutput",
+    "LoRARequest",
+    "save_lora",
+    "load_lora",
+]

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

@@ -2488,7 +2488,18 @@ class FastLlamaModel:
                 offload_embedding = False,
                 fast_inference = fast_inference,
             )
-        elif not fast_inference:
+        elif not fast_inference or (
+            os.environ.get("UNSLOTH_FAST_INFERENCE", "0") == "1"
+        ):
+            # Two callers share this branch:
+            #   * the standard HF load (``fast_inference=False``)
+            #   * the flex-inference load (``fast_inference=True`` +
+            #     ``UNSLOTH_FAST_INFERENCE=1``) -- we load the HF model the
+            #     same way and then wrap it with a ``FlexEngine`` below.
+            # ``max_batch_size`` is a FlexEngine-specific kwarg, not an HF
+            # one; stash it before ``AutoModelForCausalLM.from_pretrained``
+            # rejects it as an unexpected argument.
+            _flex_max_batch_size = kwargs.pop("max_batch_size", 32)
             model = AutoModelForCausalLM.from_pretrained(
                 model_name,
                 device_map = device_map,
@@ -2508,10 +2519,32 @@ class FastLlamaModel:
                 load_in_4bit = load_in_4bit,
                 load_in_8bit = kwargs.get("load_in_8bit", False),
                 offload_embedding = False,
-                fast_inference = False,
+                fast_inference = fast_inference,
             )
-            model.fast_generate = make_fast_generate_wrapper(model.generate)
-            model.fast_generate_batches = None
+            if fast_inference and os.environ.get("UNSLOTH_FAST_INFERENCE", "0") == "1":
+                # Snapshot the HF model BEFORE Unsloth post-patching so the
+                # flex inference copy uses the original transformers
+                # rotary / QKV layout. A later block (after the tokenizer is
+                # loaded) constructs the FlexEngine around this copy.
+                import copy as _copy
+
+                model._unsloth_flex_inference_copy = _copy.deepcopy(model)
+                model._unsloth_flex_inference_copy.eval()
+                model._unsloth_needs_flex_engine = dict(
+                    dtype = dtype,
+                    max_seq_length = max_seq_length,
+                    max_lora_rank = max_lora_rank,
+                    max_batch_size = _flex_max_batch_size,
+                    gpu_memory_utilization = gpu_memory_utilization,
+                )
+                # Provide a placeholder so downstream code that checks
+                # hasattr(model, "fast_generate") doesn't explode; we swap
+                # it out once the engine is live.
+                model.fast_generate = make_fast_generate_wrapper(model.generate)
+                model.fast_generate_batches = None
+            else:
+                model.fast_generate = make_fast_generate_wrapper(model.generate)
+                model.fast_generate_batches = None
         else:
             from unsloth_zoo.vllm_utils import (
                 load_vllm,
@@ -2584,6 +2617,32 @@ class FastLlamaModel:
             model, tokenizer, correct_dtype = dtype
         )
 
+        # UNSLOTH_FAST_INFERENCE=1 path: build the FlexEngine now that the
+        # tokenizer is available. The HF model itself is NOT flex-patched --
+        # the engine owns its own pre-patched deep-copy (captured above
+        # before Unsloth's post_patch), so `model.forward` (used by the
+        # training loop) stays intact.
+        _flex_args = getattr(model, "_unsloth_needs_flex_engine", None)
+        if _flex_args is not None:
+            del model._unsloth_needs_flex_engine
+            _inference_copy = model._unsloth_flex_inference_copy
+            del model._unsloth_flex_inference_copy
+            from unsloth.inference.flex_engine import FlexEngine
+
+            flex_engine = FlexEngine(
+                hf_model = model,
+                tokenizer = tokenizer,
+                inference_model = _inference_copy,
+                base_model = None,
+                peft_model = None,
+                **_flex_args,
+            )
+            model.vllm_engine = flex_engine
+            model.fast_generate = flex_engine.generate
+            model.fast_generate_batches = functools.partial(
+                flex_engine.generate, use_tqdm = False
+            )
+
         # Patch up QKV / O and MLP
         for idx, layer in enumerate(model.model.layers):
             layer.self_attn.apply_qkv = original_apply_qkv
@@ -3290,6 +3349,13 @@ class FastLlamaModel:
 
         patch_peft_fast_inference(model)
 
+        # Hand the training-side PEFT wrapper to the flex engine so that
+        # state_dict() reads the current LoRA weights and the inference
+        # copy can mirror them.
+        _engine = getattr(model, "vllm_engine", None)
+        if _engine is not None and hasattr(_engine, "bind_peft_model"):
+            _engine.bind_peft_model(model)
+
         # Add for_inference and for_training
         model.for_training = functools.partial(FastLlamaModel.for_training, model)
         model.for_inference = functools.partial(FastLlamaModel.for_inference, model)

unsloth/models/loader.py

@@ -353,7 +353,15 @@ class FastLanguageModel(FastLlamaModel):
             or dtype == torch.float32
         )
 
-        if fast_inference:
+        _use_flex_fast_inference = (
+            os.environ.get("UNSLOTH_FAST_INFERENCE", "0") == "1"
+        )
+        if fast_inference and _use_flex_fast_inference:
+            # Flex backend path: skip the vLLM import gate entirely. The
+            # actual engine is attached further down in
+            # ``FastLlamaModel.from_pretrained``.
+            pass
+        elif fast_inference:
             if importlib.util.find_spec("vllm") is None:
                 raise ImportError(
                     "Unsloth: Please install vLLM before enabling `fast_inference`!\n"
@@ -979,7 +987,14 @@ class FastModel(FastBaseModel):
                 )
             load_in_4bit = False
 
-        if fast_inference:
+        _use_flex_fast_inference = (
+            os.environ.get("UNSLOTH_FAST_INFERENCE", "0") == "1"
+        )
+        if fast_inference and _use_flex_fast_inference:
+            # Flex backend path: skip the vLLM import gate. The engine is
+            # attached further down in the ``FastBaseModel`` loader.
+            pass
+        elif fast_inference:
             if importlib.util.find_spec("vllm") is None:
                 raise ImportError(
                     "Unsloth: Please install vLLM before enabling `fast_inference`!\n"

unsloth/models/vision.py

@@ -606,7 +606,16 @@ class FastBaseModel:
         vllm_enable_lora = True
 
         if is_vlm and fast_inference:
-            if not any(arch in VLLM_SUPPORTED_VLM for arch in model_types):
+            # The UNSLOTH_FAST_INFERENCE=1 flex backend ships with text-only
+            # support for Gemma-4-E2B-it today. Allow gemma4 through the
+            # vision-model path when the flex backend is selected; the
+            # vLLM-only compat list below still gates the default path.
+            _use_flex = os.environ.get("UNSLOTH_FAST_INFERENCE", "0") == "1"
+            _flex_allowed = {"gemma4"}
+            if not (
+                any(arch in VLLM_SUPPORTED_VLM for arch in model_types)
+                or (_use_flex and any(arch in _flex_allowed for arch in model_types))
+            ):
                 raise RuntimeError(
                     f"Unsloth: Fast inference is only supported for Language models and Qwen2.5-VL, Gemma3 among vision models. "
                     f"Found architectures: {', '.join(model_types)}!"
@@ -912,9 +921,15 @@ class FastBaseModel:
         verify_fp8_support_if_applicable(model_config)
 
         raise_handler = RaiseUninitialized()
-        if not fast_inference:
+        _use_flex_fast_inference = (
+            fast_inference and os.environ.get("UNSLOTH_FAST_INFERENCE", "0") == "1"
+        )
+        if (not fast_inference) or _use_flex_fast_inference:
+            # Shared by the standard HF load and the flex-inference load.
             # Prevent load_in_fp8 from being forwarded into HF internal model loading
             load_in_fp8 = kwargs.pop("load_in_fp8", None)
+            # ``max_batch_size`` is a FlexEngine kwarg, not an HF one.
+            _flex_max_batch_size = kwargs.pop("max_batch_size", 32)
             # Transformers 5.x @strict config classes reject unexpected kwargs.
             # Move config-level attributes onto the config object directly.
             _num_labels = kwargs.pop("num_labels", None)
@@ -944,6 +959,16 @@ class FastBaseModel:
                 fast_inference = fast_inference,
             )
             if hasattr(model, "generate"):
+                if _use_flex_fast_inference:
+                    # Defer engine construction until tokenizer/processor is
+                    # loaded; see the wiring block further down.
+                    model._unsloth_needs_flex_engine = dict(
+                        dtype = dtype,
+                        max_seq_length = max_seq_length,
+                        max_lora_rank = max_lora_rank,
+                        max_batch_size = _flex_max_batch_size,
+                        gpu_memory_utilization = gpu_memory_utilization,
+                    )
                 model.fast_generate = make_fast_generate_wrapper(model.generate)
                 model.fast_generate_batches = error_out_no_vllm
             if offload_embedding:
@@ -1223,6 +1248,22 @@ class FastBaseModel:
                 raise _patch_err
         model = post_patch_loss_function(model)
 
+        # UNSLOTH_FAST_INFERENCE=1 path: build the FlexEngine here so that
+        # the tokenizer/processor is available. Keeps the training model's
+        # forward intact — the engine carries its own deep-copy.
+        _flex_args = getattr(model, "_unsloth_needs_flex_engine", None)
+        if _flex_args is not None:
+            del model._unsloth_needs_flex_engine
+            from unsloth.inference.flex_engine import load_flex
+
+            _tok_for_flex = getattr(tokenizer, "tokenizer", tokenizer)
+            flex_engine = load_flex(model, _tok_for_flex, **_flex_args)
+            model.vllm_engine = flex_engine
+            model.fast_generate = flex_engine.generate
+            model.fast_generate_batches = functools.partial(
+                flex_engine.generate, use_tqdm = False
+            )
+
         # Log Unsloth version for future fastpaths for inference
         if hasattr(model, "config"):
             model.config.update({"unsloth_version": __version__})
@@ -1512,6 +1553,10 @@ class FastBaseModel:
         patch_saving_functions(model, vision = True)
         patch_peft_fast_inference(model)
 
+        _engine = getattr(model, "vllm_engine", None)
+        if _engine is not None and hasattr(_engine, "bind_peft_model"):
+            _engine.bind_peft_model(model)
+
         # Add for_inference and for_training
         model.for_training = functools.partial(FastBaseModel.for_training, model)
         model.for_inference = functools.partial(FastBaseModel.for_inference, model)