执行摘要

• 一句话：修复 XPU FlashAttention 接口参数缺失
• 推荐动作：此 PR 为一次性接口同步修复，变更清晰简单，无需精读。但对于维护 XPU 后端的工程师，建议后续关注 CUDA 端对该两参数的实现，并在 XPU kernel 中提供对应支持。

功能与动机

PR #42175 引入了新参数 mask_mod 和 aux_tensors 到 flash_attn_varlen_func 函数，但 XPU 路径的接口定义未同步更新，导致 XPU 设备上调用该函数时出现参数不匹配错误。PR 作者在 body 中明确说明："xpu path broke after https://github.com/vllm-project/vllm/pull/42175, which introduce a new paramter in flash_attn_varlen_func"。

实现拆解

1. 新增导入：在 vllm/_xpu_ops.py 文件顶部添加 from collections.abc import Callable 导入，为新增的 mask_mod: Callable | None 参数提供类型支持。
2. 补充方法签名：在 AttentionOp.flash_attn_varlen_func 静态方法最后两个参数位置添加 mask_mod: Callable | None = None 和 aux_tensors: list | None = None，保持与 CUDA 端接口签名一致。
3. 向下转发：这两个参数在 XPU 内核中实际未使用（已有注释说明），仅保持 API 兼容性；因此本变更仅涉及接口声明，不涉及底层 kernel 实现改动。

关键文件：

• vllm/_xpu_ops.py（模块 XPU 算子；类别 source；类型 dependency-wiring）: XPU 注意力操作接口定义处，新增了 mask_mod 和 aux_tensors 两个参数以对齐 CUDA 端接口。

关键符号：AttentionOp.flash_attn_varlen_func

关键源码片段

vllm/_xpu_ops.py

XPU 注意力操作接口定义处，新增了 mask_mod 和 aux_tensors 两个参数以对齐 CUDA 端接口。

# vllm/_xpu_ops.py
# 新增导入，支持 Callable 类型注解
from collections.abc import Callable
​
# ... 其他导入和代码 ...
​
class AttentionOp:
    @staticmethod
    def flash_attn_varlen_func(
        q: torch.Tensor,
        k: torch.Tensor,
        v: torch.Tensor,
        # ... 其他参数 ...
        return_softmax_lse: bool | None = False,
        s_aux: torch.Tensor | None = None,
        return_attn_probs: bool | None = False,
        # 新参数：与 CUDA 端接口对齐，保持 API 兼容性
        # XPU kernel 当前未使用这两个参数，但调用方可能传入
        mask_mod: Callable | None = None,
        aux_tensors: list | None = None,
    ):
        # ... 原有断言和逻辑 ...
        return flash_attn_varlen_func(
            # ... 原有参数转发 ...
            # 注意：mask_mod 和 aux_tensors 未传递给底层 kernel
        )

评论区精华

该 PR 没有产生 review 讨论。从审核记录看，reviewer bigPYJ1151 直接批准，说明变更清晰无争议。

• 暂无高价值评论线程

风险与影响

• 风险：风险极低：本 PR 仅添加两个可选参数，默认值为 None，不影响已有调用逻辑。但需要注意，若 CUDA 端未来在调用处强制传递非 None 值给这两个参数，而 XPU 端底层 kernel 未实现对应功能，则可能出现静默错误。
• 影响：影响范围很小，仅涉及 XPU 硬件平台上的注意力计算路径。该 PR 修复了因接口不匹配导致的 XPU 设备上 FlashAttention 调用失败问题，恢复 XPU 上使用 PR #42175 新功能的能力。由于参数在 XPU 端暂未实现实质功能，性能无影响。
• 风险标记：新参数在 XPU 内核中未实现, 跨平台接口一致性

关联脉络

• PR #42175 [Core][Model] Gemma4: Unified FA4 for all layers + FlashAttention mm_prefix support: 本 PR 修复的就是由 PR #42175 引入的接口不兼容问题，PR #42175 为 CUDA 端 flash_attn_varlen_func 新增了 mask_mod 和 aux_tensors 参数。