chore: import upstream snapshot with attribution

tests/v1/distributed/test_async_llm_dp.py

@@ -0,0 +1,576 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import asyncio
+import os
+import time
+from contextlib import ExitStack
+from dataclasses import dataclass
+from typing import Any
+
+import pytest
+
+from vllm import SamplingParams
+from vllm.config import VllmConfig
+from vllm.engine.arg_utils import AsyncEngineArgs
+from vllm.inputs import PromptType
+from vllm.outputs import RequestOutput
+from vllm.platforms import current_platform
+from vllm.sampling_params import RequestOutputKind
+from vllm.v1.engine.async_llm import AsyncLLM
+from vllm.v1.engine.core_client import DPAsyncMPClient
+from vllm.v1.metrics.loggers import StatLoggerBase
+from vllm.v1.metrics.stats import IterationStats, MultiModalCacheStats, SchedulerStats
+
+DP_SIZE = int(os.getenv("DP_SIZE", 2))
+
+
+async def generate(
+    engine: AsyncLLM,
+    request_id: str,
+    prompt: PromptType,
+    output_kind: RequestOutputKind,
+    max_tokens: int,
+    prompt_logprobs: int | None = None,
+    data_parallel_rank: int | None = None,
+) -> tuple[int, str]:
+    # Ensure generate doesn't complete too fast for cancellation test.
+    await asyncio.sleep(0.2)
+
+    count = 0
+    sampling_params = SamplingParams(
+        max_tokens=max_tokens,
+        ignore_eos=True,
+        output_kind=output_kind,
+        temperature=0,
+        prompt_logprobs=prompt_logprobs,
+    )
+    async for out in engine.generate(
+        request_id=request_id,
+        prompt=prompt,
+        sampling_params=sampling_params,
+        data_parallel_rank=data_parallel_rank,
+    ):
+        num_tokens = len(out.outputs[0].token_ids)
+        if output_kind == RequestOutputKind.DELTA:
+            count += num_tokens
+        else:
+            count = num_tokens
+
+        await asyncio.sleep(0.0)
+
+    return count, request_id
+
+
+@pytest.mark.parametrize(
+    "model",
+    [
+        "ibm-research/PowerMoE-3b",
+        "hmellor/tiny-random-LlamaForCausalLM",
+    ],
+)
+@pytest.mark.parametrize(
+    "output_kind",
+    [
+        RequestOutputKind.DELTA,
+        RequestOutputKind.FINAL_ONLY,
+    ],
+)
+@pytest.mark.parametrize("data_parallel_backend", ["mp", "ray"])
+@pytest.mark.parametrize("async_scheduling", [True, False])
+@pytest.mark.asyncio
+async def test_load(
+    model: str,
+    output_kind: RequestOutputKind,
+    data_parallel_backend: str,
+    async_scheduling: bool,
+):
+    if async_scheduling and data_parallel_backend == "ray":
+        # TODO(NickLucche) Re-enable when async scheduling is supported
+        pytest.skip("Async scheduling is not supported with ray")
+    elif data_parallel_backend == "ray" and current_platform.is_rocm():
+        pytest.skip(
+            "Ray as the distributed executor backend is not supported with ROCm."
+        )
+    stats_loggers = {}
+
+    @dataclass
+    class SimpleStatsLogger(StatLoggerBase):
+        init_count: int = 0
+        finished_req_count: int = 0
+
+        def __init__(self, vllm_config: VllmConfig, engine_index: int = 0):
+            stats_loggers[engine_index] = self
+
+        def record(
+            self,
+            scheduler_stats: SchedulerStats | None,
+            iteration_stats: IterationStats | None,
+            mm_cache_stats: MultiModalCacheStats | None = None,
+            engine_idx: int = 0,
+        ):
+            if iteration_stats:
+                self.finished_req_count += len(iteration_stats.finished_requests)
+
+        def log_engine_initialized(self):
+            self.init_count += 1
+
+    with ExitStack() as after:
+        prompt = "This is a test of data parallel"
+
+        engine_args = AsyncEngineArgs(
+            model=model,
+            enforce_eager=True,
+            tensor_parallel_size=int(os.getenv("TP_SIZE", 1)),
+            data_parallel_size=DP_SIZE,
+            data_parallel_backend=data_parallel_backend,
+            async_scheduling=async_scheduling,
+        )
+        engine = AsyncLLM.from_engine_args(
+            engine_args, stat_loggers=[SimpleStatsLogger]
+        )
+        after.callback(engine.shutdown)
+
+        NUM_REQUESTS = 100
+        NUM_EXPECTED_TOKENS = 10
+
+        request_ids = [f"request-{i}" for i in range(NUM_REQUESTS)]
+
+        # Create concurrent requests.
+        tasks = []
+        for request_id in request_ids:
+            tasks.append(
+                asyncio.create_task(
+                    generate(
+                        engine, request_id, prompt, output_kind, NUM_EXPECTED_TOKENS
+                    )
+                )
+            )
+            # Short sleep to ensure that requests are distributed.
+            await asyncio.sleep(0.01)
+        # Confirm that we got all the EXPECTED tokens from the requests.
+        done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION)
+        for task in pending:
+            task.cancel()
+        for task in done:
+            num_generated_tokens, request_id = await task
+            assert num_generated_tokens == NUM_EXPECTED_TOKENS, (
+                f"{request_id} generated {num_generated_tokens} but "
+                f"expected {NUM_EXPECTED_TOKENS}"
+            )
+
+        assert not engine.output_processor.has_unfinished_requests()
+
+        # testing internals here which may break
+        core_client: DPAsyncMPClient = engine.engine_core
+        # the engines only synchronize stopping every N steps so
+        # allow a small amount of time here.
+        for _ in range(10):
+            if not core_client.engines_running:
+                break
+            await asyncio.sleep(0.5)
+
+        assert not core_client.engines_running
+        assert not core_client.reqs_in_flight
+
+        # Check that requests were distributed between the engines
+        print(f"Stats loggers after test: {stats_loggers}")
+        assert len(stats_loggers) == DP_SIZE
+        assert stats_loggers[0].init_count == 1
+
+        for sl in stats_loggers.values():
+            slogger: SimpleStatsLogger = sl
+
+            assert slogger.finished_req_count > NUM_REQUESTS // (DP_SIZE + 1), (
+                f"requests are imbalanced: {stats_loggers}"
+            )
+
+
+@pytest.mark.parametrize("prefill_schedule_interval", [1, 4])
+@pytest.mark.asyncio
+async def test_dp_prefill_schedule_interval(prefill_schedule_interval: int):
+    """Throttling new prefills to every Nth step (DP balancing) must not break
+    generation: a stream of staggered requests should still all complete with
+    the expected number of tokens.
+
+    The throttle only engages in the DP MoE/EP engine-core path
+    (`DPEngineCoreProc`), so this uses an MoE model with expert parallel.
+    """
+    with ExitStack() as after:
+        prompt = "This is a test of data parallel"
+
+        engine_args = AsyncEngineArgs(
+            model="ibm-research/PowerMoE-3b",
+            enforce_eager=True,
+            tensor_parallel_size=int(os.getenv("TP_SIZE", 1)),
+            data_parallel_size=DP_SIZE,
+            data_parallel_backend="mp",
+            enable_expert_parallel=True,
+            prefill_schedule_interval=prefill_schedule_interval,
+        )
+        engine = AsyncLLM.from_engine_args(engine_args)
+        after.callback(engine.shutdown)
+
+        NUM_REQUESTS = 50
+        NUM_EXPECTED_TOKENS = 10
+
+        request_ids = [f"request-{i}" for i in range(NUM_REQUESTS)]
+
+        # Create requests with a small stagger so they arrive across many
+        # steps and (with interval > 1) accumulate in the waiting queue
+        # before being admitted together on cadence-aligned steps.
+        tasks = []
+        for request_id in request_ids:
+            tasks.append(
+                asyncio.create_task(
+                    generate(
+                        engine,
+                        request_id,
+                        prompt,
+                        RequestOutputKind.DELTA,
+                        NUM_EXPECTED_TOKENS,
+                    )
+                )
+            )
+            await asyncio.sleep(0.01)
+
+        done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION)
+        for task in pending:
+            task.cancel()
+        for task in done:
+            num_generated_tokens, request_id = await task
+            assert num_generated_tokens == NUM_EXPECTED_TOKENS, (
+                f"{request_id} generated {num_generated_tokens} but "
+                f"expected {NUM_EXPECTED_TOKENS}"
+            )
+
+        assert not engine.output_processor.has_unfinished_requests()
+
+
+# =============================================================================
+# DP Pause/Resume Tests
+# =============================================================================
+# When expert_parallel=False: uses non-MoE model (DP replicas as separate engines).
+# When expert_parallel=True: uses MoE model + EP (DPEngineCoreProc, sync pause path).
+
+DP_PAUSE_MODEL = "hmellor/tiny-random-LlamaForCausalLM"
+DP_PAUSE_MODEL_MOE = "ibm-research/PowerMoE-3b"
+DP_PAUSE_PROMPT = "This is a test of data parallel pause"
+
+
+def _get_dp_pause_engine_args(expert_parallel: bool) -> AsyncEngineArgs:
+    """Engine args for DP pause tests: MoE+EP when expert_parallel else small Llama."""
+    model = DP_PAUSE_MODEL_MOE if expert_parallel else DP_PAUSE_MODEL
+    return AsyncEngineArgs(
+        model=model,
+        enforce_eager=True,
+        tensor_parallel_size=int(os.getenv("TP_SIZE", 1)),
+        data_parallel_size=DP_SIZE,
+        data_parallel_backend="mp",
+        enable_expert_parallel=expert_parallel,
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("expert_parallel", [False, True])
+async def test_dp_pause_resume_basic(expert_parallel: bool):
+    """Pausing from the client (one call) pauses all DP ranks; resume clears it."""
+    with ExitStack() as after:
+        engine_args = _get_dp_pause_engine_args(expert_parallel)
+        engine = AsyncLLM.from_engine_args(engine_args)
+        after.callback(engine.shutdown)
+
+        assert not await engine.is_paused()
+        await engine.pause_generation(mode="abort")
+        assert await engine.is_paused()
+        await engine.resume_generation()
+        assert not await engine.is_paused()
+
+        # Engine still works after resume
+        sampling_params = SamplingParams(max_tokens=5)
+        async for out in engine.generate(
+            request_id="after-resume",
+            prompt=DP_PAUSE_PROMPT,
+            sampling_params=sampling_params,
+        ):
+            pass
+        assert out.finished
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("expert_parallel", [False, True])
+async def test_dp_pause_abort(expert_parallel: bool):
+    """Pause with abort from one client aborts in-flight requests on all DP ranks."""
+    with ExitStack() as after:
+        engine_args = _get_dp_pause_engine_args(expert_parallel)
+        engine = AsyncLLM.from_engine_args(engine_args)
+        after.callback(engine.shutdown)
+
+        # Start several requests so they are distributed across ranks
+        sampling_params = SamplingParams(max_tokens=500, ignore_eos=True)
+        num_requests = 4
+        outputs_by_id: dict[str, list[RequestOutput]] = {}
+
+        async def gen(rid: str):
+            out_list: list[RequestOutput] = []
+            outputs_by_id[rid] = out_list
+            async for out in engine.generate(
+                request_id=rid,
+                prompt=DP_PAUSE_PROMPT,
+                sampling_params=sampling_params,
+            ):
+                out_list.append(out)
+            return out_list[-1] if out_list else None
+
+        tasks = [asyncio.create_task(gen(f"req-{i}")) for i in range(num_requests)]
+        # Wait for some tokens on at least one request
+        while not any(len(o) >= 2 for o in outputs_by_id.values()):
+            await asyncio.sleep(0.02)
+
+        await engine.pause_generation(mode="abort")
+
+        finals = await asyncio.gather(*tasks)
+        for i, final in enumerate(finals):
+            assert final is not None, f"req-{i} had no output"
+            assert final.finished
+            assert final.outputs[0].finish_reason == "abort"
+
+        assert await engine.is_paused()
+        await engine.resume_generation()
+        assert not await engine.is_paused()
+
+        # New request completes after resume
+        async for out in engine.generate(
+            request_id="after-abort",
+            prompt=DP_PAUSE_PROMPT,
+            sampling_params=SamplingParams(max_tokens=5),
+        ):
+            pass
+        assert out.finished
+        assert not engine.output_processor.has_unfinished_requests()
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("expert_parallel", [False, True])
+async def test_dp_pause_keep_then_resume(expert_parallel: bool):
+    """Start generation, pause after a few tokens (keep mode), resume; verify gap."""
+
+    pause_duration = 2.0
+    min_tokens_before_pause = 3
+
+    with ExitStack() as after:
+        engine_args = _get_dp_pause_engine_args(expert_parallel)
+        engine = AsyncLLM.from_engine_args(engine_args)
+        after.callback(engine.shutdown)
+
+        sampling_params = SamplingParams(max_tokens=15, ignore_eos=True)
+        token_times: list[tuple[int, float]] = []
+        pause_token_idx = 0
+
+        async def generator_task():
+            nonlocal pause_token_idx
+            out = None
+            async for output in engine.generate(
+                request_id="keep-resume-req",
+                prompt=DP_PAUSE_PROMPT,
+                sampling_params=sampling_params,
+            ):
+                token_count = len(output.outputs[0].token_ids)
+                token_times.append((token_count, time.monotonic()))
+                out = output
+            return out
+
+        async def controller_task():
+            nonlocal pause_token_idx
+            while len(token_times) < min_tokens_before_pause:
+                await asyncio.sleep(0.01)
+            await engine.pause_generation(mode="keep")
+            await asyncio.sleep(pause_duration)
+            pause_token_idx = len(token_times)
+            await engine.resume_generation()
+
+        gen_task = asyncio.create_task(generator_task())
+        ctrl_task = asyncio.create_task(controller_task())
+        final_output, _ = await asyncio.gather(gen_task, ctrl_task)
+
+        assert final_output is not None and final_output.finished
+        assert await engine.is_paused() is False
+        assert pause_token_idx >= min_tokens_before_pause
+        if pause_token_idx > 0 and pause_token_idx < len(token_times):
+            pause_gap = (
+                token_times[pause_token_idx][1] - token_times[pause_token_idx - 1][1]
+            )
+            assert pause_gap >= pause_duration * 0.8, (
+                f"Expected gap ~{pause_duration}s after pause, got {pause_gap:.3f}s"
+            )
+
+
+@pytest.mark.asyncio
+async def test_dp_pause_keep_race_staggered_engines():
+    """Race: send pause(keep) to engine 0, then add two requests,
+    then pause(keep) to engine 1. Ensures no deadlock when pause
+    requests are staggered and requests arrive in between."""
+    if DP_SIZE != 2:
+        pytest.skip("test_dp_pause_keep_race_staggered_engines requires DP_SIZE=2")
+
+    with ExitStack() as after:
+        engine_args = _get_dp_pause_engine_args(expert_parallel=True)
+        engine = AsyncLLM.from_engine_args(engine_args)
+        after.callback(engine.shutdown)
+
+        client = engine.engine_core
+
+        original_call_utility = client.call_utility_async
+        mid_pause_tasks: list[asyncio.Task] = []
+
+        async def staggered_pause_keep(method: str, *args) -> Any:
+            if method != "pause_scheduler" or not args or args[0] != "keep":
+                return await original_call_utility(method, *args)
+            # Fire pause(keep) to engine 0 (don't await — with DP
+            # two-phase pause, consensus requires all ranks).
+            pause_0 = asyncio.create_task(
+                client._call_utility_async(method, *args, engine=client.core_engines[0])
+            )
+            # Let the event loop send the message to engine 0.
+            await asyncio.sleep(0.5)
+            # In the middle: send two requests (race window)
+            sp = SamplingParams(max_tokens=5, ignore_eos=True)
+
+            async def consume_gen(req_id: str) -> None:
+                async for _ in engine.generate(
+                    request_id=req_id,
+                    prompt=DP_PAUSE_PROMPT,
+                    sampling_params=sp,
+                ):
+                    pass
+
+            t1 = asyncio.create_task(consume_gen("race-1"))
+            t2 = asyncio.create_task(consume_gen("race-2"))
+            mid_pause_tasks.extend([t1, t2])
+            await asyncio.sleep(3)
+            # Fire pause(keep) to engine 1, then await both so
+            # consensus can be reached.
+            pause_1 = asyncio.create_task(
+                client._call_utility_async(method, *args, engine=client.core_engines[1])
+            )
+            results = await asyncio.gather(pause_0, pause_1)
+            return results[0]
+
+        client.call_utility_async = staggered_pause_keep
+
+        await engine.pause_generation(mode="keep")
+        assert await engine.is_paused()
+        await engine.resume_generation()
+        assert not await engine.is_paused()
+        # Let the two requests we sent mid-pause complete
+        await asyncio.gather(*mid_pause_tasks)
+
+
+@pytest.mark.asyncio
+async def test_dp_pause_barrier_request_deadlock():
+    """
+    Test that start_dp_wave is ignored while paused.
+
+    Sequence:
+      1. Pause all engines (PAUSED_ALL).
+      2. Send barrier to engine 0 only — blocks in dist.barrier(dp_group).
+      3. Send a request routed to engine 1.
+      4. Wait for any (buggy) START_DP_WAVE propagation.
+      5. Send barrier to engine 1 — completes in fixed code, deadlocks
+         in buggy code because engine 1 is stuck in EP all-to-all.
+    """
+    if DP_SIZE != 2:
+        pytest.skip("requires DP_SIZE=2")
+
+    with ExitStack() as after:
+        engine_args = _get_dp_pause_engine_args(expert_parallel=True)
+        engine = AsyncLLM.from_engine_args(engine_args)
+        after.callback(engine.shutdown)
+
+        client = engine.engine_core
+
+        # Cache get_supported_tasks so that generate() won't need to
+        # send a utility call to all engines (which would hang once
+        # engine 0 is blocked in the barrier).
+        await engine.get_supported_tasks()
+
+        # Pause all engines normally — no staggering.
+        await engine.pause_generation(mode="keep")
+        assert await engine.is_paused()
+
+        original_call_utility = client.call_utility_async
+        mid_barrier_tasks: list[asyncio.Task] = []
+
+        async def staggered_barrier(method: str, *args) -> Any:
+            if method != "barrier":
+                return await original_call_utility(method, *args)
+
+            # Send barrier to engine 0 only — it blocks in
+            # dist.barrier(dp_group) waiting for engine 1.
+            barrier_0 = asyncio.create_task(
+                client._call_utility_async(method, *args, engine=client.core_engines[0])
+            )
+            await asyncio.sleep(1)
+
+            # While engine 0 is blocked, send a request routed
+            # specifically to engine 1.
+            sp = SamplingParams(max_tokens=5, ignore_eos=True)
+
+            engine_1 = client.core_engines[1]
+            original_get_engine = client.get_core_engine_for_request
+
+            def route_to_engine_1(req):
+                client.reqs_in_flight[req.request_id] = engine_1
+                return engine_1
+
+            client.get_core_engine_for_request = route_to_engine_1
+
+            async def consume_gen(req_id: str) -> None:
+                async for _ in engine.generate(
+                    request_id=req_id,
+                    prompt=DP_PAUSE_PROMPT,
+                    sampling_params=sp,
+                ):
+                    pass
+
+            t1 = asyncio.create_task(consume_gen("race-1"))
+            mid_barrier_tasks.append(t1)
+
+            # Yield so generate() preprocessing completes and
+            # add_request_async is called (which, in buggy code,
+            # would send FIRST_REQ and wake engine 1).
+            for _ in range(200):
+                await asyncio.sleep(0)
+
+            client.get_core_engine_for_request = original_get_engine
+
+            # Wait for any START_DP_WAVE to propagate and for
+            # engine 1 to potentially enter execute_dummy_batch.
+            await asyncio.sleep(5)
+
+            # Now send barrier to engine 1.  In buggy code engine 1
+            # is stuck in execute_dummy_batch (EP all-to-all) while
+            # engine 0 is stuck in dist.barrier(dp_group) — deadlock.
+            result = await client._call_utility_async(
+                method, *args, engine=client.core_engines[1]
+            )
+            await barrier_0
+            return result
+
+        client.call_utility_async = staggered_barrier
+
+        # Drive the staggered barrier.  Old code deadlocks here.
+        try:
+            await asyncio.wait_for(client.call_utility_async("barrier"), timeout=30)
+        except asyncio.TimeoutError:
+            for t in mid_barrier_tasks:
+                t.cancel()
+            pytest.fail(
+                "Staggered barrier deadlocked — FIRST_REQ sent while "
+                "paused caused collective-op mismatch between engines"
+            )
+
+        await engine.resume_generation()
+        assert not await engine.is_paused()
+        # Let the two requests we sent mid-barrier complete.
+        await asyncio.gather(*mid_barrier_tasks)

tests/v1/distributed/test_dbo.py

@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Test Dual Batch Overlap (DBO) with Data Parallelism + Expert Parallelism.
+
+DBO is specifically designed for DP+EP scenarios to hide communication latency
+by overlapping computation of two batches. This test validates that DBO works
+correctly with the DeepSeek-V2-Lite model using GSM8K evaluation.
+"""
+
+import pytest
+import torch
+
+from tests.evals.gsm8k.gsm8k_eval import evaluate_gsm8k
+from tests.utils import RemoteOpenAIServer
+from vllm.utils.import_utils import has_deep_ep
+
+# Detect Blackwell / B200 (compute capability 10.x)
+try:
+    if torch.cuda.is_available():
+        cap = torch.cuda.get_device_capability(0)
+        IS_BLACKWELL = cap[0] >= 10
+    else:
+        IS_BLACKWELL = False
+except Exception:
+    # Be conservative: if we can't detect, don't xfail by default
+    IS_BLACKWELL = False
+
+MODEL_NAME = "deepseek-ai/DeepSeek-V2-Lite-Chat"
+DP_SIZE = 2
+
+# GSM8K eval configuration
+NUM_QUESTIONS = 256  # Fast eval for CI; but must be large enough to hit dbo thresholds
+NUM_SHOTS = 5  # Few-shot examples
+MIN_ACCURACY = 0.62  # Expected 0.64 with 2% buffer (based on vLLM test data)
+
+# Increase max_num_seqs to trigger DBO for decode batches
+# With 64 seqs, decode batches should exceed the 32 token threshold
+MAX_NUM_SEQS = 64  # Increased from 16 to trigger decode DBO
+
+# DeepEP backends to test
+DEEPEP_BACKENDS = [
+    "deepep_low_latency",
+    "deepep_high_throughput",
+]
+
+
+@pytest.mark.skipif(not has_deep_ep(), reason="These tests require deep_ep to run")
+@pytest.mark.parametrize("all2all_backend", DEEPEP_BACKENDS)
+@pytest.mark.xfail(
+    IS_BLACKWELL,
+    reason=(
+        "Temporary: DBO accuracy unstable on Blackwell "
+        "(doesn't meet expectation of MIN_ACCURACY = 0.62)"
+    ),
+)
+def test_dbo_dp_ep_gsm8k(all2all_backend: str, num_gpus_available):
+    """
+    Test DBO with DP+EP using GSM8K evaluation.
+    """
+    required_gpus = DP_SIZE
+
+    if num_gpus_available < required_gpus:
+        pytest.skip(f"Need at least {required_gpus} GPUs (DP={DP_SIZE})")
+
+    # Server arguments for DBO + DP + EP
+    server_args = [
+        "--max-model-len",
+        "4096",
+        "--max-num-seqs",
+        str(MAX_NUM_SEQS),  # Use larger batch to trigger decode DBO
+        "--trust-remote-code",
+        # Note: Not using --enforce-eager to test DBO's alternate CUDA graph dispatching
+        "--data-parallel-size",
+        str(DP_SIZE),
+        "--enable-expert-parallel",
+        "--enable-dbo",
+        # Fix threshold so we know we trigger DBO
+        "--dbo-decode-token-threshold",
+        "16",
+        "--dbo-prefill-token-threshold",
+        "256",
+        "--all2all-backend",
+        all2all_backend,
+    ]
+
+    with RemoteOpenAIServer(
+        MODEL_NAME,
+        server_args,
+        max_wait_seconds=600,  # Allow time for model loading with DP+EP
+    ) as remote_server:
+        # Use host and port directly from RemoteOpenAIServer
+        host = f"http://{remote_server.host}"
+        port = remote_server.port
+
+        # Run GSM8K evaluation
+        results = evaluate_gsm8k(
+            num_questions=NUM_QUESTIONS,
+            num_shots=NUM_SHOTS,
+            host=host,
+            port=port,
+        )
+
+        # Validate accuracy is reasonable
+        accuracy = results["accuracy"]
+        assert accuracy >= MIN_ACCURACY, (
+            f"DBO+DP+EP accuracy too low ({all2all_backend}): "
+            f"{accuracy:.3f} < {MIN_ACCURACY:.3f} "
+        )

tests/v1/distributed/test_eagle_dp.py

@@ -0,0 +1,109 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import asyncio
+import os
+from contextlib import AsyncExitStack
+from dataclasses import replace
+
+import pytest
+
+from vllm import SamplingParams
+from vllm.engine.arg_utils import AsyncEngineArgs
+from vllm.platforms import current_platform
+from vllm.sampling_params import RequestOutputKind
+from vllm.v1.engine.async_llm import AsyncLLM
+
+DP_SIZE = int(os.getenv("DP_SIZE", 2))
+
+if current_platform.is_rocm():
+    ATTN_BACKENDS = ["ROCM_ATTN", "TRITON_ATTN", "FLEX_ATTENTION"]
+else:
+    ATTN_BACKENDS = ["FLASH_ATTN"]
+
+# On SM<90 (e.g., L4), batch invariance does not support CUDA graphs.
+# See https://github.com/vllm-project/vllm/pull/30018 and
+# tests/v1/determinism/utils.py for the documented limitation.
+IS_DEVICE_CAPABILITY_BELOW_90 = not current_platform.has_device_capability(90)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("attn_backend", ATTN_BACKENDS)
+@pytest.mark.xfail(
+    current_platform.is_rocm(),
+    reason="Test may fail on ROCm until batch invariance is enabled. "
+    "See: https://github.com/vllm-project/vllm/issues/27433",
+    strict=False,
+)
+async def test_run_eagle_dp(monkeypatch: pytest.MonkeyPatch, attn_backend: str):
+    if not current_platform.is_rocm() and not current_platform.is_xpu():
+        # This test checks that running a model with and without eagle
+        # leads to identical tokens.
+        #
+        # NOTE: This is only true in batch invariant mode
+        # (because the target model verifies all draft tokens in one big
+        # forward pass)
+        #
+        # TODO[ROCm]: Test is passing on ROCm CI but may break in future.
+        # Enable batch invariance for ROCm when possible. See:
+        # https://github.com/vllm-project/vllm/issues/27433
+
+        monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1")
+
+    target_model = "meta-llama/Llama-3.1-8B-Instruct"
+    draft_model = "yuhuili/EAGLE-LLaMA3.1-Instruct-8B"
+
+    engine_args = AsyncEngineArgs(
+        model=target_model,
+        tokenizer_mode="auto",
+        enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
+        tensor_parallel_size=int(os.getenv("TP_SIZE", 1)),
+        data_parallel_size=DP_SIZE,
+        data_parallel_backend="mp",  # ray takes more time
+        trust_remote_code=True,
+        max_model_len=16384,
+        attention_config={"backend": attn_backend},
+    )
+
+    eagle_engine_args = replace(
+        engine_args,
+        speculative_config={
+            "model": draft_model,
+            "method": "eagle",
+            "num_speculative_tokens": 3,
+        },
+    )
+
+    prompt = "This is a test of data parallel with eagle"
+    num_expected_tokens = 100
+    sampling_params = SamplingParams(
+        max_tokens=num_expected_tokens,
+        ignore_eos=True,
+        output_kind=RequestOutputKind.FINAL_ONLY,
+        temperature=0,
+    )
+
+    async def generate_with_timeout(given_engine: AsyncLLM):
+        async for out in given_engine.generate(
+            request_id="test-eagle-dp", prompt=prompt, sampling_params=sampling_params
+        ):
+            token_ids = out.outputs[0].token_ids
+            assert len(token_ids) == num_expected_tokens
+            return token_ids
+
+    async def engine_create_and_generate(engine_args: AsyncEngineArgs):
+        async with AsyncExitStack() as after:
+            engine = AsyncLLM.from_engine_args(engine_args)
+            after.callback(engine.shutdown)
+
+            token_ids = await asyncio.wait_for(
+                generate_with_timeout(engine), timeout=30
+            )
+
+            assert not engine.output_processor.has_unfinished_requests()
+        return token_ids
+
+    token_ids_with_eagle = await engine_create_and_generate(eagle_engine_args)
+    token_ids_no_eagle = await engine_create_and_generate(engine_args)
+
+    # Test for correctness
+    assert token_ids_with_eagle == token_ids_no_eagle

tests/v1/distributed/test_external_lb_dp.py

@@ -0,0 +1,358 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import asyncio
+import os
+import threading
+import time
+from contextlib import AsyncExitStack
+
+import openai  # use the official client for correctness check
+import pytest
+import pytest_asyncio
+import requests
+
+from tests.utils import RemoteOpenAIServer
+from vllm.platforms import current_platform
+
+MODEL_NAME = os.getenv("MODEL_NAME", "ibm-research/PowerMoE-3b")
+
+# Number of data parallel ranks for external LB testing
+DP_SIZE = int(os.getenv("DP_SIZE", "2"))
+# Default tensor parallel size to use
+TP_SIZE = int(os.getenv("TP_SIZE", "1"))
+
+
+class ExternalLBServerManager:
+    """Manages data parallel vLLM server instances for external
+    load balancer testing."""
+
+    def __init__(
+        self,
+        model_name: str,
+        dp_size: int,
+        api_server_count: int,
+        base_server_args: list,
+        tp_size: int = TP_SIZE,
+    ):
+        self.model_name = model_name
+        self.dp_size = dp_size
+        self.tp_size = tp_size
+        self.api_server_count = api_server_count
+        self.base_server_args = base_server_args
+        self.servers: list[tuple[RemoteOpenAIServer, list[str]]] = []
+        self.server_threads: list[threading.Thread] = []
+
+    def __enter__(self) -> list[tuple[RemoteOpenAIServer, list[str]]]:
+        """Start all server instances for external LB mode."""
+        for rank in range(self.dp_size):
+            # Create server args for this specific rank
+            server_args = self.base_server_args.copy()
+
+            # Add external LB specific arguments
+            server_args.extend(
+                [
+                    "--data-parallel-size",
+                    str(self.dp_size),
+                    "--data-parallel-rank",
+                    str(rank),
+                    "--data-parallel-size-local",
+                    "1",
+                    "--tensor-parallel-size",
+                    str(self.tp_size),
+                    "--port",
+                    str(8000 + rank),  # Different port for each rank
+                    "--api-server-count",
+                    str(self.api_server_count),
+                ]
+            )
+
+            # Use a thread to start each server to allow parallel initialization
+            def start_server(r: int, sargs: list[str]):
+                try:
+                    # Start the server
+                    server = RemoteOpenAIServer(
+                        self.model_name,
+                        sargs,
+                        auto_port=False,
+                        env_dict={
+                            "VLLM_SERVER_DEV_MODE": "1",
+                            current_platform.device_control_env_var: ",".join(
+                                str(current_platform.device_id_to_physical_device_id(i))
+                                for i in range(r * TP_SIZE, (r + 1) * TP_SIZE)
+                            ),
+                        },
+                    )
+                    server.__enter__()
+                    print(
+                        f"Server rank {r} started successfully with "
+                        f"{self.api_server_count} API servers"
+                    )
+                    self.servers.append((server, sargs))
+                except Exception as e:
+                    print(f"Failed to start server rank {r}: {e}")
+                    raise
+
+            thread = threading.Thread(target=start_server, args=(rank, server_args))
+            thread.start()
+
+            self.server_threads.append(thread)
+
+        # Wait for all servers to start
+        for thread in self.server_threads:
+            thread.join()
+
+        # Give servers additional time to fully initialize and coordinate
+        time.sleep(2)
+
+        if len(self.servers) != self.dp_size:
+            raise Exception("Servers failed to start")
+
+        return self.servers
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Stop all server instances."""
+        servers = [s for s, _ in self.servers]
+        self.servers.clear()
+        try:
+            RemoteOpenAIServer.shutdown_many(servers)
+        except Exception as e:
+            print(f"Error stopping servers: {e}")
+
+
+@pytest.fixture(scope="module")
+def default_server_args():
+    return [
+        # use half precision for speed and memory savings in CI environment
+        "--dtype",
+        "bfloat16",
+        "--max-model-len",
+        "2048",
+        "--max-num-seqs",
+        "128",
+        "--enforce-eager",
+    ]
+
+
+@pytest.fixture(scope="module", params=[1, 4])
+def server_manager(request, default_server_args):
+    api_server_count = request.param
+    server_manager = ExternalLBServerManager(
+        MODEL_NAME, DP_SIZE, api_server_count, default_server_args
+    )
+
+    with server_manager:
+        yield server_manager
+
+
+@pytest.fixture
+def servers(server_manager):
+    return server_manager.servers
+
+
+@pytest_asyncio.fixture
+async def clients(servers: list[tuple[RemoteOpenAIServer, list[str]]]):
+    # Create a client for each server
+    async with AsyncExitStack() as stack:
+        yield [
+            await stack.enter_async_context(server.get_async_client())
+            for server, _ in servers
+        ]
+
+
+def _get_parallel_config(server: RemoteOpenAIServer):
+    response = requests.get(server.url_for("server_info?config_format=json"))
+    response.raise_for_status()
+
+    vllm_config = response.json()["vllm_config"]
+    return vllm_config["parallel_config"]
+
+
+def test_external_lb_server_info(server_manager):
+    servers = server_manager.servers
+    api_server_count = server_manager.api_server_count
+
+    for i, (server, _) in enumerate(servers):
+        print(f"Testing {i=}")
+
+        # Each request will hit one of the API servers
+        # `n_reqs` is set so that there is a good chance each server
+        # receives at least one request
+        n_reqs = 2 * api_server_count * api_server_count
+        parallel_configs = [_get_parallel_config(server) for _ in range(n_reqs)]
+        api_process_counts = [c["_api_process_count"] for c in parallel_configs]
+        api_process_ranks = [c["_api_process_rank"] for c in parallel_configs]
+
+        assert all(c == api_server_count for c in api_process_counts), (
+            api_process_counts
+        )
+        assert all(0 <= r < api_server_count for r in api_process_ranks), (
+            api_process_ranks
+        )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "model_name",
+    [MODEL_NAME],
+)
+async def test_external_lb_single_completion(
+    clients: list[openai.AsyncOpenAI],
+    servers: list[tuple[RemoteOpenAIServer, list[str]]],
+    model_name: str,
+) -> None:
+    async def make_request(client: openai.AsyncOpenAI):
+        completion = await client.completions.create(
+            model=model_name, prompt="Hello, my name is", max_tokens=10, temperature=1.0
+        )
+
+        assert completion.id is not None
+        assert completion.choices is not None and len(completion.choices) == 1
+
+        choice = completion.choices[0]
+        # The exact number of tokens can vary slightly with temperature=1.0,
+        # so we check for a reasonable minimum length.
+        assert len(choice.text) >= 1
+        # Finish reason might not always be 'length' if the model finishes early
+        # or due to other reasons, especially with high temperature.
+        # So, we'll accept 'length' or 'stop'.
+        assert choice.finish_reason in ("length", "stop")
+
+        # Token counts can also vary, so we check they are positive.
+        assert completion.usage.completion_tokens > 0
+        assert completion.usage.prompt_tokens > 0
+        assert completion.usage.total_tokens > 0
+        return completion
+
+    # Test single request to each server
+    for i, client in enumerate(clients):
+        result = await make_request(client)
+        assert result is not None
+        print(f"Server {i} handled single completion request successfully")
+
+    await asyncio.sleep(0.5)
+
+    # Send requests to all servers in round-robin fashion
+    num_requests_per_server = 25  # Total 50 requests across 2 servers
+    all_tasks = []
+
+    for i, client in enumerate(clients):
+        tasks = [make_request(client) for _ in range(num_requests_per_server)]
+        all_tasks.extend(tasks)
+
+    results = await asyncio.gather(*all_tasks)
+    assert len(results) == num_requests_per_server * len(clients)
+    assert all(completion is not None for completion in results)
+
+    await asyncio.sleep(0.5)
+
+    # Second burst of requests
+    all_tasks = []
+    for i, client in enumerate(clients):
+        tasks = [make_request(client) for _ in range(num_requests_per_server)]
+        all_tasks.extend(tasks)
+
+    results = await asyncio.gather(*all_tasks)
+    assert len(results) == num_requests_per_server * len(clients)
+    assert all(completion is not None for completion in results)
+
+    _, server_args = servers[0]
+    api_server_count = (
+        server_args.count("--api-server-count")
+        and server_args[server_args.index("--api-server-count") + 1]
+        or 1
+    )
+    print(
+        f"Successfully completed external LB test with {len(clients)} servers "
+        f"(API server count: {api_server_count})"
+    )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "model_name",
+    [MODEL_NAME],
+)
+async def test_external_lb_completion_streaming(
+    clients: list[openai.AsyncOpenAI],
+    servers: list[tuple[RemoteOpenAIServer, list[str]]],
+    model_name: str,
+) -> None:
+    prompt = "What is an LLM?"
+
+    async def make_streaming_request(client: openai.AsyncOpenAI):
+        # Perform a non-streaming request to get the expected full output
+        single_completion = await client.completions.create(
+            model=model_name,
+            prompt=prompt,
+            max_tokens=5,
+            temperature=0.0,
+        )
+        single_output = single_completion.choices[0].text
+
+        # Perform the streaming request
+        stream = await client.completions.create(
+            model=model_name, prompt=prompt, max_tokens=5, temperature=0.0, stream=True
+        )
+        chunks: list[str] = []
+        finish_reason_count = 0
+        last_chunk = None
+        async for chunk in stream:
+            chunks.append(chunk.choices[0].text)
+            if chunk.choices[0].finish_reason is not None:
+                finish_reason_count += 1
+            last_chunk = chunk  # Keep track of the last chunk
+
+        # finish reason should only return in the last block for OpenAI API
+        assert finish_reason_count == 1, "Finish reason should appear exactly once."
+        assert last_chunk is not None, "Stream should have yielded at least one chunk."
+        assert last_chunk.choices[0].finish_reason == "length", (
+            "Finish reason should be 'length'."
+        )
+        # Check that the combined text matches the non-streamed version.
+        assert "".join(chunks) == single_output, (
+            "Streamed output should match non-streamed output."
+        )
+        return True  # Indicate success for this request
+
+    # Test single request to each server
+    for i, client in enumerate(clients):
+        result = await make_streaming_request(client)
+        assert result is not None
+        print(f"Server {i} handled single streaming request successfully")
+
+    await asyncio.sleep(0.5)
+
+    # Send streaming requests to all servers in round-robin fashion
+    num_requests_per_server = 25  # Total 50 requests across 2 servers
+    all_tasks = []
+
+    for i, client in enumerate(clients):
+        tasks = [make_streaming_request(client) for _ in range(num_requests_per_server)]
+        all_tasks.extend(tasks)
+
+    results = await asyncio.gather(*all_tasks)
+    assert len(results) == num_requests_per_server * len(clients)
+    assert all(results), "Not all streaming requests completed successfully."
+
+    await asyncio.sleep(0.5)
+
+    # Second burst of streaming requests
+    all_tasks = []
+    for i, client in enumerate(clients):
+        tasks = [make_streaming_request(client) for _ in range(num_requests_per_server)]
+        all_tasks.extend(tasks)
+
+    results = await asyncio.gather(*all_tasks)
+    assert len(results) == num_requests_per_server * len(clients)
+    assert all(results), "Not all streaming requests completed successfully."
+
+    _, server_args = servers[0]
+    api_server_count = (
+        server_args.count("--api-server-count")
+        and server_args[server_args.index("--api-server-count") + 1]
+        or 1
+    )
+    print(
+        f"Successfully completed external LB streaming test with "
+        f"{len(clients)} servers (API server count: {api_server_count})"
+    )

tests/v1/distributed/test_hybrid_lb_dp.py

@@ -0,0 +1,399 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import asyncio
+import os
+import threading
+import time
+from contextlib import AsyncExitStack
+
+import openai  # use the official client for correctness check
+import pytest
+import pytest_asyncio
+import requests
+
+from tests.utils import RemoteOpenAIServer
+from tests.v1.utils import check_request_balancing
+from vllm.platforms import current_platform
+
+MODEL_NAME = "ibm-research/PowerMoE-3b"
+
+# Number of data parallel ranks for hybrid LB testing (4 total)
+DP_SIZE = int(os.getenv("DP_SIZE", "4"))
+# Default tensor parallel size to use
+TP_SIZE = int(os.getenv("TP_SIZE", "1"))
+
+# Number of nodes (2 nodes, each with 2 DP ranks)
+NUM_NODES = 2
+DP_SIZE_LOCAL = DP_SIZE // NUM_NODES  # 2 ranks per node
+
+
+class HybridLBServerManager:
+    """Manages hybrid data parallel vLLM server instances where each node
+    runs a single logical API server that balances requests only to the
+    DP engines running on that same node."""
+
+    def __init__(
+        self,
+        model_name: str,
+        dp_size: int,
+        api_server_count: int,
+        base_server_args: list,
+        dp_size_local: int = DP_SIZE_LOCAL,
+        tp_size: int = TP_SIZE,
+    ):
+        self.model_name = model_name
+        self.dp_size = dp_size
+        self.dp_size_local = dp_size_local
+        self.tp_size = tp_size
+        self.api_server_count = api_server_count
+        self.base_server_args = base_server_args
+        self.servers: list[tuple[RemoteOpenAIServer, list[str]]] = []
+        self.server_threads: list[threading.Thread] = []
+        self.num_nodes = dp_size // dp_size_local
+
+    def __enter__(self) -> list[tuple[RemoteOpenAIServer, list[str]]]:
+        """Start all server instances for hybrid LB mode."""
+        for node_id in range(self.num_nodes):
+            # Create server args for this specific node
+            server_args = self.base_server_args.copy()
+
+            # Calculate start rank for this node
+            start_rank = node_id * self.dp_size_local
+
+            # Add hybrid LB specific arguments
+            server_args.extend(
+                [
+                    "--data-parallel-size",
+                    str(self.dp_size),
+                    "--data-parallel-size-local",
+                    str(self.dp_size_local),
+                    "--data-parallel-start-rank",
+                    str(start_rank),
+                    "--data-parallel-hybrid-lb",  # Enable hybrid LB mode
+                    "--tensor-parallel-size",
+                    str(self.tp_size),
+                    "--port",
+                    str(8000 + node_id),  # Different port for each node
+                    "--api-server-count",
+                    str(self.api_server_count),
+                    "--data-parallel-address",
+                    "127.0.0.1",
+                    "--data-parallel-rpc-port",
+                    "13345",
+                ]
+            )
+
+            # Use a thread to start each server to allow parallel initialization
+            def start_server(node: int, sargs: list[str]):
+                try:
+                    # Calculate GPU devices for this node
+                    gpus_per_node = self.dp_size_local * self.tp_size
+                    gpu_start = node * gpus_per_node
+                    gpu_end = gpu_start + gpus_per_node
+
+                    # Start the server
+                    server = RemoteOpenAIServer(
+                        self.model_name,
+                        sargs,
+                        auto_port=False,
+                        env_dict={
+                            "VLLM_SERVER_DEV_MODE": "1",
+                            current_platform.device_control_env_var: ",".join(
+                                str(current_platform.device_id_to_physical_device_id(i))
+                                for i in range(gpu_start, gpu_end)
+                            ),
+                        },
+                    )
+                    server.__enter__()
+                    print(
+                        f"Hybrid LB node {node} started successfully with "
+                        f"{self.dp_size_local} local DP ranks and "
+                        f"{self.api_server_count} API servers"
+                    )
+                    self.servers.append((server, sargs))
+                except Exception as e:
+                    print(f"Failed to start hybrid LB node {node}: {e}")
+                    raise
+
+            thread = threading.Thread(target=start_server, args=(node_id, server_args))
+            thread.start()
+
+            self.server_threads.append(thread)
+
+        # Wait for all servers to start
+        for thread in self.server_threads:
+            thread.join()
+
+        # Give servers additional time to fully initialize and coordinate
+        time.sleep(3)
+
+        if len(self.servers) != self.num_nodes:
+            raise Exception("Servers failed to start")
+
+        return self.servers
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Stop all server instances."""
+        servers = [s for s, _ in self.servers]
+        self.servers.clear()
+        try:
+            RemoteOpenAIServer.shutdown_many(servers)
+        except Exception as e:
+            print(f"Error stopping servers: {e}")
+
+
+@pytest.fixture(scope="module")
+def default_server_args():
+    return [
+        # use half precision for speed and memory savings in CI environment
+        "--dtype",
+        "bfloat16",
+        "--max-model-len",
+        "2048",
+        "--max-num-seqs",
+        "128",
+        "--enforce-eager",
+    ]
+
+
+@pytest.fixture(scope="module", params=[1, 4])
+def server_manager(request, default_server_args):
+    api_server_count = request.param
+    server_manager = HybridLBServerManager(
+        MODEL_NAME,
+        DP_SIZE,
+        api_server_count,
+        default_server_args,
+        DP_SIZE_LOCAL,
+        TP_SIZE,
+    )
+
+    with server_manager:
+        yield server_manager
+
+
+@pytest.fixture
+def servers(server_manager):
+    return server_manager.servers
+
+
+@pytest_asyncio.fixture
+async def clients(servers: list[tuple[RemoteOpenAIServer, list[str]]]):
+    # Create a client for each node (each node has its own API endpoint)
+    async with AsyncExitStack() as stack:
+        yield [
+            await stack.enter_async_context(server.get_async_client())
+            for server, _ in servers
+        ]
+
+
+def _get_parallel_config(server: RemoteOpenAIServer):
+    response = requests.get(server.url_for("server_info?config_format=json"))
+    response.raise_for_status()
+
+    vllm_config = response.json()["vllm_config"]
+    return vllm_config["parallel_config"]
+
+
+def test_hybrid_dp_server_info(server_manager):
+    servers = server_manager.servers
+    api_server_count = server_manager.api_server_count
+
+    for i, (server, _) in enumerate(servers):
+        print(f"Testing {i=}")
+
+        # Each request will hit one of the API servers
+        # `n_reqs` is set so that there is a good chance each server
+        # receives at least one request
+        n_reqs = 2 * api_server_count * api_server_count
+        parallel_configs = [_get_parallel_config(server) for _ in range(n_reqs)]
+        api_process_counts = [c["_api_process_count"] for c in parallel_configs]
+        api_process_ranks = [c["_api_process_rank"] for c in parallel_configs]
+
+        assert all(c == api_server_count for c in api_process_counts), (
+            api_process_counts
+        )
+        assert all(0 <= r < api_server_count for r in api_process_ranks), (
+            api_process_ranks
+        )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "model_name",
+    [MODEL_NAME],
+)
+async def test_hybrid_lb_completion(
+    clients: list[openai.AsyncOpenAI],
+    servers: list[tuple[RemoteOpenAIServer, list[str]]],
+    model_name: str,
+) -> None:
+    async def make_request(client: openai.AsyncOpenAI):
+        completion = await client.completions.create(
+            model=model_name, prompt="Hello, my name is", max_tokens=5, temperature=1.0
+        )
+
+        assert completion.id is not None
+        assert completion.choices is not None and len(completion.choices) == 1
+
+        choice = completion.choices[0]
+        # The exact number of tokens can vary slightly with temperature=1.0,
+        # so we check for a reasonable minimum length.
+        assert len(choice.text) >= 1
+        # Finish reason might not always be 'length' if the model finishes early
+        # or due to other reasons, especially with high temperature.
+        # So, we'll accept 'length' or 'stop'.
+        assert choice.finish_reason in ("length", "stop")
+
+        # Token counts can also vary, so we check they are positive.
+        assert completion.usage.completion_tokens > 0
+        assert completion.usage.prompt_tokens > 0
+        assert completion.usage.total_tokens > 0
+        return completion
+
+    # Test single request to each node
+    for i, client in enumerate(clients):
+        result = await make_request(client)
+        assert result is not None
+        print(f"Hybrid LB node {i} handled single completion request successfully")
+
+    await asyncio.sleep(0.5)
+
+    # Send requests to all nodes - each should balance within its local DP ranks
+    num_requests = 200  # Total 200 requests across 2 nodes
+    all_tasks = []
+    for i in range(num_requests):
+        client = clients[i % len(clients)]
+        all_tasks.append(asyncio.create_task(make_request(client)))
+        await asyncio.sleep(0.01)
+
+    results = await asyncio.gather(*all_tasks)
+    assert len(results) == num_requests
+    assert all(completion is not None for completion in results)
+
+    await asyncio.sleep(0.5)
+
+    # Second burst of requests
+    all_tasks = []
+    for i in range(num_requests):
+        client = clients[i % len(clients)]
+        all_tasks.append(asyncio.create_task(make_request(client)))
+        await asyncio.sleep(0.01)
+
+    results = await asyncio.gather(*all_tasks)
+    assert len(results) == num_requests
+    assert all(completion is not None for completion in results)
+
+    _, server_args = servers[0]
+    api_server_count = (
+        server_args.count("--api-server-count")
+        and server_args[server_args.index("--api-server-count") + 1]
+        or 1
+    )
+    print(
+        f"Successfully completed hybrid LB test with {len(clients)} nodes "
+        f"({DP_SIZE_LOCAL} DP ranks each, API server count: {api_server_count})"
+    )
+
+    # Check request balancing within each node
+    for i, (server, _) in enumerate(servers):
+        print(f"Checking request balancing for node {i}")
+        check_request_balancing(server, DP_SIZE_LOCAL)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "model_name",
+    [MODEL_NAME],
+)
+async def test_hybrid_lb_completion_streaming(
+    clients: list[openai.AsyncOpenAI],
+    servers: list[tuple[RemoteOpenAIServer, list[str]]],
+    model_name: str,
+) -> None:
+    prompt = "What is an LLM?"
+
+    async def make_streaming_request(client: openai.AsyncOpenAI):
+        # Perform a non-streaming request to get the expected full output
+        single_completion = await client.completions.create(
+            model=model_name,
+            prompt=prompt,
+            max_tokens=5,
+            temperature=0.0,
+        )
+        single_output = single_completion.choices[0].text
+
+        # Perform the streaming request
+        stream = await client.completions.create(
+            model=model_name, prompt=prompt, max_tokens=5, temperature=0.0, stream=True
+        )
+        chunks: list[str] = []
+        finish_reason_count = 0
+        last_chunk = None
+        async for chunk in stream:
+            chunks.append(chunk.choices[0].text)
+            if chunk.choices[0].finish_reason is not None:
+                finish_reason_count += 1
+            last_chunk = chunk  # Keep track of the last chunk
+
+        # finish reason should only return in the last block for OpenAI API
+        assert finish_reason_count == 1, "Finish reason should appear exactly once."
+        assert last_chunk is not None, "Stream should have yielded at least one chunk."
+        assert last_chunk.choices[0].finish_reason == "length", (
+            "Finish reason should be 'length'."
+        )
+        # Check that the combined text matches the non-streamed version.
+        assert "".join(chunks) == single_output, (
+            "Streamed output should match non-streamed output."
+        )
+        return True  # Indicate success for this request
+
+    # Test single request to each node
+    for i, client in enumerate(clients):
+        result = await make_streaming_request(client)
+        assert result is not None
+        print(f"Hybrid LB node {i} handled single streaming request successfully")
+
+    await asyncio.sleep(0.5)
+
+    # Send streaming requests to all nodes
+    num_requests = 200  # Total 200 requests across 2 nodes
+    all_tasks = []
+    for i in range(num_requests):
+        client = clients[i % len(clients)]
+        all_tasks.append(asyncio.create_task(make_streaming_request(client)))
+        await asyncio.sleep(0.01)
+
+    results = await asyncio.gather(*all_tasks)
+    assert len(results) == num_requests
+    assert all(results), "Not all streaming requests completed successfully."
+
+    await asyncio.sleep(0.5)
+
+    # Second burst of streaming requests
+    all_tasks = []
+    for i in range(num_requests):
+        client = clients[i % len(clients)]
+        all_tasks.append(asyncio.create_task(make_streaming_request(client)))
+        await asyncio.sleep(0.01)
+
+    results = await asyncio.gather(*all_tasks)
+    assert len(results) == num_requests
+    assert all(results), "Not all streaming requests completed successfully."
+
+    _, server_args = servers[0]
+    api_server_count = (
+        server_args.count("--api-server-count")
+        and server_args[server_args.index("--api-server-count") + 1]
+        or 1
+    )
+    print(
+        f"Successfully completed hybrid LB streaming test with "
+        f"{len(clients)} nodes ({DP_SIZE_LOCAL} DP ranks each, "
+        f"API server count: {api_server_count})"
+    )
+
+    # Check request balancing within each node
+    for i, (server, _) in enumerate(servers):
+        print(f"Checking streaming request balancing for node {i}")
+        check_request_balancing(server, DP_SIZE_LOCAL)

tests/v1/distributed/test_internal_lb_dp.py

@@ -0,0 +1,729 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import asyncio
+import os
+import threading
+import time
+import traceback
+from typing import cast
+
+import openai  # use the official client for correctness check
+import pytest
+import pytest_asyncio
+import requests
+
+from tests.utils import ROCM_ENV_OVERRIDES, RemoteOpenAIServer
+from tests.v1.utils import check_request_balancing
+from vllm.platforms import current_platform
+
+MODEL_NAME = "ibm-research/PowerMoE-3b"
+
+# Number of data parallel ranks for multi-node internal LB testing
+DP_SIZE = int(os.getenv("DP_SIZE", "2"))
+# Default tensor parallel size to use
+TP_SIZE = int(os.getenv("TP_SIZE", "1"))
+
+# Number of nodes to simulate
+NUM_NODES = 2
+
+
+async def _make_completion_request(
+    client: openai.AsyncOpenAI,
+    model_name: str,
+) -> openai.types.Completion:
+    """Make a single completion request and validate the response.
+
+    Uses temperature=1.0 to ensure diverse outputs across concurrent
+    requests for realistic load balancer testing.
+    """
+    completion = await client.completions.create(
+        model=model_name,
+        prompt="Hello, my name is",
+        max_tokens=5,
+        temperature=1.0,
+    )
+
+    assert completion.id is not None, (
+        f"Expected non-None completion id. usage={completion.usage!r}"
+    )
+    assert completion.choices is not None and len(completion.choices) == 1, (
+        f"Expected 1 choice, got "
+        f"{len(completion.choices) if completion.choices else 'None'}"
+    )
+
+    choice = completion.choices[0]
+    # With temperature=1.0, the model may emit a stop token immediately,
+    # producing empty text with finish_reason='stop'. This is valid
+    # model behavior - the test's purpose is load balancing, not output
+    # quality.
+    assert choice.finish_reason in ("length", "stop"), (
+        f"Expected finish_reason 'length' or 'stop', "
+        f"got {choice.finish_reason!r}. text={choice.text!r}"
+    )
+    if choice.finish_reason == "length":
+        assert len(choice.text) >= 1, (
+            f"Expected non-empty text with finish_reason='length', got {choice.text!r}"
+        )
+
+    assert completion.usage.prompt_tokens > 0, (
+        f"Expected positive prompt_tokens, got {completion.usage.prompt_tokens}"
+    )
+    assert completion.usage.total_tokens > 0, (
+        f"Expected positive total_tokens, got {completion.usage.total_tokens}"
+    )
+    return completion
+
+
+async def _run_request_bursts(
+    client: openai.AsyncOpenAI,
+    model_name: str,
+    num_requests: int = 200,
+    num_bursts: int = 2,
+):
+    """Send multiple bursts of completion requests and validate all succeed."""
+    for burst in range(num_bursts):
+        all_tasks = []
+        for _ in range(num_requests):
+            all_tasks.append(
+                asyncio.create_task(_make_completion_request(client, model_name))
+            )
+            await asyncio.sleep(0.01)
+
+        results = await asyncio.gather(*all_tasks, return_exceptions=True)
+        assert len(results) == num_requests, (
+            f"Burst {burst}: expected {num_requests} results, got {len(results)}"
+        )
+
+        for result in results:
+            if isinstance(result, BaseException):
+                raise result
+
+        assert all(completion is not None for completion in results), (
+            f"Burst {burst}: some completions were None"
+        )
+
+        await asyncio.sleep(0.5)
+
+
+class MultinodeInternalLBServerManager:
+    """Manages multi-node data parallel vLLM server instances for internal
+    load balancer testing using --headless mode."""
+
+    def __init__(
+        self,
+        model_name: str,
+        dp_size: int,
+        api_server_count: int,
+        base_server_args: list,
+        dp_per_node: int = 1,
+        tp_size: int = TP_SIZE,
+    ):
+        self.model_name = model_name
+        self.dp_size = dp_size
+        self.dp_per_node = dp_per_node
+        self.tp_size = tp_size
+        self.api_server_count = api_server_count
+        self.base_server_args = base_server_args
+        self.servers: list[tuple[RemoteOpenAIServer, list[str]] | None] = [None] * (
+            dp_size // dp_per_node
+        )
+        self.server_threads: list[threading.Thread] = []
+
+    def __enter__(self) -> list[tuple[RemoteOpenAIServer, list[str]]]:
+        """Start all server instances for multi-node internal LB mode."""
+        for server_idx, rank in enumerate(range(0, self.dp_size, self.dp_per_node)):
+            # Create server args for this specific rank
+            server_args = self.base_server_args.copy()
+
+            if rank == 0:
+                # Head node - runs API server and first DP rank
+                server_args.extend(
+                    [
+                        "--data-parallel-size",
+                        str(self.dp_size),
+                        "--data-parallel-size-local",
+                        str(self.dp_per_node),
+                        "--tensor-parallel-size",
+                        str(self.tp_size),
+                        "--port",
+                        "8000",  # Single endpoint for all requests
+                        "--api-server-count",
+                        str(self.api_server_count),
+                        "--data-parallel-address",
+                        "127.0.0.1",
+                        "--data-parallel-rpc-port",
+                        "13345",
+                    ]
+                )
+            else:
+                # Secondary nodes - run in headless mode
+                server_args.extend(
+                    [
+                        "--headless",
+                        "--data-parallel-size",
+                        str(self.dp_size),
+                        "--data-parallel-size-local",
+                        str(self.dp_per_node),
+                        "--data-parallel-start-rank",
+                        str(rank),
+                        "--tensor-parallel-size",
+                        str(self.tp_size),
+                        "--data-parallel-address",
+                        "127.0.0.1",
+                        "--data-parallel-rpc-port",
+                        "13345",
+                    ]
+                )
+
+            # Use a thread to start each server to allow parallel initialization
+            def start_server(sidx: int, r: int, sargs: list[str]):
+                gpus_per_node = self.tp_size * self.dp_per_node
+                try:
+                    # Start the server
+                    server = RemoteOpenAIServer(
+                        self.model_name,
+                        sargs,
+                        auto_port=False,
+                        env_dict={
+                            "VLLM_SERVER_DEV_MODE": "1",
+                            **ROCM_ENV_OVERRIDES,
+                            current_platform.device_control_env_var: ",".join(
+                                str(current_platform.device_id_to_physical_device_id(i))
+                                for i in range(r, r + gpus_per_node)
+                            ),
+                        },
+                    )
+                    server.__enter__()
+                    if r == 0:
+                        print(
+                            f"Head node (rank {r}) started successfully with "
+                            f"{self.api_server_count} API servers"
+                        )
+                    else:
+                        print(f"Headless node (rank {r}) started successfully")
+                    self.servers[sidx] = (server, sargs)
+                except Exception as e:
+                    print(f"Failed to start server rank {r}: {e}")
+                    traceback.print_exc()
+                    raise
+
+            thread = threading.Thread(
+                target=start_server, args=(server_idx, rank, server_args)
+            )
+            thread.start()
+
+            self.server_threads.append(thread)
+
+        # Wait for all servers to start
+        for thread in self.server_threads:
+            thread.join()
+
+        # Give servers additional time to fully initialize and coordinate
+        time.sleep(3)
+
+        if not all(self.servers):
+            raise Exception("Servers failed to start")
+
+        return cast(list[tuple[RemoteOpenAIServer, list[str]]], self.servers)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Stop all server instances."""
+        servers = [entry[0] for entry in self.servers if entry is not None]
+        self.servers.clear()
+        try:
+            RemoteOpenAIServer.shutdown_many(servers)
+        except Exception as e:
+            print(f"Error stopping servers: {e}")
+            traceback.print_exc()
+
+
+class APIOnlyServerManager:
+    """Manages API-only server (Node 0) and headless engines server (Node 1)
+    for testing separated API server and engine configuration."""
+
+    def __init__(
+        self,
+        model_name: str,
+        dp_size: int,
+        api_server_count: int,
+        base_server_args: list,
+        tp_size: int = TP_SIZE,
+    ):
+        self.model_name = model_name
+        self.dp_size = dp_size
+        self.tp_size = tp_size
+        self.api_server_count = api_server_count
+        self.base_server_args = base_server_args
+        self.servers: list[tuple[RemoteOpenAIServer, list[str]] | None] = [None] * 2
+        self.server_threads: list[threading.Thread] = []
+
+    def __enter__(self) -> list[tuple[RemoteOpenAIServer, list[str]]]:
+        """Start API-only server and headless engines server."""
+
+        # Start API-only server (Node 0) - no engines, only API server
+        api_server_args = self.base_server_args.copy()
+        api_server_args.extend(
+            [
+                "--data-parallel-size",
+                str(self.dp_size),
+                "--data-parallel-size-local",
+                "0",  # No engines on this node
+                "--tensor-parallel-size",
+                str(self.tp_size),
+                "--port",
+                "8000",
+                "--api-server-count",
+                str(self.api_server_count),
+                "--data-parallel-address",
+                "127.0.0.1",
+                "--data-parallel-rpc-port",
+                "13345",
+            ]
+        )
+
+        # Start headless engines server (Node 1) - all engines, no API server
+        engines_server_args = self.base_server_args.copy()
+        engines_server_args.extend(
+            [
+                "--headless",
+                "--data-parallel-size",
+                str(self.dp_size),
+                "--data-parallel-size-local",
+                str(self.dp_size),  # All engines on this node
+                "--tensor-parallel-size",
+                str(self.tp_size),
+                "--data-parallel-address",
+                "127.0.0.1",
+                "--data-parallel-rpc-port",
+                "13345",
+            ]
+        )
+
+        # Use threads to start both servers in parallel
+        def start_api_server():
+            try:
+                server = RemoteOpenAIServer(
+                    self.model_name,
+                    api_server_args,
+                    auto_port=False,
+                    env_dict={
+                        "VLLM_SERVER_DEV_MODE": "1",
+                        **ROCM_ENV_OVERRIDES,
+                        # No GPUs needed for API-only server
+                    },
+                )
+                server.__enter__()
+                print(
+                    f"API-only server started successfully with "
+                    f"{self.api_server_count} API servers"
+                )
+                self.servers[0] = (server, api_server_args)
+            except Exception as e:
+                print(f"Failed to start API-only server: {e}")
+                raise
+
+        def start_engines_server():
+            try:
+                server = RemoteOpenAIServer(
+                    self.model_name,
+                    engines_server_args,
+                    auto_port=False,
+                    env_dict={
+                        **ROCM_ENV_OVERRIDES,
+                        current_platform.device_control_env_var: ",".join(
+                            str(current_platform.device_id_to_physical_device_id(i))
+                            for i in range(self.dp_size * self.tp_size)
+                        ),
+                    },
+                )
+                server.__enter__()
+                print(
+                    f"Headless engines server started successfully with "
+                    f"{self.dp_size} engines"
+                )
+                self.servers[1] = (server, engines_server_args)
+            except Exception as e:
+                print(f"Failed to start headless engines server: {e}")
+                raise
+
+        # Start API server first
+        api_thread = threading.Thread(target=start_api_server)
+        api_thread.start()
+        self.server_threads.append(api_thread)
+
+        # Start engines server second
+        engines_thread = threading.Thread(target=start_engines_server)
+        engines_thread.start()
+        self.server_threads.append(engines_thread)
+
+        # Wait for both servers to start
+        for thread in self.server_threads:
+            thread.join()
+
+        # Give servers additional time to fully initialize and coordinate
+        time.sleep(3)
+
+        if not all(self.servers):
+            raise Exception("Both servers failed to start")
+
+        return cast(list[tuple[RemoteOpenAIServer, list[str]]], self.servers)
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Stop both server instances."""
+        servers = [entry[0] for entry in self.servers if entry is not None]
+        self.servers.clear()
+        try:
+            RemoteOpenAIServer.shutdown_many(servers)
+        except Exception as e:
+            print(f"Error stopping servers: {e}")
+            traceback.print_exc()
+
+
+@pytest.fixture(scope="module")
+def default_server_args():
+    return [
+        # use half precision for speed and memory savings in CI environment
+        "--dtype",
+        "bfloat16",
+        "--max-model-len",
+        "2048",
+        "--max-num-seqs",
+        "128",
+        "--enforce-eager",
+    ]
+
+
+@pytest.fixture(scope="module", params=[1, 4])
+def server_manager(request, default_server_args):
+    api_server_count = request.param
+    server_manager = MultinodeInternalLBServerManager(
+        MODEL_NAME,
+        DP_SIZE,
+        api_server_count,
+        default_server_args,
+        DP_SIZE // NUM_NODES,
+        TP_SIZE,
+    )
+
+    with server_manager:
+        yield server_manager
+
+
+@pytest.fixture
+def servers(server_manager):
+    return server_manager.servers
+
+
+@pytest.fixture(scope="module", params=[1, 4])
+def api_only_servers(request, default_server_args):
+    """Fixture for API-only server + headless engines configuration."""
+    api_server_count = request.param
+    with APIOnlyServerManager(
+        MODEL_NAME, DP_SIZE, api_server_count, default_server_args, TP_SIZE
+    ) as server_list:
+        yield server_list
+
+
+@pytest_asyncio.fixture
+async def client(servers: list[tuple[RemoteOpenAIServer, list[str]]]):
+    # For internal LB, we only connect to the head node (rank 0)
+    # which provides the single API endpoint
+    head_server = servers[0][0]
+    async with head_server.get_async_client() as client:
+        yield client
+
+
+@pytest_asyncio.fixture
+async def api_only_client(api_only_servers: list[tuple[RemoteOpenAIServer, list[str]]]):
+    """Client fixture for API-only server configuration."""
+    # Connect to the API-only server (first server in the list)
+    api_server = api_only_servers[0][0]
+    async with api_server.get_async_client() as client:
+        yield client
+
+
+def _get_parallel_config(server: RemoteOpenAIServer):
+    response = requests.get(server.url_for("server_info?config_format=json"))
+    response.raise_for_status()
+
+    vllm_config = response.json()["vllm_config"]
+    return vllm_config["parallel_config"]
+
+
+def test_multinode_dp_server_info(server_manager):
+    head_server = server_manager.servers[0][0]
+    api_server_count = server_manager.api_server_count
+
+    # Each request will hit one of the API servers
+    # `n_reqs` is set so that there is a good chance each server
+    # receives at least one request
+    n_reqs = 2 * api_server_count * api_server_count
+    parallel_configs = [_get_parallel_config(head_server) for _ in range(n_reqs)]
+    api_process_counts = [c["_api_process_count"] for c in parallel_configs]
+    api_process_ranks = [c["_api_process_rank"] for c in parallel_configs]
+
+    assert all(c == api_server_count for c in api_process_counts), api_process_counts
+    assert all(0 <= r < api_server_count for r in api_process_ranks), api_process_ranks
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "model_name",
+    [MODEL_NAME],
+)
+async def test_multinode_dp_completion(
+    client: openai.AsyncOpenAI,
+    servers: list[tuple[RemoteOpenAIServer, list[str]]],
+    model_name: str,
+) -> None:
+    # Test single request
+    result = await _make_completion_request(client, model_name)
+    assert result is not None
+    print("Multi-node internal LB handled single completion request successfully")
+
+    await asyncio.sleep(0.5)
+
+    # Send multiple bursts - internal LB should distribute across DP ranks
+    await _run_request_bursts(client, model_name)
+
+    _, server_args = servers[0]
+    api_server_count = (
+        server_args.count("--api-server-count")
+        and server_args[server_args.index("--api-server-count") + 1]
+        or 1
+    )
+    print(
+        f"Successfully completed multi-node internal LB test with "
+        f"{len(servers)} DP ranks (API server count: {api_server_count})"
+    )
+
+    # Check request balancing via Prometheus metrics
+    head_server = servers[0][0]
+    check_request_balancing(head_server, DP_SIZE)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "model_name",
+    [MODEL_NAME],
+)
+async def test_multinode_dp_completion_streaming(
+    client: openai.AsyncOpenAI,
+    servers: list[tuple[RemoteOpenAIServer, list[str]]],
+    model_name: str,
+) -> None:
+    prompt = "What is an LLM?"
+
+    async def make_streaming_request():
+        # Perform a non-streaming request to get the expected full output
+        single_completion = await client.completions.create(
+            model=model_name,
+            prompt=prompt,
+            max_tokens=5,
+            temperature=0.0,
+        )
+        single_output = single_completion.choices[0].text
+
+        # Perform the streaming request
+        stream = await client.completions.create(
+            model=model_name, prompt=prompt, max_tokens=5, temperature=0.0, stream=True
+        )
+        chunks: list[str] = []
+        finish_reason_count = 0
+        last_chunk = None
+        async for chunk in stream:
+            chunks.append(chunk.choices[0].text)
+            if chunk.choices[0].finish_reason is not None:
+                finish_reason_count += 1
+            last_chunk = chunk  # Keep track of the last chunk
+
+        # finish reason should only return in the last block for OpenAI API
+        assert finish_reason_count == 1, "Finish reason should appear exactly once."
+        assert last_chunk is not None, "Stream should have yielded at least one chunk."
+        assert last_chunk.choices[0].finish_reason == "length", (
+            "Finish reason should be 'length'."
+        )
+        # Check that the combined text matches the non-streamed version.
+        assert "".join(chunks) == single_output, (
+            "Streamed output should match non-streamed output."
+        )
+        return True  # Indicate success for this request
+
+    # Test single streaming request
+    result = await make_streaming_request()
+    assert result is not None
+    print("Multi-node internal LB handled single streaming request successfully")
+
+    await asyncio.sleep(0.5)
+
+    # Send multiple streaming requests - internal LB should distribute across
+    # DP ranks
+    num_requests = 200
+    all_tasks = []
+    for _ in range(num_requests):
+        all_tasks.append(asyncio.create_task(make_streaming_request()))
+        await asyncio.sleep(0.01)
+
+    results = await asyncio.gather(*all_tasks)
+    assert len(results) == num_requests
+    assert all(results), "Not all streaming requests completed successfully."
+
+    await asyncio.sleep(0.5)
+
+    # Second burst of streaming requests
+    all_tasks = []
+    for _ in range(num_requests):
+        all_tasks.append(asyncio.create_task(make_streaming_request()))
+        await asyncio.sleep(0.01)
+
+    results = await asyncio.gather(*all_tasks)
+    assert len(results) == num_requests
+    assert all(results), "Not all streaming requests completed successfully."
+
+    _, server_args = servers[0]
+    api_server_count = (
+        server_args.count("--api-server-count")
+        and server_args[server_args.index("--api-server-count") + 1]
+        or 1
+    )
+    print(
+        f"Successfully completed multi-node internal LB streaming test with "
+        f"{len(servers)} DP ranks (API server count: {api_server_count})"
+    )
+
+    # Check request balancing via Prometheus metrics
+    head_server = servers[0][0]
+    check_request_balancing(head_server, DP_SIZE)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "model_name",
+    [MODEL_NAME],
+)
+async def test_api_only_multinode_dp_completion(
+    api_only_client: openai.AsyncOpenAI,
+    api_only_servers: list[tuple[RemoteOpenAIServer, list[str]]],
+    model_name: str,
+) -> None:
+    """Test API-only server with all engines on separate headless server."""
+
+    # Test single request
+    result = await _make_completion_request(api_only_client, model_name)
+    assert result is not None
+    print("API-only server handled single completion request successfully")
+
+    await asyncio.sleep(0.5)
+
+    # Send multiple bursts - should be distributed across engines on
+    # headless server
+    await _run_request_bursts(api_only_client, model_name)
+
+    api_server, api_server_args = api_only_servers[0]
+    api_server_count = (
+        api_server_args.count("--api-server-count")
+        and api_server_args[api_server_args.index("--api-server-count") + 1]
+        or 1
+    )
+    print(
+        f"Successfully completed API-only multi-node test with {DP_SIZE} "
+        f"engines on headless server (API server count: {api_server_count})"
+    )
+
+    # Check request balancing via Prometheus metrics
+    check_request_balancing(api_server, DP_SIZE)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "model_name",
+    [MODEL_NAME],
+)
+async def test_api_only_multinode_dp_completion_streaming(
+    api_only_client: openai.AsyncOpenAI,
+    api_only_servers: list[tuple[RemoteOpenAIServer, list[str]]],
+    model_name: str,
+) -> None:
+    """Test API-only server streaming with all engines on separate
+    headless server."""
+    prompt = "What is an LLM?"
+
+    async def make_streaming_request():
+        # Perform a non-streaming request to get the expected full output
+        single_completion = await api_only_client.completions.create(
+            model=model_name,
+            prompt=prompt,
+            max_tokens=5,
+            temperature=0.0,
+        )
+        single_output = single_completion.choices[0].text
+
+        # Perform the streaming request
+        stream = await api_only_client.completions.create(
+            model=model_name, prompt=prompt, max_tokens=5, temperature=0.0, stream=True
+        )
+        chunks: list[str] = []
+        finish_reason_count = 0
+        last_chunk = None
+        async for chunk in stream:
+            chunks.append(chunk.choices[0].text)
+            if chunk.choices[0].finish_reason is not None:
+                finish_reason_count += 1
+            last_chunk = chunk  # Keep track of the last chunk
+
+        # finish reason should only return in the last block for OpenAI API
+        assert finish_reason_count == 1, "Finish reason should appear exactly once."
+        assert last_chunk is not None, "Stream should have yielded at least one chunk."
+        assert last_chunk.choices[0].finish_reason == "length", (
+            "Finish reason should be 'length'."
+        )
+        # Check that the combined text matches the non-streamed version.
+        assert "".join(chunks) == single_output, (
+            "Streamed output should match non-streamed output."
+        )
+        return True  # Indicate success for this request
+
+    # Test single streaming request
+    result = await make_streaming_request()
+    assert result is not None
+    print("API-only server handled single streaming request successfully")
+
+    await asyncio.sleep(0.5)
+
+    # Send multiple streaming requests - should be distributed across engines
+    num_requests = 200
+    all_tasks = []
+    for _ in range(num_requests):
+        all_tasks.append(asyncio.create_task(make_streaming_request()))
+        await asyncio.sleep(0.01)
+
+    results = await asyncio.gather(*all_tasks)
+    assert len(results) == num_requests
+    assert all(results), "Not all streaming requests completed successfully."
+
+    await asyncio.sleep(0.5)
+
+    # Second burst of streaming requests
+    all_tasks = []
+    for _ in range(num_requests):
+        all_tasks.append(asyncio.create_task(make_streaming_request()))
+        await asyncio.sleep(0.01)
+
+    results = await asyncio.gather(*all_tasks)
+    assert len(results) == num_requests
+    assert all(results), "Not all streaming requests completed successfully."
+
+    _, api_server_args = api_only_servers[0]
+    api_server_count = (
+        api_server_args.count("--api-server-count")
+        and api_server_args[api_server_args.index("--api-server-count") + 1]
+        or 1
+    )
+    print(
+        f"Successfully completed API-only streaming test with {DP_SIZE} "
+        f"engines on headless server (API server count: {api_server_count})"
+    )
+
+    # Check request balancing via Prometheus metrics
+    api_server = api_only_servers[0][0]
+    check_request_balancing(api_server, DP_SIZE)

tests/v1/distributed/test_pp_dp_v2.py

@@ -0,0 +1,178 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""V2 ModelRunner + pipeline parallel + data parallel integration tests.
+
+Covers the interaction between the V2 model runner's PP sampled-token
+broadcast and the DP per-step all-reduce across a few concurrency
+regimes. Requires 4 GPUs (DP=2, PP=2, TP=1) on CUDA.
+"""
+
+import asyncio
+import contextlib
+import os
+from contextlib import ExitStack
+
+import pytest
+
+from vllm import SamplingParams
+from vllm.engine.arg_utils import AsyncEngineArgs
+from vllm.platforms import current_platform
+from vllm.sampling_params import RequestOutputKind
+from vllm.v1.engine.async_llm import AsyncLLM
+
+PP_DP_MODEL = "ibm-research/PowerMoE-3b"  # smallest cached MoE that supports PP
+PROMPT = "This is a test of data parallel and pipeline parallel together"
+
+
+def _gpu_skip_reason() -> str | None:
+    if not current_platform.is_cuda():
+        return "requires CUDA"
+    n = current_platform.device_count()
+    if n < 4:
+        return f"requires 4 GPUs, got {n}"
+    return None
+
+
+_GPU_SKIP = _gpu_skip_reason()
+
+pytestmark = [
+    pytest.mark.skipif(
+        os.environ.get("VLLM_USE_V2_MODEL_RUNNER", "0") != "1",
+        reason="VLLM_USE_V2_MODEL_RUNNER=1 required",
+    ),
+    pytest.mark.skipif(_GPU_SKIP is not None, reason=_GPU_SKIP or ""),
+]
+
+
+def _engine_args(async_scheduling: bool) -> AsyncEngineArgs:
+    return AsyncEngineArgs(
+        model=PP_DP_MODEL,
+        pipeline_parallel_size=2,
+        data_parallel_size=2,
+        data_parallel_backend="mp",
+        tensor_parallel_size=1,
+        max_model_len=4096,
+        max_num_batched_tokens=2048,
+        max_num_seqs=256,
+        async_scheduling=async_scheduling,
+        enable_prefix_caching=False,
+        enforce_eager=False,
+        enable_expert_parallel=False,
+    )
+
+
+async def _generate(engine: AsyncLLM, prompt: str, max_tokens: int) -> int:
+    """Run one streaming completion and return the number of tokens it yielded."""
+    sampling_params = SamplingParams(
+        max_tokens=max_tokens,
+        ignore_eos=True,
+        output_kind=RequestOutputKind.DELTA,
+        temperature=0.0,
+    )
+    request_id = f"req-{id(prompt):x}-{max_tokens}"
+    total = 0
+    async for out in engine.generate(
+        request_id=request_id, prompt=prompt, sampling_params=sampling_params
+    ):
+        total += len(out.outputs[0].token_ids)
+    return total
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("async_scheduling", [True, False])
+async def test_pp_dp_v2_low_concurrency(async_scheduling: bool):
+    """A single in-flight request at a time, repeated, to exercise the
+    PP slot ring under empty batches between decodes."""
+    with ExitStack() as after:
+        engine = AsyncLLM.from_engine_args(_engine_args(async_scheduling))
+        after.callback(engine.shutdown)
+
+        for _ in range(4):
+            n = await _generate(engine, PROMPT, max_tokens=16)
+            assert n == 16
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("async_scheduling", [True, False])
+async def test_pp_dp_v2_mid_concurrency(async_scheduling: bool):
+    """64 concurrent requests, staggered, to exercise the steady-state
+    DP all-reduce + PP slot-ring path."""
+    with ExitStack() as after:
+        engine = AsyncLLM.from_engine_args(_engine_args(async_scheduling))
+        after.callback(engine.shutdown)
+
+        async def _one(i: int) -> int:
+            await asyncio.sleep(0.01 * i)  # stagger so DP load-balances
+            return await _generate(engine, f"{PROMPT} {i}", max_tokens=64)
+
+        results = await asyncio.gather(*[_one(i) for i in range(64)])
+        assert all(n == 64 for n in results), results
+
+
+@pytest.mark.asyncio
+async def test_pp_dp_v2_abort_mid_decode():
+    """Cancel half the in-flight requests mid-stream and confirm the
+    engine survives the abort storm."""
+
+    with ExitStack() as after:
+        engine = AsyncLLM.from_engine_args(_engine_args(async_scheduling=True))
+        after.callback(engine.shutdown)
+
+        async def _maybe_cancel(i: int):
+            sampling_params = SamplingParams(
+                max_tokens=64,
+                ignore_eos=True,
+                output_kind=RequestOutputKind.DELTA,
+                temperature=0.0,
+            )
+            request_id = f"abort-req-{i}"
+            count = 0
+            cancel_at = 4 if i % 2 == 0 else 64
+            async for out in engine.generate(
+                request_id=request_id,
+                prompt=f"{PROMPT} {i}",
+                sampling_params=sampling_params,
+            ):
+                count += len(out.outputs[0].token_ids)
+                if count >= cancel_at:
+                    break
+            return count, i
+
+        results = await asyncio.gather(*[_maybe_cancel(i) for i in range(32)])
+        for count, i in results:
+            if i % 2 == 0:
+                assert count >= 4
+            else:
+                assert count == 64
+
+        # Engine must still serve after the abort storm.
+        final = await _generate(engine, "post-abort warmup", max_tokens=8)
+        assert final == 8
+
+
+@pytest.mark.asyncio
+async def test_pp_dp_v2_pause_resume():
+    """Pause an engine with a request in flight, then resume and confirm
+    new requests still work."""
+
+    with ExitStack() as after:
+        engine = AsyncLLM.from_engine_args(_engine_args(async_scheduling=True))
+        after.callback(engine.shutdown)
+
+        # Start a long-running generation, let some decoding happen, then
+        # pause (abort mode) and confirm the in-flight task terminates.
+        inflight = asyncio.create_task(_generate(engine, PROMPT, max_tokens=128))
+        await asyncio.sleep(0.5)
+
+        assert not await engine.is_paused()
+        await engine.pause_generation(mode="abort")
+        assert await engine.is_paused()
+
+        with contextlib.suppress(Exception):
+            await inflight
+
+        await engine.resume_generation()
+        assert not await engine.is_paused()
+
+        n = await _generate(engine, PROMPT, max_tokens=8)
+        assert n == 8