[Bugfix][Models] Register sleep-managed runtime buffers (#53507)
Signed-off-by: Ronald1995 <ronaldautomobile@163.com> Signed-off-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: aoshen02 <aoshen@inferact.ai>
Ronald committed
Aug 29, 2026 at 10:32 UTC
16d6c376bc55e2e44ce012d26104c8af11e27c25
6 files changed
+427
-19
tests/basic_correctness/test_mem.py
+38
@@ -18,6 +18,22 @@ from ..utils import create_new_process_for_each_test, requires_fp8
18
DEVICE_TYPE = current_platform.device_type
19
20
21
+def _wake_up_with_poisoned_mappings(allocator, byte_value: int = 0xA5) -> None:
22
+ """Wake discarded allocations with deterministic nonzero contents."""
23
+ original_create_and_map = cumem.create_and_map
24
+
25
+ def create_and_map_with_poison(handle) -> None:
26
+ original_create_and_map(handle)
27
+ _, size, ptr, _ = handle
28
+ cumem.libcudart.cudaMemset(ptr, byte_value, size)
29
+
30
+ cumem.create_and_map = create_and_map_with_poison
31
+ try:
32
+ allocator.wake_up()
33
+ finally:
34
+ cumem.create_and_map = original_create_and_map
35
+
36
+
37
@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
38
def test_python_error():
39
"""
@@ -110,6 +126,28 @@ def test_discard_tags():
126
assert torch.allclose(weights, torch.ones_like(weights))
127
128
129
+@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
130
+@pytest.mark.skipif(current_platform.is_xpu(), reason="Uses the CuMem allocator")
131
+def test_level2_discards_ordinary_tensor_with_weights_tag():
132
+ """Reproduce the level-2 variant for an ordinary tensor in weights."""
133
+ allocator = get_mem_allocator_instance()
134
+
135
+ with allocator.use_memory_pool("weights"):
136
+ fake_weight = torch.full((4096,), 0x44, dtype=torch.uint8, device=DEVICE_TYPE)
137
+ ordinary_tensor = torch.full(
138
+ (4096,), 0x55, dtype=torch.uint8, device=DEVICE_TYPE
139
+ )
140
+
141
+ pointers = (fake_weight.data_ptr(), ordinary_tensor.data_ptr())
142
+ allocator.sleep(offload_tags=())
143
+ _wake_up_with_poisoned_mappings(allocator)
144
+ torch.accelerator.synchronize()
145
+
146
+ assert (fake_weight.data_ptr(), ordinary_tensor.data_ptr()) == pointers
147
+ assert torch.all(fake_weight == 0xA5)
148
+ assert torch.all(ordinary_tensor == 0xA5)
149
+
150
+
151
@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")
152
@pytest.mark.skipif(current_platform.is_xpu(), reason="CUDA graph not supported on XPU")
153
def test_cumem_with_cudagraph():
tests/model_executor/test_sleep_mode_tensor_ownership.py
new
+339
@@ -0,0 +1,339 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from dataclasses import dataclass
5
+from functools import partial
6
+from pathlib import Path
7
+
8
+import pytest
9
+import torch
10
+from transformers import Gemma3nTextConfig
11
+
12
+from vllm import LLM, SamplingParams
13
+from vllm.inputs import TokensPrompt
14
+from vllm.platforms import current_platform
15
+
16
+from ..models.utils import dummy_hf_overrides
17
+from ..utils import create_new_process_for_each_test
18
+
19
+
20
+@dataclass(frozen=True)
21
+class SleepModelCase:
22
+ name: str
23
+ architecture: str
24
+ model: str | None
25
+ tensor_names: tuple[str, ...]
26
+ modalities: tuple[str, ...] = ()
27
+ revision: str | None = None
28
+ trust_remote_code: bool = False
29
+ tokenizer_mode: str | None = None
30
+
31
+
32
+SLEEP_MODEL_CASES = (
33
+ SleepModelCase(
34
+ name="fireredasr2",
35
+ architecture="FireRedASR2ForConditionalGeneration",
36
+ model="allendou/FireRedASR2-LLM-vllm",
37
+ tensor_names=("pe",),
38
+ trust_remote_code=True,
39
+ ),
40
+ SleepModelCase(
41
+ name="gemma3n",
42
+ architecture="Gemma3nForCausalLM",
43
+ model=None,
44
+ tensor_names=(
45
+ "router_input_scale",
46
+ "embed_scale",
47
+ "embed_scale_per_layer",
48
+ "per_layer_input_scale",
49
+ "per_layer_projection_scale",
50
+ ),
51
+ ),
52
+ SleepModelCase(
53
+ name="voxtral",
54
+ architecture="VoxtralForConditionalGeneration",
55
+ model="mistralai/Voxtral-Mini-3B-2507",
56
+ tensor_names=("mel_filters",),
57
+ revision="3060fe34b35ba5d44202ce9ff3c097642914f8f3",
58
+ tokenizer_mode="mistral",
59
+ ),
60
+ SleepModelCase(
61
+ name="ernie45_vl",
62
+ architecture="Ernie4_5_VLMoeForConditionalGeneration",
63
+ model="baidu/ERNIE-4.5-VL-28B-A3B-PT",
64
+ tensor_names=("inv_freq", "_visual_token_ids_tensor_cache"),
65
+ revision="refs/pr/17",
66
+ trust_remote_code=True,
67
+ ),
68
+)
69
+
70
+
71
+def _create_local_model_config(case: SleepModelCase, model_dir: Path) -> str:
72
+ if case.name == "gemma3n":
73
+ config = Gemma3nTextConfig(
74
+ architectures=[case.architecture],
75
+ vocab_size=128,
76
+ hidden_size=64,
77
+ intermediate_size=[128, 128, 128],
78
+ num_hidden_layers=3,
79
+ num_attention_heads=2,
80
+ num_key_value_heads=2,
81
+ head_dim=32,
82
+ max_position_embeddings=256,
83
+ sliding_window=64,
84
+ layer_types=[
85
+ "sliding_attention",
86
+ "full_attention",
87
+ "sliding_attention",
88
+ ],
89
+ vocab_size_per_layer_input=128,
90
+ hidden_size_per_layer_input=16,
91
+ num_kv_shared_layers=1,
92
+ laurel_rank=16,
93
+ activation_sparsity_pattern=[0.0, 0.0, 0.0],
94
+ )
95
+ else:
96
+ raise AssertionError(f"No local config builder for {case.name}")
97
+
98
+ model_dir.mkdir()
99
+ config.save_pretrained(model_dir)
100
+ return str(model_dir)
101
+
102
+
103
+def _sleep_test_hf_overrides(hf_config, *, model_arch: str):
104
+ hf_config = dummy_hf_overrides(hf_config, model_arch=model_arch)
105
+ if model_arch == "FireRedASR2ForConditionalGeneration":
106
+ audio_encoder_conf = getattr(hf_config, "audio_encoder_conf", {})
107
+ audio_encoder_conf.update(
108
+ {
109
+ "idim": 80,
110
+ "n_layers_enc": 1,
111
+ "n_head": 1,
112
+ "d_model": 32,
113
+ "kernel_size": 3,
114
+ # MM profiling uses the repository's real feature extractor
115
+ # and can produce roughly 750 encoder frames. Keep the real
116
+ # positional-encoding capacity instead of shrinking it with
117
+ # the number of test decoder tokens.
118
+ "pe_maxlen": 5000,
119
+ }
120
+ )
121
+ hf_config.audio_encoder_conf = audio_encoder_conf
122
+ if (
123
+ model_arch == "Ernie4_5_VLMoeForConditionalGeneration"
124
+ and getattr(hf_config, "im_patch_id", None) is None
125
+ ):
126
+ # Some revisions expose this id as image_token_id, while the native
127
+ # model intentionally caches it under the older im_patch_id name.
128
+ hf_config.im_patch_id = getattr(hf_config, "image_token_id", 100295)
129
+ return hf_config
130
+
131
+
132
+def _get_sleep_tensor_targets(
133
+ model_name: str,
134
+) -> tuple[tuple[type[torch.nn.Module], tuple[str, ...]], ...]:
135
+ if model_name == "fireredasr2":
136
+ from vllm.model_executor.models.conformer_encoder import (
137
+ RelPositionalEncoding,
138
+ )
139
+
140
+ return ((RelPositionalEncoding, ("pe",)),)
141
+ if model_name == "gemma3n":
142
+ from vllm.model_executor.models.gemma3n import (
143
+ Gemma3nAltUp,
144
+ Gemma3nSelfDecoder,
145
+ )
146
+
147
+ return (
148
+ (Gemma3nAltUp, ("router_input_scale",)),
149
+ (
150
+ Gemma3nSelfDecoder,
151
+ (
152
+ "embed_scale",
153
+ "embed_scale_per_layer",
154
+ "per_layer_input_scale",
155
+ "per_layer_projection_scale",
156
+ ),
157
+ ),
158
+ )
159
+ if model_name == "voxtral":
160
+ from vllm.model_executor.models.voxtral import VoxtralEncoderModel
161
+
162
+ return ((VoxtralEncoderModel, ("mel_filters",)),)
163
+ if model_name == "ernie45_vl":
164
+ from vllm.model_executor.models.ernie45_vl import (
165
+ Ernie4_5_VisionRotaryEmbedding,
166
+ Ernie4_5_VLMoeForConditionalGeneration,
167
+ )
168
+
169
+ return (
170
+ (Ernie4_5_VisionRotaryEmbedding, ("inv_freq",)),
171
+ (
172
+ Ernie4_5_VLMoeForConditionalGeneration,
173
+ ("_visual_token_ids_tensor_cache",),
174
+ ),
175
+ )
176
+ raise AssertionError(f"No sleep tensor targets for {model_name}")
177
+
178
+
179
+def _snapshot_sleep_tensors(
180
+ model: torch.nn.Module,
181
+ *,
182
+ model_name: str,
183
+ tensor_names: tuple[str, ...],
184
+) -> dict[str, tuple[int, torch.Tensor]]:
185
+ targets = _get_sleep_tensor_targets(model_name)
186
+ snapshot: dict[str, tuple[int, torch.Tensor]] = {}
187
+ found_tensor_names = set[str]()
188
+
189
+ for module_path, module in model.named_modules():
190
+ for module_type, target_names in targets:
191
+ if not isinstance(module, module_type):
192
+ continue
193
+ for tensor_name in target_names:
194
+ tensor = getattr(module, tensor_name, None)
195
+ if tensor is None:
196
+ continue
197
+ assert isinstance(tensor, torch.Tensor)
198
+ qualified_name = (
199
+ f"{module_path}.{tensor_name}" if module_path else tensor_name
200
+ )
201
+ snapshot[qualified_name] = (
202
+ tensor.data_ptr(),
203
+ tensor.detach().cpu().clone(),
204
+ )
205
+ found_tensor_names.add(tensor_name)
206
+
207
+ assert found_tensor_names == set(tensor_names), (
208
+ f"{model_name}: expected tensors {tensor_names}, "
209
+ f"found {tuple(sorted(found_tensor_names))}"
210
+ )
211
+ return snapshot
212
+
213
+
214
+def _save_dummy_weights(worker) -> None:
215
+ model = worker.model_runner.model
216
+ worker._sleep_test_dummy_weights = [
217
+ (name, parameter.detach().cpu().clone())
218
+ for name, parameter in model.named_parameters()
219
+ ]
220
+
221
+
222
+def _reload_dummy_weights(worker) -> None:
223
+ with torch.no_grad():
224
+ worker.reload_weights(
225
+ weights_iterator=iter(worker._sleep_test_dummy_weights),
226
+ is_checkpoint_format=False,
227
+ )
228
+
229
+
230
+def _generation_signature(output) -> tuple[tuple[int, ...], tuple[float, ...]]:
231
+ completion = output[0].outputs[0]
232
+ token_ids = tuple(completion.token_ids)
233
+ assert completion.logprobs is not None
234
+ chosen_logprobs = tuple(
235
+ step[token_id].logprob
236
+ for token_id, step in zip(token_ids, completion.logprobs, strict=True)
237
+ )
238
+ return token_ids, chosen_logprobs
239
+
240
+
241
+def _assert_tensor_snapshots_equal(
242
+ before: dict[str, tuple[int, torch.Tensor]],
243
+ after: dict[str, tuple[int, torch.Tensor]],
244
+ model_name: str,
245
+) -> None:
246
+ assert after.keys() == before.keys(), model_name
247
+ for name, (before_ptr, before_value) in before.items():
248
+ after_ptr, after_value = after[name]
249
+ assert after_ptr == before_ptr, f"{model_name}: {name} address changed"
250
+ torch.testing.assert_close(
251
+ after_value,
252
+ before_value,
253
+ msg=lambda msg, name=name: (f"{model_name}: {name} was corrupted: {msg}"),
254
+ )
255
+
256
+
257
+@create_new_process_for_each_test("spawn")
258
+@pytest.mark.slow_test
259
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA sleep mode test")
260
+@pytest.mark.parametrize("case", SLEEP_MODEL_CASES, ids=lambda case: case.name)
261
+def test_static_model_tensors_survive_level2_restore(
262
+ monkeypatch: pytest.MonkeyPatch,
263
+ tmp_path: Path,
264
+ case: SleepModelCase,
265
+) -> None:
266
+ monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
267
+ monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
268
+
269
+ if case.name == "fireredasr2":
270
+ # FireRedASR2 has a regular decoder and compute_logits implementation,
271
+ # but advertises only the transcription task. Enable its text path so
272
+ # this ownership test can exercise the common LLM.generate lifecycle
273
+ # without constructing an unrelated audio fixture.
274
+ from vllm.model_executor.models.fireredasr2 import (
275
+ FireRedASR2ForConditionalGeneration,
276
+ )
277
+
278
+ monkeypatch.setattr(
279
+ FireRedASR2ForConditionalGeneration,
280
+ "supports_transcription_only",
281
+ False,
282
+ )
283
+
284
+ model = case.model or _create_local_model_config(case, tmp_path / case.name)
285
+ llm_kwargs = {
286
+ "model": model,
287
+ "load_format": "dummy",
288
+ "enable_sleep_mode": True,
289
+ "enforce_eager": True,
290
+ "attention_backend": "TRITON_ATTN",
291
+ "max_model_len": 128,
292
+ "max_num_seqs": 1,
293
+ "gpu_memory_utilization": 0.5,
294
+ "seed": 0,
295
+ "enable_prefix_caching": False,
296
+ "disable_log_stats": True,
297
+ "trust_remote_code": case.trust_remote_code,
298
+ }
299
+ if case.revision is not None:
300
+ llm_kwargs["revision"] = case.revision
301
+ if case.name not in ("fireredasr2", "voxtral", "ernie45_vl"):
302
+ llm_kwargs["skip_tokenizer_init"] = True
303
+ if case.tokenizer_mode is not None:
304
+ llm_kwargs["tokenizer_mode"] = case.tokenizer_mode
305
+ if case.model is not None:
306
+ llm_kwargs["hf_overrides"] = partial(
307
+ _sleep_test_hf_overrides, model_arch=case.architecture
308
+ )
309
+ if case.modalities:
310
+ llm_kwargs["limit_mm_per_prompt"] = {
311
+ modality: 0 for modality in case.modalities
312
+ }
313
+
314
+ llm = LLM(**llm_kwargs)
315
+ prompt = TokensPrompt(prompt_token_ids=[2, 5, 9, 12, 7, 3])
316
+ sampling_params = SamplingParams(temperature=0.0, max_tokens=8, logprobs=1)
317
+
318
+ before_generation = _generation_signature(llm.generate(prompt, sampling_params))
319
+ snapshot_fn = partial(
320
+ _snapshot_sleep_tensors,
321
+ model_name=case.name,
322
+ tensor_names=case.tensor_names,
323
+ )
324
+ (before_tensors,) = llm.apply_model(snapshot_fn)
325
+ llm.collective_rpc(_save_dummy_weights)
326
+
327
+ llm.sleep(level=2)
328
+ llm.wake_up(tags=["weights"])
329
+ (after_tensors,) = llm.apply_model(snapshot_fn)
330
+ _assert_tensor_snapshots_equal(before_tensors, after_tensors, case.name)
331
+
332
+ llm.collective_rpc(_reload_dummy_weights)
333
+ llm.wake_up(tags=["kv_cache"])
334
+ after_generation = _generation_signature(llm.generate(prompt, sampling_params))
335
+
336
+ assert after_generation[0] == before_generation[0], case.name
337
+ assert after_generation[1] == pytest.approx(
338
+ before_generation[1], rel=1e-5, abs=1e-5
339
+ ), case.name
vllm/model_executor/models/conformer_encoder.py
+3
-1
@@ -68,7 +68,9 @@ class RelPositionalEncoding(nn.Module):
68
69
pe_positive = torch.flip(pe_positive, [0]).unsqueeze(0)
70
pe_negative = pe_negative[1:].unsqueeze(0)
71
- self.pe = torch.cat([pe_positive, pe_negative], dim=1)
71
+ self.register_buffer(
72
+ "pe", torch.cat([pe_positive, pe_negative], dim=1), persistent=False
73
+ )
74
75
def forward(self, x: torch.Tensor) -> torch.Tensor:
76
# Tmax = 2 * max_len - 1
vllm/model_executor/models/ernie45_vl.py
+13
-4
@@ -343,8 +343,12 @@ class Ernie4_5_VisionPatchEmbed(nn.Module):
343
class Ernie4_5_VisionRotaryEmbedding(nn.Module):
344
def __init__(self, dim: int, theta: float = 10000.0) -> None:
345
super().__init__()
346
- self.inv_freq = 1.0 / theta ** (
347
- torch.arange(start=0, end=dim, step=2, dtype=torch.float32) / dim
346
+ self.register_buffer(
347
+ "inv_freq",
348
+ 1.0
349
+ / theta
350
+ ** (torch.arange(start=0, end=dim, step=2, dtype=torch.float32) / dim),
351
+ persistent=False,
352
)
353
354
def forward(self, seqlen: int) -> torch.Tensor:
@@ -1452,11 +1456,16 @@ class Ernie4_5_VLMoeForConditionalGeneration(
1456
]
1457
if token_id is not None
1458
]
1455
- self._visual_token_ids_tensor_cache = torch.tensor(
1459
+ visual_token_ids_tensor_cache = torch.tensor(
1460
visual_token_ids, dtype=torch.long
1461
)
1462
else:
1459
- self._visual_token_ids_tensor_cache = None
1463
+ visual_token_ids_tensor_cache = None
1464
+ self.register_buffer(
1465
+ "_visual_token_ids_tensor_cache",
1466
+ visual_token_ids_tensor_cache,
1467
+ persistent=False,
1468
+ )
1469
1470
def compute_logits(
1471
self,
vllm/model_executor/models/gemma3n.py
+29
-13
@@ -114,8 +114,10 @@ class Gemma3nAltUp(nn.Module):
114
hidden_size=hidden_size,
115
eps=rms_norm_eps,
116
)
117
- self.router_input_scale = torch.tensor(
118
- hidden_size**-1.0, dtype=self.modality_router.weight.dtype
117
+ self.register_buffer(
118
+ "router_input_scale",
119
+ torch.tensor(hidden_size**-1.0, dtype=self.modality_router.weight.dtype),
120
+ persistent=False,
121
)
122
self.correct_output_scale = nn.Parameter(
123
torch.zeros(hidden_size, dtype=torch.float32)
@@ -597,9 +599,13 @@ class Gemma3nSelfDecoder(nn.Module):
599
quant_config=quant_config,
600
prefix=f"{prefix}.embed_tokens",
601
)
600
- self.embed_scale = torch.tensor(
601
- config.hidden_size**0.5,
602
- dtype=self.embed_tokens.weight.dtype,
602
+ self.register_buffer(
603
+ "embed_scale",
604
+ torch.tensor(
605
+ config.hidden_size**0.5,
606
+ dtype=self.embed_tokens.weight.dtype,
607
+ ),
608
+ persistent=False,
609
)
610
# Additional per-layer embeddings (PLE)
611
self.embed_tokens_per_layer = VocabParallelEmbedding(
@@ -608,9 +614,13 @@ class Gemma3nSelfDecoder(nn.Module):
614
quant_config=quant_config,
615
prefix=f"{prefix}.per_layer_embed_tokens",
616
)
611
- self.embed_scale_per_layer = torch.tensor(
612
- config.hidden_size_per_layer_input**0.5,
613
- dtype=self.embed_tokens.weight.dtype,
617
+ self.register_buffer(
618
+ "embed_scale_per_layer",
619
+ torch.tensor(
620
+ config.hidden_size_per_layer_input**0.5,
621
+ dtype=self.embed_tokens.weight.dtype,
622
+ ),
623
+ persistent=False,
624
)
625
self.per_layer_model_projection = ColumnParallelLinear(
626
config.hidden_size,
@@ -625,12 +635,18 @@ class Gemma3nSelfDecoder(nn.Module):
635
hidden_size=config.hidden_size_per_layer_input,
636
eps=config.rms_norm_eps,
637
)
628
- self.per_layer_input_scale = torch.rsqrt(torch.tensor(2.0)).to(
629
- self.embed_tokens.weight.dtype
638
+ self.register_buffer(
639
+ "per_layer_input_scale",
640
+ torch.rsqrt(torch.tensor(2.0)).to(self.embed_tokens.weight.dtype),
641
+ persistent=False,
642
)
631
- self.per_layer_projection_scale = torch.tensor(
632
- config.hidden_size**0.5,
633
- dtype=self.embed_tokens.weight.dtype,
643
+ self.register_buffer(
644
+ "per_layer_projection_scale",
645
+ torch.tensor(
646
+ config.hidden_size**0.5,
647
+ dtype=self.embed_tokens.weight.dtype,
648
+ ),
649
+ persistent=False,
650
)
651
self.altup_projections = nn.ModuleList(
652
[
vllm/model_executor/models/voxtral.py
+5
-1
@@ -765,7 +765,11 @@ class VoxtralEncoderModel(nn.Module):
765
max_frequency=8000.0,
766
sampling_rate=self.config.sampling_rate,
767
)
768
- self.mel_filters = torch.tensor(mel_filters, dtype=torch.float32)
768
+ self.register_buffer(
769
+ "mel_filters",
770
+ torch.tensor(mel_filters, dtype=torch.float32),
771
+ persistent=False,
772
+ )
773
774
def compute_whisper_melspec(
775
self,