[FlyDSL] Add full-expert DeepSeek V4.1 Flash MegaMoE tuning for gfx950

aiter/ops/flydsl/kernels/mega_moe/dispatch.py

@@ -575,7 +575,7 @@ def emit_direct_fixed_slot_finalize(
     *, fz_npes, fz_epr, fz_cap, fz_mtpr, fz_rank, fz_tile_m, n_tiles, addr_disp, parity, expected
 ):
     """Finalize local fixed slots as soon as every source publishes this destination."""
-    assert 0 < fz_epr <= 64, "direct fixed-slot finalize requires 1..64 experts per rank"
+    assert 0 < fz_epr <= 256
     dispatch_table = ptr_buf_tensor(addr_disp, fx.Int64)
 
     def dp(i):
@@ -612,23 +612,27 @@ def dp(i):
             comm_ops.wait_i32_until_equals(a_source_done + fx.Int64(done_index) * fx.Int64(4), expected)
         comm_ops.fence_system_acquire()
 
-        valid_expert = lane < fx.Int32(fz_epr)
-        safe_expert = valid_expert.select(lane, fx.Int32(0))
-        count = running[safe_expert]
-        count = valid_expert.select(count, fx.Int32(0))
-        overflow_flag = (count > fx.Int32(fz_cap)).select(fx.Int32(1), fx.Int32(0))
-        overflow_prefix = _wave_inclusive_scan_i32(overflow_flag, lane)
-        overflow_count = fx.Int32(fx.rocdl.readlane(T.i32, overflow_prefix, fz_epr - 1))
-        no_overflow = overflow_count == fx.Int32(0)
-        safe_count = (count <= fx.Int32(fz_cap)).select(count, fx.Int32(0))
-        num_expert_tiles = (safe_count + fx.Int32(fz_tile_m - 1)) // fx.Int32(fz_tile_m)
-        max_expert_tiles = _wave_reduce_max_i32(num_expert_tiles, lane)
-        inclusive_tiles = _wave_inclusive_scan_i32(num_expert_tiles, lane)
-        metadata_base = inclusive_tiles - num_expert_tiles
-        total_tiles = fx.Int32(fx.rocdl.readlane(T.i32, inclusive_tiles, fz_epr - 1))
-
-        if valid_expert:
-            if no_overflow:
+        total_tiles = fx.Int32(0)
+        overflow_count = fx.Int32(0)
+        max_expert_tiles = fx.Int32(0)
+        for chunk in range_constexpr((fz_epr + 63) // 64):
+            expert = lane + fx.Int32(chunk * 64)
+            valid_expert = expert < fx.Int32(fz_epr)
+            safe_expert = valid_expert.select(expert, fx.Int32(0))
+            count = valid_expert.select(running[safe_expert], fx.Int32(0))
+            overflow_flag = (count > fx.Int32(fz_cap)).select(fx.Int32(1), fx.Int32(0))
+            overflow_prefix = _wave_inclusive_scan_i32(overflow_flag, lane)
+            last_lane = min(63, fz_epr - chunk * 64 - 1)
+            overflow_count += fx.Int32(fx.rocdl.readlane(T.i32, overflow_prefix, last_lane))
+            safe_count = (count <= fx.Int32(fz_cap)).select(count, fx.Int32(0))
+            num_expert_tiles = (safe_count + fx.Int32(fz_tile_m - 1)) // fx.Int32(fz_tile_m)
+            chunk_max = _wave_reduce_max_i32(num_expert_tiles, lane)
+            max_expert_tiles = (chunk_max > max_expert_tiles).select(chunk_max, max_expert_tiles)
+            inclusive_tiles = _wave_inclusive_scan_i32(num_expert_tiles, lane)
+            metadata_base = total_tiles + inclusive_tiles - num_expert_tiles
+            total_tiles += fx.Int32(fx.rocdl.readlane(T.i32, inclusive_tiles, last_lane))
+
+            if valid_expert:
                 global_expert = fx.Int32(fz_rank * fz_epr) + safe_expert
                 payload_base = safe_expert * fx.Int32(fz_cap)
                 for tile in range(fx.Int32(0), num_expert_tiles, 1):
@@ -646,10 +650,9 @@ def dp(i):
                         fz_npes * fz_mtpr
                     )
                 expert_tile_end[safe_expert] = metadata_base + num_expert_tiles
-            else:
-                expert_tile_end[safe_expert] = fx.Int32(0)
-            running[safe_expert] = fx.Int32(0)
+                running[safe_expert] = fx.Int32(0)
 
+        no_overflow = overflow_count == fx.Int32(0)
         if lane == fx.Int32(0):
             num_valid = no_overflow.select(total_tiles * fx.Int32(fz_tile_m), fx.Int32(0))
             ready_work = no_overflow.select(total_tiles * fx.Int32(n_tiles), fx.Int32(0))

aiter/ops/flydsl/kernels/mega_moe/envs.py

@@ -18,10 +18,6 @@ def _flag(value: str) -> bool:
 
 # name -> (parser, default)
 environment_variables: dict[str, tuple[Callable[[str], Any], str]] = {
-    # Largest MTPR still dispatched fixed-slot (direct expert slots, no count
-    # exchange). 511 / 1023 / 2047 admit the 256 / 512 / 1024-token capacities,
-    # validated on the EP8 v4_pro layout. Read once, at import.
-    "AITER_MEGA_FIXED_SLOT_MAX_MTPR": (int, "255"),
     # Default of forward(mask_invalid_slots=...): combine skips top-k slots with
     # id -1. Off because the id check costs prefill time; callers that pad with
     # -1 ask for it per call.

aiter/ops/flydsl/kernels/mega_moe/gemm1.py

@@ -38,7 +38,7 @@ def do_tile(m_tile, n_tile_base, expert, sched, a_gather, a_s2r, b_loader, b_sca
 # fmt: on
     N_ACC = M_REPEAT * NUM_ACC_N
     NUM_B_SCALE = NUM_ACC_N // _PACK
-    NUM_A_SCALE = M_REPEAT // _PACK
+    NUM_A_SCALE = (M_REPEAT + _PACK - 1) // _PACK
     B_STATE_END = (
         N_ACC + NUM_ACC_N * _PACK
     )
@@ -378,7 +378,7 @@ def compile_gemm1(
     n_tiles = (2 * inter_dim) // tile_n
     m_repeat = sort_block_m // 16
     num_acc_n = n_per_wave // 16
-    assert num_acc_n % 2 == 0 and m_repeat % 2 == 0
+    assert num_acc_n % 2 == 0 and sort_block_m in (16, 32, 64, 128)
 
     a_k_step_bytes = tile_k
     k_iters = model_dim // tile_k

aiter/ops/flydsl/kernels/mega_moe/gemm2.py

@@ -186,6 +186,7 @@ def gemm2_compute_v2(
     explicit_m_row=None,
     explicit_n_block=None,
     explicit_expert=None,
+    dense_bm16_scales=False,
 ):
     """Run GEMM2, optionally using an explicitly selected expert row/tile."""
     # SBM is the sort padding unit; BM is the compute tile and must divide SBM.
@@ -324,7 +325,9 @@ def issue_a_ds_read(slot):
 
     asc_per_mb = fx.Int32(kScaleSubBlocks) * kAS_per_chunk_dw * fx.Int32(4)
     asc_num = fx.Int64(i32_max_m_blocks) * fx.Int64(asc_per_mb)
-    scale_chunk0 = m_block_idx if const_expr(is_bm16) else m_row // 32
+    scale_chunk0 = (
+        m_block_idx if const_expr(is_bm16 and not dense_bm16_scales) else m_row // 32
+    )
 
     def make_ascale_view(sub):
         base_dw = (scale_chunk0 + fx.Int32(sub)) * kAS_per_chunk_dw
@@ -355,7 +358,12 @@ def load_a_scale_tile(kt):
                 ascale_views[sub][lane_div_16, lane_mod_16, chunk_kt, None],
                 saf,
             )
-            out.append(Vec(saf.load())[0])
+            scale_word = Vec(saf.load())[0]
+            if const_expr(is_bm16 and dense_bm16_scales):
+                scale_word = scale_word.shrui(
+                    ((m_row // fx.Int32(16)) & fx.Int32(1)) * fx.Int32(8)
+                )
+            out.append(scale_word)
         return out
 
     # Stream B weights and scales through registers so use_nt reaches the ISA cache policy.

aiter/ops/flydsl/kernels/mega_moe/gemm_util.py

@@ -355,7 +355,7 @@ def __init__(
         self._rsrc = scale_rsrc
         self._n_scale = model_dim // 32
         self._lane = fx.thread_idx.x % 64
-        self._n_groups = m_repeat // _PACK
+        self._n_groups = (m_repeat + _PACK - 1) // _PACK
         self._sort_block_m = sort_block_m
         self._total_threads = total_threads
         self._tx = fx.thread_idx.x
@@ -440,7 +440,7 @@ def load_step(self, lds_ascale, kstep_i32):
         out = []
         for g in range_constexpr(self._n_groups):
             r0 = fx.Int32(g * 32) + lane_row
-            r1 = r0 + fx.Int32(16)
+            r1 = r0 + fx.Int32(16 if self._sort_block_m >= 32 else 0)
             b = []
             for ksub in range_constexpr(_PACK):
                 for rr in (r0, r1):