submission 750364

  ⋯ 1 unchanged lines
  #!POPCORN gpu MI355X
   
  """
- v173: Optimal combined — best configs from all experiments.
- - Shapes 2,3: Switch from CKTile ksplit to CK+FlyDSL (v171/v172 discoveries)
- - Shape 2: CK_S1_M128 + FlyDSL t16x128x128, block_m=64 (154µs, was 172µs)
- - Shape 3: CK_S1_M128 + FlyDSL t16x128x128, block_m=64, NT=True (197µs, was 231µs)
- - All other shapes: same as v168 (proven optimal)
- Expected GM: ~115µs (target: <120µs)
+ v184: Best hybrid CKTile-ksplit + CK+FlyDSL.
+ - Shape 1 (E=257,bs=16): CKTile ksplit=10 (fused quant + split-K)
+ - Shape 4 (E=33,bs=16): CKTile ksplit=2 (optimal for this shape)
+ - Shapes 2,3,5,6,7: CK stage1 M128 + FlyDSL t16x128x128_atomic
  """
  import os
  import functools
  ⋯ 25 unchanged lines
          "MPerBlock": tile_m,
      }
   
-  
- # Register FlyDSL tile_k=128 kernels
  for _name, _tm, _tn in (
      ("flydsl_moe2_afp4_wfp4_bf16_t32x128x128_atomic", 32, 128),
      ("flydsl_moe2_afp4_wfp4_bf16_t32x256x128_atomic", 32, 256),
  ⋯ 2 unchanged lines
  ):
      _register_flydsl_kernel(_name, _tm, _tn)
   
- # Shape configs
  _CUSTOM_CONFIGS = {}
   
-  
  def _add_cfg(token, inter_dim, expert, block_m, ksplit, kernelName1="", kernelName2="", use_non_temporal_load=None):
      cfg = {
          "block_m": block_m,
  ⋯ 18 unchanged lines
  _4WG_STAGE1_M32 = "moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
  _FLYDSL_STAGE2_M16_K128 = "flydsl_moe2_afp4_wfp4_bf16_t16x128x128_atomic"
   
- # E=33 shapes (same as v168 — proven optimal)
- _add_cfg(16, 512, 33, block_m=32, ksplit=2)
- _add_cfg(128, 512, 33, block_m=64, ksplit=0, kernelName1=_4WG_STAGE1_M128, kernelName2=_FLYDSL_STAGE2_M16_K128)
- _add_cfg(512, 512, 33, block_m=64, ksplit=0, kernelName1=_4WG_STAGE1_M128, kernelName2=_FLYDSL_STAGE2_M16_K128)
- _add_cfg(512, 2048, 33, block_m=64, ksplit=0, kernelName1=_4WG_STAGE1_M128, kernelName2=_FLYDSL_STAGE2_M16_K128)
+ # E=33 shapes
+ _add_cfg(16, 512, 33, block_m=32, ksplit=2)   # shape 4: CKTile ksplit=2 (fused quant + split-K)
+ _add_cfg(128, 512, 33, block_m=64, ksplit=0, kernelName1=_4WG_STAGE1_M128, kernelName2=_FLYDSL_STAGE2_M16_K128)   # shape 5
+ _add_cfg(512, 512, 33, block_m=64, ksplit=0, kernelName1=_4WG_STAGE1_M128, kernelName2=_FLYDSL_STAGE2_M16_K128)   # shape 6
+ _add_cfg(512, 2048, 33, block_m=64, ksplit=0, kernelName1=_4WG_STAGE1_M128, kernelName2=_FLYDSL_STAGE2_M16_K128)  # shape 7
   
  # E=257 shapes
- _add_cfg(16, 256, 257, block_m=16, ksplit=2)  # shape 1: CKTile ksplit=2 (proven best)
- _add_cfg(128, 256, 257, block_m=64, ksplit=0, kernelName1=_4WG_STAGE1_M128, kernelName2=_FLYDSL_STAGE2_M16_K128)  # shape 2: CK+FlyDSL (v172: 154µs vs 172µs)
- _add_cfg(
-     512,
-     256,
-     257,
-     block_m=64,  # shape 3: CK_S1_M128 + block_m=64 (v171: 197µs vs 231µs)
-     ksplit=0,
-     kernelName1=_4WG_STAGE1_M128,
-     kernelName2=_FLYDSL_STAGE2_M16_K128,
-     use_non_temporal_load=True,
- )
+ _add_cfg(16, 256, 257, block_m=32, ksplit=10)  # shape 1: CKTile ksplit=10 (fused quant + split-K)
+ _add_cfg(128, 256, 257, block_m=64, ksplit=0, kernelName1=_4WG_STAGE1_M128, kernelName2=_FLYDSL_STAGE2_M16_K128, use_non_temporal_load=True)  # shape 2
+ _add_cfg(512, 256, 257, block_m=64, ksplit=0, kernelName1=_4WG_STAGE1_M128, kernelName2=_FLYDSL_STAGE2_M16_K128, use_non_temporal_load=True)  # shape 3
   
  # Pre-allocated buffer cache
  _buffer_cache = {}
  ⋯ 2 unchanged lines
      key = ("sort", M, E, topk, model_dim, block_size_M)
      if key in _buffer_cache:
          return _buffer_cache[key]
-  
      max_num_tokens_padded = int(M * topk + E * block_size_M - topk)
      max_num_m_blocks = int((max_num_tokens_padded + block_size_M - 1) // block_size_M)
-  
      bufs = {
          "sorted_ids": torch.empty(max_num_tokens_padded, dtype=dtypes.i32, device=device),
          "sorted_weights": torch.empty(max_num_tokens_padded, dtype=dtypes.fp32, device=device),
  ⋯ 16 unchanged lines
      MXFP4_QUANT_BLOCK_SIZE = 32
      BLOCK_SIZE_M, BLOCK_SIZE_N = 32, 8
      BLOCK_SIZE_M_u32, BLOCK_SIZE_N_u32 = 16, 4
-  
      key = ("quant", M, N, sorted_ids_len, topk)
      if key in _buffer_cache:
          return _buffer_cache[key]
-  
      x_fp4 = torch.empty((M, N // 2), dtype=torch.uint8, device=device)
      scaleN = triton.cdiv(N, MXFP4_QUANT_BLOCK_SIZE)
      M_o = sorted_ids_len
      N_o = scaleN
-  
      blockscale_e8m0_sorted = torch.empty(
-         (
-             triton.cdiv(M_o, BLOCK_SIZE_M),
-             triton.cdiv(N_o, BLOCK_SIZE_N),
-             BLOCK_SIZE_N_u32,
-             BLOCK_SIZE_M_u32,
-             4,
-         ),
-         dtype=torch.uint8,
-         device=device,
+         (triton.cdiv(M_o, BLOCK_SIZE_M), triton.cdiv(N_o, BLOCK_SIZE_N), BLOCK_SIZE_N_u32, BLOCK_SIZE_M_u32, 4),
+         dtype=torch.uint8, device=device,
      )
-  
      bufs = {"x_fp4": x_fp4, "blockscale": blockscale_e8m0_sorted}
      _buffer_cache[key] = bufs
      return bufs
  ⋯ 3 unchanged lines
      MXFP4_QUANT_BLOCK_SIZE = 32
      BLOCK_SIZE_Mx = 128
      BLOCK_SIZE_M, BLOCK_SIZE_N = 32, 8
-  
      scaleN = triton.cdiv(N, MXFP4_QUANT_BLOCK_SIZE)
      M_i, N_i = M, scaleN
      M_o = sorted_ids.shape[0]
-  
      qbufs = _get_or_alloc_quant_buffers(M, N, M_o, topk, device)
      x_fp4 = qbufs["x_fp4"]
      blockscale_e8m0_sorted = qbufs["blockscale"]
-  
-     num_pid = triton.cdiv(M, BLOCK_SIZE_Mx) * scaleN + triton.cdiv(
-         M_o, BLOCK_SIZE_M
-     ) * triton.cdiv(N_i, BLOCK_SIZE_N)
-  
+     num_pid = triton.cdiv(M, BLOCK_SIZE_Mx) * scaleN + triton.cdiv(M_o, BLOCK_SIZE_M) * triton.cdiv(N_i, BLOCK_SIZE_N)
      _fused_dynamic_mxfp4_quant_moe_sort_kernel[(num_pid,)](
-         x,
-         x_fp4,
-         sorted_ids,
-         num_valid_ids,
-         blockscale_e8m0_sorted,
-         M,
-         N,
-         scaleN,
-         *x.stride(),
-         *x_fp4.stride(),
-         *blockscale_e8m0_sorted.stride(),
-         token_num=token_num,
-         M_i=M_i,
-         N_i=N_i,
+         x, x_fp4, sorted_ids, num_valid_ids, blockscale_e8m0_sorted,
+         M, N, scaleN, *x.stride(), *x_fp4.stride(), *blockscale_e8m0_sorted.stride(),
+         token_num=token_num, M_i=M_i, N_i=N_i,
          MXFP4_QUANT_BLOCK_SIZE=MXFP4_QUANT_BLOCK_SIZE,
          BLOCK_SIZE_Mx=BLOCK_SIZE_Mx,
          BLOCK_SIZE_M=BLOCK_SIZE_M // 2,
          BLOCK_SIZE_N=BLOCK_SIZE_N // 2,
          TOPK=topk,
      )
-  
      return (
          x_fp4.view(dtypes.fp4x2),
          blockscale_e8m0_sorted.view(dtypes.fp8_e8m0).view(-1, scaleN),
      )
   
-  
  def _run_stage2(metadata, a2_or_quant, w1, w2, sorted_ids, sorted_expert_ids, num_valid_ids, moe_out, topk, w2_scale, a2_scale, block_size_M, sorted_weights):
      metadata.stage2(
-         a2_or_quant,
-         w1,
-         w2,
-         sorted_ids,
-         sorted_expert_ids,
-         num_valid_ids,
-         moe_out,
-         topk,
-         w2_scale=w2_scale,
-         a2_scale=a2_scale,
-         block_m=block_size_M,
-         sorted_weights=sorted_weights,
+         a2_or_quant, w1, w2, sorted_ids, sorted_expert_ids, num_valid_ids, moe_out, topk,
+         w2_scale=w2_scale, a2_scale=a2_scale, block_m=block_size_M, sorted_weights=sorted_weights,
      )
   
-  
  _injected = False
   
  def _inject_configs():
  ⋯ 1 unchanged lines
      if _injected:
          return
      _injected = True
-  
      if _fused_moe_module.cfg_2stages is None:
          import pandas as pd
          from aiter.jit.core import AITER_CONFIGS
  ⋯ 10 unchanged lines
              _fused_moe_module.cfg_2stages = df.set_index(_INDEX_COLS).to_dict("index")
          else:
              _fused_moe_module.cfg_2stages = {}
-  
      _fused_moe_module.cfg_2stages.update(_CUSTOM_CONFIGS)
-  
      _original_get_2stage_cfgs = _fused_moe_module.get_2stage_cfgs
   
      @functools.lru_cache(maxsize=2048)
  ⋯ 24 unchanged lines
                      new_kw['use_non_temporal_load'] = nt
                      metadata = _fused_moe_module.MOEMetadata(
                          functools.partial(old_s1.func, **{k: v for k, v in new_kw.items()}),
-                         metadata.stage2,
-                         metadata.block_m,
-                         metadata.ksplit,
-                         metadata.run_1stage,
-                         metadata.has_bias,
-                         nt,
+                         metadata.stage2, metadata.block_m, metadata.ksplit,
+                         metadata.run_1stage, metadata.has_bias, nt,
                      )
                      old_s2 = metadata.stage2
                      if old_s2 and hasattr(old_s2, 'keywords') and 'use_non_temporal_load' in (old_s2.keywords or {}):
  ⋯ 2 unchanged lines
                          metadata = _fused_moe_module.MOEMetadata(
                              metadata.stage1,
                              functools.partial(old_s2.func, **{k: v for k, v in new_kw2.items()}),
-                             metadata.block_m,
-                             metadata.ksplit,
-                             metadata.run_1stage,
-                             metadata.has_bias,
-                             nt,
+                             metadata.block_m, metadata.ksplit,
+                             metadata.run_1stage, metadata.has_bias, nt,
                          )
          return metadata
-  
      _fused_moe_module.get_2stage_cfgs = _patched_get_2stage_cfgs
   
-  
  def custom_kernel(data: input_t) -> output_t:
      (
          hidden_states, gate_up_weight, down_weight,
  ⋯ 2 unchanged lines
          gate_up_weight_scale_shuffled, down_weight_scale_shuffled,
          topk_weights, topk_ids, config,
      ) = data
-  
      _inject_configs()
-  
      hidden_pad = config["d_hidden_pad"] - config["d_hidden"]
      intermediate_pad = config["d_expert_pad"] - config["d_expert"]
-  
      M = hidden_states.shape[0]
      topk = topk_ids.shape[1]
      device = topk_ids.device
      w1 = gate_up_weight_shuffled
      w2 = down_weight_shuffled
      E, model_dim, inter_dim = get_inter_dim(w1.shape, w2.shape)
-  
      padded_M = get_padded_M(M)
      metadata = get_2stage_cfgs(
          padded_M, model_dim, inter_dim, E, topk,
  ⋯ 1 unchanged lines
          QuantType.per_1x32, True, ActivationType.Silu,
          False, hidden_pad, intermediate_pad, True,
      )
-  
      block_size_M = int(metadata.block_m)
-  
      bufs = _get_or_alloc_sorting_buffers(M, E, topk, model_dim, block_size_M, device)
      sorted_ids = bufs["sorted_ids"]
      sorted_weights = bufs["sorted_weights"]
      sorted_expert_ids = bufs["sorted_expert_ids"]
      num_valid_ids = bufs["num_valid_ids"]
      moe_out = bufs["moe_buf"]
-  
      aiter.moe_sorting_fwd(
          topk_ids, topk_weights,
          sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, moe_out,
          E, int(block_size_M), None, None, 0,
      )
-  
      token_num = M
      w1_scale_view = gate_up_weight_scale_shuffled.view(dtypes.fp8_e8m0)
      w2_scale_view = down_weight_scale_shuffled.view(dtypes.fp8_e8m0)
  ⋯ 1 unchanged lines
   
      if metadata.ksplit > 1:
          a1 = hidden_states.to(torch.bfloat16)
-         a1_scale = None
-  
          a2 = metadata.stage1(
-             a1, w1, w2,
-             sorted_ids, sorted_expert_ids, num_valid_ids,
-             a2_buf,
-             topk,
-             block_m=block_size_M,
-             a1_scale=a1_scale,
-             w1_scale=w1_scale_view,
-             sorted_weights=None,
+             a1, w1, w2, sorted_ids, sorted_expert_ids, num_valid_ids,
+             a2_buf, topk, block_m=block_size_M, a1_scale=None, w1_scale=w1_scale_view, sorted_weights=None,
          )
-  
-         _run_stage2(
-             metadata,
-             a2,
-             w1,
-             w2,
-             sorted_ids,
-             sorted_expert_ids,
-             num_valid_ids,
-             moe_out,
-             topk,
-             w2_scale=w2_scale_view,
-             a2_scale=None,
-             block_size_M=block_size_M,
-             sorted_weights=sorted_weights,
-         )
+         _run_stage2(metadata, a2, w1, w2, sorted_ids, sorted_expert_ids, num_valid_ids,
+             moe_out, topk, w2_scale=w2_scale_view, a2_scale=None, block_size_M=block_size_M, sorted_weights=sorted_weights)
      else:
-         a1, a1_scale = _quant_prealloc(
-             hidden_states, sorted_ids, num_valid_ids,
-             token_num, 1, block_size_M, device,
-         )
-  
+         a1, a1_scale = _quant_prealloc(hidden_states, sorted_ids, num_valid_ids, token_num, 1, block_size_M, device)
          a2 = metadata.stage1(
-             a1, w1, w2,
-             sorted_ids, sorted_expert_ids, num_valid_ids,
-             a2_buf, topk,
-             block_m=block_size_M,
-             a1_scale=a1_scale,
-             w1_scale=w1_scale_view,
-             sorted_weights=None,
+             a1, w1, w2, sorted_ids, sorted_expert_ids, num_valid_ids,
+             a2_buf, topk, block_m=block_size_M, a1_scale=a1_scale, w1_scale=w1_scale_view, sorted_weights=None,
          )
-  
          a2_flat = a2.view(-1, inter_dim)
-         a2_quant, a2_scale = _quant_prealloc(
-             a2_flat, sorted_ids, num_valid_ids,
-             token_num, topk, block_size_M, device,
-         )
+         a2_quant, a2_scale = _quant_prealloc(a2_flat, sorted_ids, num_valid_ids, token_num, topk, block_size_M, device)
          a2_quant = a2_quant.view(token_num, topk, -1)
-  
-         _run_stage2(
-             metadata,
-             a2_quant,
-             w1,
-             w2,
-             sorted_ids,
-             sorted_expert_ids,
-             num_valid_ids,
-             moe_out,
-             topk,
-             w2_scale=w2_scale_view,
-             a2_scale=a2_scale,
-             block_size_M=block_size_M,
-             sorted_weights=sorted_weights,
-         )
-  
+         _run_stage2(metadata, a2_quant, w1, w2, sorted_ids, sorted_expert_ids, num_valid_ids,
+             moe_out, topk, w2_scale=w2_scale_view, a2_scale=a2_scale, block_size_M=block_size_M, sorted_weights=sorted_weights)
      return moe_out