submission 685904

  ⋯ 8 unchanged lines
  from aiter.fused_moe import fused_moe
  import aiter.fused_moe as _fm
   
- # ── EXP-110: True best per shape after re-analysis ──
- # Shape 1,2: CK 256x32 ksplit=4 (proven)
- # Shape 3: FlyDSL t32 reduce (252 < 256 default < 272 ksplit4)
- # Shape 4,5: CK 256x32 ksplit=4 (60.6/105 < 64.0/108 ksplit2)
- # Shape 6: CK 256x64 + FlyDSL GEMM2 t64 reduce (proven)
- # Shape 7: CK 256x64 + FlyDSL GEMM2 t32 atomic (proven -6%)
+ # ── EXP-121 v4: Hybrid — 1-stage for Shape 3 + best 2-stage for rest ──
+ #
+ # Findings from v3 (all 1-stage):
+ #   S1: 146µs (-4% vs ref) but worse than EXP-112's 2-stage ksplit=4
+ #   S2: 221µs (-7.5% vs ref) same as EXP-112
+ #   S3: 275µs (-18.3% vs ref) ← HUGE WIN, was weakest shape at ~308µs
+ #   S4: 114µs (+7.3% vs ref) WORSE
+ #   S5: 139µs (-1.5% vs ref) close
+ #   S6: 214µs (-4.9% vs ref) but worse than EXP-112's FlyDSL (~178µs)
+ #   S7: 507µs (+33.2% vs ref) TERRIBLE
+ #
+ # Strategy: Use 1-stage only for Shape 3 where it clearly wins.
+ # Keep proven EXP-112 2-stage configs for shapes 1,2,4,5,6,7.
+ # Use custom dict to make test shapes use 1-stage (avoids CKTile build).
+ #
+ # JIT budget: sorting(25)+quant(25)+asm(31)+ck2stage(108) = 189s ← fits 540s
   
   
+ class _FallbackDict(dict):
+     """Dict that returns 1-stage config for any unknown key.
+     This ensures random test shapes use 1-stage (no CKTile build),
+     while benchmark shapes use explicitly injected configs."""
+  
+     _DEFAULT_1STAGE = {
+         'block_m': 32, 'ksplit': 0,
+         'kernelName1': '', 'kernelName2': '',
+         'run_1stage': 1,
+         'us': 0.0, 'us1': 0.0, 'us2': 0.0,
+         'err1': '0', 'err2': '0',
+         'tflops': 0.0, 'bw': 0.0, '_tag': float('nan'),
+     }
+  
+     def get(self, key, default=None):
+         result = super().get(key, None)
+         if result is not None:
+             return result
+         return dict(self._DEFAULT_1STAGE)
+  
+  
  def _preload_and_inject():
      try:
          cu = _fm.get_cu_count() if hasattr(_fm, 'get_cu_count') else 256
-         act = ActivationType.Silu
-         dtype_t = torch.bfloat16
-         q_a_t = dtypes.fp4x2
-         q_w_t = dtypes.fp4x2
-         q_type_t = QuantType.per_1x32
+         act_s = str(ActivationType.Silu)
+         dtype_s = str(torch.bfloat16)
+         q_a = str(dtypes.fp4x2)
+         q_w = str(dtypes.fp4x2)
+         q_type = str(QuantType.per_1x32)
   
-         _fm.get_2stage_cfgs(
-             16, 7168, 256, 257, 9,
-             dtype_t, q_a_t, q_w_t, q_type_t,
-             True, act, False,
-             0, 0, True
-         )
+         # Replace cfg_2stages with custom dict (fallback=1-stage for test shapes)
+         new_cfg = _FallbackDict()
+         new_cfg[('__skip_csv__',)] = {'_dummy': True}  # len>0 => skip CSV read
   
-         cfg = _fm.cfg_2stages
-         if cfg is None:
-             return
-  
-         act_s = str(act)
-         dtype_s = str(dtype_t)
-         q_a = str(q_a_t)
-         q_w = str(q_w_t)
-         q_type = str(q_type_t)
-  
+         # ── Kernel names ──
          kn1_32 = 'moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16'
          kn2_32 = 'moe_ck2stages_gemm2_256x32x128x128_1x4_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16'
          kn1_64 = 'moe_ck2stages_gemm1_256x64x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16'
-         kn2_64 = 'moe_ck2stages_gemm2_256x64x128x128_1x4_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16'
-         # FlyDSL gemm2 variants
          kn2_fly64_reduce = 'flydsl_moe2_afp4_wfp4_bf16_t64x256x256_reduce'
          kn2_fly64_atomic = 'flydsl_moe2_afp4_wfp4_bf16_t64x256x256_atomic'
-         kn2_fly32_reduce = 'flydsl_moe2_afp4_wfp4_bf16_t32x256x256_reduce'
-         kn2_fly32_atomic = 'flydsl_moe2_afp4_wfp4_bf16_t32x256x256_atomic'
   
-         def make_entry(block_m, ksplit, kn1, kn2):
+         def make_2stage(block_m, ksplit, kn1, kn2):
              return {
                  'block_m': block_m, 'ksplit': ksplit,
                  'kernelName1': kn1, 'kernelName2': kn2,
  ⋯ 3 unchanged lines
                  'tflops': 0.0, 'bw': 0.0, '_tag': float('nan'),
              }
   
-         # Shape 1 (bs=16, E=257, d=256): CK 256x32 + CKTile ksplit=4
-         cfg[(cu, 16, 7168, 256, 257, 9, act_s, dtype_s, q_a, q_w, q_type, True, False)] = \
-             make_entry(32, 4, kn1_32, kn2_32)
-         # Shape 2 (bs=128, E=257, d=256): CK 256x32 + CKTile ksplit=4
-         cfg[(cu, 128, 7168, 256, 257, 9, act_s, dtype_s, q_a, q_w, q_type, True, False)] = \
-             make_entry(32, 4, kn1_32, kn2_32)
-         # Shape 3 (bs=512, E=257, d=256): CK 256x32 + FlyDSL GEMM2 reduce
-         # FlyDSL reduce (252) beats default (256) and ksplit4 (272)
-         cfg[(cu, 512, 7168, 256, 257, 9, act_s, dtype_s, q_a, q_w, q_type, True, False)] = \
-             make_entry(32, 0, kn1_32, kn2_fly32_reduce)
-         # Shape 4 (bs=16, E=33, d=512): CK 256x32 + CKTile ksplit=4
-         cfg[(cu, 16, 7168, 512, 33, 9, act_s, dtype_s, q_a, q_w, q_type, True, False)] = \
-             make_entry(32, 4, kn1_32, kn2_32)
-         # Shape 5 (bs=128, E=33, d=512): CK 256x32 + CKTile ksplit=4
-         cfg[(cu, 128, 7168, 512, 33, 9, act_s, dtype_s, q_a, q_w, q_type, True, False)] = \
-             make_entry(32, 4, kn1_32, kn2_32)
-         # Shape 6 (bs=512, E=33, d=512): CK 256x64 + FlyDSL GEMM2 reduce (proven -22%)
-         cfg[(cu, 512, 7168, 512, 33, 9, act_s, dtype_s, q_a, q_w, q_type, True, False)] = \
-             make_entry(64, 0, kn1_64, kn2_fly64_reduce)
-         # Shape 7 (bs=512, E=33, d=2048): CK 256x64 + FlyDSL GEMM2 t32 atomic
-         # Try smaller tile_m=32 for better occupancy with large K
-         cfg[(cu, 512, 7168, 2048, 33, 9, act_s, dtype_s, q_a, q_w, q_type, True, False)] = \
-             make_entry(64, 0, kn1_64, kn2_fly32_atomic)
+         def make_1stage(block_m):
+             return {
+                 'block_m': block_m, 'ksplit': 0,
+                 'kernelName1': '', 'kernelName2': '',
+                 'run_1stage': 1,
+                 'us': 0.0, 'us1': 0.0, 'us2': 0.0,
+                 'err1': '0', 'err2': '0',
+                 'tflops': 0.0, 'bw': 0.0, '_tag': float('nan'),
+             }
   
+         def key(bs, d_expert, E, topk):
+             return (cu, bs, 7168, d_expert, E, topk,
+                     act_s, dtype_s, q_a, q_w, q_type, True, False)
+  
+         # Shape 1 (bs=16, E=257, d=256): 2-stage CK 256x32 ksplit=4
+         new_cfg[key(16, 256, 257, 9)] = make_2stage(32, 4, kn1_32, kn2_32)
+         # Shape 2 (bs=128, E=257, d=256): 2-stage CK 256x32 ksplit=4
+         new_cfg[key(128, 256, 257, 9)] = make_2stage(32, 4, kn1_32, kn2_32)
+         # Shape 3 (bs=512, E=257, d=256): ★ 1-STAGE ★ (275µs vs ~308µs 2-stage)
+         new_cfg[key(512, 256, 257, 9)] = make_1stage(32)
+         # Shape 4 (bs=16, E=33, d=512): 2-stage CK 256x32 ksplit=2
+         new_cfg[key(16, 512, 33, 9)] = make_2stage(32, 2, kn1_32, kn2_32)
+         # Shape 5 (bs=128, E=33, d=512): 2-stage CK 256x32 ksplit=2
+         new_cfg[key(128, 512, 33, 9)] = make_2stage(32, 2, kn1_32, kn2_32)
+         # Shape 6 (bs=512, E=33, d=512): CK gemm1 256x64 + FlyDSL gemm2 reduce
+         new_cfg[key(512, 512, 33, 9)] = make_2stage(64, 0, kn1_64, kn2_fly64_reduce)
+         # Shape 7 (bs=512, E=33, d=2048): CK gemm1 256x64 + FlyDSL gemm2 atomic
+         new_cfg[key(512, 2048, 33, 9)] = make_2stage(64, 0, kn1_64, kn2_fly64_atomic)
+  
+         _fm.cfg_2stages = new_cfg
+  
+         # Clear all caches
          if hasattr(_fm.get_2stage_cfgs, 'cache_clear'):
              _fm.get_2stage_cfgs.cache_clear()
          for fn_name in ['get_block_size_M', 'use_nt', 'get_ksplit']:
  ⋯ 19 unchanged lines
      hidden_pad = config["d_hidden_pad"] - config["d_hidden"]
      intermediate_pad = config["d_expert_pad"] - config["d_expert"]
   
+     # block_size_M override for shape 5 (bs=128, E=33)
+     try:
+         bsm = 32 if hidden_states.shape[0] == 128 and gate_up_weight_shuffled.shape[0] == 33 else None
+     except Exception:
+         bsm = None
+  
      return fused_moe(
          hidden_states, gate_up_weight_shuffled, down_weight_shuffled,
          topk_weights, topk_ids,
  ⋯ 2 unchanged lines
          w1_scale=gate_up_weight_scale_shuffled,
          w2_scale=down_weight_scale_shuffled,
          a1_scale=None, a2_scale=None,
-         block_size_M=32 if hidden_states.shape[0] == 128 and gate_up_weight_shuffled.shape[0] == 33 else None,
+         block_size_M=bsm,
          hidden_pad=hidden_pad, intermediate_pad=intermediate_pad,
      )