执行摘要

• 一句话：TRTLLM draft extend 使用 decode kernel
• 推荐动作：值得合并。改动量小且逻辑直观，只需确认 is_draft_extend_v2 枚举定义正确且与调度器行为一致。建议后续添加针对该分支的回归测试。

功能与动机

对于 speculative decoding 中的 draft extend 阶段，其序列长度通常较短（如 5 个 token），使用为长序列优化的 prefill kernel 存在冗余计算。改用 decode kernel 可减少不必要的计算开销，提升推理性能。

实现拆解

1. 修改注意力后端控制流：在 python/sglang/srt/layers/attention/trtllm_mha_backend.py 的 forward_extend 方法中，将选择 decode kernel 的条件从 is_target_verify() 扩展为 is_target_verify() or is_draft_extend_v2()。这样当 forward mode 为 draft extend v2 时，也会调用 flashinfer.decode.trtllm_batch_decode_with_kv_cache 而非 flashinfer.prefill.trtllm_batch_context_with_kv_cache。
2. 增加缓存刷新重试机制：在 python/sglang/test/bench_one_batch_server_internal.py 中新增 _flush_cache_with_retry 函数，对 flush_cache 或 reset_prefix_cache 请求最多重试 3 次，每次失败后等待 2 秒。将原 run_one_case 中直接的请求调用替换为此函数，提升基准测试的鲁棒性。
3. 修复 isort lint 问题：第二个 commit 调整了导入顺序以通过 CI lint 检查。

关键文件：

• python/sglang/srt/layers/attention/trtllm_mha_backend.py（模块 注意力后端；类别 source；类型 core-logic；符号 forward_extend）: 核心改动：在 forward_extend 中为 draft extend v2 模式启用 decode kernel 路径。通过简单添加 or 条件使 decode kernel 覆盖目标验证和草稿扩展两个模式。
• python/sglang/test/bench_one_batch_server_internal.py（模块 基准测试；类别 test；类型 test-coverage；符号 _flush_cache_with_retry, run_one_case）: 新增 _flush_cache_with_retry 函数，为缓存刷新操作添加重试逻辑，提升 benchmark 稳定性。

关键符号：forward_extend, _flush_cache_with_retry

关键源码片段

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

核心改动：在 forward_extend 中为 draft extend v2 模式启用 decode kernel 路径。通过简单添加 or 条件使 decode kernel 覆盖目标验证和草稿扩展两个模式。

# 文件 : python/sglang/srt/layers/attention/trtllm_mha_backend.py
# 在 forward_extend 方法中，选择 kernel 的分支逻辑
​
page_table = self._get_layer_page_table(layer, forward_batch)
​
# 关键变更：原来只对 is_target_verify 使用 decode kernel，
# 现在也对 is_draft_extend_v2 使用 decode kernel
if (
    forward_batch.forward_mode.is_target_verify()
    or forward_batch.forward_mode.is_draft_extend_v2()
):
    # 使用 decode kernel（针对短序列优化）
    o = flashinfer.decode.trtllm_batch_decode_with_kv_cache(
        query=q,
        kv_cache=kv_cache,
        workspace_buffer=self.workspace_buffer,
        block_tables=page_table,
        seq_lens=self.forward_metadata.cache_seqlens_int32,
        max_seq_len=self.max_context_len,
        bmm1_scale=bmm1_scale,
        bmm2_scale=bmm2_scale,
        window_left=layer.sliding_window_size,
        sinks=attention_sink,
        skip_softmax_threshold_scale_factor=envs.SGLANG_SKIP_SOFTMAX_DECODE_THRESHOLD_SCALE_FACTOR.get(),
        out_dtype=self.q_data_type,
        q_len_per_req=self.forward_metadata.max_seq_len_q,
    )
else:
    # 使用 prefill kernel（针对长序列优化）
    o = flashinfer.prefill.trtllm_batch_context_with_kv_cache(
        query=q,
        kv_cache=kv_cache,
        workspace_buffer=self.workspace_buffer,
        block_tables=page_table,
        seq_lens=self.forward_metadata.cache_seqlens_int32,
        max_q_len=self.forward_metadata.max_seq_len_q,
        max_kv_len=self.max_context_len,
        bmm1_scale=bmm1_scale,
        bmm2_scale=bmm2_scale,
        batch_size=self.forward_metadata.cu_seqlens_q.shape[0] - 1,
        cum_seq_lens_q=self.forward_metadata.cu_seqlens_q,
        cum_seq_lens_kv=self.forward_metadata.cu_seqlens_k,
        window_left=layer.sliding_window_size,
        sinks=attention_sink,
    )

python/sglang/test/bench_one_batch_server_internal.py

新增 _flush_cache_with_retry 函数，为缓存刷新操作添加重试逻辑，提升 benchmark 稳定性。

# 文件 : python/sglang/test/bench_one_batch_server_internal.py
# 新增函数：带重试的缓存刷新，最多重试 3 次，每次间隔 2 秒
​
def _flush_cache_with_retry(url: str, endpoint: str, max_retries: int = 3):
    """Post to a cache flush endpoint with retries on failure."""
    for attempt in range(max_retries):
        response = requests.post(url + endpoint, timeout=DEFAULT_TIMEOUT)
        if response.status_code == 200:
            return
        if attempt < max_retries - 1:
            time.sleep(2) # 等待 2 秒后重试
        else:
            response.raise_for_status() # 最后一次重试失败则抛出异常
​
# 在 run_one_case 中替换原本的直接请求调用
# 原来 :
# response = requests.post(url + "/flush_cache", timeout=DEFAULT_TIMEOUT)
# response.raise_for_status()
# 现在 :
# _flush_cache_with_retry(url, "/flush_cache")

评论区精华

该 PR 没有显著的 review 讨论。标签中添加了 blackwell 但正文未提及具体硬件优化细节，推测该优化可能与 Blackwell 架构的 TensorRT-LLM 后端相关。

• 暂无高价值评论线程

风险与影响

• 风险：

  1. 功能回归风险：改用 decode kernel 可能改变 draft extend 阶段的注意力计算行为，需确保输出精度与原逻辑一致。尽管改动很小（只加了一个或条件），但 is_draft_extend_v2 的正确性依赖上游调度逻辑。
  2. 性能退化风险：如果 draft extend 的序列长度并非总是很短，在某些场景下使用 decode kernel 可能反而效率更低。不过从常见 speculative decoding 实现看，draft token 数量通常较少，该风险较低。
  3. 测试覆盖不足：没有新增针对 is_draft_extend_v2 分支的单元测试或集成测试。现有 CI 可能只覆盖标准 forward 路径，该分支缺乏验证。
     - 影响：影响范围：仅影响使用 tensorrtllm_mha_backend 后端的 speculative decoding 场景（forward_mode = draft_extend_v2）。其他后端或 forward mode 无影响。
     影响程度：修改极小（4 行核心代码），但若 draft extend 是频繁调用的路径（尤其是在 speculative decoding 中），性能提升可能显著。基准测试的健壮性改进对所有 batch benchmark 用户有益。

• 风险标记：缺少测试覆盖

关联脉络

• 暂无明显关联 PR