feat(vllm): support mixed DP launch modes

docs/config-reference.md

@@ -422,6 +422,44 @@ Each worker leader gets a globally unique port starting at 5550:
 | decode_0  | 5552 |
 | decode_1  | 5553 |
 
+### vLLM DP launch mode
+
+vLLM data-parallel endpoints use one process per GPU by default. Set
+`dp_launch_mode: per_node` to launch one process per node and let vLLM
+manage the local DP ranks in a shared CUDA namespace:
+
+```yaml
+backend:
+  type: vllm
+  dp_launch_mode: per_node
+  decode_dp_launch_mode: per_gpu  # Optional role override
+  vllm_config:
+    prefill:
+      data-parallel-size: 8
+      data-parallel-hybrid-lb: true
+    decode:
+      data-parallel-size: 16
+      data-parallel-hybrid-lb: true
+```
+
+| Value      | Process layout                                      |
+| ---------- | --------------------------------------------------- |
+| `per_gpu`  | One process per DP rank/GPU (default)               |
+| `per_node` | One process manages all DP ranks allocated per node |
+
+`prefill_dp_launch_mode`, `decode_dp_launch_mode`, and
+`aggregated_dp_launch_mode` can override the global mode for one role. For
+example, keeping prefill on `per_node` while setting decode to `per_gpu`
+preserves a shared CUDA namespace for prefill collectives while exposing each
+decode rank as an independent Dynamo worker. Roles without an override inherit
+`dp_launch_mode`.
+
+In `per_node` mode, srtslurm derives `--data-parallel-size-local` and
+`--data-parallel-start-rank` from the allocated topology. Do not set those
+two flags manually. Set `data-parallel-hybrid-lb: true` when the frontend
+must route to an exact DP rank; without hybrid LB, non-leader node processes
+are launched with `--headless` for vLLM's internal DP load balancer.
+
 ### TRTLLM Backend
 
 When using `type: trtllm`, the backend uses TRTLLM with MPI-style launching:

src/srtctl/backends/vllm.py

@@ -41,6 +41,7 @@
 
 # Type alias for worker modes
 WorkerMode = Literal["prefill", "decode", "agg"]
+DPLaunchMode = Literal["per_gpu", "per_node"]
 
 # Filename for the mooncake-store JSON config srtslurm writes to log_dir at job
 # start. log_dir is mounted into every worker at /logs, so workers read the JSON
@@ -138,6 +139,8 @@ class VLLMProtocol:
           connector: nixl  # translated to --kv-transfer-config JSON
           allow_prefill_decode_colocation: true  # pack P/D on one node when all workers fit
           allow_prefill_decode_colocation_across_nodes: true  # continue packing on later nodes
+          dp_launch_mode: per_node  # one process manages all local DP ranks
+          decode_dp_launch_mode: per_gpu  # optional per-role override
           prefill_environment:
             PYTHONUNBUFFERED: "1"
           vllm_config:
@@ -191,14 +194,23 @@ class VLLMProtocol:
     # node pools. Defaults off to preserve the original one-node-only policy.
     allow_prefill_decode_colocation_across_nodes: bool = False
 
+    # DP process layout. Keep the existing per-GPU behavior by default;
+    # per-node lets vLLM manage local DP ranks in one CUDA namespace.
+    dp_launch_mode: DPLaunchMode = "per_gpu"
+
+    # Optional per-role overrides. Unset roles inherit dp_launch_mode.
+    prefill_dp_launch_mode: DPLaunchMode | None = None
+    decode_dp_launch_mode: DPLaunchMode | None = None
+    aggregated_dp_launch_mode: DPLaunchMode | None = None
+
     Schema: ClassVar[builtins.type[Schema]] = Schema
 
     # =========================================================================
     # BackendProtocol Implementation
     # =========================================================================
 
     def get_srun_config(self) -> SrunConfig:
-        """vLLM uses per-process launching (one srun per node)."""
+        """vLLM launches one srun step for each generated process."""
         from srtctl.backends.base import SrunConfig
 
         return SrunConfig(mpi=None, oversubscribe=False, launch_per_endpoint=False)
@@ -404,7 +416,7 @@ def _is_dp_mode(self, mode: WorkerMode) -> bool:
         """Check if this mode uses Data Parallel + Expert Parallel pattern.
 
         DP+EP mode is detected when data-parallel-size is set in the mode's config.
-        In this mode, each GPU runs its own process (rather than TP across GPUs).
+        ``dp_launch_mode`` controls whether a process owns one rank or all local ranks.
         """
         config = self.get_config_for_mode(mode)
         return config.get("data-parallel-size") is not None or config.get("data_parallel_size") is not None
@@ -414,6 +426,50 @@ def _get_dp_size(self, mode: WorkerMode) -> int | None:
         config = self.get_config_for_mode(mode)
         return config.get("data-parallel-size") or config.get("data_parallel_size")
 
+    def get_dp_launch_mode_for_mode(self, mode: WorkerMode) -> DPLaunchMode:
+        """Return the role-specific launch mode or the global default."""
+        overrides = {
+            "prefill": self.prefill_dp_launch_mode,
+            "decode": self.decode_dp_launch_mode,
+            "agg": self.aggregated_dp_launch_mode,
+        }
+        return overrides[mode] or self.dp_launch_mode
+
+    def get_expected_dynamo_worker_counts(self, processes: Sequence[Process]) -> tuple[int, int]:
+        """Return the generate-instance counts expected from a process layout.
+
+        Per-GPU processes each register one instance. Hybrid per-node processes
+        register their local rank range, while non-hybrid headless followers do
+        not register separately.
+        """
+        grouped: dict[tuple[WorkerMode, int], list[Process]] = {}
+        for process in processes:
+            key = (process.endpoint_mode, process.endpoint_index)
+            grouped.setdefault(key, []).append(process)
+
+        prefill_count = 0
+        decode_count = 0
+        for (mode, _endpoint_index), endpoint_processes in grouped.items():
+            registrations = 1
+            if self._is_dp_mode(mode):
+                launch_mode = self.get_dp_launch_mode_for_mode(mode)
+                if launch_mode == "per_gpu":
+                    registrations = len(endpoint_processes)
+                else:
+                    config = self.get_config_for_mode(mode)
+                    hybrid_lb = config.get("data-parallel-hybrid-lb", config.get("data_parallel_hybrid_lb", False))
+                    hybrid_lb_enabled = hybrid_lb is True or (
+                        isinstance(hybrid_lb, str) and hybrid_lb.strip().lower() in {"1", "true", "yes", "on"}
+                    )
+                    registrations = len(endpoint_processes) if hybrid_lb_enabled else 1
+
+            if mode == "prefill":
+                prefill_count += registrations
+            else:
+                decode_count += registrations
+
+        return prefill_count, decode_count
+
     def should_set_cuda_visible_devices(self, process: Process) -> bool:
         """Whether worker_stage should set CUDA_VISIBLE_DEVICES.
 
@@ -431,18 +487,65 @@ def endpoints_to_processes(
     ) -> list[Process]:
         """Convert endpoints to processes.
 
-        For DP+EP mode (data-parallel-size set), creates one process per GPU.
+        DP+EP mode uses the configured per-GPU or per-node process layout.
         For standard TP mode, creates one process per node.
         """
-        from srtctl.core.topology import NodePortAllocator, Process, endpoints_to_processes
+        from srtctl.core.topology import NodePortAllocator, endpoints_to_processes
 
-        # Check if any endpoint uses DP mode
-        has_dp_mode = any(self._is_dp_mode(ep.mode) for ep in endpoints)
-
-        if not has_dp_mode:
+        dp_endpoints = [endpoint for endpoint in endpoints if self._is_dp_mode(endpoint.mode)]
+        if not dp_endpoints:
             # Standard TP mode: one process per node
             return endpoints_to_processes(endpoints, base_sys_port=base_sys_port, port_allocator=port_allocator)
 
+        launch_modes = {self.get_dp_launch_mode_for_mode(endpoint.mode) for endpoint in dp_endpoints}
+        if launch_modes == {"per_node"}:
+            return self._dp_per_node_endpoints_to_processes(
+                endpoints,
+                base_sys_port=base_sys_port,
+                port_allocator=port_allocator,
+            )
+
+        if launch_modes == {"per_gpu"}:
+            return self._dp_per_gpu_endpoints_to_processes(
+                endpoints,
+                base_sys_port=base_sys_port,
+                port_allocator=port_allocator,
+            )
+
+        # Mixed role modes share one allocator so every coordination port stays
+        # unique across the per-node and per-GPU process groups.
+        processes: list[Process] = []
+        current_sys_port = base_sys_port
+        if port_allocator is None:
+            port_allocator = NodePortAllocator()
+
+        for endpoint in endpoints:
+            if self._is_dp_mode(endpoint.mode) and self.get_dp_launch_mode_for_mode(endpoint.mode) == "per_node":
+                endpoint_processes = self._dp_per_node_endpoints_to_processes(
+                    [endpoint],
+                    base_sys_port=current_sys_port,
+                    port_allocator=port_allocator,
+                )
+            else:
+                endpoint_processes = self._dp_per_gpu_endpoints_to_processes(
+                    [endpoint],
+                    base_sys_port=current_sys_port,
+                    port_allocator=port_allocator,
+                )
+            processes.extend(endpoint_processes)
+            current_sys_port += len(endpoint_processes)
+
+        return processes
+
+    def _dp_per_gpu_endpoints_to_processes(
+        self,
+        endpoints: list[Endpoint],
+        base_sys_port: int = DYN_SYSTEM_PORT_BASE,
+        port_allocator: NodePortAllocator | None = None,
+    ) -> list[Process]:
+        """Convert DP endpoints to one process per GPU."""
+        from srtctl.core.topology import NodePortAllocator, Process
+
         # DP+EP mode: one process per GPU
         processes: list[Process] = []
         current_sys_port = base_sys_port
@@ -520,6 +623,67 @@ def endpoints_to_processes(
 
         return processes
 
+    def _dp_per_node_endpoints_to_processes(
+        self,
+        endpoints: list[Endpoint],
+        base_sys_port: int = DYN_SYSTEM_PORT_BASE,
+        port_allocator: NodePortAllocator | None = None,
+    ) -> list[Process]:
+        """Convert DP endpoints to one process per node."""
+        from srtctl.core.topology import NodePortAllocator, Process, endpoints_to_processes
+
+        processes: list[Process] = []
+        current_sys_port = base_sys_port
+        if port_allocator is None:
+            port_allocator = NodePortAllocator()
+
+        for endpoint in endpoints:
+            if not self._is_dp_mode(endpoint.mode):
+                non_dp = endpoints_to_processes(
+                    [endpoint],
+                    base_sys_port=current_sys_port,
+                    port_allocator=port_allocator,
+                )
+                processes.extend(non_dp)
+                current_sys_port += len(non_dp)
+                continue
+
+            dp_size = self._get_dp_size(endpoint.mode) or endpoint.total_gpus
+            if dp_size != endpoint.total_gpus:
+                raise ValueError(
+                    f"{endpoint.mode} data-parallel-size={dp_size} does not match "
+                    f"the endpoint's {endpoint.total_gpus} allocated GPUs"
+                )
+
+            local_dp_size = len(endpoint.gpu_indices)
+            dp_rpc_port = port_allocator.next_dp_rpc_port(endpoint.leader_node)
+            nixl_base_port = port_allocator.next_nixl_port_block(dp_size)
+            dp_start_rank = 0
+
+            for node in endpoint.nodes:
+                processes.append(
+                    Process(
+                        node=node,
+                        gpu_indices=endpoint.gpu_indices,
+                        sys_port=current_sys_port,
+                        http_port=port_allocator.next_http_port(node),
+                        endpoint_mode=endpoint.mode,
+                        endpoint_index=endpoint.index,
+                        node_rank=dp_start_rank,
+                        bootstrap_port=(
+                            port_allocator.next_bootstrap_port(node) if endpoint.mode == "prefill" else None
+                        ),
+                        kv_events_port=port_allocator.next_kv_events_port_block(local_dp_size),
+                        nixl_port=nixl_base_port,
+                        dp_rpc_port=dp_rpc_port,
+                        het_group=endpoint.het_group,
+                    )
+                )
+                current_sys_port += 1
+                dp_start_rank += local_dp_size
+
+        return processes
+
     def build_worker_command(
         self,
         process: Process,
@@ -598,7 +762,38 @@ def build_worker_command(
 
         # Check if this is DP+EP mode (data-parallel-size set)
         is_dp_mode = self._is_dp_mode(mode)
-        if is_dp_mode:
+        if is_dp_mode and self.get_dp_launch_mode_for_mode(mode) == "per_node":
+            rpc_port_kebab = config.pop("data-parallel-rpc-port", None)
+            rpc_port_snake = config.pop("data_parallel_rpc_port", None)
+            config_dp_rpc_port = rpc_port_kebab or rpc_port_snake
+            dp_rpc_port = process.dp_rpc_port or config_dp_rpc_port or VLLM_DATA_PARALLEL_RPC_PORT
+
+            # These values are derived from the allocated process topology.
+            config.pop("data-parallel-size-local", None)
+            config.pop("data_parallel_size_local", None)
+            config.pop("data-parallel-start-rank", None)
+            config.pop("data_parallel_start_rank", None)
+
+            cmd.extend(
+                [
+                    "--data-parallel-size-local",
+                    str(len(process.gpu_indices)),
+                    "--data-parallel-start-rank",
+                    str(process.node_rank),
+                    "--data-parallel-address",
+                    leader_ip,
+                    "--data-parallel-rpc-port",
+                    str(dp_rpc_port),
+                ]
+            )
+
+            hybrid_lb = config.get("data-parallel-hybrid-lb", config.get("data_parallel_hybrid_lb", False))
+            hybrid_lb_enabled = hybrid_lb is True or (
+                isinstance(hybrid_lb, str) and hybrid_lb.strip().lower() in {"1", "true", "yes", "on"}
+            )
+            if process.node_rank > 0 and not hybrid_lb_enabled:
+                cmd.append("--headless")
+        elif is_dp_mode:
             # DP+EP mode: each GPU runs its own process
             # process.node_rank is the dp_rank (set in endpoints_to_processes)
             dp_rank = process.node_rank

src/srtctl/cli/mixins/benchmark_stage.py

@@ -46,7 +46,9 @@ def _vllm_data_parallel_size(config: "SrtConfig", mode: str) -> int:
     return int(mode_config.get("data-parallel-size") or mode_config.get("data_parallel_size") or 1)
 
 
-def _get_health_expectations(config: "SrtConfig") -> tuple[int, int, str, int]:
+def _get_health_expectations(
+    config: "SrtConfig", processes: list["Process"] | None = None
+) -> tuple[int, int, str, int]:
     """Compute expected health counts in the units reported by the frontend.
 
     Dynamo's /health endpoint reports registered generate instances. For vLLM
@@ -65,7 +67,10 @@ def _get_health_expectations(config: "SrtConfig") -> tuple[int, int, str, int]:
         worker_desc = f"{r.num_prefill}P + {r.num_decode}D"
 
     if config.frontend.type == "dynamo" and getattr(config.backend, "type", None) == "vllm":
-        if r.num_agg > 0:
+        registration_counter = getattr(config.backend, "get_expected_dynamo_worker_counts", None)
+        if processes is not None and callable(registration_counter):
+            n_prefill, n_decode = registration_counter(processes)
+        elif r.num_agg > 0:
             n_prefill = 0
             n_decode = logical_decode * _vllm_data_parallel_size(config, "aggregated")
         else:
@@ -131,7 +136,7 @@ def run_benchmark(
         """Run the benchmark."""
         logger.info("Waiting for workers to be ready...")
 
-        n_prefill, n_decode, count_desc, num_workers = _get_health_expectations(self.config)
+        n_prefill, n_decode, count_desc, num_workers = _get_health_expectations(self.config, self.backend_processes)
         logger.info("Waiting for server health (expecting %d health entries: %s)...", num_workers, count_desc)
 
         hc = self.config.health_check

src/srtctl/core/topology.py

@@ -98,6 +98,20 @@ def next_kv_events_port(self) -> int:
         self._next_kv_events_port += 1
         return port
 
+    def next_kv_events_port_block(self, size: int) -> int:
+        """Reserve consecutive KV-event ports and return the base port.
+
+        A vLLM hybrid-DP process opens one publisher per local DP rank, so
+        adjacent worker processes must receive non-overlapping port ranges.
+        """
+        if size < 1:
+            raise ValueError("KV-event port block size must be at least 1")
+        if self._next_kv_events_port == 0:
+            self._next_kv_events_port = self.base_kv_events_port
+        port = self._next_kv_events_port
+        self._next_kv_events_port += size
+        return port
+
     def next_nixl_port(self) -> int:
         """Get next available NIXL side channel port (globally unique across all nodes)."""
         if self._next_nixl_port == 0: