Stabilize PP8+DSpark token transport and draft execution

tests/config/test_max_concurrent_batches_cap.py

@@ -0,0 +1,86 @@
+#!/usr/bin/env python3
+"""Focused test for the gated max_concurrent_batches cap (WEDGE A/B).
+
+Contract: with async scheduling + V2 runner, the cap to pp_size applies ONLY
+when VLLM_WEDGE_AB_CAP_CONCURRENT=1 AND pp_size>1 AND speculative method is
+dspark. Every other combination preserves pp_size+1 (or the V1/non-async
+values). Plain unittest: pytest is not installed in the serving venv.
+
+Run: <venv>/bin/python tests/config/test_max_concurrent_batches_cap.py
+"""
+import os
+import sys
+import unittest
+from types import SimpleNamespace
+from unittest import mock
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(
+    os.path.abspath(__file__)))))
+
+from vllm.config.vllm import VllmConfig  # noqa: E402
+
+
+def make_cfg(pp_size, async_sched=True, use_v2=True, method="dspark"):
+    cfg = VllmConfig.__new__(VllmConfig)
+    object.__setattr__(cfg, "parallel_config",
+                       SimpleNamespace(pipeline_parallel_size=pp_size))
+    object.__setattr__(cfg, "scheduler_config",
+                       SimpleNamespace(async_scheduling=async_sched))
+    object.__setattr__(
+        cfg, "speculative_config",
+        SimpleNamespace(method=method) if method is not None else None)
+    # use_v2_model_runner may be a property on the class; bypass via __dict__
+    # patching where possible, else patch the type attribute in tests.
+    return cfg, use_v2
+
+
+class TestCap(unittest.TestCase):
+    def _mcb(self, cfg, use_v2):
+        with mock.patch.object(type(cfg), "use_v2_model_runner",
+                               new_callable=mock.PropertyMock,
+                               return_value=use_v2):
+            return cfg.max_concurrent_batches
+
+    def test_env_on_pp8_dspark_capped(self):
+        cfg, v2 = make_cfg(8)
+        with mock.patch.dict(os.environ, {"VLLM_WEDGE_AB_CAP_CONCURRENT": "1"}):
+            self.assertEqual(self._mcb(cfg, v2), 8)
+
+    def test_env_off_pp8_dspark_uncapped(self):
+        cfg, v2 = make_cfg(8)
+        with mock.patch.dict(os.environ, {}, clear=False):
+            os.environ.pop("VLLM_WEDGE_AB_CAP_CONCURRENT", None)
+            self.assertEqual(self._mcb(cfg, v2), 9)
+
+    def test_env_on_pp1_uncapped(self):
+        # PP=1 must keep pp_size+1=2 (async overlap), even with env on.
+        cfg, v2 = make_cfg(1)
+        with mock.patch.dict(os.environ, {"VLLM_WEDGE_AB_CAP_CONCURRENT": "1"}):
+            self.assertEqual(self._mcb(cfg, v2), 2)
+
+    def test_env_on_non_dspark_uncapped(self):
+        cfg, v2 = make_cfg(8, method="eagle")
+        with mock.patch.dict(os.environ, {"VLLM_WEDGE_AB_CAP_CONCURRENT": "1"}):
+            self.assertEqual(self._mcb(cfg, v2), 9)
+
+    def test_env_on_no_spec_config_uncapped(self):
+        cfg, v2 = make_cfg(8, method=None)
+        with mock.patch.dict(os.environ, {"VLLM_WEDGE_AB_CAP_CONCURRENT": "1"}):
+            self.assertEqual(self._mcb(cfg, v2), 9)
+
+    def test_v1_runner_pp_gt1_returns_pp_size(self):
+        cfg, v2 = make_cfg(8, use_v2=False)
+        self.assertEqual(self._mcb(cfg, v2), 8)
+
+    def test_v1_runner_pp1_async_returns_2(self):
+        cfg, v2 = make_cfg(1, use_v2=False)
+        self.assertEqual(self._mcb(cfg, v2), 2)
+
+    def test_no_async_returns_pp_size(self):
+        cfg, v2 = make_cfg(8, async_sched=False)
+        with mock.patch.dict(os.environ, {"VLLM_WEDGE_AB_CAP_CONCURRENT": "1"}):
+            self.assertEqual(self._mcb(cfg, v2), 8)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

tests/distributed/test_comm_ops.py

@@ -238,6 +238,38 @@ def _make_group_for_unit_test(
     return g
 
 
+def test_make_sibling_cpu_group_is_distinct_and_preserves_membership(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    from vllm.distributed import utils as distributed_utils
+
+    calls: list[tuple[list[int], str, str | None, object]] = []
+
+    def fake_new_group(
+        ranks: list[int], *, backend: str, group_desc: str | None, timeout: object
+    ) -> str:
+        calls.append((ranks, backend, group_desc, timeout))
+        return f"group-{len(calls)}"
+
+    timeout = object()
+    monkeypatch.setattr(torch.distributed, "new_group", fake_new_group)
+    monkeypatch.setattr(
+        distributed_utils, "get_cpu_distributed_timeout_or_none", lambda: timeout
+    )
+
+    group = _make_group_for_unit_test(rank_in_group=0, world_size=2)
+    group.rank = 2
+    group.group_ranks = [[0, 2], [1, 3]]
+
+    sibling = group.make_sibling_cpu_group(group_desc="control")
+
+    assert sibling == "group-1"
+    assert calls == [
+        ([0, 2], "gloo", "control", timeout),
+        ([1, 3], "gloo", "control", timeout),
+    ]
+
+
 def test_irecv_tensor_dict_send_allgather_postprocess_binds_keys(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None: