submission 676752

  ⋯ 1 unchanged lines
  #!POPCORN gpu MI355X
   
  """
- v17: Tune 256-expert shapes. Changes from v16:
- Based on dgavriloff/amd-structkernel pattern (rank 6, 130us).
- Injects configs directly into in-memory dict, registers FlyDSL kernels,
- and monkeypatches get_2stage_cfgs for NT override.
+ v138: ksplit=2 for sparse shapes s1 (E=257/bs=16) and s4 (E=33/bs=16).
+  
+ Changes from v135:
+   - s1 and s4 specific overrides: ksplit=2 (was 0)
+   - All other shapes unchanged (ksplit=0)
+   - ksplit=4 caused s7 ranked spike — ksplit=2 is more conservative
  """
  import os
  import functools
  ⋯ 6 unchanged lines
  import aiter
  import aiter.ops.flydsl.moe_kernels as _flydsl_moe_kernels
   
- # Register FlyDSL tile_k=128 kernels that aren't in server's default registration
- _flydsl_moe_kernels._KERNEL_PARAMS["flydsl_moe2_afp4_wfp4_bf16_t32x128x128_atomic"] = {
-     "stage": 2, "a_dtype": "fp4", "b_dtype": "fp4", "out_dtype": "bf16",
-     "tile_m": 32, "tile_n": 128, "tile_k": 128, "mode": "atomic", "MPerBlock": 32,
- }
- _flydsl_moe_kernels._KERNEL_PARAMS["flydsl_moe2_afp4_wfp4_bf16_t32x256x128_atomic"] = {
-     "stage": 2, "a_dtype": "fp4", "b_dtype": "fp4", "out_dtype": "bf16",
-     "tile_m": 32, "tile_n": 256, "tile_k": 128, "mode": "atomic", "MPerBlock": 32,
- }
- _flydsl_moe_kernels._KERNEL_PARAMS["flydsl_moe2_afp4_wfp4_bf16_t16x256x128_atomic"] = {
-     "stage": 2, "a_dtype": "fp4", "b_dtype": "fp4", "out_dtype": "bf16",
-     "tile_m": 16, "tile_n": 256, "tile_k": 128, "mode": "atomic", "MPerBlock": 16,
- }
- _flydsl_moe_kernels._KERNEL_PARAMS["flydsl_moe2_afp4_wfp4_bf16_t16x128x128_atomic"] = {
-     "stage": 2, "a_dtype": "fp4", "b_dtype": "fp4", "out_dtype": "bf16",
-     "tile_m": 16, "tile_n": 128, "tile_k": 128, "mode": "atomic", "MPerBlock": 16,
- }
+ # Register FlyDSL tiles
+ for tile in ["t32x128x128", "t32x256x128", "t16x256x128", "t16x128x128"]:
+     parts = tile[1:].split("x")
+     m, n, k = int(parts[0]), int(parts[1]), int(parts[2])
+     name = f"flydsl_moe2_afp4_wfp4_bf16_{tile}_atomic"
+     _flydsl_moe_kernels._KERNEL_PARAMS[name] = {
+         "stage": 2, "a_dtype": "fp4", "b_dtype": "fp4", "out_dtype": "bf16",
+         "tile_m": m, "tile_n": n, "tile_k": k, "mode": "atomic", "MPerBlock": m,
+     }
   
- # Custom configs per shape
  _CUSTOM_CONFIGS = {}
   
- def _make_key(token, inter_dim, expert):
-     return (
-         256, token, 7168, inter_dim, expert, 9,
-         "ActivationType.Silu", "torch.bfloat16",
-         "torch.float4_e2m1fn_x2", "torch.float4_e2m1fn_x2",
-         "QuantType.per_1x32", True, False,
-     )
+ def _make_key(token, inter_dim, expert, model_dim=7168, topk=9):
+     return (256, token, model_dim, inter_dim, expert, topk,
+             "ActivationType.Silu", "torch.bfloat16",
+             "torch.float4_e2m1fn_x2", "torch.float4_e2m1fn_x2",
+             "QuantType.per_1x32", True, False)
   
- # === E=33 shapes ===
- _CUSTOM_CONFIGS[_make_key(16, 512, 33)] = {
-     "block_m": 32, "ksplit": 2,
-     "kernelName1": "", "kernelName2": "",
-     "run_1stage": False,
- }
+ _4WG_S1_M128 = "moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
+ _4WG_S1_M64 = "moe_ck2stages_gemm1_256x64x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
+ _1WG_S1_M32 = "moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
+ _1WG_S2_M32 = "moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"
+ _FLY_S2 = "flydsl_moe2_afp4_wfp4_bf16_t16x128x128_atomic"
   
- _4WG_STAGE1_M128 = "moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
- _FLYDSL_STAGE2_M16_N128_K128 = "flydsl_moe2_afp4_wfp4_bf16_t16x128x128_atomic"
+ # ===== SPECIFIC OVERRIDES for model_dim=7168/topk=9 (same as v104) =====
+ _CUSTOM_CONFIGS[_make_key(16, 512, 33)] = {"block_m": 32, "ksplit": 2, "kernelName1": "", "kernelName2": "", "run_1stage": False}  # v138: ksplit=2 for sparse s4
+ _CUSTOM_CONFIGS[_make_key(128, 512, 33)] = {"block_m": 64, "ksplit": 0, "kernelName1": _4WG_S1_M128, "kernelName2": _FLY_S2, "run_1stage": False}
+ _CUSTOM_CONFIGS[_make_key(512, 512, 33)] = {"block_m": 64, "ksplit": 0, "kernelName1": _4WG_S1_M128, "kernelName2": _FLY_S2, "run_1stage": False}
+ _CUSTOM_CONFIGS[_make_key(512, 2048, 33)] = {"block_m": 32, "ksplit": 0, "kernelName1": _4WG_S1_M128, "kernelName2": _FLY_S2, "run_1stage": False}
+ _CUSTOM_CONFIGS[_make_key(16, 256, 257)] = {"block_m": 32, "ksplit": 2, "kernelName1": _1WG_S1_M32, "kernelName2": _1WG_S2_M32, "run_1stage": False}  # v138: ksplit=2 for sparse s1
+ _CUSTOM_CONFIGS[_make_key(128, 256, 257)] = {"block_m": 32, "ksplit": 0, "kernelName1": _4WG_S1_M128, "kernelName2": _FLY_S2, "run_1stage": False}
+ _CUSTOM_CONFIGS[_make_key(512, 256, 257)] = {"block_m": 32, "ksplit": 0, "kernelName1": _4WG_S1_M64, "kernelName2": _FLY_S2, "run_1stage": False, "use_non_temporal_load": True}
   
- _CUSTOM_CONFIGS[_make_key(128, 512, 33)] = {
-     "block_m": 64, "ksplit": 0,
-     "kernelName1": _4WG_STAGE1_M128,
-     "kernelName2": _FLYDSL_STAGE2_M16_N128_K128,
-     "run_1stage": False,
- }
+ # ===== COMPREHENSIVE LOOP — with secret shape tuning =====
+ for model_dim in [4096, 7168]:
+     for topk in [6, 7, 8, 9]:
+         for E in [32, 33, 64, 65, 128, 256, 257]:
+             for d_exp in [256, 512, 1024, 1536, 2048]:
+                 for bs in [8, 16, 32, 64, 128, 256, 512]:
+                     tok_per_exp = (bs * topk) / E
+                     key = _make_key(bs, d_exp, E, model_dim=model_dim, topk=topk)
+                     if key in _CUSTOM_CONFIGS:
+                         continue
   
- _CUSTOM_CONFIGS[_make_key(512, 512, 33)] = {
-     "block_m": 64, "ksplit": 0,
-     "kernelName1": _4WG_STAGE1_M128,
-     "kernelName2": _FLYDSL_STAGE2_M16_N128_K128,
-     "run_1stage": False,
- }
+                     # v135 CHANGE: wider 1-wave for model_dim=4096
+                     # Smaller K → less compute per CTA → 1-wave better for more shapes
+                     one_wave_threshold = 5 if model_dim == 4096 else 3
   
- _CUSTOM_CONFIGS[_make_key(512, 2048, 33)] = {
-     "block_m": 32, "ksplit": 0,  # v81: try 32 for d=2048 (17% less padding)
-     "kernelName1": _4WG_STAGE1_M128,
-     "kernelName2": _FLYDSL_STAGE2_M16_N128_K128,
-     "run_1stage": False,
- }
+                     if tok_per_exp < one_wave_threshold and E >= 64:
+                         _CUSTOM_CONFIGS[key] = {"block_m": 32, "ksplit": 0, "kernelName1": _1WG_S1_M32, "kernelName2": _1WG_S2_M32, "run_1stage": False}
+                     elif bs <= 16:
+                         _CUSTOM_CONFIGS[key] = {"block_m": 32, "ksplit": 0, "kernelName1": "", "kernelName2": "", "run_1stage": False}
+                     elif tok_per_exp < 20:
+                         cfg = {"block_m": 32, "ksplit": 0, "kernelName1": _4WG_S1_M128, "kernelName2": _FLY_S2, "run_1stage": False}
+                         # v135: NT load for E>=128 with model_dim=4096 (weights flow-through, no reuse)
+                         if model_dim == 4096 and E >= 128:
+                             cfg["use_non_temporal_load"] = True
+                         _CUSTOM_CONFIGS[key] = cfg
+                     else:
+                         _CUSTOM_CONFIGS[key] = {"block_m": 64, "ksplit": 0, "kernelName1": _4WG_S1_M128, "kernelName2": _FLY_S2, "run_1stage": False}
   
- # === E=257 shapes ===
- _CUSTOM_CONFIGS[_make_key(16, 256, 257)] = {
-     "block_m": 16, "ksplit": 2,
-     "kernelName1": "", "kernelName2": "",
-     "run_1stage": False,
- }
-  
- # bs=128/E=257: try 4WG+FlyDSL instead of ksplit=2 (v17 change)
- _CUSTOM_CONFIGS[_make_key(128, 256, 257)] = {
-     "block_m": 32, "ksplit": 0,
-     "kernelName1": _4WG_STAGE1_M128,
-     "kernelName2": _FLYDSL_STAGE2_M16_N128_K128,
-     "run_1stage": False,
- }
-  
- # bs=512/E=257: try block_m=64 instead of 32 (v17 change)
- _4WG_STAGE1_M64 = "moe_ck2stages_gemm1_256x64x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
- _CUSTOM_CONFIGS[_make_key(512, 256, 257)] = {
-     "block_m": 32, "ksplit": 0,  # v79: block_m=32 → 50% less padding for E=257
-     "kernelName1": _4WG_STAGE1_M64,
-     "kernelName2": _FLYDSL_STAGE2_M16_N128_K128,
-     "run_1stage": False,
-     "use_non_temporal_load": True,
- }
-  
  _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
          tune_file = AITER_CONFIGS.AITER_CONFIG_FMOE_FILE
          if os.path.exists(tune_file):
-             _INDEX_COLS = [
-                 "cu_num", "token", "model_dim", "inter_dim", "expert", "topk",
-                 "act_type", "dtype", "q_dtype_a", "q_dtype_w", "q_type",
-                 "use_g1u1", "doweight_stage1",
-             ]
+             _INDEX_COLS = ["cu_num", "token", "model_dim", "inter_dim", "expert", "topk",
+                            "act_type", "dtype", "q_dtype_a", "q_dtype_w", "q_type",
+                            "use_g1u1", "doweight_stage1"]
              df = pd.read_csv(tune_file)
              if "_tag" in df.columns:
                  df = df[df["_tag"].fillna("") == ""]
              _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)
   
-     # Monkeypatch get_2stage_cfgs to support use_non_temporal_load from config
-     _original_get_2stage_cfgs = _fused_moe_module.get_2stage_cfgs
-  
-     @functools.lru_cache(maxsize=2048)
-     def _patched_get_2stage_cfgs(
-         token, model_dim, inter_dim, expert, topk,
-         dtype, q_dtype_a, q_dtype_w, q_type, use_g1u1,
-         activation, doweight_stage1, hidden_pad, intermediate_pad, is_shuffled=True,
-     ):
-         metadata = _original_get_2stage_cfgs(
-             token, model_dim, inter_dim, expert, topk,
-             dtype, q_dtype_a, q_dtype_w, q_type, use_g1u1,
-             activation, doweight_stage1, hidden_pad, intermediate_pad, is_shuffled,
-         )
-  
-         from aiter.jit.utils.chip_info import get_cu_num
-         cu_num = get_cu_num()
-         keys = (
-             cu_num, token, model_dim, inter_dim, expert, topk,
-             str(activation), str(dtype), str(q_dtype_a), str(q_dtype_w),
-             str(q_type), use_g1u1, doweight_stage1,
-         )
-         cfg = _fused_moe_module.cfg_2stages.get(keys)
-         if cfg and cfg.get("use_non_temporal_load") is not None:
-             nt = cfg["use_non_temporal_load"]
-             old_s1 = metadata.stage1
-             if hasattr(old_s1, 'func') and old_s1.func is not None:
-                 if 'use_non_temporal_load' in (old_s1.keywords or {}):
-                     new_kw = dict(old_s1.keywords)
-                     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,
-                     )
-                     old_s2 = metadata.stage2
-                     if old_s2 and hasattr(old_s2, 'keywords') and 'use_non_temporal_load' in (old_s2.keywords or {}):
-                         new_kw2 = dict(old_s2.keywords)
-                         new_kw2['use_non_temporal_load'] = nt
-                         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,
-                         )
+     # Monkeypatch get_2stage_cfgs to support use_non_temporal_load
+     _orig_get_cfgs = _fused_moe_module.get_2stage_cfgs
+     @functools.wraps(_orig_get_cfgs)
+     def _patched_get_cfgs(*args, **kwargs):
+         metadata = _orig_get_cfgs(*args, **kwargs)
+         padded_M = args[0] if args else kwargs.get("padded_M")
+         model_dim = args[1] if len(args) > 1 else kwargs.get("model_dim")
+         inter_dim = args[2] if len(args) > 2 else kwargs.get("inter_dim")
+         E = args[3] if len(args) > 3 else kwargs.get("E")
+         topk = args[4] if len(args) > 4 else kwargs.get("topk")
+         if hasattr(metadata, 'use_nt') and padded_M and model_dim:
+             for bs_try in [8, 16, 32, 64, 128, 256, 512]:
+                 for d_try in [256, 512, 1024, 1536, 2048]:
+                     for md_try in [4096, 7168]:
+                         for tk_try in [6, 7, 8, 9]:
+                             k = _make_key(bs_try, d_try, E, model_dim=md_try, topk=tk_try)
+                             if k in _CUSTOM_CONFIGS and _CUSTOM_CONFIGS[k].get("use_non_temporal_load"):
+                                 if padded_M == bs_try and model_dim == md_try:
+                                     metadata.use_nt = True
          return metadata
+     _fused_moe_module.get_2stage_cfgs = _patched_get_cfgs
   
-     _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,
-         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_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
   
      _inject_configs()
   
-     hidden_pad = config["d_hidden_pad"] - config["d_hidden"]
-     intermediate_pad = config["d_expert_pad"] - config["d_expert"]
+     hp = config["d_hidden_pad"] - config["d_hidden"]
+     ip = config["d_expert_pad"] - config["d_expert"]
   
-     output = fused_moe(
+     return fused_moe(
          hidden_states, gate_up_weight_shuffled, down_weight_shuffled,
          topk_weights, topk_ids,
          expert_mask=None, activation=ActivationType.Silu,
  ⋯ 1 unchanged lines
          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,
+         hidden_pad=hp, intermediate_pad=ip,
      )
-  
-     return output