执行摘要

• 一句话：清理 TRTLLM MHA 后端 dead code 并重用 page-table buffer
• 推荐动作：建议精读以了解 trtllm_mha 后端的内部机制和 CUDA graph 元数据处理。清理过程展示了如何安全移除 dead code，值得学习。

功能与动机

PR #28106 引入了设备端 page-table 构建，移除了 host-side max 同步需求，但遗留了一些不再使用的字段和分支。后续清理可以降低维护成本、减少 CPU-GPU 同步点，并让代码更易理解。作者在 PR body 中明确列出了 5 项清理目标。

实现拆解

1. 移除 max_seq_len_k 字段：在 TRTLLMMHAMetadata 数据类中删除 max_seq_len_k 及其所有赋值（包括 _apply_cuda_graph_metadata 中的 3 处），因为 flashinfer 直接从 max_context_len 获取该信息。
2. 删除不可达 draft-extend 分支：_apply_cuda_graph_metadata 中 if forward_mode.is_draft_extend_v2() 内部有一个 else 分支，由于外层已经判断了 is_draft_extend_v2()，该分支永远无法执行，将其移除。
3. 去重 max_num_pages 计算：init_cuda_graph_state 中原先使用 (self.max_context_len + self.page_size - 1) // self.page_size 重新计算，现在直接复用 self.max_num_pages（已在别处初始化）。
4. 重用 eager page-table buffer：在 init_forward_metadata_out_graph 中，尝试重用按需增长的 self.eager_page_table_buffer 替代每次分配新 tensor。该提交后来被回退（原因见提交 #4）。
5. 添加断言与文档：在 build_trtllm_mha_page_table 中添加 _MHA_KV_INDEX_BLOCK_TOKENS % page_size == 0 断言，确保 page-block token span 能被 page size 整除；在 create_trtllm_mha_kv_indices_triton 的 docstring 中记录 SWA 查找假设有效 slot（无 -1 哨兵）。
6. 测试扩展：单元测试 test_trtllm_mha_page_table.py 中移除 _build_page_table_reference 未使用的 max_num_pages 参数，并将 page_size=256 加入 test_matches_reference_gather 的参数组合。

关键文件：

• python/sglang/srt/layers/attention/trtllm_mha_backend.py（模块 注意力后端；类别 source；类型 core-logic；符号 TRTLLMMHAMetadata, TRTLLMHAAttnBackend.init_cuda_graph_state, TRTLLMHAAttnBackend._apply_cuda_graph_metadata）: 核心清理：移除 max_seq_len_k、不可达分支、去重 max_num_pages 计算。
• test/registered/attention/test_trtllm_mha_page_table.py（模块 测试；类别 test；类型 test-coverage；符号 _build_page_table_reference, _run_case, TestTRTLLMMHAPageTable.test_matches_reference_gather）: 测试清理：移除未使用参数，扩展 page_size 覆盖至 256。
• python/sglang/srt/layers/attention/triton_ops/trtllm_mha_page_table.py（模块 Page-table 内核；类别 infra；类型 infrastructure；符号 build_trtllm_mha_page_table, create_trtllm_mha_kv_indices_triton）: 添加 page_size 整除断言和 SWA 文档说明。

关键符号：TRTLLMHAAttnBackend.init_cuda_graph_state, TRTLLMHAAttnBackend._apply_cuda_graph_metadata, build_trtllm_mha_page_table, create_trtllm_mha_kv_indices_triton, _build_page_table_reference

关键源码片段

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

核心清理：移除 max_seq_len_k、不可达分支、去重 max_num_pages 计算。

# 在 init_cuda_graph_state 中，直接复用 self.max_num_pages
# 替代重新计算：
# 之前：max_num_pages = (self.max_context_len + self.page_size - 1) // self.page_size
# 之后：
max_num_pages = self.max_num_pages
​
# 在 _apply_cuda_graph_metadata 中，删除了所有 max_seq_len_k 赋值
# 以及不可达的 else 分支
if forward_mode.is_decode_or_idle():
    # ...
    metadata.cache_seqlens_int32.copy_(seq_lens)
    # 已移除 : metadata.max_seq_len_k = self.max_context_len
    # ...
elif forward_mode.is_draft_extend_v2():
    # 已移除整个 else 块，因为外部已经判断了 is_draft_extend_v2()
    # ...

test/registered/attention/test_trtllm_mha_page_table.py

测试清理：移除未使用参数，扩展 page_size 覆盖至 256。

# 原签名：
# def _build_page_table_reference(req_to_token, req_pool_indices, cache_seqlens, page_size, max_num_pages, full_to_swa=None):
# 改为（移除 max_num_pages）：
def _build_page_table_reference(
    req_to_token: torch.Tensor,
    req_pool_indices: torch.Tensor,
    cache_seqlens: torch.Tensor,
    page_size: int,
    full_to_swa: Optional[torch.Tensor] = None,
) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
    # ... 实现不变
​
# test_matches_reference_gather 新增 page_size=256
for page_size in (1, 32, 64, 128, 256):
    # ...

python/sglang/srt/layers/attention/triton_ops/trtllm_mha_page_table.py

添加 page_size 整除断言和 SWA 文档说明。

# 在 build_trtllm_mha_page_table 中新增断言
assert (
    _MHA_KV_INDEX_BLOCK_TOKENS % page_size == 0
), f"page_size={page_size} must divide _MHA_KV_INDEX_BLOCK_TOKENS={_MHA_KV_INDEX_BLOCK_TOKENS}"
​
# create_trtllm_mha_kv_indices_triton 的 docstring 补充：
# The SWA lookup assumes valid (>= 0) slots, unlike
# translate_loc_from_full_to_swa's -1 sentinel handling; page-boundary
# reads stay within seq_len, so slots are always valid here.

评论区精华

核心讨论来自维护者 merrymercy 对“重用 eager page-table buffer”的疑问：他未能在最终代码中找到该变更。实际上，第4个提交（a69f1c36）已经回退了 buffer 重用，因为作者发现它会导致跨 stream 写入后读取 (WAR) 问题，所以保留了每次前向分配的做法。

• Eager page-table buffer 重用 (design): 该变更在后续提交中被回退，因为作者发现会导致跨 stream 的写入后读取 (WAR) 问题。

风险与影响

• 风险：
  1. 回归风险：减少了 max_seq_len_k 赋值，但 flashinfer 直接使用 max_context_len，因此无影响。
  2. 不可达分支删除：else 分支本不可达，删除后无风险。
  3. 断言新增：assert 在 page_size 不整除 4096 时会抛出异常，但现有配置中 page_size 均为 1/32/64/128/256，均为 4096 的因子，不会触发。
  4. 测试参数扩展：增加 page_size=256 覆盖，降低该配置下的潜在 bug 风险。
• 影响：
  • 影响范围：仅影响 trtllm_mha 后端，不影响其他注意力后端或模型。
  • 性能：无显著变化（buffer 重用被回退）。
  • 可维护性：减少 dead code，提升代码清晰度。
  • 兼容性：完全向后兼容。
  • 风险标记：核心路径变更, 遗漏回退影响

关联脉络

• PR #28106 [基础 PR] 引入设备端 page-table 构建: 本 PR 是对 #28106 的后续清理，直接依赖其引入的机制。