submission 690901

  ⋯ 8 unchanged lines
  from aiter.fused_moe import fused_moe
  import aiter.fused_moe as _fm
   
- # ── 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
+ # ── EXP-133: S7 FlyDSL t32x128x256 atomic ──
+ # Single variable from the proven baseline:
+ #   Shape 7 stage2: t64x256x256_atomic -> t32x128x256_atomic
+ # Rationale:
+ # - EXP-126 showed t32 atomic improves local S7 vs t64 atomic.
+ # - EXP-132 showed tile_n=128 is valid on runner and trims S7 further.
+ # Combine both only on S7, keep all other shapes on the baseline path.
   
   
- 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_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)
+         act = ActivationType.Silu
+         dtype_t = torch.bfloat16
+         q_a_t = dtypes.fp4x2
+         q_w_t = dtypes.fp4x2
+         q_type_t = QuantType.per_1x32
   
-         # 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
+         _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
+         )
   
-         # ── Kernel names ──
+         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)
+  
          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_fly64_reduce = 'flydsl_moe2_afp4_wfp4_bf16_t64x256x256_reduce'
          kn2_fly64_atomic = 'flydsl_moe2_afp4_wfp4_bf16_t64x256x256_atomic'
+         kn2_fly32_atomic_tn128 = 'flydsl_moe2_afp4_wfp4_bf16_t32x128x256_atomic'
   
-         def make_2stage(block_m, ksplit, kn1, kn2):
+         def make_entry(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'),
              }
   
-         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'),
-             }
+         # Shape 1 (bs=16, E=257, d=256): CK 256x32 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 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): NO injection — CSV default
+         # Shape 4 (bs=16, E=33, d=512): CK 256x32 ksplit=2
+         cfg[(cu, 16, 7168, 512, 33, 9, act_s, dtype_s, q_a, q_w, q_type, True, False)] = \
+             make_entry(32, 2, kn1_32, kn2_32)
+         # Shape 5 (bs=128, E=33, d=512): CK 256x32 ksplit=2
+         cfg[(cu, 128, 7168, 512, 33, 9, act_s, dtype_s, q_a, q_w, q_type, True, False)] = \
+             make_entry(32, 2, kn1_32, kn2_32)
+         # Shape 6 (bs=512, E=33, d=512): CK gemm1 256x64 + FlyDSL gemm2 t64 reduce
+         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 gemm1 256x64 + FlyDSL gemm2 t32 atomic tn128
+         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_tn128)
   
-         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: