diff --git a/vllm/_xpu_ops.py b/vllm/_xpu_ops.py
index 1adad42f10..962efd7724 100644
--- a/vllm/_xpu_ops.py
+++ b/vllm/_xpu_ops.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from collections.abc import Callable
 from typing import TYPE_CHECKING
 
 import torch
@@ -783,6 +784,8 @@ class xpu_ops:
         return_softmax_lse: bool | None = False,
         s_aux: torch.Tensor | None = None,
         return_attn_probs: bool | None = False,
+        mask_mod: Callable | None = None,
+        aux_tensors: list | None = None,
     ):
         assert cu_seqlens_k is not None or seqused_k is not None, (
             "cu_seqlens_k or seqused_k must be provided"