submission 690372

  """
- V56: V52 + CSV priority fix (custom last) + per-shape FlyDSL tuning.
- - Custom CSV now LAST in priority chain (highest priority)
- - Shape 6 (d=512, K=512): t64x256x256_reduce (original, better for short K)
- - Shape 7 (d=2048, K=2048): t64x128x256_reduce (more N-tiles, better for long K)
+ V82: Always re-sort + fused_moe_2stages for correctness in both benchmark and leaderboard.
+   - First call: warm up with fused_moe, cache metadata + pre-allocate buffers
+   - All subsequent calls: re-sort, then call fused_moe_2stages
+   - Cktile for E=33 small-batch shapes via ksplit=2
+   - Custom CSV tuning for E=33 large-batch shapes
+   - Pre-allocated sort buffers + output buffer to save allocation overhead
  """
  #!POPCORN leaderboard amd-moe-mxfp4
  #!POPCORN gpu MI355X
  ⋯ 1 unchanged lines
  import os
  import aiter
  from aiter import ActivationType, QuantType, dtypes
- from aiter.fused_moe import fused_moe, get_2stage_cfgs
+ from aiter.fused_moe import (
+     fused_moe, fused_moe_2stages, get_2stage_cfgs, get_inter_dim, get_padded_M,
+ )
  from task import input_t, output_t
   
- CK_S1_32 = "moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
- CK_S1_256x32 = "moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
  CK_S1_64 = "moe_ck2stages_gemm1_256x64x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
- CK_S2_32_v1 = "moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"
  FLYDSL_S2_64 = "flydsl_moe2_afp4_wfp4_bf16_t64x256x256_reduce"
- FLYDSL_S2_64_128 = "flydsl_moe2_afp4_wfp4_bf16_t64x128x256_reduce"
+ FLYDSL_S2_128x256 = "flydsl_moe2_afp4_wfp4_bf16_t64x128x256_reduce"
   
  _CSV_HEADER = "cu_num,token,model_dim,inter_dim,expert,topk,act_type,dtype,q_dtype_a,q_dtype_w,q_type,use_g1u1,doweight_stage1,block_m,ksplit,us1,kernelName1,err1,us2,kernelName2,err2,us,run_1stage,tflops,bw"
   
  ⋯ 2 unchanged lines
   
  def _build_custom_csv():
      rows = [_CSV_HEADER]
-     # E=33, d=512, bs=16/128: CK small tiles (block_m=32)
-     rows.append(_row(256, 16, 7168, 512, 33, 9, 32, CK_S1_32, CK_S2_32_v1))
-     rows.append(_row(256, 128, 7168, 512, 33, 9, 32, CK_S1_32, CK_S2_32_v1))
-     # E=33, d=512, bs=512: block_m=64 + FlyDSL t64x256x256 (optimal for K=512)
      rows.append(_row(256, 512, 7168, 512, 33, 9, 64, CK_S1_64, FLYDSL_S2_64))
-     # E=33, d=2048, bs=512: block_m=64 + FlyDSL t64x128x256 (optimal for K=2048)
-     rows.append(_row(256, 512, 7168, 2048, 33, 9, 64, CK_S1_64, FLYDSL_S2_64_128))
+     rows.append(_row(256, 512, 7168, 2048, 33, 9, 64, CK_S1_64, FLYDSL_S2_128x256))
      return "\n".join(rows) + "\n"
   
+ _CKTILE_SHAPES = {
+     (16, 33),
+     (128, 33),
+ }
+  
+ _cache = {}
  _initialized = False
+ _sort_fn = None
   
+  
  def custom_kernel(data: input_t) -> output_t:
-     global _initialized
+     global _initialized, _sort_fn
   
      (
          hidden_states, gate_up_weight, down_weight,
  ⋯ 3 unchanged lines
          topk_weights, topk_ids, config,
      ) = data
   
+     M = hidden_states.shape[0]
      d_hidden = config["d_hidden"]
      d_hidden_pad = config["d_hidden_pad"]
+     d_expert = config["d_expert"]
      d_expert_pad = config["d_expert_pad"]
+     E = config["n_routed_experts"] + config["n_shared_experts"]
+     topk = config["total_top_k"]
      hidden_pad = d_hidden_pad - d_hidden
-     intermediate_pad = d_expert_pad - config["d_expert"]
+     intermediate_pad = d_expert_pad - d_expert
   
      w1s = gate_up_weight_scale_shuffled.view(dtypes.fp8_e8m0)
      w2s = down_weight_scale_shuffled.view(dtypes.fp8_e8m0)
   
      if not _initialized:
-         csv_path = "/tmp/custom_tuned_fmoe.csv"
+         csv_path = "/tmp/custom_tuned_fmoe_v82.csv"
          with open(csv_path, 'w') as f:
              f.write(_build_custom_csv())
-  
          aiter_root = os.path.dirname(os.path.abspath(aiter.__file__))
          default_csv = os.path.join(aiter_root, "configs", "tuned_fmoe.csv")
          dsv3_csv = os.path.join(aiter_root, "configs", "model_configs", "dsv3_fp4_tuned_fmoe.csv")
-  
          paths = []
          if os.path.exists(default_csv):
              paths.append(default_csv)
          if os.path.exists(dsv3_csv):
              paths.append(dsv3_csv)
-         paths.append(csv_path)  # Custom last = highest priority (last-wins)
-  
+         paths.append(csv_path)
          os.environ["AITER_CONFIG_FMOE"] = ":".join(paths)
+         os.environ["AITER_KSPLIT"] = "0"
          get_2stage_cfgs.cache_clear()
+         _sort_fn = getattr(aiter, 'moe_sorting_opus_fwd', aiter.moe_sorting_fwd)
          _initialized = True
   
-     return fused_moe(
-         hidden_states, gate_up_weight_shuffled, down_weight_shuffled,
-         topk_weights, topk_ids,
-         activation=ActivationType.Silu, quant_type=QuantType.per_1x32,
+     shape_key = (M, E, d_expert)
+     is_cktile = (M, E) in _CKTILE_SHAPES
+  
+     if shape_key not in _cache:
+         # FIRST CALL: warm up + build cache
+         if is_cktile:
+             os.environ["AITER_KSPLIT"] = "2"
+         else:
+             os.environ["AITER_KSPLIT"] = "0"
+  
+         result = fused_moe(
+             hidden_states, gate_up_weight_shuffled, down_weight_shuffled,
+             topk_weights, topk_ids,
+             activation=ActivationType.Silu, quant_type=QuantType.per_1x32,
+             w1_scale=w1s, w2_scale=w2s,
+             hidden_pad=hidden_pad, intermediate_pad=intermediate_pad,
+         )
+  
+         is_shuffled = getattr(gate_up_weight_shuffled, "is_shuffled", False)
+         E2, model_dim, inter_dim = get_inter_dim(
+             gate_up_weight_shuffled.shape, down_weight_shuffled.shape
+         )
+         metadata = get_2stage_cfgs(
+             get_padded_M(M), model_dim, inter_dim, E, topk,
+             torch.bfloat16, dtypes.fp4x2, dtypes.fp4x2,
+             QuantType.per_1x32, True, ActivationType.Silu,
+             False, hidden_pad, intermediate_pad, is_shuffled,
+         )
+         block_m = metadata.block_m
+  
+         device = hidden_states.device
+         max_tok = M * topk + E * block_m
+         max_mb = (max_tok + block_m - 1) // block_m
+  
+         sorted_ids = torch.empty(max_tok, dtype=torch.int32, device=device)
+         sorted_weights = torch.empty(max_tok, dtype=torch.float32, device=device)
+         sorted_expert_ids = torch.empty(max_mb, dtype=torch.int32, device=device)
+         num_valid_ids = torch.empty(2, dtype=torch.int32, device=device)
+         moe_buf = torch.zeros((M, d_hidden_pad), dtype=torch.bfloat16, device=device)
+  
+         _cache[shape_key] = {
+             'block_m': block_m,
+             'sorted_ids': sorted_ids,
+             'sorted_weights': sorted_weights,
+             'sorted_expert_ids': sorted_expert_ids,
+             'num_valid_ids': num_valid_ids,
+             'moe_buf': moe_buf,
+         }
+         return result
+  
+     # SUBSEQUENT CALLS: always re-sort + fused_moe_2stages
+     c = _cache[shape_key]
+  
+     # Re-sort (handles both benchmark same-data and leaderboard new-data correctly)
+     _sort_fn(topk_ids, topk_weights, c['sorted_ids'], c['sorted_weights'],
+              c['sorted_expert_ids'], c['num_valid_ids'], c['moe_buf'],
+              E, c['block_m'], None, None, 0)
+  
+     c['moe_buf'].zero_()
+  
+     fused_moe_2stages(
+         hidden_states,
+         gate_up_weight_shuffled, down_weight_shuffled,
+         topk, c['sorted_ids'], c['sorted_weights'],
+         c['sorted_expert_ids'], c['num_valid_ids'], c['moe_buf'],
+         True, c['block_m'],
+         activation=ActivationType.Silu,
+         quant_type=QuantType.per_1x32,
+         doweight_stage1=False,
+         q_dtype_a=dtypes.fp4x2,
+         q_dtype_w=dtypes.fp4x2,
          w1_scale=w1s, w2_scale=w2s,
+         a1_scale=None, a2_scale=None,
          hidden_pad=hidden_pad, intermediate_pad=intermediate_pad,
      )
+  
+     return c['moe_buf']