From ff43f7c165d9260325c87e3a7d164afd4bd38808 Mon Sep 17 00:00:00 2001 From: Hashem Hashemi Date: Thu, 23 Jul 2026 20:57:20 +0000 Subject: [PATCH 1/2] perf(expt_data): skip redundant TokenStart/TileStart stores in stage1 Each non-zero block now only stores its own BLOCK-sized chunk of the prefix-sum arrays, rather than all chunks. pid==0 still writes the full arrays for the sentinel. Reduces store traffic by ~(n_expts/BLOCK - 1)x, which benefits architectures where store pressure stalls downstream loads. --- .../moe/moe_routing/expt_data.py | 28 ++++++++++++------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/moe/moe_routing/expt_data.py b/aiter/ops/triton/_triton_kernels/moe/moe_routing/expt_data.py index 70fa3b6f887..80136bfe1f8 100644 --- a/aiter/ops/triton/_triton_kernels/moe/moe_routing/expt_data.py +++ b/aiter/ops/triton/_triton_kernels/moe/moe_routing/expt_data.py @@ -33,17 +33,25 @@ def _expt_data_compute_stage1( token_acc = tl.zeros([BLOCK], dtype=TokenStart.dtype.element_ty) tile_acc = tl.zeros([BLOCK], dtype=TileStart.dtype.element_ty) offs_n = tl.arange(0, BLOCK) + # pid==0 must run the full loop: it reads TileStart[n_expts_tot-1] + # below for the sentinel, so it must have written it first. + # All other blocks can stop once they've written their own chunk. + done = False for i in range(0, n_expts_tot, BLOCK): - mask_n = offs_n < n_expts_tot - hist_token = tl.load(Hist + offs_n, mask=mask_n, other=0) - hist_tile = _cdiv_pow2(hist_token, tile_dim_log2) - token_starts = tl.cumsum(hist_token, 0) - hist_token + token_acc - tile_starts = tl.cumsum(hist_tile, 0) - hist_tile + tile_acc - token_acc += tl.sum(hist_token, 0) - tile_acc += tl.sum(hist_tile, 0) - tl.store(TokenStart + offs_n, token_starts) - tl.store(TileStart + offs_n, tile_starts) - offs_n += BLOCK + if not done: + mask_n = offs_n < n_expts_tot + hist_token = tl.load(Hist + offs_n, mask=mask_n, other=0) + hist_tile = _cdiv_pow2(hist_token, tile_dim_log2) + token_starts = tl.cumsum(hist_token, 0) - hist_token + token_acc + tile_starts = tl.cumsum(hist_tile, 0) - hist_tile + tile_acc + token_acc += tl.sum(hist_token, 0) + tile_acc += tl.sum(hist_tile, 0) + if pid == 0 or pid < i + BLOCK: + tl.store(TokenStart + offs_n, token_starts, mask=mask_n) + tl.store(TileStart + offs_n, tile_starts, mask=mask_n) + if pid != 0 and pid < i + BLOCK: + done = True + offs_n += BLOCK if pid == 0: tl.store(TokenStart + n_expts_tot, n_gates) From 2a4b11e1c50ab58d2e941befad4123230d379051 Mon Sep 17 00:00:00 2001 From: Hashem Hashemi Date: Thu, 23 Jul 2026 22:09:09 +0000 Subject: [PATCH 2/2] fix for combined_rounting path --- .../ops/triton/_triton_kernels/moe/moe_routing/expt_data.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/moe/moe_routing/expt_data.py b/aiter/ops/triton/_triton_kernels/moe/moe_routing/expt_data.py index 80136bfe1f8..d4ab258d8eb 100644 --- a/aiter/ops/triton/_triton_kernels/moe/moe_routing/expt_data.py +++ b/aiter/ops/triton/_triton_kernels/moe/moe_routing/expt_data.py @@ -36,6 +36,7 @@ def _expt_data_compute_stage1( # pid==0 must run the full loop: it reads TileStart[n_expts_tot-1] # below for the sentinel, so it must have written it first. # All other blocks can stop once they've written their own chunk. + store_pid = pid if pid < n_expts_tot else 0 done = False for i in range(0, n_expts_tot, BLOCK): if not done: @@ -46,13 +47,12 @@ def _expt_data_compute_stage1( tile_starts = tl.cumsum(hist_tile, 0) - hist_tile + tile_acc token_acc += tl.sum(hist_token, 0) tile_acc += tl.sum(hist_tile, 0) - if pid == 0 or pid < i + BLOCK: + if store_pid == 0 or store_pid < i + BLOCK: tl.store(TokenStart + offs_n, token_starts, mask=mask_n) tl.store(TileStart + offs_n, tile_starts, mask=mask_n) - if pid != 0 and pid < i + BLOCK: + if store_pid != 0 and store_pid < i + BLOCK: done = True offs_n += BLOCK - if pid == 0: tl.store(TokenStart + n_expts_tot, n_gates)