执行摘要

• 一句话：回退#26134的CUDA graph统一重构，保留SWA修复
• 推荐动作：建议尽快合并以恢复主分支稳定性，并记录回退原因；后续统一重构应充分测试并增加针对性单元测试。本PR展示了review发现深度bug的价值，值得精读review讨论。

功能与动机

26134的重构虽减少了重复代码，但review中发现两个关键bug：FlashInfer后端中use_ragged参数被硬编码为True，与动态逻辑不一致；Triton后端在replay时错误地用len(req_pool_indices)重定义bs，导致batch size错误。这些问题可能引发运行时崩溃或静默错误。为保障稳定性，作者决定回退该PR。

实现拆解

1. 执行revert：第一个commit 160bf7b 使用git revert还原#26134的提交d226f75，自动处理大部分冲突，使五个后端文件恢复到重构前的状态。
2. 重新应用SWA修复：由于#26134删除了包含#26152修复的辅助方法，第二个commit f2bc52c在回退后的代码中重新应用了相同的修复，将update_sliding_window_buffer的参数名从token_to_kv_pool_allocator改为token_to_kv_pool，并调整相关调用。
3. 涉及文件：共修改5个注意力后端文件，均为python/sglang/srt/layers/attention/下的核心源码。
4. 测试配套：无直接测试文件变更，依赖上游测试。

关键文件：

• python/sglang/srt/layers/attention/triton_backend.py（模块 Triton后端；类别 source；类型 core-logic；符号 _fill_kv_indptr_and_indices, _update_decode_kv_buffers, _update_target_verify_buffers, _update_draft_extend_buffers）: 改动最大（+330/-284），核心CUDA graph缓冲区更新逻辑，恢复_fill_kv_indptr_and_indices等辅助方法。
• python/sglang/srt/layers/attention/flashinfer_backend.py（模块 FlashInfer后端；类别 source；类型 core-logic；符号 _create_decode_wrappers, _create_prefill_wrappers, _prepare_cuda_graph_metadata, init_forward_metadata_capture_cuda_graph）: 第二重要（+150/-89），恢复工厂方法并修复review指出的use_ragged不一致问题。
• python/sglang/srt/layers/attention/wave_backend.py（模块 Wave后端；类别 source；类型 core-logic；符号 _build_cuda_graph_forward_metadata）: 恢复了_build_cuda_graph_forward_metadata方法，修正capture阶段丢失get_num_kv_splits的问题。
• python/sglang/srt/layers/attention/flashinfer_mla_backend.py（模块 MLA后端；类别 source；类型 core-logic）: 修改较小（+38/-6），分离target_verify和draft_extend分支，消除合并分支的状况。
• python/sglang/srt/layers/attention/cutlass_mla_backend.py（模块 MLA后端；类别 source；类型 core-logic）: 修改较小（+23/-17），调整capture/replay中的控制流，恢复内联实现。

关键符号：_fill_kv_indptr_and_indices, _update_decode_kv_buffers, _update_target_verify_buffers, _update_draft_extend_buffers, _build_cuda_graph_forward_metadata, update_sliding_window_buffer_cuda_graph, _create_decode_wrappers, _create_prefill_wrappers, _prepare_cuda_graph_metadata, init_forward_metadata_capture_cuda_graph, init_forward_metadata_replay_cuda_graph

关键源码片段

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

改动最大（+330/-284），核心CUDA graph缓冲区更新逻辑，恢复_fill_kv_indptr_and_indices等辅助方法。

def _update_decode_kv_buffers(
    self,
    bs: int,
    seq_lens: torch.Tensor,
    req_pool_indices: torch.Tensor,
):
    # 在 CUDA graph 捕获 / 回放时填充 decode 模式的 KV 缓存缓冲区。
    # 该函数被 #26134 内联，revert 后重新提取为独立方法，提高可读性。
    seq_lens = seq_lens[:bs]
    req_pool_indices = req_pool_indices[:bs]
    kv_indptr = self._fill_kv_indptr_and_indices(
        bs, seq_lens, req_pool_indices, self.cuda_graph_kv_indices
    )
    window_kv_indptr = self.window_kv_indptr
    window_kv_lens = None
    if self.sliding_window_size is not None and self.sliding_window_size > 0:
        # 滑动窗口缓冲更新，参数名已随 #26152 修复
        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,
            window_kv_indices=self.cuda_graph_window_kv_indices,
        )
    return kv_indptr, window_kv_indptr, window_kv_lens

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

第二重要（+150/-89），恢复工厂方法并修复review指出的use_ragged不一致问题。

def _create_decode_wrappers(self, bs: int, num_tokens: int) -> list:
    # 工厂方法：创建 FlashInfer decode wrapper 列表
    # revert 后重新独立，防止 #26134 引入的 use_ragged 硬编码问题
    return [
        BatchDecodeWithPagedKVCacheWrapper(
            self.workspace_buffer,
            "NHD",
            backend=self.decode_backend,
            use_cuda_graph=True,
            use_tensor_cores=self.decode_use_tensor_cores,
            paged_kv_indptr_buffer=self.kv_indptr[i][: num_tokens + 1],
            paged_kv_indices_buffer=self.cuda_graph_kv_indices[i],
            paged_kv_last_page_len_buffer=self.kv_last_page_len[:num_tokens],
        )
        for i in range(self.num_wrappers)
    ]

评论区精华

Reviewer gemini-code-assist[bot] 发现两个高优先级问题：

• FlashInfer后端use_ragged不一致：在is_dllm_extend模式下，PrefillMetadata的use_ragged硬编码为True，但indices_updater_prefill.update使用not self.use_paged，当self.use_paged=True时引发矛盾，可能导致崩溃。

• Triton后端bs重定义：在init_forward_metadata_replay_cuda_graph中，bs = len(req_pool_indices)错误地使用缓冲区长度而非实际batch size，影响后续indptr计算和kernel grid大小。
  这些问题直接成为回退的决策依据。

• FlashInfer后端use_ragged参数不一致 (correctness): 该 bug 是revert的直接原因之一，reviewer 明确指出启动上下文不一致。

• Triton后端replay中bs重新定义 (correctness): reviewer 指出这是严重问题，必须修复；回归到使用参数中的bs。

风险与影响

• 风险：

  1. 丢失WaveBackend修复：#26134修复了Wave后端capture阶段缺失get_num_kv_splits调用的潜在bug，回退后该bug可能重现。但当前未报告相关故障。
  2. 代码腐烂风险：回退后的代码与主分支上其他可能依赖#26134的PR存在冲突风险，合并前需确保兼容。
  3. 无测试覆盖：本次变更未新增测试，依赖已有回归测试，可能遗漏边界问题。
     - 影响：对用户：无直接影响，CUDA graph功能保持正常。
     对系统：注意力后端的CUDA graph元数据初始化逻辑恢复到重构前的独立实现模式，每个后端各自维护类似代码，可维护性降低但正确性更易保证。
     对团队：后续若再次统一需从头设计，且需小心处理本次发现的边界条件。

• 风险标记：核心路径变更, review发现潜在bug, 回退丢失部分修复, 缺乏新增测试覆盖

关联脉络

• PR #26134 [refactor] unify cuda-graph capture/replay across attention backends: 本PR revert的目标，引入统一重构但导致正确性问题。
• PR #26152 fix(swa): eliminate spurious translate_loc_from_full_to_swa warning in BCG and CG paths: 第二个commit重新应用其SWA修复，确保revert后滑动窗口功能正常。