[GG] fix(moe): safely overlap shared experts before resident grids

tests/model_executor/layers/fused_moe/test_shared_experts_stream.py

@@ -2,34 +2,151 @@
 # 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:
     shared_experts = object.__new__(SharedExperts)
     aux_stream = Mock()
     consumer_stream = Mock()
     output = Mock()
-    shared_experts_input = Mock()
+    shared_experts.enable_dbo = False
+    shared_experts._output = [output, None]
     shared_experts._stream = aux_stream
-    shared_experts._layer = Mock(return_value=output)
+    monkeypatch.setattr(shared_module, "current_stream", lambda: consumer_stream)
+
+    shared_experts._join_aux_stream()
+
+    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)
+
+
+def test_aux_stream_is_joined_before_resident_moe(monkeypatch) -> None:
+    events: list[str] = []
+    aux = Mock()
+    consumer = Mock()
+    output = Mock()
+
+    aux.wait_stream.side_effect = lambda _stream: events.append("input-ready")
+    consumer.wait_stream.side_effect = lambda _stream: events.append("join-aux")
+    monkeypatch.setattr(shared_module, "aux_stream", Mock(return_value=aux))
+    monkeypatch.setattr(shared_module, "current_stream", lambda: consumer)
+    monkeypatch.setattr(shared_module.envs, "VLLM_DISABLE_SHARED_EXPERTS_STREAM", False)
+    monkeypatch.setattr(
+        shared_module,
+        "current_platform",
+        SimpleNamespace(is_cuda=lambda: True),
+    )
 
     @contextmanager
     def use_stream(stream):
-        assert stream is aux_stream
+        assert stream is aux
         yield
 
     monkeypatch.setattr(torch.cuda, "stream", use_stream)
-    monkeypatch.setattr(shared_module, "current_stream", lambda: consumer_stream)
+    shared = SharedExperts(
+        layer=Mock(side_effect=lambda _input: events.append("shared") or output),
+        moe_config=SimpleNamespace(
+            moe_parallel_config=SimpleNamespace(
+                enable_eplb=False,
+                all2all_backend="allgather_reducescatter",
+                use_fi_nvl_two_sided_kernels=False,
+            )
+        ),
+        enable_dbo=False,
+        mk_can_overlap_shared_experts=Mock(return_value=False),
+        experts_support_aux_stream=Mock(return_value=True),
+    )
+    shared_input = Mock()
+    shared_input.shape = (1, 16)
+    shared.maybe_sync_shared_experts_stream(shared_input)
+    events.append("gate")
 
-    result = shared_experts._run_in_aux_stream(shared_experts_input)
+    runner = object.__new__(MoERunner)
+    torch.nn.Module.__init__(runner)
+    runner._shared_experts = shared
+    runner.routed_experts = SimpleNamespace(
+        quant_method=SimpleNamespace(is_monolithic=True),
+        forward_monolithic=lambda **_kwargs: events.append("routed") or object(),
+    )
+    runner._apply_quant_method(
+        hidden_states=Mock(),
+        router_logits=Mock(),
+        shared_experts_input=shared_input,
+    )
 
-    assert result is output
-    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)
+    assert events == ["input-ready", "shared", "gate", "join-aux", "routed"]
+    shared_input.record_stream.assert_called_once_with(aux)

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)

vllm/model_executor/layers/fused_moe/b12x_moe.py

@@ -453,10 +453,9 @@ def _plan_b12x_moe_execution(
 
 
 def _b12x_moe_plan_supports_aux_stream_overlap(plan: Any) -> bool:
-    # Namespaced SparkInfer does not yet publish a plan capability predicate.
-    # Some resident-grid plans use device-wide barriers, so lack of an explicit
-    # guarantee must conservatively disable auxiliary-stream overlap.
-    return False
+    from sparkinfer.moe.fused_moe import plan_supports_aux_stream_overlap
+
+    return bool(plan_supports_aux_stream_overlap(plan))
 
 
 def _b12x_scratch_nbytes(plan: Any) -> int:
@@ -802,6 +801,7 @@ def __init__(
         self._activation_amax_base_num_layers = _current_config_num_hidden_layers()
         self._activation_amax_state_key: tuple[str, str, int] | None = None
         self._activation_amax_layer_idx: int | None = None
+        self._aux_stream_overlap_by_tokens: dict[int, bool] = {}
 
     def _quant_mode(self) -> str:
         source_format = self._source_format()
@@ -826,6 +826,38 @@ def _quant_mode(self) -> str:
             return "w4a16"
         return "nvfp4" if self.quant_config.quant_dtype == "nvfp4" else "w4a16"
 
+    def supports_shared_experts_aux_stream(self, num_tokens: int) -> bool:
+        """Allow overlap only when B12X selected a barrier-free launch."""
+        num_tokens = max(int(num_tokens), 1)
+        cached = self._aux_stream_overlap_by_tokens.get(num_tokens)
+        if cached is not None:
+            return cached
+
+        prepared = self._lookup_prepared_experts()
+        if prepared is None:
+            return False
+        activation = self.moe_config.activation
+        swiglu_limit, swiglu_alpha, swiglu_beta = self._b12x_swiglu_params(activation)
+        quant_mode = self._quant_mode()
+        if (
+            not _supports_swiglu_limit(quant_mode)
+            and activation != MoEActivation.SWIGLUOAI_UNINTERLEAVE
+        ):
+            swiglu_limit = None
+        plan = _plan_b12x_moe_execution(
+            tokens=num_tokens,
+            topk=int(self.moe_config.experts_per_token),
+            device=prepared.w1_fp4.device,
+            quant_mode=quant_mode,
+            experts=prepared,
+            swiglu_limit=swiglu_limit,
+            swiglu_alpha=swiglu_alpha,
+            swiglu_beta=swiglu_beta,
+        )
+        supports_overlap = _b12x_moe_plan_supports_aux_stream_overlap(plan)
+        self._aux_stream_overlap_by_tokens[num_tokens] = supports_overlap
+        return supports_overlap
+
     def _source_format(self) -> str:
         if self.quant_config.weight_quant_dtype == "nvfp4":
             return "modelopt_nvfp4"

vllm/model_executor/layers/fused_moe/fused_moe_method_base.py

@@ -49,6 +49,20 @@ def mk_can_overlap_shared_experts(self) -> bool:
             self.moe_kernel is not None and self.moe_kernel.can_overlap_shared_experts
         )
 
+    def supports_shared_experts_aux_stream(self, num_tokens: int) -> bool:
+        """Whether the selected expert kernel can overlap another CUDA stream."""
+        if self.moe_kernel is None:
+            return True
+        experts = getattr(self.moe_kernel, "fused_experts", None)
+        supports_overlap = getattr(
+            experts,
+            "supports_shared_experts_aux_stream",
+            None,
+        )
+        if supports_overlap is None:
+            return True
+        return bool(supports_overlap(int(num_tokens)))
+
     @abstractmethod
     def create_weights(
         self,

vllm/model_executor/layers/fused_moe/runner/moe_runner.py

@@ -296,11 +296,16 @@ def __init__(
         self._shared_experts: SharedExperts | None = None
         if shared_experts is not None:
             can_overlap = lambda: self._quant_method.mk_can_overlap_shared_experts
+
+            def supports_aux_stream(num_tokens: int) -> bool:
+                return self._quant_method.supports_shared_experts_aux_stream(num_tokens)
+
             self._shared_experts = SharedExperts(
                 shared_experts,
                 moe_config=moe_config,
                 enable_dbo=enable_dbo,
                 mk_can_overlap_shared_experts=can_overlap,
+                experts_support_aux_stream=supports_aux_stream,
             )
 
         # Needed for string -> MoERunner layer lookup in custom ops.
@@ -600,6 +605,10 @@ def _apply_quant_method(
 
         if self.routed_experts.quant_method.is_monolithic:
             # Monolithic kernels: pass router_logits to routed_experts
+            self._maybe_apply_shared_experts(
+                shared_experts_input,
+                SharedExpertsOrder.MULTI_STREAM_OVERLAPPED,
+            )
             fused_out = self.routed_experts.forward_monolithic(
                 x=hidden_states,
                 router_logits=router_logits,
@@ -614,6 +623,14 @@ def _apply_quant_method(
                 input_ids=input_ids,
             )
 
+            # The auxiliary launch was queued before gate/router work. Join it
+            # only after routing, immediately before the routed kernel, to keep
+            # resident-grid plans isolated without serializing independent work.
+            self._maybe_apply_shared_experts(
+                shared_experts_input,
+                SharedExpertsOrder.MULTI_STREAM_OVERLAPPED,
+            )
+
             fused_out = self.routed_experts.forward_modular(
                 x=hidden_states,
                 topk_weights=topk_weights,
@@ -622,11 +639,6 @@ def _apply_quant_method(
                 shared_experts_input=shared_experts_input,
             )
 
-        self._maybe_apply_shared_experts(
-            shared_experts_input,
-            SharedExpertsOrder.MULTI_STREAM_OVERLAPPED,
-        )
-
         return (
             self._shared_experts.output if self._shared_experts is not None else None,
             fused_out,