submission 676297

  ⋯ 1 unchanged lines
  #!POPCORN gpu MI355X
   
  import torch
+ from typing import Dict
  from task import input_t, output_t
   
  from aiter import ActivationType, QuantType, dtypes
  from aiter.fused_moe import fused_moe
  import aiter.fused_moe as _fm
   
- # ── EXP-49: Best known production config (confirmed ceiling) ──
+ # ── 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%)
   
+  
  def _preload_and_inject():
      try:
          cu = _fm.get_cu_count() if hasattr(_fm, 'get_cu_count') else 256
  ⋯ 3 unchanged lines
          q_w_t = dtypes.fp4x2
          q_type_t = 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)
+         _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
+         )
+  
          cfg = _fm.cfg_2stages
          if cfg is None:
              return
   
-         act_s, dtype_s = str(act), str(dtype_t)
-         q_a, q_w, q_type = str(q_a_t), str(q_w_t), str(q_type_t)
+         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_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 me(bm, ks, k1, k2):
-             return {'block_m': bm, 'ksplit': ks, 'kernelName1': k1, 'kernelName2': k2,
-                     'run_1stage': 0, 'us': 0.0, 'us1': 0.0, 'us2': 0.0,
-                     'err1': '0', 'err2': '0', 'tflops': 0.0, 'bw': 0.0, '_tag': float('nan')}
+         def make_entry(block_m, ksplit, kn1, kn2):
+             return {
+                 'block_m': block_m, 'ksplit': ksplit,
+                 'kernelName1': kn1, 'kernelName2': kn2,
+                 'run_1stage': 0,
+                 'us': 0.0, 'us1': 0.0, 'us2': 0.0,
+                 'err1': '0', 'err2': '0',
+                 'tflops': 0.0, 'bw': 0.0, '_tag': float('nan'),
+             }
   
-         K = (act_s, dtype_s, q_a, q_w, q_type, True, False)
-         cfg[(cu, 16, 7168, 256, 257, 9) + K] = me(32, 4, kn1_32, kn2_32)
-         cfg[(cu, 128, 7168, 256, 257, 9) + K] = me(32, 4, kn1_32, kn2_32)
-         cfg[(cu, 16, 7168, 512, 33, 9) + K] = me(32, 4, kn1_32, kn2_32)
-         cfg[(cu, 128, 7168, 512, 33, 9) + K] = me(32, 4, kn1_32, kn2_32)
-         cfg[(cu, 512, 7168, 2048, 33, 9) + K] = me(64, 0, kn1_64, kn2_64)
+         # 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)
   
          if hasattr(_fm.get_2stage_cfgs, 'cache_clear'):
              _fm.get_2stage_cfgs.cache_clear()
  ⋯ 1 unchanged lines
              fn = getattr(_fm, fn_name, None)
              if fn and hasattr(fn, 'cache_clear'):
                  fn.cache_clear()
-     except Exception:
+  
+     except Exception as e:
          import traceback; traceback.print_exc()
   
  _preload_and_inject()
   
+  
  def custom_kernel(data: input_t) -> output_t:
-     (hidden_states, _, _, _, _,
-      w1_shuffled, w2_shuffled,
-      w1_scale_shuffled, w2_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
+  
+     hidden_pad = config["d_hidden_pad"] - config["d_hidden"]
+     intermediate_pad = config["d_expert_pad"] - config["d_expert"]
+  
      return fused_moe(
-         hidden_states, w1_shuffled, w2_shuffled,
+         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=w1_scale_shuffled,
-         w2_scale=w2_scale_shuffled,
+         w1_scale=gate_up_weight_scale_shuffled,
+         w2_scale=down_weight_scale_shuffled,
          a1_scale=None, a2_scale=None,
-         hidden_pad=config["d_hidden_pad"] - config["d_hidden"],
-         intermediate_pad=config["d_expert_pad"] - config["d_expert"])
+         block_size_M=32 if hidden_states.shape[0] == 128 and gate_up_weight_shuffled.shape[0] == 33 else None,
+         hidden_pad=hidden_pad, intermediate_pad=intermediate_pad,
+     )