XPU fused MoE allocates scratch tensors on every apply call

[ehartford]

Summary

XpuFusedMoe.apply() currently allocates several temporary XPU tensors every time the kernel path runs. In decode-heavy MoE inference, especially with small per-step batch sizes, these repeated allocations add avoidable host/runtime overhead and can make the Python wrapper more expensive than necessary.

This is particularly visible for MoE models where apply() is called once per MoE layer per decode step. The required scratch shapes are deterministic from the request shape and layer configuration, so callers should be able to reuse preallocated scratch buffers across calls instead of forcing fresh torch.empty(...) allocations inside every invocation.

Relevant code

At current main (dae50e2aa58301d3161c26d15bfd867dc5954c12), apply() exposes no way to pass reusable workspaces and forwards directly to _apply_kernel:

vllm-xpu-kernels/vllm_xpu_kernels/fused_moe_interface.py

Lines 402 to 416 in dae50e2

	self,
	output,
	hidden_states,
	topk_weights,
	topk_ids,
	expert_map=None,
	):
	if self._use_ref:
	self._apply_ref(output, hidden_states,
	topk_weights, topk_ids,
	expert_map)
	else:
	self._apply_kernel(output, hidden_states,
	topk_weights, topk_ids,
	expert_map)

Inside _apply_kernel, the following scratch tensors are allocated unconditionally on each call:

• remapped_hidden_states
  

  vllm-xpu-kernels/vllm_xpu_kernels/fused_moe_interface.py

  Lines 457 to 460 in dae50e2

  	remapped_hidden_states = torch.empty(
  	(num_rows * self.n_experts_per_token, hidden_size),
  	dtype=hidden_states.dtype,
  	device=hidden_states.device)

• gemm1_output
  

  vllm-xpu-kernels/vllm_xpu_kernels/fused_moe_interface.py

  Lines 481 to 484 in dae50e2

  	########### gemm1 ##################
  	gemm1_output = torch.empty((num_moe_inputs, 2 * self.inter_size),
  	dtype=hidden_states.dtype,
  	device=hidden_states.device)

• act_output
  

  vllm-xpu-kernels/vllm_xpu_kernels/fused_moe_interface.py

  Lines 498 to 503 in dae50e2

  	# act
  	act_output = torch.empty(
  	(num_moe_inputs, self.inter_size * self.inter_size_scale),
  	dtype=gemm1_output.dtype,
  	device=gemm1_output.device)
  	self.act_func(act_output, gemm1_output)

• gemm2_output
  

  vllm-xpu-kernels/vllm_xpu_kernels/fused_moe_interface.py

  Lines 505 to 508 in dae50e2

  	########### gemm2 ##################
  	gemm2_output = torch.empty((num_moe_inputs, hidden_size),
  	dtype=hidden_states.dtype,
  	device=hidden_states.device)

Some other per-call tensors such as rows_per_expert and unpermuted_row_to_permuted_row are also allocated each call, but the large activation/GEMM workspaces are the main concern.

The current mini test scope also does not cover EP fused MoE paths, making it easier to miss allocation-path differences in the configurations used by MoE serving:

vllm-xpu-kernels/tests/fused_moe/test_fused_moe.py

Lines 25 to 34 in dae50e2

	# override pytest parameters when enable mini pytest
	MINI_PYTEST_PARAMS = {
	"default": {
	"m,n,k": [(1, 256, 128)],
	"e": [2],
	"topk": [1],
	"dtype": [torch.bfloat16],
	"has_bias": [True],
	},
	}

Repro / measurement sketch

A simple allocation-count or timing repro can be written around XpuFusedMoe.apply():

import torch
from vllm_xpu_kernels.fused_moe_interface import XpuFusedMoe

DEVICE = "xpu"
dtype = torch.bfloat16
input_len = 1
hidden_size = 1024
intermediate_size = 2048
num_experts = 16
topk = 1

hidden_states = torch.randn((input_len, hidden_size), device=DEVICE, dtype=dtype)
w13 = torch.randn((num_experts, 2 * intermediate_size, hidden_size), device=DEVICE, dtype=dtype)
w2 = torch.randn((num_experts, hidden_size, intermediate_size), device=DEVICE, dtype=dtype)
scores = torch.randn((input_len, num_experts), device=DEVICE, dtype=torch.float32)
topk_weights, topk_ids = torch.topk(scores, k=topk, dim=-1, sorted=False)

# Match the layout expected by the current test path.
w13 = w13.transpose(-1, -2).contiguous()
w2 = w2.transpose(-1, -2).contiguous()

moe = XpuFusedMoe(
    w13=w13,
    w13_scales=None,
    w13_bias=None,
    w2=w2,
    w2_scales=None,
    w2_bias=None,
    n_experts_per_token=topk,
    activation="silu",
    num_experts=num_experts,
)

out = torch.empty_like(hidden_states)
for _ in range(1000):
    moe.apply(out, hidden_states, topk_weights, topk_ids)

Under profiling or allocation tracing, each call creates new scratch buffers for remap, GEMM1, activation, and GEMM2. This creates avoidable allocator churn in decode loops where the shapes are stable.

Expected behavior

The fused MoE wrapper should allow callers to supply reusable scratch buffers for the large temporary tensors used by the kernel path. The default API should preserve current behavior when no workspaces are provided, but optimized callers should be able to preallocate once and reuse across repeated calls.

Expected properties:

• Existing callers continue to work unchanged.
• Optional workspaces are validated for sufficient capacity.
• Reused workspace output matches the existing allocation path within the same tolerances as current fused MoE tests.
• Workspaces are not reused in a way that aliases live data before it has been consumed.
• EP and quantized fused MoE paths remain covered by tests.

Suggested tests

1. Add a focused test that compares output from the current allocation path vs. a reusable-workspace path for a small deterministic fused MoE case.
2. Extend mini test parameters to include representative EP fused MoE paths, including int4/MXFP4 variants where applicable.
3. Validate undersized workspace handling, if workspaces are exposed as an API.
4. Run the existing fused MoE test file to ensure the default allocation path remains unchanged.

Why this matters

For long-running MoE inference, repeated XPU allocations inside each layer call are unnecessary overhead. Allowing workspace reuse would reduce allocator pressure and make the Python wrapper friendlier to high-frequency decode workloads without changing the underlying kernels or default behavior.