From 0a0d0ef41f73c19fdd367d28b06bdcb5d8005c8d Mon Sep 17 00:00:00 2001 From: ColorsWind <14761584+ColorsWind@users.noreply.github.com> Date: Sun, 21 Jun 2026 17:04:52 +0000 Subject: [PATCH 1/7] flydsl mxfp4 gemm1/gemm2: raw-pointer (data_ptr i64) ABI Pass bare data_ptr() device addresses (fx.Int64) for the global buffer args instead of full memref descriptors. The kernels only need base pointers (contiguity + compile-time sizes), so the addresses pack contiguously into kernargs -> coalesced s_load prologue. ~7% faster at decode (M=4/8), converging to parity by M=64. cos preserved. --- aiter/ops/flydsl/kernels/mxfp4_gemm1.py | 164 +++++++++++++++++------- aiter/ops/flydsl/kernels/mxfp4_gemm2.py | 120 ++++++++--------- aiter/ops/flydsl/mxfp4_gemm1_kernels.py | 23 ++-- aiter/ops/flydsl/mxfp4_gemm2_kernels.py | 23 ++-- 4 files changed, 207 insertions(+), 123 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mxfp4_gemm1.py b/aiter/ops/flydsl/kernels/mxfp4_gemm1.py index f0d0bb98a1..4b9403edea 100644 --- a/aiter/ops/flydsl/kernels/mxfp4_gemm1.py +++ b/aiter/ops/flydsl/kernels/mxfp4_gemm1.py @@ -185,10 +185,15 @@ def _gep3(base_ptr, byte_off_i32): ) -def _global_base_ptr1(arg): - """One ptr<1> base for a global tensor (single memref->ptr conversion).""" - base_idx = buffer_ops.extract_base_index(arg, address_space=1) - return llvm.inttoptr(ir.Type.parse("!llvm.ptr<1>"), _raw(fx.Int64(base_idx))) +def _global_base_ptr1(addr_i64): + """ptr<1> base from a raw i64 device address. + + Global args are passed as bare ``data_ptr()`` (fx.Int64) rather than full + memref descriptors: this kernel only ever needs the base pointer (it assumes + contiguity and derives all sizes from i32_ntok / compile-time constants), so + the dynamic-memref shape/stride layout buffer was dead weight that scattered + the kernarg pointers and blocked LLVM from coalescing the scalar loads.""" + return llvm.inttoptr(ir.Type.parse("!llvm.ptr<1>"), _raw(fx.Int64(addr_i64))) def _gep1(base_ptr, byte_off_i32): @@ -240,14 +245,16 @@ def _e8m0_from_amax(amax_f32): # -- inline-quant helpers (HIP mxfp4_gemm_common.hpp:76-94) ------------------- def _pkmax_u16(a_i32, b_i32): """v_pk_max_u16 a, b -- pairwise u16 max of two packed-u16 dwords. - No FlyDSL builtin; mirror HIP inline_quant_pkmax_u16 via inline asm.""" - out = llvm.inline_asm( - res=T.i32, - operands_=[_raw(a_i32), _raw(b_i32)], - asm_string="v_pk_max_u16 $0, $1, $2", - constraints="=v,v,v", - has_side_effects=False, - ) + + No FlyDSL builtin, but no inline asm needed: bitcast each dword to + vector<2xi16> and let arith.maxui lower through MLIR -> LLVM. The AMDGPU + backend ISel-pattern-matches `umax <2 x i16>` to a single v_pk_max_u16, + so this emits the same instruction with zero inline asm.""" + _v2i16 = ir.Type.parse("vector<2xi16>") + va = llvm.BitcastOp(_v2i16, _raw(a_i32)).result + vb = llvm.BitcastOp(_v2i16, _raw(b_i32)).result + vm = arith.MaxUIOp(va, vb).result + out = llvm.BitcastOp(T.i32, vm).result return fx.Int32(out) @@ -369,26 +376,28 @@ def _gemm1_body( # max_size=False memref fallback for DLPack's dynamic-shape a_quant) would read # garbage past the logical extent into the padding rows. Pass the exact byte # count = n_tokens * K_HALF (a_quant is uint8, 1 byte/elem). + # args arrive as raw i64 device addresses (data_ptr()); build buffer resources + # straight from the address -- no memref descriptor needed (see _global_base_ptr1). aq_num_records = arith.index_cast(T.index, _raw(i32_ntok * fx.Int32(K_HALF))) - aq_rsrc = buffer_ops.create_buffer_resource( - arg_aq, max_size=False, num_records_bytes=aq_num_records + aq_rsrc = buffer_ops.create_buffer_resource_from_addr( + _raw(fx.Int64(arg_aq)), num_records_bytes=aq_num_records ) - ascale_rsrc = buffer_ops.create_buffer_resource( - arg_ascale, max_size=False, num_records_bytes=fx.Index(ASCALE_BYTES) + ascale_rsrc = buffer_ops.create_buffer_resource_from_addr( + _raw(fx.Int64(arg_ascale)), num_records_bytes=ASCALE_BYTES ) - bq_rsrc = buffer_ops.create_buffer_resource( - arg_bq, max_size=False, num_records_bytes=fx.Index(BQ_BYTES) + bq_rsrc = buffer_ops.create_buffer_resource_from_addr( + _raw(fx.Int64(arg_bq)), num_records_bytes=BQ_BYTES ) - bscale_rsrc = buffer_ops.create_buffer_resource( - arg_bscale, max_size=False, num_records_bytes=fx.Index(BSCALE_BYTES) + bscale_rsrc = buffer_ops.create_buffer_resource_from_addr( + _raw(fx.Int64(arg_bscale)), num_records_bytes=BSCALE_BYTES ) # hidden_states rsrc (inline-quant only): n_tokens*K*sizeof(bf16) bytes (HIP :92-97). # Non-inline keeps arg_hidden unused. hidden_rsrc = None if const_expr(inline_quant): hidden_num = arith.index_cast(T.index, _raw(i32_ntok * fx.Int32(K * 2))) - hidden_rsrc = buffer_ops.create_buffer_resource( - arg_hidden, max_size=False, num_records_bytes=hidden_num + hidden_rsrc = buffer_ops.create_buffer_resource_from_addr( + _raw(fx.Int64(arg_hidden)), num_records_bytes=hidden_num ) # -- LDS views (s_aq / s_asc, union-overlapping lds_acc) ------------------ @@ -608,6 +617,72 @@ def _inline_quant_core(B128_IDX, SUB, slot, kt, h_v, scale_accum): pack_byte = B128_IDX * 2 + SUB scale_accum[0] = scale_accum[0] | (e8m0 << fx.Int32(pack_byte * 8)) + def _inline_quant_core_pair(specs, slot, kt, scale_accum): + """Interleaved N-row variant of _inline_quant_core: emit each pipeline + stage (abs+pkmax, cross-lane DPP quad-reduce, e8m0) for ALL rows + back-to-back so the scheduler can hide one row's high-latency DPP / + umax behind the other row's independent work -- matching HIP gemm1 + BM16's quant interleave (flydsl's row-by-row order leaves s_nop bubbles + in the DPP latency window). specs = [(B128_IDX, SUB, h_v), ...].""" + n = len(specs) + h_dw = [ + [fx.Int32(_raw(h_v[j])) for j in range_constexpr(4)] + for (_b, _s, h_v) in specs + ] + # stage 1: |bf16| + pkmax tree -> per-row local u16 amax + la = [None] * n + for i in range_constexpr(n): + hm = [h_dw[i][j] & fx.Int32(0x7FFF7FFF) for j in range_constexpr(4)] + m01 = _pkmax_u16(hm[0], hm[1]) + m23 = _pkmax_u16(hm[2], hm[3]) + m0123 = _pkmax_u16(m01, m23) + lo = m0123 & fx.Int32(0xFFFF) + hi = m0123.shrui(fx.Int32(16)) & fx.Int32(0xFFFF) + la[i] = _umax_i32(lo, hi) + # stage 2: DPP quad-reduce, INTERLEAVED across rows (hide DPP latency) + a = [fx.Int32(_raw(la[i])) for i in range_constexpr(n)] + s1 = [ + fx.Int32(dpp_utils.update_dpp_i32(_raw(a[i]), _raw(a[i]), 0xB1, 0xF, 0xF, True)) + for i in range_constexpr(n) + ] + a = [_umax_i32(a[i], s1[i]) for i in range_constexpr(n)] + s2 = [ + fx.Int32(dpp_utils.update_dpp_i32(_raw(a[i]), _raw(a[i]), 0x4E, 0xF, 0xF, True)) + for i in range_constexpr(n) + ] + a = [_umax_i32(a[i], s2[i]) for i in range_constexpr(n)] + # stage 3: e8m0 per row + e8 = [_inline_e8m0(a[i]) for i in range_constexpr(n)] + # stage 4: cvt pack + LDS store + scale fold per row + for i in range_constexpr(n): + B128_IDX, SUB, _hv = specs[i] + qs_raw = _raw( + fx.Float32(_raw(e8[i] << fx.Int32(23)).bitcast(T.f32)) + ) + pk = _raw(fx.Int32(0)) + for j in range_constexpr(4): + src_bf16x2 = _raw( + Vec.from_elements([h_dw[i][j]], fx.Int32).bitcast(fx.BFloat16) + ) + pk = rocdl.cvt_scalef32_pk_fp4_bf16(T.i32, pk, src_bf16x2, qs_raw, j) + pk = fx.Int32(pk) + r = fx.Int32(SUB * 16) + r_in_chunk + kb_in_kt = fx.Int32(B128_IDX * 4) + lane_shr2_and3 + mask_r = _lds_swizzle_mask(r) + b_off = lib * fx.Int32(4) + aq_base = fx.Int32( + memref_dialect.extract_aligned_pointer_as_index(s_aq.get()) + ) + off = ( + fx.Int32(slot * (BM * KH_TILE)) + + r * fx.Int32(KH_TILE) + + ((kb_in_kt * fx.Int32(16)) ^ mask_r) + + b_off + ) + llvm.StoreOp(_raw(pk), _lds_ptr3(aq_base, off)) + pack_byte = B128_IDX * 2 + SUB + scale_accum[0] = scale_accum[0] | (e8[i] << fx.Int32(pack_byte * 8)) + def inline_quant_kt(B128_IDX, SUB, slot, kt, row_token, scale_accum): h_v = inline_quant_load_kt(B128_IDX, kt, row_token) _inline_quant_core(B128_IDX, SUB, slot, kt, h_v, scale_accum) @@ -774,8 +849,9 @@ def mfma_cluster(b_slot, a, a_scale, bs_slot, J, init): # Inline quant of the pre-loaded h_v0/h_v1 into write_slot (HIP :545-550). if const_expr(inline_quant): scale_accum = [fx.Int32(0)] - inline_quant_finish_kt(0, 0, write_slot, K_C, h_v0, scale_accum) - inline_quant_finish_kt(1, 0, write_slot, K_C, h_v1, scale_accum) + _inline_quant_core_pair( + [(0, 0, h_v0), (1, 0, h_v1)], write_slot, K_C, scale_accum + ) inline_quant_pack_write(K_C, scale_accum) # ---- drain: S in [0,2) (HIP 554-565) ---- @@ -1022,17 +1098,17 @@ def compile_gemm1_a4w4_port( @flyc.kernel(name=f"gemm1_a4w4_port_{name_suffix}", known_block_size=[256, 1, 1]) def gemm1_kernel( - arg_aq: fx.Tensor, - arg_ascale: fx.Tensor, - arg_bq: fx.Tensor, - arg_bscale: fx.Tensor, - arg_eids: fx.Tensor, - arg_cumsum: fx.Tensor, - arg_mind: fx.Tensor, + arg_aq: fx.Int64, + arg_ascale: fx.Int64, + arg_bq: fx.Int64, + arg_bscale: fx.Int64, + arg_eids: fx.Int64, + arg_cumsum: fx.Int64, + arg_mind: fx.Int64, i32_ntok: fx.Int32, - arg_aqout: fx.Tensor, - arg_ascaleout: fx.Tensor, - arg_hidden: fx.Tensor, + arg_aqout: fx.Int64, + arg_ascaleout: fx.Int64, + arg_hidden: fx.Int64, ): tx = gpu.thread_id("x") bx = gpu.block_id("x") @@ -1089,18 +1165,18 @@ def gemm1_kernel( @flyc.jit def launch_gemm1( - arg_aq: fx.Tensor, - arg_ascale: fx.Tensor, - arg_bq: fx.Tensor, - arg_bscale: fx.Tensor, - arg_eids: fx.Tensor, - arg_cumsum: fx.Tensor, - arg_mind: fx.Tensor, + arg_aq: fx.Int64, + arg_ascale: fx.Int64, + arg_bq: fx.Int64, + arg_bscale: fx.Int64, + arg_eids: fx.Int64, + arg_cumsum: fx.Int64, + arg_mind: fx.Int64, i32_ntok: fx.Int32, i32_grid: fx.Int32, - arg_aqout: fx.Tensor, - arg_ascaleout: fx.Tensor, - arg_hidden: fx.Tensor, + arg_aqout: fx.Int64, + arg_ascaleout: fx.Int64, + arg_hidden: fx.Int64, stream: fx.Stream, ): from flydsl.compiler.kernel_function import CompilationContext diff --git a/aiter/ops/flydsl/kernels/mxfp4_gemm2.py b/aiter/ops/flydsl/kernels/mxfp4_gemm2.py index 9ae2c68116..e23f12d0f5 100644 --- a/aiter/ops/flydsl/kernels/mxfp4_gemm2.py +++ b/aiter/ops/flydsl/kernels/mxfp4_gemm2.py @@ -63,9 +63,6 @@ # Measured: 1 workgroup/CU (grid == NUM_CU) is optimal; over-subscribing the # persistent grid only adds L2/memory-queue contention (the kernel has enough # memory-level parallelism at 1 wg/CU), so the grid is capped at NUM_CU. -# Explicit hand-tuned vmcnt for the nonatomic K-loop ds_read fence (inline asm). -# None -> let the backend choose (rocdl.barrier); an int forces s_waitcnt vmcnt(N). -_NONATOMIC_KLOOP_VMCNT = 16 # scale-layout consts (mirror gemm2_a4w4.cuh). K-independent stride: kBS_stride_k0_dw = 64 @@ -185,6 +182,16 @@ def _raw(v): return v +def _udiv(a, c): + cc = fx.Int32(c) if isinstance(c, int) else c + return a // cc + + +def _umod(a, c): + cc = fx.Int32(c) if isinstance(c, int) else c + return a % cc + + def _lds_ptr3(base_i32, byte_off_i32): """ptr<3> = inttoptr(i64(base_i32 + byte_off_i32)).""" addr_i64 = fx.Int64(base_i32 + byte_off_i32) @@ -217,10 +224,14 @@ def _s_barrier_bare(): ) -def _global_base_ptr1(arg): - """One ptr<1> base for a global tensor (single memref->ptr conversion).""" - base_idx = buffer_ops.extract_base_index(arg, address_space=1) - return llvm.inttoptr(ir.Type.parse("!llvm.ptr<1>"), _raw(fx.Int64(base_idx))) +def _global_base_ptr1(addr_i64): + """One ptr<1> base from a raw i64 device address. + + Global args are passed as bare ``data_ptr()`` (fx.Int64) rather than full + memref descriptors (ported from gemm1): the kernel only needs base pointers + (it assumes contiguity + derives sizes from compile-time consts), so raw i64 + addresses pack contiguously into kernargs -> coalesced s_load prologue.""" + return llvm.inttoptr(ir.Type.parse("!llvm.ptr<1>"), _raw(fx.Int64(addr_i64))) def _gep1(base_ptr, byte_off_i32): @@ -364,17 +375,17 @@ def compile_gemm2_a4w4_port( @flyc.kernel(name=_name, known_block_size=[256, 1, 1]) def gemm2_kernel( - arg_aq: fx.Tensor, - arg_ascale: fx.Tensor, - arg_bq: fx.Tensor, - arg_bscale: fx.Tensor, - arg_eids: fx.Tensor, - arg_cumsum: fx.Tensor, - arg_stids: fx.Tensor, - arg_sweights: fx.Tensor, + arg_aq: fx.Int64, + arg_ascale: fx.Int64, + arg_bq: fx.Int64, + arg_bscale: fx.Int64, + arg_eids: fx.Int64, + arg_cumsum: fx.Int64, + arg_stids: fx.Int64, + arg_sweights: fx.Int64, i32_M: fx.Int32, - arg_out: fx.Tensor, - arg_out_scale: fx.Tensor, # flat_out_scale (mxfp4 epilog only; dummy otherwise) + arg_out: fx.Int64, + arg_out_scale: fx.Int64, # flat_out_scale (mxfp4 epilog only; dummy otherwise) ): tx = gpu.thread_id("x") bx = gpu.block_id("x") @@ -384,8 +395,8 @@ def gemm2_kernel( lane = tx_i32 % fx.Int32(64) wave = rocdl.readfirstlane(T.i32, tx_i32 // fx.Int32(64)) # wave == wave_n - aq_rsrc = buffer_ops.create_buffer_resource( - arg_aq, max_size=False, num_records_bytes=fx.Index(_aq_bytes) + aq_rsrc = buffer_ops.create_buffer_resource_from_addr( + _raw(fx.Int64(arg_aq)), num_records_bytes=fx.Index(_aq_bytes) ) saq = SmemPtr( allocator.get_base(), lds_off, T.i8, shape=(_aStages * _slot_bytes,) @@ -450,7 +461,7 @@ def _run_tile(tile_i32): # run without it so the compiler overlaps each tile's loads with the # previous tile's epilog. cumsum0 = llvm.load(T.i32, _global_ptr1(arg_cumsum, fx.Int32(0))) - total_m_blocks = cumsum0 // fx.Int32(BM) + total_m_blocks = _udiv(cumsum0, BM) bound = total_m_blocks * fx.Int32(_num_n_blocks) grid_nb = fx.Int32(gpu.grid_dim.x) @@ -462,16 +473,16 @@ def _run_tile(tile_i32): # PORT_XCD_SWIZZLE<=0: step-1-only (mirrors xcd_remap.hpp swizzle=-1). # PORT_XCD_SWIZZLE>0 : 2-step M-major grouping (positive swizzle). _NXCD = 8 - _xq = bound // fx.Int32(_NXCD) - _xr = bound % fx.Int32(_NXCD) + _xq = _udiv(bound, _NXCD) + _xr = _umod(bound, _NXCD) _SW = PORT_XCD_SWIZZLE def _xcd(pid): - xc = pid % fx.Int32(_NXCD) + xc = _umod(pid, _NXCD) wgid = ( xc * _xq + fx.Int32(arith.minsi(_raw(xc), _raw(_xr))) - + pid // fx.Int32(_NXCD) + + _udiv(pid, _NXCD) ) if const_expr(_SW <= 0): return wgid @@ -492,7 +503,7 @@ def _xcd(pid): if bx_i32 < bound: tile = _xcd(bx_i32) - _issue_all_a_loads((tile // fx.Int32(_num_n_blocks)) * fx.Int32(BM)) + _issue_all_a_loads(_udiv(tile, _num_n_blocks) * fx.Int32(BM)) rocdl.sched_barrier(0) _run_tile(tile) @@ -507,14 +518,14 @@ def _xcd(pid): # a loop-carried iter_arg, which only MLIR values can be. setattr(saq, "_view_cache", None) tile = _xcd(wu) - _issue_all_a_loads((tile // fx.Int32(_num_n_blocks)) * fx.Int32(BM)) + _issue_all_a_loads(_udiv(tile, _num_n_blocks) * fx.Int32(BM)) _run_tile(tile) else: # One-shot grid (atomic): issue A->LDS BEFORE the cumsum load so the # A->LDS HBM latency overlaps the cumsum load + bound check (A->LDS # depends only on bx/lane). Only the first n_load_waves hold A rows # (BM16: waves 0,1), so gate on wave < n_load_waves. - m_row0 = (bx_i32 // fx.Int32(_num_n_blocks)) * fx.Int32(BM) + m_row0 = _udiv(bx_i32, _num_n_blocks) * fx.Int32(BM) if const_expr(_n_load_waves < 4): # BM16: only waves 0,1 hold A rows if wave < fx.Int32(_n_load_waves): _issue_all_a_loads(m_row0) @@ -523,7 +534,7 @@ def _xcd(pid): rocdl.sched_barrier(0) cumsum0 = llvm.load(T.i32, _global_ptr1(arg_cumsum, fx.Int32(0))) - total_m_blocks = cumsum0 // fx.Int32(BM) + total_m_blocks = _udiv(cumsum0, BM) bound = total_m_blocks * fx.Int32(_num_n_blocks) if bx_i32 < bound: @@ -531,18 +542,18 @@ def _xcd(pid): @flyc.jit def launch_gemm2( - arg_aq: fx.Tensor, - arg_ascale: fx.Tensor, - arg_bq: fx.Tensor, - arg_bscale: fx.Tensor, - arg_eids: fx.Tensor, - arg_cumsum: fx.Tensor, - arg_stids: fx.Tensor, - arg_sweights: fx.Tensor, + arg_aq: fx.Int64, + arg_ascale: fx.Int64, + arg_bq: fx.Int64, + arg_bscale: fx.Int64, + arg_eids: fx.Int64, + arg_cumsum: fx.Int64, + arg_stids: fx.Int64, + arg_sweights: fx.Int64, i32_M: fx.Int32, i32_max_m_blocks: fx.Int32, - arg_out: fx.Tensor, - arg_out_scale: fx.Tensor, # flat_out_scale (mxfp4 epilog only; dummy otherwise) + arg_out: fx.Int64, + arg_out_scale: fx.Int64, # flat_out_scale (mxfp4 epilog only; dummy otherwise) stream: fx.Stream, ): from flydsl.compiler.kernel_function import CompilationContext @@ -643,22 +654,22 @@ def _gemm2_body( b_aux = 2 if use_nt else 0 # NT: B_q loads carry aux=2 (non-temporal hint) # block -> (m_block_idx, n_block_idx) ; e = sorted_expert_ids[m_block_idx] - n_block_idx = bx_i32 % fx.Int32(_num_n_blocks) - m_block_idx = bx_i32 // fx.Int32(_num_n_blocks) + m_block_idx = _udiv(bx_i32, _num_n_blocks) + n_block_idx = bx_i32 - m_block_idx * fx.Int32(_num_n_blocks) e = llvm.load(T.i32, _global_ptr1(arg_eids, m_block_idx * fx.Int32(4))) e = rocdl.readfirstlane(T.i32, e) m_row = m_block_idx * fx.Int32(BM) # -- buffer resources (exact num_bytes) ---------------------------------- # (A_q resource + A->LDS loads are issued by the kernel before the branch.) - ascale_rsrc = buffer_ops.create_buffer_resource( - arg_ascale, max_size=False, num_records_bytes=fx.Index(_ascale_bytes) + ascale_rsrc = buffer_ops.create_buffer_resource_from_addr( + _raw(fx.Int64(arg_ascale)), num_records_bytes=fx.Index(_ascale_bytes) ) - bq_rsrc = buffer_ops.create_buffer_resource( - arg_bq, max_size=False, num_records_bytes=fx.Index(_bq_bytes) + bq_rsrc = buffer_ops.create_buffer_resource_from_addr( + _raw(fx.Int64(arg_bq)), num_records_bytes=fx.Index(_bq_bytes) ) - bscale_rsrc = buffer_ops.create_buffer_resource( - arg_bscale, max_size=False, num_records_bytes=fx.Index(_bscale_bytes) + bscale_rsrc = buffer_ops.create_buffer_resource_from_addr( + _raw(fx.Int64(arg_bscale)), num_records_bytes=fx.Index(_bscale_bytes) ) # -- LDS base ------------------------------------------------------------ @@ -858,18 +869,11 @@ def _kloop_fence(vmcnt_atomic): has_side_effects=True, ) _s_barrier_bare() - elif const_expr(_NONATOMIC_KLOOP_VMCNT is None): - # nonatomic: plain barrier (== HIP __syncthreads); the backend inserts - # the buffer_load_lds->ds_read vmcnt wait. - rocdl.barrier() else: - # nonatomic: explicit hand-tuned fence (inline asm) -- replaces the - # backend's auto waitcnt before the ds_read with a less-conservative one. - _v = _NONATOMIC_KLOOP_VMCNT llvm.inline_asm( res=None, operands_=[], - asm_string=f"s_waitcnt vmcnt({_v}) lgkmcnt(0)", + asm_string="s_waitcnt vmcnt(16) lgkmcnt(0)", constraints="", has_side_effects=True, ) @@ -1019,9 +1023,7 @@ def _flat_bf16_epilog(accm, out_base, m_row, n_block_idx, wave, lane, N_OUT, kMC vec = Vec(accm[i][J]) for v in range_constexpr(4): row = m_row + fx.Int32(i * 16) + lane_div_16 * fx.Int32(4) + fx.Int32(v) - elem = fx.Int64(row) * fx.Int64(N_OUT) + fx.Int64( - gn - ) # i64 element index + elem = fx.Int64(row) * fx.Int64(N_OUT) + fx.Int64(gn) bf = Vec.from_elements([vec[v]], fx.Float32).to(fx.BFloat16) llvm.StoreOp(_raw(bf), _gep1(out_base, elem * fx.Int64(2))) @@ -1168,12 +1170,12 @@ def _flat_mxfp4_epilog( T.i32, packed, _raw(r[6]), _raw(r[7]), qscale, 3 ) global_col = n_block_idx * fx.Int32(BN) + col0 + blk = n_block_idx * fx.Int32(NBLK) + group q_byte = fx.Int64(out_row) * fx.Int64(N_OUT // 2) + fx.Int64( global_col // fx.Int32(2) ) - llvm.StoreOp(packed, _gep1(out_q_base, q_byte), nontemporal=True) - blk = n_block_idx * fx.Int32(NBLK) + group s_byte = fx.Int64(out_row) * fx.Int64(N_OUT // 32) + fx.Int64(blk) + llvm.StoreOp(packed, _gep1(out_q_base, q_byte), nontemporal=True) if kk == fx.Int32(0): llvm.StoreOp(arith.trunci(T.i8, e8), _gep1(out_scale_base, s_byte)) diff --git a/aiter/ops/flydsl/mxfp4_gemm1_kernels.py b/aiter/ops/flydsl/mxfp4_gemm1_kernels.py index 094e1459a3..5be4390dde 100644 --- a/aiter/ops/flydsl/mxfp4_gemm1_kernels.py +++ b/aiter/ops/flydsl/mxfp4_gemm1_kernels.py @@ -102,18 +102,21 @@ def flydsl_mxfp4_gemm1( BM, use_nt, inline_quant, D_HIDDEN, D_INTER, NE, topk ) grid = gemm1_grid(n_tokens, BM, NE=NE, TOPK=topk, INTER=D_INTER) + # gemm1 only needs base pointers (it assumes contiguity + derives sizes + # from n_tokens / compile-time consts), so pass raw data_ptr() addresses + # instead of full memref descriptors -> contiguous, coalescible kernargs. launch( - a_quant, - a_scale_sorted_shuffled, - w1_u8, - w1_scale_u8, - sorted_expert_ids, - cumsum_tensor, - m_indices, + a_quant.data_ptr(), + a_scale_sorted_shuffled.data_ptr(), + w1_u8.data_ptr(), + w1_scale_u8.data_ptr(), + sorted_expert_ids.data_ptr(), + cumsum_tensor.data_ptr(), + m_indices.data_ptr(), n_tokens, grid, - inter_sorted_quant, - inter_sorted_shuffled_scale, - hidden_states, + inter_sorted_quant.data_ptr(), + inter_sorted_shuffled_scale.data_ptr(), + hidden_states.data_ptr(), torch.cuda.current_stream(), ) diff --git a/aiter/ops/flydsl/mxfp4_gemm2_kernels.py b/aiter/ops/flydsl/mxfp4_gemm2_kernels.py index e24778894b..d9956189a0 100644 --- a/aiter/ops/flydsl/mxfp4_gemm2_kernels.py +++ b/aiter/ops/flydsl/mxfp4_gemm2_kernels.py @@ -157,18 +157,21 @@ def flydsl_mxfp4_gemm2( # launch signature stays uniform. out_scale = flat_out_scale if mxfp4out else _dummy_out_scale(flat_out.device.index) + # gemm2 only needs base pointers (assumes contiguity + derives sizes from + # compile-time consts), so pass raw data_ptr() addresses instead of full + # memref descriptors -> contiguous, coalescible kernargs (ported from gemm1). launch( - inter_sorted_quant, - inter_sorted_shuffled_scale, - w2_u8, - w2_scale_u8, - sorted_expert_ids, - cumsum_tensor, - sorted_token_ids, - sorted_weights, + inter_sorted_quant.data_ptr(), + inter_sorted_shuffled_scale.data_ptr(), + w2_u8.data_ptr(), + w2_scale_u8.data_ptr(), + sorted_expert_ids.data_ptr(), + cumsum_tensor.data_ptr(), + sorted_token_ids.data_ptr(), + sorted_weights.data_ptr(), M_logical, max_m_blocks, - flat_out, - out_scale, + flat_out.data_ptr(), + out_scale.data_ptr(), torch.cuda.current_stream(), ) From 6cbec45cdfa8058eee8f0b7be1ed9228ed0b6c38 Mon Sep 17 00:00:00 2001 From: ColorsWind <14761584+ColorsWind@users.noreply.github.com> Date: Sun, 21 Jun 2026 17:05:03 +0000 Subject: [PATCH 2/7] flydsl mxfp4 gemm2: large-M (BM128) optimizations - backend-managed K-loop waitcnt: drop the hand-tuned inline-asm vmcnt that the LLVM waitcnt pass double-inserted (the large-M scheduling bubble); let rocdl.barrier() own the load->ds_read wait. - epilog output-address strength reduction: hoist the loop-invariant row*N_OUT i64 multiply out of the per-element store loop; per-element offsets become compile-time constants folded into the store address. - unsigned index division (divui/remui) for the non-negative grid/tile/count index math, dropping signed-division sign-correction SALU. M=4096 1.11 -> 1.00 (parity with HIP), M=16384 1.10 -> 1.03; faster/parity through M=2048 unchanged. cos preserved. --- aiter/ops/flydsl/kernels/mxfp4_gemm2.py | 41 +++++++++++-------------- 1 file changed, 18 insertions(+), 23 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mxfp4_gemm2.py b/aiter/ops/flydsl/kernels/mxfp4_gemm2.py index e23f12d0f5..ba166ecee0 100644 --- a/aiter/ops/flydsl/kernels/mxfp4_gemm2.py +++ b/aiter/ops/flydsl/kernels/mxfp4_gemm2.py @@ -184,12 +184,12 @@ def _raw(v): def _udiv(a, c): cc = fx.Int32(c) if isinstance(c, int) else c - return a // cc + return fx.Int32(arith.divui(_raw(a), _raw(cc))) def _umod(a, c): cc = fx.Int32(c) if isinstance(c, int) else c - return a % cc + return fx.Int32(arith.remui(_raw(a), _raw(cc))) def _lds_ptr3(base_i32, byte_off_i32): @@ -870,14 +870,9 @@ def _kloop_fence(vmcnt_atomic): ) _s_barrier_bare() else: - llvm.inline_asm( - res=None, - operands_=[], - asm_string="s_waitcnt vmcnt(16) lgkmcnt(0)", - constraints="", - has_side_effects=True, - ) - _s_barrier_bare() + # nonatomic: plain barrier (== HIP __syncthreads); the backend inserts + # the buffer_load_lds->ds_read vmcnt wait. + rocdl.barrier() if const_expr(_K_TILES_TOTAL <= kStages): # -- KIMI/DSR fast path: K_TILES_TOTAL <= 2 (D_INTER <= 512), fully @@ -1012,20 +1007,16 @@ def _flat_bf16_epilog(accm, out_base, m_row, n_block_idx, wave, lane, N_OUT, kMC breaks store coalescing, costing more than the ~37% saved padding writes).""" lane_div_16 = lane // fx.Int32(16) lane_mod_16 = lane % fx.Int32(16) + row_base = m_row + lane_div_16 * fx.Int32(4) + gn_base = n_block_idx * fx.Int32(BN) + wave * fx.Int32(BN // 4) + lane_mod_16 + byte_base = (fx.Int64(row_base) * fx.Int64(N_OUT) + fx.Int64(gn_base)) * fx.Int64(2) for i in range_constexpr(kMChunks): for J in range_constexpr(4): - gn = ( - n_block_idx * fx.Int32(BN) - + wave * fx.Int32(BN // 4) - + fx.Int32(J * 16) - + lane_mod_16 - ) vec = Vec(accm[i][J]) for v in range_constexpr(4): - row = m_row + fx.Int32(i * 16) + lane_div_16 * fx.Int32(4) + fx.Int32(v) - elem = fx.Int64(row) * fx.Int64(N_OUT) + fx.Int64(gn) + const_off = ((i * 16 + v) * N_OUT + J * 16) * 2 bf = Vec.from_elements([vec[v]], fx.Float32).to(fx.BFloat16) - llvm.StoreOp(_raw(bf), _gep1(out_base, elem * fx.Int64(2))) + llvm.StoreOp(_raw(bf), _gep1(out_base, byte_base + fx.Int64(const_off))) def _cshuffle_flat_bf16_epilog( @@ -1121,9 +1112,11 @@ def _flat_mxfp4_epilog( 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) for mr in range_constexpr(kMChunks): # BM/16 row_local = fx.Int32(mr * 16) + m_lane - out_row = m_row + row_local for half in range_constexpr(NBLK // 4): # 2 group = wave_grp + fx.Int32(half * 4) col0 = group * fx.Int32(32) + kk * fx.Int32(8) @@ -1171,10 +1164,12 @@ def _flat_mxfp4_epilog( ) global_col = n_block_idx * fx.Int32(BN) + col0 blk = n_block_idx * fx.Int32(NBLK) + group - q_byte = fx.Int64(out_row) * fx.Int64(N_OUT // 2) + fx.Int64( - global_col // fx.Int32(2) + q_byte = ( + _q_row0 + + fx.Int64(mr * 16 * (N_OUT // 2)) + + fx.Int64(global_col // fx.Int32(2)) ) - s_byte = fx.Int64(out_row) * fx.Int64(N_OUT // 32) + fx.Int64(blk) + s_byte = _s_row0 + fx.Int64(mr * 16 * (N_OUT // 32)) + fx.Int64(blk) llvm.StoreOp(packed, _gep1(out_q_base, q_byte), nontemporal=True) if kk == fx.Int32(0): llvm.StoreOp(arith.trunci(T.i8, e8), _gep1(out_scale_base, s_byte)) From b51ec857f4138cbaedb0d7fd6f9a4a5ffa9c0472 Mon Sep 17 00:00:00 2001 From: ColorsWind <14761584+ColorsWind@users.noreply.github.com> Date: Sun, 21 Jun 2026 17:05:13 +0000 Subject: [PATCH 3/7] flydsl mxfp4 gemm1: hardware exp2 in silu Use rocdl.exp2 (v_exp_f32) in silu_mul instead of the software math.exp2 expansion (matches HIP __expf). Removes ~31% of non-MFMA VALU (the v_ldexp + range-clamp v_cmp/v_cndmask the polynomial emitted); large-M ~2-6% faster (M=4096 1.07 -> 1.05). cos preserved. --- aiter/ops/flydsl/kernels/mxfp4_gemm1.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aiter/ops/flydsl/kernels/mxfp4_gemm1.py b/aiter/ops/flydsl/kernels/mxfp4_gemm1.py index 4b9403edea..c96eba8246 100644 --- a/aiter/ops/flydsl/kernels/mxfp4_gemm1.py +++ b/aiter/ops/flydsl/kernels/mxfp4_gemm1.py @@ -216,7 +216,7 @@ def _lds_swizzle_mask(row): def _silu_mul(g, u): """silu(g)*u, matching HIP silu_mul_fast (mxfp4_gemm_common.hpp:62-65): e = __expf(-g) = exp2(-g*log2e); sig = rcpf(1+e); return g*sig*u.""" - e = (g * fx.Float32(-LOG2E)).exp2() + e = fx.Float32(rocdl.exp2(T.f32, _raw(g * fx.Float32(-LOG2E)))) sig = fx.Float32(rocdl.rcp(T.f32, _raw(fx.Float32(1.0) + e))) return g * sig * u From a595f9e6d2918f37054c2f7ffe50238254580bd1 Mon Sep 17 00:00:00 2001 From: zhutaoyu Date: Mon, 22 Jun 2026 03:53:38 +0000 Subject: [PATCH 4/7] lgkmcnt --- aiter/ops/flydsl/kernels/mxfp4_gemm2.py | 38 +++++++++++++++++++------ 1 file changed, 30 insertions(+), 8 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mxfp4_gemm2.py b/aiter/ops/flydsl/kernels/mxfp4_gemm2.py index ba166ecee0..cbc355de9b 100644 --- a/aiter/ops/flydsl/kernels/mxfp4_gemm2.py +++ b/aiter/ops/flydsl/kernels/mxfp4_gemm2.py @@ -1115,15 +1115,37 @@ def _flat_mxfp4_epilog( _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) - for mr in range_constexpr(kMChunks): # BM/16 + # Each (mr, half) block needs 8 contiguous f32 from LDS = 2x ds_read_b128. + # The blocks are independent, so software-pipeline the LDS reads: issue the + # NEXT block's ds_read before computing the CURRENT block's amax/DPP/pack, so + # the ds_read latency (the epilog's top lgkmcnt(1) stall) is hidden behind the + # ~16 ALU ops of amax + 4 DPP + 4 cvt_fp4. Distance-1 prefetch keeps only one + # extra block of registers live (matters: this kernel is 1-wave / 440 VGPR). + _blocks = [(mr, half) for mr in range(kMChunks) for half in range(NBLK // 4)] + + def _issue_load(mr, half): row_local = fx.Int32(mr * 16) + m_lane - for half in range_constexpr(NBLK // 4): # 2 - group = wave_grp + fx.Int32(half * 4) - col0 = group * fx.Int32(32) + kk * fx.Int32(8) - r = [] - for e in range_constexpr(8): - idx = row_local * fx.Int32(BN) + col0 + fx.Int32(e) - r.append(llvm.load(T.f32, _gep3(lds_base, idx * fx.Int32(4)))) + group = wave_grp + fx.Int32(half * 4) + col0 = group * fx.Int32(32) + kk * fx.Int32(8) + base_idx = row_local * fx.Int32(BN) + col0 + v0 = Vec(llvm.load(T.vec(4, T.f32), _gep3(lds_base, base_idx * fx.Int32(4)))) + v1 = Vec( + llvm.load( + T.vec(4, T.f32), + _gep3(lds_base, (base_idx + fx.Int32(4)) * fx.Int32(4)), + ) + ) + return [v0[0], v0[1], v0[2], v0[3], v1[0], v1[1], v1[2], v1[3]], group, col0 + + # prologue: issue first block's loads + _r_next, _grp_next, _col0_next = _issue_load(*_blocks[0]) + for _bi in range_constexpr(len(_blocks)): + mr, half = _blocks[_bi] + r, group, col0 = _r_next, _grp_next, _col0_next + # prefetch next block's LDS reads before consuming the current block + if _bi + 1 < len(_blocks): + _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) for e in range_constexpr(1, 8): From 48ac23b31a1279d41ab887750e545792bdcaf341 Mon Sep 17 00:00:00 2001 From: ColorsWind <14761584+ColorsWind@users.noreply.github.com> Date: Mon, 22 Jun 2026 14:51:36 +0000 Subject: [PATCH 5/7] flydsl mxfp4 gemm1: prologue vmcnt relax + BM128 scale-first ds_read --- aiter/ops/flydsl/kernels/mxfp4_gemm1.py | 60 ++++++++++++++++++++----- 1 file changed, 50 insertions(+), 10 deletions(-) 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]) From 35d2a3034d5ff21efa35acab21851d73dee096b3 Mon Sep 17 00:00:00 2001 From: ColorsWind <14761584+ColorsWind@users.noreply.github.com> Date: Mon, 22 Jun 2026 14:51:36 +0000 Subject: [PATCH 6/7] flydsl mxfp4 gemm2: float fabs/maxnum amax in mxfp4out epilog --- aiter/ops/flydsl/kernels/mxfp4_gemm2.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) 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) From 63761428f541fbe58cdc72685acb5fe82a78fdf5 Mon Sep 17 00:00:00 2001 From: ColorsWind <14761584+ColorsWind@users.noreply.github.com> Date: Mon, 22 Jun 2026 15:11:52 +0000 Subject: [PATCH 7/7] flydsl mxfp4 moe: gate gemm2 mxfp4-out solely on AITER_MXFP4_INTERMEDIATE Drop the CSV `_MXFP4OUT` kernel-name flag from the dispatch gate so the mxfp4-out gemm2 path is controlled purely by AITER_MXFP4_INTERMEDIATE=1 (plus the Kimi/DSR shape guard). The CSV no longer needs a _MXFP4OUT row. --- aiter/fused_moe.py | 30 ++---------------------------- 1 file changed, 2 insertions(+), 28 deletions(-) diff --git a/aiter/fused_moe.py b/aiter/fused_moe.py index 5ae6c195ea..a3838dcb72 100644 --- a/aiter/fused_moe.py +++ b/aiter/fused_moe.py @@ -1388,15 +1388,6 @@ def mxfp4_moe_run( # sorted_weights applied inside; scatter_reduce unnecessary. out_buf = atomic_output_buf else: - # -- MXFP4-intermediate path -- CSV-driven via a `_MXFP4OUT` g2 kernel -- - # When the tuned CSV selects `..._BM128_NONATOMIC_MXFP4OUT`, gemm2 stages - # flat_out as packed fp4 + e8m0 (mxfp4-out epilog) so the scatter_reduce - # reads ~3.8x less -> ~2.25x on that kernel. The per-expert gemm2 output is - # quantized to 4-bit BEFORE the topk reduce (lossy), so the CSV only enables - # it for M buckets where the reduce win beats the gemm2 epilog overhead - # (cold full-MoE crossover ? M 8192 on Kimi -- the CSV is the M-gate). Only - # the codegen'd Kimi/DSR nonatomic shapes (NE?{257,385}, H=7168, E=512) have - # the gemm2-mxfp4out + scatter_reduce_q kernels. mxfp4out = p2.get("mxfp4out", False) _mx_shape_ok = ( BM == 128 and D_HIDDEN == 7168 and D_INTER == 512 and NE in (257, 385) @@ -1446,11 +1437,7 @@ def mxfp4_moe_run( 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 ) @@ -1860,15 +1847,6 @@ def _mxfp4_moe_run( # sorted_weights applied inside; scatter_reduce unnecessary. out_buf = atomic_output_buf else: - # -- MXFP4-intermediate path -- CSV-driven via a `_MXFP4OUT` g2 kernel -- - # When the tuned CSV selects `..._BM128_NONATOMIC_MXFP4OUT`, gemm2 stages - # flat_out as packed fp4 + e8m0 (mxfp4-out epilog) so the scatter_reduce - # reads ~3.8x less -> ~2.25x on that kernel. The per-expert gemm2 output is - # quantized to 4-bit BEFORE the topk reduce (lossy), so the CSV only enables - # it for M buckets where the reduce win beats the gemm2 epilog overhead - # (cold full-MoE crossover ? M 8192 on Kimi -- the CSV is the M-gate). Only - # the codegen'd Kimi/DSR nonatomic shapes (NE?{257,385}, H=7168, E=512) have - # the gemm2-mxfp4out + scatter_reduce_q kernels. mxfp4out = p2.get("mxfp4out", False) # mxfp4-out (lossy fp4 intermediate) is validated only for the codegen'd # Kimi/DSR shapes. It does NOT transfer to non-Kimi INTER=256: at mid-M it @@ -1922,11 +1900,7 @@ def _mxfp4_moe_run( 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 )