| 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 | ) |