submission 658848

  ⋯ 5 unchanged lines
  import torch
  from task import input_t, output_t
   
- # Import WITHOUT setting AITER_CONFIG_FMOE — let aiter use its defaults
+ # System-level env vars (set before any aiter imports trigger caching)
+ os.environ["HIP_FORCE_DEV_KERNARG"] = "1"
+ os.environ["AITER_USE_NT"] = "1"
+  
  from aiter import ActivationType, QuantType
  from aiter.fused_moe import fused_moe
  import aiter.fused_moe as _fm
   
- # Append shape 6 tuned entry (expert=33 not in DSV3 CSV)
+ # Register t16x256x128_atomic in _KERNEL_PARAMS (not in default stage2 tile_ms)
+ try:
+     from aiter.ops.flydsl.moe_kernels import _KERNEL_PARAMS
+     _t16_name = "flydsl_moe2_afp4_wfp4_bf16_t16x256x128_atomic"
+     if _t16_name not in _KERNEL_PARAMS:
+         _KERNEL_PARAMS[_t16_name] = {
+             "stage": 2, "a_dtype": "fp4", "b_dtype": "fp4", "out_dtype": "bf16",
+             "tile_m": 16, "tile_n": 256, "tile_k": 128,
+             "mode": "atomic", "MPerBlock": 16,
+         }
+         print(f"[moe] Registered {_t16_name} in _KERNEL_PARAMS", file=sys.stderr)
+ except Exception as e:
+     print(f"[moe] _KERNEL_PARAMS registration failed: {e}", file=sys.stderr)
+  
+ # CSV injection: shapes 6+7 with FlyDSL t16x256x128 stage2
  _tune_file = _fm.AITER_CONFIGS.AITER_CONFIG_FMOE_FILE
  try:
      with open(_tune_file, "a") as f:
-         # Shape 6: 32exp, M=512, inter=512, block_m=32 — both kernels 0.0% error
-         f.write("256,512,7168,512,33,9,ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,0,moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16,0.0%,0,moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16,0.0%,0,0,0,0\n")
-     print(f"[moe] Appended shape 6 entry to {_tune_file}", file=sys.stderr)
+         # Shape 6: E=33, inter=512, block_m=32 CK stage1, FlyDSL t16 stage2
+         f.write("256,512,7168,512,33,9,ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,32,0,0,moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16,0.0%,0,flydsl_moe2_afp4_wfp4_bf16_t16x256x128_atomic,0.0%,0,0,0,0\n")
+         # Shape 7: E=33, inter=2048, block_m=64 CK stage1, FlyDSL t16 stage2
+         f.write("256,512,7168,2048,33,9,ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,1,0,64,0,0,moe_ck2stages_gemm1_256x64x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16,0.0%,0,flydsl_moe2_afp4_wfp4_bf16_t16x256x128_atomic,0.0%,0,0,0,0\n")
+     print("[moe] CSV appended: shapes 6+7 FlyDSL t16x256x128", file=sys.stderr)
  except Exception as e:
-     print(f"[moe] Failed: {e}", file=sys.stderr)
+     print(f"[moe] CSV append failed: {e}", file=sys.stderr)
   
+ # Pre-compile FlyDSL stage2 MLIR for both shapes
+ try:
+     from aiter.ops.flydsl.moe_kernels import _get_compiled_stage2
+     _get_compiled_stage2(
+         model_dim=7168, inter_dim=512, experts=33, topk=9,
+         tile_m=16, tile_n=256, tile_k=128,
+         doweight=False, a_dtype="fp4", b_dtype="fp4", out_dtype="bf16",
+         accumulate=True,
+     )
+     _get_compiled_stage2(
+         model_dim=7168, inter_dim=2048, experts=33, topk=9,
+         tile_m=16, tile_n=256, tile_k=128,
+         doweight=False, a_dtype="fp4", b_dtype="fp4", out_dtype="bf16",
+         accumulate=True,
+     )
+     print("[moe] FlyDSL t16x256x128 MLIR pre-compiled for shapes 6,7", file=sys.stderr)
+ except Exception as e:
+     print(f"[moe] FlyDSL t16 MLIR pre-compile failed: {e}", file=sys.stderr)
+  
+ # GPU binary warmup: trigger @flyc.jit for t16x256x128 (shape 7, largest)
+ try:
+     from aiter.ops.flydsl.moe_kernels import flydsl_moe_stage2
+     _dev = "cuda"
+     _E = 33
+     _model_dim = 7168
+     _inter_dim = 2048
+     _inter_dim_packed = 1024
+     _tile_m = 16
+     _topk = 9
+  
+     _dummy_inter = torch.zeros((_tile_m, _topk, _inter_dim_packed), dtype=torch.uint8, device=_dev)
+     _dummy_w2 = torch.zeros((_E, _model_dim, _inter_dim_packed), dtype=torch.uint8, device=_dev)
+     _dummy_w2_scale = torch.ones((_E, _model_dim, _inter_dim // 32), dtype=torch.uint8, device=_dev)
+     _dummy_a2_scale = torch.ones((_tile_m * _topk, _inter_dim // 32), dtype=torch.uint8, device=_dev)
+     _dummy_sorted_ids = torch.arange(_tile_m, dtype=torch.int32, device=_dev)
+     _dummy_expert_ids = torch.zeros(1, dtype=torch.int32, device=_dev)
+     _dummy_num_valid = torch.tensor([_tile_m] + [0] * (_E - 1), dtype=torch.int32, device=_dev)
+  
+     flydsl_moe_stage2(
+         _dummy_inter, _dummy_w2,
+         _dummy_sorted_ids, _dummy_expert_ids, _dummy_num_valid,
+         topk=_topk, tile_m=_tile_m, tile_n=256, tile_k=128,
+         a_dtype="fp4", b_dtype="fp4", out_dtype="bf16", mode="atomic",
+         w2_scale=_dummy_w2_scale, a2_scale=_dummy_a2_scale,
+     )
+  
+     del _dummy_inter, _dummy_w2, _dummy_w2_scale, _dummy_a2_scale
+     del _dummy_sorted_ids, _dummy_expert_ids, _dummy_num_valid
+     torch.cuda.empty_cache()
+     print("[moe] FlyDSL t16x256x128 GPU warmup complete", file=sys.stderr)
+ except Exception as e:
+     print(f"[moe] FlyDSL t16 warmup failed: {e}", file=sys.stderr)
+  
  _PAD_CACHE = {}
   
   
  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
   
      M = topk_ids.shape[0]