dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

[moe] BLOCK_SIZE_M heuristic refined for BI-V100 decode workload

Browse Source 
CCCL saxpy.cu demonstrates the principle: fused operations should minimize
wasted work. The saxpy_fast (single transform) vs saxpy_slow (two transforms)
comparison shows that eliminating unnecessary memory round-trips is the
primary optimization lever for element-wise ops.

Applied to MoE: during decode, M=8 (max-num-seqs) × topk=8 = 64 tokens.
Old heuristic: numel≤64 → BLOCK_SIZE_M=32 → 2 tiles of 32, no waste.
But for smaller batches (M=1,2,4 × topk=8 = 8,16,32 tokens):
  BLOCK_SIZE_M=32 → tile padding: 24/16/0 rows wasted per tile
  BLOCK_SIZE_M=16 → tile padding: 8/0/0 rows wasted per tile

New heuristic adds a finer-grained tier:
  numel ≤ 16  → BLOCK_SIZE_M = 16  (zero waste for ≤2 seqs)
  numel ≤ 64  → BLOCK_SIZE_M = 32  (was: same, no change)
  numel ≤ 1024 → BLOCK_SIZE_M = 64  (was: same, no change)
  else → BLOCK_SIZE_M = 256          (was: same, no change)

ixformer only reads BLOCK_SIZE_M from the config dict. The 16→32 threshold
matters for low-batch decode on BI-V100 where 16 SMs benefit from more
tiles with less padding over fewer tiles with more padding.

Source: cccl_upstream/thrust/examples/saxpy.cu (fusion + waste minimization)
This commit is contained in:
project_6
2026-08-05 03:17:53 +00:00
parent 5379a573ac
commit 6bf73bdacb
1 changed files with 10 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
vllm/model_executor/layers/fused_moe/fused_moe.py
View File

@@ -353,7 +353,16 @@ def get_default_config(
'GROUP_SIZE_M': 1
}
numel = M * topk
if numel <= 64:
# CCCL principle from saxpy.cu: fused ops should minimize wasted padding.
# For BI-V100 decode: M=8 seqs × topk=8 experts = 64 active tokens.
# BLOCK_SIZE_M=32 → 50% padding waste (32-token tiles for 64 tokens = 2 tiles, ok)
# BLOCK_SIZE_M=16 → 0% waste for numel≤16, minimal waste for 16<numel≤64
# ixformer only reads BLOCK_SIZE_M from config — N/K/GROUP are ignored.
# Smaller BLOCK_SIZE_M = more tiles but less wasted computation per tile.
# On BI-V100 (16 SMs), more smaller tiles better saturate the SMs.
if numel <= 16:
config['BLOCK_SIZE_M'] = 16
elif numel <= 64:
config['BLOCK_SIZE_M'] = 32
elif numel <= 1024:
config['BLOCK_SIZE_M'] = 64