diff --git a/tests/model_executor/test_flashinfer_b12x_moe.py b/tests/model_executor/test_flashinfer_b12x_moe.py new file mode 100644 index 000000000000..421c9a1c5ef3 --- /dev/null +++ b/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 diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_b12x_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_b12x_moe.py index ca23e1272490..c6dfaeaf9aef 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_b12x_moe.py +++ b/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,12 +312,10 @@ 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, @@ -289,7 +323,5 @@ def apply( 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)