# PR #22809 完整报告

- 仓库：`sgl-project/sglang`
- 标题：Dual MoE CUDA graph capture for lora/nolora batches
- 合并时间：2026-04-23 05:11
- 原文链接：http://prhub.com.cn/sgl-project/sglang/pull/22809

---

# 执行摘要

- 一句话：为 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 图捕获的状态管理、图键生成和变体解析逻辑，是性能优化的主要入口。

```python
# 新增全局变量和函数来管理捕获变体
_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 时生效。

```python
# 新增全局配置变量
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 计算。

```python
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 图捕获同属核心路径性能调优。