[Test][V1] Add sleep/wake correctness regression test for hybrid GDN/… (#44972)
Signed-off-by: chun-wan <chun-wan@amd.com>
chun-wan committed
Aug 3, 2026 at 06:02 UTC
0033211c0bcf3d42d3b1f27059922c30e664b8e7
2 files changed
+100
csrc/cumem_allocator.cpp
+23
@@ -212,6 +212,29 @@ void create_and_map(unsigned long long device, ssize_t size, CUdeviceptr d_mem,
212
if (error_code != 0) {
213
return;
214
}
215
+
216
+#ifdef USE_ROCM
217
+ // ROCm-only: zero the freshly (re)mapped region.
218
+ //
219
+ // On some amdgpu driver versions the physical VRAM backing a buffer is not
220
+ // scrubbed on release, so the pages handed back here (e.g. on wake_up after a
221
+ // sleep that discarded a tag) can still contain stale data from a previously
222
+ // freed allocation -- possibly from another process. Callers of the device
223
+ // allocator expect newly mapped memory to be well defined; in particular
224
+ // hybrid Mamba / gated-delta-net (GDN) models keep a persisted conv/recurrent
225
+ // state cache, and a new sequence that reads an unzeroed slot will propagate
226
+ // garbage/NaN, collapsing generation to token id 0 (decoded as "!").
227
+ //
228
+ // This is scoped to the ROCm path on purpose: the CUDA path is unchanged.
229
+ // Use CUDA_CHECK so a failure sets the global error_code; the Python caller
230
+ // (python_create_and_map) only inspects error_code, so returning without
231
+ // setting it would make a failed memset look successful and still hand back
232
+ // unzeroed memory.
233
+ CUDA_CHECK((CUresult)hipMemset(reinterpret_cast<void*>(d_mem), 0, size));
234
+ if (error_code != 0) {
235
+ return;
236
+ }
237
+#endif
238
// std::cout << "create_and_map: device=" << device << ", size=" << size << ",
239
// d_mem=" << d_mem << ", p_memHandle=" << p_memHandle << std::endl;
240
}
tests/models/language/generation/test_gdn_sleep_wake.py
new
+77
@@ -0,0 +1,77 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Regression test for hybrid GDN/Mamba models under sleep -> wake.
4
+
5
+Hybrid Mamba / gated-delta-net (GDN) models (e.g. Qwen3-Next) keep a
6
+persisted conv + recurrent state cache. With sleep mode (the RLHF reuse
7
+pattern: ``sleep()`` -> weight update -> ``wake_up()``) the state-cache tag is
8
+discarded on sleep and its device memory is re-created on wake. If a *new*
9
+sequence's state slot is consumed before being reset, the gated-delta-rule
10
+kernel faithfully propagates whatever is in that (now non-zeroed) memory; when
11
+it contains NaN/inf the output becomes NaN and ``argmax`` collapses every token
12
+to id 0 (which decodes to ``"!"``), giving ``reward=0`` / NaN log-probs in RL
13
+training.
14
+
15
+This test sleeps and wakes a small hybrid GDN model and asserts that
16
+post-wake generation is neither degenerate (single repeated token) nor NaN.
17
+"""
18
+
19
+import pytest
20
+
21
+from vllm import LLM, SamplingParams
22
+
23
+# Small Qwen3-Next (GDN) model already used by the hybrid model test-suite.
24
+MODEL = "tiny-random/qwen3-next-moe"
25
+
26
+PROMPTS = [
27
+ "The capital of France is",
28
+ "Once upon a time,",
29
+ "1, 2, 3, 4,",
30
+ "Water is made of",
31
+]
32
+
33
+
34
+@pytest.mark.hybrid_model
35
+def test_gdn_sleep_wake_no_stale_state():
36
+ sampling_params = SamplingParams(temperature=0.0, max_tokens=32, logprobs=1)
37
+
38
+ # Keep the reserved fraction low. On some (notably ROCm/amdgpu) drivers the
39
+ # VRAM discarded by ``sleep()`` is not returned to the free pool before
40
+ # ``wake_up()`` re-creates it, so the woken allocation must coexist with the
41
+ # not-yet-reclaimed one (~2x peak). A high ``gpu_memory_utilization`` then
42
+ # OOMs in ``cuMemCreate`` on wake. The model is tiny, so a small fraction
43
+ # still leaves ample KV/state cache while keeping the sleep/wake cycle well
44
+ # within device memory.
45
+ llm = LLM(
46
+ model=MODEL,
47
+ enable_sleep_mode=True,
48
+ enforce_eager=True,
49
+ max_model_len=1024,
50
+ gpu_memory_utilization=0.4,
51
+ trust_remote_code=True,
52
+ )
53
+
54
+ # Warm generation before sleeping.
55
+ llm.generate(PROMPTS, sampling_params)
56
+
57
+ # Default sleep offloads weights and DISCARDS the kv / GDN state cache;
58
+ # wake_up re-creates that memory (fresh, not guaranteed zeroed).
59
+ llm.sleep()
60
+ llm.wake_up()
61
+
62
+ after = llm.generate(PROMPTS, sampling_params)
63
+
64
+ for output in after:
65
+ completion = output.outputs[0]
66
+ token_ids = list(completion.token_ids)
67
+ assert token_ids, "empty generation after wake_up"
68
+ # The bug collapses every token to a single id (e.g. 0 -> "!").
69
+ assert len(set(token_ids)) > 1, (
70
+ f"degenerate single-token output after wake_up: {token_ids[:16]}"
71
+ )
72
+ # NaN logits surface as NaN log-probs.
73
+ for step_logprobs in completion.logprobs or []:
74
+ for logprob in step_logprobs.values():
75
+ assert logprob.logprob == logprob.logprob, (
76
+ "NaN log-prob after wake_up (stale GDN state)"
77
+ )