Prhub

#52041 [Core] Skip broadcasting mm tensor data to workers for prefix-cache-covered items

原始 PR 作者 sseanliu 合并时间 2026-08-19 22:16 文件变更 4 提交数 8 评论 15 代码增减 +130 / -2

执行摘要

跳过 prefix-cache 覆盖项的 mm 张量广播,多模态 TTFT 显著下降

Issue #52040 指出:当多模态请求的图片 token 被 prefix cache 完全覆盖时,encoder 不会运行,但 EngineCore 仍会在每次请求中把这些 15.5MB 的 pixel_values 重新序列化并广播给所有 TP worker,导致 TTFT 随唯一图片数线性增长(约 19ms/图,26 张图时约 634ms)。PR body 强调这是 streaming/agentic 多模态场景(滚动窗口图片)的核心瓶颈,目标是让全缓存 TTFT 从线性变为平坦。

值得精读。该 PR 展示了一个典型的「基于前缀缓存状态做数据传输裁剪」的设计模式:通过 num_computed_tokens 判断数据是否必然不会被消费,从而在广播前剥离。重点关注 strip_covered_mm_data 的 M-RoPE 特判(keep_on_cpu)以及调度器如何传递 uses_mrope。同时建议结合 PR#52827 阅读,理解 keep_on_cpu 字段语义。对于 v1 runner 的 preemption 场景,若有后续演进可关注是否需要恢复该路径。

讨论亮点
  1. resume 路径裁剪:作者最初实现了 preemption resume 时的 re-ship 机制(CachedRequestData.resumed_mm_features),但 njhill 指出「we don't need the resume part of this though since that does not apply to the v2 model runner, and we don't intend to make these kinds of optimizations to the v1 model runner at this point」,最终 commit 83108c6 删除该路径,仅保留 admission-time 剥离。
  2. M-RoPE 元数据保护:CI 在 Qwen2/3-VL、Gemma、LLaVA、Whisper 等 M-RoPE 模型上失败,原因是 worker 的 _init_mrope_positions 需要从 feature.data 读取 image_grid_thw/video_grid_thw 计算位置。njhill 指出「we still need to transfer certain of the tensors corresponding to cached positions (for mrope in particular)」,随后在 commit 4430aa6 中补充 keep_on_cpu 字段保留逻辑,并依赖 PR#52827 的过滤能力。
  3. 设计取舍讨论:PR body 提到一种更大胆的方案——将 mm tensor 传输完全推迟到 scheduled_encoder_inputs 阶段(仅在 encoder 运行时发送),但作者选择最小 diff 路径,先做 admission-time 剥离。njhill 认可当前方案,DarkLight1337 表示 logic 合理并请 ywang96 确认。

实现拆解

  1. 新增剥离函数:在 vllm/multimodal/utils.py 添加 strip_covered_mm_data(mm_features, num_computed_tokens, uses_mrope=False)。它遍历每个 MultiModalFeatureSpec,若 data 非空且 mm_position.offset + mm_position.length <= num_computed_tokens(span 完全落在已计算 prefix 内),则用 dataclasses.replacedata 置为 None;若 uses_mrope=True,则保留所有 field.keep_on_cpu=True 的字段(如 image_grid_thw),其余 payload 字段剥离。函数不修改原列表,num_computed_tokens=0 时直接返回原列表。
  2. 调度输出接线:在 vllm/v1/core/sched/output.pyNewRequestData.from_request 中新增 uses_mrope 参数,并调用 strip_covered_mm_data 处理 request.mm_features,替换原来的直接赋值。这样从 EngineCore 下发到 worker 的 NewRequestData 中,被覆盖项的 data 变为 None,广播时序列化成本大幅下降。
  3. 调度器传递 mrope 标记:在 vllm/v1/core/sched/scheduler.pyScheduler.__init__ 中读取 model_config.uses_mrope 存入 self.model_uses_mrope,并在 schedule() 中构造 NewRequestData 时(v1 和 v2 runner 分支)将该值传给 from_request,保证 worker 对 M-RoPE 模型能拿到位置计算所需的元数据。
  4. 测试覆盖:在 tests/v1/core/test_output.py 新增 test_strip_covered_mm_data(覆盖完全覆盖、边界、未覆盖、已 None、不修改原列表)、test_strip_covered_mm_data_zero_computed(零计算不剥离)、test_strip_covered_mm_data_mrope(M-RoPE 场景保留 keep_on_cpu 字段)。测试同时验证非数据字段(identifiermm_position)在剥离后保持不变,且原列表不被修改。
文件 模块 状态 重要度
vllm/multimodal/utils.py 多模态工具 modified 7.48
vllm/v1/core/sched/output.py 调度输出 modified 6.34
tests/v1/core/test_output.py 测试 modified 7.09

关键符号

strip_covered_mm_data maybe_strip NewRequestData.from_request Scheduler.schedule

关键源码片段

vllm/multimodal/utils.py core-logic

新增核心函数 strip_covered_mm_data,决定哪些多模态项的 tensor 数据可被剥离,并处理 M-RoPE 的 keep_on_cpu 字段保留逻辑,是整个优化的核心。

# vllm/multimodal/utils.py ( 新增函数,核心剥离逻辑 )def strip_covered_mm_data(
    mm_features: list[MultiModalFeatureSpec],
    num_computed_tokens: int,
    uses_mrope: bool = False,
) -> list[MultiModalFeatureSpec]:
    """丢弃 placeholder span 完全位于 prefix-cache 已计算区域内的多模态项的
    tensor 数据。此类项永远不会被调度 encoder 运行,因此 worker 不会消费
    payload 字段。M-RoPE 模型例外:worker 需要 CPU 侧元数据字段(如
    image_grid_thw)来计算整个 prompt 的位置,因此保留这些字段。
    scheduler 侧的 Request 仍保留完整 features。"""
​
    # 没有特征或没有前缀命中时,直接返回原列表(不复制)
    if not mm_features or num_computed_tokens == 0:
        return mm_features
​
    def maybe_strip(f: MultiModalFeatureSpec) -> MultiModalFeatureSpec:
        # data 已为 None 或 span 超出已计算区域(encoder 仍可能运行)
        # 则不做处理,保持原样
        if f.data is None or (
            f.mm_position.offset + f.mm_position.length > num_computed_tokens
        ):
            return f
​
        # 默认全部剥离;对 M-RoPE 模型保留 keep_on_cpu 字段
        #(位置计算需要在 CPU 侧读取的元数据)
        data = None
        if uses_mrope:
            data = MultiModalKwargsItem(
                {k: elem for k, elem in f.data.items() if elem.field.keep_on_cpu}
            )
        # 用 dataclasses.replace 生成新对象,保持原列表不被修改
        return replace(f, data=data)
​
    return [maybe_strip(f) for f in mm_features]
vllm/v1/core/sched/output.py dependency-wiring

修改 NewRequestData.from_request 接入了 strip_covered_mm_data,使 EngineCore 下发给 worker 的数据在构造时就完成剥离,这是广播减量的实际生效点。

# vllm/v1/core/sched/output.py (NewRequestData.from_request 修改 )@classmethod
def from_request(
    cls,
    request: Request,
    block_ids: tuple[list[int], ...],
    prefill_token_ids: list[int] | None = None,
    uses_mrope: bool = False,
) -> "NewRequestData":
    """从调度器侧 Request 构造要发送给 worker 的 NewRequestData。
    这里会根据当前 num_computed_tokens 剥离被 prefix cache 完全覆盖的
    多模态项的 tensor 数据,从而避免 EngineCore 向 TP worker 广播这些
    永远不会被消费的大张量(如 15.5MB 的 pixel_values)。
    M-RoPE 模型需要额外的元数据字段(keep_on_cpu)来计算位置,
    因此通过 uses_mrope 参数保留这些字段。"""
    return cls(
        req_id=request.request_id,
        prompt_token_ids=request.prompt_token_ids,
        # 关键改动:对 mm_features 做覆盖剥离,而非直接使用 request.mm_features
        mm_features=strip_covered_mm_data(
            request.mm_features,
            request.num_computed_tokens,
            uses_mrope=uses_mrope,
        ),
        sampling_params=request.sampling_params,
        pooling_params=request.pooling_params,
        block_ids=block_ids,
        num_computed_tokens=request.num_computed_tokens,
        lora_request=request.lora_request,
        prompt_embeds=request.prompt_embeds,
        prompt_is_token_ids=request.prompt_is_token_ids,
        prefill_token_ids=prefill_token_ids,
    )
tests/v1/core/test_output.py test-coverage

新增的单元测试覆盖了剥离函数的边界条件(完全覆盖、边界、未覆盖、已 None、零计算、M-RoPE 字段保留),保证优化逻辑的正确性与非破坏性。

# tests/v1/core/test_output.py ( 新增测试,验证剥离逻辑的边界条件 )def _mm_feature(offset: int, length: int) -> MultiModalFeatureSpec:
    # 构造一个仅含 dummy 数据的 MultiModalFeatureSpec,便于测试
    return MultiModalFeatureSpec(
        data=MultiModalKwargsItem.dummy(),
        mm_position=PlaceholderRange(offset=offset, length=length),
        identifier=f"hash_{offset}",
        modality="image",
    )
​
​
def test_strip_covered_mm_data() -> None:
    """验证:完全在已计算前缀内的项被剥离;span 恰好结束于
    num_computed_tokens 的边界项也视为覆盖并剥离;超出前缀的项保留;
    原本 data 为 None 的项保持 None;非 data 字段不受影响;
    原列表不被修改。"""
    from dataclasses import replace
​
    covered = _mm_feature(offset=0, length=100)
    boundary = _mm_feature(offset=150, length=100) # 结束位置恰为 250
    uncovered = _mm_feature(offset=300, length=100)
    already_none = replace(_mm_feature(offset=100, length=50), data=None)
​
    stripped = strip_covered_mm_data(
        [covered, boundary, uncovered, already_none], num_computed_tokens=250
    )
​
    assert stripped[0].data is None # 完全覆盖 -> 剥离
    assert stripped[1].data is None # span 结束 == computed -> 算覆盖
    assert stripped[2].data is not None # 超出前缀 -> 保留
    assert stripped[3].data is None # 原本就是 None
    # 非 data 字段保持不变
    assert stripped[0].identifier == covered.identifier
    assert stripped[0].mm_position == covered.mm_position
    # 原列表未被修改
    assert covered.data is not None
​
​
def test_strip_covered_mm_data_zero_computed() -> None:
    """没有前缀命中时,任何项都不应被剥离。"""
    features = [_mm_feature(offset=0, length=100)]
    stripped = strip_covered_mm_data(features, num_computed_tokens=0)
    assert stripped[0].data is not None

评论区精华

是否保留 preemption 恢复时的 re-ship 机制 设计

作者最初实现了 preemption resume 时的重新传输路径(CachedRequestData.resumed_mm_features),但 njhill 认为该机制不适用于 v2 model runner,且 v1 model runner 不打算做此类优化,建议删除。

结论:删除 resume 路径,仅保留 admission-time 剥离(commit 83108c6)。 · 已解决

M-RoPE 模型需要保留 keep_on_cpu 元数据字段 正确性

CI 在 Qwen2/3-VL 等 M-RoPE 模型上失败,原因是被剥离后 worker 的 _init_mrope_positions 无法获取 image_grid_thw 等字段。njhill 指出必须保留这些字段,并推送到分支修复。

结论:在 strip_covered_mm_data 中增加 uses_mrope 分支,保留 field.keep_on_cpu=True 的字段;依赖 PR#52827 的 keep_on_cpu 标记能力。 · 已解决

是否需要更大范围的重构(将 mm tensor 传输推迟到 encoder 调度时) 设计

PR body 提到另一种设计:在 scheduled_encoder_inputs 阶段再发送 tensor,但作者选择最小 diff 路径,仅做 admission-time 剥离。njhill 认可当前方案。

结论:维持最小 diff 方案,后续如需更彻底优化可另行开展。 · 已解决

风险与影响

  1. M-RoPE 回归风险:如果 uses_mrope 判断不准确或 keep_on_cpu 字段被误剥离(如 PR#52827 未合入前),worker 计算位置时会缺少 image_grid_thw 等字段,导致位置错误或崩溃。本 PR 的测试覆盖了该场景,但真实模型验证仍依赖 CI。
  2. v1 model runner 的预取恢复:虽然代码路径对 v1 runner 同样生效(else 分支也传了 uses_mrope),但 v1 runner 从 preemption 恢复时是否会重新构造 NewRequestData 并重新剥离,文档未明确。若恢复时 num_computed_tokens 因 eviction 变小而不再覆盖,data 仍然为 None,则 worker 可能拿不到数据。作者原本的 resume 路径被删除,对 v1 runner 该风险依然存在。
  3. 边界条件offset + length == num_computed_tokens 被判定为覆盖并剥离,若实际调度中该 span 最后一个 token 尚未计算(边界含混),可能导致 encoder 缺失输入。测试中 boundary 样例覆盖了相等情况,但语义上需确认 prefix cache 的包含关系。

对全缓存或高命中率的多模态工作负载(streaming 视觉助手、视频 agent、多轮视觉记忆)影响显著,实测全缓存 TTFT 从 634ms 降至接近 2 张图的水平(约 54ms),单图平均广播成本从 ~21ms 降至 ~2.9ms。对未命中 prefix cache 的普通请求无行为变化(num_computed_tokens=0 或 span 超出时不剥离)。对 M-RoPE 模型有额外保护,但需要 PR#52827 配合。影响范围覆盖所有 vLLM v1 引擎的多模态请求,涉及调度器、输出数据结构与多模态工具函数,但改动集中在 4 个文件,风险可控。

多模态核心路径变更 M-RoPE 模型需特殊处理 v1 runner preemption 路径未覆盖 依赖未合入 PR#52827

关联 Issue

#52040 [Performance]: MM input tensors are re-broadcast to all TP workers on every request, even when the prefix cache fully covers the image tokens (~19 ms per unique image)

完整报告

参与讨论