diff --git a/aiter/aot/flydsl/moe.py b/aiter/aot/flydsl/moe.py index c7e7e95adf8..5dc751d6471 100644 --- a/aiter/aot/flydsl/moe.py +++ b/aiter/aot/flydsl/moe.py @@ -27,6 +27,7 @@ import os import sys import time +from typing import Optional from aiter.aot.flydsl.common import ( collect_aot_jobs, @@ -185,7 +186,14 @@ def _precompile_to_cache( out_dtype: str = "bf16", act: str = "silu", doweight_stage1: bool = False, - waves_per_eu: int = 3, + # Must match the runtime default of ``compile_mixed_moe_gemm2`` / + # ``compile_mixed_moe_gemm1`` (``Optional[int] = None``). A scalar default + # (e.g. ``3``) would make the AOT-side ``_cache_tag`` tuple disagree with + # the runtime-side tuple for any legacy kernel that does not explicitly + # pin ``waves_per_eu`` in ``get_flydsl_stage{1,2}_kernels`` (only the + # production-variant ``_persist_async_w4_cumul3`` does), causing + # ``AOT cache miss`` at runtime even though the .pkl is present on disk. + waves_per_eu: Optional[int] = None, k_batch: int = 1, b_nt: int = 2, gate_mode: str = "separated", @@ -200,6 +208,11 @@ def _precompile_to_cache( enable_bias: bool = False, stage1_fuse_quant=None, swiglu_limit: float = 0.0, + # Stage2-only kernel tuning knobs (registered by the production-variant + # entries in `get_flydsl_stage2_kernels`). Forwarded into + # `compile_flydsl_moe_stage2` for stage 2 AOT compilation. + use_async_copy: bool = False, + cu_num_mul: int = 1, **kwargs, ): """Trigger MLIR compilation by calling the runtime stage1/stage2 entry points @@ -720,6 +733,9 @@ def _make_a_user(a_dtype_user_shape): accumulate=accumulate, persist_m=_persist_m, sort_block_m=sort_block_m, + waves_per_eu=waves_per_eu, + use_async_copy=use_async_copy, + cu_num_mul=cu_num_mul, b_nt=b_nt, xcd_swizzle=xcd_swizzle, enable_bias=enable_bias, diff --git a/aiter/configs/model_configs/dsv3_fp4_tuned_fmoe.csv b/aiter/configs/model_configs/dsv3_fp4_tuned_fmoe.csv index e6b4fd6731e..f183fd5a11c 100644 --- a/aiter/configs/model_configs/dsv3_fp4_tuned_fmoe.csv +++ b/aiter/configs/model_configs/dsv3_fp4_tuned_fmoe.csv @@ -48,6 +48,8 @@ cu_num,token,model_dim,inter_dim,expert,topk,act_type,dtype,q_dtype_a,q_dtype_w, 256,16384,7168,512,257,9,ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,1082.6229,moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16,0.0,1997.4238,moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16,0.0284,3080.0467,0,0.0,0.0,flydsl_fallback 256,32768,7168,512,257,9,ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,64,0,1633.0245,moe_ck2stages_gemm1_256x64x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16,0.0,4180.9635,moe_ck2stages_gemm2_64x64x128x128_1x1_MulABScaleExpertWeightShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16,0.0291,5813.988,0,0.0,0.0,flydsl_fallback 256,2,7168,256,257,9,ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,25.0445,moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16,0.0,9.0041,moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16,0.0144,34.0486,0,0.0,0.0,flydsl_fallback +256,32768,7168,2048,64,8,ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,128,0,4655.916200000001,moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16,0.0%,4252.6419,flydsl_moe2_afp4_wfp4_bf16_t128x128x256_atomic_persist,6.8%,8908.5581,0,2591.86,395.49, +256,32768,7168,2048,64,8,ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,128,0,4655.916200000001,moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16,0.0,6110.1531,moe_ck2stages_gemm2_256x128x128x128_1x4_MulABScaleExpertWeightShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16,0.0681,10766.0693,0,0.0,0.0,flydsl_fallback 256,1,7168,2048,32,8,ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,26.8074,flydsl_moe1_afp4_wfp4_bf16_t32x64x256_fp4,0.7%,17.7858,moe_ck2stages_gemm2_256x32x128x128_1x4_MulABScaleExpertWeightShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16,6.9%,44.5932,0,15.8,31603.64, 256,2,7168,2048,32,8,ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,33.8457,flydsl_moe1_afp4_wfp4_bf16_t32x128x256_w4_fp4,0.9%,23.9086,flydsl_moe2_afp4_wfp4_bf16_t32x128x256_reduce,2.4%,57.7543,0,24.4,24402.15, 256,4,7168,2048,32,8,ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,52.869600000000005,flydsl_moe1_afp4_wfp4_bf16_t32x64x256_w3,0.0%,33.2092,flydsl_moe2_afp4_wfp4_bf16_t32x128x256_reduce_bnt2,1.9%,86.0788,0,32.74,16373.05, diff --git a/aiter/fused_moe.py b/aiter/fused_moe.py index 131a92d3c7c..44a998cdb7a 100644 --- a/aiter/fused_moe.py +++ b/aiter/fused_moe.py @@ -814,7 +814,16 @@ def _flydsl_stage2_wrapper( model_dim_pad: int = 0, **_kwargs, ): - + # `parsed` is the static per-kernel params dict registered at + # import time by `get_flydsl_stage2_kernels` (see + # aiter/ops/flydsl/moe_kernels.py) and pre-populated into + # `_KERNEL_PARAMS` by `_register_all_configs()`. Variant-specific + # knobs that live on the kernel name (e.g. the + # `..._persist_async_w4_cumul3` production variant adds + # `use_async_copy=True / waves_per_eu=4 / cu_num_mul=3`) are + # already baked into this dict, so the `parsed.get(..., default)` + # calls below pick up the registered values for that kernel name + # rather than always falling back to defaults. parsed = aiter.ops.flydsl.moe_kernels.get_flydsl_kernel_params(kernelName) if parsed is None: raise ValueError(f"Invalid FlyDSL kernel name: {kernelName}") @@ -837,6 +846,9 @@ def _flydsl_stage2_wrapper( a2_scale=a2_scale, sorted_weights=sorted_weights, sort_block_m=parsed.get("sort_block_m", 0), + waves_per_eu=parsed.get("waves_per_eu", None), + use_async_copy=parsed.get("use_async_copy", False), + cu_num_mul=parsed.get("cu_num_mul", 1), b_nt=parsed.get("b_nt", 0), persist=parsed.get("persist", None), inter_dim_pad=inter_dim_pad, diff --git a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py index e2bc48328b5..5ec24295787 100644 --- a/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py +++ b/aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py @@ -11,6 +11,7 @@ import os from contextlib import contextmanager +from typing import Optional import flydsl.compiler as flyc import flydsl.expr as fx @@ -2760,11 +2761,28 @@ def compile_mixed_moe_gemm2( inter_dim_pad: int = 0, persist_m: int = 4, sort_block_m: int = 0, - b_nt: int = 2, + waves_per_eu: Optional[int] = None, + use_async_copy: bool = False, + cu_num_mul: int = 1, + # API parity (reviewer #3): the fp4xfp4 stage2 fold doesn't consume + # `b_nt` directly -- it's only meaningful for the stage1 kernel's B-side + # cache modifier. Accept it as an ignored kwarg so callers parsing + # `_bnt{N}` registry suffixes (unifying their stage1/stage2 dispatch) + # keep working without having to special-case the fp4xfp4 path. + b_nt: int = 0, + # XCD WG-ID swizzle (lalala-sh): remap (bx_persist, by_outer) so that + # consecutive WGs land on the same XCD and share L2, which lifts the + # stage2 hit-rate on MI355X. `0` disables the remap. xcd_swizzle: int = 0, ): """Compile stage2 kernel (`moe_gemm2`) and return the compiled executable. + waves_per_eu: + - None (default): let the backend pick. + - 1..4: hint the ROCm backend with ``rocdl.waves_per_eu`` to cap or + request a specific occupancy. Useful to trade VGPR pressure for + better latency hiding on the CU. + persist_m: - > 0: legacy mode -- each CTA processes exactly persist_m consecutive M tiles. - <= 0: **persistent mode** -- grid_y = cu_num (auto-detected), each CTA @@ -2794,22 +2812,32 @@ def compile_mixed_moe_gemm2( assumed equal to `tile_m`. When set, stage2 can use a different tile_m from sorting/stage1. Requires sort_block_m % tile_m == 0. """ + del b_nt # no-op on the fp4xfp4 stage2 path; see signature note above _sort_block_m = tile_m if sort_block_m <= 0 else sort_block_m - if _sort_block_m != tile_m and _sort_block_m % tile_m != 0: + if const_expr(_sort_block_m != tile_m and _sort_block_m % tile_m != 0): raise ValueError( f"sort_block_m ({_sort_block_m}) must be a multiple of tile_m ({tile_m})" ) + # issue async X DMA early in prologue for max overlap with B/scale VMEM. + _r139_xdma_first = bool( + use_async_copy + and tile_m == 64 + and tile_n == 128 + and tile_k == 256 + and a_dtype == "fp4" + and b_dtype == "fp4" + and bool(accumulate) + ) + gpu_arch = get_hip_arch() - allocator_pong = SmemAllocator(None, arch=gpu_arch, global_sym_name="smem0") - allocator_ping = SmemAllocator(None, arch=gpu_arch, global_sym_name="smem1") - _state = {} + allocator = SmemAllocator(None, arch=gpu_arch, global_sym_name="smem0") - if a_dtype not in ("fp8", "fp16", "int8", "fp4"): + if const_expr(a_dtype not in ("fp8", "fp16", "int8", "fp4")): raise ValueError( f"a_dtype must be one of ('fp8','fp16','int8','fp4'), got {a_dtype!r}" ) - if b_dtype not in ("fp8", "fp16", "int8", "int4", "fp4"): + if const_expr(b_dtype not in ("fp8", "fp16", "int8", "int4", "fp4")): raise ValueError( f"b_dtype must be one of ('fp8','fp16','int8','int4','fp4'), got {b_dtype!r}" ) @@ -2856,20 +2884,22 @@ def compile_mixed_moe_gemm2( _expert_b_stride = (model_dim // 16) * _b_stride_n0 # K64-byte micro-step: always 64 bytes per `ku`. For fp16, this is 32 elements (2xK16 MFMA). - if (tile_k_bytes % 64) != 0: + if const_expr((tile_k_bytes % 64) != 0): raise ValueError( f"tile_k_bytes must be divisible by 64, got tile_k_bytes={tile_k_bytes} " f"(tile_k={tile_k}, elem_bytes={a_elem_bytes})" ) out_s = str(out_dtype).strip().lower() - if out_s not in ("f16", "fp16", "half", "bf16", "bfloat16", "f32", "fp32", "float"): + if const_expr( + out_s not in ("f16", "fp16", "half", "bf16", "bfloat16", "f32", "fp32", "float") + ): raise ValueError( f"out_dtype must be 'f16', 'bf16', or 'f32', got {out_dtype!r}" ) out_is_f32 = out_s in ("f32", "fp32", "float") out_is_bf16 = out_s in ("bf16", "bfloat16") - if (not bool(accumulate)) and out_is_f32: + if const_expr((not bool(accumulate)) and out_is_f32): raise ValueError( "compile_moe_gemm2(accumulate=False) only supports out_dtype in {'f16','bf16'}" ) @@ -2878,23 +2908,23 @@ def compile_mixed_moe_gemm2( is_int8 = False mfma_i32_k32 = None - if is_int8: + if const_expr(is_int8): mfma_i32_k32 = getattr(rocdl, "mfma_i32_16x16x32i8", None) or getattr( rocdl, "mfma_i32_16x16x32_i8", None ) - if mfma_i32_k32 is None: + if const_expr(mfma_i32_k32 is None): raise AttributeError( "INT8 K32 MFMA op not found: expected `rocdl.mfma_i32_16x16x32i8` " "(or `rocdl.mfma_i32_16x16x32_i8`)." ) def _x_elem_type(): - if is_f4_b: + if const_expr(is_f4_b): return T.f8 if is_f8_a else T.i8 return T.f16 if is_f16_a else (T.i8 if is_int8 else T.f8) def _w_elem_type(): - if is_f4_b: + if const_expr(is_f4_b): return T.i8 return T.f16 if is_f16_b else (T.i8 if is_int8 else T.f8) @@ -2903,7 +2933,7 @@ def _scale_elem_type(): total_threads = 256 bytes_x_per_tile = int(tile_m) * int(tile_k) * int(a_elem_bytes) - if bytes_x_per_tile % total_threads != 0: + if const_expr(bytes_x_per_tile % total_threads != 0): raise ValueError( "tile_m*tile_k*elem_bytes must be divisible by " f"{total_threads}: tile_m={tile_m}, tile_k={tile_k}, elem_bytes={a_elem_bytes}" @@ -2920,24 +2950,17 @@ def _scale_elem_type(): pad_k = 0 if _use_lds128 else 8 lds_stride = tile_k + pad_k - if a_elem_vec_pack > 1: - _eff_lds_stride = lds_stride // a_elem_vec_pack - _eff_tile_k_bytes = tile_k_bytes // a_elem_vec_pack - else: - _eff_lds_stride = lds_stride - _eff_tile_k_bytes = tile_k_bytes - - if out_is_f32: + if const_expr(out_is_f32): # Match origin/dev_a16w4: f32 output uses scalar atomics and does NOT use the CShuffle epilogue. _use_cshuffle_epilog = ( False if use_cshuffle_epilog is None else bool(use_cshuffle_epilog) ) - if _use_cshuffle_epilog: + if const_expr(_use_cshuffle_epilog): raise ValueError( "out_dtype='f32' does not support CShuffle epilogue (set use_cshuffle_epilog=False)." ) else: - if use_cshuffle_epilog is None: + if const_expr(use_cshuffle_epilog is None): _use_cshuffle_epilog = os.environ.get("FLIR_MOE_STAGE2_CSHUFFLE", "1") in ( "1", "true", @@ -2947,7 +2970,7 @@ def _scale_elem_type(): ) else: _use_cshuffle_epilog = bool(use_cshuffle_epilog) - if not _use_cshuffle_epilog: + if const_expr(not _use_cshuffle_epilog): raise ValueError( "stage2 f16 output currently requires CShuffle epilogue (FLIR_MOE_STAGE2_CSHUFFLE=1)." ) @@ -2967,47 +2990,57 @@ def _load_bias_scalar(bias_rsrc, offset): # Dynamic-shape variant: safe to reuse across (tokens/sorted_size/size_expert_ids) at runtime. # Keep a distinct ABI tag so the compile cache never mixes with historical signatures. _persistent = persist_m <= 0 - if _persistent: + # cu_num_mul: multiplies persistent-mode grid_y (CU count). =1 byte- + # identical to legacy; >1 launches more persistent CTAs each covering + # fewer M tiles. cu_num_mul=3 is the EP4 prefill sweet spot on MI355X + # (cu_num_mul=4 regresses ~2.4%). + if const_expr(not isinstance(cu_num_mul, int) or cu_num_mul < 1): + raise ValueError(f"cu_num_mul must be int >= 1, got {cu_num_mul}") + if const_expr(_persistent): from aiter.jit.utils.chip_info import get_cu_num - _cu_num = get_cu_num() + _cu_num = get_cu_num() * int(cu_num_mul) else: _cu_num = 0 _sbm_tag = "" if _sort_block_m == tile_m else f"_sbm{_sort_block_m}" _pm_tag = f"_persist_cu{_cu_num}" if _persistent else f"_pm{persist_m}" + _wpe_tag = f"_w{waves_per_eu}" if waves_per_eu is not None else "" + if const_expr(waves_per_eu is not None and not (1 <= int(waves_per_eu) <= 10)): + raise ValueError(f"waves_per_eu must be in [1, 10] or None, got {waves_per_eu}") + _num_k_tiles_per_batch = int(inter_dim) // int(tile_k) + _async_tag = "_async" if use_async_copy else "" + _cumul_tag = f"_cumul{int(cu_num_mul)}" if int(cu_num_mul) != 1 else "" _xcd_tag = f"_xcd{xcd_swizzle}" if xcd_swizzle > 0 else "" module_name = ( f"mfma_moe2_a{a_dtype}_w{b_dtype}_{out_s}_{epilog_tag}" f"_t{tile_m}x{tile_n}x{tile_k}" - f"_vscale_fix3{_pm_tag}{_sbm_tag}{_xcd_tag}" + f"_vscale_fix3_fp4opt_v1{_pm_tag}{_sbm_tag}{_wpe_tag}{_async_tag}{_cumul_tag}{_xcd_tag}" ).replace("-", "_") # -- LDS sizing (pure Python; no MLIR Context needed) --------------------- - # Ping-pong A2 tiles via separate allocators (like stage1). - _single_x_bytes = int(tile_m) * int(_eff_lds_stride) * int(a_elem_bytes) - _cshuffle_elem_bytes_s2 = 2 # f16/bf16 = 2 bytes + # Reuse a single allocation for both: + # - ping-pong A2 tiles (2 * tile_m * lds_stride * elem_bytes bytes) + # - epilogue CShuffle tile (tile_m * tile_n f16 -> 2 * tile_m * tile_n bytes) + lds_x_bytes = 2 * int(tile_m) * int(lds_stride) * int(a_elem_bytes) lds_out_bytes = ( - _cshuffle_elem_bytes_s2 * int(tile_m) * int(tile_n) - if _use_cshuffle_epilog - else 0 - ) + 2 * int(tile_m) * int(tile_n) if _use_cshuffle_epilog else 0 + ) # f16 bytes lds_tid_bytes = int(tile_m) * 4 - _input_elems = _single_x_bytes if a_elem_bytes == 1 else (_single_x_bytes // 2) - - _pong_buffer_bytes = max(_single_x_bytes, lds_out_bytes) - _ping_buffer_bytes = _single_x_bytes + # also stage sorted topk weights into LDS once per + # M-tile so the cshuffle epilogue path can prefetch them via fast LDS + # vec4 reads instead of doing a per-mi VMEM buffer_load_dwordx4 inside + # compute_tile. Only ~tile_m*4 bytes of additional LDS budget. + lds_tw_bytes = (int(tile_m) * 4) if bool(doweight_stage2) else 0 + lds_total_bytes = max(lds_x_bytes, lds_out_bytes) + lds_tid_bytes + lds_tw_bytes + lds_total_elems = lds_total_bytes if a_elem_bytes == 1 else (lds_total_bytes // 2) def x_lds_elem(): return T.f16 if is_f16_a else (T.i8 if is_int8 else T.f8) - lds_pong_offset = allocator_pong._align(allocator_pong.ptr, 16) - allocator_pong.ptr = lds_pong_offset + _pong_buffer_bytes - _lds_tid_offset_pong = allocator_pong._align(allocator_pong.ptr, 4) - allocator_pong.ptr = _lds_tid_offset_pong + lds_tid_bytes - - lds_ping_offset = allocator_ping._align(allocator_ping.ptr, 16) - allocator_ping.ptr = lds_ping_offset + _ping_buffer_bytes + lds_alloc_bytes = int(lds_total_elems) * int(a_elem_bytes) + lds_alloc_offset = allocator._align(allocator.ptr, 16) + allocator.ptr = lds_alloc_offset + lds_alloc_bytes - if True: + if const_expr(True): @flyc.kernel(name=module_name) def moe_gemm2( @@ -3030,10 +3063,9 @@ def moe_gemm2( tokens_in = arith.index_cast(ir.IndexType.get(), i32_tokens_in.ir_value()) n_in = arith.index_cast(ir.IndexType.get(), i32_n_in.ir_value()) k_in = arith.index_cast(ir.IndexType.get(), i32_k_in.ir_value()) - size_expert_ids_in = arith.index_cast( - ir.IndexType.get(), i32_size_expert_ids_in.ir_value() - ) + size_expert_ids_in = arith.index_cast(T.index, i32_size_expert_ids_in) x_elem = T.f16 if is_f16_a else (T.i8 if is_int8 else T.f8) + # For int4, weights are stored as packed bytes (i8) and unpacked to i8 packs. f32 = T.f32 i32 = T.i32 i64 = T.i64 @@ -3076,15 +3108,28 @@ def check_c_k_valid_gate(base_k): arith, c_mn=c_n_total, c_k=c_k_orig ) + # When use_async_copy is enabled with packed-elem (fp4) A, the + # DMA (buffer_load_lds) writes packed bytes into LDS, so the + # effective stride is halved. Keep legacy (non-async) path + # bit-identical by using lds_stride directly. + if const_expr(use_async_copy and a_elem_vec_pack > 1): + _eff_lds_stride = lds_stride // a_elem_vec_pack + _eff_tile_k_bytes = tile_k_bytes // a_elem_vec_pack + else: + _eff_lds_stride = lds_stride + _eff_tile_k_bytes = tile_k_bytes + shape_lds = fx.make_shape(tile_m, _eff_lds_stride) stride_lds = fx.make_stride(_eff_lds_stride, 1) layout_lds = fx.make_layout(shape_lds, stride_lds) tx = gpu.thread_id("x") - by = gpu.block_id("x") # tile along model_dim (N-dim) + by_outer = gpu.block_id("x") bx_persist = gpu.block_id("y") # persistent WG index (M-dim) if const_expr(xcd_swizzle > 0): + # Remap (bx_persist, by_outer) so consecutive WGs hit the + # same XCD and share L2 -- mirrors the stage1 XCD remap. _NUM_XCDS_S = 8 _c1_sw = arith.constant(1, index=True) _c_tn_sw = arith.constant(tile_n, index=True) @@ -3096,7 +3141,7 @@ def check_c_k_valid_gate(base_k): _c_pm_sw = arith.constant(persist_m, index=True) _gy = (size_expert_ids_in + _c_pm_sw - _c1_sw) / _c_pm_sw - _linear_id = bx_persist * _gx + by + _linear_id = bx_persist * _gx + by_outer _num_wgs = _gx * _gy _c_xcds = arith.constant(_NUM_XCDS_S, index=True) @@ -3114,34 +3159,57 @@ def check_c_k_valid_gate(base_k): _wgid_in_group = _wgid % _num_wgid_in_group bx_persist = _first_pid_m + (_wgid_in_group % _group_size_m) - by = _wgid_in_group / _group_size_m + by_outer = _wgid_in_group / _group_size_m + + by = by_outer # XOR16 swizzle parameter (in bytes; constant, power-of-two in our configs). k_blocks16 = arith.constant(_eff_tile_k_bytes // 16, index=True) layout_tx_wave_lane = fx.make_layout((4, 64), stride=(64, 1)) layout_lane16 = fx.make_layout((4, 16), stride=(16, 1)) - base_ptr_pong = allocator_pong.get_base() - base_ptr_ping = allocator_ping.get_base() - lds_x_pong = SmemPtr( - base_ptr_pong, lds_pong_offset, x_lds_elem(), shape=(_input_elems,) - ).get() - lds_x_ping = SmemPtr( - base_ptr_ping, lds_ping_offset, x_lds_elem(), shape=(_input_elems,) - ).get() + base_ptr = allocator.get_base() + lds_x_ptr = SmemPtr( + base_ptr, + lds_alloc_offset, + x_lds_elem(), + shape=(lds_total_elems,), + ) + lds_x = lds_x_ptr.get() + # Alias the same underlying LDS bytes as f16/bf16 for epilogue shuffle. lds_out = ( SmemPtr( - base_ptr_pong, - lds_pong_offset, + base_ptr, + lds_x_ptr.byte_offset, (T.bf16 if out_is_bf16 else T.f16), shape=(tile_m * tile_n,), ).get() if _use_cshuffle_epilog else None ) + + # lds_tid: alias LDS after max(x, out) for sorted_idx preload + _lds_x_b = 2 * int(tile_m) * int(lds_stride) * int(a_elem_bytes) + _lds_out_b = 2 * int(tile_m) * int(tile_n) if _use_cshuffle_epilog else 0 + _lds_tid_off = max(_lds_x_b, _lds_out_b) lds_tid = SmemPtr( - base_ptr_pong, _lds_tid_offset_pong, T.i32, shape=(tile_m,) + base_ptr, lds_x_ptr.byte_offset + _lds_tid_off, T.i32, shape=(tile_m,) ).get() + # lds_tw aliases the LDS slot immediately after + # lds_tid for staging the sorted topk weights (f32) once per M-tile. + # Only allocated when doweight_stage2; lives in a disjoint LDS region + # from lds_x / lds_out / lds_tid so it is safe to alias here. + _lds_tw_off = _lds_tid_off + int(tile_m) * 4 + lds_tw = ( + SmemPtr( + base_ptr, + lds_x_ptr.byte_offset + _lds_tw_off, + T.f32, + shape=(tile_m,), + ).get() + if doweight_stage2 + else None + ) # Buffer resources. # For dynamic memrefs, `max_size=False` cannot infer the logical size from the memref *type*, @@ -3152,6 +3220,7 @@ def check_c_k_valid_gate(base_k): # fp8/int8: 1 byte per element -> bytes = tokens*topk * K # fp4: 2 elements per byte -> bytes = tokens*topk * K / 2 c_elem_bytes = arith.constant(int(a_elem_bytes), index=True) + c_a_pack = arith.constant(int(a_elem_vec_pack), index=True) x_nbytes_idx = _div_pow2( (tokens_in * c_topk) * k_in * c_elem_bytes, int(a_elem_vec_pack) ) @@ -3196,9 +3265,9 @@ def check_c_k_valid_gate(base_k): num_valid_idx = arith.index_cast(ir.IndexType.get(), num_valid_i32) # fp16 path ignores scales completely (implicit scale=1.0). - sx_rsrc = 1 - sw_rsrc = 1 - if const_expr(not is_f16_a): + if const_expr(is_f16_a): + sx_rsrc = None + else: if const_expr(is_f4_a or is_f8_a): # A2 microscale: e8m0 in sorted layout [sorted_size, K/32]. # Caller must pre-scatter a2_scale via moe_mxfp4_sort. @@ -3216,7 +3285,9 @@ def check_c_k_valid_gate(base_k): arg_scale_x, max_size=False, num_records_bytes=sx_nbytes_i32 ) - if const_expr(not is_f16_b): + if const_expr(is_f16_b): + sw_rsrc = None + else: # Weight microscale buffer (packed i32 holding e8m0 bytes). # Use an exact descriptor size so hardware OOB checking works. kblk_w = _div_pow2(k_in, 32) # K/32 @@ -3266,15 +3337,25 @@ def check_c_k_valid_gate(base_k): _c1_p = arith.constant(1, index=True) if const_expr(_persistent): - # Expert-phase scheduling: contiguous M-tile dispatch. - # grid_y = cu_num, each CTA handles a contiguous chunk of M-tiles: - # [bx_persist * tiles_per_block, ..., (bx_persist+1) * tiles_per_block - 1] - # Adjacent blocks process adjacent M-tiles -> same expert -> B weight L2 reuse. + # Balanced contiguous M-tile dispatch: each CTA gets + # floor(T/CU) or ceil(T/CU) tiles, distributing the remainder + # across the first (T%CU) CTAs (avoids the no-work tail band + # of the old ceil mapping). Expert-contiguous ordering kept. _c_cu = arith.constant(_cu_num, index=True) _c_tm_p = arith.constant(tile_m, index=True) _num_valid_idx = arith.index_cast(ir.IndexType.get(), num_valid_i32) _total_m_tiles = (_num_valid_idx + _c_tm_p - _c1_p) / _c_tm_p - _tiles_per_block = (_total_m_tiles + _c_cu - _c1_p) / _c_cu + _tiles_per_block_base = _total_m_tiles / _c_cu + _tiles_remainder = _total_m_tiles - (_tiles_per_block_base * _c_cu) + _has_extra_tile = arith.cmpi( + CmpIPredicate.ult, bx_persist, _tiles_remainder + ) + _extra_tile = arith.select(_has_extra_tile, _c1_p, _c0_p) + _tiles_per_block = _tiles_per_block_base + _extra_tile + _start_tail = arith.select( + _has_extra_tile, bx_persist, _tiles_remainder + ) + _persist_start_tile = bx_persist * _tiles_per_block_base + _start_tail _i1 = ir.IntegerType.get_signless(1) _init_active = arith.constant(1, type=_i1) _for_persist = scf.ForOp(_c0_p, _tiles_per_block, _c1_p, [_init_active]) @@ -3296,7 +3377,7 @@ def check_c_k_valid_gate(base_k): if const_expr(_persistent): _still_active = _for_persist.inner_iter_args[0] - bx = bx_persist * _tiles_per_block + _mi_p + bx = _persist_start_tile + _mi_p else: _prev_expert_i32 = _for_persist.inner_iter_args[0] _prev_expert_b_base = _for_persist.inner_iter_args[1] @@ -3312,7 +3393,7 @@ def check_c_k_valid_gate(base_k): expert_i32 = buffer_ops.buffer_load( expert_rsrc, sort_blk, vec_width=1, dtype=T.i32 ) - expert_idx = arith.index_cast(ir.IndexType.get(), expert_i32) + expert_idx = arith.index_cast(T.index, expert_i32) exp_valid = arith.cmpi( CmpIPredicate.ult, expert_i32, arith.constant(experts, type=T.i32) ) @@ -3447,6 +3528,20 @@ def load_x(idx_i32): ) # decode routed token once (per thread's M-slice) and build a base offset. + # Async X-copy + packed-elem (fp4): the DMA uses only + # _num_dma_loads address slots, not the full num_x_loads. + # Decode only the physically needed rows to skip 2 redundant + # VMEM+decode chains per thread. + if const_expr(use_async_copy and a_elem_vec_pack > 1): + _dma_bytes_pre = 16 + _eff_bytes_pre = ( + int(tile_m) * int(_eff_lds_stride) * int(a_elem_bytes) + ) + _num_x_addr_loads = max( + 1, _eff_bytes_pre // (total_threads * _dma_bytes_pre) + ) + else: + _num_x_addr_loads = num_x_loads x_row_base_div4 = [] x_col_local_i32 = [] x_row_local = [] @@ -3455,22 +3550,26 @@ def load_x(idx_i32): x_row_local.append(row_local) x_col_local_i32.append(col_local_i32) - sorted_row_i = bx_m + row_local - fused_i = buffer_ops.buffer_load( - sorted_rsrc, sorted_row_i, vec_width=1, dtype=T.i32 - ) - t_i32 = arith.andi(fused_i, mask24) - s_i32 = arith.shrui(fused_i, arith.constant(24)) - - t_valid = arith.cmpi(CmpIPredicate.ult, t_i32, tokens_i32) - s_valid = arith.cmpi(CmpIPredicate.ult, s_i32, topk_i32) - ts_valid = arith.andi(t_valid, s_valid) - t_safe = arith.select(ts_valid, t_i32, arith.constant(0)) - s_safe = arith.select(ts_valid, s_i32, arith.constant(0)) - row_ts_i32 = t_safe * topk_i32 + s_safe - row_ts_idx = arith.index_cast(ir.IndexType.get(), row_ts_i32) - - x_row_base_div4.append(row_ts_idx * c_k_div4) + if const_expr(i < _num_x_addr_loads): + sorted_row_i = bx_m + row_local + fused_i = buffer_ops.buffer_load( + sorted_rsrc, sorted_row_i, vec_width=1, dtype=T.i32 + ) + t_i32 = arith.andi(fused_i, mask24) + s_i32 = arith.shrui(fused_i, arith.constant(24)) + + t_valid = arith.cmpi(CmpIPredicate.ult, t_i32, tokens_i32) + s_valid = arith.cmpi(CmpIPredicate.ult, s_i32, topk_i32) + ts_valid = arith.andi(t_valid, s_valid) + t_safe = arith.select(ts_valid, t_i32, arith.constant(0)) + s_safe = arith.select(ts_valid, s_i32, arith.constant(0)) + row_ts_i32 = t_safe * topk_i32 + s_safe + row_ts_idx = arith.index_cast(T.index, row_ts_i32) + + x_row_base_div4.append(row_ts_idx * c_k_div4) + else: + # Placeholder: not used in async DMA path. + x_row_base_div4.append(arith.index(0)) def load_x_tile(base_k): base_k_div4 = _div_pow2( @@ -3523,13 +3622,8 @@ def load_x_tile(base_k): _n_scale_shift_i32 = None n_intra_list = [None] * num_acc_n n_blk_list = [None] * num_acc_n - col_g_list = [None] * num_acc_n for i in range_constexpr(num_acc_n): offset = i * 16 - col_g = by_n + n_tile_base - col_g = _div_pow2(col_g, 2) + offset - col_g = col_g + lane_mod_16 - col_g_list[i] = col_g c_offset = arith.constant(offset, index=True) global_n = by_n + n_tile_base + c_offset + lane_mod_16 n_blk_list[i] = _div_pow2(global_n, 16) @@ -3543,19 +3637,18 @@ def load_x_tile(base_k): m_repeat_packed = m_repeat // pack_M num_acc_n_packed = num_acc_n // pack_N - _K_per_ku_s2 = tile_k // k_unroll - _pad_k_elems_s2 = (inter_dim_pad % tile_k) if inter_dim_pad > 0 else 0 - _pad_ku_skip_s2 = _pad_k_elems_s2 // _K_per_ku_s2 - _tail_ku_s2 = k_unroll - _pad_ku_skip_s2 - _tail_ku_packed_s2 = ( - (_tail_ku_s2 + pack_K - 1) // pack_K - if _pad_ku_skip_s2 > 0 - else None - ) - # --- B Load Logic (K64) - shared layout with preshuffle GEMM --- - def load_b_packs_k64(base_k, ku: int, ni: int): - """Load one K64-byte B micro-step: single 16B load, split into 2x i64.""" + def load_b_packs_k64( + base_k, ku: int, ni: int, *, n_blk_p=None, n_intra_p=None + ): + """Load one K64-byte B micro-step: single 16B load, split into 2x i64. + + `n_blk_p` / `n_intra_p` allow callers to override the per-N + list. When ``None`` we fall back to the body-level defaults + (``n_blk_list`` / ``n_intra_list``). + """ + _blk = n_blk_p if n_blk_p is not None else n_blk_list + _intra = n_intra_p if n_intra_p is not None else n_intra_list base_k_bytes = base_k * arith.constant( int(b_elem_bytes), index=True ) @@ -3569,10 +3662,10 @@ def load_b_packs_k64(base_k, ku: int, ni: int): # compile-time constants -> shift/mul, no Barrett. idx_pack = ( _expert_b_base - + n_blk_list[ni] * arith.constant(_b_stride_n0, index=True) + + _blk[ni] * arith.constant(_b_stride_n0, index=True) + k0 * arith.constant(_b_stride_k0, index=True) + k1 * arith.constant(_b_stride_klane, index=True) - + n_intra_list[ni] * arith.constant(_b_stride_nlane, index=True) + + _intra[ni] * arith.constant(_b_stride_nlane, index=True) ) vec_elems = kpack_bytes // int(b_elem_bytes) @@ -3585,7 +3678,6 @@ def load_b_packs_k64(base_k, ku: int, ni: int): vec_elems=vec_elems, elem_bytes=b_elem_bytes, offset_in_bytes=(b_elem_bytes == 1), - cache_modifier=b_nt, ) b_i64x2 = vector.bitcast(vec2_i64, b16) b0 = vector.extract( @@ -3596,42 +3688,74 @@ def load_b_packs_k64(base_k, ku: int, ni: int): ) return b0, b1 - def load_b_tile(base_k, ku_limit=k_unroll): + def load_b_tile(base_k, *, n_blk_p=None, n_intra_p=None): b_tile = [] - for ku in range_constexpr(ku_limit): + for ku in range_constexpr(k_unroll): packs0 = [] packs1 = [] for ni in range_constexpr(num_acc_n): - b0, b1 = load_b_packs_k64(base_k, ku, ni) + b0, b1 = load_b_packs_k64( + base_k, + ku, + ni, + n_blk_p=n_blk_p, + n_intra_p=n_intra_p, + ) packs0.append(b0) packs1.append(b1) b_tile.append((packs0, packs1)) return b_tile _b_split_enabled = k_unroll >= 2 - _b_split_ku = k_unroll // 2 if _b_split_enabled else k_unroll + # Asymmetric b_lo/b_hi split for prod prefill: b_lo loads + # 3/4 of K (extra ku in prologue, overlaps X DMA), b_hi loads + # 1/4 (smaller burst before K-loop frees VMEM slots so the + # first MFMAs don't stall on VMEM issue). + _is_prefill_shape = ( + int(tile_m) == 64 + and int(tile_n) == 128 + and int(tile_k) == 256 + and bool(use_async_copy) + and a_elem_vec_pack > 1 + ) + if const_expr(_b_split_enabled and _is_prefill_shape and k_unroll >= 4): + _b_split_ku = k_unroll - 1 + else: + _b_split_ku = k_unroll // 2 if _b_split_enabled else k_unroll - def load_b_tile_lo(base_k): + def load_b_tile_lo(base_k, *, n_blk_p=None, n_intra_p=None): """Load first half of B tile (ku < _b_split_ku).""" b_tile = [] for ku in range_constexpr(_b_split_ku): packs0 = [] packs1 = [] for ni in range_constexpr(num_acc_n): - b0, b1 = load_b_packs_k64(base_k, ku, ni) + b0, b1 = load_b_packs_k64( + base_k, + ku, + ni, + n_blk_p=n_blk_p, + n_intra_p=n_intra_p, + ) packs0.append(b0) packs1.append(b1) b_tile.append((packs0, packs1)) return b_tile - def load_b_tile_hi(base_k): + def load_b_tile_hi(base_k, *, n_blk_p=None, n_intra_p=None): """Load second half of B tile (ku >= _b_split_ku).""" b_tile = [] for ku in range_constexpr(_b_split_ku, k_unroll): packs0 = [] packs1 = [] for ni in range_constexpr(num_acc_n): - b0, b1 = load_b_packs_k64(base_k, ku, ni) + b0, b1 = load_b_packs_k64( + base_k, + ku, + ni, + n_blk_p=n_blk_p, + n_intra_p=n_intra_p, + ) packs0.append(b0) packs1.append(b1) b_tile.append((packs0, packs1)) @@ -3659,11 +3783,14 @@ def _apply_k_shift(scale_vec, k_shift_bits): return vector.from_elements(T.vec(1, T.i32), [val]) return scale_vec - def load_b_scale_tile( - base_k, k_shift_bits=0, ku_packed_limit=k_unroll_packed - ): + def load_b_scale_tile(base_k, k_shift_bits=0, *, by_n_p=None): + """Load B-scale tile. + + `by_n_p` overrides the body-default `by_n`. + """ + _byn = by_n_p if by_n_p is not None else by_n b_scale_tile = [] - for ku in range_constexpr(ku_packed_limit): + for ku in range_constexpr(k_unroll_packed): for ni in range_constexpr(num_acc_n_packed): scale = load_scale( arg_scale_w, @@ -3673,7 +3800,7 @@ def load_b_scale_tile( ni + _div_pow2( _div_pow2( - expert_off_idx + by_n + n_tile_base, + expert_off_idx + _byn + n_tile_base, _scale_pack_n, ), 16, @@ -3683,11 +3810,9 @@ def load_b_scale_tile( b_scale_tile.append(scale) return b_scale_tile - def load_a_scale_tile( - base_k, k_shift_bits=0, ku_packed_limit=k_unroll_packed - ): + def load_a_scale_tile(base_k, k_shift_bits=0): a_scale_tile = [] - for ku in range_constexpr(ku_packed_limit): + for ku in range_constexpr(k_unroll_packed): for mi in range_constexpr(m_repeat_packed): scale = load_scale( arg_scale_x, @@ -3700,25 +3825,17 @@ def load_a_scale_tile( a_scale_tile.append(scale) return a_scale_tile - def prefetch_ab_scale_tile( - base_k, k_shift_bits=0, ku_packed_limit=k_unroll_packed - ): + def prefetch_ab_scale_tile(base_k, k_shift_bits=0, *, by_n_p=None): return [ - load_a_scale_tile( - base_k, k_shift_bits, ku_packed_limit=ku_packed_limit - ), - load_b_scale_tile( - base_k, k_shift_bits, ku_packed_limit=ku_packed_limit - ), + load_a_scale_tile(base_k, k_shift_bits), + load_b_scale_tile(base_k, k_shift_bits, by_n_p=by_n_p), ] vec8_x = T.vec(vec8_elems, x_elem) vec4_x_lds = T.vec(vec4_elems, x_elem) - # ---- Pipeline helpers: store X tile to LDS (unused in DMA path) ---- - _lds_base_zero = arith.index(0) - - def store_x_tile_to_lds(vec_x_in_parts, lds_buffer): + # ---- Pipeline helpers: store X tile to LDS with ping-pong base ---- + def store_x_tile_to_lds(vec_x_in_parts, lds_base): for i in range_constexpr(num_x_loads): row_local = x_row_local[i] col_local_i32 = x_col_local_i32[i] @@ -3726,14 +3843,14 @@ def store_x_tile_to_lds(vec_x_in_parts, lds_buffer): lds_store_16b_xor16( arith, vector, - lds_memref=lds_buffer, + lds_memref=lds_x, vec16_ty=vec16_x, layout_lds=layout_lds, row_local=row_local, col_local_i32=col_local_i32, tx_c4=arith.index(4), k_blocks16=k_blocks16, - lds_base=_lds_base_zero, + lds_base=lds_base, vec_part_i32x4=vec_x_in_parts[i], elem_bytes=elem_bytes, ) @@ -3741,14 +3858,14 @@ def store_x_tile_to_lds(vec_x_in_parts, lds_buffer): lds_store_8b_xor16( arith, vector, - lds_memref=lds_buffer, + lds_memref=lds_x, vec8_ty=vec8_x, layout_lds=layout_lds, row_local=row_local, col_local_i32=col_local_i32, tx_c4=arith.index(4), k_blocks16=k_blocks16, - lds_base=_lds_base_zero, + lds_base=lds_base, vec_part_i32x2=vec_x_in_parts[i], elem_bytes=elem_bytes, ) @@ -3756,20 +3873,86 @@ def store_x_tile_to_lds(vec_x_in_parts, lds_buffer): lds_store_4b_xor16( arith, vector, - lds_memref=lds_buffer, + lds_memref=lds_x, vec4_ty=vec4_x_lds, layout_lds=layout_lds, row_local=row_local, col_local_i32=col_local_i32, tx_c4=arith.index(4), k_blocks16=k_blocks16, - lds_base=_lds_base_zero, + lds_base=lds_base, vec_part_i32x1=vec_x_in_parts[i], elem_bytes=elem_bytes, ) + # Async X-copy (buffer_load_lds): DMA X tile gmem -> LDS + # bypassing per-thread reg staging. Mirrors stage1 but uses + # stage2's single lds_x memref with ping/pong offsets. + if const_expr(use_async_copy): + _dma_bytes = 16 + _wave_size = 64 + _eff_bytes_per_buffer = ( + int(tile_m) * int(_eff_lds_stride) * int(a_elem_bytes) + ) + _num_dma_loads = max( + 1, _eff_bytes_per_buffer // (total_threads * _dma_bytes) + ) + _c_a_elem_bytes_dma = arith.constant(int(a_elem_bytes), index=True) + _c_wave_dma_bytes = arith.constant( + _wave_size * _dma_bytes, index=True + ) + + def dma_x_tile_to_lds(base_k, lds_base): + c4_idx = arith.index(4) + base_k_div4 = ( + (base_k / c_a_pack) + * arith.constant(int(elem_bytes), index=True) + ) / arith.index(4) + + lds_ptr_i64 = None + for i in range_constexpr(_num_dma_loads): + row_local_i = x_row_local[i] + col_local_i32_i = x_col_local_i32[i] + col_local_sw = swizzle_xor16( + row_local_i, col_local_i32_i * c4_idx, k_blocks16 + ) + row_k_dw = x_row_base_div4[i] + base_k_div4 + global_byte_idx = row_k_dw * c4_idx + col_local_sw + global_offset = arith.index_cast(T.i32, global_byte_idx) + + if const_expr(i == 0): + lds_addr = ( + memref.extract_aligned_pointer_as_index(lds_x) + + lds_base * _c_a_elem_bytes_dma + + wave_id * _c_wave_dma_bytes + ) + lds_ptr_i64 = rocdl.readfirstlane( + T.i64, arith.index_cast(T.i64, lds_addr) + ) + else: + lds_ptr_i64 = lds_ptr_i64 + arith.constant( + total_threads * _dma_bytes, type=T.i64 + ) + + lds_ptr_type = ir.Type.parse("!llvm.ptr<3>") + lds_ptr = llvm.inttoptr(lds_ptr_type, lds_ptr_i64) + + rocdl.raw_ptr_buffer_load_lds( + x_rsrc, + lds_ptr, + arith.constant(_dma_bytes, type=T.i32), + global_offset, + arith.constant(0, type=T.i32), + arith.constant(0, type=T.i32), + arith.constant(0, type=T.i32), + ) + + def prefetch_x_to_lds(base_k, lds_base): + dma_x_tile_to_lds(base_k, lds_base) + # --- A LDS load helper for K64 (load 16B once, extract 2x i64 halves) --- - def lds_load_packs_k64(curr_row_a_lds, col_base, lds_buffer): + def lds_load_packs_k64(curr_row_a_lds, col_base, lds_base): + # Swizzle in bytes, then convert to element offset for memref indexing. col_base_swz_bytes = swizzle_xor16( curr_row_a_lds, col_base, k_blocks16 ) @@ -3778,8 +3961,10 @@ def lds_load_packs_k64(curr_row_a_lds, col_base, lds_buffer): if elem_bytes == 1 else (col_base_swz_bytes / arith.index(2)) ) + # Pass as list so layout_utils.crd2idx uses static arith path idx_a16 = crd2idx([curr_row_a_lds, col_base_swz], layout_lds) - loaded_a16 = vector.load_op(vec16_x, lds_buffer, [idx_a16]) + idx_a16 = idx_a16 + lds_base + loaded_a16 = vector.load_op(vec16_x, lds_x, [idx_a16]) a_i64x2 = vector.bitcast(vec2_i64, loaded_a16) a0 = vector.extract( a_i64x2, static_position=[0], dynamic_position=[] @@ -3792,7 +3977,7 @@ def lds_load_packs_k64(curr_row_a_lds, col_base, lds_buffer): def compute_tile( acc_in, b_tile_in, - lds_buffer, + lds_base, a_scale=None, b_scale=None, *, @@ -3800,8 +3985,16 @@ def compute_tile( a0_prefetch=None, a1_prefetch=None, b_hi_loader=None, - ku_count=k_unroll, + n_scale_shift_p=None, ): + # n_scale_shift_p overrides body-default `_n_scale_shift_i32` + # so the K-outer/N-inner pipeline can apply the right shift + # for each ni slot. + _nss = ( + n_scale_shift_p + if n_scale_shift_p is not None + else _n_scale_shift_i32 + ) if const_expr(b_hi_loader is not None): b_tile_full = [None] * k_unroll for i in range_constexpr(_b_split_ku): @@ -3822,27 +4015,49 @@ def compute_tile( bias.append(_load_bias_scalar(bias_rsrc, bias_offset)) tw_pf = None if const_expr(doweight_stage2): + # When lds_tw is staged in the prologue (prod + # prefill default), prefetch topk-weight tile via + # vec4 LDS reads (LDS ~10cyc vs VMEM ~500cyc), + # saving m_repeat VMEM ops from compute_tile. + # Falls back to vec4 VMEM if lds_tw unavailable. tw_pf = [] lane_div_16_mul4_pf = lane_div_16 * arith.index(4) - ii_idx_list_pf = [ - arith.constant(ii, index=True) for ii in range(4) - ] - for mi in range_constexpr(m_repeat): - mi_base_pf = arith.constant(mi * 16, index=True) - for ii in range_constexpr(4): - row_off_pf = ( - lane_div_16_mul4_pf + ii_idx_list_pf[ii] + _vec4_f32_pf = T.vec(4, f32) + if const_expr(lds_tw is not None): + for mi in range_constexpr(m_repeat): + mi_base_pf = arith.constant(mi * 16, index=True) + lds_row_pf = mi_base_pf + lane_div_16_mul4_pf + tw_v4 = vector.load_op( + _vec4_f32_pf, lds_tw, [lds_row_pf] ) - row_in_tile_pf = mi_base_pf + row_off_pf - sorted_row_pf = bx_m + row_in_tile_pf - tw_pf.append( - buffer_ops.buffer_load( - sorted_w_rsrc, - sorted_row_pf, - vec_width=1, - dtype=f32, + for ii in range_constexpr(4): + tw_pf.append( + vector.extract( + tw_v4, + static_position=[ii], + dynamic_position=[], + ) ) + else: + for mi in range_constexpr(m_repeat): + mi_base_pf = arith.constant(mi * 16, index=True) + base_row_pf = ( + bx_m + mi_base_pf + lane_div_16_mul4_pf ) + tw_v4 = buffer_ops.buffer_load( + sorted_w_rsrc, + base_row_pf, + vec_width=4, + dtype=f32, + ) + for ii in range_constexpr(4): + tw_pf.append( + vector.extract( + tw_v4, + static_position=[ii], + dynamic_position=[], + ) + ) epilogue_pf = (None, tw_pf, bias) c0_i64 = arith.constant(0, type=T.i64) @@ -3859,12 +4074,31 @@ def pack_i64x4_to_i32x8(x0, x1, x2, x3): _pack_K_shift = (pack_K - 1).bit_length() _pack_K_mask = pack_K - 1 + # Emit AMDGPU `disable_xdl_arb_stall` before b_hi prefetch + # + s_setprio(1) so the high-prio K-loop drains MFMAs back- + # to-back on gfx950 (XDL arbiter normally throttles MFMA + # issue against sibling waves on the same SIMD). + _xdl_arb_hint = ( + int(tile_m) == 64 + and int(tile_n) == 128 + and int(tile_k) == 256 + and bool(use_async_copy) + and a_elem_vec_pack > 1 + ) + if const_expr(_xdl_arb_hint): + rocdl.disable_xdl_arb_stall() + if const_expr(b_hi_loader is not None): _b_hi = b_hi_loader() for _bhi_i in range_constexpr(len(_b_hi)): b_tile_full[_b_split_ku + _bhi_i] = _b_hi[_bhi_i] - for k_idx in range_constexpr(ku_count): + # Raise wave priority around MFMA-heavy compute so this wave + # is favored on the EU vs other waves doing async copies. + # Mirrors stage1's pattern (lines ~1410/1438). + rocdl.s_setprio(1) + + for k_idx in range_constexpr(k_unroll): ku128 = k_idx >> _pack_K_shift ikxdl = k_idx & _pack_K_mask @@ -3881,6 +4115,48 @@ def pack_i64x4_to_i32x8(x0, x1, x2, x3): a_scale_val = arith.shrui( a_scale_val, _m_scale_shift_i32 ) + # Hoist A-side LDS loads + a128 pack OUT of the + # ni loop (addr depends only on k_idx,mi,imxdl). + # Previously emitted num_acc_n_packed times per + # (mi,imxdl); hoisting removes redundant IR and + # groups LDS reads at the head of the mi block. + a128_list = [None] * pack_M + for imxdl in range_constexpr(pack_M): + col_base0 = col_base + mi_idx = mi * pack_M + imxdl + mi_val = arith.constant(mi_idx * 16, index=True) + curr_row_a_lds = row_a_lds + mi_val + + if const_expr( + (a0_prefetch is not None) + and (k_idx == 0) + and (mi_idx == 0) + ): + a0, a1 = a0_prefetch + elif const_expr( + (a1_prefetch is not None) + and (k_idx == 1) + and (mi_idx == 0) + ): + a0, a1 = a1_prefetch + else: + a0, a1 = lds_load_packs_k64( + curr_row_a_lds, col_base0, lds_base + ) + + if const_expr(is_f8_a): + col_base1 = col_base + 64 + a2, a3 = lds_load_packs_k64( + curr_row_a_lds, col_base1, lds_base + ) + a128_list[imxdl] = pack_i64x4_to_i32x8( + a0, a1, a2, a3 + ) + else: + a128_list[imxdl] = pack_i64x4_to_i32x8( + a0, a1, c0_i64, c0_i64 + ) + for ni in range_constexpr(num_acc_n_packed): b_scale_i32 = b_scale[ku128 * num_acc_n_packed + ni] b_scale_val = vector.extract( @@ -3888,55 +4164,35 @@ def pack_i64x4_to_i32x8(x0, x1, x2, x3): static_position=[0], dynamic_position=[], ) - if const_expr(_n_scale_shift_i32 is not None): - b_scale_val = arith.shrui( - b_scale_val, _n_scale_shift_i32 + if const_expr(_nss is not None): + b_scale_val = arith.shrui(b_scale_val, _nss) + + # Hoist B-pack i64x4_to_i32x8 outside the + # imxdl loop — b128 is a function only of + # (ni, inxdl), not (mi, imxdl). Build once per + # (ni, inxdl) and reuse across all mi/imxdl. + b128_list = [None] * pack_N + for inxdl in range_constexpr(pack_N): + ni_idx = ni * pack_N + inxdl + b0 = b_packs0[ni_idx] + b1 = b_packs1[ni_idx] + b128_list[inxdl] = pack_i64x4_to_i32x8( + b0, b1, c0_i64, c0_i64 ) for imxdl in range_constexpr(pack_M): - col_base0 = col_base mi_idx = mi * pack_M + imxdl - mi_val = arith.constant(mi_idx * 16, index=True) - curr_row_a_lds = row_a_lds + mi_val - - if const_expr( - (a0_prefetch is not None) - and (k_idx == 0) - and (mi_idx == 0) - ): - a0, a1 = a0_prefetch - elif const_expr( - (a1_prefetch is not None) - and (k_idx == 1) - and (mi_idx == 0) - ): - a0, a1 = a1_prefetch - else: - a0, a1 = lds_load_packs_k64( - curr_row_a_lds, col_base0, lds_buffer - ) - - if const_expr(is_f8_a): - col_base1 = col_base + 64 - a2, a3 = lds_load_packs_k64( - curr_row_a_lds, col_base1, lds_buffer - ) - a128 = pack_i64x4_to_i32x8(a0, a1, a2, a3) - else: - a128 = pack_i64x4_to_i32x8( - a0, a1, c0_i64, c0_i64 - ) + a128 = a128_list[imxdl] for inxdl in range_constexpr(pack_N): ni_idx = ni * pack_N + inxdl - b0 = b_packs0[ni_idx] - b1 = b_packs1[ni_idx] - b128 = pack_i64x4_to_i32x8( - b0, b1, c0_i64, c0_i64 - ) + b128 = b128_list[inxdl] acc_idx = mi_idx * num_acc_n + ni_idx + # Removed per-MFMA sched_barrier(0): with + # s_setprio(1) around the MFMA loop, LLVM can + # interleave MFMAs with next b_hi VMEM+DSREADs. acc_list[acc_idx] = ( rocdl.mfma_scale_f32_16x16x128_f8f6f4( mfma_res_ty, @@ -3954,71 +4210,69 @@ def pack_i64x4_to_i32x8(x0, x1, x2, x3): ) ) + # Keep s_setprio(1) across the post-compute barrier and + # the VMEM-heavy prefetch_x/scale phase so this wave stays + # favored vs sibling WGs in MFMA loops. Next compute_tile's + # setprio(1) becomes a no-op; the cshuffle epilog handles + # its own prio (bumps to 3, drops to 0 at end). return acc_list, epilogue_pf # ---------------- 2-stage pipeline (ping-pong LDS + B tile prefetch) ---------------- - # ---- Async DMA: GMEM -> LDS (bypasses VGPR, like stage1) ---- - _dma_bytes = 16 - _wave_size = 64 - _eff_bytes_per_buffer = ( - int(tile_m) * int(_eff_lds_stride) * int(a_elem_bytes) - ) - _num_dma_loads = max( - 1, _eff_bytes_per_buffer // (total_threads * _dma_bytes) - ) - - def dma_x_tile_to_lds(base_k, lds_buffer): - c4_idx = arith.index(4) - base_k_div4 = _div_pow2( - _div_pow2(base_k, int(a_elem_vec_pack)) - * arith.constant(int(a_elem_bytes), index=True), - 4, - ) - - lds_ptr_i64 = None - for i in range_constexpr(_num_dma_loads): - row_local_i = x_row_local[i] - col_local_i32_i = x_col_local_i32[i] - col_local_sw = swizzle_xor16( - row_local_i, col_local_i32_i * c4_idx, k_blocks16 - ) - row_k_dw = x_row_base_div4[i] + base_k_div4 - global_byte_idx = row_k_dw * c4_idx + col_local_sw - global_offset = arith.index_cast(T.i32, global_byte_idx) - - if const_expr(i == 0): - lds_addr = memref.extract_aligned_pointer_as_index( - lds_buffer - ) + wave_id * arith.constant( - _wave_size * _dma_bytes, index=True - ) - lds_ptr_i64 = rocdl.readfirstlane( - T.i64, arith.index_cast(T.i64, lds_addr) - ) - else: - lds_ptr_i64 = lds_ptr_i64 + arith.constant( - total_threads * _dma_bytes, type=T.i64 - ) - - lds_ptr_type = ir.Type.parse("!llvm.ptr<3>") - lds_ptr = llvm.inttoptr(lds_ptr_type, lds_ptr_i64) - - rocdl.raw_ptr_buffer_load_lds( - x_rsrc, - lds_ptr, - arith.constant(_dma_bytes, type=T.i32), - global_offset, - arith.constant(0, type=T.i32), - arith.constant(0, type=T.i32), - arith.constant(0, type=T.i32), - ) - - def prefetch_x_to_lds(base_k, lds_buffer): - dma_x_tile_to_lds(base_k, lds_buffer) + lds_tile_elems = arith.constant(tile_m * lds_stride, index=True) + lds_base_cur = arith.index(0) + lds_base_nxt = lds_tile_elems rocdl.sched_barrier(0) def hot_loop_scheduler(): + # - MFMA group size per "slot": num_acc_n + # - Total MFMA per tile: (2*K32 per K64) * k_unroll * m_repeat * num_acc_n + # - We emit (mfma_group + dsrd + mfma_group) per scheduler iteration. + mfma_group = num_acc_n + mfma_total = (k_unroll * 2) * m_repeat * mfma_group + mfma_per_iter = 2 * mfma_group + sche_iters = ( + 0 if mfma_per_iter == 0 else (mfma_total // mfma_per_iter) + ) + + rocdl.sched_dsrd(2) + rocdl.sched_mfma(1) + if const_expr(tile_m == 16): + rocdl.sched_vmem(1) + rocdl.sched_mfma(1) + if const_expr(tile_m == 16): + rocdl.sched_vmem(1) + if const_expr(num_acc_n < 4): + rocdl.sched_dsrd(1) + rocdl.sched_mfma(1) + if const_expr(tile_m == 16): + rocdl.sched_vmem(1) + rocdl.sched_dsrd(1) + rocdl.sched_mfma(1) + if const_expr(tile_m == 16): + rocdl.sched_vmem(1) + rocdl.sched_mfma(1) + + # DS-write hints (match total A LDS-store micro-ops/thread): + # async path's raw_ptr_buffer_load_lds writes LDS inside the + # buffer-load, so there is no register-mediated ds_write. + # Legacy dswr_tail=num_x_loads overcounts; zero it for async. + if const_expr(use_async_copy): + dswr_tail = 0 + else: + dswr_tail = num_x_loads + if const_expr(dswr_tail > sche_iters): + dswr_tail = sche_iters + dswr_start = sche_iters - dswr_tail + + for sche_i in range_constexpr(sche_iters): + rocdl.sched_vmem(1) + rocdl.sched_mfma(mfma_group) + rocdl.sched_dsrd(1) + rocdl.sched_mfma(mfma_group) + if const_expr(dswr_tail > 0 and sche_i >= dswr_start - 1): + rocdl.sched_dswr(1) + rocdl.sched_barrier(0) def _k_shift_bits(k_py): @@ -4031,43 +4285,139 @@ def _k_base(k_py): # Preload sorted_idx into lds_tid for epilogue precompute_row # (N-independent; placed before N-tile loop so it's done once per M-tile.) + # + # For the buffer-atomic prefill path (_use_buf_atomic_pre), + # pre-compute `(t & 0xFFFFFF) * row_stride_bytes` and OOB-mask + # invalid rows here in the prologue (overlapped with B/scale/X + # VMEM); the cshuffle precompute_row then needs only an LDS load. + _row_stride_bytes_pre = int(model_dim) * int(out_elem_bytes) + _use_buf_atomic_pre = bool(accumulate) and ( + _row_stride_bytes_pre <= 16384 + ) _c_tile_m_idx = arith.constant(tile_m, index=True) _tid_in_range = arith.cmpi(CmpIPredicate.ult, tx, _c_tile_m_idx) - _if_tid = scf.IfOp(_tid_in_range) - with ir.InsertionPoint(_if_tid.then_block): - _tid_row = bx_m + tx - _tid_val = buffer_ops.buffer_load( - sorted_rsrc, _tid_row, vec_width=1, dtype=T.i32 - ) - _tid_vec1 = vector.from_elements(T.vec(1, T.i32), [_tid_val]) - vector.store(_tid_vec1, lds_tid, [tx]) - scf.YieldOp([]) + # Defer lds_tid prologue (sorted_idx/sorted_w VMEM + LDS + # stores) to AFTER the long-latency X-DMA + B/scale loads, + # hiding lds_tid's scalar VMEM behind the dominant VMEM phase. + # lds_tid is consumed only by the cshuffle epilogue; the + # post-prologue barrier provides the needed LDS visibility. + _r216_defer_tid = bool( + _r139_xdma_first + and use_async_copy + and int(tile_m) == 64 + and int(tile_n) == 128 + and int(tile_k) == 256 + ) - gpu.barrier() + def _emit_tid_lds_prologue(): + # Merge lds_tid + lds_tw prologue store guards into a + # single scf.IfOp (same `tx pong. + # Removed the post-lds_tid gpu.barrier(): lds_tid is consumed + # only by the cshuffle epilogue (precompute_row) which runs + # after many K-loop barriers that already sync lds_tid. The + # removed barrier was a per-persistent-iteration cost. + + # Prologue -- B-first. k0 = arith.index(0) + k0_bk = k0 + if const_expr(_r139_xdma_first): + prefetch_x_to_lds(k0, lds_base_cur) + rocdl.sched_barrier(0) if const_expr(_b_split_enabled): - b_cur = load_b_tile_lo(k0) + b_cur = load_b_tile_lo(k0_bk) else: - b_cur = load_b_tile(k0) + b_cur = load_b_tile(k0_bk) a_scale_pong, b_scale_pong = prefetch_ab_scale_tile( _k_base(0), _k_shift_bits(0) ) rocdl.sched_barrier(0) - prefetch_x_to_lds(k0, lds_x_pong) - rocdl.s_waitcnt(0) + if const_expr(not _r139_xdma_first): + if const_expr(use_async_copy): + prefetch_x_to_lds(k0, lds_base_cur) + else: + x_regs0 = load_x_tile(k0) + store_x_tile_to_lds(x_regs0, lds_base_cur) + elif const_expr(not use_async_copy): + x_regs0 = load_x_tile(k0) + store_x_tile_to_lds(x_regs0, lds_base_cur) + if const_expr(_r216_defer_tid): + # Deferred emit: see _r216_defer_tid above for rationale. + _emit_tid_lds_prologue() + rocdl.sched_barrier(0) gpu.barrier() acc = [acc_init] * num_acc_n * m_repeat + lds_base_pong = lds_base_cur + lds_base_ping = lds_base_nxt - # Cross-tile A0+A1 LDS prefetch from pong buffer. + # Cross-tile A0+A1 LDS prefetch: issue both ds_reads back-to-back + # (??2) so LDS bandwidth is fully utilized and the second read + # completes during MFMA #1/#2 execution. a0_prefetch_pong = lds_load_packs_k64( - row_a_lds, col_offset_base, lds_x_pong + row_a_lds, col_offset_base, lds_base_pong ) _a1_col_base = col_offset_base + 128 // a_elem_vec_pack a1_prefetch_pong = ( - lds_load_packs_k64(row_a_lds, _a1_col_base, lds_x_pong) + lds_load_packs_k64(row_a_lds, _a1_col_base, lds_base_pong) if pack_K >= 2 else None ) @@ -4077,7 +4427,7 @@ def _k_base(k_py): # IMPORTANT: for odd number of K tiles, leave **1** tail tile; for even, leave **2**. # Otherwise the 2-tile tail below would double-count the last tile when num_tiles is odd # (e.g. inter_dim=192, tile_k=64 -> 3 tiles). - num_k_tiles_py = int(inter_dim) // int(tile_k) + num_k_tiles_py = _num_k_tiles_per_batch odd_k_tiles = (num_k_tiles_py % 2) == 1 tail_tiles = 1 if odd_k_tiles else 2 k_main2_py = (num_k_tiles_py - tail_tiles) * int(tile_k) @@ -4086,7 +4436,8 @@ def _k_base(k_py): c2_tile_k = arith.constant(tile_k * 2, index=True) b_pong = b_cur - k0_pong_bk = k0 + # b_pong/b_hi loader uses packed-byte K units (see note above). + k0_pong_bk = k0_bk # Only emit the scf.for when there are actually iterations to run. # When k_main2_py == 0 the loop body is empty; emitting an scf.for @@ -4098,25 +4449,30 @@ def _make_b_hi_loader(base_k): if const_expr(k_main2_py > 0): for k_iv_py in range_constexpr(0, k_main2_py, tile_k * 2): - rocdl.sched_barrier(0) k_iv = arith.index(k_iv_py) next_k1 = k_iv + tile_k next_k1_bk = next_k1 // 2 - # DMA X(next_k1) -> ping (non-blocking, overlaps with compute) - prefetch_x_to_lds(next_k1, lds_x_ping) + # Issue scales BEFORE B-VMEM in the steady-state K-loop: + # scales are 1-dword (low VMEM latency) and feed + # `vector.extract` at the head of compute_tile's MFMA + # loop, so landing them first hides B's longer loads. + if const_expr(use_async_copy): + prefetch_x_to_lds(next_k1, lds_base_ping) + else: + x_regs_ping = load_x_tile(next_k1) + a_scale_ping, b_scale_ping = prefetch_ab_scale_tile( + _k_base(next_k1), _k_shift_bits(next_k1) + ) b_ping_lo = ( load_b_tile_lo(next_k1_bk) if _b_split_enabled else load_b_tile(next_k1_bk) ) - a_scale_ping, b_scale_ping = prefetch_ab_scale_tile( - _k_base(next_k1), _k_shift_bits(next_k1) - ) acc, _ = compute_tile( acc, b_pong, - lds_x_pong, + lds_base_pong, a_scale_pong, b_scale_pong, a0_prefetch=a0_prefetch_pong, @@ -4127,16 +4483,17 @@ def _make_b_hi_loader(base_k): else None ), ) - hot_loop_scheduler() - rocdl.s_waitcnt(0) + if const_expr(not use_async_copy): + store_x_tile_to_lds(x_regs_ping, lds_base_ping) + # hot_loop_scheduler() gpu.barrier() # Cross-tile prefetch for the ping tile we are about to compute. a0_prefetch_ping = lds_load_packs_k64( - row_a_lds, col_offset_base, lds_x_ping + row_a_lds, col_offset_base, lds_base_ping ) a1_prefetch_ping = ( - lds_load_packs_k64(row_a_lds, _a1_col_base, lds_x_ping) + lds_load_packs_k64(row_a_lds, _a1_col_base, lds_base_ping) if pack_K >= 2 else None ) @@ -4144,21 +4501,24 @@ def _make_b_hi_loader(base_k): next_k2 = k_iv + c2_tile_k next_k2_py = k_iv_py + tile_k * 2 next_k2_bk = next_k2 // 2 - # DMA X(next_k2) -> pong (non-blocking, overlaps with compute) - prefetch_x_to_lds(next_k2, lds_x_pong) + # scales-before-B (see above for rationale). + if const_expr(use_async_copy): + prefetch_x_to_lds(next_k2, lds_base_pong) + else: + x_regs_pong = load_x_tile(next_k2) + a_scale_pong, b_scale_pong = prefetch_ab_scale_tile( + _k_base(next_k2_py), _k_shift_bits(next_k2_py) + ) b_pong = ( load_b_tile_lo(next_k2_bk) if _b_split_enabled else load_b_tile(next_k2_bk) ) - a_scale_pong, b_scale_pong = prefetch_ab_scale_tile( - _k_base(next_k2_py), _k_shift_bits(next_k2_py) - ) acc, _ = compute_tile( acc, b_ping_lo, - lds_x_ping, + lds_base_ping, a_scale_ping, b_scale_ping, a0_prefetch=a0_prefetch_ping, @@ -4170,25 +4530,27 @@ def _make_b_hi_loader(base_k): ), ) k0_pong_bk = next_k2_bk - hot_loop_scheduler() + if const_expr(not use_async_copy): + store_x_tile_to_lds(x_regs_pong, lds_base_pong) + # hot_loop_scheduler() gpu.barrier() # Cross-tile prefetch for the next pong tile. a0_prefetch_pong = lds_load_packs_k64( - row_a_lds, col_offset_base, lds_x_pong + row_a_lds, col_offset_base, lds_base_pong ) a1_prefetch_pong = ( - lds_load_packs_k64(row_a_lds, _a1_col_base, lds_x_pong) + lds_load_packs_k64(row_a_lds, _a1_col_base, lds_base_pong) if pack_K >= 2 else None ) if const_expr(odd_k_tiles): - # Tail: single remaining tile (already in pong buffer). + # Tail: single remaining tile (already in `b_cur` / `lds_base_pong`). acc, epilogue_pf = compute_tile( acc, b_pong, - lds_x_pong, + lds_base_pong, a_scale_pong, b_scale_pong, a0_prefetch=a0_prefetch_pong, @@ -4197,7 +4559,6 @@ def _make_b_hi_loader(base_k): b_hi_loader=( _make_b_hi_loader(k0_pong_bk) if _b_split_enabled else None ), - ku_count=_tail_ku_s2 if _pad_ku_skip_s2 > 0 else k_unroll, ) else: @@ -4207,29 +4568,23 @@ def _make_b_hi_loader(base_k): int(inter_dim) + tile_k - 1 ) // tile_k * tile_k - tile_k k_tail1_bk = k_tail1 // 2 - # DMA tail X -> ping - prefetch_x_to_lds(k_tail1, lds_x_ping) - if const_expr(_pad_ku_skip_s2 > 0): - b_ping_lo = load_b_tile(k_tail1_bk, ku_limit=_tail_ku_s2) - a_scale_ping, b_scale_ping = prefetch_ab_scale_tile( - _k_base(k_tail1_py), - _k_shift_bits(k_tail1_py), - ku_packed_limit=_tail_ku_packed_s2, - ) + if const_expr(use_async_copy): + prefetch_x_to_lds(k_tail1, lds_base_ping) else: - b_ping_lo = ( - load_b_tile_lo(k_tail1_bk) - if _b_split_enabled - else load_b_tile(k_tail1_bk) - ) - a_scale_ping, b_scale_ping = prefetch_ab_scale_tile( - _k_base(k_tail1_py), _k_shift_bits(k_tail1_py) - ) + x_regs_ping = load_x_tile(k_tail1) + b_ping_lo = ( + load_b_tile_lo(k_tail1_bk) + if _b_split_enabled + else load_b_tile(k_tail1_bk) + ) + a_scale_ping, b_scale_ping = prefetch_ab_scale_tile( + _k_base(k_tail1_py), _k_shift_bits(k_tail1_py) + ) acc, _ = compute_tile( acc, b_pong, - lds_x_pong, + lds_base_pong, a_scale_pong, b_scale_pong, a0_prefetch=a0_prefetch_pong, @@ -4239,48 +4594,41 @@ def _make_b_hi_loader(base_k): ), ) + if const_expr(not use_async_copy): + store_x_tile_to_lds(x_regs_ping, lds_base_ping) # hot_loop_scheduler() - rocdl.s_waitcnt(0) gpu.barrier() # Epilogue tile with sw prefetch. a0_prefetch_ping = lds_load_packs_k64( - row_a_lds, col_offset_base, lds_x_ping + row_a_lds, col_offset_base, lds_base_ping ) a1_prefetch_ping = ( - lds_load_packs_k64(row_a_lds, _a1_col_base, lds_x_ping) - if pack_K >= 2 and (_pad_ku_skip_s2 == 0 or _tail_ku_s2 >= 2) + lds_load_packs_k64(row_a_lds, _a1_col_base, lds_base_ping) + if pack_K >= 2 else None ) acc, epilogue_pf = compute_tile( acc, b_ping_lo, - lds_x_ping, + lds_base_ping, a_scale_ping, b_scale_ping, a0_prefetch=a0_prefetch_ping, a1_prefetch=a1_prefetch_ping, prefetch_epilogue=True, b_hi_loader=( - None - if _pad_ku_skip_s2 > 0 - else ( - _make_b_hi_loader(k_tail1_bk) - if _b_split_enabled - else None - ) + _make_b_hi_loader(k_tail1_bk) if _b_split_enabled else None ), - ku_count=_tail_ku_s2 if _pad_ku_skip_s2 > 0 else k_unroll, ) # ---------------- Epilogue: LDS CShuffle + atomic half2 (x2) ---------------- # Reuse the shared helper so GEMM / MoE kernels share the exact same CShuffle skeleton. - sw_pf = None tw_pf = None bias_pf = None if const_expr(epilogue_pf is not None): - sw_pf, tw_pf, bias_pf = epilogue_pf + _, tw_pf, bias_pf = epilogue_pf mask24_i32 = arith.constant(0xFFFFFF) topk_i32_v = topk_i32 @@ -4313,7 +4661,7 @@ def atomic_add_f16x2(val_f16x2, byte_off_i32): _llvm_ptr_ty, arg_out ) out_base_i64 = llvm.ptrtoint(T.i64, out_base_ptr) - out_base_idx = arith.index_cast(ir.IndexType.get(), out_base_i64) + out_base_idx = arith.index_cast(T.index, out_base_i64) def write_row_to_lds( *, @@ -4326,20 +4674,6 @@ def write_row_to_lds( num_acc_n: int, lds_out, ): - # Match origin/dev_a16w4: rely on sentinel padded rows + hardware OOB behavior. - fused2 = buffer_ops.buffer_load( - sorted_rsrc, row, vec_width=1, dtype=T.i32 - ) - t2 = fused2 & mask24_i32 - s2 = fused2 >> 24 - - t_ok = arith.cmpi(CmpIPredicate.ult, t2, tokens_i32) - s_ok = arith.cmpi(CmpIPredicate.ult, s2, topk_i32_v) - ts_ok = arith.andi(t_ok, s_ok) - t2_safe = arith.select(ts_ok, t2, arith.constant(0)) - s2_safe = arith.select(ts_ok, s2, arith.constant(0)) - t2_safe * topk_i32_v + s2_safe - if const_expr(doweight_stage2): tw_idx = (mi * 4) + ii if const_expr(tw_pf is not None): @@ -4370,7 +4704,41 @@ def write_row_to_lds( vector.store(v1, lds_out, [lds_idx], alignment=2) + # Buffer-atomic eligibility for accumulate=True path: + # the byte-offset must fit in i32. The relevant max is + # max_t * model_dim * out_elem_bytes; we conservatively + # cap row_stride <= 16384 bytes (covers model_dim<=8192 bf16). + _row_stride_bytes_py = int(model_dim) * int(out_elem_bytes) + _use_buf_atomic = bool(accumulate) and (_row_stride_bytes_py <= 16384) + _row_stride_bytes_i32 = ( + arith.constant(_row_stride_bytes_py, type=T.i32) + if _use_buf_atomic + else None + ) + _out_elem_bytes_i32 = ( + arith.constant(int(out_elem_bytes), type=T.i32) + if _use_buf_atomic + else None + ) + # 0x7FFF0000 is safely beyond out_rsrc.num_records_bytes + # (~470MB max for prefill DSV3 accumulate=True) and below + # INT32_MAX with room for column offset (≤tile_n*out_elem_bytes). + # HW drops OOB buffer atomics, so routing invalid rows here + # lets us drop the per-mr scf.IfOp in the cshuffle epilogue. + _OOB_safe_i32 = ( + arith.constant(0x7FFF0000, type=T.i32) if _use_buf_atomic else None + ) + def precompute_row(*, row_local, row): + # Fast path: lds_tid was prefilled with the per-row byte + # offset (see _use_buf_atomic_pre at the prologue), so this + # is just one ds_read on the per-row critical path. + if const_expr(_use_buf_atomic): + row_byte_off_i32 = memref.load(lds_tid, [row_local]) + return ( + (None, None, row_byte_off_i32), + None, + ) # Use lds_tid (sorted_idx preloaded to LDS) instead of buffer_load # to avoid extra VMEM round-trips in the epilogue. fused2 = memref.load(lds_tid, [row_local]) @@ -4392,7 +4760,8 @@ def precompute_row(*, row_local, row): row_byte_base = out_base_idx + ts_idx * arith.constant( model_dim * out_elem_bytes, index=True ) - return ((fused2, row_byte_base), row_valid) + row_byte_off_i32 = None + return ((fused2, row_byte_base, row_byte_off_i32), row_valid) def _idx_to_llvm_ptr(idx_val, addr_space=1): """Convert an index-typed byte address to !llvm.ptr.""" @@ -4403,7 +4772,7 @@ def _idx_to_llvm_ptr(idx_val, addr_space=1): return llvm.inttoptr(ptr_ty, i64_raw) def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag): - fused, row_byte_base = row_ctx + fused, row_byte_base, row_byte_off_i32 = row_ctx if const_expr(not bool(accumulate)): # ---- 64-bit global store path (avoids i32 offset overflow) ---- col_idx = col_g0 @@ -4419,8 +4788,23 @@ def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag): alignment=_e_vec * out_elem_bytes, nontemporal=True, ) + elif const_expr(_use_buf_atomic): + # Buffer-atomic path -> BUFFER_ATOMIC_PK_ADD_BF16. + # out_rsrc keeps the 64-bit base in SGPRs so only the + # 32-bit voffset lives in VGPRs, vs llvm.AtomicRMWOp + # which materializes a full 64-bit pointer per op. + col_i32 = arith.index_cast(T.i32, col_g0) + col_byte_off_i32 = col_i32 * _out_elem_bytes_i32 + byte_off_i32 = row_byte_off_i32 + col_byte_off_i32 + rocdl.raw_ptr_buffer_atomic_fadd( + frag, + out_rsrc, + byte_off_i32, + zero_i32, + zero_i32, + ) else: - # ---- accumulate=True: 64-bit global atomic path ---- + # ---- accumulate=True (large-output): 64-bit global atomic path ---- col_idx = col_g0 byte_off_col = col_idx * arith.constant( out_elem_bytes, index=True @@ -4438,6 +4822,7 @@ def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag): ) _e_vec = 2 if accumulate else min(tile_n // 32, 8) + rocdl.s_setprio(3) c_shuffle_epilog( arith=arith, vector=vector, @@ -4463,6 +4848,7 @@ def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag): precompute_row=precompute_row, store_pair=store_pair, ) + rocdl.s_setprio(0) _all_valid = arith.andi(blk_valid, arith.andi(exp_valid, tile_has_tokens)) @@ -4508,6 +4894,8 @@ def store_pair(*, row_local, row, row_ctx, col_pair0, col_g0, frag): persist_m, _sort_block_m, _cu_num if _persistent else 0, + waves_per_eu, + use_async_copy, xcd_swizzle, ) @@ -4530,12 +4918,10 @@ def launch_mixed_moe_gemm2( stream: fx.Stream, ): _ = _cache_tag - allocator_pong.finalized = False - allocator_ping.finalized = False + allocator.finalized = False ctx = CompilationContext.get_current() with ir.InsertionPoint(ctx.gpu_module_body): - allocator_pong.finalize() - allocator_ping.finalize() + allocator.finalize() n_in = arith.index_cast(ir.IndexType.get(), i32_n_in.ir_value()) _tile_n_idx = arith.constant(tile_n, index=True) @@ -4553,7 +4939,7 @@ def launch_mixed_moe_gemm2( - arith.constant(1, index=True) ) / _c_pm_l - moe_gemm2( + _launcher = moe_gemm2( arg_out, arg_x, arg_w, @@ -4568,7 +4954,18 @@ def launch_mixed_moe_gemm2( i32_n_in, i32_k_in, i32_size_expert_ids_in, - ).launch( + ) + # Stamp rocdl.waves_per_eu on the emitted gpu.func so the backend + # respects the occupancy request. Done here (not at kernel-body time) + # because the gpu.func only exists after the kernel has been emitted. + if const_expr(waves_per_eu is not None): + _wpe = int(waves_per_eu) + for _op in ctx.gpu_module_body.operations: + if hasattr(_op, "attributes") and _op.OPERATION_NAME == "gpu.func": + _op.attributes["rocdl.waves_per_eu"] = ir.IntegerAttr.get( + T.i32, _wpe + ) + _launcher.launch( grid=(gx, gy, 1), block=(256, 1, 1), stream=stream, diff --git a/aiter/ops/flydsl/moe_kernels.py b/aiter/ops/flydsl/moe_kernels.py index 357bffdcf0e..6a79ec9b235 100644 --- a/aiter/ops/flydsl/moe_kernels.py +++ b/aiter/ops/flydsl/moe_kernels.py @@ -186,9 +186,55 @@ def get_flydsl_stage2_kernels( **base_params, "persist": True, } + _register_production_variants_stage2(kernels, a_dtype, b_dtype, out_dtype) return kernels +def _register_production_variants_stage2( + kernels: Dict[str, Dict], a_dtype: str, b_dtype: str, out_dtype: str +) -> None: + """Append hand-tuned stage2 variants to ``kernels`` in-place. + + Pulled out of the 6-deep tile/mode/bnt/xcd cartesian product in + ``get_flydsl_stage2_kernels`` so we don't pay a ``base_name == "..."`` + string special-case on every iteration. Each entry pins a specific + shape (tile/mode/dtype) and applies a hand-tuned override dict. + """ + # (a, b, out, tile_m, tile_n, tile_k, mode, suffix, overrides) + PRODUCTION_VARIANTS = ( + # EP4 DeepSeek prefill on MI355X (M=49152, model_dim=7168, inter_dim=2048): + # use_async_copy=True -- async X DMA in prologue overlaps with B/scale VMEM + # cu_num_mul=3 -- persistent grid 3x CU count fills in-flight + # slack from small per-WG M tile counts; + # cu_num_mul=4 regresses ~2.4% on the same shape + # waves_per_eu=4 -- best on EP4 prefill at cu_num_mul=3; + # wpe=5/6 underperform here + ( + "fp4", + "fp4", + "bf16", + 64, + 128, + 256, + "atomic", + "_persist_async_w4_cumul3", + { + "persist": True, + "use_async_copy": True, + "waves_per_eu": 4, + "cu_num_mul": 3, + }, + ), + ) + for pa, pb, pout, ptm, ptn, ptk, pmode, psuffix, povr in PRODUCTION_VARIANTS: + if (pa, pb, pout) != (a_dtype, b_dtype, out_dtype): + continue + _base = flydsl_kernel_name(2, pa, pb, pout, ptm, ptn, ptk, pmode) + if _base not in kernels: + continue + kernels[_base + psuffix] = {**kernels[_base], **povr} + + def get_flydsl_stage1_kernels_int4_bf16(out_dtype: str) -> Dict[str, Dict]: """Return {kernelName: params} for all supported int4_bf16 stage1 configs.""" kernels = {} @@ -378,6 +424,9 @@ def compile_flydsl_moe_stage2( accumulate: bool = True, persist_m: int = 1, sort_block_m: int = 0, + waves_per_eu: Optional[int] = None, + use_async_copy: bool = False, + cu_num_mul: int = 1, b_nt: int = 0, model_dim_pad: int = 0, inter_dim_pad: int = 0, @@ -403,10 +452,18 @@ def compile_flydsl_moe_stage2( accumulate=accumulate, persist_m=persist_m, sort_block_m=sort_block_m, + waves_per_eu=waves_per_eu, + use_async_copy=use_async_copy, + cu_num_mul=cu_num_mul, + # API parity (reviewer #3): forward `b_nt` and `xcd_swizzle` + # from the kernel-name parser. They are accepted as ignored + # kwargs on the fp4xfp4 path so callers parsing the + # `_bnt{N}` / `_xcd{N}` registry suffixes don't need + # per-dtype special cases. b_nt=b_nt, + xcd_swizzle=xcd_swizzle, model_dim_pad=model_dim_pad, inter_dim_pad=inter_dim_pad, - xcd_swizzle=xcd_swizzle, enable_bias=enable_bias, ) elif a_dtype == "bf16" and b_dtype == "int4": @@ -1035,6 +1092,9 @@ def flydsl_moe_stage2( sorted_weights: Optional[torch.Tensor] = None, sort_block_m: int = 0, persist: Optional[bool] = None, + waves_per_eu: Optional[int] = None, + use_async_copy: bool = False, + cu_num_mul: int = 1, b_nt: int = 0, model_dim_pad: int = 0, inter_dim_pad: int = 0, @@ -1070,6 +1130,12 @@ def flydsl_moe_stage2( out = alloc_fn( (token_num, model_dim), dtype=torch_out_dtype, device=inter_states.device ) + # NOTE: when ``accumulate=True`` (atomic mode), the caller is responsible + # for ensuring ``out`` is zero-initialized. In the standard ``fused_moe`` + # dispatch path this is handled by ``moe_sorting_*_fwd`` which already + # zeros ``moe_buf`` via ``moe_buf_set_zero_kernel_2d``, so an extra + # ``out.fill_(0)`` here would be a redundant ~``token_num * model_dim`` + # HBM write (~130us per call at MI355X HBM bw on EP4 prefill shape). dev = inter_states.device flat_a_scale = ( @@ -1164,6 +1230,9 @@ def flydsl_moe_stage2( accumulate=accumulate, persist_m=_persist_m, sort_block_m=sort_block_m, + waves_per_eu=waves_per_eu, + use_async_copy=use_async_copy, + cu_num_mul=cu_num_mul, b_nt=b_nt, model_dim_pad=model_dim_pad, inter_dim_pad=inter_dim_pad,