执行摘要

• 一句话：简化 TRTLLM MLA draft extend 的 CUDA graph metadata 填充
• 推荐动作：可合并。变更简洁，经过 reviewer 确认，逻辑正确。建议后续关注是否有类似冗余字段存在于其他 backend，可做统一清理。

功能与动机

优化 draft extend 路径的 CUDA graph metadata 填充性能，移除冗余的 metadata 分配和复制操作，降低 GPU 开销。

实现拆解

1. _init_cuda_graph_metadata: draft extend 分支中，将 num_tokens_per_bs = num_tokens // bs 改为 num_tokens_per_bs = self.num_draft_tokens，统一使用类成员变量而非参数计算，与 _apply_cuda_graph_metadata 保持一致。
2. _apply_cuda_graph_metadata: 移除 draft extend 分支中对 cu_seqlens_q 和 seq_lens_q 的分配与复制代码（共 10 行），因为这两个字段在 draft extend 的 CUDA graph 执行中实际未被使用（仅 decode/target-verify 需要）。
3. seq_lens_k 的赋值简化：target verify 分支中去掉了显式 .to(dtype=torch.int32)，依赖 PyTorch 隐式类型转换；draft extend 分支中删除旧的复杂 seq_lens 计算（seq_lens[:bs] - metadata.seq_lens_q[:bs] + metadata.max_seq_len_q），直接使用 seq_lens[:bs] 赋值，因为移除 seq_lens_q 后该计算不再需要。

关键文件：

• python/sglang/srt/layers/attention/trtllm_mla_backend.py（模块 注意力层；类别 source；类型 core-logic；符号 _init_cuda_graph_metadata, _apply_cuda_graph_metadata）: 核心变更文件：简化 draft extend 的 CUDA graph metadata 填充，移除未使用的 cu_seqlens_q/seq_lens_q 分配与复制，统一 num_tokens_per_bs 计算。

关键符号：_init_cuda_graph_metadata, _apply_cuda_graph_metadata

关键源码片段

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

核心变更文件：简化 draft extend 的 CUDA graph metadata 填充，移除未使用的 cu_seqlens_q/seq_lens_q 分配与复制，统一 num_tokens_per_bs 计算。

def _init_cuda_graph_metadata(self, bs, num_tokens, forward_mode, seq_lens, device):
    """Allocate persistent metadata buffers for CUDA graph capture."""
    metadata = TRTLLMMLADecodeMetadata()
    if forward_mode.is_target_verify():
        metadata.seq_lens_k = torch.zeros((bs,), dtype=torch.int32, device=device)
    elif forward_mode.is_draft_extend_v2():
        # 统一使用 self.num_draft_tokens 而非 num_tokens // bs
        num_tokens_per_bs = self.num_draft_tokens
        metadata.max_seq_len_q = num_tokens_per_bs
        metadata.sum_seq_lens_q = num_tokens_per_bs * bs
        metadata.cu_seqlens_q = torch.arange(0, bs * num_tokens_per_bs + 1,
            num_tokens_per_bs, dtype=torch.int32, device=device)
        metadata.seq_lens_q = torch.full((bs,), num_tokens_per_bs, dtype=torch.int32, device=device)
        metadata.seq_lens_k = torch.zeros((bs,), dtype=torch.int32, device=device)
    # ... 后续 block_kv_indices 设置
​
​
def _apply_cuda_graph_metadata(self, bs, req_pool_indices, seq_lens, forward_mode):
    """Shared decode / target-verify / draft-extend capture+replay body."""
    metadata = self.decode_cuda_graph_metadata[bs]
    if forward_mode.is_target_verify():
        seq_lens = seq_lens[:bs] + self.num_draft_tokens
        metadata.seq_lens_k.copy_(seq_lens) # 移除显式 .to(int32)，依赖隐式转换
    elif forward_mode.is_draft_extend_v2():
        num_tokens_per_bs = self.num_draft_tokens
        metadata.max_seq_len_q = num_tokens_per_bs
        metadata.sum_seq_lens_q = num_tokens_per_bs * bs
        # 移除 cu_seqlens_q / seq_lens_q 分配与复制（它们在此路径未使用）
        seq_lens = seq_lens[:bs]
        metadata.seq_lens_k.copy_(seq_lens)
    # 更新 block_kv_indices（Triton kernel）
    create_flashmla_kv_indices_triton[(bs, ...)](...)

评论区精华

1. 类型转换问题 (@kpham-sgl): 询问 seq_lens_k 是否需要 int32，因为 forward_batch.seq_lens 是 int64。作者回复 seq_lens_k buffer 是 int32，复制时自动完成类型转换，效果相同。
2. num_tokens_per_bs 计算正确性 (@kpham-sgl): 质疑 _init_cuda_graph_metadata 中改为 self.num_draft_tokens 是否影响 DP attention padding。作者回复 _apply_cuda_graph_metadata 原本就使用 self.num_draft_tokens，因此保持一致且无影响。

• seq_lens_k 类型转换 (correctness): 作者回复 seq_lens_k buffer 是 int32，复制时自动转换，效果相同。
• num_tokens_per_bs 计算统一 (correctness): 作者回复 _apply_cuda_graph_metadata 原本就使用 self.num_draft_tokens，因此一致且无影响。

风险与影响

• 风险：低风险。变更集中在 draft extend 的 CUDA graph metadata 填充路径，且去除了未使用的字段，不会影响 decode 和 target-verify 路径。 seq_lens_k 类型转换依赖隐式 cast，与显式 .to(int32) 行为一致（truncation），无精度风险。
• 影响：对用户无影响，仅内部性能优化。移除冗余 GPU 内存分配和 kernel 启动，可能轻微降低 draft extend 的延迟。影响范围仅限于 TRTLLM MLA backend + speculative decoding 场景。
• 风险标记：暂无

关联脉络

• PR #29078 [perf] tiny optimize select_index op for draft extend: 同为 speculative decoding 的 draft extend 性能优化，关注点相近。
• PR #29200 [Cookbook] Nemotron3-Ultra: align MTP draft depth with NVIDIA reference (num_steps 5): 关联 speculative decoding 的 MTP 参数对齐。