submission 681617

  ⋯ 1 unchanged lines
  #!POPCORN gpu MI355X
   
  import torch
+ 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
   
- # Module-level constants - avoid recomputing per call
- _ACTIVATION = ActivationType.Silu
- _QUANT_TYPE = QuantType.per_1x32
+ # ── EXP-104 best config: FlyDSL gemm2 shapes 6+7 ──
+ # Shapes 1,2: CK 256x32 ksplit=4
+ # Shape 3: CSV default (64x32)
+ # Shapes 4,5: CK 256x32 ksplit=2
+ # Shape 5: block_size_M=32 override
+ # Shape 6: CK gemm1 256x64 + FlyDSL gemm2 t64x256x256_reduce (-22%)
+ # Shape 7: CK gemm1 256x64 + FlyDSL gemm2 t64x256x256_atomic (-2.4%)
   
- # Cache for padding values keyed by config shape
- _pad_cache: dict = {}
   
+ 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
   
+         _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)
+  
+         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'
+  
+         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'),
+             }
+  
+         # Shapes 1,2,3 (E=257): NO injection — CSV default handles these
+         # 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 t64 atomic
+         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_fly64_atomic)
+  
+         # NOTE: Do NOT clear get_2stage_cfgs cache — it would re-read CSV
+         # and overwrite our E=33 injections. Only clear dependent caches.
+         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:
+         import traceback; traceback.print_exc()
+  
+ _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,
+         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
   
-     # Cache padding computation
-     cache_key = (config["d_hidden"], config["d_hidden_pad"],
-                  config["d_expert"], config["d_expert_pad"])
-     if cache_key not in _pad_cache:
-         _pad_cache[cache_key] = (
-             config["d_hidden_pad"] - config["d_hidden"],
-             config["d_expert_pad"] - config["d_expert"],
-         )
-     hidden_pad, intermediate_pad = _pad_cache[cache_key]
+     hidden_pad = config["d_hidden_pad"] - config["d_hidden"]
+     intermediate_pad = config["d_expert_pad"] - config["d_expert"]
   
+     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,
-         activation=_ACTIVATION,
-         quant_type=_QUANT_TYPE,
-         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,
-         hidden_pad=hidden_pad,
-         intermediate_pad=intermediate_pad,
+         a1_scale=None, a2_scale=None,
+         block_size_M=bsm,
+         hidden_pad=hidden_pad, intermediate_pad=intermediate_pad,
      )