[Triton/Gluon] Fix MoE routing kernel compile failure

[micah-wil]

#5038 reworked the MoE routing helpers so that _expt_data_compute_stage1 computes the per-expert tile offset and returns it as a small tensor (built from a tl.gather / tl.zeros in an if/else), which is then passed into _expt_data_compute_stage2 and _expt_data_compute_stage2_fused. Those two helpers start with an early return:

    n_tokens = tl.load(Hist + expt_id)
    if n_tokens == 0:
        return
    TileInfo += tile_start   # uses the passed-in tensor after the return

This pattern makes the routing kernels fail to compile. Both _combined_routing (num_tokens > 16) and _combined_routing_fused (num_tokens <= 16) hit it, so any MXFP4 MoE run that goes through this routing path breaks. Reproduced on gfx942 and gfx950, with Triton 3.7.1 (the default in vLLM).
The compile aborts inside Triton's ConvertTritonToTritonGPU pass:

INFO 09-14 23:16:38 [mxfp4.py:1997] Using AiterW4A16ExpertsMonolithic
/usr/local/lib/python3.12/dist-packages/aiter/ops/triton/_triton_kernels/moe/moe_routing/expt_data.py:89:0: error: failed to legalize unresolved materialization from () to ('i32') that remained live after conversion
/usr/local/lib/python3.12/dist-packages/aiter/ops/triton/_triton_kernels/moe/moe_routing/expt_data.py:90:30: note: see existing live user here: %5 = "tt.addptr"(%2, %3) : (!tt.ptr<i32>, i32) -> !tt.ptr<i32>
    n_tokens = tl.load(Hist + expt_id)
    ...
RuntimeError: PassManager::run failed

(full error log: https://buildkite.com/vllm/amd-ci/builds/12929/list?jid=01a0a157-9fe0-4e1e-b3c7-46d8d32ce602&tab=output#L1425)

This appears to fundamentally be a Triton compiler bug.

Fix: replace the early return with a guarded if n_tokens != 0: block, and add into a fresh local (tile_info = TileInfo + tile_start) instead of reassigning the TileInfo pointer argument in place. This keeps the value from having to live across the early-return control flow, which is what the conversion pass couldn't handle. The computation and results are unchanged; the previous design that avoided this (loading the offset from a TileStart buffer) confirms the behavior is equivalent.

Repro (vLLM):
pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py::test_rocm_mxfp4_moe_oracle[16-256-256-8-4-AITER_TRITON_MXFP4_BF16]

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
multigpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 5558 --add-label <label>

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to opt this PR out.

[vgokhale]

LGTM - perf looks neutral. ASM is identical

[Copilot]

🔵 Needs a closer look

Required architecture-specific regression coverage is missing.

Pull request overview

Fixes Triton compilation failures in the MoE routing stage-2 helpers.

Changes:

• Replaces early returns with guarded blocks.
• Uses local tile-info pointers to avoid compiler issues.

File summaries

File	Summary
aiter/ops/triton/_triton_kernels/moe/moe_routing/expt_data.py	Applies the compiler-safe fix to both routing helpers. Moderate issue: add gfx942 and gfx950 regression coverage for empty-expert branches.

Review details

Suppressed comments (1)

aiter/ops/triton/_triton_kernels/moe/moe_routing/expt_data.py:77

• [verified] This fixes the shared stage-2 helpers for gfx942/gfx950, but the in-tree routing tests skip gfx942 (op_tests/triton_tests/moe/test_moe_routing.py:161-162, 446-447, 515-516) and this PR adds no regression test for that architecture or the MXFP4 route. The claimed gfx942 compile failure can therefore regress unnoticed even though gfx950 coverage passes. Author must add or extend a collected routing/MXFP4 test that exercises both empty-expert branches on gfx942 as well as gfx950.

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

• Files reviewed: 1/1 changed files
• Comments generated: 0
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.