执行摘要
- 一句话：修复 FlashInfer 滑动窗口 paged prefill 掩码错误
- 推荐动作：值得合入，属于重要的正确性修复。建议开发者在涉及 FlashInfer 后端和滑动窗口注意力时采用此修复。该 PR 的设计简洁清晰，条件分支的选择值得学习。
功能与动机
FlashInfer 的滑动窗口 paged prefill 包装器在规划时未提供 window_left 参数，导致 FlashInfer 选择了 use_sliding_window=False 的内核。这样一来，每个元素的左侧窗口掩码被编译排除；虽然 forward() 会传入 window_left，但它仅驱动每个 CTA 块的 KV 起始跳过，导致窗口外的令牌仍能关注到过时的 KV，从而产生错误输出。该问题影响所有通过 paged-only prefill 路径（如 SGLANG_FLASHINFER_USE_PAGED=1、确定性推理、分段 CUDA 图、多模态模型等）服务的 SWA 模型。
实现拆解
- 修改
flashinfer_backend.py 中的 update_sliding_window 方法：在 paged-only prefill 分支（非 ragged、非 spec-verify）中，向 call_begin_forward 传递 window_left=sliding_window_size，其他路径传递 -1 以保持原有行为。
- 修改
call_begin_forward 方法：新增 window_left 参数，默认值为 -1。在 paged_plan_kwargs 中设置 window_left，从而在规划阶段就让 FlashInfer 选择正确的滑动窗口内核。
- 添加回归测试
test_flashinfer.py：新增 test_projected_swa_attention_cases_paged_mode 方法，通过环境变量 SGLANG_FLASHINFER_USE_PAGED=1 触发 paged-only 路径，并包含长序列（超过滑动窗口）的测试用例，确保修复后的正确性。测试用例包括现有 SWA 窗口边沿情形和一个无前缀的长序列情形（extend_lens=6,8,12，window=4）。
关键文件：
python/sglang/srt/layers/attention/flashinfer_backend.py（模块 注意力后端；类别 source；类型 core-logic）: 核心变更文件，在 update_sliding_window 和 call_begin_forward 方法中传递 window_left 参数，修复了 paged-only prefill 路径的 SWA 掩码问题。test/registered/attention/unittests/swa/test_flashinfer.py（模块 测试；类别 test；类型 test-coverage；符号 test_projected_swa_attention_cases_paged_mode）: 新增回归测试 test_projected_swa_attention_cases_paged_mode，覆盖 paged-only 路径下的 SWA 正确性，包括超过滑动窗口长度的情形。
关键符号：update_sliding_window, call_begin_forward
关键源码片段
python/sglang/srt/layers/attention/flashinfer_backend.py
核心变更文件，在 update_sliding_window 和 call_begin_forward 方法中传递 window_left 参数，修复了 paged-only prefill 路径的 SWA 掩码问题。
# 在 update_sliding_window 方法中，paged-only 分支传递 window_left
# paged-only 路径：wrapper_id==0, use_ragged=False, spec_info is None
window_left=(
sliding_window_size
if (wrapper_id == 0 and not use_ragged and spec_info is None)
else -1 # ragged 保持原有自定义前缀掩码，spec-verify 保持树掩码
)
# 在 call_begin_forward 方法中，将 window_left 加入 paged_plan_kwargs
# 以使 FlashInfer 在规划阶段就选择带有用滑动窗口掩码的内核
def call_begin_forward(
self,
...
window_left: int = -1, # 新增参数，默认 -1 表示不使用
):
...
paged_plan_kwargs = {
...
}
if window_left >= 0:
# 选择编译了每个元素窗口掩码的模块
paged_plan_kwargs["window_left"] = window_left
wrapper_paged.begin_forward(
qo_indptr,
kv_indptr,
...
**paged_plan_kwargs
)
test/registered/attention/unittests/swa/test_flashinfer.py
新增回归测试 test_projected_swa_attention_cases_paged_mode，覆盖 paged-only 路径下的 SWA 正确性，包括超过滑动窗口长度的情形。
# 在 TestFlashInferSWAAttentionBackendCorrectness 类中
# Paged-only prefill 没有 ragged 传递或自定义前缀掩码，因此内核必须强制窗口
# 长测试用例使令牌超出窗口
PAGED_MODE_CASES = CASES + (
DenseAttentionCase(
name="swa_extend_no_prefix_above_window_long",
backend="flashinfer",
forward_mode=ForwardMode.EXTEND,
num_heads=4,
num_kv_heads=4,
page_size=16,
prefix_lens=(0, 0, 0), # 无前缀
extend_lens=(6, 8, 12), # 超出滑动窗口 4
sliding_window_size=4,
),
)
# 新增测试方法，通过环境变量强制 paged 模式
def test_projected_swa_attention_cases_paged_mode(self):
for case in self.PAGED_MODE_CASES:
with self.subTest(case=case.name, backend=case.backend, mode="paged"):
with envs.SGLANG_FLASHINFER_USE_PAGED.override(True):
# 与精确的 SDPA 滑动窗口参考对比
run_dense_attention_case(
self,
case,
head_dim=self.HEAD_DIM,
hidden_size=self.HIDDEN_SIZE,
)
评论区精华
review 讨论较少，关键点包括：
风险与影响
关联脉络
参与讨论