Skip to content

[Triton/Gluon] Fix MoE routing kernel compile failure #5558

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Fangzhou-Ai merged 5 commits into from
Sep 16, 2026
Merged

[Triton/Gluon] Fix MoE routing kernel compile failure #5558

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
f918eb1
fix compiler pass error
micah-wil Sep 15, 2026
82f13d4
clean up
micah-wil Sep 15, 2026
d78bdfe
Add compile regression test for _combined_routing[_fused]
micah-wil Sep 15, 2026
7b736cb
Revert "Add compile regression test for _combined_routing[_fused]"
micah-wil Sep 15, 2026
96fec56
Merge branch 'main' into micah/routing-early-return-legalize
micah-wil Sep 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 12 additions & 15 deletions aiter/ops/triton/_triton_kernels/moe/moe_routing/expt_data.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,24 +71,21 @@ def _expt_data_compute_stage2(
expt_id = pid

n_tokens = tl.load(Hist + expt_id)
if n_tokens == 0:
return
BLOCK: tl.constexpr = 8
n_blocks = _cdiv_pow2(n_tokens, tile_dim_log2)
TileInfo += tile_start
if n_tokens != 0:
BLOCK: tl.constexpr = 8
n_blocks = _cdiv_pow2(n_tokens, tile_dim_log2)
tile_info = TileInfo + tile_start

n_blocks = _cdiv_pow2(n_tokens, tile_dim_log2)
block_offs = tl.arange(0, BLOCK)
for i in range(0, n_blocks, BLOCK):
data = (block_offs << 16) + expt_id
tl.store(TileInfo + block_offs, data, mask=block_offs < n_blocks)
block_offs += BLOCK
block_offs = tl.arange(0, BLOCK)
for i in range(0, n_blocks, BLOCK):
data = (block_offs << 16) + expt_id
tl.store(tile_info + block_offs, data, mask=block_offs < n_blocks)
block_offs += BLOCK


@triton.jit
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 += tile_start
tl.store(TileInfo, expt_id)
if n_tokens != 0:
tile_info = TileInfo + tile_start
tl.store(tile_info, expt_id)
Loading