[RadeonFlow] flydsl mxfp4 a4w4 MoE: improv perf

aiter/fused_moe.py

@@ -1525,11 +1525,7 @@ def _mxfp4_a4w4_stage2(
             return out
 
         # Lossy before-sum 4-bit quant (ok for gsm8k, degrades other evals): opt-in.
-        if (
-            mxfp4out
-            and _mx_shape_ok
-            and os.environ.get("AITER_MXFP4_INTERMEDIATE", "0") == "1"
-        ):
+        if _mx_shape_ok and os.environ.get("AITER_MXFP4_INTERMEDIATE", "0") == "1":
             flat_out_q = torch.empty(
                 (max_sorted, D_HIDDEN // 2), dtype=torch.uint8, device=device
             )

aiter/ops/flydsl/kernels/mxfp4_gemm1.py

@@ -221,6 +221,12 @@ def _silu_mul(g, u):
     return g * sig * u
 
 
+def _silu_mul_batch(gs, us):
+    e = [fx.Float32(rocdl.exp2(T.f32, _raw(g * fx.Float32(-LOG2E)))) for g in gs]
+    sig = [fx.Float32(rocdl.rcp(T.f32, _raw(fx.Float32(1.0) + ei))) for ei in e]
+    return [gs[i] * sig[i] * us[i] for i in range(len(gs))]
+
+
 def _fabs_f32(x):
     """fabsf via bit-mask (FlyDSL has no arith.absf): clear the sign bit."""
     bits = _raw(x).bitcast(T.i32)
@@ -793,6 +799,7 @@ def mfma_cluster(b_slot, a, a_scale, bs_slot, J, init):
             )
 
     # ---- prologue: stages 0,1 (HIP 431-463) ----
+    _relax_prologue = (BM == 128) and not inline_quant
     if const_expr(not inline_quant):
         issue_a_scale_load()
     for K_C in range_constexpr(kStages):
@@ -809,19 +816,47 @@ def mfma_cluster(b_slot, a, a_scale, bs_slot, J, init):
             inline_quant_pack_write(K_C, scale_accum)
         else:
             issue_a_load_lds(K_C, K_C)
+            if const_expr(not _relax_prologue):
+                for j in range_constexpr(4):
+                    issue_b_load_j(b[K_C], K_C, j)
+        if const_expr(not _relax_prologue):
+            issue_b_scale_load(b_scale_v[K_C], K_C)
+    if const_expr(_relax_prologue):
+        rocdl.sched_barrier(0)
+        for K_C in range_constexpr(kStages):
             for j in range_constexpr(4):
                 issue_b_load_j(b[K_C], K_C, j)
-        issue_b_scale_load(b_scale_v[K_C], K_C)
+            issue_b_scale_load(b_scale_v[K_C], K_C)
 
     # ---- main loop: OFFSET in [0,26) (HIP 465-552 non-inline else) ----
     for OFFSET in range_constexpr(kUnroll):
         K_C = kStages + OFFSET
         read_slot = OFFSET % kAStages
         write_slot = K_C % kAStages
         slot_b = OFFSET % kStages
-        gpu.barrier()  # __syncthreads (HIP 472)
-        a_cur = issue_a_ds_read(read_slot)
-        asc_cur = issue_a_scale_ds_read(K_C - kStages)
+        if const_expr(_relax_prologue and OFFSET == 0):
+            llvm.inline_asm(
+                res=None,
+                operands_=[],
+                asm_string=f"s_waitcnt vmcnt({10 * kStages})",
+                constraints="",
+                has_side_effects=True,
+            )
+            llvm.inline_asm(
+                res=None,
+                operands_=[],
+                asm_string="s_barrier",
+                constraints="",
+                has_side_effects=True,
+            )
+        else:
+            gpu.barrier()  # __syncthreads (HIP 472)
+        if const_expr(BM == 128):
+            asc_cur = issue_a_scale_ds_read(K_C - kStages)
+            a_cur = issue_a_ds_read(read_slot)
+        else:
+            a_cur = issue_a_ds_read(read_slot)
+            asc_cur = issue_a_scale_ds_read(K_C - kStages)
         if const_expr(not inline_quant):
             issue_a_load_lds(write_slot, K_C)
         # Inline path (HIP :523-528): pre-load the next K-tile's hidden into regs so
@@ -858,8 +893,12 @@ def mfma_cluster(b_slot, a, a_scale, bs_slot, J, init):
     for S in range_constexpr(kStages):
         kt = K_TILES_TOTAL - kStages + S
         gpu.barrier()
-        a_cur = issue_a_ds_read(kt % kAStages)
-        asc_cur = issue_a_scale_ds_read(kt)
+        if const_expr(BM == 128):
+            asc_cur = issue_a_scale_ds_read(kt)
+            a_cur = issue_a_ds_read(kt % kAStages)
+        else:
+            a_cur = issue_a_ds_read(kt % kAStages)
+            asc_cur = issue_a_scale_ds_read(kt)
         for J in range_constexpr(4):
             mfma_cluster(
                 b[kt % kStages], a_cur, asc_cur, b_scale_v[kt % kStages], J, init=False
@@ -906,16 +945,17 @@ def mfma_cluster(b_slot, a, a_scale, bs_slot, J, init):
         row_local = fx.Int32(mr * 16) + m_lane
 
         # read 8 gate + 8 up f32 from lds_acc (epilog 75-83).
-        result = [None] * 8
+        gate_vs = [None] * 8
+        up_vs = [None] * 8
         for ee in range_constexpr(8):
             col_in_grp = fx.Int32(8) * kk + fx.Int32(ee)
             gate_col = wave_grp * fx.Int32(32) + col_in_grp
             up_col = fx.Int32(128) + gate_col
             gate_off = (row_local * fx.Int32(BN) + gate_col) * fx.Int32(4)
             up_off = (row_local * fx.Int32(BN) + up_col) * fx.Int32(4)
-            gate_v = fx.Float32(llvm.load(T.f32, _gep3(lds_acc_base, gate_off)))
-            up_v = fx.Float32(llvm.load(T.f32, _gep3(lds_acc_base, up_off)))
-            result[ee] = _silu_mul(gate_v, up_v)
+            gate_vs[ee] = fx.Float32(llvm.load(T.f32, _gep3(lds_acc_base, gate_off)))
+            up_vs[ee] = fx.Float32(llvm.load(T.f32, _gep3(lds_acc_base, up_off)))
+        result = _silu_mul_batch(gate_vs, up_vs)
 
         # local amax over the 8 results (epilog 91-95).
         local_max = _fabs_f32(result[0])

aiter/ops/flydsl/kernels/mxfp4_gemm2.py

@@ -1111,7 +1111,6 @@ def _flat_mxfp4_epilog(
     n_lane = tid_i32 % fx.Int32(16)
     wave_grp = n_lane // fx.Int32(4)
     kk = n_lane % fx.Int32(4)
-    i7fff = _raw(fx.Int32(0x7FFFFFFF))
     _m_base = m_row + m_lane
     _q_row0 = fx.Int64(_m_base) * fx.Int64(N_OUT // 2)
     _s_row0 = fx.Int64(_m_base) * fx.Int64(N_OUT // 32)
@@ -1147,12 +1146,13 @@ def _issue_load(mr, half):
             _r_next, _grp_next, _col0_next = _issue_load(*_blocks[_bi + 1])
         if True:
             # block amax over |r[0..7]| (positive-float bits) -> bf16-bits
-            maxb = arith.andi(arith.bitcast(T.i32, _raw(r[0])), i7fff)
+            amax_f = llvm.call_intrinsic(T.f32, "llvm.fabs.f32", [_raw(r[0])], [], [])
             for e in range_constexpr(1, 8):
-                maxb = arith.maxui(
-                    maxb, arith.andi(arith.bitcast(T.i32, _raw(r[e])), i7fff)
+                abs_e = llvm.call_intrinsic(
+                    T.f32, "llvm.fabs.f32", [_raw(r[e])], [], []
                 )
-            amax = arith.shrui(maxb, _raw(fx.Int32(16)))
+                amax_f = arith.maxnumf(amax_f, abs_e)
+            amax = arith.shrui(arith.bitcast(T.i32, amax_f), _raw(fx.Int32(16)))
             # DPP quad-amax (reduce across the 4 kk-lanes of the block)
             s1 = rocdl.update_dpp(T.i32, amax, amax, 0xB1, 0xF, 0xF, True)
             a = arith.maxui(amax, s1)