Skip to content

[torch.compile][Minor Fix] Gate cudagraph_unsafe tag for torch>=2.9 #25304

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
Open

[torch.compile][Minor Fix] Gate cudagraph_unsafe tag for torch>=2.9 #25304

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
8a673cf
gate cudagraph_unsafe tag for torch-2.9
BoyuanFeng Sep 20, 2025
68ed7cc
nit
BoyuanFeng Sep 20, 2025
56d5f48
move to vllm.utils.__init__
BoyuanFeng Sep 20, 2025
99aee97
gate by is_cuda_alike
BoyuanFeng Sep 20, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions tests/compile/silly_attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
import torch
from torch.library import Library

from vllm.utils import direct_register_custom_op
from vllm.utils import direct_register_custom_op, tag_cudagraph_unsafe

# Shared library for all compilation test operations
# Using "silly" namespace to match existing test expectations
Expand Down Expand Up @@ -60,5 +60,5 @@ def silly_attention_fake(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
mutates_args=["out"],
fake_impl=silly_attention_fake,
target_lib=silly_lib,
tags=(torch._C.Tag.cudagraph_unsafe, ),
tags=tag_cudagraph_unsafe,
)
7 changes: 2 additions & 5 deletions vllm/attention/layer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,14 +25,11 @@
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
from vllm.model_executor.models.vision import get_vit_attn_backend
from vllm.platforms import _Backend, current_platform
from vllm.utils import GiB_bytes, direct_register_custom_op
from vllm.utils import (GiB_bytes, direct_register_custom_op,
tag_cudagraph_unsafe)

logger = init_logger(__name__)
USE_XFORMERS_OPS = None
try:
tag_cudagraph_unsafe = (torch._C.Tag.cudagraph_unsafe, )
except AttributeError:
tag_cudagraph_unsafe = () # type: ignore[assignment]


def check_xformers_availability():
Expand Down
8 changes: 8 additions & 0 deletions vllm/utils/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3472,3 +3472,11 @@ def length_from_prompt_token_ids_or_embeds(
f" prompt_token_ids={prompt_token_len}"
f" prompt_embeds={prompt_embeds_len}")
return prompt_token_len


if is_torch_equal_or_newer("2.9.0.dev"):
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

May you also add condition check for attribute? We are working on xpu(intel GPU) and Hpu(intel gaudi), we use torch-XPU or torch-HPU, which I am afraid will not have cudagraph related attributes or assert.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes added current_platform.is_cuda_alike()

from vllm.platforms import current_platform
tag_cudagraph_unsafe = (torch._C.Tag.cudagraph_unsafe,
) if current_platform.is_cuda_alike() else ()
else:
tag_cudagraph_unsafe = () # type: ignore[assignment]