# PR #31501 完整报告

- 仓库：`sgl-project/sglang`
- 标题：[flashinfer] Pass window_left at plan time for the SWA paged prefill wrapper
- 合并时间：2026-07-19 04:33
- 原文链接：http://prhub.com.cn/sgl-project/sglang/pull/31501

---

# 执行摘要

- 一句话：修复 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 模型。

# 实现拆解

1. **修改 `flashinfer_backend.py` 中的 `update_sliding_window` 方法**：在 paged-only prefill 分支（非 ragged、非 spec-verify）中，向 `call_begin_forward` 传递 `window_left=sliding_window_size`，其他路径传递 `-1` 以保持原有行为。
2. **修改 `call_begin_forward` 方法**：新增 `window_left` 参数，默认值为 `-1`。在 `paged_plan_kwargs` 中设置 `window_left`，从而在规划阶段就让 FlashInfer 选择正确的滑动窗口内核。
3. **添加回归测试 `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 掩码问题。

```python
# 在 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 正确性，包括超过滑动窗口长度的情形。

```python
# 在 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 讨论较少，关键点包括：
- 作者在 PR 描述中详细说明了问题根源：规划时未传递 `window_left` 导致选错内核模块，而 `forward()` 的 `window_left` 仅用于粗略的 CTA 块跳过，无法保证每个元素的掩码正确性。
- reviewer `merrymercy` 批准了该 PR。
- 用户 `vincentzed` 确认在 Olmo3 及其他模型中也观察到了类似问题，并表示感谢。

- 暂无高价值评论线程

# 风险与影响

- 风险：该 PR 仅修改了 FlashInfer 后端的 SWA paged prefill 路径，风险较低。主要风险为：
 - 确保 `window_left` 仅在正确的条件分支（paged-only）下传递，避免影响 ragged 和 spec-verify 路径。代码通过 `wrapper_id == 0 and not use_ragged and spec_info is None` 条件严格限制。
 - 新增的测试用例覆盖了无前缀且 extend 长度超过窗口的场景，可验证修复效果。
 - 需注意 paged-only 路径的额外回归测试，但功能逻辑简单（仅添加一个关键字参数），风险可控。
- 影响：**影响范围**：修复了 FlashInfer 后端在滑动窗口注意力场景下的 paged-only prefill 路径的正确性。受影响的使用场景包括：
 - 设置环境变量 `SGLANG_FLASHINFER_USE_PAGED=1`
 - 启用 `--enable-deterministic-inference`
 - 分段 CUDA 图 prefill
 - 多模态模型（如 Gemma3 VLM）
 - 多条目评分 (multi-item scoring)
**影响程度**：中等。对于使用上述路径的 SWA 模型（如 Gemma-3-1b-it），修复前输出存在显著错误（GSM8K exact_match 从 0.130 提升至 0.255，趋近 Triton 后端的 0.275），修复后 token logprob 差异从平均 2.3 降至 0.045。

- 风险标记：核心路径变更 , 缺少回归测试覆盖

# 关联脉络

- 暂无明显关联 PR