[Fix][FlyDSL] Handle remainder workgroups in MoE XCD swizzle

[Fangzhou-Ai]

Summary

• Correct the FlyDSL MoE XCD swizzle when the launch grid is not divisible by the number of XCDs.
• Distribute remainder workgroups across the leading XCD ranges so the mapping remains a complete permutation.
• Add focused coverage for all grid sizes 1-256 and the DSV4-like stage-2 grids 28x399 and 28x400.

How the mapping works

For a launch with W workgroups and X = 8 XCDs, define:

• q = W // X and r = W % X;
• x = linear_id % X, the XCD lane used by the swizzle;
• j = linear_id // X, that workgroup's local ordinal on XCD x.

Linear IDs are distributed round-robin, so XCDs 0..r-1 receive q + 1 workgroups and the remaining XCDs receive q. The previous mapping was:

mapped_id = x * q + j

That formula implicitly reserves only q logical IDs for every XCD. When r != 0, the extra workgroup on each leading XCD overlaps the next XCD's range, producing duplicate logical tiles and leaving the final r tiles unassigned.

The corrected mapping is:

mapped_id = x * q + min(x, r) + j

The term x * q + min(x, r) is the prefix sum of the workgroup counts assigned to XCDs before x. It gives each XCD a non-overlapping contiguous range: q + 1 IDs when x < r, otherwise q IDs.

For a compact example, let W = 10 and X = 4. Then q = 2 and r = 2:

XCD	Physical linear IDs	Previous mapped IDs	Corrected mapped IDs
0	0, 4, 8	0, 1, 2	0, 1, 2
1	1, 5, 9	2, 3, 4	3, 4, 5
2	2, 6	4, 5	6, 7
3	3, 7	6, 7	8, 9

The previous mapping duplicates IDs 2 and 4 and never assigns 8 or 9. The corrected mapping is a permutation of 0..9.

The affected kernel A/B below uses X = 8 and has W = 19 * 28 = 532, so q = 66 and r = 4. The previous XCD ranges overlap at 66, 132, 198, 264 and stop at 527, omitting 528..531. The corrected ranges start at 0, 67, 134, 201, 268, 334, 400, 466, have sizes 67, 67, 67, 67, 66, 66, 66, 66, and cover 0..531 exactly once.

This remapped ID then enters the existing WGM tiling logic; the fix changes tile assignment, not GEMM arithmetic. It is bypassed when XCD swizzling is disabled. For divisible grids, r = 0, so the corrected formula reduces exactly to the previous one. The correction is therefore relevant only to swizzled grids with a remainder.

Scope

This is a standalone fix targeting main. It has no prerequisites and nothing depends on it.

Earlier revisions of this PR made it the base of a review chain with #4269 (FHMoE) and #4314 (fused-MoE tuning). That coupling has been removed: every route selected by #4314 uses XCD0, and the initial FHMoE enablement in #4269 is XCD-free, so neither PR needs this mapping correction to be reviewed or landed. #4269 and #4314 now form their own two-layer stack on main, independent of this PR, and this PR can be reviewed and merged on its own in either order.

Empirical kernel A/B

A controlled same-host A/B used identical deterministic A4W4 inputs, the DSV4 model dimension of 7168, 28 N tiles, 19 sorted M blocks, and an affected 19x28 = 532 workgroup XCD4 grid (532 % 8 = 4). The expert count was reduced to fit alongside the running server; the XCD and WGM mapping is unchanged.

Output versus the torch reference:

Mapping	Relative L2	Max absolute error
Current release mapping	9.345%	0.8671875
Corrected mapping	0.3534%	0.015625

The corrected and release outputs differ by 9.235% relative L2 and are not bitwise equal. This demonstrates a kernel-output correction on a real routed workload, although robust end-to-end release evaluations may not frequently exercise affected grids or may tolerate their localized error.

Performance used 50 warmups followed by 10 repeats of 200 HIP-event-timed iterations:

Mapping	Mean	Median
Current release mapping	33.945 us	33.949 us
Corrected mapping	34.298 us	34.387 us

The correction does not provide a performance gain in this replay: it is approximately 1.0% slower by mean and 1.3% slower by median. It is proposed as a correctness prerequisite, not as a performance optimization by itself.

The A/B used the same patch, now at f8fba0211 after a rebase onto main at 4a1cc773f. The rebase changed only the parent and the import position of the new moe_common helper; intervening upstream changes do not touch the swizzle arithmetic.

Unit validation

python -m pytest \
  op_tests/flydsl_tests/test_moe_xcd_swizzle.py -q

Result on 141c6dcda, whose swizzle logic is unchanged at the current head f8fba0211: 2 passed, 258 subtests passed.

Static checks passed:

uvx ruff check aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py \
  aiter/ops/flydsl/moe_common.py \
  op_tests/flydsl_tests/test_moe_xcd_swizzle.py
python -m py_compile \
  aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py \
  aiter/ops/flydsl/moe_common.py \
  op_tests/flydsl_tests/test_moe_xcd_swizzle.py
git diff --check origin/main...HEAD

Duplicate-work check

Open-PR searches found no separate equivalent fix.

AI assistance

OpenAI Codex was used to audit the mapping, build and run the kernel A/B, run tests, organize the stacked history, and draft this description. The human submitter reviewed the change and remains responsible for understanding and defending it end-to-end.

[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 4315 --add-label <label>