[Bugfix][XPU] Run GDN attention as eager break under breakable cudagraph (#51928)
Signed-off-by: Huanxing <huanxing.shen@intel.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
+2
-1
@@ -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,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user