xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-17 07:55:02 +08:00
Code Issues Packages Projects Releases Wiki Activity

[BUG FIX][NON-CUDA]quick fix to avoid call cudagraph_unsafe in attention (#25298)

Browse Source 
Signed-off-by: Chendi Xue <Chendi.Xue@intel.com>
This commit is contained in:
Chendi.Xue 2025-09-19 23:41:23 -05:00 committed by GitHub
parent b7f186bbb3
commit 6c5f82e5aa
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 6 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
vllm/attention/layer.py
View File

@ -29,6 +29,10 @@ from vllm.utils import GiB_bytes, direct_register_custom_op
logger = init_logger(__name__) logger = init_logger(__name__)
USE_XFORMERS_OPS = None 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(): def check_xformers_availability():
@ -577,7 +581,7 @@ direct_register_custom_op(
mutates_args=[], mutates_args=[],
fake_impl=unified_attention_fake, fake_impl=unified_attention_fake,
dispatch_key=current_platform.dispatch_key, dispatch_key=current_platform.dispatch_key,
tags=(torch._C.Tag.cudagraph_unsafe, ), tags=tag_cudagraph_unsafe,
) )
@ -628,5 +632,5 @@ direct_register_custom_op(
mutates_args=["output", "output_block_scale"], mutates_args=["output", "output_block_scale"],
fake_impl=unified_attention_with_output_fake, fake_impl=unified_attention_with_output_fake,
dispatch_key=current_platform.dispatch_key, dispatch_key=current_platform.dispatch_key,
tags=(torch._C.Tag.cudagraph_unsafe, ), tags=tag_cudagraph_unsafe,
) )