use the same stream for cuda graph catpure and replay for NCCL

[Amir-19]

Purpose

With NCCL 2.28+, window registration during cuda graph capture crashes with NCCL complaining about the registration. since ncclMemAlloc is tied to a stream, warm up and cuda graph capture on separate streams causes new memory allocations and thus window registrations (when you use the mempool associated with ncclMemAlloc and ncclCommWindowRegister on different streams, if there are no available segments, you need new allocations and thus registration.). in this PR, we explicitly set the stream for cuda graph capture and forcing it to be on the same stream as warm up iterations.

Test Plan

Test Result

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request makes the CUDA stream explicit when capturing CUDA graphs, which is a good practice for clarity and correctness. I've found one area for improvement in vllm/compilation/cuda_graph.py regarding the use of torch.cuda.current_stream(), where a more performant, cached version from vLLM's utilities should be used instead.

[gemini-code-assist]

For performance reasons, it's better to use current_stream from vllm.utils.torch_utils instead of torch.cuda.current_stream(). The vLLM version is a cached version that avoids the overhead of creating a new stream object on each call, as documented in vllm/utils/torch_utils.py.

You'll need to update the import on line 20:

from vllm.utils.torch_utils import current_stream, weak_ref_tensors

Suggested change

	with torch.cuda.graph(
	cudagraph,
	pool=self.graph_pool,
	stream=torch.cuda.current_stream(),
	):
	with torch.cuda.graph(
	cudagraph,
	pool=self.graph_pool,
	stream=current_stream(),
	):

[BoyuanFeng]

iiuc, graph capture should happen on a side stream instead of current main stream?

We can add self.stream = torch.cuda.Stream() in the _init_ of CUDAGraphWrapper, and use this stream for both warmup and graph capture.

[Amir-19]

isn't the current_stream here already a non-default stream shared with warm-up iterations?

[BoyuanFeng]

do you have a code pointer to a non-default stream shared with warm-up iterations? Looks like in cuda_graph.py, there is no explicit cuda stream.

[Amir-19]

I will check the fellow again to see where a context manager sets a new stream.

[zou3519]

@Amir-19 what's the context behind this PR?

[Amir-19]

@zou3519 this came out as the result of investigating #28901. With NCCL 2.28+, window registration during cuda graph capture crashes with NCCL complaining about the registration. since ncclMemAlloc is tied to a stream, warm up and cuda graph capture on separate streams causes new memory allocations and thus window registrations. in this PR, we explicitly set the stream for cuda graph capture and forcing it to be on the same stream as warm up iterations. before this PR, cuda graph didn't have an explicit stream so it would create a new side stream, is that intentional?

[BoyuanFeng]

before this PR, cuda graph didn't have an explicit stream so it would create a new side stream, is that intentional?

yes. torch.cuda.graph(...) automatically creates a side stream, unless it is called with explicit stream torch.cuda.graph(..., stream=explicit_stream).

With NCCL 2.28+, window registration during cuda graph capture crashes with NCCL complaining about the registration.

Could you elaborate what is window registration?

we explicitly set the stream for cuda graph capture and forcing it to be on the same stream as warm up iterations

In general, warm up on one stream and graph capture on another stream is fine, except some extra memory consumption. So using the same stream for warmup and capture is an optimization.

However, using different streams should not lead to an error. Could you elaborate a bit on why it errors? e.g., the window registration is a cudagraph unsafe op?

[Amir-19]

@BoyuanFeng ncclCommWindowRegister is used to register local buffers into NCCL window which enables us to use symmetric kernels. This window registration also requires the memory to be from a VMM-based allocators like ncclMemAlloc. since memory allocated using ncclMemAlloc is tied to a stream, when you use the mempool associated with ncclMemAlloc and ncclCommWindowRegister on different streams, if there are no available segments, you need new allocations and thus registration.

intuitively, graph capture means "do this every time graph replays", so even if registration was allowed during cuda graph capture, it would have lead to creating new window handles each time which required destroying them later. This is not efficient or useful. the proper pattern is to register the window once before capture, then reuse it.

starting with NCCL 2.28, there is this restriction that ncclCommWindowRegister should not be called during graph capture which caused the failure reported in #28901

to fix this we need to make sure that the warm up and cuda graph capture are on the same side stream.

[BoyuanFeng]

what would be the issue without this change?

[Amir-19]

this benchmark would fail with NCCL 2.28 as cuda graph would create a new side stream for capture.

[BoyuanFeng]

setting context via torch.cuda.stream(stream) does not pass the stream to cudagraph capture. this is a no-op.

Would there be an issue w/o the change?

[Amir-19]

without this the call to the stream=torch.cuda.current_stream() in cuda_graph.py would return the default stream and capturing on the default stream is not allowed.

[youkaichao]

it should not be necessary after 73094c7

[BoyuanFeng]

do you have a code pointer to a non-default stream shared with warm-up iterations? Looks like in cuda_graph.py, there is no explicit cuda stream.

[BoyuanFeng]

before this PR, cuda graph didn't have an explicit stream so it would create a new side stream, is that intentional?

yes. torch.cuda.graph(...) automatically creates a side stream, unless it is called with explicit stream torch.cuda.graph(..., stream=explicit_stream).

With NCCL 2.28+, window registration during cuda graph capture crashes with NCCL complaining about the registration.

Could you elaborate what is window registration?

we explicitly set the stream for cuda graph capture and forcing it to be on the same stream as warm up iterations

In general, warm up on one stream and graph capture on another stream is fine, except some extra memory consumption. So using the same stream for warmup and capture is an optimization.

However, using different streams should not lead to an error. Could you elaborate a bit on why it errors? e.g., the window registration is a cudagraph unsafe op?

[youkaichao]

Sorry for the long delay.

Taking a closer look at the error stack at #28901 , the error happens at https://github.com/NVIDIA/nccl/blob/v2.28.3-1/src/dev_runtime.cc#L591 , and the stream nccl uses to synchonize is created by nccl itself from https://github.com/NVIDIA/nccl/blob/v2.28.3-1/src/dev_runtime.cc#L584 , which is a local variable defined at https://github.com/NVIDIA/nccl/blob/v2.28.3-1/src/dev_runtime.cc#L553 . It should not interfere with anything outside nccl.

My hypothesis is there might be some caching inside the driver, that when we create many streams, some streams will actually be the same stream, causing the trouble here. The solution is, then to create less streams to reduce the stream collision.

NOTE: I think pytorch has some stream-caching:

import torch

assert torch.cuda.is_available(), "CUDA is required for CUDA streams"

streams = [torch.cuda.Stream() for _ in range(1000)]

# Collect stream pointers
ptrs = []
for i, s in enumerate(streams):
    # cuda_stream is an integer pointer (uintptr_t)
    ptr = s.cuda_stream
    ptrs.append(ptr)

# Check for duplicates
unique_ptrs = set(ptrs)
print("\nSummary:")
print(f"Total streams created: {len(ptrs)}")
print(f"Unique stream pointers: {len(unique_ptrs)}")

if len(unique_ptrs) != len(ptrs):
    print("⚠️ Duplicate streams detected!")
else:
    print("✅ No duplicate streams detected.")

It only prints:

Summary:
Total streams created: 1000
Unique stream pointers: 32
⚠️ Duplicate streams detected!

Which means I can create at most 32 unique streams from pytorch.

If I don't use pytorch, I can create 10k unique streams from cuda directly. It still seems to be a bug somewhere.

I'll try to dig further with driver / nccl team, but for now the PR looks good to me. Reducing the number of streams created by pytorch, can reduce the chance of stream collision (since pytorch can only create up to 32 streams).

[youkaichao]

Diving deeper, after turning on CUDA_LOG_FILE=stderr (only valid for cuda 12.9+), I can get more error information:

[06:56:50.922][136234295043968][CUDA][E] API call conflicts with a stream capture sequence initiated from the calling thread

The limitation imposed by this PR "forcing cudagraph capture stream to be the same stream as warm up iterations" seems to be either improper use of nccl or a driver bug.

Nevertheless, I will stop here and let nccl/driver team continue the investigation.

The fix in this PR is simple, and we can go ahead with this workaround.

[youkaichao]

there's an error in ci:

RuntimeError: CUDA graphs must be captured on a non-default stream. (However, after capture, it's ok to replay them on the default stream.)

@Amir-19 I think we need to update

vllm/vllm/utils/torch_utils.py

Line 451 in 42826bb

def current_stream() -> torch.cuda.Stream:

too, to use a similar logic to rocm, to avoid using the default stream.

[youkaichao]

the only failing test comes from main. merging.

[pkousha]

@youkaichao Pouya from NCCL team here.

The limitation imposed by this PR "forcing cudagraph capture stream to be the same stream as warm up iterations" seems to be either improper use of nccl or a driver bug.
Nevertheless, I will stop here and let nccl/driver team continue the investigation.

I am not sure there is a bug in nccl VS a usage issue. let me explain: symmetric memory window registration is tied to a CUDA stream. Warm up is capturing that and putting things in a graph means "do this everytime graph is executed". Re-registering the same memory might be ok, we cache that at the physical layer, but we don't guarantee that the window handle written will be the same if the memory is already registered. So it literally means write over the window handle with a new window (that needs destroying later) every time. then who will free up that handle?
Do you still think there is a bug?