执行摘要

• 一句话：MoRIIO 新增逐层读 barrier，修复 READ 模式高并发解码精度退化
• 推荐动作：值得精读：该 PR 展示了异步 RDMA 完成与 CUDA 图捕获之间的真实矛盾，以及维护者在正确性与性能之间权衡后的最终取舍（尊重用户配置 + 警告）。关注点包括 per-layer 状态设计、屏障超时策略，以及 requires_piecewise_for_cudagraph override 从引入到移除的演进过程。建议跟进 issue #49643 的 FULL 图方案。

功能与动机

PR body 明确说明：wait_for_layer_load() was a no-op (pass), so the attention kernel for a layer could run before that layer's remote KV had actually landed. Under a full CUDA graph the missing barrier is silently skipped on replay, so attention races the reads and produces garbage ("salad") decode output that gets worse with concurrency. 验证数据从 0.4405 提升到 0.9644，说明该缺口直接影响高并发 P/D 场景的正确性。

实现拆解

实现拆解：

1. 数据结构升级：MoRIIOConnectorWorker._recving_transfers 从 defaultdict[ReqId, list] 改为 defaultdict[ReqId, dict]（layer_name -> status），对应修改写入点 _read_blocks 和消费点 _pop_done_transfers，使屏障可以精确等待某层而非整个请求的末尾状态。
2. 屏障实现：MoRIIOConnector.wait_for_layer_load 从 pass 改为委托 MoRIIOConnectorWorker.wait_for_layer_load；worker 内使用带超时（moriio_config.transfer_timeout）的轮询循环，逐层收集所有请求的在途状态，Succeeded 和 Failed 均视为终态，失败处理仍交给 get_finished / _pop_done_transfers 通知 prefill，超时仅告警并继续，避免单个卡住传输拖垮 EngineCore。
3. CUDA 图适配与权衡演进：最初通过 requires_piecewise_for_cudagraph 返回 True 强制所有模式使用 PIECEWISE；随后（rkhan055）改为仅 READ 模式强制；因 Minimax-M3 TP<->TP 吞吐回归，引入 VLLM_MORIIO_FORCE_PIECEWISE 环境变量；最终 review 反馈移除环境变量和 override，改为在 __init__ 中对 READ + consumer + FULL 图组合打印一次性警告，并在 wait_for_layer_load 内检测 cudagraph_runtime_mode == FULL 时直接返回。
4. 测试配套：更新 tests/v1/kv_connector/unit/test_moriio_tp_ack.py 中 _recving_transfers fixture 为 {"req": {"layer0": DoneStatus()}} 形状，新增 test_requested_cudagraph_mode_is_never_overridden 固化“不覆盖用户配置”的行为。

关键文件：

• vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py（模块 KV 连接器；类别 source；类型 core-logic；符号 wait_for_layer_load, _pop_done_transfers, _read_blocks）: 核心修复：将 _recving_transfers 改为 per-layer dict，实现带超时的 wait_for_layer_load 屏障，并在 FULL 图模式下直接返回，同时增加 READ+consumer+FULL 组合的一次性警告。
• tests/v1/kv_connector/unit/test_moriio_tp_ack.py（模块 单元测试；类别 test；类型 test-coverage；符号 test_requested_cudagraph_mode_is_never_overridden）: 测试配套：更新 _recving_transfers fixture 以匹配 per-layer dict 结构，并新增 cudagraph 模式不被覆盖的回归测试。

关键符号：MoRIIOConnector.wait_for_layer_load, MoRIIOConnectorWorker.wait_for_layer_load, MoRIIOConnectorWorker._pop_done_transfers, MoRIIOConnectorWorker._read_blocks, MoRIIOConnector.init, test_requested_cudagraph_mode_is_never_overridden

关键源码片段

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py

核心修复：将 _recving_transfers 改为 per-layer dict，实现带超时的 wait_for_layer_load 屏障，并在 FULL 图模式下直接返回，同时增加 READ+consumer+FULL 组合的一次性警告。

# MoRIIOConnectorWorker —— READ 模式逐层读完成屏障
​
# READ 模式下解码节点对每个请求逐层发起异步 RDMA 读，状态表从 list
# 升级为 per-layer dict，便于精确等待即将执行 attention 的那一层。
self._recving_transfers: defaultdict[ReqId, dict] = defaultdict(dict)
​
def wait_for_layer_load(self, layer_name: str) -> None:
    """阻塞直到本层所有 in-flight READ 到达终态，避免 attention 读到未落地的远端 KV。"""
    # Producer 或非 READ 模式没有需要等待的读操作；FULL 图重放时无法执行
    # 宿主侧阻塞等待，直接返回（正确性由用户显式配置 PIECEWISE 保证）。
    if self.is_producer or self.mode != MoRIIOMode.READ:
        return
    if get_forward_context().cudagraph_runtime_mode == CUDAGraphMode.FULL:
        return
​
    deadline = time.monotonic() + self.moriio_config.transfer_timeout
    while True:
        with self.moriio_wrapper.lock:
            # 汇总所有请求中针对该层仍在途的传输状态
            pending = [
                status_by_layer[layer_name]
                for status_by_layer in self._recving_transfers.values()
                if layer_name in status_by_layer
            ]
            if not pending:
                return
            # 每个状态到达 Succeeded / Failed 即为终态，未终态则继续阻塞
            if not any(not s.Succeeded() and not s.Failed() for s in pending):
                return
        # 超时只告警并继续，避免单个卡住传输拖垮 EngineCore；
        # 失败读最终由 get_finished / _pop_done_transfers 清理并通知 prefill。
        if time.monotonic() > deadline:
            logger.warning(
                "MoRIIO READ barrier timed out for layer %s; proceeding "
                "(request dropped via get_finished).",
                layer_name,
            )
            return
        time.sleep(0.001)

# MoRIIOConnector.__init__ —— FULL 图与 READ 模式组合时只告警不覆盖
if (
    self.mode == MoRIIOMode.READ
    and self.kv_transfer_config.is_kv_consumer
    and vllm_config.compilation_config.cudagraph_mode.has_full_cudagraphs()
):
    # warn only；屏障无法在 FULL 图内触发，是否改 PIECEWISE 由用户决定
    logger.warning_once(
        "MoRIIO READ mode is running with %s CUDA graphs: per-layer "
        "KV-read barrier can't fire inside full graph; accuracy may "
        "degrade at high concurrency. Set cudagraph_mode=PIECEWISE "
        "in --compilation-config.",
        vllm_config.compilation_config.cudagraph_mode.name,
    )

tests/v1/kv_connector/unit/test_moriio_tp_ack.py

测试配套：更新 _recving_transfers fixture 以匹配 per-layer dict 结构，并新增 cudagraph 模式不被覆盖的回归测试。

# 保证连接器从不默默覆盖用户配置的 cudagraph 模式：
# READ 模式 + PIECEWISE 时屏障正常触发；READ 模式 + FULL 图时
# 仅打印警告，是否切换到 PIECEWISE 由 --compilation-config 决定。
def test_requested_cudagraph_mode_is_never_overridden():
    assert (
        MoRIIOConnector.requires_piecewise_for_cudagraph({"read_mode": True}) is False
    )
    assert (
        MoRIIOConnector.requires_piecewise_for_cudagraph({"read_mode": False}) is False
    )

评论区精华

核心讨论围绕“正确性 vs 性能”展开：

• Rohan138 质疑 requires_piecewise_for_cudagraph 是否对 WRITE 模式也适用，rkhan055 回复已在 6096858 改为按 get_moriio_mode(extra_config) 只对 READ 模式返回 True，避免 WRITE 模式被无谓降级。
• tjtanaa 报告 TP<->TP 的 minimax m3 输出吞吐量回归 -65.53%，因为改用 PIECEWISE 禁用了 fullgraph；edwinlim0919 回应应先合入正确性修复、再通过 RFC #49643 实现 FULL 图安全的 barrier，并反问是否验证过高并发精度。
• tanpinsiang 提供双节点实测：MiniMax-M3 精度无差异但输出吞吐从 4830 降到 4038 tok/s、TPOT 从 28.97ms 升到 88.19ms；用 DeepSeek-R1-0528 复测确认精度从 93.78% 提升到 96.06%，支持屏障的正确性价值。

• tanpinsiang 还指出测试 fixture 与新的 dict 结构不匹配（1 failed 28 passed），给出最小修复 diff，已合入。

• requires_piecewise_for_cudagraph 是否应覆盖所有模式 (design): WRITE 模式没有 per-layer 屏障（wait_for_layer_load 立即返回），保留 FULL 图，避免无谓降级。

• wait_for_layer_load 在 FULL 模式下是否应直接返回 (correctness): 在 cudagraph_runtime_mode == FULL 时直接返回，避免阻塞等待被写入 CUDA 图。
• Minimax-M3 TP<->TP 输出吞吐回归 -65.53% (performance): 正确性修复优先合入，FULL 图安全的 barrier 作为后续工作。
• 测试 fixture 与新 _recving_transfers 结构不匹配 (testing): 测试已更新为 {"req": {"layer0": DoneStatus()}} 形状并保持通过。
• VLLM_MORIIO_FORCE_PIECEWISE 环境变量的引入与移除 (design): 默认尊重配置：仅当 cudagraph_mode=PIECEWISE 时屏障触发；FULL 图下只打印警告，由用户在 --compilation-config 中决策。

风险与影响

• 风险：主要风险如下：
  • 性能风险：PIECEWISE 下每层执行宿主侧轮询等待（1ms sleep），decode TPOT 可能显著上升（实测 28.97ms→88.19ms），对 TP<->TP 场景影响明显。
  • 正确性风险：FULL 图 + READ 模式默认不触发 barrier，仅启动时打印一次警告；用户若忽略警告，高并发下仍会出现精度退化。
  • 稳定性风险：超时后继续执行可能读到部分到达的数据；失败读不会立即终止请求，而是留给 get_finished 在后续迭代清理，期间可能产出坏 token。
  • 测试覆盖不足：仅单元测试且基于 FakeWrapper，无法覆盖真实 CQ 线程时序和分布式 P/D 端到端行为。
  • 兼容性风险：_recving_transfers 类型从 list 改为 dict，任何直接访问该字段的外部代码都会受影响。
• 影响：影响范围集中在 MoRIIO 连接器的 READ 模式用户：
  • 对正确性：高并发 P/D 解码的随机精度退化被消除，DeepSeek-R1-0528 的 gsm8k 指标从 0.44 量级恢复到 0.96 量级。
  • 对性能：默认尊重用户配置，只有显式选择 PIECEWISE 才会付出屏障的开销；WRITE 模式和 producer 侧完全不受影响。
  • 对团队：为后续 RFC #49643（FULL 图安全的 KV-read barrier）打下基础，也为 kv-connector 子系统其他连接器的异步语义提供参考。
  • 风险标记：PIECEWISE 性能回退, FULL 图下 barrier 失效, 轮询阻塞 decode 热路径, 缺分布式端到端测试

关联脉络

• PR #50234 [PD][PushConnector] Record last activity of remotes to allow clean up of stale ones: 同为 KV connector v1 的远端状态管理演进，反映该子系统的异步语义日趋成熟。
• PR #49644 [Feat][Core] Add disk offloading support to SimpleCPUOffloadConnector: 同属 KV connector v1 领域，说明该模块在稳定性和容量方面的持续演进背景。
• PR #51100 [Bugfix] Fix Mamba all-mode CPU offload boundary alignment: 同为 KV connector 子系统的 bugfix，体现跨连接器一致的状态管理问题。