执行摘要

• 一句话：修复 trtllm_mha 在 FROZEN_KV MTP 下的 SWA 越界崩溃
• 推荐动作：该 PR 是一个教科书式的精确 bugfix：定位清晰、修改最小、逻辑自洽、参考了既有实现（FlashInfer）。值得关注的设计决策是「读取 allocator 而非 pool」作为稳定信源的思路，以及防御性 getattr 处理。推荐精读 _resolve_swa_kv_pool 方法和相关的守卫条件调整。

功能与动机

在 hybrid-SWA 模型（如 Gemma-4-E4B-IT）上使用 --attention-backend trtllm_mha 和 --speculative-algorithm NEXTN 时，当 SWA 池使用率超过约 85% 后，会确定性崩溃，报错 CUDA illegal memory access。根本原因是 TRTLLMHAAttnBackend 在 __init__ 中从 model_runner.token_to_kv_pool 缓存了 SWA 池状态，但在 FROZEN_KV MTP 模式下，draft worker 的 token_to_kv_pool 随每次前向调用被替换为目标的 SWAKVPool，而缓存状态未更新，导致 SWA 层使用了错误的索引空间（详见 Issue #26957）。

实现拆解

该 PR 由 2 个提交组成，均为单人改动，仅修改一个核心文件 trtllm_mha_backend.py，净增 22 行、删 12 行。

1. 移除 __init__ 中基于 token_to_kv_pool 的 SWA 池缓存：删除了原本直接通过 model_runner.token_to_kv_pool 判断并缓存 use_sliding_window_kv_pool 和 _swa_kv_pool 的代码。因为该属性会在 FROZEN_KV MTP 运行时被交换，导致初始判断结果过时失效。

2. 新增静态方法 _resolve_swa_kv_pool：该方法从 model_runner.token_to_kv_pool_allocator 获取 KV 缓存实例。allocator 在模型构造期间设置且后续稳定不变，不受 KV 池交换影响。方法通过 get_kvcache() 获得实际的缓存池，再判断其是否为 SWAKVPool 类型，并针对无 get_kvcache 属性的轻量 allocator 桩做了防御性处理（使用 getattr）。

3. 将所有 SWA 相关的守卫条件从 self.use_sliding_window_kv_pool 改为 self._swa_kv_pool is None/is not None：包括 _maybe_translate_swa、_alloc_swa_page_table、_get_layer_cache_loc、_get_layer_page_table 和 _copy_swa_page_table。这样确保所有 SWA 路径都基于统一的 _swa_kv_pool 字段判断，消除因布尔值未更新导致的死代码路径问题。

4. 移除旧的布尔型属性 use_sliding_window_kv_pool，统一使用 _swa_kv_pool 对象的存在性判断。

5. 测试与文档配套：PR Body 提到了端到端复现脚本可在 B200 上验证修复，但本次变更未新增单元测试文件；CI 中已有专用测试 test_gemma4_mtp_26b_a4b_trtllm_mha_extra.py，且该测试已通过。

关键文件：

• python/sglang/srt/layers/attention/trtllm_mha_backend.py（模块 注意力后端；类别 source；类型 core-logic；符号 _resolve_swa_kv_pool）: 唯一变更文件，包含核心 bugfix：新增 _resolve_swa_kv_pool 方法，将 SWA 池的解析从 token_to_kv_pool 切换到 token_to_kv_pool_allocator，并相应修改所有 SWA 守卫条件。

关键符号：_resolve_swa_kv_pool

关键源码片段

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

唯一变更文件，包含核心 bugfix：新增 _resolve_swa_kv_pool 方法，将 SWA 池的解析从 token_to_kv_pool 切换到 token_to_kv_pool_allocator，并相应修改所有 SWA 守卫条件。

# 新增静态方法：从 allocator 解析 SWA KV 池
@staticmethod
def _resolve_swa_kv_pool(model_runner: ModelRunner) -> Optional[SWAKVPool]:
    """返回用于索引翻译的 SWAKVPool，对非 SWA 模型返回 None。
​
    从 token_to_kv_pool_allocator 读取：在 FROZEN_KV MTP 模式下，
    draft 共享目标的 SWA allocator，而其自身的 token_to_kv_pool
    直到每次前向调用交换前都是非 SWA 的。
    使用 getattr 仅为兼容注意力测试夹具中的最小 allocator 桩。
    """
    allocator = model_runner.token_to_kv_pool_allocator
    get_kvcache = getattr(allocator, "get_kvcache", None)
    kvcache = get_kvcache() if get_kvcache is not None else None
    return kvcache if isinstance(kvcache, SWAKVPool) else None
​
# 在 __init__ 中替换原来的缓存逻辑（约 5 行→1 行）：
# 原来：
# kv_pool = model_runner.token_to_kv_pool
# self.use_sliding_window_kv_pool = isinstance(kv_pool, SWAKVPool)
# self._swa_kv_pool = kv_pool if self.use_sliding_window_kv_pool else None
# 改为：
self._swa_kv_pool: Optional[SWAKVPool] = self._resolve_swa_kv_pool(model_runner)

评论区精华

该 PR 的 review 评论数量较少（0 条），但获得了两位 reviewer（pyc96 和 ispobock）的批准。评论中无争议性讨论。PR Body 和 commit 消息已清晰说明了 root cause 和修复思路。

• 暂无高价值评论线程

风险与影响

• 风险：
  • 回归风险：改动集中在 __init__ 和核心 SWA 守卫逻辑，涉及 trtllm_mha 后端的初始化路径和所有 SWA 相关方法。_resolve_swa_kv_pool 使用了 getattr 防御性设计，对非标准 allocator 桩兼容，但若 allocator 返回的 kvcache 类型不符合预期，仍可能引入问题。对于非 hybrid-SWA 模型，_swa_kv_pool 为 None，逻辑退化为原有行为，回归风险较低。
  • 性能风险：无新增计算密集型操作，仅将初始化阶段的属性判断从 O(1) 调整为 O(1) + 方法调用开销，可忽略不计。
  • 兼容性风险：移除了 use_sliding_window_kv_pool 公共属性，若其他模块直接引用该属性（当前代码库中似乎没有）会导致兼容问题。建议确认无外部依赖。
• 影响：
  • 用户影响：修复了 Gemma-4 等 hybrid-SWA 模型在 Blackwell 上使用 trtllm_mha + FROZEN_KV MTP 时的崩溃问题，使该组合可行。
  • 系统影响：改进了 TRTLLMHAAttnBackend 对动态 KV 池交换场景的鲁棒性，为未来更多 MTP 或 speculative decoding 变体铺平道路。
  • 团队影响：合并后需确保相关 CI 测试（如 test_gemma4_mtp_26b_a4b_trtllm_mha_extra.py）持续通过，并验证其他模型后端（如 FlashInfer）的一致性。
  • 风险标记：核心路径变更, 缺少直接单元测试

关联脉络

• PR #26531 prior attempt at fix: PR body 提及之前的修复尝试，但此处是更精确的修复。
• PR #26643 [DP] Fix FlashInfer dispatcher workspace sizing and set_dp_buffer_len: 同为闪电推断（FlashInfer）后端中的类似 SWA 池修复，提供了参考模式：对分配器而非池进行 instanceof 检查。