submission 642119

- # MoE v4 — Config injection for shape-specific kernel tuning
+ # MoE v19 — v17 base + v18's d=2048 fix + isolation test
  #
- # Strategy: inject custom tile configs into AITER's fused_moe config registry.
- # The actual kernel call is identical to the reference — all optimization is
- # in selecting better (block_m, ksplit, kernel) combos per shape.
+ # 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)
  #
- # Based on analysis of top competitors + AITER's CK/FlyDSL kernel system.
- # Each config was chosen based on tokens-per-expert analysis:
- #   - ksplit=2 helps when tokens/expert is very low (sparse dispatch)
- #   - 4-WG stage1 kernels help for medium/large batches (better CU utilization)
- #   - FlyDSL stage2 with tile_k=128 for better K-dimension throughput
+ # 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.
  #
- # No graph capture. No cross-call state. No banned words. No expert masking.
- # Fresh computation every call — would pass "could this run in vLLM" test.
- #
- # Test:       popcorn submit --gpu MI355X --leaderboard moe-mxfp4 --mode test moe_v4.py
- # Benchmark:  popcorn submit --gpu MI355X --leaderboard moe-mxfp4 --mode benchmark moe_v4.py
- # Leaderboard: popcorn submit --gpu MI355X --leaderboard moe-mxfp4 --mode leaderboard moe_v4.py
+ # 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
   
  import os
+ import sys
  import functools
  import torch
  from task import input_t, output_t
  ⋯ 1 unchanged lines
  from aiter.fused_moe import fused_moe
  import aiter.fused_moe as _fused_moe_module
   
- # ── Register FlyDSL tile_k=128 kernels not in server defaults ──
+ # ── 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.
  try:
      import aiter.ops.flydsl.moe_kernels as _flydsl_moe_kernels
-     # These kernel variants exist in AITER but may not be registered by default
-     for tile_m, tile_n in [(32, 128), (32, 256), (16, 256), (16, 128)]:
+     for tile_m, tile_n in [(32, 128), (32, 256), (16, 256), (16, 128), (64, 128), (64, 256)]:
          name = f"flydsl_moe2_afp4_wfp4_bf16_t{tile_m}x{tile_n}x128_atomic"
          _flydsl_moe_kernels._KERNEL_PARAMS[name] = {
              "stage": 2, "a_dtype": "fp4", "b_dtype": "fp4", "out_dtype": "bf16",
  ⋯ 1 unchanged lines
              "mode": "atomic", "MPerBlock": tile_m,
          }
  except ImportError:
-     pass  # FlyDSL not available — fall back to defaults
+     pass
   
- # ── Shape-specific configs ──
- # Key format: (cu_num, token, model_dim, inter_dim, expert, topk,
- #              act_type, dtype, q_dtype_a, q_dtype_w, q_type, use_g1u1, doweight_stage1)
- _CUSTOM_CONFIGS = {}
-  
+ # ── 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,
      )
   
- # 4-WG CK stage1 kernel names (multi-workgroup for better CU utilization)
+ # ── Kernel name constants ──
  _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"
-  
- # FlyDSL stage2 kernel names
+ _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"
   
- # ── E=33 shapes (TP=4, fewer larger experts) ──
- # bs=16: very sparse (< 2 tokens/expert avg). ksplit=2 helps.
+ _CUSTOM_CONFIGS = {}
+  
+ # ═══════════════════════════════════════════════════════════════
+ # E=33 shapes — EXACT v4 configs (proven best)
+ # ═══════════════════════════════════════════════════════════════
+  
  _CUSTOM_CONFIGS[_key(16, 512, 33)] = {
      "block_m": 32, "ksplit": 2,
      "kernelName1": "", "kernelName2": "",
      "run_1stage": False,
  }
- # bs=128: ~4 tokens/expert. 4-WG stage1 + FlyDSL stage2.
  _CUSTOM_CONFIGS[_key(128, 512, 33)] = {
      "block_m": 64, "ksplit": 0,
      "kernelName1": _4WG_M128, "kernelName2": _FLY_16x128,
      "run_1stage": False,
  }
- # bs=512/d=512: ~16 tokens/expert. 4-WG stage1 + FlyDSL, block_m=64.
  _CUSTOM_CONFIGS[_key(512, 512, 33)] = {
      "block_m": 64, "ksplit": 0,
      "kernelName1": _4WG_M128, "kernelName2": _FLY_16x128,
      "run_1stage": False,
  }
- # bs=512/d=2048: largest shape. 4-WG + FlyDSL, block_m=64.
+ # 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.
  _CUSTOM_CONFIGS[_key(512, 2048, 33)] = {
      "block_m": 64, "ksplit": 0,
      "kernelName1": _4WG_M128, "kernelName2": _FLY_16x128,
      "run_1stage": False,
  }
   
- # ── E=257 shapes (EP-off, many small experts) ──
- # bs=16: extremely sparse (~0.6 tokens/expert). ksplit=2.
+ # ═══════════════════════════════════════════════════════════════
+ # 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.
+ # ═══════════════════════════════════════════════════════════════
+  
  _CUSTOM_CONFIGS[_key(16, 256, 257)] = {
      "block_m": 16, "ksplit": 2,
      "kernelName1": "", "kernelName2": "",
      "run_1stage": False,
+     "use_non_temporal_load": True,
  }
- # bs=128: ~4.5 tokens/expert. ksplit=2.
  _CUSTOM_CONFIGS[_key(128, 256, 257)] = {
      "block_m": 16, "ksplit": 2,
      "kernelName1": "", "kernelName2": "",
      "run_1stage": False,
+     "use_non_temporal_load": True,
  }
- # bs=512: ~18 tokens/expert. 4-WG + FlyDSL + non-temporal loads.
  _CUSTOM_CONFIGS[_key(512, 256, 257)] = {
      "block_m": 32, "ksplit": 0,
      "kernelName1": _4WG_M32, "kernelName2": _FLY_16x128,
  ⋯ 1 unchanged lines
      "use_non_temporal_load": True,
  }
   
- # ── Config injection (runs once) ──
+ # ── Config injection ──
  _injected = False
   
  def _inject_configs():
  ⋯ 2 unchanged lines
          return
      _injected = True
   
-     # Load the tuning CSV if not already loaded
      if _fused_moe_module.cfg_2stages is None:
          import pandas as pd
          from aiter.jit.core import AITER_CONFIGS
  ⋯ 11 unchanged lines
          else:
              _fused_moe_module.cfg_2stages = {}
   
-     # Merge our custom configs (overrides CSV defaults)
      _fused_moe_module.cfg_2stages.update(_CUSTOM_CONFIGS)
   
-     # Monkeypatch get_2stage_cfgs to support use_non_temporal_load
      _original_get_2stage_cfgs = _fused_moe_module.get_2stage_cfgs
   
      @functools.lru_cache(maxsize=2048)
  ⋯ 19 unchanged lines
          if cfg and cfg.get("use_non_temporal_load") is not None:
              nt = cfg["use_non_temporal_load"]
              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,
-                     )
+             kw = getattr(old_s1, 'keywords', None) or {}
+             if hasattr(old_s1, 'func') and 'non_temporal_load' in kw:
+                 new_kw = dict(kw)
+                 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)
+  
          return metadata
   
      _fused_moe_module.get_2stage_cfgs = _patched_get_2stage_cfgs