执行摘要

• 一句话：为 MoE LoRA 启用双 CUDA 图捕获，无适配器批次使用无 LoRA 图以提升解码吞吐量。
• 推荐动作：该 PR 值得精读，特别是 cuda_graph_runner.py 中的变体状态管理和图键设计，展示了如何在 CUDA 图捕获中实现条件分支优化。关注 _resolve_lora_variant 如何根据批次动态选择变体，以及 lora_moe_runners.py 中捕获时的内核跳过逻辑，这些是性能提升的关键。同时，注意 review 中提到的对齐计算优化点，可作为进一步性能调优的切入点。

功能与动机

根据 PR body，当 LoRA 启用时，原有的单 CUDA 图会记录 LoRA 内核（即使适配器未启用，内核也通过早期退出实现零开销），但这仍对无 LoRA 请求造成性能惩罚。通过捕获双图，无适配器批次可以完全跳过 LoRA 内核，从而提升吞吐量，基准测试显示解码吞吐量从 31,445 tok/s 提升至 40,740 tok/s（提升 30%）。

实现拆解

1. 添加命令行配置：在 server_args.py 中新增 --record-nolora-graph 标志（默认 True），用于控制是否启用双图捕获。
2. 配置验证与访问：在 moe/utils.py 中新增 RECORD_NOLORA_GRAPH 全局变量和 should_record_nolora_graph() 函数，根据 enable_lora 和 MoE 后端（仅验证 triton 后端）决定是否启用双图捕获，并添加警告日志。
3. 核心图捕获逻辑：在 cuda_graph_runner.py 中引入 _capture_lora_variant 状态变量和 get_capture_lora_variant()、_set_capture_lora_variant() 函数来跟踪捕获变体；新增 _default_make_graph_key() 和 _make_graph_key() 用于生成带变体标签的图键；在 CudaGraphRunner 类中添加 record_nolora_graph 属性和 _resolve_lora_variant() 方法，根据批次是否有活跃 LoRA 适配器返回 "lora" 或 "nolora" 标签，从而在捕获和回放时路由到对应图。
4. LoRA 内核跳过：在 lora_moe_runners.py 的 _add_lora_gate_up_delta()、_add_lora_down_delta() 和 build_lora_hooks() 函数中，在捕获模式下检查 get_capture_lora_variant()，若为 "nolora" 则跳过 LoRA 内核记录或钩子构建，避免无谓计算。
5. 测试与部署配套：本次改动未包含直接测试文件变更，但 PR body 提及了正确性验证（MMLU 准确性和 logprobs 差异），确保功能可靠。

关键文件：

• python/sglang/srt/model_executor/cuda_graph_runner.py（模块 图执行器；类别 source；类型 core-logic；符号 get_capture_lora_variant, _set_capture_lora_variant, _default_make_graph_key, _make_graph_key）: 核心变更文件，实现了双 CUDA 图捕获的状态管理、图键生成和变体解析逻辑，是性能优化的主要入口。
• python/sglang/srt/layers/moe/utils.py（模块 MoE层；类别 source；类型 configuration；符号 should_record_nolora_graph）: 配置管理文件，新增双图捕获的启用逻辑和验证，确保仅在对 triton MoE 后端且启用 LoRA 时生效。
• python/sglang/srt/lora/lora_moe_runners.py（模块 LoRA模块；类别 source；类型 dependency-wiring）: LoRA 内核实现文件，修改了在捕获模式下根据变体跳过 LoRA 内核的逻辑，确保 nolora 图中不包含 LoRA 计算。
• python/sglang/srt/server_args.py（模块 服务器参数；类别 source；类型 entrypoint）: 命令行接口文件，新增 --record-nolora-graph 标志，允许用户控制双图捕获的启用。

关键符号：get_capture_lora_variant, _set_capture_lora_variant, _default_make_graph_key, _make_graph_key, _resolve_lora_variant, should_record_nolora_graph

关键源码片段

python/sglang/srt/model_executor/cuda_graph_runner.py

核心变更文件，实现了双 CUDA 图捕获的状态管理、图键生成和变体解析逻辑，是性能优化的主要入口。

# 新增全局变量和函数来管理捕获变体
_capture_lora_variant: Optional[str] = None # 跟踪当前捕获的变体：None（非双捕获）、"lora"（捕获含 LoRA 内核的图）、"nolora"（捕获无 LoRA 内核的图）
​
def get_capture_lora_variant() -> Optional[str]:
    """返回当前正在捕获的 LoRA 变体，若非双捕获模式则返回 None。"""
    return _capture_lora_variant
​
def _set_capture_lora_variant(variant: Optional[str]):
    """设置捕获变体，用于在捕获上下文中切换。"""
    global _capture_lora_variant
    _capture_lora_variant = variant
​
# 独立的图键生成函数，便于子类复用
def _default_make_graph_key(bs, stream_idx=None, variant_label=None):
    """根据批次大小、流索引和 LoRA 变体标签构建图字典键。"""
    key = bs if stream_idx is None else f"{stream_idx}_{bs}"
    if variant_label is not None:
        key = f"{variant_label}_{key}" # 添加变体前缀，如 "lora_" 或 "nolora_"
    return key
​
class CudaGraphRunner:
    # ... 其他初始化代码 ...
    def _make_graph_key(self, bs, stream_idx=None, variant_label=None):
        """包装 _default_make_graph_key 以供实例使用。"""
        return _default_make_graph_key(bs, stream_idx, variant_label)
​
    def _resolve_lora_variant(self, forward_batch: ForwardBatch):
        """根据批次中是否有活跃 LoRA 适配器返回变体标签，若未启用双捕获则返回 None。"""
        if not getattr(self, "record_nolora_graph", False):
            return None # 双捕获未启用
        if forward_batch.lora_ids is not None and any(uid is not None for uid in forward_batch.lora_ids):
            return "lora" # 批次包含活跃 LoRA 适配器
        return "nolora" # 批次无活跃 LoRA 适配器

python/sglang/srt/layers/moe/utils.py

配置管理文件，新增双图捕获的启用逻辑和验证，确保仅在对 triton MoE 后端且启用 LoRA 时生效。

# 新增全局配置变量
RECORD_NOLORA_GRAPH: bool = False # 控制是否启用双图捕获
​
def initialize_moe_config(server_args: ServerArgs):
    global RECORD_NOLORA_GRAPH
    # 仅当启用 LoRA 且 MoE 后端为 triton 时才启用双图捕获
    _triton_ok = MOE_RUNNER_BACKEND in (MoeRunnerBackend.TRITON, MoeRunnerBackend.TRITON_KERNELS)
    if bool(server_args.record_nolora_graph) and bool(server_args.enable_lora) and not _triton_ok:
        logger.warning(f"record_nolora_graph only validated for triton MoE backend, but moe_runner_backend={server_args.moe_runner_backend}. Disabling.")
    RECORD_NOLORA_GRAPH = bool(server_args.record_nolora_graph) and bool(server_args.enable_lora) and _triton_ok
​
def should_record_nolora_graph() -> bool:
    """返回是否应记录无 LoRA 图，供其他模块查询。"""
    return RECORD_NOLORA_GRAPH

python/sglang/srt/lora/lora_moe_runners.py

LoRA 内核实现文件，修改了在捕获模式下根据变体跳过 LoRA 内核的逻辑，确保 nolora 图中不包含 LoRA 计算。

def _add_lora_gate_up_delta(...):
    # ... 参数定义 ...
    if get_is_capture_mode():
        from sglang.srt.model_executor.cuda_graph_runner import get_capture_lora_variant
        # 仅在捕获 lora 变体时记录 LoRA 内核，nolora 变体跳过
        has_active_lora = get_capture_lora_variant() != "nolora"
    else:
        # 非捕获模式下根据适配器状态判断
        num_loras = len(lora_info.lora_ranks)
        has_active_lora = (lora_info.adapter_enabled[:num_loras] * (lora_info.lora_ranks > 0).to(lora_info.adapter_enabled.dtype)).any().item()
    if not has_active_lora or lora_info is None or lora_info.max_lora_rank == 0:
        return # 跳过 LoRA 计算
    # ... 后续 LoRA 内核逻辑 ...
​
def build_lora_hooks(...):
    if lora_info is None or lora_info.max_lora_rank == 0:
        return LoRAHooks()
    if get_is_capture_mode():
        from sglang.srt.model_executor.cuda_graph_runner import get_capture_lora_variant
        if get_capture_lora_variant() == "nolora":
            return LoRAHooks() # 捕获 nolora 变体时直接返回空钩子，避免对齐计算
    # ... 后续对齐计算和钩子构建 ...

评论区精华

review 中仅有一条来自 chatgpt-codex-connector[bot] 的评论，指出当前实现在捕获 nolora 变体时仍会计算 LoRA 对齐（如 _compute_lora_alignment），导致 nolora 图中仍包含 LoRA 路由工作，速度提升可能不完整。评论建议在 build_lora_hooks() 中更早地门控钩子构建或对齐计算，以进一步优化性能。此评论未被回复或解决，可能作为未来改进点。

• Bypass LoRA alignment when capturing the nolora variant (performance): 未在 PR 中解决，可能作为未来优化点；建议在 build_lora_hooks 中更早门控钩子构建以完全跳过对齐计算。

风险与影响

• 风险：
  1. 回归风险：双图捕获逻辑增加了代码复杂度，若 _resolve_lora_variant 或图键生成出错，可能导致错误图被使用，影响输出正确性；但 PR body 中的正确性验证（MMLU 准确性、logprobs 差异）降低了此风险。
  2. 性能风险：双图捕获会占用额外 GPU 内存存储两组图，可能在小内存设备上导致内存压力；同时，仅验证 triton MoE 后端，若在其他后端误启用可能导致未定义行为或性能下降。
  3. 兼容性风险：record_nolora_graph 默认启用，但仅当 enable_lora=True 且 MoE 后端为 triton 时才实际生效，其他后端会通过警告禁用，这确保了向后兼容性。
  4. 安全风险：无直接安全影响，但新增的命令行参数需确保文档更新。
• 影响：
  1. 用户影响：用户可通过 --record-nolora-graph 标志（默认开启）获得无 LoRA 请求的性能提升，解码吞吐量提升 30%，延迟降低，无需修改现有配置；但需注意仅适用于 triton MoE 后端。
  2. 系统影响：CUDA 图管理更复杂，内存使用略有增加，但整体系统吞吐量显著提升，尤其在高并发混合批次场景下。
  3. 团队影响：引入了双图捕获的设计模式，为未来类似优化（如其他变体图）提供参考；但 review 中未解决的性能疑虑需后续关注。
     - 风险标记：核心路径变更, 内存使用增加, 后端兼容性限制, 性能优化不完整

关联脉络

• PR #23178 [LoRA] Fix EP + per-expert MoE LoRA illegal memory access: 同样涉及 MoE 和 LoRA 的交互，修复内存访问问题，与本 PR 的性能优化互补。
• PR #22685 [CPU] [Quantization] Add GPTQ/AWQ 4bits quantization support for CPU: 涉及性能优化和内核支持，展示了仓库对多后端性能改进的关注，与本 PR 的 CUDA 图优化类似。
• PR #23327 Skip mamba_pool_idx revert for session requests in _get_new_batch_prefill_raw: 涉及调度器和缓存管理优化，与本 PR 的 CUDA 图捕获同属核心路径性能调优。