执行摘要

• 一句话：修复 SWA 翻译缓存在 BCG/CG 路径的警告
• 推荐动作：此 PR 是聚焦的 bugfix，逻辑清晰、改动量小（+14/-11），建议批准合并。虽然缺少新自动化测试，但修复已在实际模型上充分验证。值得关注的设计点是：参数改名揭示了 TokenToKVPool 和 TokenToKVAllocator 之间的职责边界——方法应定义在拥有属性的对象上，避免中间层转发。

功能与动机

PR#25824 引入 SWAKVPool.translate_loc_from_full_to_swa 时带有基于 (data_ptr, numel) 的每 batch memoisation 缓存，当 loc tensor 未先调用 invalidate_loc_cache() 就变化时会触发 warning。该 warning 在每个 BCG 和 CG 前向中都出现——说明 not a real error but noise。PR 旨在消除此 warning，避免用户困惑并降低日志噪音。

实现拆解

修复分为两个文件中的 3 个逻辑点：

1. model_runner.py — 将 self.token_to_kv_pool.invalidate_loc_cache() 调用从 can_run_graph 检查之后的 eager 路径（第 3295-3296 行）提前到 can_run_graph 检查之前（第 3269-3270 行）。这样所有前向路径（eager、常规 CG replay、PCG replay、BCG replay）在 init_forward_metadata 运行前都会先清除缓存。注意此调用被包裹在 if self.is_hybrid_swa: 条件内，只影响使用 sliding-window attention 的模型。

2. triton_backend.py — 修改 update_sliding_window_buffer 函数的签名：将参数 token_to_kv_pool_allocator 重命名为 token_to_kv_pool。这是因为 invalidate_loc_cache() 和 translate_loc_from_full_to_swa 方法都定义在 TokenToKVPool 上，而不是其子对象 allocator 上，改后可直接访问 pool 对象，消除中间代理调用。

3. triton_backend.py — 在 update_sliding_window_buffer 函数内部，在调用 translate_loc_from_full_to_swa 之前和之后各加一个 token_to_kv_pool.invalidate_loc_cache()。原因是 window_kv_indices 是一个每次新计算的 tensor（不同于稍后 set_kv_buffer 使用的 out_cache_loc），缓存会 miss 并触发 warning。前一个 invalidate_loc_cache() 防止上一轮缓存的 key 引发 warning；后一个 invalidate_loc_cache() 为后续使用 out_cache_loc 的 set_kv_buffer 调用清理缓存，使 out_cache_loc 的翻译能够正常使用缓存。

所有调用 update_sliding_window_buffer 的地方（_update_decode_kv_buffers, _update_target_verify_buffers, init_forward_metadata 中的 decode 和 target_verify 分支）都同步将 token_to_kv_pool_allocator=... 改为 token_to_kv_pool=...。

测试方面：未增加新测试文件，但报告称已在 GptOssForCausalLM（hybrid SWA, 24 交替层）上使用 --enable-breakable-cuda-graph 验证：CG capture、BCG capture（74 种 size）和实际请求服务均无 warning。

关键文件：

• python/sglang/srt/layers/attention/triton_backend.py（模块 注意力；类别 source；类型 core-logic；符号 update_sliding_window_buffer, _update_decode_kv_buffers, _update_target_verify_buffers, init_forward_metadata）: 核心修复：重命名参数并添加双向 invalidate_loc_cache 调用以消除误警告
• python/sglang/srt/model_executor/model_runner.py（模块 执行引擎；类别 source；类型 control-flow；符号 _forward_raw）: 将 invalidate_loc_cache 从 eager 路径提前到 can_run_graph 之前，覆盖 BCG/CG replay 路径

关键符号：update_sliding_window_buffer, _update_decode_kv_buffers, _update_target_verify_buffers, init_forward_metadata, _forward_raw

关键源码片段

python/sglang/srt/layers/attention/triton_backend.py

核心修复：重命名参数并添加双向 invalidate_loc_cache 调用以消除误警告

def update_sliding_window_buffer(
    window_kv_indptr, req_to_token, sliding_window_size,
    seq_lens, req_pool_indices, bs, device=None,
    token_to_kv_pool=None, # 原来叫 token_to_kv_pool_allocator，但 invalidate_loc_cache 和 translate_loc_from_full_to_swa 都定义在 pool 上
    window_kv_indices=None,
):
    # ... 前面的逻辑不变 ...
    if hasattr(token_to_kv_pool, "translate_loc_from_full_to_swa"):
        kv_last_index = window_kv_indptr[-1]
        # window_kv_indices 是 freshly-computed tensor，不同于 out_cache_loc
        # 先 flush 避免上一轮缓存的 key 触发 warning
        token_to_kv_pool.invalidate_loc_cache()
        window_kv_indices[:kv_last_index] = (
            token_to_kv_pool.translate_loc_from_full_to_swa(
                window_kv_indices[:kv_last_index]
            )
        )
        # 再 flush，为后续 set_kv_buffer 使用 out_cache_loc 清理缓存
        token_to_kv_pool.invalidate_loc_cache()
    return window_kv_indptr, window_kv_indices, window_kv_lens, window_kv_start_idx

def _update_decode_kv_buffers(self, bs, seq_lens, req_pool_indices):
    # ...
    if self.sliding_window_size is not None and self.sliding_window_size > 0:
        window_kv_indptr, _, window_kv_lens, _ = update_sliding_window_buffer(
            self.window_kv_indptr,
            self.req_to_token,
            self.sliding_window_size,
            seq_lens,
            req_pool_indices,
            bs,
            token_to_kv_pool=self.token_to_kv_pool, # 原为 token_to_kv_pool_allocator
            window_kv_indices=self.cuda_graph_window_kv_indices,
        )
    return kv_indptr, window_kv_indptr, window_kv_lens

python/sglang/srt/model_executor/model_runner.py

将 invalidate_loc_cache 从 eager 路径提前到 can_run_graph 之前，覆盖 BCG/CG replay 路径

# _forward_raw 方法内部，关键调整后顺序：
​
# (1) Hisparse coordinator 等待（不变）
if (...):
    self.hisparse_coordinator.wait_for_pending_backup()
    self.hisparse_coordinator.num_real_reqs.fill_(forward_batch.batch_size)
​
# (2) ★ 提前到此：确保 can_run_graph 之前的任何路径都清除缓存
if self.is_hybrid_swa:
    self.token_to_kv_pool.invalidate_loc_cache()
​
# (3) 原来 invalidate_loc_cache 在此位置之后，现在移到上面
if can_run_graph:
    ret = self.graph_runner.replay(...)
    return ModelRunnerOutput(...)
​
# (4) eager 路径继续 ...

评论区精华

本 PR 无公开 review 评论。但从 commit message 和 PR body 可知，作者 ch-wan 自己完成了 root cause 分析和修复设计，且修复已在实际模型上验证通过。

• 暂无高价值评论线程

风险与影响

• 风险：
  • 回归风险低：invalidate_loc_cache() 只重置两个 Python 属性为 None，不影响计算值；参数改名 token_to_kv_pool_allocator → token_to_kv_pool 只是访问路径变化，功能等价。
  • 覆盖范围：仅影响 is_hybrid_swa=True 的模型路径，不会影响标准 attention 模型。
  • 潜在风险：若第三方子类依赖旧参数名，会中断。但 update_sliding_window_buffer 属于内部函数，外部依赖概率低。
  • 测试覆盖：当前变更未包含自动化测试；仅依赖人工验证和已有的 CI 端到端测试。
• 影响：
  • 用户影响：消除 hybrid SWA 模型在 BCG/CG 路径上每次前向的 translate_loc_from_full_to_swa warning，日志更加干净。
  • 系统影响：每个 forward 前多两次 Python 属性赋 None，开销可忽略。
  • 团队影响：降低维护者排查 warning 噪音的负担。
  • 风险标记：核心路径变更, 缺少测试覆盖

关联脉络

• PR #25824 [PR#25824] likely introduced the memoisation cache that this PR fixes: 引入 translate_loc_from_full_to_swa 和缓存机制，本 PR 修复其缓存未正确失效导致的误 warning
• PR #26134 [refactor] unify cuda-graph capture/replay across attention backends: 同样涉及 attention backend 的 cuda graph 路径修改，与本 PR 有重叠的上下文