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 14, 05:30 UTC
[BUG FIX][NON-CUDA]quick fix to avoid call cudagraph_unsafe in attention (#25298)
Signed-off-by: Chendi Xue <Chendi.Xue@intel.com>
Chendi.Xue committed
Sep 19, 2025 at 23:41 UTC
6c5f82e5aa87cd73ce03ce10fc44138f75ee1aea
1 file changed
+6
-2
vllm/attention/layer.py
+6
-2
@@ -29,6 +29,10 @@ from vllm.utils import GiB_bytes, direct_register_custom_op
29
30
logger = init_logger(__name__)
31
USE_XFORMERS_OPS = None
32
+try:
33
+ tag_cudagraph_unsafe = (torch._C.Tag.cudagraph_unsafe, )
34
+except AttributeError:
35
+ tag_cudagraph_unsafe = () # type: ignore[assignment]
36
37
38
def check_xformers_availability():
@@ -577,7 +581,7 @@ direct_register_custom_op(
581
mutates_args=[],
582
fake_impl=unified_attention_fake,
583
dispatch_key=current_platform.dispatch_key,
580
- tags=(torch._C.Tag.cudagraph_unsafe, ),
584
+ tags=tag_cudagraph_unsafe,
585
)
586
587
@@ -628,5 +632,5 @@ direct_register_custom_op(
632
mutates_args=["output", "output_block_scale"],
633
fake_impl=unified_attention_with_output_fake,
634
dispatch_key=current_platform.dispatch_key,
631
- tags=(torch._C.Tag.cudagraph_unsafe, ),
635
+ tags=tag_cudagraph_unsafe,
636
)