diff --git a/aiter/fused_moe.py b/aiter/fused_moe.py index 6aec365165..b578d0001f 100644 --- a/aiter/fused_moe.py +++ b/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 ) diff --git a/aiter/ops/flydsl/kernels/mxfp4_gemm1.py b/aiter/ops/flydsl/kernels/mxfp4_gemm1.py index c96eba8246..0b144f3412 100644 --- a/aiter/ops/flydsl/kernels/mxfp4_gemm1.py +++ b/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,9 +816,17 @@ 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): @@ -819,9 +834,29 @@ def mfma_cluster(b_slot, a, a_scale, bs_slot, J, init): 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]) diff --git a/aiter/ops/flydsl/kernels/mxfp4_gemm2.py b/aiter/ops/flydsl/kernels/mxfp4_gemm2.py index cbc355de9b..1a174b9ead 100644 --- a/aiter/ops/flydsl/kernels/mxfp4_gemm2.py +++ b/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)