[Bugfix] Fix DSA crash under breakable piecewise cudagraphs (#49302)
Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Nick Hill committed
Jul 21, 2026 at 18:28 UTC
de6ec294ef0772673ade8fa3abf3f61f501880ae
3 files changed
+30
-12
tests/v1/cudagraph/test_breakable_cudagraph.py
+24
-2
@@ -6,7 +6,6 @@ Unit tests for the breakable cudagraph primitives.
6
7
from __future__ import annotations
8
9
-import os
9
import threading
10
from contextlib import nullcontext
11
from unittest.mock import patch
@@ -14,7 +13,21 @@ from unittest.mock import patch
13
import pytest
14
import torch
15
17
-os.environ["VLLM_USE_BREAKABLE_CUDAGRAPH"] = "1"
16
+
17
+@pytest.fixture(autouse=True)
18
+def _enable_breakable_cudagraph(monkeypatch: pytest.MonkeyPatch):
19
+ """Enable breakable cudagraphs for this module's tests only.
20
+
21
+ eager_break_during_capture reads the env at decoration time, which
22
+ happens inside the test bodies, so a per-test fixture suffices.
23
+ monkeypatch restores the env so other test files running in the same
24
+ pytest process are unaffected (a module-level os.environ assignment
25
+ used to leak into test_cudagraph_dispatch.py and break it).
26
+ """
27
+ import vllm.envs as envs
28
+
29
+ monkeypatch.setenv("VLLM_USE_BREAKABLE_CUDAGRAPH", "1")
30
+ envs.disable_envs_cache()
31
32
33
def test_piecewise_capture_builds_fresh_metadata_for_both_passes():
@@ -86,10 +99,19 @@ def cuda_capture_stream():
99
"""
100
if not torch.cuda.is_available():
101
pytest.skip("CUDA required")
102
+ from vllm.utils.torch_utils import _current_stream_tls
103
+
104
+ prev_stream = getattr(_current_stream_tls, "value", None)
105
stream = torch.cuda.Stream()
106
with torch.cuda.stream(stream):
107
yield stream
108
torch.cuda.current_stream().wait_stream(stream)
109
+ # Exiting torch.cuda.stream() records the default stream in vllm's
110
+ # patched set_stream cache. A later CUDAGraphWrapper capture in the same
111
+ # process would then run on the default stream, which cannot capture
112
+ # (this broke test_cudagraph_dispatch.py when run in-process after this
113
+ # file). Restore the pre-fixture value so this module leaves no trace.
114
+ _current_stream_tls.value = prev_stream
115
116
117
# ---------------------------------------------------------------------------
vllm/v1/worker/gpu/cudagraph_utils.py
+5
-5
@@ -254,12 +254,12 @@ class CudaGraphManager:
254
255
if mixed_mode:
256
# for PIECEWISE graphs there is no limit on requests when replaying
257
- # i.e. no request padding is needed
258
- # so we leave it as None
257
+ # i.e. no request padding is needed, so we leave it as None.
258
+ # For breakable PW graphs, break-point kernels read the real batch
259
+ # from the forward context; in-graph kernels handle the token padding
260
+ # themselves from the padded slot_mapping (rows with slot == -1).
261
num_reqs = None
260
- if mixed_mode == CUDAGraphMode.FULL or (
261
- mixed_mode == CUDAGraphMode.PIECEWISE and self.use_breakable_cg
262
- ):
262
+ if mixed_mode == CUDAGraphMode.FULL:
263
num_reqs = min(num_tokens, self.max_num_reqs)
264
desc = BatchExecutionDescriptor(
265
cg_mode=mixed_mode,
vllm/v1/worker/gpu/model_states/default.py
+1
-5
@@ -5,7 +5,6 @@ from typing import Any
5
import torch
6
import torch.nn as nn
7
8
-from vllm.compilation.breakable_cudagraph import is_breakable_cudagraph_enabled
8
from vllm.config import VllmConfig
9
from vllm.config.compilation import CUDAGraphMode
10
from vllm.v1.core.sched.output import NewRequestData
@@ -140,10 +139,7 @@ class DefaultModelState(ModelState):
139
kv_cache_config: KVCacheConfig,
140
for_capture: bool = False,
141
) -> dict[str, Any]:
143
- if cudagraph_mode == CUDAGraphMode.FULL or (
144
- cudagraph_mode == CUDAGraphMode.PIECEWISE
145
- and is_breakable_cudagraph_enabled()
146
- ):
142
+ if cudagraph_mode == CUDAGraphMode.FULL:
143
# Use padded sizes - padding is handled by model_runner.prepare_attn.
144
num_reqs = input_batch.num_reqs_after_padding
145
num_tokens = input_batch.num_tokens_after_padding