From 5cb5c9ff57e0a53c9f152d7a41a0d0cc2546478d Mon Sep 17 00:00:00 2001 From: shanmugamr1992 Date: Thu, 6 Aug 2026 16:40:03 -0700 Subject: [PATCH] Retune MoE token-count kernel block size for decode shapes The persistent token-count kernel used BLOCK_SIZE=1024, which makes each thread issue eight atomics back to back rather than spreading them across CTAs. At decode shape the kernel is flat in both grid size and routed-pair count, the signature of per-thread serialization rather than launch geometry or atomic volume. Sizing the block to 128 cuts the kernel from 8.08 us to 2.11 us; larger blocks only pay off once there are enough pairs to fill the machine, where the curve is flat. Measured on GB200, Qwen3-30B-A3B EP4/TP1, BS256: +2.29% decode throughput. Bit-exact. Kill switch: MCORE_COUNT_TOKENS_BLOCK. Signed-off-by: shanmugamr1992 --- megatron/core/inference/moe/permute.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/megatron/core/inference/moe/permute.py b/megatron/core/inference/moe/permute.py index f65ac5c4200..979a194a09e 100644 --- a/megatron/core/inference/moe/permute.py +++ b/megatron/core/inference/moe/permute.py @@ -8,6 +8,7 @@ - Unpermute expert outputs back to original token order """ +import os from typing import Optional from unittest.mock import MagicMock @@ -45,6 +46,15 @@ def _ceil_div(a, b): return (a + b - 1) // b +# Pairs per program for the persistent token-count kernel. Each thread issues one +# atomic per pair it owns, so an oversized block serialises those atomics inside the +# thread instead of spreading them across CTAs: at decode shape (2048 routed pairs) +# the kernel costs 8.08 us at 1024 and 2.11 us at 128. Larger blocks only pay off once +# there are enough pairs to fill the machine anyway, where the curve is flat (13.3 us +# at 128 vs 13.2 us at 512 for 131072 pairs). Measured on GB200, Qwen3-30B-A3B EP4. +_COUNT_TOKENS_BLOCK = int(os.environ.get("MCORE_COUNT_TOKENS_BLOCK", "128")) + + @triton.jit def _count_local_tokens_kernel( routing_map_ptr, # [max_tokens, topk] flattened expert assignments @@ -143,7 +153,7 @@ def compute_local_tokens_per_expert( local_expert_start, num_local_experts, num_sms, - BLOCK_SIZE=BLOCK, + BLOCK_SIZE=_COUNT_TOKENS_BLOCK, ) else: _count_local_tokens_kernel[(_ceil_div(max_pairs, BLOCK),)](