[perf] MoE intermediate cache pre-allocation: eliminate 189 CUDA mallocs per decode step

fused_experts() is called 64 times per decode step (once per MoE layer).
Each call allocated 3 intermediate tensors via torch.empty = 192 mallocs.
For decode (M=1, topk=8), all 64 calls use identical shapes.

Fix: module-level _moe_intermediate_cache dict that reuses tensors when
shapes match. First layer call allocates, subsequent 63 calls reuse.
Saves 189 CUDA mallocs per decode step = 74,655 mallocs/second at 395 TPS.

Design follows CCCL's dispatch_reduce.cuh pattern: pre-allocate temp_storage
once via alias_temporaries, reuse across kernel invocations.

No functional change — tensors are .empty() (uninitialized), overwritten
before use by ixformer kernels.
This commit is contained in:
project_6
2026-08-05 03:58:46 +00:00
parent 4e16133c7a
commit d3b110803c

View File

@@ -15,6 +15,25 @@ from vllm.platforms import current_platform
logger = init_logger(__name__)
# Module-level cache for MoE intermediate tensors.
# Eliminates 192 torch.empty (CUDA malloc) calls per decode step by reusing
# buffers across the 64 MoE layer invocations within a single forward pass.
# Key insight from CCCL dispatch_reduce.cuh: NVIDIA pre-allocates temp_storage
# once and reuses it across kernel invocations rather than re-allocating.
_moe_intermediate_cache = {}
def _get_or_alloc(name, shape, dtype, device):
"""Get a cached tensor or allocate a new one. Reuses if shape fits."""
key = (name, dtype, device)
cached = _moe_intermediate_cache.get(key)
if cached is not None and cached.shape == shape:
return cached
# Shape changed (different M, different topk) — reallocate
t = torch.empty(shape, dtype=dtype, device=device)
_moe_intermediate_cache[key] = t
return t
@triton.jit
def fused_moe_kernel(
@@ -534,15 +553,15 @@ def fused_experts(hidden_states: torch.Tensor,
config = get_config_func(M)
intermediate_cache1 = torch.empty((M, topk_ids.shape[1], N),
device=hidden_states.device,
dtype=hidden_states.dtype)
intermediate_cache2 = torch.empty((M * topk_ids.shape[1], N // 2),
device=hidden_states.device,
dtype=hidden_states.dtype)
intermediate_cache3 = torch.empty((M, topk_ids.shape[1], w2.shape[1]),
device=hidden_states.device,
dtype=hidden_states.dtype)
intermediate_cache1 = _get_or_alloc(
'moe_c1', (M, topk_ids.shape[1], N),
hidden_states.dtype, hidden_states.device)
intermediate_cache2 = _get_or_alloc(
'moe_c2', (M * topk_ids.shape[1], N // 2),
hidden_states.dtype, hidden_states.device)
intermediate_cache3 = _get_or_alloc(
'moe_c3', (M, topk_ids.shape[1], w2.shape[1]),
hidden_states.dtype, hidden_states.device)
compute_type = (tl.bfloat16
if hidden_states.dtype == torch.bfloat16 else tl.float16)