From 6bf73bdacba09aaea0ccf5de36fc2c7f987d5971 Mon Sep 17 00:00:00 2001
From: project_6 <dylan@project6.dev>
Date: Wed, 5 Aug 2026 03:17:53 +0000
Subject: [PATCH] [moe] BLOCK_SIZE_M heuristic refined for BI-V100 decode
 workload
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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)
---
 vllm/model_executor/layers/fused_moe/fused_moe.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py
index 03ffc2f2..a365fc73 100644
--- a/vllm/model_executor/layers/fused_moe/fused_moe.py
+++ b/vllm/model_executor/layers/fused_moe/fused_moe.py
@@ -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