submission 567197

+ #!POPCORN leaderboard amd-moe-mxfp4
+ #!POPCORN gpu MI355X
+  
  """
- V188: Fix force-NT arg positions + always force NT for all CK Codegen (ksplit=0) shapes.
- V176 had wrong arg mapping (args[3]=inter_dim not expert, args[4]=expert not topk).
- This caused E=33 M=512 d=2048 (slowest shape) to NOT get force-NT.
+ v103: Enable opus sorting. Opus sorting has explicit fused buffer zeroing
+ (moe_buf_set_zero_kernel_2d) which may interact better with FlyDSL atomic stage2.
+ Monkeypatch _moe_sorting_impl to use opus sorting via use_opus=True.
  """
  import os
- import gc
  import functools
- import sys
-  
- os.environ["AITER_USE_OPUS_MOE_SORTING"] = "1"
-  
  import torch
- import aiter
- import aiter.fused_moe as _fmoe
- from aiter import ActivationType, QuantType
- from aiter.fused_moe import fused_moe, get_2stage_cfgs
+ from typing import Dict
  from task import input_t, output_t
   
- gc.disable()
+ from aiter import ActivationType, QuantType
+ from aiter.fused_moe import fused_moe
+ import aiter.fused_moe as _fused_moe_module
+ import aiter
   
- # Patch get_2stage_cfgs to force non_temporal_load=True for large-M CK Codegen shapes
- _orig_get_2stage_cfgs = get_2stage_cfgs.__wrapped__
- @functools.lru_cache(maxsize=2048)
- def _patched_get_2stage_cfgs(*args, **kwargs):
-     meta = _orig_get_2stage_cfgs(*args, **kwargs)
-     if meta.ksplit == 0:
-         # Force non_temporal_load=True for ALL CK Codegen shapes
-         # Default heuristic only enables NT for small-M; benchmarks show NT helps large-M too
-         if hasattr(meta.stage1, 'keywords') and 'use_non_temporal_load' in meta.stage1.keywords:
-             meta.stage1.keywords['use_non_temporal_load'] = True
-         if hasattr(meta.stage2, 'keywords') and 'use_non_temporal_load' in meta.stage2.keywords:
-             meta.stage2.keywords['use_non_temporal_load'] = True
-     return meta
- sys.modules['aiter.fused_moe'].get_2stage_cfgs = _patched_get_2stage_cfgs
+ # Inject ksplit=2 configs for shapes that benefit from cktile_moe path
+ _CUSTOM_CONFIGS = {}
   
- _sorting_cache = {}
- def _cached_sorting_impl(topk_ids, topk_weights, num_experts, model_dim, moebuf_dtype,
-     block_size, expert_mask, num_local_tokens, dispatch_policy, use_opus):
-     M, topk = topk_ids.shape
-     max_num_tokens_padded = int(topk_ids.numel() + num_experts * block_size - topk)
-     max_num_m_blocks = int((max_num_tokens_padded + block_size - 1) // block_size)
-     device = topk_ids.device
-     key = (M, model_dim, max_num_tokens_padded, max_num_m_blocks)
-     if key not in _sorting_cache:
-         _sorting_cache[key] = (
-             torch.empty(max_num_tokens_padded, dtype=torch.int32, device=device),
-             torch.empty(max_num_tokens_padded, dtype=torch.float32, device=device),
-             torch.empty(max_num_m_blocks, dtype=torch.int32, device=device),
-             torch.empty(2, dtype=torch.int32, device=device),
-             torch.empty((M, model_dim), dtype=moebuf_dtype, device=device),
+ def _make_key(token, inter_dim, expert):
+     return (
+         256, token, 7168, inter_dim, expert, 9,
+         "ActivationType.Silu", "torch.bfloat16",
+         "torch.float4_e2m1fn_x2", "torch.float4_e2m1fn_x2",
+         "QuantType.per_1x32", True, False,
+     )
+  
+ # === E=33 shapes (from v018, proven) ===
+ # bs=16/E=33/d=512: cktile_moe gives 59.6us vs 88.7us baseline (-32.8%)
+ _CUSTOM_CONFIGS[_make_key(16, 512, 33)] = {
+     "block_m": 32,
+     "ksplit": 2,
+     "kernelName1": "",
+     "kernelName2": "",
+     "run_1stage": False,
+ }
+  
+ # bs=128/E=33/d=512: cktile_moe gives 108us vs 124us baseline (-12.9%)
+ _CUSTOM_CONFIGS[_make_key(128, 512, 33)] = {
+     "block_m": 64,
+     "ksplit": 2,
+     "kernelName1": "",
+     "kernelName2": "",
+     "run_1stage": False,
+ }
+  
+ # === E=257 shapes (NEW in v020) ===
+ # bs=16/E=257/d=256: try cktile_moe ksplit=2 (overrides tuned CSV config)
+ # With 144 token-expert pairs across 257 experts, most experts get 0-1 tokens.
+ # Skipping activation quantization + using split-K may help.
+ _CUSTOM_CONFIGS[_make_key(16, 256, 257)] = {
+     "block_m": 16,
+     "ksplit": 2,
+     "kernelName1": "",
+     "kernelName2": "",
+     "run_1stage": False,
+ }
+  
+ # bs=128/E=257/d=256: try cktile_moe ksplit=2 (overrides tuned CSV config)
+ # bs=128 has ~4.5 tokens/expert avg, similar to E=33 where ksplit=2 helped (-12.9%)
+ _CUSTOM_CONFIGS[_make_key(128, 256, 257)] = {
+     "block_m": 16,
+     "ksplit": 2,
+     "kernelName1": "",
+     "kernelName2": "",
+     "run_1stage": False,
+ }
+  
+ # === bs=512/E=33 shapes: inject 4-WG stage1 kernel ===
+ # The 256x64x128x128_1x4 kernel uses 4 workgroups per CU for better utilization.
+ # v037 showed d=2048: -3.2% (349->338µs). Now also try d=512 with same 4-WG kernel.
+ _4WG_STAGE1 = "moe_ck2stages_gemm1_256x64x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
+  
+ _FLYDSL_STAGE2 = "flydsl_moe2_afp4_wfp4_bf16_t32x128x256_atomic"
+  
+ _CUSTOM_CONFIGS[_make_key(512, 2048, 33)] = {
+     "block_m": 64,
+     "ksplit": 0,
+     "kernelName1": _4WG_STAGE1,
+     "kernelName2": _FLYDSL_STAGE2,
+     "run_1stage": False,
+ }
+  
+ # === bs=512/E=33/d=512: inject FlyDSL stage2 ===
+ # v050 tried t32x256x256_reduce -> correctness failure. Try t32x128x256_atomic.
+ # Down-GEMM: K=512, tile_k=256 -> 2 K-iterations. N=7168, tile_n=128 -> 56 N-blocks.
+ _CUSTOM_CONFIGS[_make_key(512, 512, 33)] = {
+     "block_m": 64,
+     "ksplit": 0,
+     "kernelName1": "",
+     "kernelName2": _FLYDSL_STAGE2,
+     "run_1stage": False,
+ }
+  
+ # === bs=512/E=257: inject tuned CSV kernels + NT=True ===
+ # Heuristic says NT=True for 17 tokens/expert, but tuned CSV path always sets NT=False.
+ # Inject the same kernel names as tuned CSV but add use_non_temporal_load flag.
+ _CUSTOM_CONFIGS[_make_key(512, 256, 257)] = {
+     "block_m": 32,
+     "ksplit": 0,
+     "kernelName1": "moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16",
+     "kernelName2": "moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16",
+     "run_1stage": False,
+     "use_non_temporal_load": True,  # custom flag, read by monkeypatch
+ }
+  
+ _injected = False
+  
+ def _inject_configs():
+     global _injected
+     if _injected:
+         return
+     _injected = True
+  
+     if _fused_moe_module.cfg_2stages is None:
+         import pandas as pd
+         from aiter.jit.core import AITER_CONFIGS
+         tune_file = AITER_CONFIGS.AITER_CONFIG_FMOE_FILE
+         if os.path.exists(tune_file):
+             _INDEX_COLS = [
+                 "cu_num", "token", "model_dim", "inter_dim", "expert", "topk",
+                 "act_type", "dtype", "q_dtype_a", "q_dtype_w", "q_type",
+                 "use_g1u1", "doweight_stage1",
+             ]
+             df = pd.read_csv(tune_file)
+             if "_tag" in df.columns:
+                 df = df[df["_tag"].fillna("") == ""]
+             _fused_moe_module.cfg_2stages = df.set_index(_INDEX_COLS).to_dict("index")
+         else:
+             _fused_moe_module.cfg_2stages = {}
+  
+     _fused_moe_module.cfg_2stages.update(_CUSTOM_CONFIGS)
+  
+     # Enable opus sorting for better buffer zeroing with FlyDSL atomic
+     _fused_moe_module._USE_OPUS_MOE_SORTING = True
+  
+     # Monkeypatch get_2stage_cfgs to support use_non_temporal_load from config
+     _original_get_2stage_cfgs = _fused_moe_module.get_2stage_cfgs
+  
+     @functools.lru_cache(maxsize=2048)
+     def _patched_get_2stage_cfgs(
+         token, model_dim, inter_dim, expert, topk,
+         dtype, q_dtype_a, q_dtype_w, q_type, use_g1u1,
+         activation, doweight_stage1, hidden_pad, intermediate_pad, is_shuffled=True,
+     ):
+         # Get the original metadata
+         metadata = _original_get_2stage_cfgs(
+             token, model_dim, inter_dim, expert, topk,
+             dtype, q_dtype_a, q_dtype_w, q_type, use_g1u1,
+             activation, doweight_stage1, hidden_pad, intermediate_pad, is_shuffled,
          )
-     sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, moe_buf = _sorting_cache[key]
-     fwd_fn = aiter.moe_sorting_opus_fwd if use_opus else aiter.moe_sorting_fwd
-     fwd_fn(topk_ids, topk_weights, sorted_ids, sorted_weights, sorted_expert_ids,
-            num_valid_ids, moe_buf, num_experts, int(block_size),
-            expert_mask, num_local_tokens, dispatch_policy)
-     return sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, moe_buf
- _fmoe._moe_sorting_impl = _cached_sorting_impl
   
- _stage1_cache = {}
- def _cached_cktile_stage1(hidden_states, w1, w2, sorted_token_ids, sorted_expert_ids,
-     num_valid_ids, out, topk, block_m, a1_scale, w1_scale,
-     sorted_weights=None, n_pad_zeros=0, k_pad_zeros=0,
-     bias1=None, activation=ActivationType.Silu, split_k=1, dtype=torch.bfloat16):
-     token_num = hidden_states.shape[0]
-     _, n1, k1 = w1.shape
-     _, k2, n2 = w2.shape
-     D = n2 if k2 == k1 else n2 * 2
-     if w1.dtype is torch.uint32: D = D * 8
-     key = (token_num, topk, D, n1, split_k)
-     if key not in _stage1_cache:
-         device = hidden_states.device
-         _stage1_cache[key] = (
-             torch.empty((token_num, topk, D), dtype=dtype, device=device),
-             torch.empty((token_num, topk, n1), dtype=hidden_states.dtype, device=device) if split_k > 1 else None,
+         # Check if this shape has a custom NT setting
+         from aiter.jit.utils.chip_info import get_cu_num
+         cu_num = get_cu_num()
+         keys = (
+             cu_num, token, model_dim, inter_dim, expert, topk,
+             str(activation), str(dtype), str(q_dtype_a), str(q_dtype_w),
+             str(q_type), use_g1u1, doweight_stage1,
          )
-     out_buf, tmp_buf = _stage1_cache[key]
-     if split_k > 1:
-         tmp_buf.zero_(); tmp_out = tmp_buf
-     else: tmp_out = out_buf
-     aiter.moe_cktile2stages_gemm1(hidden_states, w1, tmp_out,
-         sorted_token_ids, sorted_expert_ids, num_valid_ids,
-         topk, n_pad_zeros, k_pad_zeros, sorted_weights, a1_scale, w1_scale,
-         bias1, activation, block_m, split_k)
-     if split_k > 1:
-         if activation == ActivationType.Silu: aiter.silu_and_mul(out_buf, tmp_out)
-         else: aiter.gelu_and_mul(out_buf, tmp_out)
-     return out_buf
- _fmoe.cktile_moe_stage1 = _cached_cktile_stage1
- _current_mode = None
- def custom_kernel(data):
-     global _current_mode
-     (hidden_states, guw, dw, gus, ds, guw_s, dw_s, gus_s, ds_s, tw, ti, config) = data
-     M = hidden_states.shape[0]; E = guw_s.shape[0]
-     hp = config["d_hidden_pad"] - config["d_hidden"]
-     ip = config["d_expert_pad"] - config["d_expert"]
-     de = config["d_expert"]
+         cfg = _fused_moe_module.cfg_2stages.get(keys)
+         if cfg and cfg.get("use_non_temporal_load") is not None:
+             nt = cfg["use_non_temporal_load"]
+             # Rebuild stage1 partial with NT override
+             old_s1 = metadata.stage1
+             if hasattr(old_s1, 'func') and old_s1.func is not None:
+                 if 'use_non_temporal_load' in (old_s1.keywords or {}):
+                     new_kw = dict(old_s1.keywords)
+                     new_kw['use_non_temporal_load'] = nt
+                     metadata = _fused_moe_module.MOEMetadata(
+                         functools.partial(old_s1.func, **{k: v for k, v in new_kw.items()}),
+                         metadata.stage2,
+                         metadata.block_m,
+                         metadata.ksplit,
+                         metadata.run_1stage,
+                         metadata.has_bias,
+                         nt,
+                     )
+                     # Also patch stage2 if it's a CK kernel (not FlyDSL)
+                     old_s2 = metadata.stage2
+                     if old_s2 and hasattr(old_s2, 'keywords') and 'use_non_temporal_load' in (old_s2.keywords or {}):
+                         new_kw2 = dict(old_s2.keywords)
+                         new_kw2['use_non_temporal_load'] = nt
+                         metadata = _fused_moe_module.MOEMetadata(
+                             metadata.stage1,
+                             functools.partial(old_s2.func, **{k: v for k, v in new_kw2.items()}),
+                             metadata.block_m,
+                             metadata.ksplit,
+                             metadata.run_1stage,
+                             metadata.has_bias,
+                             nt,
+                         )
+         return metadata
   
-     if E > 64 and M <= 128: mode = "cktile_e257_k7"
-     elif E > 64: mode = "e257_csv"
-     elif E <= 64 and M <= 16: mode = "cktile_k7"
-     elif E <= 64 and M <= 128: mode = "cktile_k2"
-     else: mode = "default"
+     _fused_moe_module.get_2stage_cfgs = _patched_get_2stage_cfgs
   
-     if mode != _current_mode:
-         if mode == "cktile_e257_k7":
-             os.environ["AITER_KSPLIT"] = "7"
-             os.environ["AITER_BYPASS_TUNE_CONFIG"] = "1"
-         elif mode == "e257_csv":
-             os.environ.pop("AITER_KSPLIT", None)
-             os.environ.pop("AITER_BYPASS_TUNE_CONFIG", None)
-         elif mode == "cktile_k7":
-             os.environ["AITER_KSPLIT"] = "7"
-             os.environ.pop("AITER_BYPASS_TUNE_CONFIG", None)
-         elif mode == "cktile_k2":
-             os.environ["AITER_KSPLIT"] = "2"
-             os.environ.pop("AITER_BYPASS_TUNE_CONFIG", None)
-         else:
-             os.environ.pop("AITER_KSPLIT", None)
-             os.environ.pop("AITER_BYPASS_TUNE_CONFIG", None)
-         _patched_get_2stage_cfgs.cache_clear(); _current_mode = mode
   
-     if mode == "default" and E <= 64: bsm = 64
-     elif mode == "cktile_k2": bsm = 32
-     else: bsm = None
+ def custom_kernel(data: input_t) -> output_t:
+     (
+         hidden_states, gate_up_weight, down_weight,
+         gate_up_weight_scale, down_weight_scale,
+         gate_up_weight_shuffled, down_weight_shuffled,
+         gate_up_weight_scale_shuffled, down_weight_scale_shuffled,
+         topk_weights, topk_ids, config,
+     ) = data
   
-     act = ActivationType.Silu
-     if mode == "default" and de <= 512: act = ActivationType.Swiglu
+     _inject_configs()
   
-     return fused_moe(hidden_states, guw_s, dw_s, tw, ti, expert_mask=None,
-         activation=act, quant_type=QuantType.per_1x32, doweight_stage1=False,
-         w1_scale=gus_s, w2_scale=ds_s, a1_scale=None, a2_scale=None,
-         block_size_M=bsm, hidden_pad=hp, intermediate_pad=ip)
+     hidden_pad = config["d_hidden_pad"] - config["d_hidden"]
+     intermediate_pad = config["d_expert_pad"] - config["d_expert"]
+  
+     output = fused_moe(
+         hidden_states, gate_up_weight_shuffled, down_weight_shuffled,
+         topk_weights, topk_ids,
+         expert_mask=None, activation=ActivationType.Silu,
+         quant_type=QuantType.per_1x32, doweight_stage1=False,
+         w1_scale=gate_up_weight_scale_shuffled,
+         w2_scale=down_weight_scale_shuffled,
+         a1_scale=None, a2_scale=None,
+         hidden_pad=hidden_pad, intermediate_pad=intermediate_pad,
+     )
+  
+     return output