From 1950ef70cb8d835837840890d9445c6621667f62 Mon Sep 17 00:00:00 2001 From: Stephan Toivonen Date: Tue, 28 Jul 2026 15:21:56 -0600 Subject: [PATCH 1/2] [ROCm][Kernel] W4A16 prefill: retune small-group-size Triton tiles The group_size <= 64 branch of _triton_config (added in f47a359c96 for Qwen3-Omni) was derived from a fixed-M sweep at a single captured problem size. Two of its assumptions do not hold across the prefill batch range: - num_stages=1 was inherited from the group_size=128 branch, whose "num_stages>1 regresses (weak ROCm Triton pipeliner)" note was measured at BM=32, BN=16/64, nw=2 -- much narrower tiles than the ones this branch emits. - gemm2's (BN=64, nw=4) is not the best tile at any M swept. Tiling x M sweep on Strix Halo (gfx1151) at the Qwen3-30B-A3B-AWQ shape (group_size=32, E=128, top_k=8, alignment=128): 108 configs x 10 token counts (M_routed 48..24448) x 2 routing models -- the captured expert concentration and a uniform 128-expert spread. Configs ranked by geomean of (us / per-M best) over all 20 points, so the pick holds across the batch range rather than at the single M a fixed-M sweep sees. Score 1.000 = best at every M; "worst" is the worst-case ratio. gemm1 (K>=1024) geomean worst BM128 BN64 GM1 nw8 ns4 (new) 1.040 1.146 BM128 BN64 GM1 nw8 ns1 (old) 1.284 1.933 BM128 BN128 GM1 nw8 ns4 1.162 1.770 gemm2 (K<1024) geomean worst BM64 BN128 GM1 nw8 ns4 (new) 1.043 1.087 BM64 BN64 GM1 nw4 ns1 (old) 1.109 1.207 BM128 BN128 GM1 nw8 ns4 1.044 1.153 num_stages=4 is the one universal win -- both GEMMs, both routing models, every batch size swept. For gemm1 that is the only field that moves: BN=128 wins only near M_routed=8192 and is up to 1.77x off the best at small batches, and nw=4 ties nw=8 at 1.040 so num_warps stays. For gemm2, BM=128 ties on geomean (1.044 vs 1.043) but has a worse tail (1.153 vs 1.087), so BLOCK_SIZE_M stays at 64. These are pure tiling parameters -- no change to which kernel runs, to the alignment, or to the group_size > 64 path (byte-identical). End-to-end TTFT, vllm-bench.py on cyankiwi/Qwen3-30B-A3B-Instruct-2507- AWQ-4bit, 10 prompts, --max-num-seqs 1 --output-len 8 --no-cudagraph --dtype float16, median ms: input-len before after 128 205.11 184.84 (1.11x) 512 288.80 270.43 (1.07x) 1024 408.45 393.62 (1.04x) Improvement is consistent across mean/median/P99. Signed-off-by: Stephan Toivonen Co-authored-by: Claude --- .../layers/fused_moe/hybrid_w4a16_moe.py | 24 +++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/hybrid_w4a16_moe.py b/vllm/model_executor/layers/fused_moe/hybrid_w4a16_moe.py index d455cde02673..2683f74b58f4 100644 --- a/vllm/model_executor/layers/fused_moe/hybrid_w4a16_moe.py +++ b/vllm/model_executor/layers/fused_moe/hybrid_w4a16_moe.py @@ -353,22 +353,38 @@ def _triton_config( if self._group_size <= HybridW4A16MoEExperts.SMALL_GROUP_SIZE_THRESHOLD: # gemm2 BM=64 < alignment=128; apply()'s _expert_ids_for # repeat_interleaves expert_ids to compensate. + # + # Swept on Strix Halo (gfx1151), Qwen3-30B-A3B AWQ shapes + # (group_size=32, E=128, alignment=128): 108 tile configs x 10 + # token counts (M_routed 48..24448) x 2 routing distributions. + # Picks below are the best geomean of (us / per-M best), so they + # hold across the batch range and not just at one M. num_stages=4 + # is the common win; the "pipeliner regresses" note in the + # large-group branch below was measured at its much narrower + # BLOCK_SIZE_N=16/64, num_warps=2 tiles. if K < 1024: + # gemm2 (K=768). Chose BLOCK_SIZE_N=128, num_warps=8, + # num_stages=4: geomean 1.043 vs 1.109 for the previous + # BLOCK_SIZE_N=64, num_warps=4, num_stages=1. BLOCK_SIZE_M=128 + # ties on geomean but has a worse tail, so it stays at 64. return dict( BLOCK_SIZE_M=64, - BLOCK_SIZE_N=64, + BLOCK_SIZE_N=128, BLOCK_SIZE_K=BLOCK_SIZE_K, GROUP_SIZE_M=1, - num_warps=4, - num_stages=1, + num_warps=8, + num_stages=4, ) + # gemm1 (K=2048). Only num_stages moves: 4 scores 1.040 vs 1.284 + # for 1. BLOCK_SIZE_N=128 is up to 1.77x off the best at small + # batches and num_warps=4 ties 8, so both stay put. return dict( BLOCK_SIZE_M=128, BLOCK_SIZE_N=64, BLOCK_SIZE_K=BLOCK_SIZE_K, GROUP_SIZE_M=1, num_warps=8, - num_stages=1, + num_stages=4, ) # Per-shape sweep on Strix Halo (gfx1151) at BLOCK_M=32 (the From f86313e1bea4c53305bc3e6c506a7f114601824c Mon Sep 17 00:00:00 2001 From: Stephan Toivonen Date: Tue, 28 Jul 2026 15:33:01 -0600 Subject: [PATCH 2/2] [ROCm][Kernel] W4A16 prefill: M-aware moe_align block size for small group_size _select_block_size_m returns TRITON_BLOCK_SIZE_M_SMALL_GS for every prefill batch, with no M-awareness. For a large-E model that is very expensive at short prefill, and it dominates anything the tiling can reach. moe_align_block_size pads *every touched expert* to a full block, and fused_moe_kernel_gptq_awq's only early exit is if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded: return so an all-padding block that lands inside num_tokens_post_padded runs the entire K-contraction -- it loads and dequantizes its full slice of INT4 weights and only discards the result at the store, via token_mask = offs_token < num_valid_tokens. (gemm2 does not even discard: it is invoked with num_valid_tokens = num_slots, so its padding rows are stored.) Kernel work is therefore proportional to num_tokens_post_padded, not to M_routed, and npp = sum over touched experts of ceil(tokens_e / A) * A so halving the alignment A halves the per-expert round-up. With E=128 a short prefill touches nearly every expert with a handful of tokens each and almost all of npp is padding. Measured on real routing -- topk_ids captured from live prefills of cyankiwi/Qwen3-30B-A3B-Instruct-2507-AWQ-4bit (E=128, top_k=8, group_size=32), averaged over MoE layers 0-3: tokens experts npp@64 npp@128 npp@128 / routed rows 6 33 2112 4224 88.0x 128 110 7120 14080 13.8x 512 122 9456 16576 4.0x 1024 122 12864 18752 2.3x 2048 124 20816 25984 1.6x A 128-token prefill does 13.8x the GEMM work it needs to. Alignment x tile x M sweep on Strix Halo (gfx1151), joint gemm1+gemm2 (they share one alignment, so the pair is what matters), us: tokens A=32 A=64 A=128 best 6 1509.2 876.1 1537.3 64 (1.75x vs A=128) 128 3485.3 2619.7 4655.7 64 (1.78x) 512 4391.1 3407.5 5453.1 64 (1.60x) 1024 6264.1 4754.0 6248.9 64 (1.31x) 2048 9838.3 8256.8 8355.5 64 (1.01x) A=64 wins at every point from 6 to 1024 tokens under all three routing models swept (1.14-1.85x) and is a wash at 2048. A=16 is worse than A=32 everywhere: with BLOCK_SIZE_M = A the block count is pinned to the number of touched experts regardless of A, so the per-block INT4 weight load + dequant is roughly constant and below 64 the M-tile can no longer amortize it. 64 is the floor, not an arbitrary halving. The crossover is where the mechanism says it should be -- one full block per expert, i.e. num_tokens * topk / E >= 128 -- so the tier is selected on routed rows per expert rather than a token count, and it tracks E and top_k on other models. This is the alignment, not the tile. Decomposed at 128 tokens (gemm1+gemm2, us): 5237.2 at the pre-retune config and A=128, 4755.6 at the retuned config and A=128, 4882.9 at the new tile and A=128, 2619.7 at the new tile and A=64. The new gemm1 tile is *worse* than the previous commit's at A=128; it only pays once the alignment moves, which is why the two land together. Keeping the A=128 tile at A=64 would give back 1.10-1.76x of the win. BN=32/nw=2 is best at all 27 (tokens x routing) points at A=64, by 1.10x to 2.13x over the runner- up; gemm2's tile is unchanged (geomean 1.000, worst 1.004 at A=64). Batches of 2048+ tokens take the _LARGE tier and get byte-identical config to before this patch, so f47a359c96's Qwen3-Omni tune -- which was validated at --input-len 4096 -- is unaffected. Only sub-2048 batches change, which is the regime that tune never measured. The group_size > 64 path is untouched. Alignment is not a numerical parameter: it regroups routed rows into blocks, and each row does the same full-K contraction against the same expert weights in the same order. Verified rather than assumed -- greedy decode over 3 prompts x 48 tokens (2x 128-token and 1x 512-token prefills, both on the changed tier) produces byte-identical token ids before and after. Workspace shrinks as a side effect: num_slots at 128 tokens 17280 -> 9088, at 1024 tokens 24448 -> 16256. End-to-end TTFT, vllm-bench.py, 10 prompts, --max-num-seqs 1 --output-len 8 --no-cudagraph --dtype float16, median ms: input-len before after vs previous commit vs pre-retune 128 184.84 126.39 1.46x 1.62x 512 270.43 217.29 1.24x 1.33x 1024 393.62 348.11 1.13x 1.17x The exllama HIP path this fallback regressed against measures 121.56 ms at 128 tokens under the same settings, so the short-prefill gap is now 4% (it was 69%). At 512 tokens the Triton path is 1.65x faster than exllama and at 1024 tokens 2.10x. Decode is unaffected: median TPOT 17.80 -> 17.16 at 128 tokens and 22.97 -> 22.76 at 1024. Signed-off-by: Stephan Toivonen Co-authored-by: Claude --- .../layers/fused_moe/hybrid_w4a16_moe.py | 101 +++++++++++++----- 1 file changed, 76 insertions(+), 25 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/hybrid_w4a16_moe.py b/vllm/model_executor/layers/fused_moe/hybrid_w4a16_moe.py index 2683f74b58f4..075fe61958f8 100644 --- a/vllm/model_executor/layers/fused_moe/hybrid_w4a16_moe.py +++ b/vllm/model_executor/layers/fused_moe/hybrid_w4a16_moe.py @@ -242,7 +242,30 @@ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce: # Alignment for the group_size <= SMALL_GROUP_SIZE_THRESHOLD branch # of _triton_config. Must be lcm of the BLOCK_SIZE_M values that # branch emits (128 for gemm1, 64 for gemm2). - TRITON_BLOCK_SIZE_M_SMALL_GS = 128 + # + # moe_align_block_size pads every touched expert up to a full block, + # and the kernel's only early exit is + # `pid_m*BLOCK_SIZE_M >= num_tokens_post_padded` -- an all-padding + # block still runs the whole contraction. A short prefill on a + # large-E model touches nearly every expert with a few tokens each, + # so the padded row count tracks E*alignment, not the batch: at + # E=128, top_k=8 a 128-token prefill routes 1024 rows and pads to + # 14976 at alignment 128, 14.6x the useful work. + # + # Swept 32 / 64 / 128 on Strix Halo (gfx1151) over 6..2048 tokens x + # 3 routing distributions, gemm1+gemm2. Chose 64: it beats 128 by + # 1.31-1.85x and 32 by 1.2-1.7x. 32 loses because the block count + # is pinned to the touched-expert count either way, so the per-block + # INT4 weight load + dequant does not shrink with the alignment and + # the narrower M-tile cannot amortize it. + TRITON_BLOCK_SIZE_M_SMALL_GS = 64 + + # Alignment once routed rows per expert reach a full block: every + # expert fills its block anyway, so the padding term is gone and the + # wider M-tile wins instead. The measured crossover is at exactly + # that point (2048 tokens at E=128, top_k=8), where the two tiers are + # within 1-4% of each other in both directions. + TRITON_BLOCK_SIZE_M_SMALL_GS_LARGE = 128 # group_size <= THRESHOLD uses the small-gs _triton_config branch; # > THRESHOLD uses the default. Boundary between 32 and 128. @@ -253,12 +276,18 @@ def _select_block_size_m(self, num_tokens: int, topk: int, E: int) -> int: Decode (num_tokens <= MAX_SKINNY_BATCH_SIZE): small block sizes compatible with the wvSplitK_int4 HIP kernel. - Prefill: TRITON_BLOCK_SIZE_M (or TRITON_BLOCK_SIZE_M_SMALL_GS for small - group_size). The rdna_moe_gemm WMMA gemm1 and the Triton gemm2 share this one - alignment. + Prefill: TRITON_BLOCK_SIZE_M, or for small group_size the + padding-aware pair TRITON_BLOCK_SIZE_M_SMALL_GS{,_LARGE} chosen on + routed tokens per expert (see those constants). The rdna_moe_gemm WMMA + gemm1 and the Triton gemm2 share this one alignment. """ if num_tokens > HybridW4A16MoEExperts.MAX_SKINNY_BATCH_SIZE: if self._group_size <= HybridW4A16MoEExperts.SMALL_GROUP_SIZE_THRESHOLD: + large = HybridW4A16MoEExperts.TRITON_BLOCK_SIZE_M_SMALL_GS_LARGE + # Average routed rows per expert -- what the padding cost + # scales with. + if num_tokens * topk / E >= large: + return large return HybridW4A16MoEExperts.TRITON_BLOCK_SIZE_M_SMALL_GS return HybridW4A16MoEExperts.TRITON_BLOCK_SIZE_M if num_tokens > 1: @@ -342,31 +371,38 @@ def _triton_config( sized for the alignment block_size_m used by moe_align_block_size. apply() handles the BLOCK_M < align_block_size_m case by repeat-interleaving expert_ids; for that to be valid we require - BLOCK_SIZE_M divides align_block_size_m. Today we only emit - BLOCK_SIZE_M = align_block_size_m so this is a no-op, but the - infrastructure stays in place to enable future per-gemm BLOCK_M - tuning when paired with per-gemm alignment. + BLOCK_SIZE_M divides align_block_size_m. The small-group-size + branch below relies on that (it emits BLOCK_SIZE_M=64 for gemm2 at + either alignment) and reads align_block_size_m to pick its gemm1 + tile, since the two alignment tiers are bound by different things. """ BLOCK_SIZE_K = self._group_size # = 128 for the Qwen3.5-A3B path assert BLOCK_SIZE_K % 8 == 0 if self._group_size <= HybridW4A16MoEExperts.SMALL_GROUP_SIZE_THRESHOLD: - # gemm2 BM=64 < alignment=128; apply()'s _expert_ids_for - # repeat_interleaves expert_ids to compensate. - # # Swept on Strix Halo (gfx1151), Qwen3-30B-A3B AWQ shapes - # (group_size=32, E=128, alignment=128): 108 tile configs x 10 - # token counts (M_routed 48..24448) x 2 routing distributions. - # Picks below are the best geomean of (us / per-M best), so they - # hold across the batch range and not just at one M. num_stages=4 - # is the common win; the "pipeliner regresses" note in the - # large-group branch below was measured at its much narrower - # BLOCK_SIZE_N=16/64, num_warps=2 tiles. + # (group_size=32, E=128): tile configs x 9-10 token counts + # (M_routed 48..24448) x routing distributions, at both alignments + # _select_block_size_m emits. Picks below are the best geomean of + # (us / per-M best), so they hold across the batch range and not + # just at one M. num_stages=4 is the common win; the "pipeliner + # regresses" note in the large-group branch below was measured at + # its much narrower BLOCK_SIZE_N=16/64, num_warps=2 tiles. + # + # gemm1 needs a different tile per alignment tier, and not just a + # rescaled one: at the low alignment the padding is gone and the + # kernel is bound by the per-block INT4 weight load + dequant + # rather than by the accumulate. + align = align_block_size_m or self.TRITON_BLOCK_SIZE_M_SMALL_GS + small_align = align <= self.TRITON_BLOCK_SIZE_M_SMALL_GS if K < 1024: - # gemm2 (K=768). Chose BLOCK_SIZE_N=128, num_warps=8, - # num_stages=4: geomean 1.043 vs 1.109 for the previous - # BLOCK_SIZE_N=64, num_warps=4, num_stages=1. BLOCK_SIZE_M=128 - # ties on geomean but has a worse tail, so it stays at 64. + # gemm2 (K=768), same tile at both alignments. Chose + # BLOCK_SIZE_N=128, num_warps=8, num_stages=4: geomean 1.043 vs + # 1.109 for the previous BLOCK_SIZE_N=64, num_warps=4, + # num_stages=1, and still the pick at alignment 64 (1.000). + # BLOCK_SIZE_M stays 64 in both tiers; at alignment 128 that + # makes it smaller than the alignment, which apply()'s + # _expert_ids_for handles by repeat_interleaving expert_ids. return dict( BLOCK_SIZE_M=64, BLOCK_SIZE_N=128, @@ -375,9 +411,24 @@ def _triton_config( num_warps=8, num_stages=4, ) - # gemm1 (K=2048). Only num_stages moves: 4 scores 1.040 vs 1.284 - # for 1. BLOCK_SIZE_N=128 is up to 1.77x off the best at small - # batches and num_warps=4 ties 8, so both stay put. + if small_align: + # gemm1 (K=2048) at alignment 64. Chose BLOCK_SIZE_N=32, + # num_warps=2: best at all 27 (tokens x routing) points swept, + # by 1.10x (2048 tokens) to 2.13x (6 tokens) over the + # runner-up. Keeping the alignment-128 tile here would give + # back most of what the smaller alignment buys. + return dict( + BLOCK_SIZE_M=64, + BLOCK_SIZE_N=32, + BLOCK_SIZE_K=BLOCK_SIZE_K, + GROUP_SIZE_M=1, + num_warps=2, + num_stages=4, + ) + # gemm1 (K=2048) at alignment 128, reached only once every expert + # fills a padding block. Only num_stages moves: 4 scores 1.040 vs + # 1.284 for 1. BLOCK_SIZE_N=128 is up to 1.77x off the best at + # small batches and num_warps=4 ties 8, so both stay put. return dict( BLOCK_SIZE_M=128, BLOCK_SIZE_N=64,