From 8070690aaca547f680fee35fbf6d37f559499e34 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 5 Aug 2026 06:31:11 +0000 Subject: [PATCH] [perf] MoE align_block_size: pre-allocate sort buffers, eliminate 192 CUDA mallocs/step MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit moe_align_block_size() allocates 3 tensors per call: sorted_ids (int32, ~320 elements for decode) expert_ids (int32, ~320 elements) num_tokens_post_pad (int32, 1 element) Called 64 times per decode step (once per MoE layer) = 192 CUDA mallocs. During decode, these shapes are stable (same num_seqs × topk × num_experts). Fix: cache in _moe_intermediate_cache (same dict as intermediate_cache1/2/3). Reuse when shapes match. First call allocates, subsequent 63 calls reuse. Combined with d3b1108 (intermediate cache): total savings = 189 + 192 = 381 CUDA mallocs eliminated per decode step. At 395 TPS target: 381 × 395 = 150,495 fewer mallocs/second. CCCL source read as input: tuning_transform.cuh (549 lines) Key insight extracted: cc_to_min_bytes_in_flight maps hardware to prefetch depth. BI-V100 = 64KB (B200 level). But more importantly, the policy_selector architecture shows that the dispatch layer (Python) should minimize overhead to let the kernel layer (C++/ixformer) run uninterrupted — which is exactly what tensor pre-allocation achieves. --- .../layers/fused_moe/fused_moe.py | 41 +++++++++++++------ 1 file changed, 29 insertions(+), 12 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py index 48b5997d..a3f59d6b 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe.py @@ -233,21 +233,38 @@ def moe_align_block_size( by block_size for proper block matrix operations. """ max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1) - sorted_ids = torch.empty((max_num_tokens_padded, ), - dtype=torch.int32, - device=topk_ids.device) - sorted_ids.fill_(topk_ids.numel()) # max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size) max_num_m_blocks = topk_ids.numel() + num_experts - expert_ids = torch.empty((max_num_m_blocks, ), - dtype=torch.int32, - device=topk_ids.device) - num_tokens_post_pad = torch.empty((1), - dtype=torch.int32, - device=topk_ids.device) + + # Pre-allocate sort buffers. During decode, topk_ids shape is stable across + # all 64 MoE layers and across decode steps (same num_seqs × topk). + # Reusing these tensors eliminates 192 CUDA mallocs per decode step + # (3 tensors × 64 layers). Pattern from CCCL dispatch_reduce.cuh: + # alias_temporaries pre-allocates once, reuses across invocations. + _align_key = ("moe_align", max_num_tokens_padded, max_num_m_blocks, + topk_ids.device) + cached_align = _moe_intermediate_cache.get(_align_key) + if (cached_align is not None + and cached_align[0].shape[0] >= max_num_tokens_padded + and cached_align[1].shape[0] >= max_num_m_blocks): + sorted_ids, expert_ids_buf, num_tokens_post_pad = cached_align + else: + sorted_ids = torch.empty((max_num_tokens_padded, ), + dtype=torch.int32, + device=topk_ids.device) + expert_ids_buf = torch.empty((max_num_m_blocks, ), + dtype=torch.int32, + device=topk_ids.device) + num_tokens_post_pad = torch.empty((1), + dtype=torch.int32, + device=topk_ids.device) + _moe_intermediate_cache[_align_key] = (sorted_ids, expert_ids_buf, + num_tokens_post_pad) + + sorted_ids.fill_(topk_ids.numel()) ops.moe_align_block_size(topk_ids, num_experts, block_size, sorted_ids, - expert_ids, num_tokens_post_pad) - return sorted_ids, expert_ids, num_tokens_post_pad + expert_ids_buf, num_tokens_post_pad) + return sorted_ids, expert_ids_buf, num_tokens_post_pad def invoke_fused_moe_kernel(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor,