全模态模型的解耦推理

本指南将介绍如何在vllm-omni中为多阶段流水线配置和使用分布式连接器（vllm_omni/distributed/connectors）。

1. 概述

连接器支持在流水线阶段之间传输数据（例如，Thinker -> Talker）。当前支持的连接器以D2H2D（设备到主机到设备）模式运行：1. SharedMemoryConnector：使用系统共享内存。2. MooncakeConnector：使用Mooncake。

• SharedMemoryConnector（默认）：零拷贝（在主机上），延迟最低。最适合单节点部署。如果没有指定连接器，则自动配置。
• MooncakeConnector：基于TCP/RDMA。最适合多节点分布式部署。需要Mooncake Master服务。

2. API设计

连接器系统围绕OmniConnectorBase抽象构建，该抽象将数据传输与阶段逻辑解耦。

核心接口

class OmniConnectorBase(ABC):
    @abstractmethod
    def put(self, from_stage: str, to_stage: str, request_id: str, data: Any) -> tuple[bool, int, Optional[dict]]:
        """
        Store data.
        Returns: (success, serialized_size, metadata)
        """
        pass

    @abstractmethod
    def get(self, from_stage: str, to_stage: str, request_id: str, metadata: Optional[dict] = None) -> Optional[tuple[Any, int]]:
        """
        Retrieve data.
        Args: metadata - Transport-specific handles returned by put() (e.g., SHM name).
        Returns: (object, serialized_size)
        """
        pass

关键概念：元数据传递

与纯键值存储不同，某些连接器（如SharedMemoryConnector）在put()期间会生成临时资源（例如，共享内存块名称）。此metadata必须通过控制平面（例如，HTTP头，队列消息）从生产者阶段传递到消费者阶段，以便get()能够找到数据。

3. 后端与用例

3.1 SharedMemoryConnector

最适合：单节点，高性能IPC。

• 机制
  • 小负载（<阈值）：数据被序列化并直接“内联”在metadata字典中传递。这避免了为微小消息创建SHM块的开销。
  • 大负载（>=阈值）：数据被写入一个命名的系统共享内存块。块名称在metadata中返回。
• 配置
  • shm_threshold_bytes：从内联切换到SHM的字节大小（默认：64KB）。

3.2 MooncakeConnector

最适合：多节点分布式推理。

• 机制：使用Mooncake的分布式KVCache存储。
  • 数据平面：TCP或RDMA用于高带宽传输。
  • 控制平面：使用集中的Mooncake Master和Metadata Server。
  • 键控：基于request_id/from_stage_to_stage的确定性键。
• 要求：需要运行Mooncake Master服务。

4. 与vLLM的关系

vLLM为特定工件提供了专门的分布式机制：* KV传输（vllm.distributed.kv_transfer）：优化用于在预填充和解码实例之间传输KV缓存（使用NCCL，Mooncake等）。* EC传输（vllm.distributed.ec_transfer）：优化用于共享编码器嵌入。* 设备通信器（vllm.distributed.device_communicators）：用于张量/流水线并行化的低级原语（NCCL，SHM）。

vllm-omni通过为多模态流水线引入通用连接器抽象（OmniConnector）来补充这一点。虽然vLLM的连接器是特定于工件的，但vllm-omni

1. 统一传输：提供单一API（put/get）来传输任何阶段的工件（输入嵌入，隐藏状态，音频/图像张量，KV缓存，最终输出）在任意流水线阶段之间（例如，AudioEncoder -> LLM -> AudioGenerator）。
2. 扩展连接性：支持灵活构建复杂的DAG（有向无环图），其中阶段可以在同一进程，同一节点或跨节点运行，为每个边使用最合适的后端（SHM，Mooncake等）。
3. 封装和适配：可以内部利用vLLM的专用kv_transfer进行KV路径，同时为其他数据类型使用通用传输（SHM/Mooncake），为应用程序层提供一致的接口。

5. 安装（Mooncake）

如果使用MooncakeConnector，请先安装库

# For CUDA-enabled systems (Recommended)
pip install mooncake-transfer-engine

# For non-CUDA systems
pip install mooncake-transfer-engine-non-cuda

6. 使用MooncakeConnector

6.1 启动Mooncake Master

在您的主节点上启动Master服务

# if you use mooncake SSD storage
mkdir -p ./mc_storage

mooncake_master \
  --rpc_port=50051 \
  --enable_http_metadata_server=true \
  --http_metadata_server_host=0.0.0.0 \
  --http_metadata_server_port=8080 \
  --metrics_port=9003 \
  --root_fs_dir=./mc_storage/ \
  --cluster_id=mc-local-1 &

6.2 配置（YAML）

编辑您的阶段配置（例如，qwen2_5_omni.yaml）。

步骤1：在全局运行时中定义连接器

runtime:
  connectors:
    connector_of_mooncake:
      name: MooncakeConnector
      extra:
        host: "127.0.0.1"           # Local Worker IP
        metadata_server: "http://<MASTER_IP>:8080/metadata"
        master: "<MASTER_IP>:50051"
        segment: 512000000          # 512MB segment
        localbuf: 64000000          # 64MB buffer
        proto: "tcp"                # "tcp" or "rdma"
    ```

**Mooncake Configuration Parameters:**

*   **host**: The hostname or IP address of the local machine (worker). Mooncake uses this to register itself in the metadata server so other nodes can find it.
*   **metadata_server**: The URL of the metadata server. This is used for service discovery and connection establishment (e.g., exchanging QP information for RDMA).
*   **master**: The address of the Mooncake Master Server (e.g., `<MASTER_IP>:50051`). This is used for global state management and control plane operations.
*   **segment**: The size of the global memory segment in bytes (default: ~512MB). This defines the shared memory region accessible by Mooncake for data transfer.
*   **localbuf**: The size of the local buffer in bytes (default: ~64MB). Used for local data buffering during transfer operations.
*   **proto**: The transport protocol to use. Options:
    *   `tcp`: Standard TCP/IP (easier setup, universal compatibility).
    *   `rdma`: Remote Direct Memory Access (higher performance, requires RDMA-capable hardware).

For more details, refer to the [Mooncake Repository](https://github.com/kvcache-ai/Mooncake).

    **Step 2: Reference in Stages**

Explicitly link stages using `input_connectors` and `output_connectors`:

```yaml
stage_args:
  - stage_id: 0
    # ...
    output_connectors:
      to_stage_1: connector_of_mooncake

  - stage_id: 1
    # ...
    input_connectors:
      from_stage_0: connector_of_mooncake

7. 使用SharedMemoryConnector（自动模式）

最适合单节点。

系统将自动为任何没有显式定义连接器的流水线边创建SharedMemoryConnector。这是从以下项推断的：1. 配置中的runtime.edges列表。2. stage_args中定义的engine_input_source依赖项。

阈值配置

默认情况下，大于64KB（默认阈值）的负载通过共享内存传输，而较小的负载则使用控制队列（内联）。

要调整此阈值（例如，调整到1GB），请在您的runtime.connectors中添加以下内容

runtime:
  connectors:
    connector_of_shared_memory:
      name: SharedMemoryConnector
      extra:
        shm_threshold_bytes: 1024 # 1KB threshold

8. 总结

Use Case	推荐的连接器	配置
单节点	SharedMemoryConnector	无（自动）或自定义阈值
多节点	MooncakeConnector	显式YAML + Mooncake Master

9. 操作说明（重要）

• 快速失败的配置验证：如果任何预期的边缺少连接器，加载器将抛出错误。定义input_connectors/output_connectors或依赖于自动SHM填充；否则启动将中止。
• 丢失的负载会使阶段停止：工作节点期望连接器负载；如果元数据或连接器配置丢失，阶段将引发错误并停止。在生产之前，请验证连接器接线和元数据传播。

10. 未来路线图：设备到设备（D2D）传输

当前实现（SharedMemoryConnector，MooncakeConnector）使用D2H2D（设备到主机到设备）数据路径。张量被移动到CPU内存（主机）进行传输，这会产生PCIe开销。

如设计RFC中所述，未来版本将引入D2D（设备到设备）连接器

• 目标：直接GPU到GPU传输（通过NCCL，UCX或IPC）以最小化大张量负载的延迟。
• 机制：OmniConnector API允许put()发起传输并通过控制平面返回一个轻量级句柄（元数据），而重负载直接在设备之间流动。