submission 530820

  import os
- import functools
  import sys
+ import functools
  from pathlib import Path
   
  import pandas as pd
  import torch
  from aiter import ActivationType, QuantType
- import aiter.fused_moe as fused_moe_mod
- from aiter.fused_moe import fused_moe
  from task import input_t, output_t
   
   
  DEBUG_FMOE = os.getenv("JOSU_DEBUG_FMOE") == "1"
  FORCE_REBUILD_FMOE = DEBUG_FMOE or os.getenv("JOSU_REBUILD_FMOE") == "1"
- CFG_VERSION = 'auto_v9_force_nt_large_e33'
- FORCE_NT_MODE = 'large_e33'
- TARGET_SHAPES = (
-     (16, 7168, 256, 257, 9),
-     (128, 7168, 256, 257, 9),
-     (512, 7168, 256, 257, 9),
-     (16, 7168, 512, 33, 9),
-     (128, 7168, 512, 33, 9),
-     (512, 7168, 512, 33, 9),
-     (512, 7168, 2048, 33, 9),
- )
+ CFG_VERSION = "v85_blockm16_small"
   
   
- def _shape_mask(frame: pd.DataFrame, shape: tuple[int, int, int, int, int]) -> pd.Series:
-     token, model_dim, inter_dim, expert, topk = shape
-     return (
-         (frame["token"] == token)
-         & (frame["model_dim"] == model_dim)
-         & (frame["inter_dim"] == inter_dim)
-         & (frame["expert"] == expert)
-         & (frame["topk"] == topk)
-     )
-  
-  
- def _debug_rows(label: str, frame: pd.DataFrame) -> None:
-     if not DEBUG_FMOE:
-         return
-  
-     cols = [
-         "token",
-         "model_dim",
-         "inter_dim",
-         "expert",
-         "topk",
-         "block_m",
-         "ksplit",
-         "run_1stage",
-         "us",
-         "kernelName1",
-         "kernelName2",
-         "_tag",
-     ]
-     keep_cols = [col for col in cols if col in frame.columns]
-     print(f"[best.py] {label}", file=sys.stderr)
-     for shape in TARGET_SHAPES:
-         rows = frame.loc[_shape_mask(frame, shape), keep_cols].sort_values("us").head(6)
-         print(f"[best.py] shape={shape} rows={len(rows)}", file=sys.stderr)
-         if not rows.empty:
-             print(rows.to_string(index=False), file=sys.stderr)
-  
-  
  def _bootstrap_cfg() -> str:
      out_path = Path(f"/tmp/josu_cfg_hybrid_under150_fmoe_{CFG_VERSION}.csv")
      if out_path.exists() and out_path.stat().st_size > 0 and not FORCE_REBUILD_FMOE:
  ⋯ 42 unchanged lines
          }
      )
   
-     online_tuned_rows = [{'token': 16, 'model_dim': 7168, 'inter_dim': 256, 'expert': 257, 'topk': 9, 'block_m': 32, 'ksplit': 2, 'us': 0.0, 'us1': 0.0, 'kernelName1': '', 'err1': '0.0%', 'us2': 0.0, 'kernelName2': '', 'err2': '0.0%', 'run_1stage': 0, 'tflops': 0.0, 'bw': 0.0}, {'token': 128, 'model_dim': 7168, 'inter_dim': 256, 'expert': 257, 'topk': 9, 'block_m': 32, 'ksplit': 2, 'us': 0.0, 'us1': 0.0, 'kernelName1': '', 'err1': '0.0%', 'us2': 0.0, 'kernelName2': '', 'err2': '0.0%', 'run_1stage': 0, 'tflops': 0.0, 'bw': 0.0}, {'token': 512, 'model_dim': 7168, 'inter_dim': 512, 'expert': 33, 'topk': 9, 'block_m': 32, 'ksplit': 0, 'us': 0.0, 'us1': 0.0, 'kernelName1': 'moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16', 'err1': '0.0%', 'us2': 0.0, 'kernelName2': 'moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16', 'err2': '2.9%', 'run_1stage': 0, 'tflops': 0.0, 'bw': 0.0}, {'token': 512, 'model_dim': 7168, 'inter_dim': 2048, 'expert': 33, 'topk': 9, 'block_m': 128, 'ksplit': 0, 'us': 0.0, 'us1': 0.0, 'kernelName1': 'moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16', 'err1': '0.0%', 'us2': 0.0, 'kernelName2': 'moe_ck2stages_gemm2_256x128x128x128_1x4_MulABScaleExpertWeightShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16', 'err2': '7.1%', 'run_1stage': 0, 'tflops': 0.0, 'bw': 0.0}]
-     k2_overrides = []
-     default_2stage_rows = [{'token': 16, 'model_dim': 7168, 'inter_dim': 512, 'expert': 33, 'topk': 9, 'block_m': 32, 'ksplit': 2}, {'token': 128, 'model_dim': 7168, 'inter_dim': 512, 'expert': 33, 'topk': 9, 'block_m': 32, 'ksplit': 2}]
+     K1_SMALL = "moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
+     K1_MED = "moe_ck2stages_gemm1_256x32x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
+     K1_LARGE = "moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
+     K2_SMALL = "moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"
+     K2_LARGE = "moe_ck2stages_gemm2_256x128x128x128_1x4_MulABScaleExpertWeightShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"
   
-     rows = []
-     for row in online_tuned_rows:
-         seeded = defaults.copy()
-         seeded.update(row)
-         seeded["us"] = -1.0
-         rows.append(seeded)
+     # Yufeng98 (#3) uses blockm16. Try block_m=16 for shapes with very few tokens per expert.
+     # E=257 bs=16: ~0.6 tokens/expert → padding to 16 instead of 32 saves half the wasted compute
+     # E=257 bs=128: ~4.5 tokens/expert → padding to 16 is better fit
+     # E=33 bs=16: ~4.4 tokens/expert → padding to 16 instead of 32
+     online_tuned_rows = [
+         # E=257 bs=16: block_m=16 + ksplit=7 (was block_m=32)
+         {"token": 16, "model_dim": 7168, "inter_dim": 256, "expert": 257, "topk": 9,
+          "block_m": 16, "ksplit": 7,
+          "kernelName1": K1_SMALL, "kernelName2": K2_SMALL, "run_1stage": 0},
+         # E=257 bs=128: block_m=16 + ksplit=4 (was block_m=32)
+         {"token": 128, "model_dim": 7168, "inter_dim": 256, "expert": 257, "topk": 9,
+          "block_m": 16, "ksplit": 4,
+          "kernelName1": K1_SMALL, "kernelName2": K2_SMALL, "run_1stage": 0},
+         # E=257 bs=512: keep block_m=32 ksplit=0 (enough tokens per expert)
+         {"token": 512, "model_dim": 7168, "inter_dim": 256, "expert": 257, "topk": 9,
+          "block_m": 32, "ksplit": 0,
+          "kernelName1": K1_SMALL, "kernelName2": K2_SMALL, "run_1stage": 0},
   
-     for token, model_dim, inter_dim, expert, topk in k2_overrides:
-         mask = (
-             (merged["cu_num"] == 256)
-             & (merged["token"] == token)
-             & (merged["model_dim"] == model_dim)
-             & (merged["inter_dim"] == inter_dim)
-             & (merged["expert"] == expert)
-             & (merged["topk"] == topk)
-             & (merged["act_type"] == "ActivationType.Silu")
-             & (merged["dtype"] == "torch.bfloat16")
-             & (merged["q_dtype_a"] == "torch.float4_e2m1fn_x2")
-             & (merged["q_dtype_w"] == "torch.float4_e2m1fn_x2")
-             & (merged["q_type"] == "QuantType.per_1x32")
-             & (merged["use_g1u1"] == 1)
-             & (merged["doweight_stage1"] == 0)
-         )
-         if not mask.any():
-             continue
+         # E=33 bs=16: block_m=16 + ksplit=4 (was block_m=32)
+         {"token": 16, "model_dim": 7168, "inter_dim": 512, "expert": 33, "topk": 9,
+          "block_m": 16, "ksplit": 4,
+          "kernelName1": K1_SMALL, "kernelName2": K2_SMALL, "run_1stage": 0},
+         # E=33 bs=128: keep block_m=32 ksplit=4 (34 tokens/expert is reasonable for block_m=32)
+         {"token": 128, "model_dim": 7168, "inter_dim": 512, "expert": 33, "topk": 9,
+          "block_m": 32, "ksplit": 4,
+          "kernelName1": K1_MED, "kernelName2": K2_SMALL, "run_1stage": 0},
+         # E=33 bs=512 d=512: keep block_m=32 ksplit=0
+         {"token": 512, "model_dim": 7168, "inter_dim": 512, "expert": 33, "topk": 9,
+          "block_m": 32, "ksplit": 0,
+          "kernelName1": K1_MED, "kernelName2": K2_SMALL, "run_1stage": 0},
   
-         row = merged.loc[mask].sort_values("us").iloc[0].to_dict()
-         kernel_name_1 = row.get("kernelName1")
-         kernel_name_2 = row.get("kernelName2")
-         if pd.isna(kernel_name_1) or pd.isna(kernel_name_2) or not kernel_name_1 or not kernel_name_2:
-             continue
+         # E=33 d=2048: keep block_m=128 ksplit=0
+         {"token": 512, "model_dim": 7168, "inter_dim": 2048, "expert": 33, "topk": 9,
+          "block_m": 128, "ksplit": 0,
+          "kernelName1": K1_LARGE, "kernelName2": K2_LARGE, "run_1stage": 0},
+     ]
   
-         row["ksplit"] = 2
-         row["us"] = -1.0
-         rows.append(row)
-  
-     for row in default_2stage_rows:
+     rows = []
+     for row in online_tuned_rows:
          seeded = defaults.copy()
          seeded.update(row)
          seeded["us"] = -1.0
          rows.append(seeded)
   
      merged = pd.concat([merged, pd.DataFrame(rows)], ignore_index=True)
-     _debug_rows("pre-dedup", merged)
      merged = (
          merged.sort_values("us")
          .drop_duplicates(subset=keys, keep="first")
          .reset_index(drop=True)
      )
-     _debug_rows("post-dedup", merged)
   
      out_path.parent.mkdir(parents=True, exist_ok=True)
      merged.to_csv(out_path, index=False)
  ⋯ 2 unchanged lines
   
  os.environ["AITER_CONFIG_FMOE"] = _bootstrap_cfg()
   
+ from aiter.fused_moe import fused_moe
+ import aiter.fused_moe as _fm
   
- def _should_force_nt(token: int, inter_dim: int, expert: int) -> bool:
-     if FORCE_NT_MODE == "off":
-         return False
-     if FORCE_NT_MODE == "all":
-         return True
-     if FORCE_NT_MODE == "e33":
-         return expert == 33
-     if FORCE_NT_MODE == "large":
-         return token >= 512
-     if FORCE_NT_MODE == "large_e33":
-         return expert == 33 and token >= 512
-     if FORCE_NT_MODE == "257_only":
-         return expert == 257
-     raise ValueError("Unknown FORCE_NT_MODE=" + str(FORCE_NT_MODE))
+ _orig_ck_moe_stage1 = _fm.ck_moe_stage1
   
+ @functools.wraps(_orig_ck_moe_stage1)
+ def _nt_ck_moe_stage1(*args, **kwargs):
+     kwargs["use_non_temporal_load"] = True
+     return _orig_ck_moe_stage1(*args, **kwargs)
   
- def _install_force_nt_patch() -> None:
-     if FORCE_NT_MODE == "off":
-         return
-     if getattr(fused_moe_mod, "_josu_force_nt_patch", False):
-         return
+ _fm.ck_moe_stage1 = _nt_ck_moe_stage1
   
-     os.environ["AITER_USE_NT"] = "1"
-     original_get_2stage_cfgs = fused_moe_mod.get_2stage_cfgs
+ if hasattr(_fm, "ck_moe_stage2"):
+     _orig_ck_moe_stage2 = _fm.ck_moe_stage2
   
-     def wrapped_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,
-     ):
-         meta = 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,
-         )
-         if not _should_force_nt(token, inter_dim, expert):
-             return meta
+     @functools.wraps(_orig_ck_moe_stage2)
+     def _nt_ck_moe_stage2(*args, **kwargs):
+         kwargs["use_non_temporal_load"] = True
+         return _orig_ck_moe_stage2(*args, **kwargs)
   
-         stage1 = meta.stage1
-         stage2 = meta.stage2
-         if isinstance(stage1, functools.partial) and stage1.func is fused_moe_mod.ck_moe_stage1:
-             stage1_kwargs = dict(stage1.keywords or dict())
-             stage1_kwargs["use_non_temporal_load"] = True
-             stage1 = functools.partial(stage1.func, *(stage1.args or ()), **stage1_kwargs)
-         if (
-             isinstance(stage2, functools.partial)
-             and stage2.func is fused_moe_mod.aiter.ck_moe_stage2_fwd
-         ):
-             stage2_kwargs = dict(stage2.keywords or dict())
-             stage2_kwargs["use_non_temporal_load"] = True
-             stage2 = functools.partial(stage2.func, *(stage2.args or ()), **stage2_kwargs)
-         return fused_moe_mod.MOEMetadata(
-             stage1,
-             stage2,
-             int(meta.block_m),
-             int(meta.ksplit),
-             meta.run_1stage,
-             meta.has_bias,
-             True,
-         )
+     _fm.ck_moe_stage2 = _nt_ck_moe_stage2
   
-     fused_moe_mod.get_2stage_cfgs = wrapped_get_2stage_cfgs
-     fused_moe_mod._josu_force_nt_patch = True
+ print("[v85] block_m=16 for small shapes (E=257 bs=16/128, E=33 bs=16) + NT + v53 base", file=sys.stderr)
   
   
- _install_force_nt_patch()
-  
-  
  def custom_kernel(data: input_t) -> output_t:
      (
          hidden_states,