[Bugfix] Share FlashInfer B12x MoE workspace across layers

tests/model_executor/test_flashinfer_b12x_moe.py

@@ -0,0 +1,182 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import gc
+import sys
+import weakref
+from types import ModuleType, SimpleNamespace
+from typing import cast
+
+import pytest
+import torch
+
+import vllm.model_executor.layers.fused_moe.experts.flashinfer_b12x_moe as b12x
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEConfig,
+    FusedMoEQuantConfig,
+)
+
+pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup]
+
+
+def _make_experts(
+    activation: MoEActivation,
+    max_num_tokens: int = 16,
+) -> b12x.FlashInferB12xExperts:
+    quant_config = cast(
+        FusedMoEQuantConfig,
+        SimpleNamespace(
+            quant_dtype="nvfp4",
+            g1_alphas=torch.full((2,), 1.0),
+            g2_alphas=torch.full((2,), 2.0),
+            w1_scale=torch.full((2,), 3.0),
+            w2_scale=torch.full((2,), 4.0),
+        ),
+    )
+    moe_config = cast(
+        FusedMoEConfig,
+        SimpleNamespace(
+            in_dtype=torch.bfloat16,
+            num_experts=2,
+            experts_per_token=2,
+            num_local_experts=2,
+            hidden_dim=4,
+            intermediate_size_per_partition=8,
+            max_num_tokens=max_num_tokens,
+            dp_size=2,
+            device=torch.device("cpu"),
+            activation=activation,
+        ),
+    )
+    experts = b12x.FlashInferB12xExperts(moe_config, quant_config)
+    experts._fc2_input_scale = torch.full((2,), 5.0)
+    experts.w1_sf_mma = torch.full((2,), 6.0)
+    experts.w2_sf_mma = torch.full((2,), 7.0)
+    return experts
+
+
+@pytest.mark.parametrize(
+    ("activation", "expected_activation"),
+    [
+        (MoEActivation.SILU, "silu"),
+        (MoEActivation.RELU2_NO_MUL, "relu2"),
+    ],
+)
+def test_b12x_layers_share_owned_workspace_and_preserve_call_contract(
+    monkeypatch: pytest.MonkeyPatch,
+    activation: MoEActivation,
+    expected_activation: str,
+):
+    """Reuse one fixed-capacity wrapper while passing each layer's weights."""
+    wrappers = []
+    calls = []
+
+    class FakeB12xMoEWrapper:
+        def __init__(self, **kwargs):
+            self.config = kwargs
+            wrappers.append(self)
+
+        def run(self, **kwargs):
+            calls.append(kwargs)
+            return torch.full_like(kwargs["x"], len(calls))
+
+    fake_fused_moe = ModuleType("flashinfer.fused_moe")
+    fake_fused_moe.__dict__["B12xMoEWrapper"] = FakeB12xMoEWrapper
+    monkeypatch.setitem(sys.modules, "flashinfer.fused_moe", fake_fused_moe)
+
+    expert_instances = (_make_experts(activation), _make_experts(activation))
+    assert len(wrappers) == 1
+    assert wrappers[0].config == {
+        "num_experts": 2,
+        "top_k": 2,
+        "hidden_size": 4,
+        "intermediate_size": 8,
+        "use_cuda_graph": True,
+        "max_num_tokens": 32,
+        "num_local_experts": 2,
+        "output_dtype": torch.bfloat16,
+        "device": "cpu",
+        "activation": expected_activation,
+    }
+
+    output = torch.empty(3, 4, dtype=torch.bfloat16)
+    hidden_states = torch.empty_like(output)
+    w1 = torch.full((2, 1, 1), 11.0)
+    w2 = torch.full((2, 1, 1), 12.0)
+    topk_weights = torch.empty(3, 2)
+    topk_ids = torch.zeros(3, 2, dtype=torch.int64)
+
+    for experts in expert_instances:
+        experts.apply(
+            output=output,
+            hidden_states=hidden_states,
+            w1=w1,
+            w2=w2,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            activation=activation,
+            global_num_experts=2,
+            expert_map=None,
+            a1q_scale=None,
+            a2_scale=None,
+            workspace13=None,
+            workspace2=None,
+            expert_tokens_meta=None,
+            apply_router_weight_on_input=False,
+        )
+
+    assert len(calls) == 2
+    assert torch.equal(output, torch.full_like(output, 2.0))
+    for call, experts in zip(calls, expert_instances):
+        assert call["x"] is hidden_states
+        assert call["token_final_scales"] is topk_weights
+        assert call["token_selected_experts"].dtype == torch.int32
+        assert torch.equal(call["token_selected_experts"], topk_ids)
+        assert call["w1_weight"] is w1
+        assert call["w2_weight"] is w2
+        assert call["w1_weight_sf"] is experts.w1_sf_mma
+        assert call["w2_weight_sf"] is experts.w2_sf_mma
+        assert call["w1_alpha"] is experts.g1_alphas
+        assert call["w2_alpha"] is experts.g2_alphas
+        assert call["fc2_input_scale"] is experts._fc2_input_scale
+
+
+def test_b12x_owned_workspace_is_released_with_model(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    class FakeB12xMoEWrapper:
+        def __init__(self, **kwargs):
+            pass
+
+    fake_fused_moe = ModuleType("flashinfer.fused_moe")
+    fake_fused_moe.__dict__["B12xMoEWrapper"] = FakeB12xMoEWrapper
+    monkeypatch.setitem(sys.modules, "flashinfer.fused_moe", fake_fused_moe)
+
+    experts = _make_experts(MoEActivation.SILU, max_num_tokens=17)
+    wrapper_ref = weakref.ref(experts._wrapper)
+
+    del experts
+    gc.collect()
+
+    assert wrapper_ref() is None
+
+
+def test_b12x_workspace_is_not_shared_between_model_owners(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    class FakeB12xMoEWrapper:
+        def __init__(self, **kwargs):
+            pass
+
+    fake_fused_moe = ModuleType("flashinfer.fused_moe")
+    fake_fused_moe.__dict__["B12xMoEWrapper"] = FakeB12xMoEWrapper
+    monkeypatch.setitem(sys.modules, "flashinfer.fused_moe", fake_fused_moe)
+
+    owners = iter((object(), object()))
+    monkeypatch.setattr(b12x, "get_current_vllm_config_or_none", lambda: next(owners))
+
+    first = _make_experts(MoEActivation.SILU, max_num_tokens=18)
+    second = _make_experts(MoEActivation.SILU, max_num_tokens=18)
+
+    assert first._wrapper is not second._wrapper

vllm/model_executor/layers/fused_moe/experts/flashinfer_b12x_moe.py

@@ -1,11 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from dataclasses import dataclass
+from threading import Lock
 from typing import Any
+from weakref import WeakValueDictionary
 
 import torch
 
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm.config import get_current_vllm_config_or_none
 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.config import (
     FusedMoEConfig,
@@ -27,6 +31,46 @@
 )
 
 
+@dataclass(frozen=True)
+class _B12xWrapperKey:
+    owner_id: int
+    num_experts: int
+    top_k: int
+    hidden_size: int
+    intermediate_size: int
+    max_num_tokens: int
+    num_local_experts: int
+    output_dtype: torch.dtype
+    device: str
+    activation: str
+
+
+_B12X_WRAPPERS: WeakValueDictionary[_B12xWrapperKey, Any] = WeakValueDictionary()
+_B12X_WRAPPERS_LOCK = Lock()
+
+
+def _get_b12x_wrapper(key: _B12xWrapperKey) -> Any:
+    with _B12X_WRAPPERS_LOCK:
+        wrapper = _B12X_WRAPPERS.get(key)
+        if wrapper is None:
+            from flashinfer.fused_moe import B12xMoEWrapper
+
+            wrapper = B12xMoEWrapper(
+                num_experts=key.num_experts,
+                top_k=key.top_k,
+                hidden_size=key.hidden_size,
+                intermediate_size=key.intermediate_size,
+                use_cuda_graph=True,
+                max_num_tokens=key.max_num_tokens,
+                num_local_experts=key.num_local_experts,
+                output_dtype=key.output_dtype,
+                device=key.device,
+                activation=key.activation,
+            )
+            _B12X_WRAPPERS[key] = wrapper
+        return wrapper
+
+
 class FlashInferB12xExperts(mk.FusedMoEExpertsModular):
     """FlashInfer CuteDSL fused MoE expert for SM12x (SM120/SM121,
     RTX Pro 6000 / DGX Spark).
@@ -59,22 +103,19 @@ def __init__(
         )
         self.out_dtype = moe_config.in_dtype
         self.num_local_experts = moe_config.num_local_experts
-        self.ep_rank = moe_config.moe_parallel_config.ep_rank
-        # FC2 input scale tensor bound in process_weights_after_loading: the
-        # calibrated (now-zeroed) a2_gscale for static-quant checkpoints, or
-        # a synthesized uniform-1.0 tensor for W4A16 checkpoints that lack
-        # one. Holding it on the instance keeps apply() alloc-free.
-        self._fc2_input_scale: torch.Tensor | None = None
-
-        # Shape params for B12xMoEWrapper construction.
         self.global_num_experts = moe_config.num_experts
         self.topk = moe_config.experts_per_token
         self.hidden_dim = moe_config.hidden_dim
         self.intermediate_size_per_partition = (
             moe_config.intermediate_size_per_partition
         )
-        self.max_num_tokens = moe_config.max_num_tokens
-        self.local_expert_offset = self.ep_rank * self.num_local_experts
+        self.max_num_tokens = moe_config.max_num_tokens * moe_config.dp_size
+        self.device = str(torch.device(moe_config.device))
+        # FC2 input scale tensor bound in process_weights_after_loading: the
+        # calibrated (now-zeroed) a2_gscale for static-quant checkpoints, or
+        # a synthesized uniform-1.0 tensor for W4A16 checkpoints that lack
+        # one. Holding it on the instance keeps apply() alloc-free.
+        self._fc2_input_scale: torch.Tensor | None = None
 
         activation = moe_config.activation
         if activation not in self._ACTIVATION_MAP:
@@ -85,8 +126,21 @@ def __init__(
             )
         self._activation_str = self._ACTIVATION_MAP[activation]
 
-        # Lazily created on first apply() call.
-        self._wrapper: Any | None = None
+        self._wrapper = _get_b12x_wrapper(
+            _B12xWrapperKey(
+                owner_id=id(get_current_vllm_config_or_none()),
+                num_experts=self.global_num_experts,
+                top_k=self.topk,
+                hidden_size=self.hidden_dim,
+                intermediate_size=self.intermediate_size_per_partition,
+                max_num_tokens=self.max_num_tokens,
+                num_local_experts=self.num_local_experts,
+                output_dtype=self.out_dtype,
+                device=self.device,
+                activation=self._activation_str,
+            )
+        )
+
         self.w1_sf_mma: torch.Tensor | None = None
         self.w2_sf_mma: torch.Tensor | None = None
 
@@ -214,7 +268,7 @@ def workspace_shapes(
         expert_tokens_meta: mk.ExpertTokensMetadata | None,
         activation: MoEActivation,
     ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
-        # b12x_fused_moe manages its own internal workspace.
+        # B12xMoEWrapper manages its own internal workspace.
         workspace1 = (1,)
         workspace2 = (0,)
         output_shape = (M, K)
@@ -227,24 +281,6 @@ def expects_unquantized_inputs(self) -> bool:
         # from pre-quantizing activations.
         return True
 
-    def _ensure_wrapper(self) -> None:
-        """Lazily create B12xMoEWrapper on first use."""
-        if self._wrapper is not None:
-            return
-
-        from flashinfer.fused_moe import B12xMoEWrapper
-
-        self._wrapper = B12xMoEWrapper(
-            num_experts=self.global_num_experts,
-            top_k=self.topk,
-            hidden_size=self.hidden_dim,
-            intermediate_size=self.intermediate_size_per_partition,
-            use_cuda_graph=True,
-            max_num_tokens=self.max_num_tokens,
-            num_local_experts=self.num_local_experts,
-            activation=self._activation_str,
-        )
-
     def apply(
         self,
         output: torch.Tensor,
@@ -276,20 +312,16 @@ def apply(
             "process_weights_after_loading must run before FlashInferB12xExperts.apply"
         )
 
-        self._ensure_wrapper()
-        wrapper = self._wrapper
-        assert wrapper is not None
-
-        wrapper_output = wrapper.run(
+        wrapper_output = self._wrapper.run(
             x=hidden_states,
+            token_selected_experts=topk_ids.to(torch.int32),
+            token_final_scales=topk_weights,
             w1_weight=w1,
             w1_weight_sf=self.w1_sf_mma,
             w1_alpha=self.g1_alphas,
             fc2_input_scale=self._fc2_input_scale,
             w2_weight=w2,
             w2_weight_sf=self.w2_sf_mma,
             w2_alpha=self.g2_alphas,
-            token_selected_experts=topk_ids.to(torch.int32),
-            token_final_scales=topk_weights,
         )
         output.copy_(wrapper_output)