replace with torch.cuda.stream()

[xinyu-intel]

Purpose

Part of #30679

In PyTorch, torch.Stream objects implement enter / exit and it should work well with context manager. So, replace with torch.cuda.stream(astream): with with astream:. For cases which astream can be None, use torch. accelerator.current_stream() instead.

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 refactors the codebase to replace the deprecated torch.cuda.stream() with the modern context manager usage on torch.Stream objects. The changes are generally correct and improve code clarity. However, I've identified a potential performance issue in one of the files where a new torch.Stream() is created in a loop, which can be inefficient. I've provided a suggestion to address this.

[gemini-code-assist]

Creating a new torch.Stream() on every forward pass when self.shared_experts_stream is None is inefficient and can lead to resource exhaustion. The original code with torch.cuda.stream(None) was a no-op that ran on the current stream, which is more efficient.

To preserve the original behavior while using the new context manager syntax, you can use torch.cuda.current_stream() as a fallback. This avoids creating a new stream in a loop.

Suggested change

	shared_experts_stream = self.shared_experts_stream or torch.Stream()
	shared_experts_stream = self.shared_experts_stream or torch.cuda.current_stream()

[hickeyma]

Agree

[xinyu-intel]

fixed.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @xinyu-intel.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

Documentation preview: https://vllm--33225.org.readthedocs.build/en/33225/

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @xinyu-intel.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

Hi @xinyu-intel, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[hmellor]

Please fix DCO

[hmellor]

You should also update the pattern in tools/pre_commit/check_torch_cuda.py to include stream

[mergify]

Hi @xinyu-intel, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @xinyu-intel.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @xinyu-intel.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork