Prhub

#35536 [Model Runner V2][Bugfix] Fix MRV2 LoRA warmup

原始 PR 作者 jeejeelee 合并时间 2026-06-16 04:17 文件变更 5 提交数 31 评论 27 代码增减 +227 / -63

执行摘要

修复 MRV2 LoRA warmup 及 CUDA Graph 集成

PR Body 明确指出需要修复 LoRA warmup 并完成 LoRA Verification。在 V2 Model Runner 中,LoRA 的预热(warmup)和 CUDA Graph 捕获步骤未能正确处理活跃 LoRA 数量,导致 dummy run 时可能设置错误的适配器,影响推理正确性与性能。此 PR 旨在通过新的工具函数和集成点,使 LoRA 状态在 warmup 和 capture 阶段被正确追踪。

  • 值得精读:特别是 _build_lora_dispatch_map 的预计算策略,可以推广到其他维度(如 token 数量、batch 大小)的 dispatch 优化。
  • 关注后续重构LoRAModelRunnerMixin 的重构计划值得跟踪,可能进一步简化 LoRA 集成。
讨论亮点
  • 参数类型错误(critical)gemini-code-assist[bot] 发现 _dummy_run 中调用 maybe_dummy_run_with_loranum_scheduled_tokens 被错误写为 num_scheduled_tokens == np.array(...),导致传递的是布尔值而非数组。作者立即修正并表示感谢。
  • 性能优化建议WoosukKwon 建议放弃运行时 bisect 而预计算 dispatch 映射,作者实现了 _build_lora_dispatch_map 使 dispatch 变为 O(1) 查表。
  • 设计取舍讨论WoosukKwon 指出 create_lora_capture_hook 中通过 maybe_select_dummy_loras 设置 dummy LoRA 的方式偏向 'hack',作者承认并计划在未来 PR 中重构 LoRAModelRunnerMixin

实现拆解

步骤 1:提取 LoRA 工具函数
vllm/v1/worker/gpu/lora_utils.py 中新增 get_lora_capture_cases(计算 CUDA Graph 捕获所需的活跃 LoRA 数量列表)、get_num_active_loras_for_dispatch(根据当前请求计算实际活跃 LoRA 数)、create_lora_capture_hook(创建在每次捕获前设置 dummy LoRA 的回调)三个顶层函数,并将 LoraState.make_lora_inputs 中的内联循环抽取为 get_activate_loras 方法,便于复用。

步骤 2:扩展 CUDA Graph 描述符
vllm/v1/worker/gpu/cudagraph_utils.py 中为 BatchExecutionDescriptor 添加 num_active_loras 字段,修改 _is_compatible 使其参与匹配;在 CudaGraphManager 中新增 lora_capture_cases 参数,并实现 _build_lora_dispatch_map 预计算从实际活跃 LoRA 数量到捕获 case 的映射,将 dispatch 从二分查找降级为字典查表。

步骤 3:集成到 Model Runner
vllm/v1/worker/gpu/model_runner.py 中初始化时调用 get_lora_capture_cases 计算捕获案例;在 _dummy_run 中使用 maybe_dummy_run_with_lora context manager 包裹模型执行;在 capture_model 时传入 lora_capture_hook;在 execute_model 中调用 get_num_active_loras_for_dispatch 获取当前批次的活跃 LoRA 数并传递给 dispatch 函数。

步骤 4:适配数据并行
vllm/v1/worker/gpu/dp_utils.py 中为 dispatch_cg_and_sync_dpsync_cudagraph_and_dp_padding 添加 num_active_loras 参数,在构造 BatchExecutionDescriptor 时透传该值,确保 DP 各 rank 的 LoRA 信息一致(但不需要跨 rank 同步)。

步骤 5:测试配套
tests/lora/test_qwen3_with_multi_loras.py 中添加 set_mrv2_env fixture 强制设置 VLLM_USE_V2_MODEL_RUNNER=1,并移除所有测试用例的 enforce_eager=True,使得测试在启用 CUDA Graph 的情况下运行,验证 warmup 修复。

文件 模块 状态 重要度
vllm/v1/worker/gpu/lora_utils.py LoRA 工具 modified 8.34
vllm/v1/worker/gpu/cudagraph_utils.py CG 管理 modified 7.97
vllm/v1/worker/gpu/model_runner.py 模型执行器 modified 7.42
vllm/v1/worker/gpu/dp_utils.py 数据并行 modified 5.89
tests/lora/test_qwen3_with_multi_loras.py 测试 modified 5.43

关键符号

get_lora_capture_cases get_num_active_loras_for_dispatch create_lora_capture_hook LoraState.get_activate_loras _build_lora_dispatch_map _is_compatible dispatch_cg_and_sync_dp sync_cudagraph_and_dp_padding

关键源码片段

vllm/v1/worker/gpu/cudagraph_utils.py core-logic

核心变更:BatchExecutionDescriptor 增加 num_active_loras 字段,CudaGraphManager 集成 lora_capture_cases 并实现预计算分发表。

@dataclass(frozen=True)
class BatchExecutionDescriptor:
    cg_mode: CUDAGraphMode
    num_tokens: int
    num_reqs: int | None
    uniform_token_count: int | None = None
    num_active_loras: int = 0 # 新增字段,用于 LoRA 感知的匹配# ... existing code ...class CudaGraphManager:
    def __init__(
        self,
        vllm_config: VllmConfig,
        device: torch.device,
        cudagraph_mode: CUDAGraphMode,
        decode_query_len: int,
        lora_capture_cases: list[int] | None = None, # 新增参数
    ):
        # ... 原有初始化 ...
        self.lora_capture_cases = lora_capture_cases or [0]
        # 预计算实际 active LoRA 数量到 captured case 的映射
        self._lora_dispatch_map, self._max_lora_case = self._build_lora_dispatch_map()
​
    def _build_lora_dispatch_map(self) -> tuple[dict[int, int], int]:
        """Precompute actual num_active_loras -> effective captured case.        Mirrors the num_tokens candidate expansion: every possible active-LoRA
        count is mapped ahead of time to the smallest captured case that can
        serve it, so dispatch is a plain dict lookup instead of a per-call
        bisect.
        """
        captured_with_lora = sorted(c for c in self.lora_capture_cases if c > 0)
        if not captured_with_lora:
            return {}, 0
        dispatch_map: dict[int, int] = {}
        case_idx = 0
        for n in range(1, captured_with_lora[-1] + 1):
            while captured_with_lora[case_idx] < n:
                case_idx += 1
            dispatch_map[n] = captured_with_lora[case_idx]
        return dispatch_map, captured_with_lora[-1]

评论区精华

参数传递错误:num_scheduled_tokens 错用比较运算符 正确性

gemini-code-assist[bot] 在 review 中发现 `_dummy_run` 中 `maybe_dummy_run_with_lora` 的 `num_scheduled_tokens` 参数被写作 `num_scheduled_tokens == np.array(num_tokens_per_request, dtype=np.int32)`,导致传递布尔值而非数组,会引发 TypeError。

结论:作者立即采纳建议并修正。 · 已解决

预计算 dispatch 映射以避免二分查找 性能

WoosukKwon 建议放弃运行时 `bisect` 而预计算 dispatch 映射,将 dispatch 从二分查找优化为字典查表。作者实现 `_build_lora_dispatch_map`。

结论:已实现并合并。 · 已解决

create_lora_capture_hook 的 hack 性质 设计

WoosukKwon 指出 `hook` 实现依赖 `maybe_select_dummy_loras` 是一种 hack,建议未来重构 `LoRAModelRunnerMixin`。作者同意并将后续处理。

结论:未解决,推迟至后续 PR。 · 待处理

风险与影响

  • 核心路径变更风险model_runner.pycudagraph_utils.py 是推理执行的核心,任何逻辑错误都可能导致推理崩溃或静默错误。本次新增的 LoRA 状态耦合需要仔细验证。
  • 配置依赖风险lora_capture_cases 的计算依赖 compilation_config.cudagraph_specialize_loralora_config.max_loras 等配置,若配置体系未来变化可能引入不一致。
  • DP 一致性假设num_active_loras 不跨 DP rank 同步,假设各 rank 的 LoRA 请求分布相同,若分布式部署时请求分布不均可能导致 dispatch 不匹配。
  • 测试覆盖局限:仅测试了 Qwen3 模型,其他模型(如 DeepSeek、LLaMA)在 MRV2 + LoRA 场景下可能存在未覆盖的问题。
  • 用户影响:启用 VLLM_USE_V2_MODEL_RUNNER=1 并使用 LoRA 的用户将得到正确的 CUDA Graph warmup,首次推理延迟可能降低,吞吐量提升。未使用 LoRA 或 V1 Runner 的用户不受影响。
  • 系统影响:增加了内存和计算开销(预计算 dispatch 映射),但总量很小。代码复杂度略有上升,但功能模块化程度更高。
  • 团队影响:此 PR 弥合了 V2 Model Runner 的一个重要功能缺口,推动其生产化进程。create_lora_capture_hook 的设计可能为后续重构提供参考。
核心路径变更 配置依赖 DP 一致性假设 测试覆盖不足

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论