perf(expt_data): skip redundant TokenStart/TileStart stores in stage1

[amd-hhashemi]

In _expt_data_compute_stage1, every block previously computed the full
prefix-sum over all experts and wrote the entire TokenStart and TileStart
arrays to global memory. This meant n_expts_tot blocks were all storing the
same values to the same addresses — pure redundant traffic.

This change gates the stores so each non-zero block only writes the
BLOCK-sized chunk that contains its own expert. pid==0 still runs the full
loop and writes all chunks, as it must — it reads TileStart[n_expts_tot-1]
afterwards to compute the sentinel value.

The reduction in store traffic unblocks the loads in stage2 on architectures
where the memory subsystem serializes stores and loads to overlapping cache
lines. On MI450 this gives a 2x end-to-end speedup to _combined_routing.

Changes

• _expt_data_compute_stage1: add done flag and per-chunk store gate
  (pid == 0 or pid < i + BLOCK) so non-zero blocks exit after writing
  their own chunk
• No change to _expt_data_compute_stage2, kernel signatures, or call sites
• Correctness validated across 108 scenarios (varied expert counts, token
  counts, topk, seeds) covering the EQUAL_BLOCK path and all three remainder
  cases of the unrolled loop

[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:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
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 4361 --add-label <label>

[Copilot]

Pull request overview

Optimizes the MoE routing “expert data” stage1 Triton helper by gating global TokenStart/TileStart stores (and early-exiting non-zero programs) to reduce redundant memory traffic and improve _combined_routing performance.

Changes:

• Add a done flag to allow non-zero programs to stop stage1 once their relevant chunk is written.
• Gate TokenStart/TileStart stores so non-zero programs only write the chunk that contains their expert.
• Apply masked stores for the tail chunk in the non-EQUAL_BLOCK path.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 1 out of 1 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (1)

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

• In _combined_routing[(blocks1a + blocks1b,)] / _combined_routing_fused, stage1 is executed by routing program instances too (pids >= blocks1a). blocks1a is set to n_expts_tot in _compute_expt_data_internal (aiter/ops/triton/moe/moe_routing/routing.py:260-262). Those routing pids then immediately run _routing_compute_indx*, which loads TokensStart + expert for arbitrary experts (aiter/ops/triton/_triton_kernels/moe/moe_routing/routing.py:83-85 / 99-101). With the new pid < i + BLOCK + done gating, pid >= n_expts_tot blocks only write the last chunk (or nothing) and can read TokenStart entries they never initialized; there is no global barrier in _combined_routing to ensure pid0 ran first.

To preserve correctness without cross-program synchronization, routing pids should keep the old behavior (store all chunks / no early-exit) and only expert pids (pid < n_expts_tot) should use the per-chunk store + early-exit optimization.

                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

[Copilot]

Pull request overview

Copilot reviewed 1 out of 1 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (1)

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

• The comment above store_pid says “All other blocks can stop once they've written their own chunk”, but in _combined_routing this helper runs for pid >= n_expts_tot routing blocks too. Those blocks are intentionally treated like pid==0 (store_pid becomes 0) and therefore do not early-exit. Updating the comment to mention the routing-block case will prevent future confusion about why pid >= n_expts_tot writes all chunks.

        # 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