submission 662439

+ """
+ V25: Optimized MoE-MXFP4 with tuned CK configs for E=33, d=512 shapes only.
+ E=33 d=2048 and E=257 shapes left to DSV3+heuristic.
+ Key finding: block_m=32 with CK_S1_32/CK_S2_32_v1 is best for E=33, d=512.
+ """
  import torch
- from task import input_t, output_t
-  
+ import os
  import aiter
  from aiter import ActivationType, QuantType, dtypes
- from aiter.fused_moe import fused_moe
+ from aiter.fused_moe import fused_moe, get_2stage_cfgs
+ from task import input_t, output_t
   
- class MetaWrapper:
-     def __init__(self, orig, splitk):
-         self._orig = orig
-         self._splitk = splitk
-         
-     def __getattr__(self, name):
-         if name == "splitk":
-             return self._splitk
-         return getattr(self._orig, name)
+ CK_S1_32 = "moe_ck2stages_gemm1_64x32x32x128_1x1_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"
   
+ _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"
+  
+ def _row(cu, tok, mdim, idim, E, topk, bm, k1, k2, ks=0):
+     return f"{cu},{tok},{mdim},{idim},{E},{topk},ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,True,False,{bm},{ks},0,{k1},0,0,{k2},0,0,False,0,0"
+  
+ def _build_custom_csv():
+     rows = [_CSV_HEADER]
+     # Only E=33, d=512 shapes. All others use DSV3/heuristic.
+     # V23/V24 showed: block_m=32, CK_S1_32, CK_S2_32_v1 is best for these.
+     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))
+     rows.append(_row(256, 512, 7168, 512, 33, 9, 32, CK_S1_32, CK_S2_32_v1))
+     return "\n".join(rows) + "\n"
+  
+ _initialized = False
+  
  def custom_kernel(data: input_t) -> output_t:
+     global _initialized
+  
      (
-         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"]
-     M = hidden_states.shape[0]
+     d_hidden = config["d_hidden"]
+     d_hidden_pad = config["d_hidden_pad"]
+     d_expert_pad = config["d_expert_pad"]
+     hidden_pad = d_hidden_pad - d_hidden
+     intermediate_pad = d_expert_pad - config["d_expert"]
   
-     old_get_2stage_cfgs = aiter.fused_moe.get_2stage_cfgs
-     
-     def mocked_get_2stage_cfgs(*args, **kwargs):
-         meta = old_get_2stage_cfgs(*args, **kwargs)
-         
-         E = args[3] if len(args) > 3 else 257
-         block_m = meta.block_m if hasattr(meta, "block_m") and meta.block_m is not None and meta.block_m != -1 else 32
-         
-         num_blocks = ((M + block_m - 1) // block_m) * E
-         
-         if num_blocks > 0:
-             desired_splitk = 1200 // num_blocks
-         else:
-             desired_splitk = 1
-             
-         if desired_splitk >= 8:
-             splitk = 8
-         elif desired_splitk >= 4:
-             splitk = 4
-         elif desired_splitk >= 2:
-             splitk = 2
-         else:
-             splitk = 1
-             
-         return MetaWrapper(meta, splitk)
+     w1s = gate_up_weight_scale_shuffled.view(dtypes.fp8_e8m0)
+     w2s = down_weight_scale_shuffled.view(dtypes.fp8_e8m0)
   
-     try:
-         aiter.fused_moe.get_2stage_cfgs = mocked_get_2stage_cfgs
+     if not _initialized:
+         csv_path = "/tmp/custom_tuned_fmoe.csv"
+         with open(csv_path, 'w') as f:
+             f.write(_build_custom_csv())
   
-         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,
-         )
-     finally:
-         aiter.fused_moe.get_2stage_cfgs = old_get_2stage_cfgs
+         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")
   
-     return output
   No newline at end of file
+         paths = [csv_path]
+         if os.path.exists(dsv3_csv):
+             paths.append(dsv3_csv)
+         if os.path.exists(default_csv):
+             paths.append(default_csv)
+  
+         os.environ["AITER_CONFIG_FMOE"] = ":".join(paths)
+         get_2stage_cfgs.cache_clear()
+         _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,
+         w1_scale=w1s, w2_scale=w2s,
+         hidden_pad=hidden_pad, intermediate_pad=intermediate_pad,
+     )