[GG] fix(concurrency): isolate graph state and gate B12X MoE overlap

tests/distributed/test_b12x_fused_all_reduce.py

@@ -7,6 +7,7 @@
 import torch
 import torch.distributed as dist
 
+import vllm.envs as envs
 import vllm.ir.ops
 from tests.compile.backend import TestBackend
 from vllm.compilation.passes.fusion import allreduce_rms_fusion
@@ -26,6 +27,7 @@
 from vllm.distributed import tensor_model_parallel_all_reduce
 from vllm.distributed.device_communicators.custom_all_reduce import (
     CustomAllreduce,
+    _b12x_pcie_oneshot_limits,
     get_b12x_pcie_allreduce,
 )
 from vllm.distributed.parallel_state import get_tp_group, graph_capture
@@ -131,6 +133,14 @@ def test_b12x_fused_allreduce_zero_cutoff_disables_support() -> None:
     assert not custom_allreduce.supports_fused_add_rms_norm()
 
 
+def test_b12x_oneshot_defaults_to_stream_isolation(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.delenv("VLLM_PCIE_ONESHOT_SINGLE_CHANNEL", raising=False)
+
+    assert not envs.environment_variables["VLLM_PCIE_ONESHOT_SINGLE_CHANNEL"]()
+
+
 def test_b12x_channel_checkpoint_delegates_to_runtime() -> None:
     custom_allreduce, runtime = make_b12x_custom_allreduce(
         allreduce_max_size=64,
@@ -146,6 +156,22 @@ def test_b12x_channel_checkpoint_delegates_to_runtime() -> None:
     runtime.rollback_channels.assert_called_once_with(checkpoint)
 
 
+def test_b12x_oneshot_buffer_tracks_dispatch_limits(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setenv("VLLM_PCIE_ONESHOT_ALLREDUCE_MAX_SIZE", "64KB")
+    monkeypatch.setenv(
+        "VLLM_PCIE_ONESHOT_FUSED_ADD_RMS_NORM_MAX_SIZE",
+        "84KB",
+    )
+
+    assert _b12x_pcie_oneshot_limits() == (
+        64 * 1024,
+        84 * 1024,
+        84 * 1024,
+    )
+
+
 def test_b12x_fused_custom_op_dispatch(monkeypatch) -> None:
     custom_allreduce = MagicMock()
     custom_allreduce.try_fused_add_rms_norm.return_value = True

tests/distributed/test_dcp_a2a.py

@@ -10,6 +10,7 @@
 
 import importlib.util
 import math
+from contextlib import contextmanager
 from typing import Any
 
 import multiprocess as mp
@@ -542,6 +543,86 @@ def fake_all_reduce(tensor, *, op, group):
     assert captured["op"] == dist.ReduceOp.MAX
 
 
+def test_b12x_pool_uses_independent_stream_channels(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    from vllm.v1.attention.ops import dcp_alltoall
+
+    captured: dict[str, Any] = {}
+
+    class _FakePool:
+        @classmethod
+        def from_exchange_group(cls, **kwargs):
+            captured.update(kwargs)
+            return cls()
+
+        def for_stream(self):
+            captured["warmed"] = True
+
+    group = _FakeCPGroup(2, object())  # type: ignore[arg-type]
+    monkeypatch.setattr(dcp_alltoall, "_B12X_DCP_A2A_POOLS", {})
+    monkeypatch.setattr(dcp_alltoall, "_B12X_DCP_A2A_DISABLED", set())
+    monkeypatch.setattr(dcp_alltoall, "_load_b12x_dcp_a2a_pool", lambda: _FakePool)
+    monkeypatch.setattr(dcp_alltoall, "_b12x_dcp_init_failed", lambda *args: False)
+    monkeypatch.setattr(
+        dcp_alltoall.torch.cuda,
+        "is_current_stream_capturing",
+        lambda: False,
+    )
+
+    pool = dcp_alltoall._get_b12x_dcp_a2a_pool(
+        group,  # type: ignore[arg-type]
+        device=torch.device("cuda:0"),
+        total_heads=64,
+        head_dim=512,
+        query_head_dim=576,
+        max_batch_size=64,
+    )
+
+    assert pool is not None
+    assert captured["single_channel"] is False
+    assert captured["warmed"] is True
+
+
+def test_b12x_dcp_capture_selects_only_current_group_pools(monkeypatch):
+    from vllm.v1.attention.ops import dcp_alltoall
+
+    events = []
+
+    class _FakePool:
+        def __init__(self, name):
+            self.name = name
+
+        @contextmanager
+        def capture(self, *, stream):
+            events.append(("enter", self.name, stream))
+            try:
+                yield
+            finally:
+                events.append(("exit", self.name, stream))
+
+    device_group = object()
+    group = _FakeCPGroup(2, device_group)  # type: ignore[arg-type]
+    stream = object()
+    pools = {
+        (id(device_group), 0, 64, 512, 576, 64): _FakePool("output"),
+        (id(device_group), 0, 64, 576, 576, 64): _FakePool("query"),
+        (id(object()), 0, 64, 512, 576, 64): _FakePool("foreign"),
+    }
+    monkeypatch.setattr(dcp_alltoall, "_B12X_DCP_A2A_POOLS", pools)
+
+    with dcp_alltoall.capture_b12x_dcp_a2a(group, stream):  # type: ignore[arg-type]
+        events.append(("body", None, stream))
+
+    assert events == [
+        ("enter", "output", stream),
+        ("enter", "query", stream),
+        ("body", None, stream),
+        ("exit", "query", stream),
+        ("exit", "output", stream),
+    ]
+
+
 def test_b12x_dcp_channel_rollback_restores_existing_and_closes_new_pools(
     monkeypatch,
 ):
@@ -634,6 +715,44 @@ def __init__(self, *, world_size, communicator=None):
         ("rollback-dcp", "dcp-checkpoint"),
         ("rollback-tp", "tp-checkpoint"),
     ]
+
+
+def test_global_graph_capture_enters_b12x_dcp_pool(monkeypatch):
+    from vllm.distributed import parallel_state
+    from vllm.v1.attention.ops import dcp_alltoall
+
+    events = []
+
+    class _FakeGroup:
+        world_size = 2
+
+        @contextmanager
+        def graph_capture(self, context):
+            yield context
+
+    tp_group = _FakeGroup()
+    pp_group = _FakeGroup()
+    dcp_group = _FakeGroup()
+    stream = object()
+    context = parallel_state.GraphCaptureContext(stream)  # type: ignore[arg-type]
+
+    @contextmanager
+    def fake_b12x_capture(group, selected_stream):
+        events.append((group, selected_stream))
+        yield
+
+    monkeypatch.setattr(parallel_state, "_DCP", dcp_group)
+    monkeypatch.setattr(parallel_state, "get_tp_group", lambda: tp_group)
+    monkeypatch.setattr(parallel_state, "get_pp_group", lambda: pp_group)
+    monkeypatch.setattr(parallel_state, "get_dcp_group", lambda: dcp_group)
+    monkeypatch.setattr(dcp_alltoall, "capture_b12x_dcp_a2a", fake_b12x_capture)
+
+    with parallel_state.graph_capture(torch.device("cpu"), context) as actual:
+        assert actual is context
+
+    assert events == [(dcp_group, stream)]
+
+
 @pytest.mark.skipif(torch.accelerator.device_count() < 1, reason="CUDA is required.")
 def test_b12x_lse_reduce_honors_token_cap(monkeypatch: pytest.MonkeyPatch):
     from vllm.v1.attention.ops import dcp_alltoall

tests/model_executor/layers/fused_moe/test_shared_experts_stream.py

@@ -2,12 +2,53 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from contextlib import contextmanager
-from unittest.mock import Mock
+from types import SimpleNamespace
+from unittest.mock import Mock, call
 
 import torch
 
 import vllm.model_executor.layers.fused_moe.runner.shared_experts as shared_module
-from vllm.model_executor.layers.fused_moe.runner.shared_experts import SharedExperts
+from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner
+from vllm.model_executor.layers.fused_moe.runner.shared_experts import (
+    SharedExperts,
+    SharedExpertsOrder,
+)
+
+
+def test_aux_stream_respects_expert_kernel_capability(monkeypatch) -> None:
+    stream = Mock()
+    supports_aux_stream = Mock(side_effect=lambda num_tokens: num_tokens <= 8)
+    monkeypatch.setattr(shared_module, "aux_stream", Mock(return_value=stream))
+    monkeypatch.setattr(shared_module.envs, "VLLM_DISABLE_SHARED_EXPERTS_STREAM", False)
+    monkeypatch.setattr(
+        shared_module,
+        "current_platform",
+        SimpleNamespace(is_cuda=lambda: True),
+    )
+    moe_config = SimpleNamespace(
+        moe_parallel_config=SimpleNamespace(
+            enable_eplb=False,
+            all2all_backend="allgather_reducescatter",
+            use_fi_nvl_two_sided_kernels=False,
+        )
+    )
+    shared_experts = SharedExperts(
+        layer=Mock(),
+        moe_config=moe_config,
+        enable_dbo=False,
+        mk_can_overlap_shared_experts=Mock(return_value=False),
+        experts_support_aux_stream=supports_aux_stream,
+    )
+
+    assert (
+        shared_experts._determine_shared_experts_order(torch.empty(8, 16))
+        == SharedExpertsOrder.MULTI_STREAM_OVERLAPPED
+    )
+    assert (
+        shared_experts._determine_shared_experts_order(torch.empty(16, 16))
+        == SharedExpertsOrder.NO_OVERLAP
+    )
+    assert supports_aux_stream.call_args_list == [call(8), call(16)]
 
 
 def test_aux_stream_output_lifetime_extends_to_consumer(monkeypatch) -> None:
@@ -33,3 +74,30 @@ def use_stream(stream):
     shared_experts._layer.assert_called_once_with(shared_experts_input)
     consumer_stream.wait_stream.assert_called_once_with(aux_stream)
     output.record_stream.assert_called_once_with(consumer_stream)
+
+
+def test_aux_shared_experts_are_enqueued_before_resident_moe() -> None:
+    runner = object.__new__(MoERunner)
+    events: list[object] = []
+    fused_out = object()
+    shared_out = object()
+
+    runner._maybe_apply_shared_experts = lambda _input, order: events.append(order)
+    runner.routed_experts = SimpleNamespace(
+        quant_method=SimpleNamespace(is_monolithic=True),
+        forward_monolithic=lambda **_kwargs: events.append("routed") or fused_out,
+    )
+    runner._shared_experts = SimpleNamespace(output=shared_out)
+
+    result = runner._apply_quant_method(
+        hidden_states=Mock(),
+        router_logits=Mock(),
+        shared_experts_input=Mock(),
+    )
+
+    assert events == [
+        SharedExpertsOrder.NO_OVERLAP,
+        SharedExpertsOrder.MULTI_STREAM_OVERLAPPED,
+        "routed",
+    ]
+    assert result == (shared_out, fused_out)

tests/model_executor/layers/test_b12x_moe_warmup.py

@@ -89,6 +89,7 @@ def _make_fake_b12x_experts() -> b12x_moe.B12xExperts:
     experts._activation_amax_base_num_layers = None
     experts._activation_amax_state_key = None
     experts._activation_amax_layer_idx = None
+    experts._aux_stream_overlap_by_tokens = {}
     return experts
 
 
@@ -131,6 +132,39 @@ def test_non_b12x_moe_runner_keeps_generic_custom_op(monkeypatch) -> None:
     assert forward_entry._qualified_op_name == "vllm::moe_forward"
 
 
+def test_b12x_shared_expert_overlap_follows_launch_plan_capability(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    for name in (
+        "B12X_MOE_FORCE_A8",
+        "B12X_FORCE_MOE_A8",
+        "B12X_MOE_FORCE_A16",
+    ):
+        monkeypatch.delenv(name, raising=False)
+    plans = []
+    plan = object()
+
+    def fake_plan(**kwargs):
+        plans.append(kwargs)
+        return plan
+
+    monkeypatch.setattr(b12x_moe, "_plan_b12x_moe_execution", fake_plan)
+    monkeypatch.setattr(
+        b12x_moe,
+        "_b12x_moe_plan_supports_aux_stream_overlap",
+        lambda candidate: candidate is plan,
+    )
+    experts = _make_fake_b12x_experts()
+
+    assert experts.supports_shared_experts_aux_stream(8)
+    assert experts.supports_shared_experts_aux_stream(8)
+
+    assert len(plans) == 1
+    assert plans[0]["tokens"] == 8
+    assert plans[0]["topk"] == 4
+    assert plans[0]["quant_mode"] == "nvfp4"
+
+
 def test_b12x_moe_custom_op_matches_generic_mutation_contract() -> None:
     b12x_schema = str(torch.ops.vllm.b12x_moe_forward.default._schema)
     b12x_shared_schema = str(torch.ops.vllm.b12x_moe_forward_shared.default._schema)

tests/quantization/test_nvfp4_nf3_hybrid.py

@@ -8,6 +8,7 @@
 from vllm.model_executor.layers.quantization import get_quantization_config
 from vllm.model_executor.layers.quantization.nvfp4_nf3_hybrid import (
     NvFp4Nf3HybridConfig,
+    NvFp4Nf3HybridMoEMethod,
     _combined_tier_local_descriptors,
     _read_hybrid_keys,
     _unpack_nf3_codes,
@@ -109,3 +110,10 @@ def test_grid188_tier_descriptors_encode_exact_partition():
 def test_grid188_tier_descriptors_reject_incomplete_partition():
     with pytest.raises(ValueError, match="does not cover all 256"):
         _combined_tier_local_descriptors({0: (0, 0)})
+
+
+@pytest.mark.parametrize("num_tokens", [1, 8, 256, 3072])
+def test_hybrid_moe_rejects_shared_expert_aux_stream(num_tokens):
+    method = object.__new__(NvFp4Nf3HybridMoEMethod)
+
+    assert not method.supports_shared_experts_aux_stream(num_tokens)