[Bugfix][MoE] Size the permute scratch from its inputs

benchmarks/kernels/benchmark_moe_permute_unpermute.py

@@ -56,13 +56,9 @@ def benchmark_permute(
         qhidden_states, input_gating, topk, False
     )
     scratch = MoEPermuteScratch(
-        max_num_tokens=num_tokens,
-        topk=topk,
         num_experts=num_experts,
         num_local_experts=num_experts,
         device=qhidden_states.device,
-        hidden_size=hidden_size,
-        hidden_dtype=qhidden_states.dtype,
     )
 
     def prepare(i: int):
@@ -135,13 +131,9 @@ def benchmark_unpermute(
         qhidden_states, input_gating, topk, False
     )
     scratch = MoEPermuteScratch(
-        max_num_tokens=num_tokens,
-        topk=topk,
         num_experts=num_experts,
         num_local_experts=num_experts,
         device=qhidden_states.device,
-        hidden_size=hidden_size,
-        hidden_dtype=qhidden_states.dtype,
     )
 
     def prepare():

tests/kernels/moe/test_moe.py

@@ -1437,6 +1437,7 @@ def test_humming_global_valid_shape_m(
 def test_humming_permute_scratch_is_shared_by_config(
     monkeypatch: pytest.MonkeyPatch,
 ):
+    """Layers share one scratch whose key is independent of batch shape."""
     from types import SimpleNamespace
     from unittest.mock import Mock
 
@@ -1447,34 +1448,19 @@ def test_humming_permute_scratch_is_shared_by_config(
     manager = workspace.WorkspaceManager(torch.device("cpu"))
     monkeypatch.setattr(workspace, "_manager", manager)
 
-    scratch_topk6 = Mock()
-    scratch_topk1 = Mock()
-    scratch_type = Mock(side_effect=[scratch_topk6, scratch_topk1])
+    scratch = Mock()
+    scratch_type = Mock(side_effect=[scratch])
     monkeypatch.setattr(humming, "moe_permute_unpermute_supported", lambda: True)
     monkeypatch.setattr(permute, "MoEPermuteScratch", scratch_type)
     moe_config = make_dummy_moe_config(max_num_tokens=512, experts_per_token=6)
-    moe_config.moe_parallel_config.dp_size = 2
     experts = SimpleNamespace(moe_config=moe_config)
     other_layer = SimpleNamespace(moe_config=moe_config)
 
-    assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6
-    assert (
-        humming.HummingExpertsBase._get_permute_scratch(other_layer, 6) is scratch_topk6
-    )
-    assert humming.HummingExpertsBase._get_permute_scratch(experts, 1) is scratch_topk1
-    assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6
-
-    assert scratch_type.call_count == 2
-    first_call, second_call = scratch_type.call_args_list
-    assert first_call.kwargs["max_num_tokens"] == 1024
-    assert first_call.kwargs["topk"] == 6
-    assert second_call.kwargs["max_num_tokens"] == 6144
-    assert second_call.kwargs["topk"] == 1
-
+    assert humming.HummingExpertsBase._get_permute_scratch(experts) is scratch
+    assert humming.HummingExpertsBase._get_permute_scratch(other_layer) is scratch
     manager.lock()
-    assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6
-    with pytest.raises(AssertionError, match="was not allocated during warmup"):
-        humming.HummingExpertsBase._get_permute_scratch(experts, 2)
+    assert humming.HummingExpertsBase._get_permute_scratch(experts) is scratch
+    assert scratch_type.call_count == 1
 
 
 def test_humming_delegates_to_instance_activation():
@@ -1531,7 +1517,7 @@ def test_humming_grouped_apply_forwards_valid_prefix(
         num_experts=2,
         estimate_local_valid_shape_m=lambda _: 6,
         prepare_buffers=lambda *_: buffers,
-        _get_permute_scratch=lambda _, *, indices_only=False: object(),
+        _get_permute_scratch=lambda: object(),
         process_input=Mock(
             side_effect=lambda _, **kwargs: (
                 kwargs["inputs"],

tests/kernels/moe/test_moe_permute_unpermute.py

@@ -169,13 +169,9 @@ def test_moe_permute_unpermute(
     scratch = None
     if use_scratch:
         scratch = get_moe_permute_scratch(
-            max_num_tokens=n_token,
-            topk=topk,
             num_experts=n_expert,
             num_local_experts=n_local_expert,
             device=hidden_states.device,
-            hidden_size=n_hidden,
-            hidden_dtype=dtype,
         )
 
     (
@@ -247,13 +243,9 @@ def test_moe_permute_reuses_scratch_buffers(
     _, topk_ids, _ = fused_topk(hidden_states, gating_output, topk, False)
 
     scratch_config = dict(
-        max_num_tokens=n_token,
-        topk=topk,
         num_experts=n_expert,
         num_local_experts=n_expert,
         device=hidden_states.device,
-        hidden_size=n_hidden,
-        hidden_dtype=hidden_states.dtype,
     )
     scratch = get_moe_permute_scratch(**scratch_config)
 
@@ -300,15 +292,10 @@ def test_moe_permute_reuses_scratch_buffers(
         permuted_idx_2,
     ) = second
 
-    torch.testing.assert_close(permuted_hidden_states_1, permuted_hidden_states_2)
     torch.testing.assert_close(expert_first_token_offset_1, expert_first_token_offset_2)
     torch.testing.assert_close(inv_permuted_idx_1, inv_permuted_idx_2)
     torch.testing.assert_close(permuted_idx_1, permuted_idx_2)
 
-    assert (
-        permuted_hidden_states_1.untyped_storage().data_ptr()
-        == permuted_hidden_states_2.untyped_storage().data_ptr()
-    )
     assert (
         expert_first_token_offset_1.untyped_storage().data_ptr()
         == expert_first_token_offset_2.untyped_storage().data_ptr()
@@ -330,14 +317,11 @@ def test_moe_permute_scratch_reused_across_graph_sizes(workspace_init) -> None:
 
     device = torch.device("cuda:0")
     scratch = get_moe_permute_scratch(
-        max_num_tokens=128,
-        topk=6,
         num_experts=16,
         num_local_experts=16,
         device=device,
-        hidden_size=2048,
-        hidden_dtype=torch.bfloat16,
     )
+    scratch.reserve(128 * 6)
     current_workspace_manager().lock()
     runs = []
     for n_token in (128, 1, 33):
@@ -372,13 +356,9 @@ def test_moe_permute_scratch_isolated_across_execution_slots(monkeypatch) -> Non
     ubatch = 0
     monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: ubatch)
     scratch_config = dict(
-        max_num_tokens=64,
-        topk=6,
         num_experts=16,
         num_local_experts=16,
         device=device,
-        hidden_size=2048,
-        hidden_dtype=torch.bfloat16,
     )
     runs = []
     for i, n_token in enumerate((1, 7, 33, 64)):
@@ -396,7 +376,7 @@ def test_moe_permute_scratch_isolated_across_execution_slots(monkeypatch) -> Non
             runs.append((stream, graph, hidden, topk_ids, output, scratch))
         torch.cuda.current_stream().wait_stream(stream)
 
-    assert len({run[-1].permuted_hidden_states.data_ptr() for run in runs}) == 4
+    assert len({run[-1].permuted_idx.data_ptr() for run in runs}) == 4
     manager.lock()
     for _ in range(3):
         for i, (stream, graph, hidden, topk_ids, _, scratch) in enumerate(runs):
@@ -423,15 +403,15 @@ def test_moe_permute_scratch_without_manager(monkeypatch) -> None:
 
     monkeypatch.setattr(workspace, "_manager", None)
     config = dict(
-        max_num_tokens=4,
-        topk=2,
         num_experts=4,
         num_local_experts=4,
         device=torch.device("cuda"),
     )
     first = get_moe_permute_scratch(**config)
     second = get_moe_permute_scratch(**config)
 
+    first.reserve(8)
+    second.reserve(8)
     assert (
         first.token_expert_indices.data_ptr() != second.token_expert_indices.data_ptr()
     )
@@ -450,13 +430,9 @@ def test_moe_permute_ignores_invalid_expert_ids_with_scratch() -> None:
     topk_ids = torch.tensor([[0], [-1], [1], [4], [2]], device="cuda")
     expert_map = torch.tensor([0, 1, -1, -1], dtype=torch.int32, device="cuda")
     scratch = MoEPermuteScratch(
-        max_num_tokens=5,
-        topk=1,
         num_experts=4,
         num_local_experts=2,
         device=hidden_states.device,
-        hidden_size=16,
-        hidden_dtype=hidden_states.dtype,
     )
 
     permuted, _, expert_offsets, inverse, _ = moe_permute(
@@ -498,3 +474,39 @@ def test_moe_permute_ignores_invalid_expert_ids_with_scratch() -> None:
     expected_offsets = torch.tensor([0, 5, 5], dtype=torch.int64, device="cuda")
     torch.testing.assert_close(indices.flatten(), expected_inverse)
     torch.testing.assert_close(expert_offsets, expected_offsets)
+
+
+def test_moe_permute_scratch_grows_until_workspace_lock(workspace_init) -> None:
+    """Capacity follows the inputs seen in warmup, whatever their topk or size."""
+    if not moe_permute_unpermute_supported():
+        pytest.skip("moe_permute_unpermute is not supported on this platform.")
+
+    device = torch.device("cuda:0")
+    scratch = get_moe_permute_scratch(
+        num_experts=16, num_local_experts=16, device=device
+    )
+
+    def run(n_token: int, topk: int, dtype: torch.dtype) -> None:
+        hidden = torch.randn((n_token, 64), device=device).to(dtype)
+        topk_ids = torch.randint(16, (n_token, topk), device=device)
+        actual = moe_permute(hidden, None, topk_ids, 16, scratch=scratch)
+        expected = moe_permute(hidden, None, topk_ids, 16)
+        for a, e in zip(actual, expected):
+            if a is not None:
+                torch.testing.assert_close(a.float(), e.float())
+
+    run(8, 6, torch.bfloat16)
+    run(64, 6, torch.bfloat16)
+    run(384, 1, torch.float8_e4m3fn)
+    assert scratch.max_expanded_rows == 384
+
+    current_workspace_manager().lock()
+    run(64, 6, torch.bfloat16)
+    with pytest.raises(AssertionError, match="Scratch growth is not allowed"):
+        run(65, 6, torch.bfloat16)
+
+    # Experts rebuilt after the lock (e.g. a weight reload) keep the capacity.
+    rebuilt = get_moe_permute_scratch(
+        num_experts=16, num_local_experts=16, device=device
+    )
+    assert rebuilt.max_expanded_rows == 384

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

@@ -20,6 +20,7 @@
 )
 from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
     MoEPermuteScratch,
+    get_moe_permute_scratch,
     moe_permute,
     moe_permute_unpermute_supported,
     moe_unpermute,
@@ -304,7 +305,6 @@ def __init__(
         self.ab_strides2 = ab_strides2
         self.c_strides1 = c_strides1
         self.c_strides2 = ab_strides1_c_strides2
-        self._permute_scratch: MoEPermuteScratch | None = None
 
     @staticmethod
     def _supports_current_device() -> bool:
@@ -335,15 +335,13 @@ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
         return TopKWeightAndReduceDelegate()
 
     def _get_permute_scratch(self) -> MoEPermuteScratch | None:
-        if self._permute_scratch is None and moe_permute_unpermute_supported():
-            self._permute_scratch = MoEPermuteScratch(
-                max_num_tokens=self.moe_config.max_num_tokens,
-                topk=self.moe_config.experts_per_token,
-                num_experts=self.moe_config.num_experts,
-                num_local_experts=self.moe_config.num_local_experts,
-                device=torch.device(self.moe_config.device),
-            )
-        return self._permute_scratch
+        if not moe_permute_unpermute_supported():
+            return None
+        return get_moe_permute_scratch(
+            num_experts=self.moe_config.num_experts,
+            num_local_experts=self.moe_config.num_local_experts,
+            device=torch.device(self.moe_config.device),
+        )
 
     def apply(
         self,
@@ -1296,7 +1294,6 @@ def __init__(
         self.s_strides2[:, 0] = k
 
         self.group_size = group_size
-        self._permute_scratch: MoEPermuteScratch | None = None
 
     @staticmethod
     def activation_format() -> mk.FusedMoEActivationFormat:
@@ -1357,15 +1354,13 @@ def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype:
         return self.out_dtype if self.out_dtype is not None else act_dtype
 
     def _get_permute_scratch(self) -> MoEPermuteScratch | None:
-        if self._permute_scratch is None and moe_permute_unpermute_supported():
-            self._permute_scratch = MoEPermuteScratch(
-                max_num_tokens=self.moe_config.max_num_tokens,
-                topk=self.moe_config.experts_per_token,
-                num_experts=self.moe_config.num_experts,
-                num_local_experts=self.moe_config.num_local_experts,
-                device=torch.device(self.moe_config.device),
-            )
-        return self._permute_scratch
+        if not moe_permute_unpermute_supported():
+            return None
+        return get_moe_permute_scratch(
+            num_experts=self.moe_config.num_experts,
+            num_local_experts=self.moe_config.num_local_experts,
+            device=torch.device(self.moe_config.device),
+        )
 
     def workspace_shapes(
         self,

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

@@ -258,25 +258,14 @@ def humming_forward(
             **kwargs,
         )
 
-    def _get_permute_scratch(
-        self, topk: int, indices_only: bool = False
-    ) -> MoEPermuteScratch | None:
+    def _get_permute_scratch(self) -> MoEPermuteScratch | None:
         if not moe_permute_unpermute_supported():
             return None
 
-        max_expanded_rows = (
-            self.moe_config.max_num_tokens
-            * self.moe_config.dp_size
-            * self.moe_config.experts_per_token
-        )
         return get_moe_permute_scratch(
-            max_num_tokens=math.ceil(max_expanded_rows / topk),
-            topk=topk,
             num_experts=self.moe_config.num_experts,
             num_local_experts=self.moe_config.num_local_experts,
             device=torch.device(self.moe_config.device),
-            hidden_size=None if indices_only else self.moe_config.hidden_dim,
-            hidden_dtype=None if indices_only else self.moe_config.in_dtype,
         )
 
     def get_global_valid_shape_m(self, topk_ids: torch.Tensor):
@@ -911,22 +900,27 @@ def apply(
         )
         buffers["output"] = output
 
+        scratch = self._get_permute_scratch()
         if a1q_scale is None:
-            scratch = self._get_permute_scratch(topk_ids.size(1), indices_only=True)
             assert scratch is not None
             scatter_metadata = moe_prepare_scatter(topk_ids, expert_map, scratch)
             expert_offsets, scatter_idx = scatter_metadata
             inv_perm = scatter_idx.flatten()
             num_valid_tokens = expert_offsets[-1:] if expert_map is not None else None
         else:
+            # Prequantized inputs skip w13 quantization, so its buffer is free
+            # to hold the permuted hidden states.
+            permuted_hidden_states = buffers["quanted_gate_up_input"]
+            assert permuted_hidden_states.dtype == hidden_states.dtype
             hidden_states, a1q_scale, expert_offsets, inv_perm, _ = moe_permute(
                 hidden_states=hidden_states,
                 a1q_scale=a1q_scale,
                 topk_ids=topk_ids,
                 n_expert=global_num_experts,
                 n_local_expert=self.num_experts,
                 expert_map=expert_map,
-                scratch=self._get_permute_scratch(topk_ids.size(1)),
+                permuted_hidden_states=permuted_hidden_states,
+                scratch=scratch,
             )
             scatter_idx = None
             num_valid_tokens = None