submission 682681

  #!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
+  
+ import gc
  import os
  import sys
  import torch
  from task import input_t, output_t
   
- # System-level env vars (set before any aiter imports trigger caching)
  os.environ["HIP_FORCE_DEV_KERNARG"] = "1"
  os.environ["AITER_USE_NT"] = "1"
+ os.environ["HSA_ENABLE_INTERRUPT"] = "0"
   
  from aiter import ActivationType, QuantType
  from aiter.fused_moe import fused_moe
  import aiter.fused_moe as _fm
   
- # Import separated quant + sort functions
  from aiter.ops.quant import per_1x32_f4_quant_hip as _quant_hip
  from aiter.utility import fp4_utils as _fp4_utils
   
- # Monkey-patch: replace fused quant+sort with separated HIP quant + Triton sort
  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(
  ⋯ 3 unchanged lines
      return x_fp4, sorted_scale
   
  _fm.fused_dynamic_mxfp4_quant_moe_sort = _split_quant_moe_sort
- print("[moe] Monkey-patched fused_quant+sort → split HIP quant + Triton sort", file=sys.stderr)
+ print("[moe] Split quant patched + HSA polling mode", file=sys.stderr)
   
- # Register t16x256x128_atomic in _KERNEL_PARAMS
  try:
      from aiter.ops.flydsl.moe_kernels import _KERNEL_PARAMS
      _t16_name = "flydsl_moe2_afp4_wfp4_bf16_t16x256x128_atomic"
  ⋯ 3 unchanged lines
              "tile_m": 16, "tile_n": 256, "tile_k": 128,
              "mode": "atomic", "MPerBlock": 16,
          }
-         print(f"[moe] Registered {_t16_name}", file=sys.stderr)
  except Exception as e:
      print(f"[moe] _KERNEL_PARAMS registration failed: {e}", file=sys.stderr)
   
- # CSV injection: shapes 5+6+7 with FlyDSL t16 stage2
- # Shape 5 (NEW): M=128, E=33, inter=512, block_m=64 — fp4 2-stage instead of BYPASS
  _tune_file = _fm.AITER_CONFIGS.AITER_CONFIG_FMOE_FILE
  try:
      with open(_tune_file, "a") as f:
-         # Shape 5: M=128, E=33, inter=512, block_m=64
          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: M=512, E=33, inter=512, block_m=128 (v67: -6.7% from 120→112µs)
          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: M=512, E=33, inter=2048, block_m=64
          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: shapes 5+7 bm=64, shape 6 bm=128, all FlyDSL t16", file=sys.stderr)
  except Exception as e:
      print(f"[moe] CSV append failed: {e}", file=sys.stderr)
   
- # Pre-compile FlyDSL stage2 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,
-     )
-     print("[moe] FlyDSL t16x256x128 MLIR pre-compiled for shapes 5,6,7", file=sys.stderr)
- except Exception as e:
-     print(f"[moe] FlyDSL t16 MLIR pre-compile failed: {e}", file=sys.stderr)
+     _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:
+     pass
   
- # GPU binary warmup for BOTH inter_dim variants
- # Shape 5 runs before shape 6 in benchmark, so we must warmup inter_dim=512 here
  try:
      from aiter.ops.flydsl.moe_kernels import flydsl_moe_stage2
      _dev = "cuda"
-     _E = 33
-     _model_dim = 7168
-     _topk = 9
-     _tile_m = 16
-  
      for _inter_dim in (512, 2048):
-         _inter_dim_packed = _inter_dim // 2
-         _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
-  
+         _ip = _inter_dim // 2
+         _di = torch.zeros((16, 9, _ip), dtype=torch.uint8, device=_dev)
+         _dw = torch.zeros((33, 7168, _ip), dtype=torch.uint8, device=_dev)
+         _dws = torch.ones((33, 7168, _inter_dim // 32), dtype=torch.uint8, device=_dev)
+         _das = torch.ones((144, _inter_dim // 32), dtype=torch.uint8, device=_dev)
+         _dsi = torch.arange(16, dtype=torch.int32, device=_dev)
+         _dei = torch.zeros(1, dtype=torch.int32, device=_dev)
+         _dnv = torch.tensor([16] + [0] * 32, dtype=torch.int32, device=_dev)
+         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()
-     print("[moe] FlyDSL t16x256x128 GPU warmup complete (inter=512,2048)", file=sys.stderr)
- except Exception as e:
-     print(f"[moe] FlyDSL t16 warmup failed: {e}", file=sys.stderr)
+ except:
+     pass
   
- # Warmup the HIP quant kernel
  try:
-     _dummy_hs = torch.randn(512, 7168, dtype=torch.bfloat16, device="cuda")
-     _dummy_fp4, _dummy_scale = _quant_hip(_dummy_hs)
-     del _dummy_hs, _dummy_fp4, _dummy_scale
+     _d = torch.randn(512, 7168, dtype=torch.bfloat16, device="cuda")
+     _quant_hip(_d)
+     del _d
      torch.cuda.empty_cache()
-     print("[moe] HIP quant kernel warmed up", file=sys.stderr)
- except Exception as e:
-     print(f"[moe] HIP quant warmup failed: {e}", file=sys.stderr)
+ except:
+     pass
   
+ # Disable GC to prevent pauses during benchmark
+ gc.disable()
+ print("[moe] GC disabled", file=sys.stderr)
+  
  _PAD_CACHE = {}
  _LAST_MODE = [None]
  _environ = os.environ
  ⋯ 20 unchanged lines
          intermediate_pad = config["d_expert_pad"] - config["d_expert"]
          _PAD_CACHE[cfg_id] = (hidden_pad, intermediate_pad)
   
-     # Per-shape mode selection:
-     # - M<=16: always BYPASS (shapes 1,4)
-     # - M<=128 + E>64: BYPASS (shape 2 with E=257)
-     # - M<=128 + E<=64: 2-stage fp4 with CSV (shape 5 with E=33)
-     # - M>128: 2-stage fp4 with CSV (shapes 3,6,7)
      if M <= 16 or (M <= 128 and E > 64):
-         mode = 1  # BYPASS
+         mode = 1
      else:
-         mode = 0  # 2-stage
+         mode = 0
   
      if mode != _LAST_MODE[0]:
          _LAST_MODE[0] = mode