A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 22:45 UTC
use the same stream for cuda graph catpure and replay for NCCL (#29207)
Signed-off-by: Amir Samani <asamani@nvidia.com> Signed-off-by: youkaichao <youkaichao@gmail.com> Co-authored-by: youkaichao <youkaichao@gmail.com>
Amir Samani committed
Dec 25, 2025 at 03:10 UTC
030fc4491465d361e4bed626d76c184f8a7d8a07
4 files changed
+23
-27
benchmarks/kernels/benchmark_device_communicators.py
+1
-1
@@ -293,7 +293,7 @@ class CommunicatorBenchmark:
293
graph = torch.cuda.CUDAGraph()
294
graph_pool = torch.cuda.graph_pool_handle()
295
set_graph_pool_id(graph_pool)
296
- with torch.cuda.graph(graph, pool=graph_pool):
296
+ with torch.cuda.graph(graph, pool=graph_pool, stream=stream):
297
for _ in range(CUDA_GRAPH_CAPTURE_CYCLES):
298
allreduce_fn(graph_input)
299
tests/utils_/test_torch_utils.py
+9
-21
@@ -99,30 +99,18 @@ def _test_stream_thread(main_expected_stream: torch.cuda.Stream):
99
100
101
def test_current_stream_multithread():
102
- from vllm.platforms import current_platform
103
-
102
if not torch.cuda.is_available():
103
pytest.skip("CUDA not available")
104
107
- if current_platform.is_rocm():
108
- main_dedicated_stream = current_stream()
109
-
110
- assert main_dedicated_stream.cuda_stream != 0, (
111
- "ROCm should create a dedicated stream, not use default stream (0x0)"
112
- )
113
-
114
- main_stream_again = current_stream()
115
- assert main_stream_again == main_dedicated_stream, (
116
- "Multiple calls to current_stream should return the same dedicated stream"
117
- )
105
+ main_dedicated_stream = current_stream()
106
119
- _test_stream_thread(main_dedicated_stream)
120
- else:
121
- main_default_stream = torch.cuda.default_stream()
122
- main_initial_stream = current_stream()
107
+ assert main_dedicated_stream.cuda_stream != 0, (
108
+ "ROCm/CUDA should create a dedicated stream, not use default stream (0x0)"
109
+ )
110
124
- assert main_initial_stream == main_default_stream, (
125
- "First call to current_stream should return default stream on CUDA"
126
- )
111
+ main_stream_again = current_stream()
112
+ assert main_stream_again == main_dedicated_stream, (
113
+ "Multiple calls to current_stream should return the same dedicated stream"
114
+ )
115
128
- _test_stream_thread(main_default_stream)
116
+ _test_stream_thread(main_dedicated_stream)
vllm/compilation/cuda_graph.py
+6
-2
@@ -18,7 +18,7 @@ from vllm.distributed.device_communicators.pynccl_allocator import set_graph_poo
18
from vllm.forward_context import BatchDescriptor, get_forward_context
19
from vllm.logger import init_logger
20
from vllm.platforms import current_platform
21
-from vllm.utils.torch_utils import weak_ref_tensors
21
+from vllm.utils.torch_utils import current_stream, weak_ref_tensors
22
23
logger = init_logger(__name__)
24
@@ -263,7 +263,11 @@ class CUDAGraphWrapper:
263
else:
264
set_graph_pool_id(current_platform.graph_pool_handle())
265
# mind-exploding: carefully manage the reference and memory.
266
- with torch.cuda.graph(cudagraph, pool=self.graph_pool):
266
+ with torch.cuda.graph(
267
+ cudagraph,
268
+ pool=self.graph_pool,
269
+ stream=current_stream(),
270
+ ):
271
# `output` is managed by pytorch's cudagraph pool
272
output = self.runnable(*args, **kwargs)
273
if self.cudagraph_options.weak_ref_output:
vllm/utils/torch_utils.py
+7
-3
@@ -465,9 +465,13 @@ def current_stream() -> torch.cuda.Stream:
465
# when this function is called before any stream is set,
466
# we return the default stream.
467
# On ROCm using the default 0 stream in combination with RCCL
468
- # is hurting performance. Therefore creating a dedicated stream
469
- # per process
470
- if current_platform.is_rocm():
468
+ # is hurting performance.
469
+ # On CUDA, we capture and replay cudagraph on the same stream,
470
+ # so we need to avoid using the default stream as well. The default
471
+ # stream cannot be used for cudagraph capture, see
472
+ # https://github.com/pytorch/pytorch/blob/42ad9edfb754743fdae3276ade43de000beb4f60/aten/src/ATen/cuda/CUDAGraph.cpp#L77
473
+ # for more details. Therefore, we create a dedicated stream per process.
474
+ if current_platform.is_rocm() or current_platform.is_cuda():
475
# torch.cuda.set_stream here is the alias of _pathed_set_stream
476
torch.cuda.set_stream(torch.cuda.Stream())
477
elif current_platform.is_cpu():