[MoE] Allow reusable workspaces in XPU fused MoE

[ehartford]

Summary

Fixes #390.

XpuFusedMoe.apply() currently allocates several large scratch tensors on every kernel-path call:

• remapped hidden states
• GEMM1 output
• activation output
• GEMM2 output

In decode-heavy MoE inference this creates avoidable allocator churn, since the required workspace shapes are deterministic from hidden_states.shape, n_experts_per_token, hidden_size, and inter_size.

This PR adds optional caller-provided workspaces to the XPU fused MoE kernel path. Existing callers keep the current allocation behavior when workspaces are omitted, while optimized callers can preallocate scratch buffers and reuse them across repeated decode steps.

Changes

• Add optional workspace13 and workspace2 parameters to XpuFusedMoe.apply() / _apply_kernel().
• Use validated views into those buffers for the large scratch tensors when provided.
• Preserve the existing allocation path when no workspaces are passed.
• Add a regression test comparing reusable-workspace output against the existing allocation path.
• Expand mini fused MoE test coverage to include representative EP variants.
• Add a guard for invalid num_experts <= 0 in the kernel path.

Compatibility

The API change is additive. Existing callers do not need to pass any new arguments and should see the same behavior as before.

The workspace buffers are scratch only. The implementation reuses each workspace only after the previous view backed by that workspace has been consumed:

• workspace13: remapped hidden states, then activation output
• workspace2: GEMM1 output, then GEMM2 output

Test Plan

Built this branch independently from origin/main in a clean worktree with:

source ~/.venv/bin/activate
PATH=/opt/intel/oneapi/compiler/2026.0/bin:$PATH \
CMPLR_ROOT=/opt/intel/oneapi/compiler/2026.0 \
VLLM_PAGED_DECODE_CONFIG=paged_decode_default.conf \
VLLM_CHUNK_PREFILL_CONFIG=chunk_prefill_default.conf \
python setup.py build_ext --inplace

Import smoke:

source ~/.venv/bin/activate
python - <<'PY'
import vllm_xpu_kernels._C
import vllm_xpu_kernels._xpu_C
import vllm_xpu_kernels._moe_C
print('kernel imports ok')
PY

Focused fused MoE tests:

source ~/.venv/bin/activate
pytest -q tests/fused_moe/test_fused_moe.py --tb=short

Test Result

kernel imports ok
201 passed in 52.92s

Related Work

There is nearby MoE refactor work in #376. This PR is intentionally narrow: it does not refactor the MoE class hierarchy or grouped GEMM signatures; it only adds optional workspace reuse while preserving the default allocation path.

[Copilot]

Pull request overview

This PR adds optional caller-provided scratch workspaces to the XPU fused MoE kernel path (XpuFusedMoe.apply() / _apply_kernel()) to avoid repeated large temporary allocations during decode-heavy inference, while preserving the existing allocation behavior when workspaces are omitted.

Changes:

• Extends XpuFusedMoe.apply() / _apply_kernel() with optional workspace13 and workspace2 and uses views into them for major scratch tensors when provided.
• Adds a kernel-path guard for invalid num_experts <= 0.
• Expands fused MoE tests (including EP mini-scope entries) and adds a regression test comparing workspace reuse vs the allocation path.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File	Description
vllm_xpu_kernels/fused_moe_interface.py	Adds optional workspace parameters and workspace-backed views for scratch tensors in the kernel MoE path.
tests/fused_moe/test_fused_moe.py	Adds a regression test for workspace reuse equivalence and extends mini-scope parameters for EP variants.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[mayuyuace]

Num_moe_inputs is a dynamic number.
So I think the code change may raise some problems.
And please provide comparative data both in kernels e2e and models.

[xinyu-intel]

@mayuyuace can you investigate whether the memory can be re-used across layers in the single iteration?

[mayuyuace]

@xinyu-intel
Good question!
Without this PR, every layer can re-use memory by allocate and release memory, so total model layers do not need allocate (numbers of layer * single layer memory).
With this PR, every layer would need allocate it's own memory for this.
So, this is indeed something I overlooked: this PR will also lead to higher memory usage.

[xinyu-intel]

@xinyu-intel Good question! Without this PR, every layer can re-use memory by allocate and release memory, so total model layers do not need allocate (numbers of layer * single layer memory). With this PR, every layer would need allocate it's own memory for this. So, this is indeed something I overlooked: this PR will also lead to higher memory usage.

is it possible to allocate/release workspace memory once per iteration instead of per layer?

[mayuyuace]

@xinyu-intel
Torch employs a memory cache mechanism, so there are no actual allocation or release operations on the device; consequently, this only affects the host-side time required for Torch to manage its own memory cache pool.

[mayuyuace]

@xinyu-intel Good question! Without this PR, every layer can re-use memory by allocate and release memory, so total model layers do not need allocate (numbers of layer * single layer memory). With this PR, every layer would need allocate it's own memory for this. So, this is indeed something I overlooked: this PR will also lead to higher memory usage.

is it possible to allocate/release workspace memory once per iteration instead of per layer?

I think no.
In a real-world production environment, the input length is a dynamic value; consequently, the sizes of all intermediate tensors cannot be determined upfront. A fixed size can be maintained only during the decoding phase in a testing environment.

[xinyu-intel]

@xinyu-intel Torch employs a memory cache mechanism, so there are no actual allocation or release operations on the device; consequently, this only affects the host-side time required for Torch to manage its own memory cache pool.

you're right, we do see aten::fill host/device call in the trace of each moe layer. Though the device time is very limited, the host call and device bubble cause the decode overhead.

[xinyu-intel]

@xinyu-intel Good question! Without this PR, every layer can re-use memory by allocate and release memory, so total model layers do not need allocate (numbers of layer * single layer memory). With this PR, every layer would need allocate it's own memory for this. So, this is indeed something I overlooked: this PR will also lead to higher memory usage.

is it possible to allocate/release workspace memory once per iteration instead of per layer?

I think no. In a real-world production environment, the input length is a dynamic value; consequently, the sizes of all intermediate tensors cannot be determined upfront. A fixed size can be maintained only during the decoding phase in a testing environment.

I mean per iteration, a.k.a (num_tokens, hidden_size), it should be known at the beginning of the iteration.

[mayuyuace]

@xinyu-intel Good question! Without this PR, every layer can re-use memory by allocate and release memory, so total model layers do not need allocate (numbers of layer * single layer memory). With this PR, every layer would need allocate it's own memory for this. So, this is indeed something I overlooked: this PR will also lead to higher memory usage.

is it possible to allocate/release workspace memory once per iteration instead of per layer?

I think no. In a real-world production environment, the input length is a dynamic value; consequently, the sizes of all intermediate tensors cannot be determined upfront. A fixed size can be maintained only during the decoding phase in a testing environment.

I mean per iteration, a.k.a (num_tokens, hidden_size), it should be known at the beginning of the iteration.

I see what you mean: you want to allocate a block of memory before each step starts and have every layer reuse it. I think this is feasible, but the change would need to be implemented within vLLM itself; it cannot be achieved solely at the kernels level.

[xinyu-intel]

@xinyu-intel Good question! Without this PR, every layer can re-use memory by allocate and release memory, so total model layers do not need allocate (numbers of layer * single layer memory). With this PR, every layer would need allocate it's own memory for this. So, this is indeed something I overlooked: this PR will also lead to higher memory usage.

is it possible to allocate/release workspace memory once per iteration instead of per layer?

I think no. In a real-world production environment, the input length is a dynamic value; consequently, the sizes of all intermediate tensors cannot be determined upfront. A fixed size can be maintained only during the decoding phase in a testing environment.

I mean per iteration, a.k.a (num_tokens, hidden_size), it should be known at the beginning of the iteration.

I see what you mean: you want to allocate a block of memory before each step starts and have every layer reuse it. I think this is feasible, but the change would need to be implemented within vLLM itself; it cannot be achieved solely at the kernels level.

exactly. to check if vLLM modular MOE design can support this.

[ehartford]

just try running qwen3.6-35b with and without this patch. shrug not sure what the debate is about.