fix(flydsl): handle remainder workgroups in MoE XCD swizzle

Distribute remainder workgroups across XCD ranges so non-divisible grids remain a complete permutation. The previous floor-division mapping duplicated workgroups and omitted the tail of affected grids.

Co-authored-by: OpenAI Codex <codex@openai.com>
Signed-off-by: Fangzhou Ai <fangzhouai@gmail.com>

Author: Fangzhou-Ai

aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py

@@ -64,6 +64,7 @@ def supports_bf16_global_atomics(arch: str) -> bool:
     tile_chunk_coord_i32,
     unpack_b_mxfp4_bf16,
 )
+from aiter.ops.flydsl.moe_common import xcd_swizzle_workgroup_id
 
 
 @contextmanager
@@ -563,8 +564,15 @@ def ptr_buffer_resource(ptr, num_records_bytes):
                 num_wgs = gx * gy
 
                 c_xcds = arith.constant(NUM_XCDS_S1, index=True)
-                wgs_per_xcd = num_wgs // c_xcds
-                wgid = (linear_id % c_xcds) * wgs_per_xcd + (linear_id // c_xcds)
+                wgid = xcd_swizzle_workgroup_id(
+                    linear_id,
+                    num_wgs,
+                    c_xcds,
+                    divide=lambda lhs, rhs: lhs // rhs,
+                    minimum=lambda lhs, rhs: arith.select(
+                        arith.cmpi(CmpIPredicate.ult, lhs, rhs), lhs, rhs
+                    ),
+                )
 
                 WGM_S1 = xcd_swizzle
                 c_wgm = arith.constant(WGM_S1, index=True)
@@ -3366,8 +3374,15 @@ def check_c_k_valid_gate(base_k):
                 num_wgs = gx * gy
 
                 c_xcds = arith.constant(NUM_XCDS_S, index=True)
-                wgs_per_xcd = num_wgs // c_xcds
-                wgid = (linear_id % c_xcds) * wgs_per_xcd + (linear_id // c_xcds)
+                wgid = xcd_swizzle_workgroup_id(
+                    linear_id,
+                    num_wgs,
+                    c_xcds,
+                    divide=lambda lhs, rhs: lhs // rhs,
+                    minimum=lambda lhs, rhs: arith.select(
+                        arith.cmpi(CmpIPredicate.ult, lhs, rhs), lhs, rhs
+                    ),
+                )
 
                 WGM_S = xcd_swizzle
                 c_wgm = arith.constant(WGM_S, index=True)

aiter/ops/flydsl/moe_common.py

@@ -1,11 +1,20 @@
 # SPDX-License-Identifier: MIT
 # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
 
-"""Common types shared across MoE FlyDSL kernel modules."""
+"""Common helpers and types shared across MoE FlyDSL kernel modules."""
 
 from enum import Enum
 
 
+def xcd_swizzle_workgroup_id(linear_id, num_workgroups, num_xcds, divide, minimum):
+    """Map round-robin workgroup IDs into balanced contiguous XCD ranges."""
+    workgroups_per_xcd = divide(num_workgroups, num_xcds)
+    remainder = num_workgroups % num_xcds
+    xcd_id = linear_id % num_xcds
+    local_id = divide(linear_id, num_xcds)
+    return xcd_id * workgroups_per_xcd + minimum(xcd_id, remainder) + local_id
+
+
 class GateMode(str, Enum):
     """Gate/Up computation strategy for stage1 GEMM.
 

op_tests/flydsl_tests/test_moe_xcd_swizzle.py

@@ -0,0 +1,47 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+
+import importlib.util
+import operator
+import unittest
+from pathlib import Path
+
+_MOE_COMMON_PATH = (
+    Path(__file__).resolve().parents[2] / "aiter/ops/flydsl/moe_common.py"
+)
+_SPEC = importlib.util.spec_from_file_location("moe_common", _MOE_COMMON_PATH)
+_MOE_COMMON = importlib.util.module_from_spec(_SPEC)
+_SPEC.loader.exec_module(_MOE_COMMON)
+xcd_swizzle_workgroup_id = _MOE_COMMON.xcd_swizzle_workgroup_id
+
+
+class TestMoeXcdSwizzle(unittest.TestCase):
+    def assert_permutation(self, num_workgroups):
+        mapped_ids = [
+            xcd_swizzle_workgroup_id(
+                linear_id,
+                num_workgroups,
+                8,
+                divide=operator.floordiv,
+                minimum=min,
+            )
+            for linear_id in range(num_workgroups)
+        ]
+        self.assertTrue(
+            all(0 <= mapped_id < num_workgroups for mapped_id in mapped_ids)
+        )
+        self.assertEqual(len(set(mapped_ids)), num_workgroups)
+
+    def test_divisible_and_remainder_grids_are_permutations(self):
+        for num_workgroups in range(1, 257):
+            with self.subTest(num_workgroups=num_workgroups):
+                self.assert_permutation(num_workgroups)
+
+    def test_dsv4_stage2_grids_are_permutations(self):
+        for num_workgroups in (28 * 400, 28 * 399):
+            with self.subTest(num_workgroups=num_workgroups):
+                self.assert_permutation(num_workgroups)
+
+
+if __name__ == "__main__":
+    unittest.main()