From ea4ce21685c32e0c343c2f0fa6edd97a7412ebf0 Mon Sep 17 00:00:00 2001
From: guoliang <ghu@amd.com>
Date: Sun, 2 Aug 2026 06:50:50 +0000
Subject: [PATCH 1/4] fix hang problem and code clean

---
 .../flydsl_dispatch_combine_intranode_op.py   |   2 +-
 kernels/mega_moe/dispatch.py                  | 136 ++++----------
 kernels/mega_moe/gemm1.py                     |  46 ++---
 kernels/mega_moe/gemm2.py                     |  23 ++-
 kernels/mega_moe/gemm_util.py                 | 153 +++++++++-------
 kernels/mega_moe/mega_moe.py                  |  15 +-
 kernels/mega_moe/mega_moe_config.py           | 169 ++++++++++++++++--
 kernels/mega_moe/mega_moe_stage1.py           | 134 ++++++++------
 kernels/mega_moe/mega_moe_stage2.py           |   8 +-
 tests/kernels/test_mega_moe_v2.py             |   1 +
 tests/unit/test_mega_moe_config.py            | 114 +++++++++++-
 11 files changed, 517 insertions(+), 284 deletions(-)

diff --git a/kernels/comm/flydsl_dispatch_combine_intranode_op.py b/kernels/comm/flydsl_dispatch_combine_intranode_op.py
index 880bf569b..f083682a8 100644
--- a/kernels/comm/flydsl_dispatch_combine_intranode_op.py
+++ b/kernels/comm/flydsl_dispatch_combine_intranode_op.py
@@ -269,7 +269,7 @@ def __post_init__(self):
 
     @property
     def is_fp4(self):
-        return self.data_type == torch.float4_e2m1fn_x2
+        return (self.data_type or self.dispatch_dtype) == torch.float4_e2m1fn_x2
 
     @property
     def elem_size(self):
diff --git a/kernels/mega_moe/dispatch.py b/kernels/mega_moe/dispatch.py
index b4e941536..f58a26852 100644
--- a/kernels/mega_moe/dispatch.py
+++ b/kernels/mega_moe/dispatch.py
@@ -43,10 +43,11 @@ class DispatchSlot(IntEnum):
     WORK_HEAD = 25
     WORK_TAIL = 26
     EXPERT_TILE_END = 27
-    ACTIVE_EXPERTS = 28
-    ACTIVE_COUNT = 29
-    RUNNING = 30
-    P2P_RUNNING = 31
+    GROUP_DONE = 28
+    RUNNING = 29
+    P2P_RUNNING = 30
+    LAUNCH_READY = 31
+    P2P_LAUNCH_READY = 32
 
 
 DISPATCH_TABLE_SIZE = max(DispatchSlot) + 1
@@ -89,12 +90,13 @@ def dp(i):
     p_sm = dp(DispatchSlot.P2P_SRCMAP)
     p_running = dp(DispatchSlot.P2P_RUNNING)
     p_source_done = dp(DispatchSlot.P2P_COUNT_DONE)
-    a_producer_done = dp(DispatchSlot.ACTIVE_COUNT)
+    a_producer_done = dp(DispatchSlot.GROUP_DONE)
 
     tid = fx.thread_idx.x
     lane = tid & fx.Int32(63)
     warp = tid >> fx.Int32(6)
     destination_groups = 2
+    assert dispatch_blocks % destination_groups == 0, "direct fixed-slot dispatch needs even producer groups"
     producers_per_group = dispatch_blocks // destination_groups
     producer_group = producer_slot % fx.Int32(destination_groups)
     group_slot = producer_slot // fx.Int32(destination_groups)
@@ -260,7 +262,7 @@ def dp(i):
 @flyc.jit
 def emit_dispatch_plan(
     *, num_waves, fz_npes, fz_epr, fz_k, fz_mtpr, fz_rank, fz_tile_m, fz_total_experts, addr_disp,
-    i32_cur_tok, addr_in_idx, parity, expected, active_expert_producer, external_grouping, external_counting,
+    i32_cur_tok, addr_in_idx, parity, expected, external_grouping, external_counting,
     dispatch_blocks,
 ):
 # fmt: on
@@ -288,9 +290,7 @@ def dp(i):
     a_pair_ready = dp(DispatchSlot.PAIR_READY)
     a_pair_order_ready = dp(DispatchSlot.PAIR_ORDER_READY)
     a_expert_tile_end = dp(DispatchSlot.EXPERT_TILE_END)
-    a_active_experts = dp(DispatchSlot.ACTIVE_EXPERTS)
-    a_active_count = dp(DispatchSlot.ACTIVE_COUNT)
-    p_payload_ready = dp(DispatchSlot.P2P_PAYLOAD_READY)
+    a_group_done = dp(DispatchSlot.GROUP_DONE)
 
     tid = fx.thread_idx.x
     lane = tid & fx.Int32(63)
@@ -306,16 +306,11 @@ def dp(i):
     r_pair_base = crfa(a_pair_base)
     r_pair = crfa(a_pair_order)
     r_lc = crfa(a_lc)
-    if const_expr(active_expert_producer):
-        if tid == fx.Int32(0):
-            buffer_ops.buffer_store(fx.Int32(0), crfa(a_active_count), fx.Int32(0))
-        fx.rocdl.s_waitcnt(0)
-        fx.barrier()
     if const_expr(external_counting):
         if tid == fx.Int32(0):
-            mori_shmem.int32_wait_until_equals(a_active_count, fx.Int32(dispatch_blocks))
+            mori_shmem.int32_wait_until_equals(a_group_done, fx.Int32(dispatch_blocks))
             comm_ops.fence_agent_acquire()
-            buffer_ops.buffer_store(fx.Int32(0), crfa(a_active_count), fx.Int32(0))
+            buffer_ops.buffer_store(fx.Int32(0), crfa(a_group_done), fx.Int32(0))
             fx.rocdl.s_waitcnt(0)
             comm_ops.fence_agent_release()
     else:
@@ -376,7 +371,6 @@ def dp(i):
             ge = fx.Int32(fz_rank * fz_epr + local_expert)
             source_counts = []
             total_count = fx.Int32(0)
-            zero_sources = fx.Int32(0)
             for source in range_constexpr(fz_npes):
                 source_count = buffer_ops.buffer_load(
                     r_bc, fx.Int32(source * fz_epr) + safe_expert, vec_width=1, dtype=fx.Int32
@@ -384,16 +378,6 @@ def dp(i):
                 source_count = valid_expert.select(source_count, fx.Int32(0))
                 source_counts.append(source_count)
                 total_count = total_count + source_count
-                zero_sources = zero_sources + (source_count == fx.Int32(0)).select(fx.Int32(1), fx.Int32(0))
-            if const_expr(active_expert_producer):
-                if valid_expert & (zero_sources > fx.Int32(0)):
-                    local_payload_ready = buffer_ops.buffer_load(
-                        crfa(p_payload_ready), fx.Int32(fz_rank), vec_width=1, dtype=fx.Int64
-                    )
-                    ready_index = parity * fx.Int32(fz_epr) + local_expert
-                    comm_ops.store_i32_system(
-                        local_payload_ready, ready_index, expected - fx.Int32(fz_npes) + zero_sources
-                    )
             num_tiles = (total_count + fx.Int32(fz_tile_m - 1)) // fx.Int32(fz_tile_m)
             padded_rows = num_tiles * fx.Int32(fz_tile_m)
             inclusive_rows = _wave_inclusive_scan_i32(padded_rows, lane)
@@ -442,13 +426,6 @@ def dp(i):
             safe_ge = valid_ge.select(ge, fx.Int32(0))
             source_count = buffer_ops.buffer_load(r_lh, safe_ge, vec_width=1, dtype=fx.Int32)
             source_count = valid_ge.select(source_count, fx.Int32(0))
-            if const_expr(active_expert_producer):
-                if valid_ge & (source_count > fx.Int32(0)):
-                    active_slot = fx.Int32(comm_ops.atomic_add_agent(a_active_count, fx.Int32(1)))
-                    buffer_ops.buffer_store(
-                        (ge % fx.Int32(fz_epr)) * fx.Int32(fz_npes) + ge // fx.Int32(fz_epr),
-                        crfa(a_active_experts), active_slot,
-                    )
             lane_counts.append(source_count)
             lane_total = lane_total + source_count
         lane_prefix = _wave_inclusive_scan_i32(lane_total, lane) - lane_total
@@ -487,7 +464,7 @@ def dp(i):
     fx.barrier()
     if tid == fx.Int32(0):
         if const_expr(external_grouping):
-            mori_shmem.int32_wait_until_equals(a_active_count, fx.Int32(dispatch_blocks))
+            mori_shmem.int32_wait_until_equals(a_group_done, fx.Int32(dispatch_blocks))
             comm_ops.fence_agent_acquire()
         comm_ops.fence_agent_release()
         buffer_ops.buffer_store(expected, crfa(a_pair_order_ready), parity)
@@ -512,7 +489,7 @@ def dp(i):
     a_local_hist = dp(DispatchSlot.LOCAL_HIST)
     a_local_cursor = dp(DispatchSlot.LOCAL_CURSOR)
     a_pair_order = dp(DispatchSlot.PAIR_ORDER)
-    a_group_done = dp(DispatchSlot.ACTIVE_COUNT)
+    a_group_done = dp(DispatchSlot.GROUP_DONE)
     r_idx = crfa(addr_in_idx)
     r_pair = crfa(a_pair_order)
     tid = fx.thread_idx.x
@@ -560,7 +537,7 @@ def dp(i):
 def emit_dispatch_payload(
     *, num_waves, fz_epr, fz_k, fz_mtpr, fz_rank, fz_total_experts, fz_nbytes, fz_n_i32, fz_safe_end_i32,
     fz_scale_n_i32, fz_enable_scales, addr_disp, addr_in_tok, addr_in_wts, addr_in_sc, dispatch_blocks,
-    producer_slot, parity, expected, active_expert_producer, cooperative_payload_copy,
+    producer_slot, parity, expected,
 ):
 # fmt: on
     """Produce independently publishable expert payloads from a compact plan."""
@@ -580,8 +557,6 @@ def dp(i):
     p_payload_ready = dp(DispatchSlot.P2P_PAYLOAD_READY)
     a_pair_order = dp(DispatchSlot.PAIR_ORDER)
     a_plan_ready = dp(DispatchSlot.PLAN_READY)
-    a_active_experts = dp(DispatchSlot.ACTIVE_EXPERTS)
-    a_active_count = dp(DispatchSlot.ACTIVE_COUNT)
 
     tid = fx.thread_idx.x
     lane = tid & fx.Int32(63)
@@ -591,18 +566,11 @@ def dp(i):
     r_mb = crfa(a_mb)
     r_pair = crfa(a_pair_order)
     r_wts = crfa(addr_in_wts)
-    if const_expr(active_expert_producer):
-        task_limit = buffer_ops.buffer_load(crfa(a_active_count), fx.Int32(0), vec_width=1, dtype=fx.Int32)
-    else:
-        task_limit = fx.Int32(fz_total_experts)
+    task_limit = fx.Int32(fz_total_experts)
     task0 = producer_slot
     task_stride = fx.Int32(dispatch_blocks)
-    if const_expr(cooperative_payload_copy):
-        row0 = fx.Int32(0)
-        row_stride = fx.Int32(1)
-    else:
-        row0 = warp
-        row_stride = fx.Int32(num_waves)
+    row0 = warp
+    row_stride = fx.Int32(num_waves)
 
     def _publish_task(destination, local_expert, ge):
         comm_ops.fence_system_release()
@@ -613,27 +581,15 @@ def _publish_task(destination, local_expert, ge):
 
     num_destinations = fz_total_experts // fz_epr
     hoist_remote_resources = fz_mtpr >= 1024
-    if const_expr(not active_expert_producer):
-        producer_destination = producer_slot % fx.Int32(num_destinations)
-        ready_index = parity * fx.Int32(num_destinations) + producer_destination
-        if tid == fx.Int32(0):
-            mori_shmem.int32_wait_until_equals(a_plan_ready + fx.Int64(ready_index) * fx.Int64(4), expected)
-            comm_ops.fence_system_acquire()
-        fx.barrier()
+    producer_destination = producer_slot % fx.Int32(num_destinations)
+    ready_index = parity * fx.Int32(num_destinations) + producer_destination
+    if tid == fx.Int32(0):
+        mori_shmem.int32_wait_until_equals(a_plan_ready + fx.Int64(ready_index) * fx.Int64(4), expected)
+        comm_ops.fence_system_acquire()
+    fx.barrier()
     for task_index in range(task0, task_limit, task_stride):
-        if const_expr(active_expert_producer):
-            task = buffer_ops.buffer_load(crfa(a_active_experts), task_index, vec_width=1, dtype=fx.Int32)
-            local_expert = task // fx.Int32(num_destinations)
-            destination = task % fx.Int32(num_destinations)
-            ready_index = parity * fx.Int32(num_destinations) + destination
-            if tid == fx.Int32(0):
-                mori_shmem.int32_wait_until_equals(a_plan_ready + fx.Int64(ready_index) * fx.Int64(4), expected)
-                comm_ops.fence_system_acquire()
-            fx.barrier()
-        else:
-            task = task_index
-            local_expert = task // fx.Int32(num_destinations)
-            destination = producer_destination
+        local_expert = task_index // fx.Int32(num_destinations)
+        destination = producer_destination
         ge = destination * fx.Int32(fz_epr) + local_expert
         source_count_lane = fx.Int32(0)
         source_base_lane = fx.Int32(0)
@@ -673,18 +629,11 @@ def _copy_route_header():
                     srcmap_remote = buffer_ops.buffer_load(crfa(p_sm), destination, vec_width=1, dtype=fx.Int64)
                     buffer_ops.buffer_store(source_encoding, crfa(srcmap_remote), destination_row)
 
-            if const_expr(not cooperative_payload_copy):
-                if lane == fx.Int32(0):
-                    _copy_route_header()
-            else:
-                if tid == fx.Int32(0):
-                    _copy_route_header()
+            if lane == fx.Int32(0):
+                _copy_route_header()
 
             if const_expr(fz_enable_scales):
-                if const_expr(not cooperative_payload_copy):
-                    scale_lane = lane
-                else:
-                    scale_lane = tid
+                scale_lane = lane
                 if const_expr(fz_scale_n_i32 % 4 == 0):
                     scale_offset = scale_lane * fx.Int32(4)
                     if scale_offset < fx.Int32(fz_scale_n_i32):
@@ -725,24 +674,17 @@ def _copy_route_header():
             else:
                 row_token_remote = buffer_ops.buffer_load(crfa(p_rx), destination, vec_width=1, dtype=fx.Int64)
                 destination_rsrc = crfa(row_token_remote + fx.Int64(destination_row) * fx.Int64(fz_nbytes))
-            if const_expr(not cooperative_payload_copy):
-                lane_offset = lane * fx.Int32(4)
-                if const_expr(fz_safe_end_i32 > 0):
-                    for column in range(lane_offset, fz_safe_end_i32, 512):
-                        value0 = buffer_ops.buffer_load(source_rsrc, column, vec_width=4, dtype=fx.Int32)
-                        value1 = buffer_ops.buffer_load(
-                            source_rsrc, column + fx.Int32(256), vec_width=4, dtype=fx.Int32
-                        )
-                        buffer_ops.buffer_store(value0, destination_rsrc, column)
-                        buffer_ops.buffer_store(value1, destination_rsrc, column + fx.Int32(256))
-                if const_expr(fz_safe_end_i32 < fz_n_i32):
-                    for column in range(lane_offset + fz_safe_end_i32, fz_n_i32, 256):
-                        value = buffer_ops.buffer_load(source_rsrc, column, vec_width=4, dtype=fx.Int32)
-                        buffer_ops.buffer_store(value, destination_rsrc, column)
-            else:
-                thread_offset = tid * fx.Int32(4)
-                thread_stride = fx.Int32(num_waves * 64 * 4)
-                for column in range(thread_offset, fz_n_i32, thread_stride):
+            lane_offset = lane * fx.Int32(4)
+            if const_expr(fz_safe_end_i32 > 0):
+                for column in range(lane_offset, fz_safe_end_i32, 512):
+                    value0 = buffer_ops.buffer_load(source_rsrc, column, vec_width=4, dtype=fx.Int32)
+                    value1 = buffer_ops.buffer_load(
+                        source_rsrc, column + fx.Int32(256), vec_width=4, dtype=fx.Int32
+                    )
+                    buffer_ops.buffer_store(value0, destination_rsrc, column)
+                    buffer_ops.buffer_store(value1, destination_rsrc, column + fx.Int32(256))
+            if const_expr(fz_safe_end_i32 < fz_n_i32):
+                for column in range(lane_offset + fz_safe_end_i32, fz_n_i32, 256):
                     value = buffer_ops.buffer_load(source_rsrc, column, vec_width=4, dtype=fx.Int32)
                     buffer_ops.buffer_store(value, destination_rsrc, column)
 
diff --git a/kernels/mega_moe/gemm1.py b/kernels/mega_moe/gemm1.py
index 117a7d8a5..c4df0f630 100644
--- a/kernels/mega_moe/gemm1.py
+++ b/kernels/mega_moe/gemm1.py
@@ -8,7 +8,6 @@
 import flydsl.expr as fx
 from flydsl.expr import const_expr, range_constexpr, rocdl
 from flydsl.expr.typing import Vector as Vec
-from kernels.common import buffer_ops as _buffer_ops
 from kernels.common.tensor_shim import _run_compiled
 
 from .gemm_util import (
@@ -21,6 +20,8 @@
     MfmaScaleGU,
     SiluQuantEpilogue,
     TileScheduler,
+    _buffer_load,
+    _make_buffer,
     wait_lds_barrier,
 )
 
@@ -44,7 +45,7 @@ def do_tile(m_tile, n_tile_base, expert, sched, a_gather, a_s2r, b_loader, b_sca
     )
     SB_STATE_END = B_STATE_END + NUM_B_SCALE
     last = fx.Int32(K_ITERS - 1)
-    tile_row_base = _buffer_ops.buffer_load(trb_rsrc, m_tile, vec_width=1, dtype=fx.Int32)
+    tile_row_base = _buffer_load(trb_rsrc, m_tile, fx.Int32)
     b_row = sched.gate_base_row(expert) + n_tile_base
     a_gather.for_tile(tile_row_base)
     if const_expr(pipe_weights):
@@ -264,8 +265,8 @@ def a_load(mi, ks, _base=cur_off):
 
 
 # fmt: off
-def build_fused_gemm1(*, x_rsrc, x_base_addr, x_tensor, w_rsrc, sw_rsrc, sx_rsrc,
-    out_rsrc, os_rsrc, trb_rsrc, expert_rsrc, out_base_addr, a_buf, a_scale_lds, c_tile,
+def build_fused_gemm1(*, x_tensor, w_rsrc, sw_rsrc, sx_rsrc,
+    out_rsrc, os_rsrc, trb_rsrc, expert_rsrc, out_tensor, a_buf, a_scale_lds, c_tile,
     model_dim, inter_dim, sort_block_m, tile_n, num_waves, n_per_wave, wave_id,
     m_repeat, num_acc_n, a_k_step_bytes, total_threads, k_iters, a_lds_i32, n_tiles,
     expert_offset, b_cache_modifier, swizzle_a, pipe_weights, mfma_amajor, async_a_copy,
@@ -280,10 +281,9 @@ def build_fused_gemm1(*, x_rsrc, x_base_addr, x_tensor, w_rsrc, sw_rsrc, sx_rsrc
     n_wave_base = wave_id * fx.Int32(n_per_wave)
 
     # fmt: off
-    a_gather = ATileLoader(x_rsrc=x_rsrc, row_bytes=model_dim, sort_block_m=sort_block_m,
+    a_gather = ATileLoader(row_bytes=model_dim, sort_block_m=sort_block_m,
         k_step_bytes=a_k_step_bytes, total_threads=total_threads, swizzle=swizzle_a,
-        x_base_addr=x_base_addr, x_tensor=x_tensor,
-        async_copy=async_a_copy)
+        x_tensor=x_tensor, async_copy=async_a_copy)
     # fmt: on
     a_s2r = AS2RLoader(k_step_bytes=a_k_step_bytes, swizzle=swizzle_a)
     b_loader = BWeightLoader(
@@ -305,7 +305,7 @@ def build_fused_gemm1(*, x_rsrc, x_base_addr, x_tensor, w_rsrc, sw_rsrc, sx_rsrc
     epi = SiluQuantEpilogue(out_rsrc=out_rsrc, out_scale_rsrc=os_rsrc, sorted_rsrc=trb_rsrc, tokens=0,
         inter_dim=inter_dim, m_repeat=m_repeat, num_acc_n=num_acc_n, sort_block_m=sort_block_m, tile_n=tile_n,
         num_waves=num_waves, lds_out=c_tile, always_valid=True,
-        out_base_addr=out_base_addr if use_tile_resource else None)
+        out_tensor=out_tensor if use_tile_resource else None)
     # fmt: on
 
     def _decode(flat):
@@ -380,30 +380,30 @@ def kernel(
         a_scale_lds = lds.A_scale
         c_tile = _LdsF32View(fx.recast_iter(fx.Float32, lds.pool.ptr))
 
-        x_rsrc = _buffer_ops.create_buffer_resource(x, max_size=True)
-        x_base_addr = fx.Int64(_buffer_ops.extract_base_index(x, address_space=1))
-        w_rsrc = _buffer_ops.create_buffer_resource(w, max_size=True)
-        sx_rsrc = _buffer_ops.create_buffer_resource(scale_x, max_size=True)
-        sw_rsrc = _buffer_ops.create_buffer_resource(scale_w, max_size=True)
-        trb_rsrc = _buffer_ops.create_buffer_resource(tile_row_base, max_size=True)
-        expert_rsrc = _buffer_ops.create_buffer_resource(expert_ids, max_size=True)
-        out_base_addr = fx.Int64(_buffer_ops.extract_base_index(out, address_space=1))
+        w_rsrc = _make_buffer(w, fx.Int32, 4)
+        sx_rsrc = _make_buffer(scale_x, fx.Int32, 4)
+        sw_rsrc = _make_buffer(scale_w, fx.Int32)
+        trb_rsrc = _make_buffer(tile_row_base, fx.Int32)
+        expert_rsrc = _make_buffer(expert_ids, fx.Int32)
         if const_expr(use_tile_resource):
-            out_rsrc = _buffer_ops.create_buffer_resource(out, max_size=True)
+            out_rsrc = None
         else:
-            out_rsrc = _buffer_ops.create_buffer_resource(
-                out, max_size=False, num_records_bytes=num_valid * fx.Int32(inter_dim)
+            out_rsrc = _make_buffer(
+                out, fx.Int16, max_size=False, num_records_bytes=num_valid * fx.Int32(inter_dim)
             )
         scale_cols = (inter_dim // 32 + 7) // 8 * 8
-        os_rsrc = _buffer_ops.create_buffer_resource(
-            out_scale, max_size=False, num_records_bytes=num_valid * fx.Int32(scale_cols) + fx.Int32(8192)
+        os_rsrc = _make_buffer(
+            out_scale,
+            fx.Int8,
+            max_size=False,
+            num_records_bytes=num_valid * fx.Int32(scale_cols) + fx.Int32(8192),
         )
         wave_id = fx.thread_idx.x // 64
 
         _, run_tile = build_fused_gemm1(
-            x_rsrc=x_rsrc, x_base_addr=x_base_addr, x_tensor=x, w_rsrc=w_rsrc, sw_rsrc=sw_rsrc,
+            x_tensor=x, w_rsrc=w_rsrc, sw_rsrc=sw_rsrc,
             sx_rsrc=sx_rsrc, out_rsrc=out_rsrc, os_rsrc=os_rsrc, trb_rsrc=trb_rsrc,
-            expert_rsrc=expert_rsrc, out_base_addr=out_base_addr, a_buf=a_buf,
+            expert_rsrc=expert_rsrc, out_tensor=out, a_buf=a_buf,
             a_scale_lds=a_scale_lds, c_tile=c_tile, model_dim=model_dim, inter_dim=inter_dim,
             sort_block_m=sort_block_m, tile_n=tile_n, num_waves=num_waves, n_per_wave=n_per_wave,
             wave_id=wave_id, m_repeat=m_repeat, num_acc_n=num_acc_n, a_k_step_bytes=a_k_step_bytes,
diff --git a/kernels/mega_moe/gemm2.py b/kernels/mega_moe/gemm2.py
index ee80b1385..17bbb0e62 100644
--- a/kernels/mega_moe/gemm2.py
+++ b/kernels/mega_moe/gemm2.py
@@ -4,7 +4,6 @@
 
 import flydsl.compiler as flyc
 import flydsl.expr as fx
-from flydsl.expr import _to_raw as _raw
 from flydsl.expr import const_expr, gpu, range_constexpr, rocdl
 from flydsl.expr.typing import (
     Float4E2M1FN,
@@ -30,7 +29,7 @@
 
 def scale_view(arg_scale, base_dw, K_TILES_TOTAL, k0_stride_dw=64, num_records_bytes=None):
     """View one e8m0 scale word, optionally bounded to the real buffer extent."""
-    base_dw = rocdl.readfirstlane(T.i32, _raw(base_dw))
+    base_dw = rocdl.readfirstlane(T.i32, base_dw)
     i32_ptr_ty = fx.PointerType.get(T.i32, address_space=fx.AddressSpace.Global, alignment=4)
     off_i64 = fx.Int64(base_dw)
     base_iter = fx.inttoptr(i32_ptr_ty, fx.Int64(arg_scale) + off_i64 * fx.Int64(4))
@@ -211,11 +210,11 @@ def gemm2_compute_v2(
     n_block_idx = bx_i32 - m_block_idx * num_n_blocks
     eids_ptr = global_typed_ptr(arg_eids, T.i32)
     if const_expr(SBM == BM):
-        e = rocdl.readfirstlane(T.i32, _raw(eids_ptr[m_block_idx]))
+        e = rocdl.readfirstlane(T.i32, eids_ptr[m_block_idx])
         m_row = m_block_idx * BM
     else:
         m_row = m_block_idx * BM
-        e = rocdl.readfirstlane(T.i32, _raw(eids_ptr[m_row // fx.Int32(SBM)]))
+        e = rocdl.readfirstlane(T.i32, eids_ptr[m_row // fx.Int32(SBM)])
     if const_expr(expert_offset != 0):
         e = e - fx.Int32(expert_offset)
 
@@ -319,7 +318,7 @@ def load_a_scale_tile(kt):
                 ascale_views[sub][lane_div_16, lane_mod_16, chunk_kt, None],
                 saf,
             )
-            out.append(_raw(Vec(saf.load())[0]))
+            out.append(Vec(saf.load())[0])
         return out
 
     # Stream B weights and scales through registers so use_nt reaches the ISA cache policy.
@@ -391,12 +390,12 @@ def shift_scale_word(scale, kt_rt):
         if const_expr(tilesPerScaleChunk == 1):
             return scale
         scale_shift = (kt_rt % fx.Int32(tilesPerScaleChunk)) * fx.Int32(16)
-        return _raw(fx.Int32(scale).shrui(scale_shift))
+        return fx.Int32(scale).shrui(scale_shift)
 
     def mfma_cluster(bqf, bsf, sa, kt_rt):
         # opsel (no gate/up split): mni=J//2, in_b=J%2; sa is a per-32-row-chunk list.
         sa = [shift_scale_word(sa[sub], kt_rt) for sub in range_constexpr(kScaleSubBlocks)]
-        sb_words = [shift_scale_word(_raw(Vec(bsf[mni].load())[0]), kt_rt) for mni in range_constexpr(nPairs)]
+        sb_words = [shift_scale_word(Vec(bsf[mni].load())[0], kt_rt) for mni in range_constexpr(nPairs)]
         for J in range_constexpr(numAccN):
             mni, in_b = J // 2, J % 2
             sb = sb_words[mni]
@@ -449,7 +448,7 @@ def store_c_carry(state):
         return n
 
     if const_expr(BM == 64 and BN == 256):
-        # BM64/BN256 uses the 1-stage B path unconditionally; do not depend on env knobs.
+        # BM64/BN256 uses the 1-stage B path unconditionally.
         for kt_iv, state in range(
             fx.Int32(0),
             K_TILES_RT,
@@ -529,7 +528,7 @@ def issue_a_scale_load_into(saf, kt_rt):
             # A-scale vmem load(s) for K-tile kt_rt into the given (per-stage) fragment(s).
             sa = load_a_scale_tile(kt_rt)
             for sub in range_constexpr(kScaleSubBlocks):
-                saf[sub].store(sa[sub])
+                saf[sub].store(Vec.filled(1, sa[sub], Int32))
 
         def load_carry():
             return load_c_carry() + load_b_carry()
@@ -581,7 +580,7 @@ def prefetch_next_b(kt_rt):
                 issue_a_load_lds(nxt_a % fx.Int32(aStages), nxt_a)
             # A-scale from the prefetch carry (g2_ascale_pf) or loaded synchronously here.
             if const_expr(g2_ascale_pf):
-                sa = [_raw(Vec(cur_saf[sub].load())[0]) for sub in range_constexpr(kScaleSubBlocks)]
+                sa = [Vec(cur_saf[sub].load())[0] for sub in range_constexpr(kScaleSubBlocks)]
             else:
                 sa = load_a_scale_tile(kt_rt)
             if const_expr(not g2_bhoist):
@@ -614,10 +613,10 @@ def _spart_output_tile_index(block_1d_id, M0, N0, group_num, m01):
     group_id_y = block_1d_id // gn
     group_id_x = block_1d_id - group_id_y * gn
 
-    # remap = group_id_x <= big_group_num ? gx*gs + gy : gx*gs + big - gx + gy
+    # remap = group_id_x < big_group_num ? gx*gs + gy : gx*gs + big - gx + gy
     remap_a = group_id_x * group_size + group_id_y
     remap_b = group_id_x * group_size + big_group_num - group_id_x + group_id_y
-    remap = (group_id_x <= big_group_num).select(remap_a, remap_b)
+    remap = (group_id_x < big_group_num).select(remap_a, remap_b)
 
     idx_M0 = remap // n0
     idx_N0 = remap - idx_M0 * n0
diff --git a/kernels/mega_moe/gemm_util.py b/kernels/mega_moe/gemm_util.py
index b38124cdd..6aa8a50dc 100644
--- a/kernels/mega_moe/gemm_util.py
+++ b/kernels/mega_moe/gemm_util.py
@@ -7,11 +7,41 @@
 from flydsl.expr import const_expr, gpu, range_constexpr, rocdl
 from flydsl.expr.typing import T
 from flydsl.expr.typing import Vector as Vec
-from kernels.common import buffer_ops as _buffer_ops
 
 _PACK = 2  # fp4 micro-scale pack (per-32 E8M0): pack_M = pack_N = pack_K = 2
 
 
+def _make_buffer(tensor, elem_ty, width=1, *, max_size=True, num_records_bytes=None):
+    alignment = max(1, elem_ty.width * width // 8)
+    ptr_ty = fx.PointerType.get(elem_ty.ir_type, fx.AddressSpace.Global, alignment)
+    base = fx.inttoptr(ptr_ty, fx.Int64(fx.ptrtoint(fx.get_iter(tensor))))
+    view = fx.Tensor(fx.make_view(base, fx.make_layout((width, 1), (1, width))))
+    return fx.rocdl.make_buffer_tensor(view, max_size=max_size, num_records_bytes=num_records_bytes)
+
+
+def _make_buffer_from_addr(addr, elem_ty, width=1, *, num_records_bytes=None):
+    alignment = max(1, elem_ty.width * width // 8)
+    ptr_ty = fx.PointerType.get(elem_ty.ir_type, fx.AddressSpace.Global, alignment)
+    base = fx.inttoptr(ptr_ty, fx.Int64(addr))
+    view = fx.Tensor(fx.make_view(base, fx.make_layout((width, 1), (1, width))))
+    return fx.rocdl.make_buffer_tensor(view, num_records_bytes=num_records_bytes)
+
+
+def _buffer_load(buffer, group_index, elem_ty, width=1, cache_modifier=0):
+    atom = fx.make_copy_atom(fx.rocdl.BufferCopy(elem_ty.width * width, cache_modifier), elem_ty)
+    fragment = fx.make_rmem_tensor(width, elem_ty)
+    fx.copy(atom, fx.slice(buffer, (None, group_index)), fragment)
+    value = Vec(fragment.load())
+    return value[0] if width == 1 else value
+
+
+def _buffer_store(buffer, group_index, value, elem_ty, width=1, cache_modifier=0):
+    atom = fx.make_copy_atom(fx.rocdl.BufferCopy(elem_ty.width * width, cache_modifier), elem_ty)
+    fragment = fx.make_rmem_tensor(width, elem_ty)
+    fragment.store(Vec.from_elements([value], elem_ty) if width == 1 else Vec(value))
+    fx.copy(atom, fragment, fx.slice(buffer, (None, group_index)))
+
+
 def wait_lds_barrier(vmcnt=63):
     """Drain LDS writes and optionally older VMEM while preserving newer loads."""
     waitcnt = (vmcnt & 0xF) | ((vmcnt & 0x30) << 10) | (7 << 4)
@@ -29,7 +59,7 @@ def __init__(self, *, expert_rsrc, inter_dim, expert_offset=0):
         self._expert_offset = int(expert_offset)
 
     def expert_of(self, m_tile_i32):
-        g = _buffer_ops.buffer_load(self._expert_rsrc, m_tile_i32, vec_width=1, dtype=fx.Int32)
+        g = _buffer_load(self._expert_rsrc, m_tile_i32, fx.Int32)
         if const_expr(self._expert_offset != 0):
             return g - fx.Int32(self._expert_offset)
         return g
@@ -44,18 +74,14 @@ class ATileLoader:
     def __init__(
         self,
         *,
-        x_rsrc,
         row_bytes,
         sort_block_m,
         k_step_bytes,
         total_threads,
         swizzle=False,
-        x_base_addr=None,
         x_tensor=None,
         async_copy=False,
     ):
-        self._x_rsrc = x_rsrc
-        self._x_base_addr = x_base_addr
         self._sort_block_m = sort_block_m
         self._k_step_bytes = k_step_bytes
         self._total_threads = total_threads
@@ -65,8 +91,8 @@ def __init__(
         self._wave = self._tx // 64
         self._x_tensor = x_tensor
         self._async_copy = bool(async_copy)
+        assert x_tensor is not None
         if const_expr(self._async_copy):
-            assert x_tensor is not None
             assert total_threads % 64 == 0
             assert (sort_block_m * 16) % total_threads == 0
             assert row_bytes % 16 == 0 and k_step_bytes % 16 == 0
@@ -77,26 +103,19 @@ def __init__(
 
     def for_tile(self, tile_row_base_i32):
         """Precompute LDS and tile-local global offsets for one M tile."""
-        if self._x_base_addr is not None:
-            tile_base_addr = self._x_base_addr + fx.Int64(tile_row_base_i32) * fx.Int64(self._row_bytes)
-            self._tile_rsrc = _buffer_ops.create_buffer_resource_from_addr(
-                tile_base_addr,
-                num_records_bytes=self._sort_block_m * self._row_bytes,
-            )
+        tile_iter = fx.add_offset(
+            fx.get_iter(self._x_tensor),
+            fx.Int64(tile_row_base_i32) * fx.Int64(self._row_bytes),
+        )
+        tile_view = fx.Tensor(fx.make_view(tile_iter, fx.make_layout(self._sort_block_m * self._row_bytes, 1)))
+        self._tile_rsrc = _make_buffer(
+            tile_view,
+            fx.Int32,
+            4,
+            max_size=False,
+            num_records_bytes=self._sort_block_m * self._row_bytes,
+        )
         if const_expr(self._async_copy):
-            tile_iter = fx.add_offset(
-                fx.get_iter(self._x_tensor),
-                fx.Int64(tile_row_base_i32) * fx.Int64(self._row_bytes),
-            )
-            tile_view = fx.Tensor(
-                fx.make_view(
-                    tile_iter,
-                    fx.make_layout(
-                        self._sort_block_m * self._row_bytes,
-                        1,
-                    ),
-                )
-            )
             tile_buffer = fx.rocdl.make_buffer_tensor(
                 tile_view,
                 max_size=False,
@@ -114,10 +133,7 @@ def for_tile(self, tile_row_base_i32):
             lin = fx.Int32(c) + fx.Int32(self._tx)
             row = lin // fx.Int32(chunks_per_row)
             chunk = lin % fx.Int32(chunks_per_row)
-            if self._x_base_addr is not None:
-                row_byte = row * fx.Int32(self._row_bytes)
-            else:
-                row_byte = (tile_row_base_i32 + row) * fx.Int32(self._row_bytes)
+            row_byte = row * fx.Int32(self._row_bytes)
             if const_expr(self._swizzle):
                 col_i32 = chunk * fx.Int32(4)
                 swz = row * fx.Int32(row_stride_i32) + (col_i32 ^ ((row & fx.Int32(15)) << fx.Int32(2)))
@@ -131,17 +147,21 @@ def load_regs(self, k_step_byte_off):
         koff = fx.Int32(k_step_byte_off)
         regs = []
         for lds_byte, chunk_base in self._chunks:
-            g_i32 = (chunk_base + koff) // fx.Int32(4)
-            x_rsrc = self._tile_rsrc if self._x_base_addr is not None else self._x_rsrc
-            regs.append((lds_byte, _buffer_ops.buffer_load(x_rsrc, g_i32, vec_width=4, dtype=fx.Int32)))
+            group = (chunk_base + koff) // fx.Int32(16)
+            regs.append((lds_byte, _buffer_load(self._tile_rsrc, group, fx.Int32, 4)))
         return regs
 
     def store(self, lds_dst, regs, base_i32=0):
         """Scatter loaded chunks into LDS via ds_write (precomputed lds_byte incl. swizzle). base_i32 = ping/pong."""
         base_bytes = fx.Int32(base_i32) * fx.Int32(4)
         for lds_byte, v in regs:
-            ptr = fx.recast_iter(fx.Uint8, fx.add_offset(lds_dst.ptr, fx.make_int_tuple(base_bytes + lds_byte)))
-            fx.ptr_store(Vec(v).bitcast(fx.Uint8), ptr)
+            dst = fx.make_view(
+                fx.add_offset(fx.recast_iter(fx.Int32, lds_dst.ptr), (base_bytes + lds_byte) // fx.Int32(4)),
+                fx.make_layout(4, 1),
+            )
+            fragment = fx.make_rmem_tensor(4, fx.Int32)
+            fragment.store(Vec(v))
+            fx.copy(fx.make_copy_atom(fx.UniversalCopy128b(), fx.Int32), fragment, dst)
 
     def prefetch_to_lds(self, k_step_byte_off, lds_dst, base_i32=0):
         """Issue swizzled direct global-to-LDS copies with a wave-uniform LDS base."""
@@ -234,15 +254,13 @@ def _load_pack(self, row_base_i32, ni, kstep_i32, ksub):
             + lane_k * fx.Int32(self._stride_klane)
             + lane_row * fx.Int32(self._stride_nlane)
         )
-        i32_off = byte // fx.Int32(4)
-        v = _buffer_ops.buffer_load(
+        return _buffer_load(
             self._w_rsrc,
-            i32_off,
-            vec_width=4,
-            dtype=fx.Int32,
-            cache_modifier=self._cache_modifier,
+            byte // fx.Int32(16),
+            fx.Int32,
+            4,
+            self._cache_modifier,
         )
-        return Vec(v)
 
     def load_step(self, row_base_i32, kstep_i32):
         """list[num_acc_n] of [ksub0_i32x4, ksub1_i32x4] for this K-step."""
@@ -271,7 +289,7 @@ def load_step(self, base_row_i32, kstep_i32):
         out = []
         for g in range_constexpr(self._n_groups):
             off = (base_group + fx.Int32(g)) * fx.Int32(self._row_stride) + kterm + lane
-            out.append(_buffer_ops.buffer_load(self._rsrc, off, vec_width=1, dtype=fx.Int32))
+            out.append(_buffer_load(self._rsrc, off, fx.Int32))
         return out
 
 
@@ -297,21 +315,19 @@ def stage(self, lds_ascale, tile_row_base_i32):
         @flyc.jit
         def copy_chunk(lin: fx.Int32):
             if lin < fx.Int32(n16):
-                g_i32 = (base + lin * fx.Int32(16)) // fx.Int32(4)
-                v = _buffer_ops.buffer_load(
+                v = _buffer_load(
                     self._rsrc,
-                    g_i32,
-                    vec_width=4,
-                    dtype=fx.Int32,
+                    (base + lin * fx.Int32(16)) // fx.Int32(16),
+                    fx.Int32,
+                    4,
                 )
-                ptr = fx.recast_iter(
-                    fx.Uint8,
-                    fx.add_offset(
-                        lds_ascale.ptr,
-                        fx.make_int_tuple(lin * fx.Int32(16)),
-                    ),
+                dst = fx.make_view(
+                    fx.add_offset(fx.recast_iter(fx.Int32, lds_ascale.ptr), lin * fx.Int32(4)),
+                    fx.make_layout(4, 1),
                 )
-                fx.ptr_store(Vec(v).bitcast(fx.Uint8), ptr)
+                fragment = fx.make_rmem_tensor(4, fx.Int32)
+                fragment.store(v)
+                fx.copy(fx.make_copy_atom(fx.UniversalCopy128b(), fx.Int32), fragment, dst)
 
         for c in range_constexpr(0, n16, self._total_threads):
             lin = fx.Int32(c) + fx.Int32(self._tx)
@@ -478,7 +494,7 @@ class SiluQuantEpilogue:
 
     # fmt: off
     def __init__(self, *, out_rsrc, out_scale_rsrc, sorted_rsrc, tokens, inter_dim, m_repeat, num_acc_n,
-        sort_block_m, tile_n, num_waves, lds_out, always_valid=False, out_base_addr=None):
+        sort_block_m, tile_n, num_waves, lds_out, always_valid=False, out_tensor=None):
     # fmt: on
         self._out_rsrc = out_rsrc
         self._out_scale_rsrc = out_scale_rsrc
@@ -492,7 +508,7 @@ def __init__(self, *, out_rsrc, out_scale_rsrc, sorted_rsrc, tokens, inter_dim,
         self._num_waves = num_waves
         self._lds_out = lds_out
         self._always_valid = always_valid
-        self._out_base_addr = out_base_addr
+        self._out_tensor = out_tensor
         self._lane = fx.thread_idx.x % 64
         self._sorted_scale_cols_i32 = (inter_dim // 32 + 7) // 8 * 8
 
@@ -518,11 +534,16 @@ def _silu_mul(self, gate_v4, up_v4):
     def store(self, acc, tile_i32, tile_row_base_i32, n_tile_base_i32):
         combined = self._combine(acc)
         n_per = len(combined) // self._m_repeat
-        if self._out_base_addr is not None:
-            out_rsrc = _buffer_ops.create_buffer_resource_from_addr(
-                self._out_base_addr
-                + fx.Int64(tile_i32)
-                * fx.Int64(self._sort_block_m * self._inter_dim),
+        if self._out_tensor is not None:
+            tile_iter = fx.add_offset(
+                fx.get_iter(self._out_tensor),
+                fx.Int64(tile_i32) * fx.Int64(self._sort_block_m * self._inter_dim),
+            )
+            tile_view = fx.Tensor(fx.make_view(tile_iter, fx.make_layout(1, 1)))
+            out_rsrc = _make_buffer(
+                tile_view,
+                fx.Int16,
+                max_size=False,
                 num_records_bytes=self._sort_block_m * self._inter_dim,
             )
         else:
@@ -568,11 +589,11 @@ def store(self, acc, tile_i32, tile_row_base_i32, n_tile_base_i32):
                 valid = fx.Boolean(True)
                 out_row_base = (
                     row * fx.Int32(self._inter_dim)
-                    if self._out_base_addr is not None
+                    if self._out_tensor is not None
                     else row_g * fx.Int32(self._inter_dim)
                 )
             else:
-                tok = _buffer_ops.buffer_load(self._sorted_rsrc, slot, vec_width=1, dtype=fx.Int32)
+                tok = _buffer_load(self._sorted_rsrc, slot, fx.Int32)
                 valid = tok < fx.Int32(self._tokens)
                 out_row_base = slot * fx.Int32(self._inter_dim)
             for nr in range_constexpr(n_reps):
@@ -599,7 +620,7 @@ def store(self, acc, tile_i32, tile_row_base_i32, n_tile_base_i32):
                 short_raw = fx.Int32(packed).to(fx.Int16)
                 out_byte = out_row_base + gcol
                 out_byte = valid.select(out_byte, fx.Int32(0x40000000))
-                _buffer_ops.buffer_store(short_raw, out_rsrc, out_byte, offset_is_bytes=True)
+                _buffer_store(out_rsrc, out_byte // fx.Int32(2), short_raw, fx.Int16)
 
                 col_s = gcol >> fx.Int32(5)
                 is_writer = (gcol & fx.Int32(31)) == fx.Int32(0)
@@ -612,5 +633,5 @@ def store(self, acc, tile_i32, tile_row_base_i32, n_tile_base_i32):
                 byte_off = d0 * n32 + d3 * fx.Int32(256) + d5 * fx.Int32(64) + d2 * fx.Int32(4) + d4 * fx.Int32(2) + d1
                 byte_off = is_writer.select(byte_off, fx.Int32(0x40000000))
                 e8m0_i8 = e8m0_v.to(fx.Int8)
-                _buffer_ops.buffer_store(e8m0_i8, self._out_scale_rsrc, byte_off, offset_is_bytes=True)
+                _buffer_store(self._out_scale_rsrc, byte_off, e8m0_i8, fx.Int8)
         wait_lds_barrier()
diff --git a/kernels/mega_moe/mega_moe.py b/kernels/mega_moe/mega_moe.py
index 8cfbfe823..e8505c303 100644
--- a/kernels/mega_moe/mega_moe.py
+++ b/kernels/mega_moe/mega_moe.py
@@ -111,20 +111,21 @@ def _allocate_dispatch_workspace(self, op):
             "work_head": torch.zeros(8 * 16, dtype=torch.int32, device=self.dev),
             "work_tail": torch.zeros(1, dtype=torch.int32, device=self.dev),
             "expert_tile_end": torch.empty(self.epr, dtype=torch.int32, device=self.dev),
-            "active_experts": torch.empty(total_experts, dtype=torch.int32, device=self.dev),
-            "active_count": torch.zeros(self.world_size, dtype=torch.int32, device=self.dev),
+            "group_done": torch.zeros(1, dtype=torch.int32, device=self.dev),
         }
         workspace["bigcnt"] = op._sym((self.world_size * self.epr,), torch.int32)
         workspace["count_done"] = op._sym((2 * self.world_size,), torch.int32)
         workspace["my_base"] = op._sym((total_experts,), torch.int32)
         workspace["plan_ready"] = op._sym((2 * self.world_size,), torch.int32)
         workspace["payload_ready"] = op._sym((2 * self.epr,), torch.int32)
+        workspace["launch_ready"] = op._sym((self.world_size,), torch.int32)
         ms.shmem_barrier_all()
         workspace["p2p_bigcnt"] = op._p2p_table(workspace["bigcnt"])
         workspace["p2p_count_done"] = op._p2p_table(workspace["count_done"])
         workspace["p2p_my_base"] = op._p2p_table(workspace["my_base"])
         workspace["p2p_plan_ready"] = op._p2p_table(workspace["plan_ready"])
         workspace["p2p_payload_ready"] = op._p2p_table(workspace["payload_ready"])
+        workspace["p2p_launch_ready"] = op._p2p_table(workspace["launch_ready"])
         self._s1_dispatch_workspace = workspace
 
     def _build_v2_disp_table(self):
@@ -159,10 +160,11 @@ def _build_v2_disp_table(self):
         table[DispatchSlot.WORK_HEAD] = workspace["work_head"].data_ptr()
         table[DispatchSlot.WORK_TAIL] = workspace["work_tail"].data_ptr()
         table[DispatchSlot.EXPERT_TILE_END] = workspace["expert_tile_end"].data_ptr()
-        table[DispatchSlot.ACTIVE_EXPERTS] = workspace["active_experts"].data_ptr()
-        table[DispatchSlot.ACTIVE_COUNT] = workspace["active_count"].data_ptr()
+        table[DispatchSlot.GROUP_DONE] = workspace["group_done"].data_ptr()
         table[DispatchSlot.RUNNING] = op.running.data_ptr()
         table[DispatchSlot.P2P_RUNNING] = op.p2p_running.data_ptr()
+        table[DispatchSlot.LAUNCH_READY] = workspace["launch_ready"].data_ptr()
+        table[DispatchSlot.P2P_LAUNCH_READY] = workspace["p2p_launch_ready"].data_ptr()
         self._s1_disp = torch.tensor(table, dtype=torch.int64, device=self.dev)
 
     def _select_config(self, tokens: int) -> MegaMoEConfig:
@@ -207,9 +209,8 @@ def _run_fused_stage1(self, x, wts, scales, topk_ids, stream=None, config: Stage
             sort_block_m=config.sort_block_m, tile_n=config.tile_n, tile_k=config.tile_k,
             num_waves=config.num_waves, grid_mult=config.grid_mult, pipe_weights=config.pipe_weights,
             mfma_amajor=config.mfma_amajor, swizzle_a=config.swizzle_a,
-            async_a_copy=config.async_a_copy, active_expert_producer=config.active_expert_producer,
-            cooperative_payload_copy=config.cooperative_payload_copy,
-            num_dispatch_cu=config.num_dispatch_cu, use_tile_resource=config.use_tile_resource,
+            async_a_copy=config.async_a_copy, num_dispatch_cu=config.num_dispatch_cu,
+            use_tile_resource=config.use_tile_resource,
             waves_per_eu_hint=config.waves_per_eu_hint, b_nt=config.b_nt,
             work_shards=config.work_shards, external_grouping=config.external_grouping,
             external_counting=config.external_counting)
diff --git a/kernels/mega_moe/mega_moe_config.py b/kernels/mega_moe/mega_moe_config.py
index 6c787d882..94932975e 100644
--- a/kernels/mega_moe/mega_moe_config.py
+++ b/kernels/mega_moe/mega_moe_config.py
@@ -6,8 +6,24 @@
 from dataclasses import dataclass, replace
 from functools import lru_cache
 
-TOKEN_BUCKETS = (1, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768)
-P2P_FP8_MIN_TOKENS = 1024
+TOKEN_BUCKETS = (
+    1,
+    4,
+    8,
+    16,
+    32,
+    64,
+    128,
+    256,
+    512,
+    1024,
+    2048,
+    4096,
+    8192,
+    16384,
+    32768,
+)
+P2P_FP8_MIN_MTPR = 1024
 FIXED_SLOT_MAX_MTPR = 255
 
 
@@ -26,8 +42,6 @@ class Stage1Config:
     tile_k: int = 256
     pipe_weights: bool = True
     swizzle_a: bool = True
-    active_expert_producer: bool = False
-    cooperative_payload_copy: bool = False
     work_shards: int = 8
     external_grouping: bool = False
     external_counting: bool = False
@@ -75,7 +89,63 @@ def __post_init__(self):
     128: (3, 224, False, 2, 3),
 }
 
-_COMPACT_SMALL_DISPATCH_CU = {1: 224, 4: 128, 8: 192, 16: 64, 32: 128, 64: 192, 128: 128}
+_COMPACT_SMALL_DISPATCH_CU = {
+    1: 224,
+    4: 128,
+    8: 192,
+    16: 64,
+    32: 128,
+    64: 192,
+    128: 128,
+}
+
+_OVERSIZED_SMALL_DISPATCH_CU = {
+    4096: {4: 224, 8: 128},
+    16384: {1: 128, 4: 224, 8: 64, 64: 64},
+    32768: {4: 224, 8: 128, 128: 64},
+}
+
+# work shards, external grouping, external counting
+_OVERSIZED_SMALL_PROTOCOL = {
+    2048: {
+        4: (8, False, False),
+        32: (8, False, False),
+        128: (8, False, False),
+    },
+    4096: {
+        1: (1, False, False),
+        32: (8, False, False),
+        64: (1, False, False),
+        128: (8, False, False),
+    },
+    8192: {
+        1: (1, False, False),
+        4: (1, True, False),
+        8: (1, False, False),
+        16: (1, False, False),
+        32: (1, False, False),
+        64: (4, False, False),
+        128: (4, False, False),
+    },
+    16384: {
+        1: (1, True, False),
+        4: (8, False, False),
+        8: (2, True, False),
+        16: (2, True, False),
+        32: (4, False, False),
+        64: (1, False, False),
+        128: (4, False, False),
+    },
+    32768: {
+        1: (1, False, False),
+        4: (1, True, False),
+        8: (1, True, False),
+        16: (1, False, False),
+        32: (2, False, False),
+        64: (4, False, False),
+        128: (4, False, False),
+    },
+}
 
 
 def nearest_token_bucket(tokens: int) -> int:
@@ -165,16 +235,90 @@ def _select_stage1(bucket: int, fixed_slot: bool, mtpr: int) -> Stage1Config:
             use_tile_resource=bucket >= 16384,
             b_nt=0,
         )
+
+    oversized_capacity = not fixed_slot and mtpr > bucket
+    if oversized_capacity:
+        dispatch_cu = _OVERSIZED_SMALL_DISPATCH_CU.get(mtpr, {}).get(bucket)
+        if dispatch_cu is not None:
+            config = replace(config, num_dispatch_cu=dispatch_cu)
+        elif bucket == 32:
+            config = replace(config, num_dispatch_cu=192 if mtpr in (4096, 16384) else 64)
+        elif bucket == 64:
+            config = replace(config, num_dispatch_cu=160)
+        elif bucket == 128:
+            config = replace(config, num_dispatch_cu=192)
+        elif bucket == 256 and mtpr >= 16384:
+            config = replace(config, sort_block_m=32, grid_mult=1, num_dispatch_cu=64)
+        elif bucket == 512:
+            config = replace(
+                config,
+                sort_block_m=64,
+                grid_mult=2 if mtpr >= 32768 else 1,
+                num_dispatch_cu=64,
+                use_tile_resource=mtpr != 4096,
+                b_nt=3 if mtpr >= 16384 else 0,
+            )
+        elif bucket == 1024:
+            config = replace(
+                config,
+                sort_block_m=64,
+                grid_mult=1,
+                num_dispatch_cu=64,
+                use_tile_resource=True,
+                b_nt=0,
+            )
+        elif bucket == 2048 and mtpr >= 4096:
+            config = replace(
+                config,
+                grid_mult=1 if mtpr >= 16384 else 2,
+                num_dispatch_cu=64,
+                async_a_copy=mtpr != 8192,
+            )
+        elif bucket == 4096:
+            config = replace(
+                config,
+                grid_mult=1 if mtpr == 8192 else 2,
+                num_dispatch_cu=64 if mtpr == 8192 else 32,
+                use_tile_resource=mtpr >= 16384,
+            )
+
+    if mtpr >= 16384:
+        config = replace(config, use_tile_resource=True)
     external_grouping = not fixed_slot and mtpr >= 2048
+    work_shards = 4 if mtpr >= 8192 or (bucket == 1024 and mtpr >= 4096) else 8
+    external_counting = external_grouping and (mtpr >= 8192 or (bucket == 1024 and mtpr >= 4096))
+    if oversized_capacity and bucket <= 128:
+        work_shards, external_grouping, external_counting = _OVERSIZED_SMALL_PROTOCOL.get(mtpr, {}).get(
+            bucket, (work_shards, external_grouping, external_counting)
+        )
+    if bucket == 2048 and mtpr == 8192:
+        work_shards = 2
+    elif bucket == 4096 and mtpr == 16384:
+        work_shards, external_counting = 1, False
     return replace(
         config,
-        work_shards=4 if mtpr >= 8192 else 8,
+        work_shards=work_shards,
         external_grouping=external_grouping,
-        external_counting=external_grouping and mtpr >= 8192,
+        external_counting=external_counting,
     )
 
 
-def _select_stage2(bucket: int, fixed_slot: bool) -> Stage2Config:
+def _select_stage2(bucket: int, fixed_slot: bool, mtpr: int, sort_block_m: int) -> Stage2Config:
+    if not fixed_slot and mtpr > bucket:
+        if bucket == 128 and mtpr == 4096:
+            persist_cu = 128
+        elif bucket == 256 and mtpr >= 32768:
+            persist_cu = 256
+        else:
+            persist_cu = {1024: 224, 2048: 256, 8192: 256, 16384: 192}.get(bucket, 240)
+        return Stage2Config(
+            block_m=64 if sort_block_m == 128 else 32,
+            block_n=128 if bucket == 256 and sort_block_m == 64 else 256,
+            persist=True,
+            persist_cu=persist_cu,
+            use_nt=bucket <= 128,
+            persist_strided=bucket in (512, 1024, 2048),
+        )
     block_m = 64 if bucket >= 4096 else 32
     block_n = 256 if bucket in (1, 4, 64) or bucket >= 1024 or (not fixed_slot and bucket < 128) else 128
     persist = bucket >= 128
@@ -192,10 +336,12 @@ def _select_stage2(bucket: int, fixed_slot: bool) -> Stage2Config:
 
 
 @lru_cache(maxsize=None)
-def _select_bucket_config(bucket: int, mtpr: int, p2p_quant: str) -> MegaMoEConfig:
+def _select_bucket_config(bucket: int, mtpr: int) -> MegaMoEConfig:
     fixed_slot = mtpr <= FIXED_SLOT_MAX_MTPR
     stage1 = _select_stage1(bucket, fixed_slot, mtpr)
-    stage2 = _select_stage2(bucket, fixed_slot)
+    stage2 = _select_stage2(bucket, fixed_slot, mtpr, stage1.sort_block_m)
+    # MTPR is rank-invariant; local token counts need not be.
+    p2p_quant = "fp8_blockwise_1x32" if mtpr > P2P_FP8_MIN_MTPR else "none"
     return MegaMoEConfig(stage1=stage1, stage2=stage2, p2p_quant=p2p_quant)
 
 
@@ -208,5 +354,4 @@ def select_mega_moe_config(tokens: int, mtpr: int) -> MegaMoEConfig:
     fixed_slot = mtpr <= FIXED_SLOT_MAX_MTPR
     if fixed_slot and bucket not in _FIXED_GEOMETRY:
         raise ValueError(f"fixed-slot does not support token bucket {bucket}")
-    p2p_quant = "fp8_blockwise_1x32" if tokens > P2P_FP8_MIN_TOKENS else "none"
-    return _select_bucket_config(bucket, mtpr, p2p_quant)
+    return _select_bucket_config(bucket, mtpr)
diff --git a/kernels/mega_moe/mega_moe_stage1.py b/kernels/mega_moe/mega_moe_stage1.py
index 9c072bd89..e804d0550 100644
--- a/kernels/mega_moe/mega_moe_stage1.py
+++ b/kernels/mega_moe/mega_moe_stage1.py
@@ -9,10 +9,10 @@
 import flydsl.compiler as flyc
 import flydsl.expr as fx
 from flydsl.expr import const_expr, range_constexpr
+from flydsl.expr.typing import T
 from flydsl.expr.typing import Vector as Vec
 from flydsl.runtime.device import get_rocm_arch
 from kernels.comm import communication_ops_utils as comm_ops
-from kernels.common import buffer_ops as _buffer_ops
 from kernels.common.tensor_shim import _run_compiled
 from kernels.gemm.fp8_gemm_utils import ceildiv
 
@@ -25,6 +25,7 @@
     emit_dispatch_plan,
 )
 from .gemm1 import _LdsF32View, build_fused_gemm1
+from .gemm_util import _buffer_load, _buffer_store, _make_buffer, _make_buffer_from_addr
 
 _SC0_CACHE = 1
 _BUFFER_OFFSET_ABI_BYTES = 1 << 32
@@ -54,8 +55,7 @@ def compile_mega_moe_stage1(
     fuse_cap: int, fuse_mtpr: int, fuse_scale_dim: int, fixed_slot_dispatch: bool, sort_block_m: int = 32,
     tile_n: int = 256, tile_k: int = 256, num_waves: int = 4, grid_mult: int = 8,
     pipe_weights: bool = True, mfma_amajor: bool = False, swizzle_a: bool = True,
-    async_a_copy: bool = False, active_expert_producer: bool = False,
-    cooperative_payload_copy: bool = False, use_tile_resource: bool = True,
+    async_a_copy: bool = False, use_tile_resource: bool = True,
     waves_per_eu_hint: int = 2, num_cu: int = 256, num_dispatch_cu: int = 32, b_nt: int = -1,
     work_shards: int | None = None, external_grouping: bool | None = None,
     external_counting: bool | None = None,
@@ -106,8 +106,12 @@ def compile_mega_moe_stage1(
 
     fz_npes, fz_epr, fz_k = int(fuse_npes), int(experts_per_rank), int(fuse_topk)
     fz_cap, fz_mtpr, fz_rank = int(fuse_cap), int(fuse_mtpr), int(rank)
+    if fz_npes * fz_mtpr > 1 << 24:
+        raise ValueError("MegaMoE v2 source-token encoding exceeds 24 bits")
+    if fz_k > 1 << 8:
+        raise ValueError("MegaMoE v2 top-k slot encoding exceeds 8 bits")
     if external_grouping is None:
-        external_grouping = fz_mtpr >= 2048 and fz_npes == 8 and fz_epr == 48 and not active_expert_producer
+        external_grouping = fz_mtpr >= 2048 and fz_npes == 8 and fz_epr == 48
     if external_counting is None:
         external_counting = external_grouping and fz_mtpr >= 8192
     assert not external_counting or external_grouping
@@ -122,6 +126,8 @@ def compile_mega_moe_stage1(
     fz_n_i32, fz_nbytes = model_dim // 4, model_dim
     fz_scale_bytes = int(fuse_scale_dim)
     fz_scale_n_i32 = (fz_scale_bytes + 3) // 4 if fz_scale_bytes > 0 else 0
+    if direct_fixed_slot and fz_scale_n_i32 > 64:
+        raise ValueError("direct fixed-slot dispatch supports at most 64 packed scale columns")
     fz_enable_scales = fz_scale_bytes > 0
     fz_safe_end_i32 = (fz_n_i32 // 512) * 512
     _validate_dispatch_capacity(
@@ -138,7 +144,7 @@ class SharedStorage:
         f"megamoe_stage1_{dispatch_path}_t{sort_block_m}x{tile_n}x{tile_k}"
         f"_w{NUM_WAVES}_gm{grid_mult}"
         f"_dcu{dispatch_blocks}_pw{int(pipe_weights)}ma{int(mfma_amajor)}sw{int(swizzle_a)}"
-        f"aa{int(async_a_copy)}_aep{int(active_expert_producer)}cpc{int(cooperative_payload_copy)}"
+        f"aa{int(async_a_copy)}"
         f"_tr{int(use_tile_resource)}wpe{waves_per_eu_hint}_bnt{b_cache_modifier}_ws{WORK_SHARDS}"
     )
 
@@ -155,18 +161,21 @@ def kernel(
         a_buf = lds.pool
         a_scale_lds = lds.A_scale
         c_tile = _LdsF32View(fx.recast_iter(fx.Float32, lds.pool.ptr))
-        disp_rsrc = _buffer_ops.create_buffer_resource_from_addr(addr_disp)
-        parity_rsrc = _buffer_ops.create_buffer_resource_from_addr(addr_parity)
-        expected_rsrc = _buffer_ops.create_buffer_resource_from_addr(addr_expected)
+        disp_rsrc = _make_buffer_from_addr(addr_disp, fx.Int64)
+        parity_rsrc = _make_buffer_from_addr(addr_parity, fx.Int32)
+        expected_rsrc = _make_buffer_from_addr(addr_expected, fx.Int32)
+
         def _disp_ptr(slot):
-            return _buffer_ops.buffer_load(disp_rsrc, fx.Int32(int(slot)), vec_width=1, dtype=fx.Int64)
+            return _buffer_load(disp_rsrc, fx.Int32(int(slot)), fx.Int64)
 
         a_entry_count = _disp_ptr(DispatchSlot.ENTRY_COUNT)
         a_epoch_gate = _disp_ptr(DispatchSlot.EPOCH_GATE)
         a_pair_order_ready = _disp_ptr(DispatchSlot.PAIR_ORDER_READY)
         a_work_head = _disp_ptr(DispatchSlot.WORK_HEAD)
         a_work_tail = _disp_ptr(DispatchSlot.WORK_TAIL)
-        a_group_done = _disp_ptr(DispatchSlot.ACTIVE_COUNT)
+        a_group_done = _disp_ptr(DispatchSlot.GROUP_DONE)
+        a_launch_ready = _disp_ptr(DispatchSlot.LAUNCH_READY)
+        p_launch_ready = _disp_ptr(DispatchSlot.P2P_LAUNCH_READY)
 
         ticket_scratch = fx.recast_iter(fx.Int64, a_buf.ptr)
         ticket_view = fx.make_view(ticket_scratch, fx.make_layout(1, 1))
@@ -186,24 +195,43 @@ def _disp_ptr(slot):
         producer_slot = ticket - fx.Int32(1)
 
         if compact_owner:
+            next_parity_lane = fx.Int32(0)
+            launch_epoch_lane = fx.Int32(0)
             if tid == fx.Int32(0):
-                old_parity = _buffer_ops.buffer_load(parity_rsrc, fx.Int32(0), vec_width=1, dtype=fx.Int32)
-                next_parity = old_parity ^ fx.Int32(1)
-                previous_expected = _buffer_ops.buffer_load(expected_rsrc, next_parity, vec_width=1, dtype=fx.Int32)
+                old_parity = _buffer_load(parity_rsrc, fx.Int32(0), fx.Int32)
+                next_parity_lane = old_parity ^ fx.Int32(1)
+                previous_expected = _buffer_load(expected_rsrc, next_parity_lane, fx.Int32)
                 next_expected = previous_expected + fx.Int32(fz_npes)
-                _buffer_ops.buffer_store(next_expected, expected_rsrc, next_parity)
-                fx.rocdl.s_waitcnt(0)
-                comm_ops.fence_agent_release()
-                _buffer_ops.buffer_store(next_parity, parity_rsrc, fx.Int32(0))
-                work_head_rsrc = _buffer_ops.create_buffer_resource_from_addr(a_work_head)
+                _buffer_store(expected_rsrc, next_parity_lane, next_expected, fx.Int32)
+                launch_epoch_lane = (
+                    (next_expected // fx.Int32(fz_npes)) * fx.Int32(2) - next_parity_lane
+                )
+            next_parity = fx.Int32(fx.rocdl.readfirstlane(T.i32, next_parity_lane))
+            launch_epoch = fx.Int32(fx.rocdl.readfirstlane(T.i32, launch_epoch_lane))
+            if tid < fx.Int32(fz_npes):
+                peer = (tid + fx.Int32(fz_rank)) % fx.Int32(fz_npes)
+                comm_ops.fence_system_release()
+                launch_ready_table = _make_buffer_from_addr(p_launch_ready, fx.Int64)
+                remote_launch_ready = _buffer_load(launch_ready_table, peer, fx.Int64)
+                comm_ops.store_i32_system(remote_launch_ready, fx.Int32(fz_rank), launch_epoch)
+                mori_shmem.int32_wait_until_greater_than(
+                    a_launch_ready + fx.Int64(peer) * fx.Int64(4), launch_epoch - fx.Int32(1)
+                )
+                comm_ops.fence_system_acquire()
+            if tid == fx.Int32(0):
+                work_head_rsrc = _make_buffer_from_addr(a_work_head, fx.Int32)
                 for shard in range_constexpr(8):
-                    _buffer_ops.buffer_store(fx.Int32(0), work_head_rsrc, fx.Int32(shard * 16))
-                _buffer_ops.buffer_store(fx.Int32(0), _buffer_ops.create_buffer_resource_from_addr(a_work_tail),
-                                         fx.Int32(0))
+                    _buffer_store(work_head_rsrc, fx.Int32(shard * 16), fx.Int32(0), fx.Int32)
+                _buffer_store(_make_buffer_from_addr(a_work_tail, fx.Int32), fx.Int32(0), fx.Int32(0), fx.Int32)
                 if const_expr(external_grouping or direct_fixed_slot):
-                    group_done_rsrc = _buffer_ops.create_buffer_resource_from_addr(a_group_done)
+                    group_done_rsrc = _make_buffer_from_addr(a_group_done, fx.Int32)
                     for destination in range_constexpr(fz_npes if direct_fixed_slot else 1):
-                        _buffer_ops.buffer_store(fx.Int32(0), group_done_rsrc, fx.Int32(destination))
+                        _buffer_store(group_done_rsrc, fx.Int32(destination), fx.Int32(0), fx.Int32)
+            fx.barrier()
+            if tid == fx.Int32(0):
+                fx.rocdl.s_waitcnt(0)
+                comm_ops.fence_agent_release()
+                _buffer_store(parity_rsrc, fx.Int32(0), next_parity, fx.Int32)
                 fx.rocdl.s_waitcnt(0)
                 comm_ops.fence_agent_release()
                 comm_ops.store_i32_system(gate_addr, fx.Int32(0), gate_epoch)
@@ -215,12 +243,8 @@ def _disp_ptr(slot):
                 comm_ops.fence_agent_acquire()
             fx.barrier()
 
-        payload_parity = _buffer_ops.buffer_load(
-            parity_rsrc, fx.Int32(0), vec_width=1, dtype=fx.Int32, cache_modifier=_SC0_CACHE
-        )
-        payload_expected = _buffer_ops.buffer_load(
-            expected_rsrc, payload_parity, vec_width=1, dtype=fx.Int32, cache_modifier=_SC0_CACHE
-        )
+        payload_parity = _buffer_load(parity_rsrc, fx.Int32(0), fx.Int32, cache_modifier=_SC0_CACHE)
+        payload_expected = _buffer_load(expected_rsrc, payload_parity, fx.Int32, cache_modifier=_SC0_CACHE)
 
         if compact_owner:
             if const_expr(not direct_fixed_slot):
@@ -228,8 +252,8 @@ def _disp_ptr(slot):
                     num_waves=NUM_WAVES, fz_npes=fz_npes, fz_epr=fz_epr, fz_k=fz_k, fz_mtpr=fz_mtpr,
                     fz_rank=fz_rank, fz_tile_m=fz_tile_m, fz_total_experts=fz_total_experts, addr_disp=addr_disp,
                     i32_cur_tok=i32_cur_tok, addr_in_idx=addr_in_idx, parity=payload_parity,
-                    expected=payload_expected, active_expert_producer=active_expert_producer,
-                    external_grouping=external_grouping, external_counting=external_counting,
+                    expected=payload_expected, external_grouping=external_grouping,
+                    external_counting=external_counting,
                     dispatch_blocks=dispatch_blocks,
                 )
 
@@ -265,8 +289,6 @@ def _disp_ptr(slot):
                     fz_enable_scales=fz_enable_scales, addr_disp=addr_disp, addr_in_tok=addr_in_tok,
                     addr_in_wts=addr_in_wts, addr_in_sc=addr_in_sc, dispatch_blocks=dispatch_blocks,
                     producer_slot=producer_slot, parity=payload_parity, expected=payload_expected,
-                    active_expert_producer=active_expert_producer,
-                    cooperative_payload_copy=cooperative_payload_copy,
                 )
         if const_expr(direct_fixed_slot):
             if compact_owner:
@@ -276,35 +298,31 @@ def _disp_ptr(slot):
                     expected=payload_expected,
                 )
         else:
-            payload_table = _buffer_ops.buffer_load(
-                disp_rsrc, fx.Int32(int(DispatchSlot.P2P_PAYLOAD_READY)), vec_width=1, dtype=fx.Int64)
-            addr_payload_ready = _buffer_ops.buffer_load(
-                _buffer_ops.create_buffer_resource_from_addr(payload_table), fx.Int32(fz_rank), vec_width=1,
-                dtype=fx.Int64)
+            payload_table = _buffer_load(disp_rsrc, fx.Int32(int(DispatchSlot.P2P_PAYLOAD_READY)), fx.Int64)
+            addr_payload_ready = _buffer_load(
+                _make_buffer_from_addr(payload_table, fx.Int64), fx.Int32(fz_rank), fx.Int64
+            )
         wave_id = fx.thread_idx.x // 64
 
-        x_rsrc = _buffer_ops.create_buffer_resource(x, max_size=True)
-        x_base_addr = fx.Int64(_buffer_ops.extract_base_index(x, address_space=1))
-        w_rsrc = _buffer_ops.create_buffer_resource(w, max_size=True)
-        sx_rsrc = _buffer_ops.create_buffer_resource(scale_x, max_size=True)
-        sw_rsrc = _buffer_ops.create_buffer_resource(scale_w, max_size=True)
-        trb_rsrc = _buffer_ops.create_buffer_resource(sorted_token_ids, max_size=True)
-        expert_rsrc = _buffer_ops.create_buffer_resource(expert_ids, max_size=True)
-        nv_rsrc = _buffer_ops.create_buffer_resource(num_valid_ids, max_size=True)
+        w_rsrc = _make_buffer(w, fx.Int32, 4)
+        sx_rsrc = _make_buffer(scale_x, fx.Int32, 4)
+        sw_rsrc = _make_buffer(scale_w, fx.Int32)
+        trb_rsrc = _make_buffer(sorted_token_ids, fx.Int32)
+        expert_rsrc = _make_buffer(expert_ids, fx.Int32)
+        nv_rsrc = _make_buffer(num_valid_ids, fx.Int32)
         scale_cols = (inter_dim // 32 + 7) // 8 * 8
         os_nbytes = tokens * fx.Int32(scale_cols) + fx.Int32(8192)
-        out_base_addr = fx.Int64(_buffer_ops.extract_base_index(out, address_space=1))
         if const_expr(use_tile_resource):
-            out_rsrc = _buffer_ops.create_buffer_resource(out, max_size=True)
+            out_rsrc = None
         else:
             out_nbytes = tokens * fx.Int32(inter_dim)
-            out_rsrc = _buffer_ops.create_buffer_resource(out, max_size=False, num_records_bytes=out_nbytes)
-        os_rsrc = _buffer_ops.create_buffer_resource(out_scale, max_size=False, num_records_bytes=os_nbytes)
+            out_rsrc = _make_buffer(out, fx.Int16, max_size=False, num_records_bytes=out_nbytes)
+        os_rsrc = _make_buffer(out_scale, fx.Int8, max_size=False, num_records_bytes=os_nbytes)
 
         expert_of_flat, _do_scheduled_tile = build_fused_gemm1(
-            x_rsrc=x_rsrc, x_base_addr=x_base_addr, x_tensor=x, w_rsrc=w_rsrc,
+            x_tensor=x, w_rsrc=w_rsrc,
             sw_rsrc=sw_rsrc, sx_rsrc=sx_rsrc, out_rsrc=out_rsrc, os_rsrc=os_rsrc,
-            trb_rsrc=trb_rsrc, expert_rsrc=expert_rsrc, out_base_addr=out_base_addr,
+            trb_rsrc=trb_rsrc, expert_rsrc=expert_rsrc, out_tensor=out,
             a_buf=a_buf, a_scale_lds=a_scale_lds, c_tile=c_tile,
             model_dim=model_dim, inter_dim=inter_dim, sort_block_m=sort_block_m,
             tile_n=tile_n, num_waves=NUM_WAVES, n_per_wave=n_per_wave, wave_id=wave_id,
@@ -316,15 +334,14 @@ def _disp_ptr(slot):
         )
 
         if tid == fx.Int32(0):
-            local_plan_ready = _buffer_ops.buffer_load(
-                disp_rsrc, fx.Int32(int(DispatchSlot.PLAN_READY)), vec_width=1, dtype=fx.Int64)
+            local_plan_ready = _buffer_load(disp_rsrc, fx.Int32(int(DispatchSlot.PLAN_READY)), fx.Int64)
             ready_index = payload_parity * fx.Int32(fz_npes) + fx.Int32(fz_rank)
             mori_shmem.int32_wait_until_equals(
                 local_plan_ready + fx.Int64(ready_index) * fx.Int64(4), payload_expected)
             comm_ops.fence_agent_acquire()
         fx.barrier()
 
-        num_valid = _buffer_ops.buffer_load(nv_rsrc, fx.Int32(0), vec_width=1, dtype=fx.Int32)
+        num_valid = _buffer_load(nv_rsrc, fx.Int32(0), fx.Int32)
         num_m_tiles = ceildiv(num_valid, fx.Int32(sort_block_m))
         total_work = num_m_tiles * fx.Int32(N_TILES)
 
@@ -387,8 +404,8 @@ def run_mega_moe_stage1(out, x, w, scale_x, scale_w, sorted_token_ids, expert_id
     addr_parity, addr_expected, stream, *, model_dim, inter_dim, rank, experts_per_rank, fuse_npes,
     fuse_topk, fuse_cap, fuse_mtpr, fuse_scale_dim, fixed_slot_dispatch, num_cu,
     sort_block_m=32, tile_n=256, tile_k=256, num_waves=4, grid_mult=4, pipe_weights=True,
-    mfma_amajor=False, swizzle_a=True, async_a_copy=False, active_expert_producer=False,
-    cooperative_payload_copy=False, num_dispatch_cu=32, use_tile_resource=True, waves_per_eu_hint=2,
+    mfma_amajor=False, swizzle_a=True, async_a_copy=False, num_dispatch_cu=32,
+    use_tile_resource=True, waves_per_eu_hint=2,
     b_nt=-1, work_shards=None, external_grouping=None, external_counting=None):
     launch = compile_mega_moe_stage1(
         model_dim=model_dim, inter_dim=inter_dim, rank=rank, experts_per_rank=experts_per_rank,
@@ -396,8 +413,7 @@ def run_mega_moe_stage1(out, x, w, scale_x, scale_w, sorted_token_ids, expert_id
         fuse_scale_dim=fuse_scale_dim, fixed_slot_dispatch=fixed_slot_dispatch,
         sort_block_m=sort_block_m, tile_n=tile_n, tile_k=tile_k, num_waves=num_waves,
         grid_mult=grid_mult, pipe_weights=pipe_weights, mfma_amajor=mfma_amajor, swizzle_a=swizzle_a,
-        async_a_copy=async_a_copy, active_expert_producer=active_expert_producer,
-        cooperative_payload_copy=cooperative_payload_copy, use_tile_resource=use_tile_resource,
+        async_a_copy=async_a_copy, use_tile_resource=use_tile_resource,
         waves_per_eu_hint=waves_per_eu_hint, num_cu=num_cu, num_dispatch_cu=num_dispatch_cu,
         b_nt=b_nt, work_shards=work_shards, external_grouping=external_grouping,
         external_counting=external_counting,
diff --git a/kernels/mega_moe/mega_moe_stage2.py b/kernels/mega_moe/mega_moe_stage2.py
index 0c3db3614..a9dd8c0b2 100644
--- a/kernels/mega_moe/mega_moe_stage2.py
+++ b/kernels/mega_moe/mega_moe_stage2.py
@@ -25,6 +25,8 @@
     kStages,
 )
 
+_BUFFER_OFFSET_ABI_BYTES = 1 << 31
+
 
 @flyc.jit
 def _fp8_scale_for_leader(is_leader, local_max):
@@ -291,12 +293,12 @@ def compile_mega_moe_stage2(*, model_dim: int, inter_dim: int, experts: int, top
     log2_max_tok = max_tok.bit_length() - 1
     mask_max_tok = max_tok - 1
     N_OUT = model_dim
-    # scatter path uses the f32 cshuffle slab (no bf16 LDS); knobs env-defaulted (spart 402 etc.).
+    # The scatter path uses the f32 CShuffle slab rather than BF16 LDS.
     g2_bhoist, g2_ascale_pf, g2_spart, g2_group_num, g2_m01, _g2_bf16_lds = _resolve_g2_knobs(
         g2_bhoist, g2_ascale_pf, g2_spart, False, False
     )
     is_f8 = a_dtype == "fp8"
-    aStages = 2 if g2_bf16_lds else 3
+    aStages = kStages + 1
     KH_TILE_A = BK // (1 if is_f8 else 2)
     compute_lds_bytes = _stage2_lds_bytes(BM, BN, BK, a_dtype, aStages, g2_bf16_lds)
     lds_packed_off = compute_lds_bytes
@@ -306,6 +308,8 @@ def compile_mega_moe_stage2(*, model_dim: int, inter_dim: int, experts: int, top
     _recv_cap = npes * max_tok if recv_cap is None else int(recv_cap)
     _row_nbytes = N_OUT + N_OUT // 32 if p2p_quant_type == "fp8_blockwise_1x32" else N_OUT * 2
     _comb_inp_nbytes = max_tok * topk * _row_nbytes if comb_inp_nbytes is None else int(comb_inp_nbytes)
+    if not 0 < _comb_inp_nbytes < _BUFFER_OFFSET_ABI_BYTES:
+        raise ValueError("MegaMoE v2 stage2 P2P buffer exceeds the 32-bit buffer-resource ABI")
     _expert_offset = rank * experts
 
     @fx.struct
diff --git a/tests/kernels/test_mega_moe_v2.py b/tests/kernels/test_mega_moe_v2.py
index c168ad43c..fd6c0c67d 100644
--- a/tests/kernels/test_mega_moe_v2.py
+++ b/tests/kernels/test_mega_moe_v2.py
@@ -750,6 +750,7 @@ def _atom_body():
         scale_type_size=1,
         enable_std_moe=False,
     )
+    assert cfg_fp8.is_fp4 == _is_fp4
     dcf = FlyDSLDispatchCombineIntraNodeOp(cfg_fp8)
     torch.cuda.synchronize()
     ms.shmem_barrier_all()
diff --git a/tests/unit/test_mega_moe_config.py b/tests/unit/test_mega_moe_config.py
index 69d3bc7b9..33ca8ba07 100644
--- a/tests/unit/test_mega_moe_config.py
+++ b/tests/unit/test_mega_moe_config.py
@@ -3,7 +3,11 @@
 
 import pytest
 
-from kernels.mega_moe.mega_moe_config import nearest_token_bucket, select_mega_moe_config
+from kernels.mega_moe.mega_moe_config import (
+    TOKEN_BUCKETS,
+    nearest_token_bucket,
+    select_mega_moe_config,
+)
 
 _STANDARD_PROFILES = {
     1: (32, 256, 4, 1, 64, 0, 1, 2, 32, 256, 0, 0, 0, "none"),
@@ -58,7 +62,6 @@ def test_standard_profiles_match_tuned_artifacts(tokens, expected):
     assert stage1.external_grouping == (tokens >= 2048)
     assert stage1.external_counting == (tokens >= 8192)
     assert stage1.pipe_weights and stage1.swizzle_a
-    assert not stage1.active_expert_producer and not stage1.cooperative_payload_copy
     assert stage2.use_nt == (tokens <= 128)
     assert stage2.b_hoist and stage2.ascale_prefetch
     assert stage2.spatial_partition == 402 and not stage2.bf16_lds
@@ -66,7 +69,15 @@ def test_standard_profiles_match_tuned_artifacts(tokens, expected):
 
 @pytest.mark.parametrize(
     "tokens,bucket",
-    [(2, 1), (3, 4), (6, 8), (16300, 16384), (16400, 16384), (24576, 32768), (65536, 32768)],
+    [
+        (2, 1),
+        (3, 4),
+        (6, 8),
+        (16300, 16384),
+        (16400, 16384),
+        (24576, 32768),
+        (65536, 32768),
+    ],
 )
 def test_nearest_token_bucket_prefers_larger_on_ties(tokens, bucket):
     assert nearest_token_bucket(tokens) == bucket
@@ -76,13 +87,106 @@ def test_mtpr_selects_fixed_or_compact_configs():
     fixed = select_mega_moe_config(128, 128)
     compact = select_mega_moe_config(128, 8192)
 
-    assert (fixed.stage1.tile_n, fixed.stage1.num_waves, fixed.stage1.num_dispatch_cu) == (128, 4, 224)
-    assert (compact.stage1.tile_n, compact.stage1.num_waves, compact.stage1.num_dispatch_cu) == (512, 8, 128)
+    assert (
+        fixed.stage1.tile_n,
+        fixed.stage1.num_waves,
+        fixed.stage1.num_dispatch_cu,
+    ) == (128, 4, 224)
+    assert (
+        compact.stage1.tile_n,
+        compact.stage1.num_waves,
+        compact.stage1.num_dispatch_cu,
+    ) == (512, 8, 192)
     for tokens in (8, 16, 32):
         assert select_mega_moe_config(tokens, 128).stage2.block_n == 128
         assert select_mega_moe_config(tokens, 8192).stage2.block_n == 256
 
 
+@pytest.mark.parametrize(
+    "tokens,mtpr,stage1,stage2",
+    [
+        (8, 8192, (32, 1, 192, False, 3, 1), (32, 256, 240, False)),
+        (256, 8192, (64, 1, 160, True, 3, 4), (32, 128, 240, False)),
+        (512, 8192, (64, 1, 64, True, 0, 4), (32, 256, 240, True)),
+        (1024, 32768, (64, 1, 64, True, 0, 4), (32, 256, 224, True)),
+        (2048, 16384, (64, 1, 64, True, 0, 4), (32, 256, 256, True)),
+        (4096, 8192, (128, 1, 64, False, 0, 4), (64, 256, 240, False)),
+    ],
+)
+def test_oversized_capacity_profiles_match_tuned_rules(tokens, mtpr, stage1, stage2):
+    config = select_mega_moe_config(tokens, mtpr)
+    s1 = config.stage1
+    s2 = config.stage2
+
+    assert (
+        s1.sort_block_m,
+        s1.grid_mult,
+        s1.num_dispatch_cu,
+        s1.use_tile_resource,
+        s1.b_nt,
+        s1.work_shards,
+    ) == stage1
+    assert (s2.block_m, s2.block_n, s2.persist_cu, s2.persist_strided) == stage2
+    assert s2.persist
+
+
+@pytest.mark.parametrize("mtpr", [16384, 32768])
+@pytest.mark.parametrize("tokens", [1, 128, 512, 1024, 2048, 4096])
+def test_large_capacity_uses_safe_tile_resource_addressing(tokens, mtpr):
+    assert select_mega_moe_config(tokens, mtpr).stage1.use_tile_resource
+
+
+@pytest.mark.parametrize("mtpr", [2048, 4096, 8192, 16384, 32768])
+def test_requested_oversized_capacity_matrix_is_valid(mtpr):
+    for tokens in (bucket for bucket in TOKEN_BUCKETS if bucket <= mtpr // 2):
+        config = select_mega_moe_config(tokens, mtpr)
+
+        assert config.stage2.block_m <= config.stage1.sort_block_m
+        assert config.stage1.sort_block_m % config.stage2.block_m == 0
+        assert config.p2p_quant == "fp8_blockwise_1x32"
+
+
+@pytest.mark.parametrize(
+    "tokens,mtpr,expected",
+    [
+        (4, 2048, (128, 8, False, False)),
+        (1, 4096, (224, 1, False, False)),
+        (64, 4096, (160, 1, False, False)),
+        (1, 8192, (224, 1, False, False)),
+        (8, 8192, (192, 1, False, False)),
+        (64, 8192, (160, 4, False, False)),
+        (4, 16384, (224, 8, False, False)),
+        (32, 16384, (192, 4, False, False)),
+        (16, 32768, (64, 1, False, False)),
+        (32, 32768, (64, 2, False, False)),
+    ],
+)
+def test_oversized_small_tail_matches_tuned_rules(tokens, mtpr, expected):
+    stage1 = select_mega_moe_config(tokens, mtpr).stage1
+
+    assert (
+        stage1.num_dispatch_cu,
+        stage1.work_shards,
+        stage1.external_grouping,
+        stage1.external_counting,
+    ) == expected
+
+
+@pytest.mark.parametrize(
+    "mtpr,expected",
+    [
+        (128, "none"),
+        (1024, "none"),
+        (2048, "fp8_blockwise_1x32"),
+        (8192, "fp8_blockwise_1x32"),
+    ],
+)
+def test_p2p_quant_is_rank_invariant_for_an_mtpr(mtpr, expected):
+    configs = [select_mega_moe_config(tokens, mtpr) for tokens in TOKEN_BUCKETS if tokens <= mtpr]
+
+    assert {config.p2p_quant for config in configs} == {expected}
+
+
 def test_nearby_tokens_share_the_bucket_config():
     assert select_mega_moe_config(500, 512) is select_mega_moe_config(512, 512)
 

From 9513a2ab6c08b39bef0e0038c370864480235b85 Mon Sep 17 00:00:00 2001
From: guoliang <ghu@amd.com>
Date: Tue, 4 Aug 2026 03:26:29 +0000
Subject: [PATCH 2/4] fix deadlock bug

---
 kernels/mega_moe/dispatch.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/kernels/mega_moe/dispatch.py b/kernels/mega_moe/dispatch.py
index f58a26852..2d8f259a8 100644
--- a/kernels/mega_moe/dispatch.py
+++ b/kernels/mega_moe/dispatch.py
@@ -467,7 +467,7 @@ def dp(i):
             mori_shmem.int32_wait_until_equals(a_group_done, fx.Int32(dispatch_blocks))
             comm_ops.fence_agent_acquire()
         comm_ops.fence_agent_release()
-        buffer_ops.buffer_store(expected, crfa(a_pair_order_ready), parity)
+        comm_ops.store_i32_system(a_pair_order_ready, parity, expected)
 
 
 # fmt: off

From 52f335e00ea57a1d84d2cc6d5a73794a61e419b6 Mon Sep 17 00:00:00 2001
From: guoliang <ghu@amd.com>
Date: Tue, 4 Aug 2026 05:31:17 +0000
Subject: [PATCH 3/4] ci: rerun


From 9407d8e1ec9fff6f1ac1d5c159ae1bf2d08fe93b Mon Sep 17 00:00:00 2001
From: guoliang <ghu@amd.com>
Date: Tue, 4 Aug 2026 07:31:20 +0000
Subject: [PATCH 4/4] fix gemm1 clamp

---
 kernels/mega_moe/gemm1.py           | 10 ++++---
 kernels/mega_moe/gemm_util.py       | 14 +++++++--
 kernels/mega_moe/mega_moe.py        |  7 +++--
 kernels/mega_moe/mega_moe_stage1.py |  9 ++++--
 tests/kernels/test_mega_moe_v2.py   | 44 ++++++++++++++++++++---------
 5 files changed, 58 insertions(+), 26 deletions(-)

diff --git a/kernels/mega_moe/gemm1.py b/kernels/mega_moe/gemm1.py
index c4df0f630..797602ac4 100644
--- a/kernels/mega_moe/gemm1.py
+++ b/kernels/mega_moe/gemm1.py
@@ -270,7 +270,7 @@ def build_fused_gemm1(*, x_tensor, w_rsrc, sw_rsrc, sx_rsrc,
     model_dim, inter_dim, sort_block_m, tile_n, num_waves, n_per_wave, wave_id,
     m_repeat, num_acc_n, a_k_step_bytes, total_threads, k_iters, a_lds_i32, n_tiles,
     expert_offset, b_cache_modifier, swizzle_a, pipe_weights, mfma_amajor, async_a_copy,
-    use_tile_resource):
+    use_tile_resource, swiglu_limit=0.0):
     # fmt: on
     """Build the GEMM1 atoms and return its expert resolver and tile runner."""
     sched = TileScheduler(
@@ -304,7 +304,7 @@ def build_fused_gemm1(*, x_tensor, w_rsrc, sw_rsrc, sx_rsrc,
     # fmt: off
     epi = SiluQuantEpilogue(out_rsrc=out_rsrc, out_scale_rsrc=os_rsrc, sorted_rsrc=trb_rsrc, tokens=0,
         inter_dim=inter_dim, m_repeat=m_repeat, num_acc_n=num_acc_n, sort_block_m=sort_block_m, tile_n=tile_n,
-        num_waves=num_waves, lds_out=c_tile, always_valid=True,
+        num_waves=num_waves, lds_out=c_tile, swiglu_limit=swiglu_limit, always_valid=True,
         out_tensor=out_tensor if use_tile_resource else None)
     # fmt: on
 
@@ -339,6 +339,7 @@ def compile_gemm1(
     tile_n: int = 256, tile_k: int = 256, num_waves: int = 4, pipe_weights: bool = True,
     mfma_amajor: bool = False, swizzle_a: bool = True, async_a_copy: bool = False,
     use_tile_resource: bool = True, waves_per_eu_hint: int = 2, b_cache_modifier: int = 0,
+    swiglu_limit: float = 0.0,
 ):
     # fmt: on
     """Compile standalone group GEMM1 from the fused Stage1 compute body."""
@@ -410,7 +411,7 @@ def kernel(
             total_threads=total_threads, k_iters=k_iters, a_lds_i32=a_lds_i32, n_tiles=n_tiles,
             expert_offset=expert_offset, b_cache_modifier=b_cache_modifier, swizzle_a=swizzle_a,
             pipe_weights=pipe_weights, mfma_amajor=mfma_amajor, async_a_copy=async_a_copy,
-            use_tile_resource=use_tile_resource,
+            use_tile_resource=use_tile_resource, swiglu_limit=swiglu_limit,
         )
         total_work = (num_valid // fx.Int32(sort_block_m)) * fx.Int32(n_tiles)
         for flat in range(fx.block_idx.x, total_work, grid_x):
@@ -440,7 +441,7 @@ def gemm1_kernel(
     tile_n: int = 256, tile_k: int = 256, num_waves: int = 4, grid_mult: int = 4,
     pipe_weights: bool = True, mfma_amajor: bool = False, swizzle_a: bool = True,
     async_a_copy: bool = False, use_tile_resource: bool = True, waves_per_eu_hint: int = 2,
-    num_cu: int = 256, b_cache_modifier: int = 0,
+    num_cu: int = 256, b_cache_modifier: int = 0, swiglu_limit: float = 0.0,
 ):
     # fmt: on
     """Run standalone MegaMoEV2 group GEMM1 and return ``(out, out_scale)``."""
@@ -458,6 +459,7 @@ def gemm1_kernel(
         pipe_weights=pipe_weights, mfma_amajor=mfma_amajor, swizzle_a=swizzle_a,
         async_a_copy=async_a_copy, use_tile_resource=use_tile_resource,
         waves_per_eu_hint=waves_per_eu_hint, b_cache_modifier=b_cache_modifier,
+        swiglu_limit=swiglu_limit,
     )
     _run_compiled(
         launch, out, x, w, scale_x, scale_w, tile_row_base, expert_ids, out_scale,
diff --git a/kernels/mega_moe/gemm_util.py b/kernels/mega_moe/gemm_util.py
index 6aa8a50dc..8f353ac5b 100644
--- a/kernels/mega_moe/gemm_util.py
+++ b/kernels/mega_moe/gemm_util.py
@@ -490,11 +490,11 @@ def call_pipe_am_final(self, a_load, b_prev, acc, sa, sb):
 
 
 class SiluQuantEpilogue:
-    """silu(gate)*up -> fp8 + per-32 E8M0 out-scale (aiter/CK swizzled), via inline .ptr cshuffle."""
+    """SwiGLU followed by FP8 quantization and per-32 E8M0 output scales."""
 
     # fmt: off
     def __init__(self, *, out_rsrc, out_scale_rsrc, sorted_rsrc, tokens, inter_dim, m_repeat, num_acc_n,
-        sort_block_m, tile_n, num_waves, lds_out, always_valid=False, out_tensor=None):
+        sort_block_m, tile_n, num_waves, lds_out, swiglu_limit=0.0, always_valid=False, out_tensor=None):
     # fmt: on
         self._out_rsrc = out_rsrc
         self._out_scale_rsrc = out_scale_rsrc
@@ -507,6 +507,7 @@ def __init__(self, *, out_rsrc, out_scale_rsrc, sorted_rsrc, tokens, inter_dim,
         self._tile_n = tile_n
         self._num_waves = num_waves
         self._lds_out = lds_out
+        self._swiglu_limit = float(swiglu_limit)
         self._always_valid = always_valid
         self._out_tensor = out_tensor
         self._lane = fx.thread_idx.x % 64
@@ -528,7 +529,14 @@ def _combine(self, acc):
     def _silu_mul(self, gate_v4, up_v4):
         gv = Vec(gate_v4)
         uv = Vec(up_v4)
-        elems = [self._silu(gv[i]) * uv[i] for i in range_constexpr(4)]
+        if self._swiglu_limit <= 0:
+            elems = [self._silu(gv[i]) * uv[i] for i in range_constexpr(4)]
+            return Vec.from_elements(elems, fx.Float32)
+        limit = fx.Float32(self._swiglu_limit)
+        elems = [
+            self._silu(-(-gv[i]).maximumf(-limit)) * fx.clampf(uv[i], -limit, limit)
+            for i in range_constexpr(4)
+        ]
         return Vec.from_elements(elems, fx.Float32)
 
     def store(self, acc, tile_i32, tile_row_base_i32, n_tile_base_i32):
diff --git a/kernels/mega_moe/mega_moe.py b/kernels/mega_moe/mega_moe.py
index e8505c303..2a729d57d 100644
--- a/kernels/mega_moe/mega_moe.py
+++ b/kernels/mega_moe/mega_moe.py
@@ -24,7 +24,7 @@ class MegaMoEV2:
     # fmt: off
     def __init__(self, *, rank: int, world_size: int, model_dim: int, inter_dim: int, experts: int, topk: int,
         quant: str, w1: torch.Tensor, w1_scale: torch.Tensor, w2: torch.Tensor, w2_scale: torch.Tensor,
-        max_tok_per_rank: int, mega_scheme: str = "fixedslot"):
+        max_tok_per_rank: int, mega_scheme: str = "fixedslot", swiglu_limit: float = 0.0):
     # fmt: on
         if quant != "a8w4":
             raise ValueError("MegaMoEV2 currently supports quant='a8w4' only")
@@ -40,6 +40,9 @@ def __init__(self, *, rank: int, world_size: int, model_dim: int, inter_dim: int
         self.epr = int(experts // world_size)
         self.topk = int(topk)
         self.mtpr = int(max_tok_per_rank)
+        self.swiglu_limit = float(swiglu_limit)
+        if self.swiglu_limit < 0:
+            raise ValueError("swiglu_limit must be non-negative")
         self.dev = torch.device("cuda", rank)
         self.max_recv = self.world_size * self.mtpr
         compact = self.mtpr > FIXED_SLOT_MAX_MTPR
@@ -213,7 +216,7 @@ def _run_fused_stage1(self, x, wts, scales, topk_ids, stream=None, config: Stage
             use_tile_resource=config.use_tile_resource,
             waves_per_eu_hint=config.waves_per_eu_hint, b_nt=config.b_nt,
             work_shards=config.work_shards, external_grouping=config.external_grouping,
-            external_counting=config.external_counting)
+            external_counting=config.external_counting, swiglu_limit=self.swiglu_limit)
         # fmt: on
         self._s1_active_tile_m = config.sort_block_m
         return self._s1_active_tile_m
diff --git a/kernels/mega_moe/mega_moe_stage1.py b/kernels/mega_moe/mega_moe_stage1.py
index e804d0550..3c2efbe85 100644
--- a/kernels/mega_moe/mega_moe_stage1.py
+++ b/kernels/mega_moe/mega_moe_stage1.py
@@ -58,7 +58,7 @@ def compile_mega_moe_stage1(
     async_a_copy: bool = False, use_tile_resource: bool = True,
     waves_per_eu_hint: int = 2, num_cu: int = 256, num_dispatch_cu: int = 32, b_nt: int = -1,
     work_shards: int | None = None, external_grouping: bool | None = None,
-    external_counting: bool | None = None,
+    external_counting: bool | None = None, swiglu_limit: float = 0.0,
 ):
     arch = str(get_rocm_arch() or "")
     if not arch.startswith("gfx95"):
@@ -140,12 +140,14 @@ class SharedStorage:
         A_scale: fx.Array[fx.Int8, n_scale_bytes, 16]
 
     dispatch_path = "fixedslot" if fixed_slot_dispatch else "compact"
+    swiglu_suffix = "" if swiglu_limit <= 0 else f"_sl{str(float(swiglu_limit)).replace('.', 'p')}"
     kernel_name = (
         f"megamoe_stage1_{dispatch_path}_t{sort_block_m}x{tile_n}x{tile_k}"
         f"_w{NUM_WAVES}_gm{grid_mult}"
         f"_dcu{dispatch_blocks}_pw{int(pipe_weights)}ma{int(mfma_amajor)}sw{int(swizzle_a)}"
         f"aa{int(async_a_copy)}"
         f"_tr{int(use_tile_resource)}wpe{waves_per_eu_hint}_bnt{b_cache_modifier}_ws{WORK_SHARDS}"
+        f"{swiglu_suffix}"
     )
 
     @flyc.kernel(name=kernel_name, known_block_size=[TOTAL_THREADS, 1, 1])
@@ -331,6 +333,7 @@ def _disp_ptr(slot):
             n_tiles=N_TILES, expert_offset=fz_rank * fz_epr, b_cache_modifier=b_cache_modifier,
             swizzle_a=swizzle_a, pipe_weights=pipe_weights, mfma_amajor=mfma_amajor,
             async_a_copy=async_a_copy, use_tile_resource=use_tile_resource,
+            swiglu_limit=swiglu_limit,
         )
 
         if tid == fx.Int32(0):
@@ -406,7 +409,7 @@ def run_mega_moe_stage1(out, x, w, scale_x, scale_w, sorted_token_ids, expert_id
     sort_block_m=32, tile_n=256, tile_k=256, num_waves=4, grid_mult=4, pipe_weights=True,
     mfma_amajor=False, swizzle_a=True, async_a_copy=False, num_dispatch_cu=32,
     use_tile_resource=True, waves_per_eu_hint=2,
-    b_nt=-1, work_shards=None, external_grouping=None, external_counting=None):
+    b_nt=-1, work_shards=None, external_grouping=None, external_counting=None, swiglu_limit=0.0):
     launch = compile_mega_moe_stage1(
         model_dim=model_dim, inter_dim=inter_dim, rank=rank, experts_per_rank=experts_per_rank,
         fuse_npes=fuse_npes, fuse_topk=fuse_topk, fuse_cap=fuse_cap, fuse_mtpr=fuse_mtpr,
@@ -416,7 +419,7 @@ def run_mega_moe_stage1(out, x, w, scale_x, scale_w, sorted_token_ids, expert_id
         async_a_copy=async_a_copy, use_tile_resource=use_tile_resource,
         waves_per_eu_hint=waves_per_eu_hint, num_cu=num_cu, num_dispatch_cu=num_dispatch_cu,
         b_nt=b_nt, work_shards=work_shards, external_grouping=external_grouping,
-        external_counting=external_counting,
+        external_counting=external_counting, swiglu_limit=swiglu_limit,
     )
     _run_compiled(
         launch, out, x, w, scale_x, scale_w, sorted_token_ids, expert_ids, num_valid_ids, out_scale,
diff --git a/tests/kernels/test_mega_moe_v2.py b/tests/kernels/test_mega_moe_v2.py
index fd6c0c67d..8ac7cf9db 100644
--- a/tests/kernels/test_mega_moe_v2.py
+++ b/tests/kernels/test_mega_moe_v2.py
@@ -46,7 +46,7 @@
 NETWORKS = {
     "r1_v3": dict(model_dim=7168, inter_dim=2048, experts=256, topk=8),
     "v4_flash": dict(model_dim=4096, inter_dim=2048, experts=256, topk=6),
-    "v4_pro": dict(model_dim=7168, inter_dim=3072, experts=384, topk=6),
+    "v4_pro": dict(model_dim=7168, inter_dim=3072, experts=384, topk=6, swiglu_limit=10.0),
 }
 
 # batch-size sweeps for --matrix / --full-bs.
@@ -274,6 +274,13 @@ def _rmsnorm(x, eps=1e-6):
     return n.to(x.dtype)
 
 
+def _swiglu(gate, up, limit):
+    if limit > 0:
+        gate = gate.clamp(max=limit)
+        up = up.clamp(-limit, limit)
+    return torch.nn.functional.silu(gate) * up
+
+
 def _calc_diff(x, y):
     """Return one minus FP64 cosine similarity."""
     x, y = x.double(), y.double()
@@ -297,9 +304,10 @@ def _make_layer_routings(n_layers, tokens, experts, topk, dev, seed, rank):
 class RefModel:
     """Pure PyTorch FP32 reference for chained MoE residual layers."""
 
-    def __init__(self, w1_f32, w2_f32, inter_dim, dev, sw1=None, sw2=None):
+    def __init__(self, w1_f32, w2_f32, inter_dim, dev, swiglu_limit=0.0, sw1=None, sw2=None):
         self.w1_f32, self.w2_f32 = w1_f32, w2_f32  # full-precision [E, 2I, H], [E, H, I]
         self.inter_dim = inter_dim
+        self.swiglu_limit = float(swiglu_limit)
         self.sw1, self.sw2 = sw1, sw2  # optional dense shared experts (None here)
         self.dev = dev
         self._cache = {}
@@ -314,12 +322,9 @@ def _expert(self, g):
             )
         return wd
 
-    @staticmethod
-    def _ffn(x, w1d, w2d):
-        import torch.nn.functional as _F
-
+    def _ffn(self, x, w1d, w2d):
         gate, up = (x @ w1d.t()).chunk(2, dim=-1)
-        return (_F.silu(gate) * up) @ w2d.t()
+        return _swiglu(gate, up, self.swiglu_limit) @ w2d.t()
 
     def _shared(self, x):
         if self.sw1 is None:
@@ -460,6 +465,7 @@ def _run_full_e2e(
     experts,
     epr,
     topk,
+    swiglu_limit,
     run_tokens,
     mtpr,
     a_dtype,
@@ -474,7 +480,6 @@ def _run_full_e2e(
 ):
     """Compare MegaMoEV2 with FP8- and BF16-dispatch ATOM pipelines."""
     import numpy as _np
-    import torch.nn.functional as _F
 
     from kernels.mega_moe import MegaMoEV2
     from kernels.mega_moe.quant import mxfp4_moe_scale_sort, per_1x32_mx_quant
@@ -562,6 +567,7 @@ def _cg_time(body, dc_op):
         w2=w2_kernel,
         w2_scale=w2_scale_1d,
         max_tok_per_rank=mtpr,
+        swiglu_limit=swiglu_limit,
     )
     torch.cuda.synchronize()
     ms.shmem_barrier_all()
@@ -810,7 +816,9 @@ def _atom_fp8_body():
         w1e = _dequant_mx_to_f32(w1_all[e], "fp4")  # [2*inter_dim, model_dim]
         w2e = _dequant_mx_to_f32(w2_all[e], "fp4")  # [model_dim, inter_dim]
         xr = x32[rows]
-        _a1 = _F.silu(xr @ w1e[:inter_dim].t()) * (xr @ w1e[inter_dim : 2 * inter_dim].t())
+        gate = xr @ w1e[:inter_dim].t()
+        up = xr @ w1e[inter_dim : 2 * inter_dim].t()
+        _a1 = _swiglu(gate, up, swiglu_limit)
         oracle_w[rows] += w_e[:, None] * (_a1 @ w2e.t())
         del w1e, w2e
     orw = oracle_w.cpu().numpy()
@@ -827,7 +835,7 @@ def _atom_fp8_body():
     _oracle_broken = _ra_w > _floor
     # Allow the expected FP4 quantization divergence while requiring no material regression.
     _match_ok = (_rma < 5e-2) or (_rm_w <= _ra8_w + 2e-2)
-    ok = _match_ok and (_oracle_broken or (_mega_ok and _atom8_ok))
+    ok = _mega_ok if swiglu_limit > 0 else _match_ok and (_oracle_broken or (_mega_ok and _atom8_ok))
 
     # Gate on the worst expert shard across ranks.
     _rm_w_max = _all_max(dev, _rm_w)
@@ -955,6 +963,7 @@ def _run_mega_only(
     experts,
     epr,
     topk,
+    swiglu_limit,
     run_tokens,
     mtpr,
     quant,
@@ -973,7 +982,6 @@ def _run_mega_only(
 ):
     """Run the aiter-free MegaMoEV2 accuracy and performance CI path."""
     import numpy as _np
-    import torch.nn.functional as _F
 
     from kernels.mega_moe import MegaMoEV2
 
@@ -1030,6 +1038,7 @@ def _relL2(a, b):
         w2=w2_kernel,
         w2_scale=w2_scale_1d,
         max_tok_per_rank=mtpr,
+        swiglu_limit=swiglu_limit,
     )
     torch.cuda.synchronize()
     ms.shmem_barrier_all()
@@ -1102,7 +1111,9 @@ def _body():
             rows = torch.nonzero(eids_s1 == int(expert), as_tuple=False).flatten()
             w1e = _dequant_mx_to_f32(w_ref_local[expert - rank * epr], "fp4")
             xr = inp_s1[rows]
-            ref_s1[rows] = _F.silu(xr @ w1e[:inter_dim].t()) * (xr @ w1e[inter_dim:].t())
+            gate = xr @ w1e[:inter_dim].t()
+            up = xr @ w1e[inter_dim:].t()
+            ref_s1[rows] = _swiglu(gate, up, swiglu_limit)
         if not torch.isfinite(got_s1).all() or not torch.isfinite(ref_s1).all():
             got_bad = int((~torch.isfinite(got_s1)).sum().item())
             ref_bad = int((~torch.isfinite(ref_s1)).sum().item())
@@ -1163,7 +1174,7 @@ def _body():
             torch.cuda.synchronize()
             ms.shmem_barrier_all()
             out_dev = xd[:run_tokens].float()
-            out_ref = RefModel(w1_all, w2_all, inter_dim, dev).run(x_in, routings).float()
+            out_ref = RefModel(w1_all, w2_all, inter_dim, dev, swiglu_limit).run(x_in, routings).float()
             _acc_metric = _calc_diff(out_ref, out_dev)  # 1 - cosine (fp64), end-to-end accumulated
             _acc_floor = _CHAIN_TOL
             _acc_label = f"cos_diff(chain x{n_layers})"
@@ -1193,7 +1204,9 @@ def _body():
                 w1e = _dequant_mx_to_f32(w1_all[local_e], "fp4")
                 w2e = _dequant_mx_to_f32(w2_all[local_e], "fp4")
                 xr = x32[rows]
-                _a1 = _F.silu(xr @ w1e[:inter_dim].t()) * (xr @ w1e[inter_dim : 2 * inter_dim].t())
+                gate = xr @ w1e[:inter_dim].t()
+                up = xr @ w1e[inter_dim : 2 * inter_dim].t()
+                _a1 = _swiglu(gate, up, swiglu_limit)
                 oracle[rows] += w_e[:, None] * (_a1 @ w2e.t())
                 del w1e, w2e
             if local_experts_only:
@@ -1327,6 +1340,7 @@ def _stage2_body():
 def run_one(args, rank, world, dev):
     net = NETWORKS[args.network]
     model_dim, inter_dim, experts = net["model_dim"], net["inter_dim"], net["experts"]
+    swiglu_limit = float(net.get("swiglu_limit", 0.0))
     # topk: --topk>0 overrides; else use the network's native topk (r1_v3=8, v4_*=6).
     topk = int(args.topk) if int(args.topk) > 0 else int(net["topk"])
     run_tokens = max(int(args.tokens), 1)  # allow bs=1 (1 token/rank); routing still reaches all ranks
@@ -1369,6 +1383,7 @@ def run_one(args, rank, world, dev):
             experts=experts,
             epr=epr,
             topk=topk,
+            swiglu_limit=swiglu_limit,
             run_tokens=run_tokens,
             mtpr=mtpr,
             quant=args.quant,
@@ -1397,6 +1412,7 @@ def run_one(args, rank, world, dev):
         experts=experts,
         epr=epr,
         topk=topk,
+        swiglu_limit=swiglu_limit,
         run_tokens=run_tokens,
         mtpr=mtpr,
         a_dtype=a_dtype,