执行摘要

• 一句话：修复 FA3 varlen 注意力 out 参数传递错误
• 推荐动作：建议批准合并。该 PR 修复了一个实际运行时的 bug，影响面小，逻辑清晰。后续可考虑为 _call_fa3_kernel 添加单元测试。

功能与动机

flash_attn_varlen_func 未通过 _call_fa3_kernel 辅助函数，而是直接传递 out=out，导致不支持 out 参数的内核抛出 flash_attn_varlen_func() got an unexpected keyword argument 'out' 错误。

实现拆解

1. 在 python/sglang/jit_kernel/flash_attention_v3.py 中，将 _call_fa3_kernel 函数签名从 (kernel, *args, out=None) 修改为 (kernel, *args, out=None, **kwargs)，使函数能接收任意关键字参数。
2. 函数内部的所有 kernel()`` 调用均通过**kwargs` 传递额外参数，保持一致性。
3. 将 flash_attn_varlen_func() 函数中原本直接调用的内核改为通过 _call_fa3_kernel() 辅助函数调用，实现与 flash_attn_with_kvcache 相同的 out 参数 fallback 逻辑。
4. 在 .github/workflows/pr-test.yml 中调整 CI 触发器路径，将 python/sglang/jit_kernel/diffusion/** 扩大为 python/sglang/jit_kernel/**，确保本次变更（位于 flash_attention_v3.py）能被 CI 正确触发。

关键文件：

• python/sglang/jit_kernel/flash_attention_v3.py（模块 JIT 内核；类别 source；类型 core-logic；符号 _call_fa3_kernel）: 核心修复文件：修改 _call_fa3_kernel 支持 kwargs，并将 flash_attn_varlen_func 调用路由到该辅助函数。
• .github/workflows/pr-test.yml（模块 CI 配置；类别 infra；类型 infrastructure）: CI 配置调整：扩大触发器 glob 路径，确保 jit_kernel 目录下变更能触发正确 CI 流程。

关键符号：_call_fa3_kernel

关键源码片段

python/sglang/jit_kernel/flash_attention_v3.py

核心修复文件：修改 _call_fa3_kernel 支持 kwargs，并将 flash_attn_varlen_func 调用路由到该辅助函数。

def _call_fa3_kernel(kernel, *args, out=None, **kwargs):
    # 支持额外关键字参数，使 varlen 和 kvcache 分支共用同一 fallback 逻辑
    if out is None:
        return kernel(*args, **kwargs)
    try:
        return kernel(*args, **kwargs, out=out)
    except TypeError as exc:
        if "unexpected keyword argument 'out'" not in str(exc):
            raise
        # 内核不支持 out 参数时静默回退
        return kernel(*args, **kwargs)

评论区精华

仅有一条来自 gemini-code-assist[bot] 的自动审查评论，确认变更合理性，无额外讨论。

• 暂无高价值评论线程

风险与影响

• 风险：风险较低：变更集中在一个辅助函数和一个调用点，逻辑简单且已有成熟的 kvcache 分支作为参考。但缺少单元测试覆盖该 fallback 路径，若未来内核接口变化可能导致静默回退行为不符合预期。
• 影响：影响范围为使用 FA3 后端的扩散模型推理，修复了因 out 参数导致的内核错误。由于是运行时 fallback，对正常使用 out 参数的内核无影响。
• 风险标记：缺少测试覆盖

关联脉络

• 暂无明显关联 PR