[MoE][Qwen3.8] Replace gfx942 prefill inverse routing with Opus aux sort

aiter/fused_moe.py

@@ -79,6 +79,8 @@ def is_flydsl_available():
 #     tuned-CSV rows routed to the port (exactly the codegen'd shapes).
 #   opus / ck -> never use adaptive (legacy; ck still needs AITER_USE_CK_MOE_SORTING)
 _MOE_SORT_BACKEND = os.environ.get("AITER_MOE_SORT_BACKEND", "auto").lower()
+# Request Opus auxiliary maps without entering the shape-specific adaptive sort.
+AUX_SORT_OPUS = "opus"
 _ACT_TYPE_DISABLED_KEY = "__ignore__"
 _SWIGLU_MXFP4_BF16_BOUND = int(os.environ.get("GPTOSS_SWIGLU_MXFP4_BF16_BOUND", "256"))
 _MOE_A8W4_BYPASS_QUANT = os.environ.get("AITER_MOE_A8W4_BYPASS_QUANT", "0") == "1"
@@ -197,7 +199,11 @@ def _moe_sorting_impl(
     device = topk_ids.device
     M, topk = topk_ids.shape
 
-    if output_aux and _MOE_SORT_BACKEND not in ("opus", "ck"):
+    if (
+        output_aux
+        and output_aux != AUX_SORT_OPUS
+        and _MOE_SORT_BACKEND not in ("opus", "ck")
+    ):
         # adaptive (fused) sort emits the a4w4 extras (m_indices + reverse_sorted)
         # plus the atomic zero-init; opus single-pass aux is the env-gated fallback.
         return _adaptive_moe_sort(
@@ -239,7 +245,10 @@ def _moe_sorting_impl(
     aux_reverse_sorted = None
     if output_aux:
         use_opus = True
-        dispatch_policy = 1
+        # Opus aux callers keep their policy, including auto/multi-phase prefill.
+        # Boolean aux callers retain the legacy single-kernel path.
+        if output_aux != AUX_SORT_OPUS:
+            dispatch_policy = 1
         aux_m_indices = torch.empty(
             max_num_tokens_padded, dtype=dtypes.i32, device=device
         )

aiter/ops/flydsl/fused_moe_gfx942.py

@@ -11,13 +11,12 @@
 
 import aiter
 from aiter import ActivationType, QuantType
-from aiter.fused_moe import moe_sorting
+from aiter.fused_moe import AUX_SORT_OPUS, moe_sorting
 from aiter.fused_moe_registry import FusedMoeRequest
 from aiter.jit.utils.chip_info import get_gfx
 from aiter.ops.flydsl.kernels.moe_gemm_2stage_gfx942 import (
     flydsl_absmax,
     flydsl_quant_per_tensor,
-    invert_sorted_ids,
     sorted_sum,
 )
 from aiter.ops.flydsl.kernels.moe_gemm_2stage_gfx942.common import (
@@ -210,8 +209,7 @@ def unsupported_reason(self, problem: "_Problem") -> str | None:
         else:
             if problem.inter_dim not in (192, 256, 320, 384, 512, 640):
                 return (
-                    "8x1 paths require inter_dim in "
-                    "{192, 256, 320, 384, 512, 640}"
+                    "8x1 paths require inter_dim in " "{192, 256, 320, 384, 512, 640}"
                 )
             expected_block_m = 256 if self.down_path == "8x1" else 64
             if self.BLOCK_M != expected_block_m or gateup_block_m != 64:
@@ -688,9 +686,7 @@ def compile_kernel(*, stage, alg, block_m, block_n, **kwargs):
             )
 
             metadata_capacity = (
-                batch * topk + experts * config.BLOCK_M - topk
-                + config.BLOCK_M
-                - 1
+                batch * topk + experts * config.BLOCK_M - topk + config.BLOCK_M - 1
             ) // config.BLOCK_M
             full_capacity, tail_capacity = task_capacities(
                 metadata_capacity,
@@ -708,16 +704,6 @@ def compile_kernel(*, stage, alg, block_m, block_n, **kwargs):
             )
         compile_launcher(down, *down_args, 0)
 
-        invert = invert_sorted_ids(topk)
-        compile_launcher(
-            invert._flydsl_executable,
-            _ptr(int32),
-            _ptr(int32),
-            _ptr(int32),
-            fx.Uint32(batch * topk),
-            fx.Uint32(batch),
-            0,
-        )
         reduction = sorted_sum(topk, model_dim, config.down_output_padding_bytes)
         compile_launcher(
             reduction._flydsl_executable,
@@ -843,7 +829,16 @@ def _run_prefill(
     situ_linear_beta: float,
 ):
     gateup_block_m = config.GATEUP_BLOCK_M or config.BLOCK_M
-    sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, cur_out = moe_sorting(
+    # Opus also returns each route's sorted row; merged repeated expert IDs get -1.
+    (
+        sorted_ids,
+        sorted_weights,
+        sorted_expert_ids,
+        num_valid_ids,
+        cur_out,
+        _,
+        loc_ids,
+    ) = moe_sorting(
         topk_ids,
         topk_weight,
         problem.experts,
@@ -853,6 +848,7 @@ def _run_prefill(
         None,
         None,
         0,
+        output_aux=AUX_SORT_OPUS,
     )
     weight_dtype_str = "bf16" if w1.dtype == torch.bfloat16 else "fp8"
     act_quant_type_str = problem.quant_type
@@ -931,16 +927,11 @@ def _run_prefill(
         if output_padding_bytes is not None
         else 0
     )
-    # Opus may coalesce duplicate expert IDs in zero-weight HIP padding rows.
-    # Their omitted route slots are not written by invert_sorted_ids. Map those
-    # slots to a dedicated zero row instead of dereferencing uninitialized IDs.
-    zero_route_row = sorted_expert_ids.shape[0] * config.BLOCK_M
     gemm2_out = torch.empty(
-        [zero_route_row + 1, output_row_size],
+        [sorted_expert_ids.shape[0] * config.BLOCK_M, output_row_size],
         dtype=hidden_states.dtype,
         device=hidden_states.device,
     )
-    gemm2_out[zero_route_row].zero_()
     down_tile_n = {
         "default": 128,
         "1x4_64x256": 256,
@@ -1000,19 +991,6 @@ def _run_prefill(
     else:
         _launch(down_kernel, *down_args)
 
-    loc_ids = torch.full(
-        [problem.batch, problem.topk],
-        zero_route_row,
-        dtype=torch.int32,
-        device=hidden_states.device,
-    )
-    invert_sorted_ids(problem.topk)(
-        sorted_ids,
-        loc_ids,
-        num_valid_ids,
-        sorted_ids.shape[0],
-        problem.batch,
-    )
     sorted_sum(problem.topk, problem.model_dim, output_padding_bytes)(
         loc_ids, gemm2_out, cur_out, problem.batch
     )

aiter/ops/flydsl/kernels/moe_gemm_2stage_gfx942/__init__.py

@@ -12,7 +12,6 @@
 from .gemm2 import compile_moe_gemm2
 from .moe_reduce import (
     compile_moe_reduction,
-    invert_sorted_ids,
     sorted_sum,
 )
 from .quant import flydsl_absmax, flydsl_quant_per_tensor
@@ -24,6 +23,5 @@
     "compile_moe_reduction",
     "flydsl_absmax",
     "flydsl_quant_per_tensor",
-    "invert_sorted_ids",
     "sorted_sum",
 ]

aiter/ops/flydsl/kernels/moe_gemm_2stage_gfx942/moe_reduce.py

@@ -35,6 +35,8 @@ def sorted_sum_kernel(loc_ids: fx.Pointer, A: fx.Pointer, B: fx.Pointer):
         # preload all TOPK locations
         loc_ids += batch * TOPK
         token_locs = [loc_ids[topk] for topk in fx.range_constexpr(TOPK)]
+        # Sorting gives -1 to routes merged as repeated expert IDs of a token.
+        valid = [loc >= 0 for loc in token_locs]
 
         copy_bits = 128
         copy_atom = fx.make_copy_atom(fx.UniversalCopy(copy_bits), A.dtype)
@@ -49,7 +51,7 @@ def sorted_sum_kernel(loc_ids: fx.Pointer, A: fx.Pointer, B: fx.Pointer):
         )
 
         token_ptrs = [
-            (A + fx.Int64(token_locs[topk]) * source_row_stride)
+            (A + fx.Int64(valid[topk].select(token_locs[topk], 0)) * source_row_stride)
             for topk in fx.range_constexpr(TOPK)
         ]
 
@@ -63,9 +65,9 @@ def load_atoms(off):
             return [load_atom(topk, off) for topk in fx.range_constexpr(TOPK)]
 
         def reduce_store(dst, frag):
-            vec_sum = frag[0].load().to(fx.Float32)
+            vec_sum = valid[0].select(frag[0].load().to(fx.Float32), 0.0)
             for m in fx.range_constexpr(1, TOPK):
-                vec_sum += frag[m].load().to(fx.Float32)
+                vec_sum += valid[m].select(frag[m].load().to(fx.Float32), 0.0)
             vec_sum = vec_sum.to(dst.dtype)
             out_frag = fx.make_fragment_like(dst)
             out_frag.store(vec_sum)
@@ -120,75 +122,3 @@ def sorted_sum(TOPK, N, row_padding_bytes=None):
 def compile_moe_reduction(*, topk, model_dim, row_padding_bytes=None):
     """Build and cache the route reduction callable for one static configuration."""
     return sorted_sum(topk, model_dim, row_padding_bytes)
-
-
-@functools.cache
-def _invert_sorted_ids_cached(device_cache_key, TOPK):
-    del device_cache_key
-    num_threads = 64
-
-    @flyc.kernel(known_block_size=[num_threads, 1, 1])
-    def invert_sorted_ids_kernel(
-        sorted_ids: fx.Pointer,
-        invert: fx.Pointer,
-        p_num_valid: fx.Pointer,
-        num_ids: fx.Uint32,
-        batch_size: fx.Uint32,
-    ):
-        batch = fx.block_idx.x
-        tid = fx.thread_idx.x
-        slot = batch * num_threads + tid
-        # Scan only the down-written region [0, num_valid): the tail of sorted_ids
-        # (>= num_valid) is uninitialized, and its garbage would racily map real
-        # tokens onto unwritten gemm2_out rows. Read the bound on-device (no host sync).
-        num_valid = p_num_valid[0].to(fx.Uint32)
-        if slot < num_valid:
-            sid = sorted_ids[slot].to(fx.Uint32)
-            tok_id = sid & 0xFFFFFF
-            top_id = sid >> 24
-            idx = tok_id * TOPK + top_id
-            if top_id < TOPK and tok_id < batch_size:
-                invert[idx] = fx.Uint32(slot)
-
-    @flyc.jit
-    def launch(
-        sorted_ids: fx.Pointer,
-        invert: fx.Pointer,
-        p_num_valid: fx.Pointer,
-        num_ids: fx.Uint32,
-        batch_size: fx.Uint32,
-        stream: fx.Stream,
-    ):
-        grid_size = fxh.div_up(num_ids, num_threads)
-        invert_sorted_ids_kernel(
-            sorted_ids, invert, p_num_valid, num_ids, batch_size
-        ).launch(grid=(grid_size, 1, 1), block=(num_threads, 1, 1), stream=stream)
-
-    def callable(
-        sorted_ids: torch.Tensor,
-        invert: torch.Tensor,
-        num_valid: torch.Tensor,
-        num_ids: int,
-        batch_size: int,
-    ):
-        stream = torch.cuda.current_stream()
-        _run_compiled(
-            launch,
-            _ptr(sorted_ids),
-            _ptr(invert),
-            _ptr(num_valid),
-            fx.Uint32(num_ids),
-            fx.Uint32(batch_size),
-            stream,
-        )
-
-    callable._flydsl_executable = launch
-    return callable
-
-
-def invert_sorted_ids(TOPK):
-    return _invert_sorted_ids_cached(get_device_cache_key(), TOPK)
-
-
-invert_sorted_ids.cache_clear = _invert_sorted_ids_cached.cache_clear
-invert_sorted_ids.cache_info = _invert_sorted_ids_cached.cache_info