submission 649696

  ⋯ 4 unchanged lines
  from typing import Dict
  from task import input_t, output_t
   
- from aiter import ActivationType, QuantType
+ from aiter import ActivationType, QuantType, dtypes
  from aiter.fused_moe import fused_moe
+ import aiter.fused_moe as _fm
   
+ # ── EXP-16: Selective tuned config injection ──
+ # Only inject for shapes 4 and 5 (E=33, bs=16/128, d_expert=512)
+ # where EXP-15 showed -31% and -9% improvement.
+ # Leave all other shapes at baseline default.
   
+ _patched = False
+  
+ def _inject_winning_configs():
+     global _patched
+     if _patched:
+         return
+     _patched = True
+  
+     try:
+         cfg = _fm.cfg_2stages
+         if cfg is None:
+             return
+  
+         cu = _fm.get_cu_count() if hasattr(_fm, 'get_cu_count') else 256
+  
+         act = str(ActivationType.Silu)
+         dtype_s = str(torch.bfloat16)
+         q_a = str(dtypes.fp4x2)
+         q_w = str(dtypes.fp4x2)
+         q_type = str(QuantType.per_1x32)
+  
+         # Large tile kernels that won for shapes 4+5
+         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'
+  
+         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'),
+             }
+  
+         # Only shapes 4 and 5 (E=33, d_expert=512, bs=16 and bs=128)
+         cfg[(cu, 16, 7168, 512, 33, 9, act, dtype_s, q_a, q_w, q_type, True, False)] = \
+             make_entry(32, 4, kn1, kn2)
+         cfg[(cu, 128, 7168, 512, 33, 9, act, dtype_s, q_a, q_w, q_type, True, False)] = \
+             make_entry(32, 4, kn1, kn2)
+  
+         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 as e:
+         pass
+  
+  
  def custom_kernel(data: input_t) -> output_t:
+     global _patched
+  
      (
-         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,
+         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"]
   
+     if not _patched:
+         output = fused_moe(
+             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=gate_up_weight_scale_shuffled,
+             w2_scale=down_weight_scale_shuffled,
+             a1_scale=None, a2_scale=None,
+             hidden_pad=hidden_pad, intermediate_pad=intermediate_pad,
+         )
+         _inject_winning_configs()
+         return output
+  
      output = fused_moe(
-         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,
+         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=gate_up_weight_scale_shuffled,
          w2_scale=down_weight_scale_shuffled,
-         a1_scale=None,
-         a2_scale=None,
-         hidden_pad=hidden_pad,
-         intermediate_pad=intermediate_pad,
+         a1_scale=None, a2_scale=None,
+         hidden_pad=hidden_pad, intermediate_pad=intermediate_pad,
      )
   
      return output