[analysis] MoE execution path: 640+ kernel launches/step, 192 CUDA mallocs/step

Critical finding from reading fused_moe.py end-to-end:

The real decode bottleneck is NOT tuning parameters. It's:
1. 640+ ixformer kernel launches per decode step (64 MoE layers ×
   ~10 ops each). At target 395 TPS = 253K launches/second.
2. 192 torch.empty calls per step (3 intermediate caches × 64 layers).
3. Python-level dispatch overhead for each of these calls.

The BLOCK_SIZE_M heuristic is already reasonable (16 for decode).
The fused_moe Triton kernel is dead code — ixformer's C++ kernel
is called instead.

Actionable optimization: pre-allocate intermediate caches outside the
layer loop to eliminate 192 CUDA mallocs per decode step.

Source: vllm/model_executor/layers/fused_moe/fused_moe.py
        vllm/_custom_ops.py (ixf_F.vllm_invoke_fused_moe_kernel)
        cccl_upstream/cub/cub/device/dispatch/tuning/tuning_batch_memcpy.cuh
This commit is contained in:
project_6
2026-08-05 03:58:06 +00:00
parent afb5d23b12
commit 4e16133c7a

View File

@@ -0,0 +1,60 @@
# MoE Execution Path Analysis
> Source: vllm/model_executor/layers/fused_moe/fused_moe.py + vllm/_custom_ops.py
> Read: 2026-08-04
## The Real Bottleneck
Qwen3.6-35B-A3B has 64 MoE layers, each with:
- 256 experts, top-8 routing
- Gate up projection (w1): hidden_dim → intermediate_dim
- SiLU activation
- Down projection (w2): intermediate_dim → hidden_dim
- Weighted sum of 8 expert outputs
### Per-decode-step kernel launches
| Operation | Count | Implementation |
|---|---|---|
| fused_moe_kernel (w1) | 64 | ixf_F.vllm_invoke_fused_moe_kernel |
| silu_and_mul | 64 | ixf_F.silu_and_mul |
| fused_moe_kernel (w2) | 64 | ixf_F.vllm_invoke_fused_moe_kernel |
| topk_softmax | 64 | ixf_F.vllm_moe_topk_softmax |
| moe_align_block_size | 64 | ixf_F.vllm_moe_align_block_size |
| torch.sum (expert merge) | 64 | PyTorch |
| paged_attention_v1 | 1 | ixf_F.vllm_single_query_cached_kv_attention |
| rms_norm | 128 | ixf_F.rms_norm |
| fused_add_rms_norm | 64 | ixf_F.fused_add_rms_norm |
| rotary_embedding | 64 | ixf_F.vllm_rotary_embedding_neox |
| **Total** | **~640+** | |
640+ kernel launches per decode step. At target Output TPS ≥ 395,
that's 395 × 640 = 253,000 kernel launches per second.
### Memory allocation per step
```python
# Inside fused_experts, called 64 times per step:
intermediate_cache1 = torch.empty((M, topk, N)) # 64 × alloc
intermediate_cache2 = torch.empty((M * topk, N // 2)) # 64 × alloc
intermediate_cache3 = torch.empty((M, topk, w2_shape[1])) # 64 × alloc
```
192 torch.empty calls per decode step = 192 CUDA mallocs.
At 395 TPS = 75,840 mallocs/second.
### What we can actually change
1. **BLOCK_SIZE_M** (passed to ixformer): 16 for decode (numel=8, M=1×topk=8)
- Already optimized: 16 for ≤16 tokens, 32 for ≤64, 64 for ≤1024
- ixformer may or may not respect N/K/GROUP values
2. **Intermediate cache pre-allocation**: move torch.empty outside the layer loop
- Allocate once, reuse across 64 layers
- Saves 192 CUDA mallocs per decode step
3. **torch.sum → ixformer?**: the expert merge `torch.sum(dim=1)` is PyTorch,
could potentially be fused into the second fused_moe_kernel call
4. **Chunk size**: VLLM_FUSED_MOE_CHUNK_SIZE controls batching.
For decode M=1, chunking adds overhead for no benefit.