[Refactor][MoE] Share CUTLASS permute scratch across layers

tests/kernels/moe/test_cutlass_moe.py

@@ -66,6 +66,37 @@
 vllm_config = VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
 
 
+@pytest.mark.parametrize("experts_cls", [CutlassExpertsFp8, CutlassExpertsW4A8Fp8])
+def test_cutlass_permute_scratch_shared_across_layers(monkeypatch, experts_cls):
+    from types import SimpleNamespace
+
+    import vllm.model_executor.layers.fused_moe.experts.cutlass_moe as cutlass
+    import vllm.model_executor.layers.fused_moe.moe_permute_unpermute as permute
+    import vllm.v1.worker.workspace as workspace
+
+    monkeypatch.setattr(
+        workspace,
+        "_manager",
+        workspace.WorkspaceManager(torch.device("cpu"), num_lanes=2),
+    )
+    monkeypatch.setattr(cutlass, "moe_permute_unpermute_supported", lambda: True)
+    factory = lambda **kwargs: object()
+    monkeypatch.setattr(permute, "MoEPermuteScratch", factory)
+    monkeypatch.setattr(cutlass, "MoEPermuteScratch", factory)
+    config = make_dummy_moe_config()
+    first = SimpleNamespace(moe_config=config, _permute_scratch=None)
+    second = SimpleNamespace(moe_config=config, _permute_scratch=None)
+    get_scratch = experts_cls._get_permute_scratch
+
+    scratch = get_scratch(first)
+    assert get_scratch(second) is scratch
+    with workspace.use_workspace_lane(1):
+        other = get_scratch(first)
+        assert other is not scratch
+        assert get_scratch(second) is other
+    assert get_scratch(first) is scratch
+
+
 @pytest.mark.parametrize(
     "experts_cls",
     [

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,15 @@ 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(
+            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),
+        )
 
     def apply(
         self,
@@ -1296,7 +1296,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 +1356,15 @@ 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(
+            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),
+        )
 
     def workspace_shapes(
         self,