inference: cuMem sleep/wake + lazy FlexEngine batch sizing

Follow-up to 35231d4f (initial flex backend wiring).

Sleep/wake:
- sleep_mode.py adds kv_cache_pool / weight_pool context managers
  backed by vLLM's CuMemAllocator. Activated when UNSLOTH_VLLM_STANDBY=1
  and vLLM is importable; no-op otherwise so TRL's unconditional
  sleep / wake_up calls stay valid.
- FlexEngine routes the inference deep-copy + per-layer PagedKVCache
  through the pools. Captured CUDA graphs survive a sleep -> wake
  round-trip because cuMem keeps the GPU virtual addresses stable.
- 4-bit single-copy path drops only the KV cache; level 2 warns and
  falls back to level 1.
- tests/flex_sleep_mode_smoke.py covers sleep / wake memory deltas,
  captured-graph survival, and the no-op path.

Lazy batch sizing:
- FlexEngine's max_batch_size drives fixed-shape page tables, the
  input_pos_buffer, the block_mask_logical build, and the CUDA-graph
  bucket list at __init__ time. There is no post-init resize, so
  picking it at from_pretrained time forces over- or under-shoot.
- build_flex_engine() defers construction until the GRPO rollout shape
  is known. install_flex_sentinel() attaches a _LazyFlexEngineSentinel
  to model.vllm_engine so hasattr(model, "vllm_engine") keeps working
  between from_pretrained and the first build; fast_generate triggers
  a floor build on first call.
- rl.py injects _build_flex_from_args(model, args) before both
  self.llm = model.vllm_engine rewrite sites (pre-TRL-0.18
  sampling_params prefix and >=0.18 colocate LLM replacement). Sizes
  the engine from max(pdbs * spg, pdbs * spg * ngen) derived from the
  GRPO args. No-op on non-flex models, so the injection is safe for
  every TRL backend.
- Precedence: user's max_batch_size kwarg is a floor; the GRPO target
  overrides only when strictly larger, with a warning naming both.
- First-build only: post-build growth raises RuntimeError pointing the
  user back to max_batch_size= in from_pretrained. The pristine
  inference deep-copy is consumed on first build and gemma4 shell
  extraction mutates its module tree, so a safe rebuild would require
  a second deep-copy.
- tests/flex_lazy_batch_smoke.py (unit, stubbed FlexEngine) covers
  default, GRPO bump, user-floor, and post-build-refused cases.
  tests/flex_lazy_live_smoke.py exercises sentinel + build against a
  live Qwen3-0.6B-Base.

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

@@ -0,0 +1,299 @@
+# SPDX-License-Identifier: GNU Affero General Public License v3.0
+# Copyright 2023-present the Unsloth team. All rights reserved.
+
+"""Smoke-test :meth:`FlexEngine.sleep` and :meth:`FlexEngine.wake_up`.
+
+Invoked as:
+    CUDA_VISIBLE_DEVICES=3 \
+    UNSLOTH_FAST_INFERENCE=1 UNSLOTH_VLLM_STANDBY=1 \
+    python tests/flex_sleep_mode_smoke.py --model unsloth/Qwen3-4B-Base
+
+What it checks:
+  1. ``model.vllm_engine._sleep_mode_enabled`` is True when vLLM is
+     importable.
+  2. Captured CUDA graphs survive a sleep/wake round-trip: a second
+     ``fast_generate`` call after ``sleep`` -> ``wake_up`` does not
+     re-capture and still produces the same token ids.
+  3. ``torch.cuda.memory_allocated()`` drops on sleep and returns close
+     to the pre-sleep value on wake. With ``--no-standby``, memory
+     should be identical across the three probes (no-op path).
+  4. Hardens the regression guard by running the same workflow with
+     ``UNSLOTH_VLLM_STANDBY`` unset (``--no-standby``); ``sleep`` /
+     ``wake_up`` must be exact no-ops.
+
+This is a smoke test, not the full verification matrix from the plan
+(that lives under ``scripts/benchmarks``).
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import os
+import sys
+import time
+from pathlib import Path
+
+_REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(_REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(_REPO_ROOT))
+
+
+def _gb(n: int) -> float:
+    return round(n / 1e9, 3)
+
+
+def _probe(label: str) -> dict:
+    import torch
+
+    gc.collect()
+    torch.cuda.synchronize()
+    # ``memory_allocated`` / ``memory_reserved`` only track torch's
+    # caching allocator and ignore cuMem-backed pools, so they do NOT
+    # drop on sleep even though cuMem has unmapped the pages.
+    # ``mem_get_info()`` asks the CUDA runtime directly, so it sees
+    # cuMem unmaps and is the right probe for sleep / wake verification.
+    free_bytes, total_bytes = torch.cuda.mem_get_info()
+    return {
+        "label": label,
+        "allocated_gb": _gb(torch.cuda.memory_allocated()),
+        "reserved_gb": _gb(torch.cuda.memory_reserved()),
+        "cuda_free_gb": _gb(free_bytes),
+        "cuda_used_gb": _gb(total_bytes - free_bytes),
+    }
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--model", default = "unsloth/Qwen3-4B-Base")
+    p.add_argument("--dtype", choices = ["bf16", "fp16"], default = "bf16")
+    p.add_argument("--load_in_4bit", action = "store_true")
+    p.add_argument("--max_new_tokens", type = int, default = 32)
+    p.add_argument("--max_seq_length", type = int, default = 1024)
+    p.add_argument("--prompt", default = "The quick brown fox jumps over")
+    p.add_argument(
+        "--no-standby",
+        action = "store_true",
+        help = "Force UNSLOTH_VLLM_STANDBY=0 to validate the no-op regression path.",
+    )
+    p.add_argument(
+        "--cycles",
+        type = int,
+        default = 1,
+        help = "Number of sleep / wake / generate cycles after warmup. "
+        ">1 exercises the repeated-cycle regression (run #6).",
+    )
+    args = p.parse_args()
+
+    os.environ.setdefault("UNSLOTH_FAST_INFERENCE", "1")
+    if args.no_standby:
+        os.environ["UNSLOTH_VLLM_STANDBY"] = "0"
+    else:
+        os.environ.setdefault("UNSLOTH_VLLM_STANDBY", "1")
+    standby = os.environ.get("UNSLOTH_VLLM_STANDBY", "0") == "1"
+    print(
+        f"[sleep-smoke] UNSLOTH_FAST_INFERENCE="
+        f"{os.environ.get('UNSLOTH_FAST_INFERENCE')} "
+        f"UNSLOTH_VLLM_STANDBY={os.environ.get('UNSLOTH_VLLM_STANDBY')}"
+    )
+
+    import torch
+
+    import unsloth
+    from unsloth import FastLanguageModel
+
+    print(f"[sleep-smoke] unsloth={unsloth.__file__}")
+
+    dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
+
+    t0 = time.perf_counter()
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name = args.model,
+        max_seq_length = args.max_seq_length,
+        dtype = dtype,
+        load_in_4bit = args.load_in_4bit,
+        fast_inference = True,
+    )
+    print(
+        f"[sleep-smoke] loaded {args.model} in "
+        f"{time.perf_counter() - t0:.1f}s; dtype={model.dtype}"
+    )
+
+    engine = model.vllm_engine
+    print(f"[sleep-smoke] engine type: {type(engine).__name__}")
+    sleep_enabled = getattr(engine, "_sleep_mode_enabled", None)
+    print(f"[sleep-smoke] engine._sleep_mode_enabled: {sleep_enabled}")
+    if standby:
+        if sleep_enabled is not True:
+            # Most likely vLLM is not importable in this environment.
+            print(
+                "[sleep-smoke] WARNING: UNSLOTH_VLLM_STANDBY=1 was set "
+                "but engine._sleep_mode_enabled is False (vLLM missing?)"
+            )
+    else:
+        assert sleep_enabled is False, (
+            f"Expected sleep mode to be disabled with UNSLOTH_VLLM_STANDBY=0, "
+            f"got {sleep_enabled}"
+        )
+
+    ll_cfg = engine.llm_engine.vllm_config.model_config
+    print(
+        f"[sleep-smoke] llm_engine.vllm_config.model_config.enable_sleep_mode: "
+        f"{getattr(ll_cfg, 'enable_sleep_mode', None)}"
+    )
+    assert getattr(ll_cfg, "enable_sleep_mode", None) == bool(sleep_enabled), (
+        "_LLMEngineStub.model_config.enable_sleep_mode must mirror "
+        "engine._sleep_mode_enabled"
+    )
+
+    from unsloth.inference.vllm_shim import LoRARequest  # noqa: F401
+
+    prompts = [args.prompt]
+
+    # ----- warmup -----
+    t0 = time.perf_counter()
+    out1 = engine.generate(
+        prompts,
+        sampling_params = type(
+            "SP", (), {"max_tokens": args.max_new_tokens, "temperature": 0.0},
+        )(),
+    )
+    print(
+        f"[sleep-smoke] warmup generate: {time.perf_counter() - t0:.2f}s; "
+        f"tok_ids[:10]={out1[0].outputs[0].token_ids[:10]}"
+    )
+    pre_tokens = list(out1[0].outputs[0].token_ids)
+
+    probe_pre = _probe("pre-sleep")
+    print(f"[sleep-smoke] {probe_pre}")
+
+    # Diagnostic: checksum the inference-model weights so we can detect
+    # if cuMem's sleep/wake round-trip corrupts any parameter.
+    def _checksum_params(mod, limit = 16):
+        import torch as _t
+
+        out = []
+        for i, (name, p) in enumerate(mod.named_parameters()):
+            if i >= limit:
+                break
+            t = p.detach()
+            out.append((name, list(t.shape), float(t.float().abs().sum().item())))
+        return out
+
+    pre_sums = _checksum_params(engine._inference_model)
+    print("[sleep-smoke] pre-sleep first-16 param |sum|:")
+    for n, s, v in pre_sums:
+        print(f"  {n} {s} {v:.4f}")
+
+    # Repeated sleep / wake / generate cycle test (plan matrix run #6).
+    # Each cycle validates that the engine does not drift: tokens remain
+    # bitwise identical, memory returns to baseline, weights round-trip
+    # cleanly. A bug that only surfaces on the second or third cycle
+    # (stale Python state, double-wake, leaked handles) fails here.
+    for cycle in range(args.cycles):
+        if args.cycles > 1:
+            print(f"[sleep-smoke] --- cycle {cycle + 1}/{args.cycles} ---")
+
+        # ----- sleep -----
+        t0 = time.perf_counter()
+        engine.sleep(level = 1)
+        t_sleep = time.perf_counter() - t0
+        probe_post = _probe(f"post-sleep[{cycle + 1}]")
+        print(
+            f"[sleep-smoke] sleep(level=1) took {t_sleep:.3f}s; {probe_post}"
+        )
+
+        if sleep_enabled:
+            drop = probe_pre["cuda_used_gb"] - probe_post["cuda_used_gb"]
+            print(
+                f"[sleep-smoke] process-level VRAM drop on sleep: "
+                f"{drop:+.3f} GB (cuMem-managed; not visible in "
+                f"torch.memory_allocated)"
+            )
+            if not getattr(engine, "_single_copy_mode", False):
+                # 16-bit path: both weights + kv_cache pools are dropped.
+                assert drop >= 1.0, (
+                    f"Expected multi-GB drop in process-level VRAM on "
+                    f"sleep(level=1); got {drop:+.3f} GB"
+                )
+            else:
+                # 4-bit single-copy: only KV cache drops; weights stay.
+                assert drop > 0.0, (
+                    f"Expected KV-cache drop in process-level VRAM on "
+                    f"sleep(level=1); got {drop:+.3f} GB"
+                )
+        else:
+            # With sleep mode off, the sleep() call must not free VRAM.
+            # Process-level jitter is allowed (shared GPU); torch-owned
+            # allocations must be untouched.
+            assert (
+                probe_post["allocated_gb"] == probe_pre["allocated_gb"]
+            ), (
+                "With sleep mode disabled, torch.memory_allocated must "
+                "be unchanged by sleep()"
+            )
+
+        # ----- wake -----
+        t0 = time.perf_counter()
+        engine.wake_up()
+        t_wake = time.perf_counter() - t0
+        probe_wake = _probe(f"post-wake[{cycle + 1}]")
+        print(
+            f"[sleep-smoke] wake_up() took {t_wake:.3f}s; {probe_wake}"
+        )
+
+        post_sums = _checksum_params(engine._inference_model)
+        diffs = []
+        for (n1, s1, v1), (n2, s2, v2) in zip(pre_sums, post_sums):
+            delta = abs(v1 - v2)
+            if delta > 0.0:
+                diffs.append((n1, v1, v2, delta))
+        print(
+            f"[sleep-smoke] post-wake weight diff: "
+            f"{len(diffs)}/{len(pre_sums)} params changed "
+            f"(bitwise-exact restore expected)"
+        )
+        for n, v1, v2, d in diffs[:8]:
+            print(f"  diff {n}: pre={v1:.4f} post={v2:.4f} delta={d:.4f}")
+        assert len(diffs) == 0, (
+            f"Weight corruption on sleep / wake (cycle {cycle + 1}): "
+            f"{len(diffs)}/{len(pre_sums)} first-layer params changed"
+        )
+
+        # ----- verify we can still generate -----
+        t0 = time.perf_counter()
+        out2 = engine.generate(
+            prompts,
+            sampling_params = type(
+                "SP", (), {"max_tokens": args.max_new_tokens, "temperature": 0.0},
+            )(),
+        )
+        t_regen = time.perf_counter() - t0
+        post_tokens = list(out2[0].outputs[0].token_ids)
+        match = pre_tokens == post_tokens
+        print(
+            f"[sleep-smoke] post-wake generate: {t_regen:.2f}s; "
+            f"tok_ids[:10]={post_tokens[:10]}; matches_pre={match}"
+        )
+        assert match, (
+            f"Pre-sleep / post-wake token ids must match exactly "
+            f"(cycle {cycle + 1}).\n"
+            f"Pre:  {pre_tokens}\nPost: {post_tokens}"
+        )
+
+        if sleep_enabled:
+            delta = probe_wake["cuda_used_gb"] - probe_pre["cuda_used_gb"]
+            print(
+                f"[sleep-smoke] post-wake vs pre-sleep process-level "
+                f"VRAM delta: {delta:+.3f} GB (tolerance: +/- 1.5 GB)"
+            )
+            assert abs(delta) < 1.5, (
+                f"Post-wake VRAM diverged from pre-sleep "
+                f"(cycle {cycle + 1}): delta={delta:+.3f} GB"
+            )
+
+    print(f"[sleep-smoke] PASS ({args.cycles} cycle(s))")
+
+
+if __name__ == "__main__":
+    main()

unsloth/inference/__init__.py

@@ -14,7 +14,12 @@ 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 .flex_engine import (
+    FlexEngine,
+    build_flex_engine,
+    install_flex_sentinel,
+    load_flex,
+)
 from .vllm_shim import (
     CompletionOutput,
     LoRARequest,
@@ -26,6 +31,8 @@ from .vllm_shim import (
 __all__ = [
     "FlexEngine",
     "load_flex",
+    "build_flex_engine",
+    "install_flex_sentinel",
     "LoRARequest",
     "RequestOutput",
     "CompletionOutput",

unsloth/inference/flex_engine.py

@@ -32,6 +32,7 @@ from __future__ import annotations
 
 import copy
 import functools
+import gc
 import importlib.util
 import os
 import types
@@ -49,6 +50,12 @@ from .flex_qwen3_llama import (
     refresh_lora_merge_from_pristine,
 )
 from .flex_gemma4 import FlexGemma4Inference
+from .sleep_mode import (
+    _get_cumem_allocator,
+    kv_cache_pool,
+    sleep_mode_enabled,
+    weight_pool,
+)
 from .vllm_shim import CompletionOutput, LoRARequest, RequestOutput
 
 
@@ -242,8 +249,13 @@ class _LLMEngineStub:
     ``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())
+    def __init__(self, sleep_enabled: bool = False):
+        self.vllm_config = types.SimpleNamespace(
+            lora_config = types.SimpleNamespace(),
+            model_config = types.SimpleNamespace(
+                enable_sleep_mode = bool(sleep_enabled),
+            ),
+        )
         self.model_executor = types.SimpleNamespace(
             driver_worker = types.SimpleNamespace(
                 model_runner = types.SimpleNamespace(
@@ -309,6 +321,22 @@ class FlexEngine:
 
         self.device = hf_model.device
 
+        # Sleep-mode setup. When ``UNSLOTH_VLLM_STANDBY=1`` is set AND
+        # vLLM is importable, we route the engine's heavy allocations
+        # (the inference deep-copies + per-layer PagedKVCache buffers)
+        # through cuMem-backed pools so ``FlexEngine.sleep`` can offload
+        # weights to pinned CPU memory without destroying captured CUDA
+        # graphs. The allocator assigns stable GPU virtual addresses, so
+        # unmapping on sleep and re-mapping on wake preserves pointer
+        # validity. ``expandable_segments:True`` on
+        # ``PYTORCH_CUDA_ALLOC_CONF`` is incompatible with cuMem; if the
+        # user has it set, ``_get_cumem_allocator`` returns None and
+        # sleep stays a no-op.
+        self._sleep_mode_enabled = sleep_mode_enabled()
+        self._cumem_allocator = (
+            _get_cumem_allocator() if self._sleep_mode_enabled else None
+        )
+
         # Colocate pattern (mirrors vLLM's "colocate" mode): the engine
         # runs on its own deep-copy of the HF model so that flex attention
         # patching + KV-cache attachment does not mutate the training
@@ -325,11 +353,18 @@ class FlexEngine:
         # 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()
+            with weight_pool(self._cumem_allocator):
+                inference_model = copy.deepcopy(hf_model)
+                inference_model.eval()
         self._pristine_base = base_model  # None until bind_peft_model runs
         self._inference_model = inference_model
         self._inference_peft = peft_model  # filled in by bind_peft_model
+        # ``_inference_model is hf_model`` is the 4-bit / single-copy
+        # fallback path (see ``bind_peft_model``); flex skips the second
+        # deep-copy there because bnb-4bit packed weights can't be
+        # in-place refreshed. In that mode there is no CPU-backup step
+        # to do on ``sleep`` — only the KV cache gets dropped.
+        self._single_copy_mode = inference_model is hf_model
 
         # Autocast wrapping (applied by the impl via a context manager).
         fa4_prefill, prefill_kernel_options, decode_kernel_options = (
@@ -353,6 +388,14 @@ class FlexEngine:
             inference_model = _extract_gemma4_text_shell(inference_model)
             self._inference_model = inference_model
         Impl = FlexGemma4Inference if arch == "gemma4" else FlexInference
+        # Pass the cuMem allocator through so the impl can wrap ONLY
+        # the paged-KV allocations (``PageTable`` + per-layer
+        # ``PagedKVCache``) in the ``kv_cache`` pool. Everything else
+        # the impl creates (``input_pos_buffer``, ``block_mask_logical``,
+        # captured CUDA-graph scratch / graph_vars) stays in torch's
+        # default allocator — those buffers are tiny and, critically,
+        # the captured CUDA graphs reference block_mask indices by
+        # address, so they must survive sleep / wake unchanged.
         self._impl = Impl(
             inference_model,
             tokenizer,
@@ -366,8 +409,11 @@ class FlexEngine:
             fa4_prefill = fa4_prefill,
             base_model = self._pristine_base,
             peft_model = peft_model,
+            cumem_allocator = self._cumem_allocator,
+        )
+        self._llm_engine_stub = _LLMEngineStub(
+            sleep_enabled = self._sleep_mode_enabled,
         )
-        self._llm_engine_stub = _LLMEngineStub()
 
     # ----- configuration helpers -----
 
@@ -605,20 +651,74 @@ class FlexEngine:
             self._warned_sampling = True
         return int(max_tokens), {}
 
-    # ----- sleep-mode stubs -----
+    # ----- sleep-mode -----
     #
     # 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.
+    # training can use the freed VRAM. The flex backend implements level 1
+    # (weights offloaded to pinned CPU, KV cache dropped and re-zeroed on
+    # wake) via :class:`vllm.device_allocator.cumem.CuMemAllocator`.
+    # Captured CUDA graphs survive the round-trip because cuMem keeps the
+    # GPU virtual addresses stable across sleep / wake.
+    #
+    # Sleep activates only when ``UNSLOTH_VLLM_STANDBY=1`` is set AND
+    # vLLM is importable (evaluated at ``__init__`` time). Otherwise
+    # ``sleep`` / ``wake_up`` are no-ops so code that unconditionally
+    # calls the API (TRL's GRPO trainer) stays correct.
 
-    def sleep(self, level: int = 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)."""
+    def sleep(self, level: int = 1):
+        """Offload inference weights to pinned CPU memory (level 1).
+
+        ``level=2`` is not implemented on the flex backend (it would
+        require rebuilding the inference deep-copy from the training
+        model on wake); requesting it emits a warning and falls back to
+        level 1.
+
+        In the 4-bit single-copy fallback path the inference model
+        shares storage with the training model, so only the KV cache is
+        dropped on sleep; the weights stay resident.
+        """
+        if not self._sleep_mode_enabled or self._cumem_allocator is None:
+            return None
+        if level not in (1, 2):
+            raise ValueError(
+                f"FlexEngine.sleep: level must be 1 or 2, got {level}"
+            )
+        if level == 2:
+            warnings.warn(
+                "FlexEngine.sleep(level=2) is not implemented on the "
+                "flex backend; falling back to level=1 (CPU-pinned "
+                "weight offload).",
+                RuntimeWarning,
+                stacklevel = 2,
+            )
+        if self._single_copy_mode:
+            # Weights are shared with the training model (4-bit path);
+            # only the kv_cache pool is ours to drop.
+            self._cumem_allocator.sleep(offload_tags = ())
+        else:
+            self._cumem_allocator.sleep(offload_tags = ("weights",))
+        gc.collect()
+        # NOTE: we deliberately do NOT call torch.cuda.empty_cache() here.
+        # Captured CUDA graphs may retain scratch / workspace tensors in
+        # torch's default caching allocator at fixed addresses; emptying
+        # the cache between sleep and wake can invalidate those
+        # addresses and cause the next graph replay to read freed
+        # memory. The cuMem pools have already released their physical
+        # pages; there is no additional VRAM to reclaim via empty_cache.
         return None
 
     def wake_up(self, tags: Optional[list] = None):
-        """No-op companion to :meth:`sleep`."""
+        """Re-map cuMem handles and restore offloaded weights.
+
+        ``tags=None`` wakes everything; TRL's GRPO trainer calls
+        ``wake_up(tags=["kv_cache"])`` and then ``wake_up(tags=["weights"])``
+        on consecutive steps to stagger the VRAM reclaim.
+        """
+        if not self._sleep_mode_enabled or self._cumem_allocator is None:
+            return None
+        if tags is not None and not isinstance(tags, list):
+            tags = list(tags)
+        self._cumem_allocator.wake_up(tags = tags)
         return None
 
     @property
@@ -654,8 +754,9 @@ class FlexEngine:
             # 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()
+            with weight_pool(self._cumem_allocator):
+                self._pristine_base = copy.deepcopy(self._inference_model)
+                self._pristine_base.eval()
             self._impl.base_model = self._pristine_base
 
         if self._inference_peft is None:
@@ -666,8 +767,11 @@ class FlexEngine:
                 # 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()
+                with weight_pool(self._cumem_allocator):
+                    self._inference_peft = _get_peft_model(
+                        self._inference_model, peft_cfg,
+                    )
+                    self._inference_peft.eval()
             except Exception as e:
                 warnings.warn(
                     f"FlexEngine.bind_peft_model: could not build an "
@@ -724,4 +828,211 @@ def load_flex(
     )
 
 
-__all__ = ["FlexEngine", "load_flex"]
+# ---------------------------------------------------------------------------
+# Lazy engine construction
+#
+# FlexEngine's ``max_batch_size`` drives fixed-shape GPU page tables, the
+# ``input_pos_buffer``, the ``block_mask_logical`` build, and the CUDA-graph
+# bucket list. These are allocated inside ``FlexEngine.__init__`` and there is
+# no post-init resize path. Picking the wrong value at ``from_pretrained``
+# time forces users to either overshoot (wasted KV pages + slower graph
+# capture) or undershoot (``PageTable.can_reserve`` stalls the rollout).
+#
+# ``build_flex_engine`` defers the construction until the real rollout batch
+# size is known: GRPOTrainer's ``__init__`` patch in ``unsloth/models/rl.py``
+# calls :func:`_build_flex_from_args` to pass
+# ``per_device_train_batch_size * steps_per_generation * num_generations``
+# through; the plain ``model.fast_generate`` path falls back to the
+# ``max_batch_size`` kwarg originally passed to ``from_pretrained``.
+# ---------------------------------------------------------------------------
+
+
+class _LazyFlexEngineSentinel:
+    """Placeholder for ``model.vllm_engine`` before the FlexEngine is built.
+
+    Forwards attribute access to the real engine, triggering construction
+    (with the stashed ``max_batch_size`` floor) on first access. This keeps
+    ``hasattr(model, "vllm_engine")`` True between ``from_pretrained`` and
+    the first build, which matters for ``rl.py``'s ``args.use_vllm`` setter.
+    """
+
+    __slots__ = ("_model",)
+
+    def __init__(self, model):
+        object.__setattr__(self, "_model", model)
+
+    def _resolve(self):
+        engine = getattr(self._model, "_flex_engine_instance", None)
+        if engine is None:
+            engine = build_flex_engine(self._model)
+        return engine
+
+    def __getattr__(self, name):
+        if name == "_model":
+            raise AttributeError(name)
+        return getattr(self._resolve(), name)
+
+    def __bool__(self):
+        return True
+
+    def __repr__(self):
+        engine = getattr(self._model, "_flex_engine_instance", None)
+        if engine is None:
+            return "<LazyFlexEngine (not built)>"
+        return repr(engine)
+
+
+def install_flex_sentinel(model, tokenizer):
+    """Wire the lazy ``vllm_engine`` / ``fast_generate`` placeholders.
+
+    Called from ``unsloth/models/llama.py`` and ``unsloth/models/vision.py``
+    in place of the eager ``FlexEngine(...)`` construction. The real build
+    happens inside :func:`build_flex_engine`, triggered either by the
+    GRPOTrainer patch (``_build_flex_from_args``) or by the first
+    ``model.fast_generate`` call.
+    """
+    model._unsloth_flex_tokenizer = tokenizer
+    model.vllm_engine = _LazyFlexEngineSentinel(model)
+
+    def _lazy_fast_generate(prompts = None, *gen_args, **gen_kwargs):
+        engine = getattr(model, "_flex_engine_instance", None)
+        if engine is None:
+            engine = build_flex_engine(model)
+        return engine.generate(prompts, *gen_args, **gen_kwargs)
+
+    def _lazy_fast_generate_batches(prompts = None, *gen_args, **gen_kwargs):
+        engine = getattr(model, "_flex_engine_instance", None)
+        if engine is None:
+            engine = build_flex_engine(model)
+        gen_kwargs.setdefault("use_tqdm", False)
+        return engine.generate(prompts, *gen_args, **gen_kwargs)
+
+    model.fast_generate = _lazy_fast_generate
+    model.fast_generate_batches = _lazy_fast_generate_batches
+
+
+def _construct_and_attach(model, max_batch_size: int):
+    """Construct the FlexEngine with the given batch size and wire it up."""
+    pending = model._unsloth_needs_flex_engine
+    tokenizer = getattr(model, "_unsloth_flex_tokenizer", None)
+    inference_copy = getattr(model, "_unsloth_flex_inference_copy", None)
+
+    engine_kwargs = dict(pending)
+    engine_kwargs["max_batch_size"] = int(max_batch_size)
+
+    engine = FlexEngine(
+        hf_model = model,
+        tokenizer = tokenizer,
+        inference_model = inference_copy,
+        base_model = None,
+        peft_model = None,
+        **engine_kwargs,
+    )
+    model._flex_engine_instance = engine
+
+    # Drop the sentinel (plain attribute; setting replaces it).
+    model.vllm_engine = engine
+    model.fast_generate = engine.generate
+    model.fast_generate_batches = functools.partial(engine.generate, use_tqdm = False)
+
+    # Consume the one-shot stashes: the deep-copy is owned by the engine now,
+    # and the needs-dict's role as "build spec" is over.
+    for _attr in (
+        "_unsloth_flex_inference_copy",
+        "_unsloth_flex_tokenizer",
+        "_unsloth_needs_flex_engine",
+    ):
+        if hasattr(model, _attr):
+            try:
+                delattr(model, _attr)
+            except AttributeError:
+                pass
+    return engine
+
+
+def build_flex_engine(model, max_batch_size: Optional[int] = None):
+    """Construct or return the FlexEngine attached to ``model``.
+
+    Called lazily: by the RL patch (:func:`_build_flex_from_args`) once
+    ``GRPOTrainer.args`` is resolved, or by ``model.fast_generate`` on
+    first use.
+
+    ``max_batch_size`` resolution (first build):
+      - ``floor = model._unsloth_needs_flex_engine['max_batch_size']``
+      - ``effective = max(floor, max_batch_size or 0)``
+      - A warning is emitted when ``effective > floor`` so users see the
+        GRPO-driven bump.
+
+    Once the engine is built, it is the sole source of truth for the
+    batch-size dimension. Subsequent calls are idempotent when the
+    requested size fits; requesting a larger size raises
+    :class:`RuntimeError` because the engine's fixed-shape GPU buffers
+    and captured CUDA graphs cannot be grown in place.
+    """
+    existing = getattr(model, "_flex_engine_instance", None)
+    pending = getattr(model, "_unsloth_needs_flex_engine", None)
+
+    # Non-flex model (plain HF or plain vLLM). No-op so ``rl.py``'s patch
+    # stays unconditional.
+    if existing is None and pending is None:
+        return None
+
+    requested = int(max_batch_size) if max_batch_size else 0
+
+    if existing is not None:
+        if requested <= existing.max_batch_size:
+            return existing
+        raise RuntimeError(
+            f"Unsloth: FlexEngine was built at max_batch_size="
+            f"{existing.max_batch_size}; cannot grow to {requested} after "
+            f"construction (fixed-shape GPU page tables + CUDA graphs). "
+            f"Pass max_batch_size={requested} to "
+            f"FastLanguageModel.from_pretrained before the first "
+            f"fast_generate / GRPOTrainer call."
+        )
+
+    floor = int(pending["max_batch_size"])
+    target = max(floor, requested)
+    if target > floor:
+        warnings.warn(
+            f"Unsloth: increasing FlexEngine max_batch_size {floor} -> "
+            f"{target} to fit the GRPO rollout batch "
+            f"(per_device_train_batch_size * steps_per_generation * "
+            f"num_generations). Pass max_batch_size={target} to "
+            f"FastLanguageModel.from_pretrained to silence this warning.",
+            stacklevel = 2,
+        )
+    return _construct_and_attach(model, target)
+
+
+def _build_flex_from_args(model, args):
+    """Helper used by the ``rl.py`` GRPOTrainer-init patch.
+
+    Reads the rollout batch size from the TRL args and triggers the
+    FlexEngine build. No-op when ``model`` wasn't loaded through the
+    flex-inference path (plain vLLM / plain HF).
+    """
+    if not hasattr(model, "_unsloth_needs_flex_engine") and not hasattr(
+        model, "_flex_engine_instance"
+    ):
+        return None
+    pdbs = int(getattr(args, "per_device_train_batch_size", 1) or 1)
+    spg = int(
+        getattr(args, "steps_per_generation", None)
+        or getattr(args, "gradient_accumulation_steps", 1)
+        or 1
+    )
+    ngen = int(getattr(args, "num_generations", 1) or 1)
+    # Written as ``max(A, B)`` for reviewer clarity. Reduces to the second
+    # term whenever ``num_generations >= 1`` (always).
+    grpo_target = max(pdbs * spg, pdbs * spg * ngen)
+    return build_flex_engine(model, max_batch_size = grpo_target)
+
+
+__all__ = [
+    "FlexEngine",
+    "load_flex",
+    "build_flex_engine",
+    "install_flex_sentinel",
+    "_build_flex_from_args",
+]

unsloth/inference/flex_gemma4.py

@@ -430,6 +430,7 @@ class FlexGemma4Inference:
         fa4_prefill = None,
         base_model = None,
         peft_model = None,
+        cumem_allocator = None,
     ):
         assert max_seq_length % page_size == 0
         self.model = model
@@ -477,14 +478,19 @@ class FlexGemma4Inference:
             base_prefill_opts["BACKEND"] = "FLASH"
         self.prefill_kernel_options = base_prefill_opts
 
-        self.page_table = PageTable(
-            n_pages = n_pages,
-            page_size = page_size,
-            max_batch_size = max_batch_size,
-            device = self.device.type,
-        )
+        # See ``FlexInference.__init__`` for why only the paged-KV
+        # allocations go through the cuMem ``kv_cache`` pool.
+        from .sleep_mode import kv_cache_pool as _kv_cache_pool
 
-        patch_gemma4_attention_forwards(model, self.page_table)
+        with _kv_cache_pool(cumem_allocator):
+            self.page_table = PageTable(
+                n_pages = n_pages,
+                page_size = page_size,
+                max_batch_size = max_batch_size,
+                device = self.device.type,
+            )
+
+            patch_gemma4_attention_forwards(model, self.page_table)
 
         self.input_pos_buffer = torch.zeros(
             max_batch_size, dtype = torch.int32, device = self.device

unsloth/inference/flex_qwen3_llama.py

@@ -499,6 +499,7 @@ class FlexInference:
         fa4_prefill = None,
         base_model = None,
         peft_model = None,
+        cumem_allocator = None,
     ):
         assert max_seq_length % page_size == 0
         self.model = model
@@ -558,13 +559,22 @@ class FlexInference:
             base_prefill_opts["BACKEND"] = "FLASH"
         self.prefill_kernel_options = base_prefill_opts
 
-        self.page_table = PageTable(
-            n_pages = n_pages,
-            page_size = page_size,
-            max_batch_size = max_batch_size,
-            device = self.device.type,
-        )
-        patch_model_attention_forwards(model, self.page_table)
+        # Route the paged-KV allocations through the cuMem "kv_cache"
+        # pool when sleep mode is active, so FlexEngine.sleep can drop
+        # them and wake_up can re-map fresh zeroed pages at the same
+        # virtual addresses. The block_mask / input_pos scratch below
+        # stays in the default allocator so captured CUDA graphs that
+        # reference them stay valid across sleep / wake.
+        from .sleep_mode import kv_cache_pool as _kv_cache_pool
+
+        with _kv_cache_pool(cumem_allocator):
+            self.page_table = PageTable(
+                n_pages = n_pages,
+                page_size = page_size,
+                max_batch_size = max_batch_size,
+                device = self.device.type,
+            )
+            patch_model_attention_forwards(model, self.page_table)
 
         # Pre-allocated decode state.
         self.input_pos_buffer = torch.zeros(

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/models/llama.py

@@ -2526,10 +2526,28 @@ class FastLlamaModel:
                 # 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,
+                )
 
-                model._unsloth_flex_inference_copy = _copy.deepcopy(model)
-                model._unsloth_flex_inference_copy.eval()
+                _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,
@@ -2617,31 +2635,19 @@ 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 --
+        # 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.
-        _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
+        if hasattr(model, "_unsloth_needs_flex_engine"):
+            from unsloth.inference.flex_engine import install_flex_sentinel
 
-            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
-            )
+            install_flex_sentinel(model, tokenizer)
 
         # Patch up QKV / O and MLP
         for idx, layer in enumerate(model.model.layers):

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

@@ -1248,21 +1248,17 @@ 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
+        # 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.
-        _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
+        if hasattr(model, "_unsloth_needs_flex_engine"):
+            from unsloth.inference.flex_engine import install_flex_sentinel
 
             _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
-            )
+            install_flex_sentinel(model, _tok_for_flex)
 
         # Log Unsloth version for future fastpaths for inference
         if hasattr(model, "config"):