[Bugfix][MoE] Size the permute scratch from its inputs (#60447)

Signed-off-by: mgoin <mgoin64@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>

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,8 +403,6 @@ 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"),
@@ -432,6 +410,8 @@ def test_moe_permute_scratch_without_manager(monkeypatch) -> None:
     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.config import (
 )
 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 @@ class CutlassExpertsFp8Base(mk.FusedMoEExpertsModular):
         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 @@ class CutlassExpertsFp8Base(mk.FusedMoEExpertsModular):
         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 @@ class CutlassExpertsW4A8Fp8(mk.FusedMoEExpertsModular):
         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 @@ class CutlassExpertsW4A8Fp8(mk.FusedMoEExpertsModular):
         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 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
             **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,14 +900,18 @@ class HummingGroupedExperts(HummingExpertsBase):
         )
         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,
@@ -926,7 +919,8 @@ class HummingGroupedExperts(HummingExpertsBase):
                 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

vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py

@@ -13,90 +13,70 @@ from vllm.v1.worker.workspace import (
 
 @dataclass
 class MoEPermuteScratch:
-    # Reused metadata buffers for repeated grouped-MoE permutes.
-    max_num_tokens: int
-    topk: int
+    """Reused metadata buffers for repeated grouped-MoE permutes.
+
+    Capacity follows the largest input seen, like the shared workspace: it grows
+    during warmup and is fixed once the workspace manager is locked.
+    """
+
     num_experts: int
     num_local_experts: int
     device: torch.device
-    hidden_size: int | None = None
-    hidden_dtype: torch.dtype | None = None
+    max_expanded_rows: int = field(init=False, default=0)
     token_expert_indices: torch.Tensor = field(init=False)
     expert_first_token_offset: torch.Tensor = field(init=False)
     permuted_idx: torch.Tensor = field(init=False)
     inv_permuted_idx: torch.Tensor = field(init=False)
-    permuted_hidden_states: torch.Tensor | None = field(init=False, default=None)
     sort_workspace: torch.Tensor = field(init=False)
     permuted_experts_id: torch.Tensor = field(init=False)
     sorted_row_idx: torch.Tensor = field(init=False)
     topk_ids_int32: torch.Tensor = field(init=False)
     topk_ids_for_sort: torch.Tensor = field(init=False)
-    max_expanded_rows: int = field(init=False)
 
     def __post_init__(self) -> None:
-        assert self.max_num_tokens > 0
-        assert self.topk > 0
         assert self.num_experts > 0
         assert self.num_local_experts > 0
-        if self.hidden_size is None:
-            assert self.hidden_dtype is None
-        else:
-            assert self.hidden_dtype is not None
-
-        self.max_expanded_rows = self.max_num_tokens * self.topk
-        self.token_expert_indices = torch.arange(
-            self.max_expanded_rows, dtype=torch.int32, device=self.device
-        )
         self.expert_first_token_offset = torch.empty(
             self.num_local_experts + 1, dtype=torch.int64, device=self.device
         )
-        self.permuted_idx = torch.empty(
-            self.max_expanded_rows, dtype=torch.int32, device=self.device
-        )
-        self.inv_permuted_idx = torch.empty(
-            self.max_expanded_rows, dtype=torch.int32, device=self.device
-        )
-        if self.hidden_size is not None:
-            hidden_numel = self.max_expanded_rows * self.hidden_size
-            self.permuted_hidden_states = torch.empty(
-                hidden_numel, dtype=self.hidden_dtype, device=self.device
+        # torch.device("cuda") in config, after initialized,
+        # will be changed to cuda:{index}, so we need to refresh here.
+        self.device = self.expert_first_token_offset.device
+
+    def reserve(self, expanded_rows: int) -> None:
+        """Ensure capacity for ``expanded_rows`` (num_tokens * topk) rows."""
+        if expanded_rows <= self.max_expanded_rows:
+            return
+        if (
+            is_workspace_manager_initialized()
+            and current_workspace_manager().is_locked()
+        ):
+            raise AssertionError(
+                f"Workspace is locked but MoE permute scratch requires "
+                f"{expanded_rows} rows, current capacity is "
+                f"{self.max_expanded_rows} rows. Scratch growth is not allowed "
+                "after locking."
             )
-        self.permuted_experts_id = torch.empty(
-            self.max_expanded_rows, dtype=torch.int32, device=self.device
-        )
-        self.sorted_row_idx = torch.empty(
-            self.max_expanded_rows, dtype=torch.int32, device=self.device
-        )
-        self.topk_ids_int32 = torch.empty(
-            self.max_expanded_rows, dtype=torch.int32, device=self.device
-        )
-        self.topk_ids_for_sort = torch.empty(
-            self.max_expanded_rows, dtype=torch.int32, device=self.device
+
+        def empty_rows() -> torch.Tensor:
+            return torch.empty(expanded_rows, dtype=torch.int32, device=self.device)
+
+        self.token_expert_indices = torch.arange(
+            expanded_rows, dtype=torch.int32, device=self.device
         )
+        self.permuted_idx = empty_rows()
+        self.inv_permuted_idx = empty_rows()
+        self.permuted_experts_id = empty_rows()
+        self.sorted_row_idx = empty_rows()
+        self.topk_ids_int32 = empty_rows()
+        self.topk_ids_for_sort = empty_rows()
         sorter_size = torch.ops._moe_C.moe_permute_sort_workspace_size(
-            self.max_expanded_rows, self.num_experts
+            expanded_rows, self.num_experts
         )
         self.sort_workspace = torch.empty(
             sorter_size, dtype=torch.int8, device=self.device
         )
-        # torch.device("cuda") in config, after initialized,
-        # will be changed to cuda:{index}, so we need to refresh here.
-        self.device = self.token_expert_indices.device
-
-    def validate(self, hidden_states: torch.Tensor, topk_ids: torch.Tensor) -> None:
-        n_token, n_hidden = hidden_states.shape
-        assert hidden_states.device == self.device
-        assert topk_ids.device == self.device
-        assert n_token <= self.max_num_tokens
-        assert topk_ids.size(1) == self.topk
-        assert topk_ids.size(0) == n_token
-        if self.hidden_size is not None:
-            assert n_hidden == self.hidden_size
-            assert hidden_states.dtype == self.hidden_dtype
-            assert self.permuted_hidden_states is not None
-
-    def token_expert_indices_view(self, n_token: int) -> torch.Tensor:
-        return self.token_expert_indices[: n_token * self.topk].view(n_token, self.topk)
+        self.max_expanded_rows = expanded_rows
 
     def prepare_topk_ids(self, topk_ids: torch.Tensor) -> torch.Tensor:
         if topk_ids.dtype == torch.int32:
@@ -114,17 +94,16 @@ def moe_prepare_scatter(
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Generate expert offsets and shared scatter/unpermute destination indices."""
     assert topk_ids.device == scratch.device
-    assert topk_ids.size(0) <= scratch.max_num_tokens
-    assert topk_ids.size(1) == scratch.topk
     n_token, topk = topk_ids.shape
     expanded_rows = topk_ids.numel()
+    scratch.reserve(expanded_rows)
     inverse = scratch.inv_permuted_idx[:expanded_rows].view(n_token, topk)
     if expanded_rows == 0:
         scratch.expert_first_token_offset.zero_()
         return scratch.expert_first_token_offset, inverse
     torch.ops._moe_C.moe_prepare_scatter(
         scratch.prepare_topk_ids(topk_ids),
-        scratch.token_expert_indices_view(n_token),
+        scratch.token_expert_indices[:expanded_rows].view(n_token, topk),
         expert_map,
         scratch.num_experts,
         scratch.num_local_experts,
@@ -140,13 +119,9 @@ def moe_prepare_scatter(
 
 def get_moe_permute_scratch(
     *,
-    max_num_tokens: int,
-    topk: int,
     num_experts: int,
     num_local_experts: int,
     device: torch.device,
-    hidden_size: int | None = None,
-    hidden_dtype: torch.dtype | None = None,
 ) -> MoEPermuteScratch:
     """Share scratch across sequential layers in the current ubatch and lane.
 
@@ -155,28 +130,15 @@ def get_moe_permute_scratch(
 
     def create_scratch() -> MoEPermuteScratch:
         return MoEPermuteScratch(
-            max_num_tokens=max_num_tokens,
-            topk=topk,
             num_experts=num_experts,
             num_local_experts=num_local_experts,
             device=device,
-            hidden_size=hidden_size,
-            hidden_dtype=hidden_dtype,
         )
 
     if not is_workspace_manager_initialized():
         return create_scratch()
 
-    key = (
-        MoEPermuteScratch,
-        max_num_tokens,
-        topk,
-        num_experts,
-        num_local_experts,
-        device,
-        hidden_size,
-        hidden_dtype,
-    )
+    key = (MoEPermuteScratch, num_experts, num_local_experts, device)
     return current_workspace_manager().get_persistent_resource(key, create_scratch)
 
 
@@ -204,9 +166,8 @@ def moe_permute(
             parallel shard.
         permuted_hidden_states (Optional[torch.Tensor]): Optional output tensor.
             If None, the output tensor will be created in this function.
-        scratch (Optional[MoEPermuteScratch]): Optional preallocated scratch
-            buffers. Validated against hidden_states and topk_ids when given,
-            otherwise the buffers are allocated in this function.
+        scratch (Optional[MoEPermuteScratch]): Optional reusable metadata
+            buffers. Otherwise they are allocated in this function.
 
     Returns:
     - permuted_hidden_states (torch.Tensor): permuted activation.
@@ -227,20 +188,11 @@ def moe_permute(
     if n_local_expert == -1:
         n_local_expert = n_expert
     if permuted_hidden_states is None:
-        if scratch is None:
-            permuted_hidden_states = torch.empty(
-                (permuted_row_size, n_hidden),
-                dtype=hidden_states.dtype,
-                device=hidden_states.device,
-            )
-        else:
-            scratch.validate(hidden_states, topk_ids)
-            hidden_numel = permuted_row_size * n_hidden
-            scratch_hidden_states = scratch.permuted_hidden_states
-            assert scratch_hidden_states is not None
-            permuted_hidden_states = scratch_hidden_states[:hidden_numel].view(
-                permuted_row_size, n_hidden
-            )
+        permuted_hidden_states = torch.empty(
+            (permuted_row_size, n_hidden),
+            dtype=hidden_states.dtype,
+            device=hidden_states.device,
+        )
     assert permuted_hidden_states.size() == (permuted_row_size, n_hidden), (
         f"Expected permuted hidden states to be {(permuted_row_size, n_hidden)}"
         f" but got {permuted_hidden_states.size()}"
@@ -278,10 +230,15 @@ def moe_permute(
             permuted_idx,
         )
     else:
-        scratch.validate(hidden_states, topk_ids)
+        assert hidden_states.device == scratch.device
+        assert topk_ids.device == scratch.device
+        assert topk_ids.size(0) == n_token
         assert n_expert == scratch.num_experts
         assert n_local_expert == scratch.num_local_experts
-        token_expert_indices = scratch.token_expert_indices_view(n_token)
+        scratch.reserve(permuted_row_size)
+        token_expert_indices = scratch.token_expert_indices[:permuted_row_size].view(
+            n_token, topk
+        )
         expert_first_token_offset = scratch.expert_first_token_offset
         permuted_idx = scratch.permuted_idx[:permuted_row_size]
         permuted_idx.fill_(permuted_row_size)

vllm/v1/worker/workspace.py

@@ -139,7 +139,8 @@ class WorkspaceManager:
         AssertionError.
 
         Keys must identify the resource type and configuration. Callers must not
-        resize or replace tensor storage held by a cached resource. Uses of the
+        resize or replace tensor storage held by a cached resource once the
+        manager is locked. Uses of the
         same resource must not overlap; include the CUDA stream in the key when
         streams can execute concurrently.