diff --git a/flashinfer/fused_moe/cute_dsl/b12x_moe.py b/flashinfer/fused_moe/cute_dsl/b12x_moe.py index e13839ebdd4..db9d509156e 100644 --- a/flashinfer/fused_moe/cute_dsl/b12x_moe.py +++ b/flashinfer/fused_moe/cute_dsl/b12x_moe.py @@ -256,6 +256,12 @@ class B12xMoEWrapper: When set, this selects the backend and internal workspace family. source_format: Source weight format for quant_mode="w4a16". Supports "modelopt" and "compressed_tensors". Default: "modelopt". + shared_static_workspace, shared_dynamic_workspace, shared_output: + Optional externally-allocated buffers reused instead of fresh + allocations. Callers running many identically-shaped wrappers + (e.g. one per MoE layer) can share a single set, since layers + execute sequentially. Shapes must match this wrapper's config. + Only valid with ``use_cuda_graph=True``. Example: >>> moe = B12xMoEWrapper(num_experts=256, top_k=8, ...) @@ -283,6 +289,9 @@ def __init__( activation_precision: str = "fp4", quant_mode: Optional[str] = None, source_format: str = "modelopt", + shared_static_workspace: Optional[object] = None, + shared_dynamic_workspace: Optional[object] = None, + shared_output: Optional[torch.Tensor] = None, ): r"""Configure the b12x fused-MoE wrapper. @@ -328,6 +337,17 @@ def __init__( source_format : str Source weight format for ``quant_mode="w4a16"`` — ``"modelopt"`` (default) or ``"compressed_tensors"``. + shared_static_workspace, shared_dynamic_workspace : Optional[object] + Externally allocated workspaces reused instead of fresh + allocations. Callers running many identically-shaped wrappers + (e.g. one per MoE layer) can share a single set, since layers + execute sequentially. Shapes must match this wrapper's config. + Only valid with ``use_cuda_graph=True``. + shared_output : Optional[torch.Tensor] + Externally allocated output buffer, reused like the workspaces. + Must be 2-D with shape ``(>= max_num_tokens, hidden_size)``, + ``output_dtype``, and the same device as this wrapper. + Only valid with ``use_cuda_graph=True``. """ from ...jit.cpp_ext import get_cuda_version from .blackwell_sm12x.moe_dispatch import ( @@ -378,6 +398,9 @@ def __init__( # Pre-allocated objects. Both workspace slots may be populated so # run() can pick per-call; without this, the backend would be locked # to whichever workspace was allocated at init time. + self._shared_static_workspace = shared_static_workspace + self._shared_dynamic_workspace = shared_dynamic_workspace + self._shared_output = shared_output self._static_workspace: object = None self._dynamic_workspace: object = None self._weight_views: object = None @@ -388,35 +411,85 @@ def __init__( self._folded_w1_alpha: Optional[torch.Tensor] = None self._folded_w1_alpha_key: Optional[Tuple] = None + if not use_cuda_graph and any( + resource is not None + for resource in ( + shared_static_workspace, + shared_dynamic_workspace, + shared_output, + ) + ): + raise ValueError( + "shared_static_workspace / shared_dynamic_workspace / " + "shared_output require use_cuda_graph=True; without " + "pre-allocated buffers, run() ignores shared resources." + ) + + if shared_output is not None: + expected_device = torch.device(self.device) + if expected_device.type == "cuda" and expected_device.index is None: + # An index-less "cuda" means the current device; resolve it so + # the comparison covers the device index, not just the type. + expected_device = torch.device("cuda", torch.cuda.current_device()) + if ( + shared_output.dim() != 2 + or shared_output.shape[0] < self.max_num_tokens + or shared_output.shape[1] != self.hidden_size + or shared_output.dtype != self.output_dtype + or shared_output.device != expected_device + ): + raise ValueError( + "shared_output must have shape (>=max_num_tokens, " + f"hidden_size)=({self.max_num_tokens}, {self.hidden_size}), " + f"dtype {self.output_dtype}, and device " + f"{expected_device}; got shape " + f"{tuple(shared_output.shape)}, dtype {shared_output.dtype}, " + f"device {shared_output.device}." + ) + if use_cuda_graph: self._allocate_buffers() def _allocate_buffers(self) -> None: - """Pre-allocate buffers for CUDA graph compatibility.""" + """Pre-allocate buffers for CUDA graph compatibility. + + When shared buffers are injected (``shared_static_workspace`` / + ``shared_dynamic_workspace`` / ``shared_output``), they are used + instead of fresh allocations. MoE layers execute strictly + sequentially, so identically-shaped wrappers (e.g. one per layer in + vLLM) can safely share a single set of workspaces instead of paying + the memory cost per wrapper. + """ from .blackwell_sm12x.moe_dispatch import ( allocate_sm120_moe_workspace, select_sm120_moe_backend, _get_static_compact_cutover_pairs, ) + self._static_workspace = self._shared_static_workspace + self._dynamic_workspace = self._shared_dynamic_workspace + self._moe_output = self._shared_output + max_routed_rows = self.max_num_tokens * self.top_k if self.quant_mode == "w4a16": - self._static_workspace = allocate_sm120_moe_workspace( - state_E=self.num_local_experts, - weight_E=self.num_experts, - routed_rows=max_routed_rows, - k=self.hidden_size, - n=self.intermediate_size, - num_topk=self.top_k, - device=torch.device(self.device), - quant_mode=self.quant_mode, - activation=self.activation, - ) - self._moe_output = torch.empty( - (self.max_num_tokens, self.hidden_size), - dtype=self.output_dtype, - device=self.device, - ) + if self._static_workspace is None: + self._static_workspace = allocate_sm120_moe_workspace( + state_E=self.num_local_experts, + weight_E=self.num_experts, + routed_rows=max_routed_rows, + k=self.hidden_size, + n=self.intermediate_size, + num_topk=self.top_k, + device=torch.device(self.device), + quant_mode=self.quant_mode, + activation=self.activation, + ) + if self._moe_output is None: + self._moe_output = torch.empty( + (self.max_num_tokens, self.hidden_size), + dtype=self.output_dtype, + device=self.device, + ) return # Allocate a dynamic workspace alongside the static one when @@ -446,20 +519,21 @@ def _allocate_buffers(self) -> None: if needs_dynamic else max_routed_rows ) - self._static_workspace = allocate_sm120_moe_workspace( - state_E=self.num_local_experts, - weight_E=self.num_experts, - max_rows=max(1, static_max_rows), - k=self.hidden_size, - n=self.intermediate_size, - num_topk=self.top_k, - device=torch.device(self.device), - quant_mode=self.quant_mode, - backend="static", - activation=self.activation, - ) + if self._static_workspace is None: + self._static_workspace = allocate_sm120_moe_workspace( + state_E=self.num_local_experts, + weight_E=self.num_experts, + max_rows=max(1, static_max_rows), + k=self.hidden_size, + n=self.intermediate_size, + num_topk=self.top_k, + device=torch.device(self.device), + quant_mode=self.quant_mode, + backend="static", + activation=self.activation, + ) - if needs_dynamic: + if needs_dynamic and self._dynamic_workspace is None: self._dynamic_workspace = allocate_sm120_moe_workspace( state_E=self.num_local_experts, weight_E=self.num_experts, @@ -475,11 +549,12 @@ def _allocate_buffers(self) -> None: # Allocated after arch-specific buffers to preserve memory layout # that the autotuner's CUDA graph profiling is sensitive to. - self._moe_output = torch.empty( - (self.max_num_tokens, self.hidden_size), - dtype=self.output_dtype, - device=self.device, - ) + if self._moe_output is None: + self._moe_output = torch.empty( + (self.max_num_tokens, self.hidden_size), + dtype=self.output_dtype, + device=self.device, + ) @flashinfer_api(trace=b12x_moe_wrapper_run_trace) def run( diff --git a/tests/moe/test_b12x_fused_moe.py b/tests/moe/test_b12x_fused_moe.py index 390dbc151cb..de7c9bcb5c6 100644 --- a/tests/moe/test_b12x_fused_moe.py +++ b/tests/moe/test_b12x_fused_moe.py @@ -3121,6 +3121,131 @@ def test_relu2_cuda_graph(self): assert not (output == 0).all(), "All zeros after ReLU2 CUDA graph replay" +def _make_cpu_wrapper(monkeypatch, use_cuda_graph=True, **shared): + """Build a B12xMoEWrapper without a GPU: fake CUDA 13, CPU buffers.""" + from flashinfer.fused_moe.cute_dsl import b12x_moe as b12x_moe_mod + from flashinfer.jit import cpp_ext + + monkeypatch.setattr(cpp_ext, "get_cuda_version", _fake_cuda_13_version) + return b12x_moe_mod.B12xMoEWrapper( + num_experts=8, + top_k=1, + hidden_size=256, + intermediate_size=128, + use_cuda_graph=use_cuda_graph, + max_num_tokens=1024, # crosses the static/dynamic cutover + device="cpu", + **shared, + ) + + +def test_wrapper_defaults_allocate_own_buffers(monkeypatch): + """Without sharing, each wrapper allocates its own workspaces and output.""" + from flashinfer.fused_moe.cute_dsl.blackwell_sm12x import moe_dispatch + + allocs = [] + monkeypatch.setattr( + moe_dispatch, + "allocate_sm120_moe_workspace", + lambda **kw: allocs.append(kw) or object(), + ) + + first = _make_cpu_wrapper(monkeypatch) + second = _make_cpu_wrapper(monkeypatch) + + assert len(allocs) == 4 # static + dynamic per wrapper + assert second._static_workspace is not first._static_workspace + assert second._dynamic_workspace is not first._dynamic_workspace + assert second._moe_output is not first._moe_output + + +def test_wrapper_shared_buffers_are_reused(monkeypatch): + """Injected shared buffers are reused instead of fresh allocations.""" + from flashinfer.fused_moe.cute_dsl.blackwell_sm12x import moe_dispatch + + allocs = [] + monkeypatch.setattr( + moe_dispatch, + "allocate_sm120_moe_workspace", + lambda **kw: allocs.append(kw) or object(), + ) + + first = _make_cpu_wrapper(monkeypatch) + assert len(allocs) == 2 # static + dynamic + + # Spy on the output factory: sharing must skip the output allocation too. + real_empty = torch.empty + output_allocs = [] + monkeypatch.setattr( + torch, + "empty", + lambda *a, **kw: output_allocs.append((a, kw)) or real_empty(*a, **kw), + ) + + second = _make_cpu_wrapper( + monkeypatch, + shared_static_workspace=first._static_workspace, + shared_dynamic_workspace=first._dynamic_workspace, + shared_output=first._moe_output, + ) + + assert len(allocs) == 2 # nothing new allocated + assert not output_allocs + assert second._static_workspace is first._static_workspace + assert second._dynamic_workspace is first._dynamic_workspace + assert second._moe_output is first._moe_output + + +def test_wrapper_shared_output_is_validated(monkeypatch): + """Incompatible shared_output buffers are rejected at construction.""" + from flashinfer.fused_moe.cute_dsl.blackwell_sm12x import moe_dispatch + + monkeypatch.setattr( + moe_dispatch, "allocate_sm120_moe_workspace", lambda **kw: object() + ) + bad = torch.zeros((4, 64), dtype=torch.float32) # too small, wrong dtype + with pytest.raises(ValueError, match="shared_output"): + _make_cpu_wrapper(monkeypatch, shared_output=bad) + + +def test_wrapper_shared_buffers_require_cuda_graph(monkeypatch): + """Shared buffers are rejected when CUDA graph mode is disabled.""" + from flashinfer.fused_moe.cute_dsl.blackwell_sm12x import moe_dispatch + + monkeypatch.setattr( + moe_dispatch, "allocate_sm120_moe_workspace", lambda **kw: object() + ) + output = torch.zeros((1024, 256), dtype=torch.bfloat16) + with pytest.raises(ValueError, match="use_cuda_graph"): + _make_cpu_wrapper(monkeypatch, use_cuda_graph=False, shared_output=output) + + +@sm120_required +@cuda_13_required +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires two GPUs") +def test_wrapper_shared_output_rejects_other_gpu(): + """A shared_output on a different CUDA device is rejected. + + Regression test: the device check must compare the full device, including + the index — with index-less ``device="cuda"`` resolved to the current + device — instead of only the device type. + """ + from flashinfer.fused_moe.cute_dsl import b12x_moe as b12x_moe_mod + + other_gpu = torch.zeros((16, 256), dtype=torch.bfloat16, device="cuda:1") + with torch.cuda.device(0), pytest.raises(ValueError, match="shared_output"): + b12x_moe_mod.B12xMoEWrapper( + num_experts=8, + top_k=1, + hidden_size=256, + intermediate_size=128, + use_cuda_graph=True, + max_num_tokens=16, + device="cuda", # index-less: resolves to cuda:0 + shared_output=other_gpu, + ) + + # --------------------------------------------------------------------------- # Dense-kernel borrow contract (CPU-only, pure source scan). #