From 4a1ca458088cb186c94dae0c825920678b113b77 Mon Sep 17 00:00:00 2001 From: Lukasz Burzawa Date: Wed, 26 Aug 2026 21:07:59 +0000 Subject: [PATCH 1/3] optimize stores --- .../moe/moe_routing/expt_data.py | 86 ++++++------------- .../moe/moe_routing/routing.py | 15 ++-- 2 files changed, 31 insertions(+), 70 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 509f86fc2f0..fc2bc621106 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 @@ -21,36 +21,32 @@ def _expt_data_compute_stage1( BLOCK: tl.constexpr, EQUAL_BLOCK: tl.constexpr, ): + offs_n = tl.arange(0, BLOCK) if EQUAL_BLOCK: - offs_n = tl.arange(0, BLOCK) hist_token = tl.load(Hist + offs_n) - hist_tile = _cdiv_pow2(hist_token, tile_dim_log2) - token_starts = tl.cumsum(hist_token, 0) - hist_token - tile_starts = tl.cumsum(hist_tile, 0) - hist_tile - tl.store(TokenStart + offs_n, token_starts) - tl.store(TileStart + offs_n, tile_starts) else: - 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) - 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) + 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 + tile_starts = tl.cumsum(hist_tile, 0) - hist_tile + + if pid < n_expts_tot: + expt_id = tl.zeros([1], tl.int32) + pid + tile_start = tl.gather(tile_starts, expt_id, 0) + token_start = tl.gather(token_starts, expt_id, 0) + tl.store(TileStart + expt_id, tile_start) + tl.store(TokenStart + expt_id, token_start) + else: + tile_start = tl.zeros([1], tl.int32) + if EQUAL_BLOCK: tl.store(TokenStart + offs_n, token_starts) - tl.store(TileStart + offs_n, tile_starts) - offs_n += BLOCK + else: + tl.store(TokenStart + offs_n, token_starts, mask=mask_n) if pid == 0: tl.store(TokenStart + n_expts_tot, n_gates) - - hist_tok_last = tl.load(Hist + n_expts_tot - 1) - hist_tile_last = _cdiv_pow2(hist_tok_last, tile_dim_log2) - tile_off_last = tl.load(TileStart + n_expts_tot - 1) + hist_tile_last + tile_off_last = tl.sum(hist_tile, 0) tl.store(TileStart + n_expts_tot, tile_off_last) MEMSET_BLOCK: tl.constexpr = 16 @@ -60,10 +56,12 @@ def _expt_data_compute_stage1( MDTileInfo + block_offs, 0xFFFFFFFF, mask=block_offs < max_num_tiles ) + return tile_start + @triton.jit def _expt_data_compute_stage2( - pid, Hist, TileStart, TileInfo, tile_dim_log2: tl.constexpr + pid, Hist, tile_start, TileInfo, tile_dim_log2: tl.constexpr ): expt_id = pid @@ -73,7 +71,7 @@ def _expt_data_compute_stage2( return BLOCK: tl.constexpr = 8 n_blocks = _cdiv_pow2(n_tokens, tile_dim_log2) - TileInfo += tl.load(TileStart + expt_id) + TileInfo += tile_start n_blocks = _cdiv_pow2(n_tokens, tile_dim_log2) block_offs = tl.arange(0, BLOCK) @@ -84,43 +82,9 @@ def _expt_data_compute_stage2( @triton.jit -def _expt_data_compute_stage2_fused(expt_id, Hist, TileStart, TileInfo): +def _expt_data_compute_stage2_fused(expt_id, Hist, tile_start, TileInfo): n_tokens = tl.load(Hist + expt_id) if n_tokens == 0: return - TileInfo += tl.load(TileStart + expt_id) + TileInfo += tile_start tl.store(TileInfo, expt_id) - - -@triton.jit -def _expt_data_only_kernel( - Hist, - n_expts_tot, - TokenStart, - TileStart, - MDTileInfo, - max_num_tiles, - n_gates, - tile_dim_log2: tl.constexpr, - BLOCK: tl.constexpr, - EQUAL_BLOCK: tl.constexpr, -): - """Standalone stage1+stage2 launch — builds ExptData from a precomputed - histogram with no memset. Grid: (n_expts_tot,).""" - pid = tl.program_id(0) - - _expt_data_compute_stage1( - pid, - Hist, - n_expts_tot, - TokenStart, - TileStart, - MDTileInfo, - max_num_tiles, - n_gates, - tile_dim_log2, - BLOCK, - EQUAL_BLOCK, - ) - - _expt_data_compute_stage2(pid, Hist, TileStart, MDTileInfo, tile_dim_log2) diff --git a/aiter/ops/triton/_triton_kernels/moe/moe_routing/routing.py b/aiter/ops/triton/_triton_kernels/moe/moe_routing/routing.py index 39e36d32cb1..553ddd14387 100644 --- a/aiter/ops/triton/_triton_kernels/moe/moe_routing/routing.py +++ b/aiter/ops/triton/_triton_kernels/moe/moe_routing/routing.py @@ -211,7 +211,7 @@ def _combined_routing( pid = tl.program_id(0) - _expt_data_compute_stage1( + tile_start = _expt_data_compute_stage1( pid, ExpertHist, n_expts_tot, @@ -226,7 +226,9 @@ def _combined_routing( ) if pid < blocks1a: - _expt_data_compute_stage2(pid, ExpertHist, TileStart, MDTileInfo, tile_dim_log2) + _expt_data_compute_stage2( + pid, ExpertHist, tile_start, MDTileInfo, tile_dim_log2 + ) else: pid -= blocks1a _routing_compute_indx( @@ -294,12 +296,7 @@ def _combined_routing_fused( tl.debug_barrier() - if pid != 0 and pid < blocks1a: - n_tokens = tl.load(ExpertHist + pid) - if n_tokens == 0: - return - - _expt_data_compute_stage1( + tile_start = _expt_data_compute_stage1( pid, ExpertHist, N_EXPTS_TOT, @@ -314,7 +311,7 @@ def _combined_routing_fused( ) if pid < blocks1a: - _expt_data_compute_stage2_fused(pid, ExpertHist, TileStart, MDTileInfo) + _expt_data_compute_stage2_fused(pid, ExpertHist, tile_start, MDTileInfo) else: _routing_compute_indx_fused( GatherIndx, From 1bedf998fdd44a8af162540a9996868a96af2c9f Mon Sep 17 00:00:00 2001 From: Lukasz Burzawa Date: Thu, 27 Aug 2026 01:43:46 +0000 Subject: [PATCH 2/3] optimize --- .../_triton_kernels/moe/moe_routing/expt_data.py | 15 ++++++++++----- .../_triton_kernels/moe/moe_routing/routing.py | 5 +++++ 2 files changed, 15 insertions(+), 5 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 fc2bc621106..119aab6fa97 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 @@ -27,24 +27,29 @@ def _expt_data_compute_stage1( else: 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 - tile_starts = tl.cumsum(hist_tile, 0) - hist_tile if pid < n_expts_tot: + hist_tile = _cdiv_pow2(hist_token, tile_dim_log2) + tile_starts = tl.cumsum(hist_tile, 0) - hist_tile expt_id = tl.zeros([1], tl.int32) + pid tile_start = tl.gather(tile_starts, expt_id, 0) - token_start = tl.gather(token_starts, expt_id, 0) tl.store(TileStart + expt_id, tile_start) - tl.store(TokenStart + expt_id, token_start) else: tile_start = tl.zeros([1], tl.int32) + token_starts = tl.cumsum(hist_token, 0) - hist_token if EQUAL_BLOCK: tl.store(TokenStart + offs_n, token_starts) else: tl.store(TokenStart + offs_n, token_starts, mask=mask_n) if pid == 0: + hist_tile = _cdiv_pow2(hist_token, tile_dim_log2) + tile_starts = tl.cumsum(hist_tile, 0) - hist_tile + if EQUAL_BLOCK: + tl.store(TileStart + offs_n, tile_starts) + else: + tl.store(TileStart + offs_n, tile_starts, mask=mask_n) + tl.store(TokenStart + n_expts_tot, n_gates) tile_off_last = tl.sum(hist_tile, 0) tl.store(TileStart + n_expts_tot, tile_off_last) diff --git a/aiter/ops/triton/_triton_kernels/moe/moe_routing/routing.py b/aiter/ops/triton/_triton_kernels/moe/moe_routing/routing.py index 553ddd14387..27f7acc7865 100644 --- a/aiter/ops/triton/_triton_kernels/moe/moe_routing/routing.py +++ b/aiter/ops/triton/_triton_kernels/moe/moe_routing/routing.py @@ -296,6 +296,11 @@ def _combined_routing_fused( tl.debug_barrier() + if pid != 0 and pid < blocks1a: + n_tokens = tl.load(ExpertHist + pid) + if n_tokens == 0: + return + tile_start = _expt_data_compute_stage1( pid, ExpertHist, From 97057b61bf41ddd419bd0d69c80d153502ba6cd5 Mon Sep 17 00:00:00 2001 From: Lukasz Burzawa Date: Thu, 27 Aug 2026 01:57:35 +0000 Subject: [PATCH 3/3] remove extra store --- aiter/ops/triton/_triton_kernels/moe/moe_routing/expt_data.py | 1 - 1 file changed, 1 deletion(-) 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 119aab6fa97..010fe1ffdf8 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,7 +33,6 @@ def _expt_data_compute_stage1( tile_starts = tl.cumsum(hist_tile, 0) - hist_tile expt_id = tl.zeros([1], tl.int32) + pid tile_start = tl.gather(tile_starts, expt_id, 0) - tl.store(TileStart + expt_id, tile_start) else: tile_start = tl.zeros([1], tl.int32) token_starts = tl.cumsum(hist_token, 0) - hist_token