Skip to content

[Perf][MoE] Share FlashInfer B12x MoE workspaces across layers - #53081

Draft
lucifer1004 wants to merge 1 commit into
vllm-project:mainfrom
lucifer1004:b12x-shared-workspaces
Draft

[Perf][MoE] Share FlashInfer B12x MoE workspaces across layers#53081
lucifer1004 wants to merge 1 commit into
vllm-project:mainfrom
lucifer1004:b12x-shared-workspaces

Conversation

@lucifer1004

Copy link
Copy Markdown
Contributor

What

Share the FlashInfer B12x MoE workspaces across layers within each worker process.

Today every MoE layer constructs its own B12xMoEWrapper with its own static workspace, dynamic workspace, and output buffer. For large-expert-count models this is a substantial per-layer GPU memory cost (e.g. ~640 MiB per layer on a 384-expert NVFP4 shape — 60+ GiB across layers). The workspaces are identically shaped for every layer and layers execute sequentially, so all layers can share the buffers allocated by the first wrapper.

Dependency

Depends on flashinfer-ai/flashinfer#4603, which adds the shared_static_workspace / shared_dynamic_workspace / shared_output parameters to B12xMoEWrapper. This PR is opened as a draft until that lands and the FlashInfer version pin is updated; with an older FlashInfer the new keyword arguments would raise TypeError at model init.

Notes

  • Sharing is keyed on (global_num_experts, topk, hidden_dim, intermediate_size_per_partition, max_num_tokens, activation, device); layers with a different MoE config allocate their own buffers.
  • No numerical change: the wrapper writes into the shared buffers on every call, and layers do not execute concurrently within a worker.

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>
@lucifer1004
lucifer1004 force-pushed the b12x-shared-workspaces branch from 9c6feca to 790ab3a Compare August 20, 2026 09:01
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

Status: No status

Development

Successfully merging this pull request may close these issues.

1 participant