# PR #32415 完整报告

- 仓库：`sgl-project/sglang`
- 标题：[VLM] Split multimodal scheduling from mm_utils
- 合并时间：2026-08-05 20:24
- 原文链接：http://prhub.com.cn/sgl-project/sglang/pull/32415

---

# 执行摘要

- 一句话：多模态调度逻辑从 mm_utils 拆分为独立模块，纯重构
- 推荐动作：值得快速浏览而非精读：关注 mm_schedule.py 的模块边界划分、mm_utils.py 的兼容 re-export 手法，以及测试中猴补丁隔离 CUDA-IPC ack 的做法。建议后续跟踪 mm_schedule.py 上陆续出现的调度优化 PR，并留意仓库是否会在后续清理 mm_utils 的 re-export。

# 功能与动机

mm_utils.py 已增长到 2,090 行，混合了通用多模态辅助函数与 cross-request ViT batching 引入的延迟敏感调度路径。PR body 明确说明动机：“给该路径一个独立模块可以使后续调度优化和 review 更聚焦，而不改变运行时行为。” 这是为后续调度优化铺路的模块边界重构。

# 实现拆解

1. 新建 `python/sglang/srt/managers/mm_schedule.py`（+658 行）：迁入 16 个调度相关定义，包括模块级 `embedding_cache`、`init_mm_embedding_cache`、`get_embedding_chunk`、`_get_precomputed_embedding`、`_flatten_embedding_result`、`_can_skip_pre_embed_feature_move`、`_move_items_to_device`、`_acknowledge_deferred_cuda_ipc_cache_hits`、`_get_chunked_embedding_full`、`_get_chunked_embedding_by_item`、`_batch_encode_per_image_misses`、`PerImageRequestInfo`、`get_embedding_and_mask`、`DataEmbeddingFunc` 等。模块只保留调度路径所需依赖（`schedule_batch`、`multimodal_cache`、`runtime_context`、`utils`），并用 docstring 明确职责为“Multimodal embedding scheduling and cache coordination”。
2. 精简 `python/sglang/srt/managers/mm_utils.py`（-654 行）：删除上述定义及不再使用的 import（如 `dataclass/field`、`MultimodalDataItem`、`EmbeddingResult`、`MultiModalStaticCache`、`EVSEmbeddingResult`、`get_parallel`、`get_schedule`、`is_hip`、`is_npu`），保留 `pad_input_tokens`、`init_feature_buffer`、`embed_mm_inputs`、`TensorTransportMode` 等通用工具。兼容性处理是文件头部新增 `from mm_schedule import (init_mm_embedding_cache as init_mm_embedding_cache)` 以及 `DataEmbeddingFunc`、`get_embedding_and_mask` 的 re-export，保证外部调用 `mm_utils.init_mm_embedding_cache(...)` 仍可用；但只 re-export 了 3 个符号，旧代码若直接引用 `mm_utils._get_chunked_embedding_full`、`mm_utils.embedding_cache` 等会断裂。
3. 更新缓存初始化调用点 `python/sglang/srt/mem_cache/kv_cache_builder.py`：单行 import 从 `mm_utils` 改为 `mm_schedule`，确保 KV cache builder 初始化时引用新模块持有的 `embedding_cache` 全局对象。
4. 测试配套 `test/registered/chunked_prefill/test_mm_chunked_embedding_unit.py`：所有用例从 `mm_utils` 切换到 `mm_schedule`，并新增 autouse fixture `_skip_cuda_ipc_acknowledgement`，用 monkeypatch 把 `_acknowledge_deferred_cuda_ipc_cache_hits` 替换为空操作。原因是 CPU-only 测试不初始化 tensor-parallel 运行时状态，而该函数会调用 `get_parallel()`，在无并行状态的 CPU 环境下会失败；PR body 指出当前 origin/main 上该测试同样需要这个隔离才能通过。
5. 验证手段：16 个被迁移定义与 origin/main 做 AST 等价性检查；pre-commit、compileall、`git diff --check` 通过；远程 NVIDIA devbox 上运行该测试 5 passed。未附带性能数据，因为本 PR 只修改模块归属，不影响运行时行为。

关键文件：
- `python/sglang/srt/managers/mm_schedule.py`（模块 调度模块；类别 source；类型 module-creation；符号 init_mm_embedding_cache, get_embedding_chunk, _get_precomputed_embedding, _flatten_embedding_result）: 本 PR 的核心产物：新建模块迁入全部 16 个多模态调度定义，是后续跨请求 ViT 批量编码等调度优化的主战场，模块边界的合理与否直接决定重构价值。
- `python/sglang/srt/managers/mm_utils.py`（模块 通用工具；类别 source；类型 compat-re-export；符号 init_mm_embedding_cache, get_embedding_chunk, _get_precomputed_embedding, _flatten_embedding_result）: 删除 654 行并保留兼容 re-export，决定外部引用是否断裂；是本次重构的风险集中点。
- `test/registered/chunked_prefill/test_mm_chunked_embedding_unit.py`（模块 单元测试；类别 test；类型 test-coverage；符号 _skip_cuda_ipc_acknowledgement）: 测试全部切换到 mm_schedule，并新增 autouse fixture 隔离 CUDA-IPC ack 路径，修复 CPU-only 环境下对 TP 运行时状态的依赖。
- `python/sglang/srt/mem_cache/kv_cache_builder.py`（模块 缓存构建；类别 source；类型 dependency-wiring）: KV cache 构建入口的 cache 初始化 import 从 mm_utils 切换到 mm_schedule，验证了对外接口重指向。

关键符号：init_mm_embedding_cache, get_embedding_chunk, _get_precomputed_embedding, _flatten_embedding_result, _can_skip_pre_embed_feature_move, _move_items_to_device, _acknowledge_deferred_cuda_ipc_cache_hits, _get_chunked_embedding_full, _get_chunked_embedding_by_item, _batch_encode_per_image_misses, get_embedding_and_mask

## 关键源码片段

### `python/sglang/srt/managers/mm_schedule.py`

本 PR 的核心产物：新建模块迁入全部 16 个多模态调度定义，是后续跨请求 ViT 批量编码等调度优化的主战场，模块边界的合理与否直接决定重构价值。

```python
"""Multimodal embedding scheduling and cache coordination.

本模块由 mm_utils.py 拆分而来：集中放置多模态 embedding 缓存协调与
chunked-prefill 的延迟敏感调度路径（跨请求 ViT 批量编码、chunk 组装与对齐）。
本 PR 仅做模块级代码移动，16 个定义已与 origin/main 做 AST 等价校验。
"""

from dataclasses import dataclass, field
from typing import Callable, Dict, List, Optional, Tuple

import torch

from sglang.srt.managers.schedule_batch import MultimodalDataItem
from sglang.srt.mem_cache.multimodal_cache import EmbeddingResult, MultiModalStaticCache
from sglang.srt.multimodal.evs import EVSEmbeddingResult
from sglang.srt.runtime_context import get_parallel, get_schedule
from sglang.srt.utils import is_hip, is_npu
from sglang.utils import logger

_is_hip = is_hip()
_is_npu = is_npu()

# 模块级 embedding 缓存；原来是 mm_utils.embedding_cache，搬迁后全局对象
# 归 mm_schedule 所有，旧模块通过 re-export 保持兼容
embedding_cache: Optional[MultiModalStaticCache] = None


def init_mm_embedding_cache(max_size: int = 0):
    global embedding_cache
    embedding_cache = MultiModalStaticCache(max_size)


def _acknowledge_deferred_cuda_ipc_cache_hits(items: List[MultimodalDataItem]) -> None:
    """缓存命中跳过 ViT 时，释放惰性的 Kimi IPC 切片。

    encoder-DP miss 时只有一个 rank 拷贝图像并对整个 TP 组确认；
    缓存命中时没有任何 rank 拷贝，因此由 rank 0 执行等价的单次确认，
    以维持固定池的生命周期，同时避免引入不必要的 GPU-to-GPU 拷贝。
    """
    parallel = get_parallel()
    if parallel.attn_tp_rank != 0:
        return
    # 池的回收器按整个 TP 组计数，确认次数必须匹配，即使注意力子组更小
    consumer_count = max(parallel.tp_size, 1)
    for item in items:
        item.acknowledge_deferred_cuda_ipc_feature(consumer_count)

```

### `python/sglang/srt/managers/mm_utils.py`

删除 654 行并保留兼容 re-export，决定外部引用是否断裂；是本次重构的风险集中点。

```python
# 保留既有初始化 import，供下游调用者继续使用（兼容性重导出）。
# 注意：只有 DataEmbeddingFunc、get_embedding_and_mask、init_mm_embedding_cache
# 三个符号被 re-export；若外部代码仍从 mm_utils 引用 _get_chunked_embedding_full
# 或 embedding_cache，会直接 ImportError 或拿到 None。
from sglang.srt.managers.mm_schedule import (
    DataEmbeddingFunc,
    get_embedding_and_mask,
)
from sglang.srt.managers.mm_schedule import (
    init_mm_embedding_cache as init_mm_embedding_cache,
)

```

# 评论区精华

本 PR 没有任何 review 评论；两条 issue 评论分别是 gemini-code-assist bot 的停用通知和作者自己的 /tag-and-rerun-ci 命令。核心“讨论”来自 PR body 的三点设计声明：一是用 AST 等价校验保证纯移动的无行为变化；二是刻意保留 `mm_utils.init_mm_embedding_cache` 兼容导入；三是 CPU 测试通过 monkeypatch 隔离 CUDA-IPC ack，并指出该测试在当前 origin/main 上同样会失败，属于对测试改动的主动说明。

- 暂无高价值评论线程

# 风险与影响

- 风险：兼容性风险（主要）：mm_utils 只 re-export 了 `DataEmbeddingFunc`、`get_embedding_and_mask`、`init_mm_embedding_cache` 三个符号；仓库内若有其他模块直接 `from mm_utils import _get_chunked_embedding_full / embedding_cache / get_embedding_chunk`，会立即 ImportError 或拿到 None。PR 只改了 4 个文件，未提供全仓引用扫描证据，CI 是主要防线。
全局对象迁移风险：`embedding_cache` 的所有权从 mm_utils 转到 mm_schedule，任何仍从 mm_utils 读取该全局对象的代码都不会再看到缓存实例，可能表现为缓存命中率下降或 AttributeError。
测试盲区：monkeypatch 屏蔽了缓存命中时的 CUDA-IPC ack 生命周期逻辑，CPU 单测无法覆盖该路径；这里仍依赖 GPU e2e 测试保障，回归后 CPU 测试不会暴露。
运行时与性能风险：按 PR 声明（AST 等价 + 无性能数字），运行时行为与性能均无变化。

- 影响：对用户与系统：无运行时行为、无性能影响，属于纯代码组织结构变更。
对代码结构：mm_utils.py 从约 2,090 行缩减（删除 654 行），新的 mm_schedule.py 成为多模态调度路径的唯一归属模块，chunked-prefill 分类、跨请求 ViT 批量编码、embedding 对齐等后续优化将集中在此进行。
对团队：调度相关 PR 的 diff 和 review 上下文更聚焦，模块边界更清晰；但一次性移动大量代码也可能让“旧符号引用”问题转移到下游 PR 中解决。
对测试：CPU 单测通过 monkeypatch 摆脱 TP 运行时依赖，稳定性和独立运行能力提升。

- 风险标记：纯移动重构，AST 等价校验通过 , mm_utils 仅 re-export 三个符号 , embedding_cache 全局对象迁移 , CUDA-IPC ack 逻辑被测试 monkeypatch 屏蔽

# 关联脉络

- PR #32895 fix mm-chunk-embedding test suite: 与本 PR 修改同一测试文件 test_mm_chunked_embedding_unit.py，且同属 CPU CI 稳定性方向；本 PR 新增的 _skip_cuda_ipc_acknowledgement fixture 是该方向的延续。
- PR #30683 [Diffusion] Batch GLM-Image AR requests: 多模态请求批量调度与动态批入队方向的演进背景，与本 PR 将调度路径独立成模块的动机吻合；后续调度优化预计集中在 mm_schedule.py 上展开。