@tduncan / Nvidia-Optimization-Stack-VLLM / commits / e8e7b592d1

[Kernel][MoE] Tune block-FP8 fused MoE for low-batch decode (#46642)

Signed-off-by: mgoin <mgoin64@gmail.com>

Michael Goin committed Jun 25, 2026 at 13:38 UTC e8e7b592d11d3ea5d5cdb78ee0f1ab4c5b440667
10 files changed +110 -88
benchmarks/kernels/benchmark_moe.py
+10 -7
@@ -391,16 +391,19 @@ def get_configs_compute_bound(use_fp16, block_quant_shape) -> list[dict[str, int
391 config = dict(zip(keys, config_values))
392 configs.append(config)
393
394 - # Remove configs that are not compatible with fp8 block quantization
395 - # BLOCK_SIZE_K must be a multiple of block_k
396 - # BLOCK_SIZE_N must be a multiple of block_n
394 + # Drop configs incompatible with fp8 block quantization. A tile must align
395 + # to the quant-block scale grid, i.e. tile and block must divide one
396 + # another. The kernel indexes scales per element (offs_bn // group_n,
397 + # k_start // group_k), so a tile narrower than the block (e.g. N=64 with
398 + # block_n=128) is valid -- and often faster at small batch. An exact
399 + # multiple was required before, which dropped those smaller tiles entirely.
400 if block_quant_shape is not None and not use_fp16:
401 block_n, block_k = block_quant_shape[0], block_quant_shape[1]
402 for config in configs[:]:
400 - if (
401 - config["BLOCK_SIZE_K"] % block_k != 0
402 - or config["BLOCK_SIZE_N"] % block_n != 0
403 - ):
403 + bn, bk = config["BLOCK_SIZE_N"], config["BLOCK_SIZE_K"]
404 + n_aligned = bn % block_n == 0 or block_n % bn == 0
405 + k_aligned = bk % block_k == 0 or block_k % bk == 0
406 + if not (n_aligned and k_aligned):
407 configs.remove(config)
408 return configs
409
vllm/model_executor/layers/fused_moe/configs/E=128,N=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json
+10 -10
@@ -2,35 +2,35 @@
2 "triton_version": "3.5.0",
3 "1": {
4 "BLOCK_SIZE_M": 16,
5 - "BLOCK_SIZE_N": 128,
5 + "BLOCK_SIZE_N": 64,
6 "BLOCK_SIZE_K": 128,
7 "GROUP_SIZE_M": 1,
8 "num_warps": 4,
9 "num_stages": 4
10 },
11 "2": {
12 - "BLOCK_SIZE_M": 16,
13 - "BLOCK_SIZE_N": 128,
12 + "BLOCK_SIZE_M": 8,
13 + "BLOCK_SIZE_N": 64,
14 "BLOCK_SIZE_K": 128,
15 "GROUP_SIZE_M": 1,
16 "num_warps": 4,
17 "num_stages": 4
18 },
19 "4": {
20 - "BLOCK_SIZE_M": 16,
21 - "BLOCK_SIZE_N": 128,
20 + "BLOCK_SIZE_M": 8,
21 + "BLOCK_SIZE_N": 64,
22 "BLOCK_SIZE_K": 128,
23 - "GROUP_SIZE_M": 64,
23 + "GROUP_SIZE_M": 1,
24 "num_warps": 4,
25 - "num_stages": 4
25 + "num_stages": 3
26 },
27 "8": {
28 "BLOCK_SIZE_M": 16,
29 - "BLOCK_SIZE_N": 128,
29 + "BLOCK_SIZE_N": 64,
30 "BLOCK_SIZE_K": 128,
31 - "GROUP_SIZE_M": 64,
31 + "GROUP_SIZE_M": 16,
32 "num_warps": 4,
33 - "num_stages": 4
33 + "num_stages": 3
34 },
35 "16": {
36 "BLOCK_SIZE_M": 16,
vllm/model_executor/layers/fused_moe/configs/E=128,N=768,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json
+9 -9
@@ -1,17 +1,17 @@
1 {
2 "1": {
3 - "BLOCK_SIZE_M": 64,
3 + "BLOCK_SIZE_M": 16,
4 "BLOCK_SIZE_N": 64,
5 "BLOCK_SIZE_K": 128,
6 "GROUP_SIZE_M": 1,
7 "num_warps": 4,
8 - "num_stages": 3
8 + "num_stages": 5
9 },
10 "2": {
11 - "BLOCK_SIZE_M": 64,
12 - "BLOCK_SIZE_N": 128,
11 + "BLOCK_SIZE_M": 16,
12 + "BLOCK_SIZE_N": 64,
13 "BLOCK_SIZE_K": 128,
14 - "GROUP_SIZE_M": 32,
14 + "GROUP_SIZE_M": 1,
15 "num_warps": 4,
16 "num_stages": 4
17 },
@@ -25,11 +25,11 @@
25 },
26 "8": {
27 "BLOCK_SIZE_M": 16,
28 - "BLOCK_SIZE_N": 128,
28 + "BLOCK_SIZE_N": 64,
29 "BLOCK_SIZE_K": 128,
30 - "GROUP_SIZE_M": 16,
30 + "GROUP_SIZE_M": 1,
31 "num_warps": 4,
32 - "num_stages": 5
32 + "num_stages": 3
33 },
34 "16": {
35 "BLOCK_SIZE_M": 16,
@@ -143,4 +143,4 @@
143 "num_warps": 4,
144 "num_stages": 4
145 }
146 -}
\ No newline at end of file
146 +}
vllm/model_executor/layers/fused_moe/configs/E=160,N=640,device_name=NVIDIA_H100,dtype=fp8_w8a8,block_shape=[128,128].json
+7 -7
@@ -1,7 +1,7 @@
1 {
2 "1": {
3 "BLOCK_SIZE_M": 16,
4 - "BLOCK_SIZE_N": 128,
4 + "BLOCK_SIZE_N": 64,
5 "BLOCK_SIZE_K": 128,
6 "GROUP_SIZE_M": 1,
7 "num_warps": 4,
@@ -9,11 +9,11 @@
9 },
10 "2": {
11 "BLOCK_SIZE_M": 16,
12 - "BLOCK_SIZE_N": 128,
12 + "BLOCK_SIZE_N": 64,
13 "BLOCK_SIZE_K": 128,
14 "GROUP_SIZE_M": 1,
15 "num_warps": 4,
16 - "num_stages": 3
16 + "num_stages": 4
17 },
18 "4": {
19 "BLOCK_SIZE_M": 16,
@@ -24,12 +24,12 @@
24 "num_stages": 3
25 },
26 "8": {
27 - "BLOCK_SIZE_M": 16,
28 - "BLOCK_SIZE_N": 128,
27 + "BLOCK_SIZE_M": 32,
28 + "BLOCK_SIZE_N": 64,
29 "BLOCK_SIZE_K": 128,
30 - "GROUP_SIZE_M": 16,
30 + "GROUP_SIZE_M": 1,
31 "num_warps": 4,
32 - "num_stages": 5
32 + "num_stages": 3
33 },
34 "16": {
35 "BLOCK_SIZE_M": 16,
vllm/model_executor/layers/fused_moe/configs/E=256,N=256,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json
+13 -13
@@ -1,35 +1,35 @@
1 {
2 "1": {
3 - "BLOCK_SIZE_M": 16,
4 - "BLOCK_SIZE_N": 128,
3 + "BLOCK_SIZE_M": 8,
4 + "BLOCK_SIZE_N": 32,
5 "BLOCK_SIZE_K": 128,
6 - "GROUP_SIZE_M": 16,
6 + "GROUP_SIZE_M": 1,
7 "num_warps": 4,
8 "num_stages": 4
9 },
10 "2": {
11 - "BLOCK_SIZE_M": 16,
12 - "BLOCK_SIZE_N": 128,
11 + "BLOCK_SIZE_M": 8,
12 + "BLOCK_SIZE_N": 32,
13 "BLOCK_SIZE_K": 128,
14 - "GROUP_SIZE_M": 16,
14 + "GROUP_SIZE_M": 1,
15 "num_warps": 4,
16 - "num_stages": 4
16 + "num_stages": 5
17 },
18 "4": {
19 - "BLOCK_SIZE_M": 16,
20 - "BLOCK_SIZE_N": 128,
19 + "BLOCK_SIZE_M": 8,
20 + "BLOCK_SIZE_N": 64,
21 "BLOCK_SIZE_K": 128,
22 - "GROUP_SIZE_M": 64,
22 + "GROUP_SIZE_M": 1,
23 "num_warps": 4,
24 "num_stages": 4
25 },
26 "8": {
27 "BLOCK_SIZE_M": 16,
28 - "BLOCK_SIZE_N": 128,
28 + "BLOCK_SIZE_N": 64,
29 "BLOCK_SIZE_K": 128,
30 - "GROUP_SIZE_M": 16,
30 + "GROUP_SIZE_M": 32,
31 "num_warps": 4,
32 - "num_stages": 4
32 + "num_stages": 3
33 },
34 "16": {
35 "BLOCK_SIZE_M": 16,
vllm/model_executor/layers/fused_moe/configs/E=256,N=384,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128,128].json
+7 -7
@@ -2,19 +2,19 @@
2 "triton_version": "3.5.0",
3 "1": {
4 "BLOCK_SIZE_M": 16,
5 - "BLOCK_SIZE_N": 128,
5 + "BLOCK_SIZE_N": 64,
6 "BLOCK_SIZE_K": 128,
7 "GROUP_SIZE_M": 1,
8 "num_warps": 4,
9 - "num_stages": 4
9 + "num_stages": 5
10 },
11 "2": {
12 - "BLOCK_SIZE_M": 16,
13 - "BLOCK_SIZE_N": 128,
12 + "BLOCK_SIZE_M": 8,
13 + "BLOCK_SIZE_N": 64,
14 "BLOCK_SIZE_K": 128,
15 "GROUP_SIZE_M": 1,
16 "num_warps": 4,
17 - "num_stages": 4
17 + "num_stages": 5
18 },
19 "4": {
20 "BLOCK_SIZE_M": 16,
@@ -34,11 +34,11 @@
34 },
35 "16": {
36 "BLOCK_SIZE_M": 16,
37 - "BLOCK_SIZE_N": 128,
37 + "BLOCK_SIZE_N": 64,
38 "BLOCK_SIZE_K": 128,
39 "GROUP_SIZE_M": 1,
40 "num_warps": 4,
41 - "num_stages": 5
41 + "num_stages": 3
42 },
43 "24": {
44 "BLOCK_SIZE_M": 16,
vllm/model_executor/layers/fused_moe/configs/E=256,N=512,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json
+7 -7
@@ -1,20 +1,20 @@
1 {
2 "triton_version": "3.6.0",
3 "1": {
4 - "BLOCK_SIZE_M": 16,
5 - "BLOCK_SIZE_N": 128,
4 + "BLOCK_SIZE_M": 8,
5 + "BLOCK_SIZE_N": 64,
6 "BLOCK_SIZE_K": 128,
7 "GROUP_SIZE_M": 1,
8 "num_warps": 4,
9 - "num_stages": 3
9 + "num_stages": 4
10 },
11 "2": {
12 "BLOCK_SIZE_M": 16,
13 - "BLOCK_SIZE_N": 128,
13 + "BLOCK_SIZE_N": 64,
14 "BLOCK_SIZE_K": 128,
15 - "GROUP_SIZE_M": 16,
16 - "num_warps": 8,
17 - "num_stages": 4
15 + "GROUP_SIZE_M": 1,
16 + "num_warps": 4,
17 + "num_stages": 3
18 },
19 "4": {
20 "BLOCK_SIZE_M": 16,
vllm/model_executor/layers/fused_moe/configs/E=384,N=128,device_name=NVIDIA_H200,dtype=fp8_w8a8,block_shape=[128,128].json
+12 -12
@@ -1,23 +1,23 @@
1 {
2 "1": {
3 "BLOCK_SIZE_M": 16,
4 - "BLOCK_SIZE_N": 128,
4 + "BLOCK_SIZE_N": 64,
5 "BLOCK_SIZE_K": 128,
6 "GROUP_SIZE_M": 1,
7 "num_warps": 4,
8 "num_stages": 3
9 },
10 "2": {
11 - "BLOCK_SIZE_M": 16,
12 - "BLOCK_SIZE_N": 128,
11 + "BLOCK_SIZE_M": 8,
12 + "BLOCK_SIZE_N": 64,
13 "BLOCK_SIZE_K": 128,
14 "GROUP_SIZE_M": 1,
15 "num_warps": 4,
16 - "num_stages": 4
16 + "num_stages": 5
17 },
18 "4": {
19 - "BLOCK_SIZE_M": 16,
20 - "BLOCK_SIZE_N": 128,
19 + "BLOCK_SIZE_M": 8,
20 + "BLOCK_SIZE_N": 64,
21 "BLOCK_SIZE_K": 128,
22 "GROUP_SIZE_M": 1,
23 "num_warps": 4,
@@ -25,19 +25,19 @@
25 },
26 "8": {
27 "BLOCK_SIZE_M": 16,
28 - "BLOCK_SIZE_N": 128,
28 + "BLOCK_SIZE_N": 64,
29 "BLOCK_SIZE_K": 128,
30 - "GROUP_SIZE_M": 1,
30 + "GROUP_SIZE_M": 32,
31 "num_warps": 4,
32 - "num_stages": 4
32 + "num_stages": 3
33 },
34 "16": {
35 "BLOCK_SIZE_M": 16,
36 - "BLOCK_SIZE_N": 128,
36 + "BLOCK_SIZE_N": 64,
37 "BLOCK_SIZE_K": 128,
38 - "GROUP_SIZE_M": 1,
38 + "GROUP_SIZE_M": 16,
39 "num_warps": 4,
40 - "num_stages": 4
40 + "num_stages": 3
41 },
42 "24": {
43 "BLOCK_SIZE_M": 16,
vllm/model_executor/layers/fused_moe/configs/E=512,N=256,device_name=NVIDIA_H100_80GB_HBM3,dtype=fp8_w8a8,block_shape=[128,128].json
+9 -9
@@ -2,27 +2,27 @@
2 "triton_version": "3.4.0",
3 "1": {
4 "BLOCK_SIZE_M": 16,
5 - "BLOCK_SIZE_N": 128,
5 + "BLOCK_SIZE_N": 64,
6 "BLOCK_SIZE_K": 128,
7 "GROUP_SIZE_M": 1,
8 "num_warps": 4,
9 - "num_stages": 4
9 + "num_stages": 5
10 },
11 "2": {
12 - "BLOCK_SIZE_M": 16,
13 - "BLOCK_SIZE_N": 128,
12 + "BLOCK_SIZE_M": 8,
13 + "BLOCK_SIZE_N": 64,
14 "BLOCK_SIZE_K": 128,
15 "GROUP_SIZE_M": 1,
16 "num_warps": 4,
17 "num_stages": 4
18 },
19 "4": {
20 - "BLOCK_SIZE_M": 16,
21 - "BLOCK_SIZE_N": 128,
20 + "BLOCK_SIZE_M": 8,
21 + "BLOCK_SIZE_N": 64,
22 "BLOCK_SIZE_K": 128,
23 "GROUP_SIZE_M": 1,
24 "num_warps": 4,
25 - "num_stages": 3
25 + "num_stages": 4
26 },
27 "8": {
28 "BLOCK_SIZE_M": 16,
@@ -36,9 +36,9 @@
36 "BLOCK_SIZE_M": 16,
37 "BLOCK_SIZE_N": 128,
38 "BLOCK_SIZE_K": 128,
39 - "GROUP_SIZE_M": 1,
39 + "GROUP_SIZE_M": 32,
40 "num_warps": 4,
41 - "num_stages": 4
41 + "num_stages": 3
42 },
43 "24": {
44 "BLOCK_SIZE_M": 16,
vllm/model_executor/layers/fused_moe/fused_moe.py
+26 -7
@@ -1257,19 +1257,38 @@ def get_default_config(
1257 num_stages_rocm = 2
1258
1259 if dtype == "fp8_w8a8" and block_shape is not None:
1260 - # Block-wise quant: tile sizes are constrained by block_shape.
1261 - # Use a small M tile for decode-like batches where tokens are
1262 - # spread thin across experts. Larger batches benefit from
1263 - # GROUP_SIZE_M > 1 because the per-block scales add memory
1264 - # traffic that benefits from L2 tile reuse.
1260 + # Block-wise quant. Use a small M tile for decode-like batches where
1261 + # tokens are spread thin across experts. Larger batches benefit from
1262 + # GROUP_SIZE_M > 1 because the per-block scales add memory traffic
1263 + # that benefits from L2 tile reuse.
1264 + #
1265 + # BLOCK_SIZE_N need not equal block_shape[0]: the kernel indexes block
1266 + # scales per N element (offs_bn // group_n), so any N tile dividing the
1267 + # quant block is valid. At decode a 128-wide N tile leaves the gate-up
1268 + # GEMM SM-bound; a 64-wide tile exposes ~2x the thread blocks, and the
1269 + # swap-AB kernel keeps it efficient on Hopper down to the smallest
1270 + # batches, so prefer N=64 through low batch sizes. CUDA only (validated
1271 + # on NVIDIA); ROCm keeps its prior tile/pipeline sizes.
1272 + if current_platform.is_rocm():
1273 + block_n = block_shape[0]
1274 + num_stages = num_stages_rocm
1275 + elif M <= 8 and block_shape[0] % 64 == 0:
1276 + block_n = 64
1277 + # The smallest batches are memory-latency bound, so a deeper
1278 + # pipeline hides the weight loads; by M=8 it turns occupancy/SMEM
1279 + # bound and the extra stages hurt.
1280 + num_stages = 4 if M <= 4 else 3
1281 + else:
1282 + block_n = block_shape[0]
1283 + num_stages = 3
1284 config = {
1285 "BLOCK_SIZE_M": 16 if M <= 64 else 64,
1267 - "BLOCK_SIZE_N": block_shape[0],
1286 + "BLOCK_SIZE_N": block_n,
1287 "BLOCK_SIZE_K": block_shape[1],
1288 "GROUP_SIZE_M": 1 if M <= 16 else 32,
1289 "SPLIT_K": 1,
1290 "num_warps": 4,
1272 - "num_stages": 3 if not current_platform.is_rocm() else num_stages_rocm,
1291 + "num_stages": num_stages,
1292 }
1293 elif dtype in ["int4_w4a16", "int8_w8a16"] and block_shape is not None:
1294 # moe wna16 kernels