[Perf][MoE] Share FlashInfer B12x MoE workspaces across layers

Every MoE layer allocated its own B12xMoEWrapper workspaces, which is
a large per-layer GPU memory cost. The workspaces are identical for
every layer and layers execute sequentially, so all layers in a worker
process can share the buffers allocated by the first wrapper.

Depends on flashinfer-ai/flashinfer#4603, which adds the
shared_static_workspace / shared_dynamic_workspace / shared_output
parameters to B12xMoEWrapper.

Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>

Author: lucifer1004

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

@@ -49,6 +49,10 @@ class FlashInferB12xExperts(mk.FusedMoEExpertsModular):
         MoEActivation.RELU2_NO_MUL: "relu2",
     }
 
+    # Per-process (per-worker) buffers shared across layers; see
+    # _ensure_wrapper.
+    _shared_buffers: dict[tuple, tuple] = {}
+
     def __init__(
         self,
         moe_config: FusedMoEConfig,
@@ -239,6 +243,21 @@ def _ensure_wrapper(self) -> None:
 
         from flashinfer.fused_moe import B12xMoEWrapper
 
+        # Workspaces are identical for every MoE layer and layers execute
+        # sequentially, so all layers share the buffers allocated by the
+        # first wrapper instead of paying the (large) per-layer cost.
+        key = (
+            self.global_num_experts,
+            self.topk,
+            self.hidden_dim,
+            self.intermediate_size_per_partition,
+            self.max_num_tokens,
+            self._activation_str,
+            torch.accelerator.current_device_index(),
+        )
+        shared = FlashInferB12xExperts._shared_buffers.get(key)
+        static_ws, dynamic_ws, output = shared if shared else (None, None, None)
+
         self._wrapper = B12xMoEWrapper(
             num_experts=self.global_num_experts,
             top_k=self.topk,
@@ -248,7 +267,16 @@ def _ensure_wrapper(self) -> None:
             max_num_tokens=self.max_num_tokens,
             num_local_experts=self.num_local_experts,
             activation=self._activation_str,
+            shared_static_workspace=static_ws,
+            shared_dynamic_workspace=dynamic_ws,
+            shared_output=output,
         )
+        if shared is None:
+            FlashInferB12xExperts._shared_buffers[key] = (
+                self._wrapper._static_workspace,
+                self._wrapper._dynamic_workspace,
+                self._wrapper._moe_output,
+            )
 
     def apply(
         self,