submission 704012

  #!POPCORN leaderboard amd-moe-mxfp4
  #!POPCORN gpu MI355X
   
- # v76: gc.disable() + HSA polling mode (keep split quant as control)
- # HSA_ENABLE_INTERRUPT=0 → CPU polls for kernel completion instead of interrupt
- # This reduces kernel completion→next launch latency by ~0.5-1µs per transition
- # gc.disable() prevents Python GC pauses during benchmark
+ # v90: Fix shape 5 bypass regression — restore per-E routing from v55/v67
+ # Shape 5 (M=128, E=32) must use 2-stage FlyDSL, NOT bypass (~15µs faster)
   
  import gc
  import os
  import sys
- import torch
- from task import input_t, output_t
   
  os.environ["HIP_FORCE_DEV_KERNARG"] = "1"
  os.environ["AITER_USE_NT"] = "1"
  os.environ["HSA_ENABLE_INTERRUPT"] = "0"
   
+ import torch
+ from task import input_t, output_t
+  
  from aiter import ActivationType, QuantType
- from aiter.fused_moe import fused_moe
+ from aiter.fused_moe import fused_moe, moe_sorting
  import aiter.fused_moe as _fm
   
  from aiter.ops.quant import per_1x32_f4_quant_hip as _quant_hip
  from aiter.utility import fp4_utils as _fp4_utils
   
- def _split_quant_moe_sort(x, sorted_ids, num_valid_ids, token_num, topk, block_size=32, scaling_mode="even"):
-     x_fp4, scale = _quant_hip(x)
-     sorted_scale = _fp4_utils.moe_mxfp4_sort(
-         scale, sorted_ids=sorted_ids, num_valid_ids=num_valid_ids,
-         token_num=token_num, block_size=block_size,
-     )
-     return x_fp4, sorted_scale
+ # Adaptive quant: split HIP for large tensors, fused Triton for small
+ _original_fused_quant = _fm.fused_dynamic_mxfp4_quant_moe_sort
   
- _fm.fused_dynamic_mxfp4_quant_moe_sort = _split_quant_moe_sort
- print("[moe] Split quant patched + HSA polling mode", file=sys.stderr)
+ def _adaptive_quant_moe_sort(x, sorted_ids, num_valid_ids, token_num, topk, block_size=32, scaling_mode="even"):
+     if x.numel() > 1_000_000:
+         x_fp4, scale = _quant_hip(x)
+         sorted_scale = _fp4_utils.moe_mxfp4_sort(
+             scale, sorted_ids=sorted_ids, num_valid_ids=num_valid_ids,
+             token_num=token_num, block_size=block_size,
+         )
+         return x_fp4, sorted_scale
+     else:
+         return _original_fused_quant(x, sorted_ids, num_valid_ids, token_num, topk, block_size, scaling_mode)
   
+ _fm.fused_dynamic_mxfp4_quant_moe_sort = _adaptive_quant_moe_sort
+ print("[moe] Adaptive quant: split for numel>1M, fused for numel<=1M", file=sys.stderr)
+  
+ # Register FlyDSL t16x256x128 kernel params
  try:
      from aiter.ops.flydsl.moe_kernels import _KERNEL_PARAMS
      _t16_name = "flydsl_moe2_afp4_wfp4_bf16_t16x256x128_atomic"
  ⋯ 6 unchanged lines
  except Exception as e:
      print(f"[moe] _KERNEL_PARAMS registration failed: {e}", file=sys.stderr)
   
+ # CSV append for shapes 5,6,7 with tuned CK stage1 + FlyDSL t16 stage2
  _tune_file = _fm.AITER_CONFIGS.AITER_CONFIG_FMOE_FILE
  try:
      with open(_tune_file, "a") as f:
+         # Shape 5: block_m=64, CK 256x64 stage1
          f.write("256,128,7168,512,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")
+         # Shape 6: block_m=128, CK 256x128 stage1
          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,128,0,0,moe_ck2stages_gemm1_256x128x128x128_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: block_m=64, CK 256x64 stage1
          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")
  except Exception as e:
      print(f"[moe] CSV append failed: {e}", file=sys.stderr)
   
+ # Pre-compile FlyDSL MLIR for both inter_dim variants
  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)
- except:
+ except Exception:
      pass
   
+ # GPU binary warmup: trigger @flyc.jit compilation with dummy tensors
  try:
      from aiter.ops.flydsl.moe_kernels import flydsl_moe_stage2
      _dev = "cuda"
  ⋯ 9 unchanged lines
          flydsl_moe_stage2(_di, _dw, _dsi, _dei, _dnv, topk=9, tile_m=16, tile_n=256, tile_k=128, a_dtype="fp4", b_dtype="fp4", out_dtype="bf16", mode="atomic", w2_scale=_dws, a2_scale=_das)
          del _di, _dw, _dws, _das, _dsi, _dei, _dnv
      torch.cuda.empty_cache()
- except:
+ except Exception:
      pass
   
+ # HIP quant kernel warmup
  try:
      _d = torch.randn(512, 7168, dtype=torch.bfloat16, device="cuda")
      _quant_hip(_d)
      del _d
      torch.cuda.empty_cache()
- except:
+ except Exception:
      pass
   
+ # moe_sorting kernel warmup (trigger JIT compile for both E configs)
+ try:
+     for _E in (33, 257):
+         _ids = torch.zeros((16, 9), dtype=torch.int32, device="cuda")
+         _wts = torch.ones((16, 9), dtype=torch.float32, device="cuda")
+         moe_sorting(_ids, _wts, _E, 7168, torch.bfloat16, 32)
+         del _ids, _wts
+     torch.cuda.empty_cache()
+     print("[moe] moe_sorting warmed up", file=sys.stderr)
+ except Exception as e:
+     print(f"[moe] moe_sorting warmup failed: {e}", file=sys.stderr)
+  
  # Disable GC to prevent pauses during benchmark
  gc.disable()
- print("[moe] GC disabled", file=sys.stderr)
+ print("[moe] v90: fix shape5 routing + warmup + gc.disable", file=sys.stderr)
   
  _PAD_CACHE = {}
- _LAST_MODE = [None]
- _environ = os.environ
   
   
  def custom_kernel(data: input_t) -> output_t:
  ⋯ 6 unchanged lines
      ) = data
   
      M = topk_ids.shape[0]
-     E = gate_up_weight_shuffled.shape[0]
   
      cfg_id = id(config)
      cached = _PAD_CACHE.get(cfg_id)
  ⋯ 4 unchanged lines
          intermediate_pad = config["d_expert_pad"] - config["d_expert"]
          _PAD_CACHE[cfg_id] = (hidden_pad, intermediate_pad)
   
+     E = gate_up_weight_shuffled.shape[0]  # E+1 (includes shared expert)
+  
+     # Per-shape routing (v55/v67 condition):
+     # - Shapes 1,2 (M<=128, E=257): BYPASS (cktile faster for large E)
+     # - Shape 4 (M=16, E=33): BYPASS (too small for 2-stage overhead)
+     # - Shape 5 (M=128, E=33): 2-STAGE (FlyDSL t16 is ~15µs faster than bypass)
+     # - Shapes 3,6,7 (M>128): 2-STAGE (always)
      if M <= 16 or (M <= 128 and E > 64):
-         mode = 1
+         os.environ["AITER_KSPLIT"] = "2"
+         os.environ["AITER_BYPASS_TUNE_CONFIG"] = "1"
      else:
-         mode = 0
+         os.environ["AITER_KSPLIT"] = "0"
+         os.environ["AITER_BYPASS_TUNE_CONFIG"] = "0"
   
-     if mode != _LAST_MODE[0]:
-         _LAST_MODE[0] = mode
-         if mode:
-             _environ["AITER_KSPLIT"] = "2"
-             _environ["AITER_BYPASS_TUNE_CONFIG"] = "1"
-         else:
-             _environ["AITER_KSPLIT"] = "0"
-             _environ["AITER_BYPASS_TUNE_CONFIG"] = "0"
-  
      return fused_moe(
          hidden_states, gate_up_weight_shuffled, down_weight_shuffled,
          topk_weights, topk_ids,