replica.py — Rollout 副本管理

文件概述

定义 RolloutReplica 抽象基类和 RolloutReplicaRegistry 注册表，管理推理引擎的多副本部署。

核心概念

Rollout Mode（推理模式）

class RolloutMode(Enum):
    HYBRID = "hybrid"          # 混合模式：训练和推理共用 GPU
    COLOCATED = "colocated"    # 共置模式：同节点但独立进程
    STANDALONE = "standalone"  # 独立模式：专用 GPU

TokenOutput

推理输出的数据结构：

@dataclass
class TokenOutput:
    token_ids: list[int]        # 生成的 token IDs
    log_probs: list[float]      # 对应的 log 概率
    routed_experts: list = None  # MoE 路由结果（Router Replay 用）
    stop_reason: str = None     # 停止原因
    extra_info: dict = None     # 额外信息

核心类

RolloutReplica(ABC)

class RolloutReplica(ABC):
    """管理推理引擎在一个副本（可能跨多节点）上的部署

    一个 replica 包含:
    - 一组 worker 进程（对应 GPU）
    - 一个或多个 HTTP server（每个节点一个）
    """
    def __init__(self, replica_rank, config, model_config, gpus_per_node, ...):
        self.replica_rank = replica_rank
        self.world_size = config.tensor_model_parallel_size * config.data_parallel_size
        self.nnodes = self.world_size // gpus_per_node

    @abstractmethod
    async def launch_servers(self):
        """启动推理服务器"""
        pass

    async def sleep(self):
        """释放显存"""
        pass

    async def wake_up(self):
        """恢复显存"""
        pass

RolloutReplicaRegistry

class RolloutReplicaRegistry:
    """注册表：根据引擎名称获取对应的 Replica 类"""
    _registry = {}

    @classmethod
    def register(cls, name):
        def decorator(replica_cls):
            cls._registry[name] = replica_cls
            return replica_cls
        return decorator

副本部署示例

2 个副本, 每个副本 TP=4:

Replica 0:              Replica 1:
┌────────────────┐      ┌────────────────┐
│ GPU0 GPU1      │      │ GPU4 GPU5      │
│ GPU2 GPU3      │      │ GPU6 GPU7      │
│ HTTP Server    │      │ HTTP Server    │
└────────────────┘      └────────────────┘

与其他模块的关系

• 被 vllm_rollout/vllm_async_server.py 的 vLLMReplica 继承
• 被 sglang_rollout/async_sglang_server.py 的 SGLangReplica 继承
• 被 trtllm_rollout/trtllm_async_server.py 的 TRTLLMReplica 继承

小结

RolloutReplica 管理推理引擎的多副本部署，支持混合/共置/独立三种模式，是推理引擎水平扩展的基础。