submission 662595

  ⋯ 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: Import-time cfg injection + shape 7 block_m=64 ──
- # Pre-load cfg_2stages by calling get_2stage_cfgs at import time,
- # then inject our tuned entries before any benchmark call.
- # Shapes 1,2,4,5: block_m=32, ksplit=4, 256x32 tile (EXP-32)
- # Shape 7: block_m=64, ksplit=0, 256x64 tile (EXP-49)
+ # ── EXP-49: Best known production config (confirmed ceiling) ──
   
  def _preload_and_inject():
-     """Trigger cfg_2stages loading and inject tuned configs at import time."""
      try:
-         # Call get_2stage_cfgs with shape 1 args to trigger CSV loading
-         # This only loads config (no JIT), so it's fast
          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
   
-         # Trigger loading with shape 1 args (will load CSV + cache result)
-         _fm.get_2stage_cfgs(
-             16, 7168, 256, 257, 9,       # token, model_dim, inter_dim, expert, topk
-             dtype_t, q_a_t, q_w_t, q_type_t,  # dtype, q_dtype_a, q_dtype_w, q_type
-             True, act, False,              # use_g1u1, activation, doweight_stage1
-             0, 0, True                     # hidden_pad, intermediate_pad, is_shuffled
-         )
-  
-         # Now cfg_2stages is loaded. Inject our entries.
+         _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 = 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)
+         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)
   
-         kn1 = 'moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16'
-         kn2 = 'moe_ck2stages_gemm2_256x32x128x128_1x4_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16'
+         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'
   
-         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'),
-             }
+         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')}
   
-         # Shapes 1,2 (E=257, bs=16/128) — 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, kn2)
-         cfg[(cu, 128, 7168, 256, 257, 9, act_s, dtype_s, q_a, q_w, q_type, True, False)] = \
-             make_entry(32, 4, kn1, kn2)
+         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)
   
-         # Shapes 4,5 (E=33, bs=16/128) — 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, kn2)
-         cfg[(cu, 128, 7168, 512, 33, 9, act_s, dtype_s, q_a, q_w, q_type, True, False)] = \
-             make_entry(32, 4, kn1, kn2)
-  
-         # Shapes 3,6,7: NOT injected — default CK path optimal
-         # Except: Shape 7 (bs=512, E=33, d=2048) — try block_m=64 (default=128)
-         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_64)
-  
-         # Clear ALL caches so next calls use our injected entries
          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']:
              fn = getattr(_fm, fn_name, None)
              if fn and hasattr(fn, 'cache_clear'):
                  fn.cache_clear()
+     except Exception:
+         import traceback; traceback.print_exc()
   
-     except Exception as e:
-         pass
-  
- # Execute injection at import time
  _preload_and_inject()
   
-  
  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_pad = config["d_hidden_pad"] - config["d_hidden"]
-     intermediate_pad = config["d_expert_pad"] - config["d_expert"]
-  
+     (hidden_states, _, _, _, _,
+      w1_shuffled, w2_shuffled,
+      w1_scale_shuffled, w2_scale_shuffled,
+      topk_weights, topk_ids, config) = data
      return fused_moe(
-         hidden_states, gate_up_weight_shuffled, down_weight_shuffled,
+         hidden_states, w1_shuffled, w2_shuffled,
          topk_weights, topk_ids,
          expert_mask=None, activation=ActivationType.Silu,
          quant_type=QuantType.per_1x32, doweight_stage1=False,
-         w1_scale=gate_up_weight_scale_shuffled,
-         w2_scale=down_weight_scale_shuffled,
+         w1_scale=w1_scale_shuffled,
+         w2_scale=w2_scale_shuffled,
          a1_scale=None, a2_scale=None,
-         hidden_pad=hidden_pad, intermediate_pad=intermediate_pad,
-     )
+         hidden_pad=config["d_hidden_pad"] - config["d_hidden"],
+         intermediate_pad=config["d_expert_pad"] - config["d_expert"])