[Bugfix] Fix sleep-mode tensor ownership and recovery

[Ronald1995]

Purpose

Addresses the ownership and recovery design in #53343.

CuMem allocator tags describe the allocation site, but sleep mode treats them as lifecycle policy. A persistent runtime tensor allocated while the kv_cache or weights pool is active can therefore be discarded and remapped without its semantic value, while its virtual address remains stable. This can produce silent wrong results, including through an already captured CUDA graph.

This PR:

• narrows the kv_cache pool in MRV1 and MRV2 to the backing KV-cache allocation only;
• keeps builders, bind/alias setup, connector initialization, and other runtime setup outside that pool;
• registers model and kernel runtime state with an owning module where appropriate;
• restores or resets derived MoE state, Humming locks, and MiniMax Lamport workspace state after weight reload/wake while preserving captured storage addresses;
• adds focused allocation-boundary tests and deterministic poisoned-remap reproductions for ordinary tensors and CUDA graphs.

Related and duplicate work

Searched all vLLM issues and PRs for sleep mode tensor ownership and sleep mode kv cache tensor tag before submission.

• [Bugfix] Narrow kv_cache mempool context to prevent sleep mode regressions #35956 proposed only the KV-cache scope narrowing, but was closed unmerged after conflicts. This PR rebases that boundary fix on current main and adds recovery coverage.
• [Bugfix][Attention] Preserve post-load tensors across weight reloads #48251 is merged and fixes specific MLA and FlashInfer post-load tensors; this PR is complementary and covers other audited owners.
• [Bugfix] CuMemAllocator: zero discarded-tag pages on wake_up #45542 proposed zeroing discarded pages and was closed unmerged. Zeroing makes bytes deterministic but cannot restore tensor semantics.

No open PR implementing this combined scope and recovery design was found.

Root-cause categories and potentially affected models

The affected surface is determined by both model architecture and the selected runtime backend. In particular, the MoE fixes below are not Gemma3n-specific: Gemma3n exercises model-owned static scaling tensors, while the MoE changes cover quantization backends shared by many sparse models.

Root cause	Failure during sleep/wake	Potentially affected models or paths	Main trigger
The kv_cache allocation context covered initialization work beyond the backing cache allocation	Persistent metadata or helper tensors can inherit the KV-cache tag and be discarded without semantic recovery	Any attention model with complex KV initialization; higher-risk examples include Gemma3n KV sharing, hybrid/Mamba/GDN or MLA caches, FlashInfer attention, and KV connector/disaggregated-serving setup	Primarily level-1 sleep, because KV-tagged allocations are discarded rather than CPU-backed
Model-owned constants were created as plain tensor attributes while the weights pool was active	Level-2 wake reloads checkpoint parameters, but cannot reconstruct non-parameter tensor attributes	Confirmed audited paths include Gemma3n scaling tensors, Voxtral mel filters, ERNIE 4.5 VL visual-token caches, OpenPangu attention-sink values, and relative positional encodings used by Conformer-based encoders	Level-2 sleep
MoE weight post-processing created derived scales, strides, or GEMM control tensors without registering them with an owning module	The checkpoint reload restores expert weights but not these derived runtime values	Backend-dependent examples include DeepSeek-V3/R1, Kimi-K2/K3, Qwen3-MoE, Mixtral, Llama 4, GPT-OSS, and Nemotron/ModelOpt MoE checkpoints	Level-2 sleep only when the selected backend is one of the affected FlashInfer CUTLASS/B12x, FlashInfer TRT-LLM FP8/MXFP4, CUTLASS W4A8, or Humming paths
A derived buffer was recreated during reload while a helper or captured graph still referenced its old storage	Python attributes can point at new tensors while expert helpers or CUDA graphs continue reading the old address	The same backend-selected MoE families above, especially configurations using CUDA Graph capture	Weight reload/wake; storage must be restored in place and helper references rebound
Synchronization/workspace tensors require protocol reset rather than byte-for-byte restoration	Restoring a stale lock, flag, or workspace state can leave the next invocation in an invalid communication state	Humming linear/MoE kernels; MiniMax Text01/M2 tensor-parallel Lamport RMS-QK path	Level-1 or level-2 wake, when the corresponding optimized kernel is enabled

Backend-dependent MoE exposure

A model name alone is insufficient to determine exposure. The same Qwen3-MoE or DeepSeek checkpoint can use Triton, CUTLASS, FlashInfer TRT-LLM, B12x, or Humming depending on its quantization format, GPU architecture, expert-parallel configuration, installed kernels, and --moe-backend.

The concrete backend classes audited by this PR include:

• CutlassExpertsW4A8Fp8: W4A8 MoE checkpoints using INT4 weights and FP8 activations.
• FlashInferExperts: FlashInfer CUTLASS paths used by compatible unquantized, FP8, MXFP4, or NVFP4 MoE checkpoints.
• TrtLlmFp8Experts family: compatible FP8 MoE checkpoints, including possible DeepSeek, Kimi, Qwen3-MoE, and Mixtral deployments.
• TrtLlmMxfp4Experts family: MXFP4 deployments such as GPT-OSS, compatible Qwen3-MoE checkpoints, and the NVIDIA Kimi-K3 MXFP4 path.
• FlashInferB12xExperts: compatible NVFP4 MoE checkpoints, including Qwen3-MoE and ModelOpt/Nemotron-style variants when this backend is selected.
• HummingMoEMethod: compatible compressed-tensors/ModelOpt MoE checkpoints when Humming is installed and selected.

These are conditional examples, not claims that every checkpoint in each family is affected. Dense variants, deployments without sleep mode, and MoE models selecting an unaffected backend do not exercise the corresponding MoE recovery path.

Sleep-level distinction

• Level 1: weight-tagged allocations are CPU-backed and restored as raw bytes, while KV-tagged allocations are discarded. Accidental KV ownership is therefore the primary concern; stateful locks/workspaces may still require an explicit reset.
• Level 2: both KV and weight pools are discarded, and model weights are subsequently reloaded. Plain model constants and derived MoE tensors that are absent from the checkpoint are the primary concern.

Test Plan

pytest -q tests/v1/worker/test_kv_cache_allocation_scope.py
pytest -q tests/model_executor/test_sleep_mode_tensor_ownership.py
pytest -q tests/basic_correctness/test_mem.py -k "tagged_ordinary_tensor or level2_discards or cudagraph_replays"
pytest -q tests/kernels/core/test_minimax_reduce_rms.py -k lamport_workspace_reset

Test Result

• uvx ruff check passed for all changed Python files.
• python -m compileall passed for all changed vLLM Python files.
• Three focused CPU ownership checks passed through the local vLLM environment.
• Full pytest collection was not available on the Windows development host because the vLLM test configuration requires Linux-only uvloop.
• CUDA poisoned-remap and CUDA Graph tests require a supported Linux CUDA/ROCm vLLM runtime and were not executed locally; they are included for CI/GPU validation.

e2e case for gemma3n

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Natural-remap E2E reproducer for Gemma3n level-2 sleep corruption."""

import os

import pytest
import torch

from tests.utils import create_new_process_for_each_test
from vllm import LLM, SamplingParams
from vllm.inputs import TokensPrompt
from vllm.platforms import current_platform

SLEEP_WAKE_CYCLES = int(os.getenv("VLLM_SLEEP_WAKE_REPRO_CYCLES", "3"))

# These tensors are created during model construction and participate in
# Gemma3n text forward passes.
SLEEP_TENSOR_NAMES = (
    "embed_scale",
    "embed_scale_per_layer",
    "per_layer_input_scale",
    "per_layer_projection_scale",
)


def _create_tiny_gemma3n_config(model_dir) -> None:
    """Create a small local Gemma3n config without downloading a model."""
    from transformers import Gemma3nTextConfig

    config = Gemma3nTextConfig(
        architectures=["Gemma3nForCausalLM"],
        vocab_size=128,
        hidden_size=64,
        intermediate_size=[128, 128, 128],
        num_hidden_layers=3,
        num_attention_heads=4,
        num_key_value_heads=2,
        head_dim=32,
        max_position_embeddings=256,
        sliding_window=64,
        layer_types=[
            "sliding_attention",
            "full_attention",
            "sliding_attention",
        ],
        vocab_size_per_layer_input=128,
        hidden_size_per_layer_input=16,
        num_kv_shared_layers=1,
        laurel_rank=16,
        activation_sparsity_pattern=[0.0, 0.0, 0.0],
    )
    config.save_pretrained(model_dir)


def _gemma3n_runtime_tensor_snapshot(
    model,
) -> dict[str, tuple[int, float]]:
    """Return addresses and values of sleep-sensitive runtime tensors."""
    from vllm.model_executor.models.gemma3n import Gemma3nSelfDecoder

    decoder = next(
        (
            module
            for module in model.modules()
            if isinstance(module, Gemma3nSelfDecoder)
        ),
        None,
    )
    assert decoder is not None

    return {
        name: (tensor.data_ptr(), tensor.float().item())
        for name in SLEEP_TENSOR_NAMES
        if (tensor := getattr(decoder, name, None)) is not None
    }


def _save_dummy_weights(worker) -> None:
    """Save the initial random parameters in the worker process."""
    model = worker.model_runner.model
    worker._gemma3n_dummy_weights = [
        (name, parameter.detach().cpu().clone())
        for name, parameter in model.named_parameters()
    ]


def _reload_dummy_weights(worker) -> None:
    """Restore exactly the same random parameters after level-2 sleep."""
    with torch.no_grad():
        worker.reload_weights(
            weights_iterator=iter(worker._gemma3n_dummy_weights),
            is_checkpoint_format=False,
        )


def _generation_signature(
    outputs,
) -> tuple[tuple[int, ...], tuple[float, ...]]:
    """Return generated IDs and selected-token log probabilities."""
    completion = outputs[0].outputs[0]
    token_ids = tuple(completion.token_ids)

    assert completion.logprobs is not None
    assert len(completion.logprobs) == len(token_ids)

    chosen_logprobs = tuple(
        step_logprobs[token_id].logprob
        for token_id, step_logprobs in zip(
            token_ids,
            completion.logprobs,
            strict=True,
        )
    )
    return token_ids, chosen_logprobs


@create_new_process_for_each_test()
@pytest.mark.slow_test
@pytest.mark.skipif(
    not current_platform.is_cuda(),
    reason="Reproduces the CUDA CuMemAllocator level-2 sleep path",
)
def test_gemma3n_level2_sleep_wake_preserves_generation(
    monkeypatch,
    tmp_path,
) -> None:
    """Detect natural Gemma3n runtime-tensor corruption after level-2 sleep."""
    # Needed for collective_rpc(callable).
    monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")

    # Keep the worker in the test process and avoid unnecessary multiprocess
    # setup for this single-GPU reproducer.
    monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")

    _create_tiny_gemma3n_config(tmp_path)

    llm = LLM(
        model=str(tmp_path),
        load_format="dummy",
        skip_tokenizer_init=True,
        enable_sleep_mode=True,
        enforce_eager=True,
        attention_backend="TRITON_ATTN",
        max_model_len=256,
        max_num_seqs=1,
        seed=0,
        enable_prefix_caching=False,
        disable_log_stats=True,
        # Optional: limit KV-cache reservation if supported by the branch.
        # kv_cache_memory_bytes=64 * 1024 * 1024,
    )

    # No tokenizer is initialized, so provide token IDs directly.
    prompt = TokensPrompt(
        prompt_token_ids=[2, 5, 9, 12, 7, 3],
    )
    sampling_params = SamplingParams(
        temperature=0.0,
        max_tokens=16,
        logprobs=1,
    )

    # Establish the baseline output.
    baseline_output = llm.generate(prompt, sampling_params)
    baseline_token_ids, baseline_logprobs = _generation_signature(
        baseline_output
    )

    # Verify that repeated awake inference is deterministic before involving
    # sleep mode.
    control_output = llm.generate(prompt, sampling_params)
    control_token_ids, control_logprobs = _generation_signature(
        control_output
    )

    assert control_token_ids == baseline_token_ids
    assert control_logprobs == pytest.approx(
        baseline_logprobs,
        rel=1e-5,
        abs=1e-5,
    )

    baseline_tensors = llm.apply_model(
        _gemma3n_runtime_tensor_snapshot
    )
    assert all(
        set(snapshot) == set(SLEEP_TENSOR_NAMES)
        for snapshot in baseline_tensors
    )

    # DummyModelLoader cannot reload weights from disk. Preserve the randomly
    # initialized parameters once and restore exactly the same values after
    # every level-2 sleep.
    llm.collective_rpc(_save_dummy_weights)

    for cycle in range(1, SLEEP_WAKE_CYCLES + 1):
        # No cudaMemset, allocator monkeypatch, memory pressure allocation, or
        # other test-side GPU write occurs between sleep and wake.
        llm.sleep(level=2)

        llm.wake_up(tags=["weights"])
        llm.collective_rpc(_reload_dummy_weights)

        after_tensors = llm.apply_model(
            _gemma3n_runtime_tensor_snapshot
        )

        # Check corruption before another forward: invalid runtime scales may
        # make generate raise and hide the direct ownership failure.
        assert after_tensors == baseline_tensors, (
            f"cycle {cycle}: Gemma3n runtime tensors changed after "
            f"sleep/wake; before={baseline_tensors}, "
            f"after={after_tensors}"
        )

        llm.wake_up(tags=["kv_cache"])

        after_output = llm.generate(prompt, sampling_params)
        after_token_ids, after_logprobs = _generation_signature(
            after_output
        )

        assert after_token_ids == baseline_token_ids, (
            f"cycle {cycle}: generated tokens changed after sleep/wake; "
            f"expected={baseline_token_ids}, actual={after_token_ids}"
        )

        assert after_logprobs == pytest.approx(
            baseline_logprobs,
            rel=1e-5,
            abs=1e-5,
        ), (
            f"cycle {cycle}: selected-token logprobs changed "
            "after sleep/wake"
        )

AI assistance: OpenAI Codex was used to investigate, implement, and test this change. The submitter reviewed the design, code, tests, and commits and takes responsibility for the contribution.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

Documentation preview: https://vllm--53344.org.readthedocs.build/en/53344/

[Ronald1995]

test result for gemma3n

before:

>                       raise original_exception
E                       AssertionError: cycle 1: Gemma3n runtime tensors changed after sleep/wake; before=[{'embed_scale': (140380362670080, 8.0), 'embed_scale_per_layer': (140380362682880, 4.0), 'per_layer_input_scale': (140380362690560, 0.70703125), 'per_layer_projection_scale': (140380362690048, 8.0)}], after=[{'embed_scale': (140380362670080, 0.0), 'embed_scale_per_layer': (140380362682880, 0.0), 'per_layer_input_scale': (140380362690560, 0.0), 'per_layer_projection_scale': (140380362690048, 0.0)}]
E                       assert [{'embed_scal...690048, 0.0)}] == [{'embed_scal...690048, 8.0)}]
E
E                         At index 0 diff: {'embed_scale': (140380362670080, 0.0), 'embed_scale_per_layer': (140380362682880, 0.0), 'per_layer_input_scale': (140380362690560, 0.0), 'per_layer_projection_scale': (140380362690048, 0.0)} != {'
E
E                         ...Full output truncated (34 lines hidden), use '-vv' to show

after

PASSED

=============================================================================================== warnings summary ================================================================================================
.venv/lib/python3.12/site-packages/torch/jit/_script.py:365: 14 warnings
  /home/liurong/vllm/.venv/lib/python3.12/site-packages/torch/jit/_script.py:365: DeprecationWarning: `torch.jit.script_method` is deprecated. Please switch to `torch.compile` or `torch.export`.
    warnings.warn(

test2.py::test_gemma3n_level2_sleep_wake_preserves_generation
  /home/liurong/vllm/tests/utils.py:1759: DeprecationWarning: This process (pid=2078323) is multi-threaded, use of fork() may lead to deadlocks in the child.
    pid = os.fork()

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
================================================================================== 1 passed, 15 warnings in 118.44s (0:01:58) ===================================================================================

[aoshen02]

No need to fix mrv1, let's focus on mrv2 first.

[Ronald1995]

ok， i'll focus on mrv2.

[aoshen02]

let's just split this huge pr into several pieces for example

1. sleep scope: only kvcache
2. register buffer
3. ...

[Ronald1995]

ok，i will split it into several prs

[Ronald1995]

Superseded by the focused implementation series tracked in #53343: #53508, #53507, #53509, #53510, and #53511. The dependency graph and recommended merge order are documented in the RFC.