fix(inference): map external engine ranks

src/prime_rl/inference/vllm/ranks.py

@@ -0,0 +1,39 @@
+def global_inference_rank(
+    *,
+    rank_offset: int,
+    data_parallel_index: int,
+    data_parallel_size: int,
+    worker_rank: int,
+    tensor_parallel_size: int,
+    pipeline_parallel_size: int,
+    inference_world_size: int,
+    prefill_context_parallel_size: int = 1,
+    engine_world_size: int | None = None,
+) -> int:
+    """Map one vLLM worker to its rank in Prime's inference NCCL group."""
+    model_parallel_size = tensor_parallel_size * pipeline_parallel_size * prefill_context_parallel_size
+    logical_data_parallel_size = data_parallel_size
+    if engine_world_size is not None:
+        if engine_world_size <= 0 or engine_world_size % model_parallel_size:
+            raise ValueError(
+                f"engine world size {engine_world_size} is not divisible by model parallel size {model_parallel_size}"
+            )
+        logical_data_parallel_size = engine_world_size // model_parallel_size
+        # Dense vLLM EngineCore processes retain their global DP index but rewrite
+        # data_parallel_size to one. MoE EngineCore processes preserve the logical
+        # size, so keep validating that value rather than masking bad discovery.
+        if data_parallel_size != 1 and data_parallel_size != logical_data_parallel_size:
+            raise ValueError(
+                f"data parallel size {data_parallel_size} does not match engine-derived size "
+                f"{logical_data_parallel_size}"
+            )
+    if not 0 <= data_parallel_index < logical_data_parallel_size:
+        raise ValueError(
+            f"data parallel index {data_parallel_index} is outside logical data parallel size "
+            f"{logical_data_parallel_size}"
+        )
+
+    rank = rank_offset + data_parallel_index * model_parallel_size + worker_rank % model_parallel_size
+    if not 0 <= rank < inference_world_size:
+        raise ValueError(f"calculated inference rank {rank} is outside inference world size {inference_world_size}")
+    return rank

src/prime_rl/inference/vllm/server.py

@@ -143,11 +143,21 @@ async def init_broadcaster(request: Request):
     timeout = data.get("timeout")
     rank_offset = data.get("rank_offset")
     inference_world_size = data.get("inference_world_size")
+    engine_world_size = data.get("engine_world_size")
     quantize_in_weight_transfer = data.get("quantize_in_weight_transfer", False)
     session_id = data.get("session_id", "default")
     await engine_client(request).collective_rpc(
         "init_broadcaster",
-        args=(host, port, rank_offset, inference_world_size, timeout, quantize_in_weight_transfer, session_id),
+        args=(
+            host,
+            port,
+            rank_offset,
+            inference_world_size,
+            timeout,
+            quantize_in_weight_transfer,
+            session_id,
+            engine_world_size,
+        ),
     )
     return {"status": "ok"}
 

src/prime_rl/inference/vllm/worker/nccl.py

@@ -7,6 +7,7 @@
 from vllm.distributed.utils import StatelessProcessGroup
 from vllm.logger import init_logger
 
+from prime_rl.inference.vllm.ranks import global_inference_rank
 from prime_rl.inference.vllm.worker.weight_transfer import (
     load_weights_checkpoint_layerwise,
     load_weights_kernel,
@@ -101,27 +102,37 @@ def init_broadcaster(
         timeout: int,
         quantize_in_weight_transfer: bool = False,
         session_id: str = "default",
+        engine_world_size: int | None = None,
     ) -> None:
         """Initialize the NCCL broadcast receiver.
 
         Args:
             rank_offset: Starting GPU offset for this server in the global inference group.
             inference_world_size: Total number of inference GPUs across all servers.
+            engine_world_size: Number of inference GPUs assigned to this server.
         """
         del session_id
         self.quantize_in_weight_transfer = quantize_in_weight_transfer
-        # Use the worker's device index directly as the local rank.
-        # The previous dp_group-based computation broke in vLLM v1 multiprocess
-        # DP mode where each worker is a separate process with a singleton
-        # DP group (rank_in_group is always 0).
-        local_rank = self.device.index
-        global_rank_inference = rank_offset + local_rank
+        if engine_world_size is None:
+            global_rank_inference = rank_offset + self.device.index
+        else:
+            parallel_config = self.parallel_config
+            global_rank_inference = global_inference_rank(
+                rank_offset=rank_offset,
+                data_parallel_index=parallel_config.data_parallel_index,
+                data_parallel_size=parallel_config.data_parallel_size,
+                worker_rank=self.rank,
+                tensor_parallel_size=parallel_config.tensor_parallel_size,
+                pipeline_parallel_size=parallel_config.pipeline_parallel_size,
+                prefill_context_parallel_size=parallel_config.prefill_context_parallel_size,
+                inference_world_size=inference_world_size,
+                engine_world_size=engine_world_size,
+            )
 
         logger.info(
-            f"Worker [local_rank={local_rank} rank_offset={rank_offset}] "
+            f"Worker [worker_rank={self.rank} rank_offset={rank_offset}] "
             f"-> [global_rank={global_rank_inference} inference_world_size={inference_world_size}]"
         )
-
         self.nccl_broadcast_receiver = NCCLWeightBroadcastReceiver(
             host=host,
             port=port,

src/prime_rl/inference/vllm/worker/nixl.py

@@ -18,6 +18,7 @@
 from vllm.config import set_current_vllm_config
 from vllm.logger import init_logger
 
+from prime_rl.inference.vllm.ranks import global_inference_rank
 from prime_rl.inference.vllm.worker.weight_transfer import update_mla_absorbed_weights
 from prime_rl.trainer.rl.broadcast.nixl.agent import MemDesc, NixlAgent, make_agent_name, set_ucx_env_defaults
 from prime_rl.trainer.rl.broadcast.nixl.cuda_malloc_memory import (
@@ -96,9 +97,24 @@ def init_broadcaster(
         timeout: int,
         quantize_in_weight_transfer: bool = False,
         session_id: str = "default",
+        engine_world_size: int | None = None,
     ) -> None:
-        del inference_world_size, quantize_in_weight_transfer
-        global_rank = rank_offset + self.device.index
+        del quantize_in_weight_transfer
+        if engine_world_size is None:
+            global_rank = rank_offset + self.device.index
+        else:
+            parallel_config = self.parallel_config
+            global_rank = global_inference_rank(
+                rank_offset=rank_offset,
+                data_parallel_index=parallel_config.data_parallel_index,
+                data_parallel_size=parallel_config.data_parallel_size,
+                worker_rank=self.rank,
+                tensor_parallel_size=parallel_config.tensor_parallel_size,
+                pipeline_parallel_size=parallel_config.pipeline_parallel_size,
+                prefill_context_parallel_size=parallel_config.prefill_context_parallel_size,
+                inference_world_size=inference_world_size,
+                engine_world_size=engine_world_size,
+            )
         server_url = f"{host}:{port}"
         set_ucx_env_defaults()
         self.nixl_agent = NixlAgent(make_agent_name("inference", global_rank))

src/prime_rl/utils/client.py

@@ -515,6 +515,8 @@ async def init_nccl_broadcast(
     timeout: int,
     inference_world_size: int | None = None,
     quantize_in_weight_transfer: bool = False,
+    *,
+    engine_world_sizes: list[int] | None = None,
 ) -> None:
     """Initialize NCCL broadcast on all inference servers.
 
@@ -524,32 +526,46 @@ async def init_nccl_broadcast(
     """
     logger = get_logger()
 
+    has_explicit_engine_world_sizes = engine_world_sizes is not None
     if inference_world_size is None:
-        inference_world_size = len(admin_clients)
+        if engine_world_sizes is not None:
+            inference_world_size = sum(engine_world_sizes)
+        else:
+            inference_world_size = len(admin_clients)
         logger.warning(
             f"inference_world_size not provided, defaulting to {inference_world_size} (one GPU per admin client)"
         )
 
-    gpus_per_server = inference_world_size // len(admin_clients)
+    if engine_world_sizes is None:
+        if inference_world_size % len(admin_clients) != 0:
+            raise ValueError("inference_world_size must be divisible by the number of admin clients")
+        engine_world_sizes = [inference_world_size // len(admin_clients)] * len(admin_clients)
+    if len(engine_world_sizes) != len(admin_clients):
+        raise ValueError("one engine world size is required for each admin client")
+    rank_offsets = _rank_offsets(engine_world_sizes, inference_world_size)
 
     logger.info(
         f"Initializing NCCL broadcast: {len(admin_clients)} servers, "
-        f"inference_world_size={inference_world_size}, gpus_per_server={gpus_per_server}"
+        f"inference_world_size={inference_world_size}, engine_world_sizes={engine_world_sizes}"
     )
 
-    async def _init_nccl_broadcast(admin_client: AsyncClient, rank_offset: int) -> None:
+    async def _init_nccl_broadcast(
+        admin_client: AsyncClient,
+        rank_offset: int,
+        engine_world_size: int,
+    ) -> None:
+        payload = {
+            "host": host,
+            "port": port,
+            "rank_offset": rank_offset,
+            "inference_world_size": inference_world_size,
+            "timeout": timeout,
+            "quantize_in_weight_transfer": quantize_in_weight_transfer,
+        }
+        if has_explicit_engine_world_sizes:
+            payload["engine_world_size"] = engine_world_size
         try:
-            response = await admin_client.post(
-                "/init_broadcaster",
-                json={
-                    "host": host,
-                    "port": port,
-                    "rank_offset": rank_offset,
-                    "inference_world_size": inference_world_size,
-                    "timeout": timeout,
-                    "quantize_in_weight_transfer": quantize_in_weight_transfer,
-                },
-            )
+            response = await admin_client.post("/init_broadcaster", json=payload)
             response.raise_for_status()
         except httpx.HTTPStatusError as e:
             if e.response.status_code == 404:
@@ -558,8 +574,10 @@ async def _init_nccl_broadcast(admin_client: AsyncClient, rank_offset: int) -> N
 
     await asyncio.gather(
         *[
-            _init_nccl_broadcast(admin_client, client_num * gpus_per_server)
-            for client_num, admin_client in enumerate(admin_clients)
+            _init_nccl_broadcast(admin_client, rank_offset, engine_world_size)
+            for admin_client, rank_offset, engine_world_size in zip(
+                admin_clients, rank_offsets, engine_world_sizes, strict=True
+            )
         ]
     )
 
@@ -571,31 +589,61 @@ async def init_nixl_broadcast(
     timeout: int,
     inference_world_size: int,
     session_id: str,
+    *,
+    engine_world_sizes: list[int] | None = None,
 ) -> None:
     """Configure every vLLM worker for NIXL + ModelExpress pulls."""
-    workers_per_server = inference_world_size // len(admin_clients)
-
-    async def initialize(admin_client: AsyncClient, rank_offset: int) -> None:
+    has_explicit_engine_world_sizes = engine_world_sizes is not None
+    if engine_world_sizes is None:
+        if inference_world_size % len(admin_clients) != 0:
+            raise ValueError("inference_world_size must be divisible by the number of admin clients")
+        engine_world_sizes = [inference_world_size // len(admin_clients)] * len(admin_clients)
+    if len(engine_world_sizes) != len(admin_clients):
+        raise ValueError("one engine world size is required for each admin client")
+    rank_offsets = _rank_offsets(engine_world_sizes, inference_world_size)
+
+    async def initialize(admin_client: AsyncClient, rank_offset: int, engine_world_size: int) -> None:
+        payload = {
+            "host": host,
+            "port": port,
+            "rank_offset": rank_offset,
+            "inference_world_size": inference_world_size,
+            "timeout": timeout,
+            "quantize_in_weight_transfer": False,
+            "session_id": session_id,
+        }
+        if has_explicit_engine_world_sizes:
+            payload["engine_world_size"] = engine_world_size
         await _admin_post(
             admin_client,
             "/init_broadcaster",
             timeout_s=max(ADMIN_TIMEOUT_S, timeout),
-            json={
-                "host": host,
-                "port": port,
-                "rank_offset": rank_offset,
-                "inference_world_size": inference_world_size,
-                "timeout": timeout,
-                "quantize_in_weight_transfer": False,
-                "session_id": session_id,
-            },
+            json=payload,
         )
 
     await asyncio.gather(
-        *[initialize(admin_client, index * workers_per_server) for index, admin_client in enumerate(admin_clients)]
+        *[
+            initialize(admin_client, rank_offset, engine_world_size)
+            for admin_client, rank_offset, engine_world_size in zip(
+                admin_clients, rank_offsets, engine_world_sizes, strict=True
+            )
+        ]
     )
 
 
+def _rank_offsets(engine_world_sizes: list[int], inference_world_size: int) -> list[int]:
+    if not engine_world_sizes or any(isinstance(size, bool) or size <= 0 for size in engine_world_sizes):
+        raise ValueError("engine world sizes must be positive integers")
+    if sum(engine_world_sizes) != inference_world_size:
+        raise ValueError("engine world sizes do not match inference_world_size")
+    offsets: list[int] = []
+    offset = 0
+    for world_size in engine_world_sizes:
+        offsets.append(offset)
+        offset += world_size
+    return offsets
+
+
 async def prefill_logprobs(openai: AsyncOpenAI, model: str, token_ids: list[int]) -> list[float]:
     """Prefill-score ``token_ids`` under ``model`` via ``/inference/v1/generate``
     + ``prompt_logprobs`` (the prime-rl server-side extension in