diff --git a/vllm/_xpu_ops.py b/vllm/_xpu_ops.py index 288dbdae9b..aea1785897 100644 --- a/vllm/_xpu_ops.py +++ b/vllm/_xpu_ops.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING import torch from vllm_xpu_kernels.flash_attn_interface import flash_attn_varlen_func +from vllm.compilation.breakable_cudagraph import eager_break_during_capture from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.triton_utils import tl, triton @@ -1257,7 +1258,7 @@ class xpu_ops: direct_register_custom_op( op_name="gdn_attention_core_xpu", - op_func=_gdn_attention_core_xpu_impl, + op_func=eager_break_during_capture(_gdn_attention_core_xpu_impl), mutates_args=["core_attn_out", "z"], fake_impl=_gdn_attention_core_xpu_fake, )