submission 643447

- # MoE v19 — v17 base + v18's d=2048 fix + isolation test
+ # MoE v30 — Cherry-picked wins from v29 applied to v19 base
  #
- # Key insight from v17/v18 comparison:
- #   v17 registered (64,128) and (64,256) FlyDSL kernels → sparse E=257 improved
- #   v18 removed those registrations → sparse E=257 regressed
- #   BUT v17's (64,128) registration caused d=2048 to use t64x128 instead of t16x128
- #   → d=2048 regressed (235us vs v4's 186us)
+ # v29 tested 4 config changes. Results:
+ #   bs=512 E=257: 4WG_M64 stage1 → 208→180us (13% WIN) ← KEEP
+ #   bs=128 E=33:  block_m=32      → 91.4→89.6us (2% win) ← KEEP
+ #   bs=16  E=33:  ksplit=1        → 60→90us (50% WORSE) ← REVERT
+ #   bs=512 E=33 d=2048: FLY_64x128 → 191→237us (24% WORSE) ← REVERT
  #
- # Strategy for v19:
- #   - Keep v17's full FlyDSL registration set (including 64x128, 64x256)
- #   - Keep v17's E=257 sparse configs (proven: 89.7, 172, 208)
- #   - Keep v4's E=33 configs (proven: 63.6, 91.5, 106, 186)
- #   - For d=2048: DON'T specify a stage2 kernel name — let AITER's default
- #     chooser pick, but with block_m=64. v18 proved that when we explicitly
- #     specify _FLY_16x128, it works (192us). But v17 showed that with the
- #     64x128 registration, AITER overrides our choice. So we try: remove
- #     the explicit kernelName2 and see what AITER picks naturally.
- #   - Add detailed stderr logging for all shapes.
+ # v30 = v19 + these two proven improvements. Nothing else changes.
  #
- # Test:       popcorn submit --gpu MI355X --leaderboard moe-mxfp4 --mode test moe_v19.py
- # Benchmark:  popcorn submit --gpu MI355X --leaderboard moe-mxfp4 --mode benchmark moe_v19.py
- # Leaderboard: popcorn submit --gpu MI355X --leaderboard moe-mxfp4 --mode leaderboard moe_v19.py
+ # Expected geomean improvement:
+ #   v19: (89.7 × 172 × 208 × 60.1 × 91.4 × 108 × 191)^(1/7) ≈ 122us
+ #   v30: (89.7 × 172 × 180 × 60.1 × 89.6 × 108 × 191)^(1/7) ≈ 118us
+ #
+ # Test:       popcorn submit --gpu MI355X --leaderboard moe-mxfp4 --mode test moe_v30.py
+ # Benchmark:  popcorn submit --gpu MI355X --leaderboard moe-mxfp4 --mode benchmark moe_v30.py
+ # Leaderboard: popcorn submit --gpu MI355X --leaderboard moe-mxfp4 --mode leaderboard moe_v30.py
   
  import os
  import sys
  ⋯ 4 unchanged lines
  from aiter.fused_moe import fused_moe
  import aiter.fused_moe as _fused_moe_module
   
- # ── Register FlyDSL tile_k=128 kernels ──
- # CRITICAL: the (64,128) and (64,256) registrations influence AITER's
- # internal kernel selection for the sparse E=257 shapes, even though
- # we don't explicitly assign them. Removing them causes regression.
+ # ── FlyDSL registration (identical to v19) ──
  try:
      import aiter.ops.flydsl.moe_kernels as _flydsl_moe_kernels
      for tile_m, tile_n in [(32, 128), (32, 256), (16, 256), (16, 128), (64, 128), (64, 256)]:
  ⋯ 6 unchanged lines
  except ImportError:
      pass
   
- # ── Config key helper ──
  def _key(token, inter_dim, expert):
      return (
          256, token, 7168, inter_dim, expert, 9,
  ⋯ 2 unchanged lines
          "QuantType.per_1x32", True, False,
      )
   
- # ── Kernel name constants ──
+ # ── Kernel names ──
  _4WG_M128 = "moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
+ _4WG_M64  = "moe_ck2stages_gemm1_256x64x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
  _4WG_M32  = "moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
  _FLY_16x128 = "flydsl_moe2_afp4_wfp4_bf16_t16x128x128_atomic"
   
  _CUSTOM_CONFIGS = {}
   
  # ═══════════════════════════════════════════════════════════════
- # E=33 shapes — EXACT v4 configs (proven best)
+ # E=33 shapes
  # ═══════════════════════════════════════════════════════════════
   
+ # bs=16: EXACT v19 (ksplit=2 proven, ksplit=1 was 50% worse in v29)
  _CUSTOM_CONFIGS[_key(16, 512, 33)] = {
      "block_m": 32, "ksplit": 2,
      "kernelName1": "", "kernelName2": "",
      "run_1stage": False,
  }
+  
+ # bs=128: CHANGED from v19 — block_m=32 instead of 64 (v29: 91.4→89.6, 2% win)
  _CUSTOM_CONFIGS[_key(128, 512, 33)] = {
-     "block_m": 64, "ksplit": 0,
+     "block_m": 32, "ksplit": 0,  # CHANGED: block_m 64→32
      "kernelName1": _4WG_M128, "kernelName2": _FLY_16x128,
      "run_1stage": False,
  }
+  
+ # bs=512 d=512: EXACT v19
  _CUSTOM_CONFIGS[_key(512, 512, 33)] = {
      "block_m": 64, "ksplit": 0,
      "kernelName1": _4WG_M128, "kernelName2": _FLY_16x128,
      "run_1stage": False,
  }
- # d=2048: v17 registered (64,128) which caused AITER to pick t64x128
- # instead of our explicit _FLY_16x128. v18 fixed this by removing the
- # (64,128) registration but that hurt E=257 shapes.
- # Strategy: keep (64,128) registered but use empty kernelName2 to let
- # AITER's natural dispatch handle it. If AITER picks t64x128, that's
- # what we observed in v17 at 235us. If it picks something else, we learn.
- # Alternative: try block_m=128 which might pair better with t64x128.
+  
+ # bs=512 d=2048: EXACT v19 (FLY_64x128 was 24% worse in v29)
  _CUSTOM_CONFIGS[_key(512, 2048, 33)] = {
      "block_m": 64, "ksplit": 0,
      "kernelName1": _4WG_M128, "kernelName2": _FLY_16x128,
  ⋯ 1 unchanged lines
  }
   
  # ═══════════════════════════════════════════════════════════════
- # E=257 shapes — v17 configs (proven best: 89.7, 172, 208)
- # The improvement comes from the (64,128)/(64,256) FlyDSL
- # registrations influencing AITER's internal chooser, combined
- # with NT loads for sparse dispatch.
+ # E=257 shapes
  # ═══════════════════════════════════════════════════════════════
   
+ # bs=16: EXACT v19
  _CUSTOM_CONFIGS[_key(16, 256, 257)] = {
      "block_m": 16, "ksplit": 2,
      "kernelName1": "", "kernelName2": "",
      "run_1stage": False,
      "use_non_temporal_load": True,
  }
+  
+ # bs=128: EXACT v19
  _CUSTOM_CONFIGS[_key(128, 256, 257)] = {
      "block_m": 16, "ksplit": 2,
      "kernelName1": "", "kernelName2": "",
      "run_1stage": False,
      "use_non_temporal_load": True,
  }
+  
+ # bs=512: CHANGED from v19 — 4WG_M64 instead of 4WG_M32 (v29: 208→180, 13% win!)
  _CUSTOM_CONFIGS[_key(512, 256, 257)] = {
      "block_m": 32, "ksplit": 0,
-     "kernelName1": _4WG_M32, "kernelName2": _FLY_16x128,
+     "kernelName1": _4WG_M64,  # CHANGED: M32→M64
+     "kernelName2": _FLY_16x128,
      "run_1stage": False,
      "use_non_temporal_load": True,
  }
   
- # ── Config injection ──
+ # ── Config injection (identical to v19) ──
  _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",
-             ]
+             _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)
   
-     _original_get_2stage_cfgs = _fused_moe_module.get_2stage_cfgs
-  
+     _original = _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,
-     ):
-         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,
-         )
-  
+     def _patched(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):
+         metadata = _original(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)
          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,
-         )
+         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)
          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"]
  ⋯ 4 unchanged lines
                  new_kw['non_temporal_load'] = nt
                  metadata = _fused_moe_module.MOEMetadata(
                      functools.partial(old_s1.func, **new_kw),
-                     metadata.stage2,
-                     metadata.block_m,
-                     metadata.ksplit,
-                     metadata.run_1stage,
-                     metadata.has_bias,
-                     nt,
-                 )
-  
-         # Log actual dispatch for diagnosis
-         s1_name = getattr(getattr(metadata.stage1, 'keywords', {}), 'get', lambda k, d=None: d)('kernelName', '?')
-         print(f"[v19] token={token} inter={inter_dim} expert={expert} "
-               f"block_m={metadata.block_m} ksplit={metadata.ksplit} "
-               f"nt={getattr(metadata, 'use_non_temporal_load', 'N/A')} "
-               f"stage2={metadata.stage2}",
-               file=sys.stderr)
-  
+                     metadata.stage2, metadata.block_m, metadata.ksplit,
+                     metadata.run_1stage, metadata.has_bias, nt)
          return metadata
+     _fused_moe_module.get_2stage_cfgs = _patched
   
-     _fused_moe_module.get_2stage_cfgs = _patched_get_2stage_cfgs
   
-  
  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
+     (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
   
      _inject_configs()