submission 673796

- # HIP MoE v64 -- Direct dispatch + NT load fix
+ # HIP MoE v72 -- v65 base + _4WG128 for E=257 bs=512
  #
- # v64 bypasses fused_moe's Python dispatch chain (fused_moe → fused_moe_ →
- # fused_moe_2stages, 3+ nested functions, enum conversions, conditional logic)
- # and calls the GPU kernels directly. Saves ~8-12µs of Python overhead per call.
+ # v71 proved CSV configs are MUCH WORSE for E=257 shapes:
+ #   bs=16: 90.6us (ours) vs 140us (CSV) — CSV 54% slower
+ #   bs=128: 175us vs 219us — CSV 25% slower
+ #   bs=512: 196us vs 252us — CSV 29% slower
+ # CSV uses 1WG CK kernels (64x32) which are slower than our CKtile/4WG+FlyDSL.
+ # NEVER load CSV for E=257.
  #
- # Also fixes TWO bugs in NT load patching that existed since v19:
- #   1. Wrong key construction: a[:13] doesn't include cu_num, so _CC.get(k)
- #      never matched. Fixed: reconstruct keys matching cfg_2stages format.
- #   2. Wrong keyword name: 'non_temporal_load' vs 'use_non_temporal_load'.
- #      Fixed: use correct keyword name.
+ # v72 experiment: E=257 bs=512 stage1 _4WG64 → _4WG128
+ #   - _4WG64: MPerBlock=64, processes 2 sorting blocks (64/32) per WG
+ #   - _4WG128: MPerBlock=128, processes 4 sorting blocks (128/32) per WG
+ #   - Both fit in 1 CU round for 257 experts, but _4WG128 does more work/WG
+ #   - E=33 shapes already use _4WG128 successfully
  #
- # Added NT loads for bs=128 E=33 (heuristic: tokens_per_expert=35 < 64).
- #
- # Flow:
- #   1st call per shape → fused_moe (populates metadata cache, verified correct)
- #   2nd+ calls → direct dispatch (sorting → quant → stage1 → re-quant → stage2)
- #   Any error → falls back to fused_moe permanently for that shape
+ # All other shapes IDENTICAL to v65.
   
  from task import input_t, output_t
  import torch
  ⋯ 33 unchanged lines
  _F16 = "flydsl_moe2_afp4_wfp4_bf16_t16x128x128_atomic"
   
  _CC = {}
- # E=33, TP=4 shapes
+ # E=33, TP=4 shapes (identical to v65)
  _CC[_key(16, 512, 33)] = {
      "block_m": 16, "ksplit": 2,
      "kernelName1": "", "kernelName2": "",
  ⋯ 2 unchanged lines
  _CC[_key(128, 512, 33)] = {
      "block_m": 32, "ksplit": 0,
      "kernelName1": _4WG128, "kernelName2": _F16,
-     "run_1stage": False, "use_non_temporal_load": True,  # tokens/expert=35<64
+     "run_1stage": False, "use_non_temporal_load": True,
  }
  _CC[_key(512, 512, 33)] = {
      "block_m": 64, "ksplit": 0,
  ⋯ 16 unchanged lines
      "kernelName1": "", "kernelName2": "",
      "run_1stage": False, "use_non_temporal_load": True,
  }
+ # CHANGED: _4WG64 → _4WG128 for E=257 bs=512
  _CC[_key(512, 256, 257)] = {
      "block_m": 32, "ksplit": 0,
-     "kernelName1": _4WG64, "kernelName2": _F16,
-     "run_1stage": False, "use_non_temporal_load": True,  # tokens/expert=18<64
+     "kernelName1": _4WG128, "kernelName2": _F16,
+     "run_1stage": False, "use_non_temporal_load": True,
  }
   
  # =========================================================================== #
- # Injection + metadata capture with FIXED NT load patching
+ # Injection + metadata capture with NT load patching
  # =========================================================================== #
  _injected = False
- _captured_meta = {}  # (padded_M, model_dim, inter_dim, E, topk) → MOEMetadata
+ _captured_meta = {}
  _cu = None
   
   
  ⋯ 15 unchanged lines
          global _cu
          m = orig(*a)
   
-         # FIX #1: Reconstruct keys matching cfg_2stages format
-         # get_2stage_cfgs builds 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)
-         # But function args a = (token[0], model_dim[1], inter_dim[2],
-         #   expert[3], topk[4], dtype[5], q_dtype_a[6], q_dtype_w[7],
-         #   q_type[8], use_g1u1[9], activation[10], doweight_stage1[11], ...)
          if _cu is None:
              try:
                  from aiter.jit.utils.chip_info import get_cu_num
  ⋯ 9 unchanged lines
          if c and c.get("use_non_temporal_load"):
              s = m.stage1
              kw = getattr(s, 'keywords', None) or {}
-             # FIX #2: correct keyword is 'use_non_temporal_load', not 'non_temporal_load'
              if hasattr(s, 'func') and 'use_non_temporal_load' in kw:
                  nk = dict(kw)
                  nk['use_non_temporal_load'] = True
  ⋯ 2 unchanged lines
                      m.stage2, m.block_m, m.ksplit,
                      m.run_1stage, m.has_bias, True)
   
-         # Cache metadata for direct dispatch
          _captured_meta[(a[0], a[1], a[2], a[3], a[4])] = m
          return m
   
  ⋯ 3 unchanged lines
  # =========================================================================== #
  # Per-shape cache for direct dispatch
  # =========================================================================== #
- _shape_cache = {}  # shape_key → cache dict or None (fallback)
+ _shape_cache = {}
   
   
  def _init_shape(data, shape_key):
  ⋯ 9 unchanged lines
      ip = cfg['d_expert_pad'] - cfg['d_expert']
      device = hs.device
   
-     # Run fused_moe once to populate metadata cache and return correct output
      result = fused_moe(
          hs, guw_sh, dw_sh, tw, ti,
          expert_mask=None, activation=ActivationType.Silu,
  ⋯ 3 unchanged lines
          hidden_pad=hp, intermediate_pad=ip,
      )
   
-     # Find captured metadata
      _, model_dim, inter_dim = _fm.get_inter_dim(guw_sh.shape, dw_sh.shape)
      padded_M = _fm.get_padded_M(M)
      meta_key = (padded_M, model_dim, inter_dim, E, topk)
   
      metadata = _captured_meta.get(meta_key)
      if metadata is None:
-         print(f"[V64] No metadata for {shape_key}, using fallback",
+         print(f"[V72] No metadata for {shape_key}, using fallback",
                file=sys.stderr, flush=True)
          _shape_cache[shape_key] = None
          return result
   
      block_m = int(metadata.block_m)
+     ksplit = int(metadata.ksplit)
      max_tok = M * topk + E * block_m - topk
      max_blk = (max_tok + block_m - 1) // block_m
   
      _shape_cache[shape_key] = {
          'meta': metadata,
          'block_m': block_m,
+         'ksplit': ksplit,
          'topk': topk,
          'E': E,
          'M': M,
  ⋯ 1 unchanged lines
          'inter_dim': inter_dim,
          'hp': hp,
          'ip': ip,
-         # Pre-allocated sorting buffers (reused across calls)
          'sid': torch.empty(max_tok, dtype=torch.int32, device=device),
          'sw': torch.empty(max_tok, dtype=torch.float32, device=device),
          'seid': torch.empty(max_blk, dtype=torch.int32, device=device),
          'nvi': torch.empty(2, dtype=torch.int32, device=device),
      }
   
-     # Log metadata details
      s1 = metadata.stage1
      s1k = getattr(s1, 'keywords', {}) if hasattr(s1, 'func') else {}
      nt_val = s1k.get('use_non_temporal_load', 'N/A')
      s1name = s1.func.__name__ if hasattr(s1, 'func') else str(s1)[:40]
-     print(f"[V64] Init {shape_key}: block_m={block_m} inter={inter_dim} "
-           f"model={model_dim} nt={nt_val} stage1={s1name}",
+     print(f"[V72] Init {shape_key}: block_m={block_m} ksplit={ksplit} "
+           f"inter={inter_dim} model={model_dim} nt={nt_val} stage1={s1name}",
            file=sys.stderr, flush=True)
   
      return result
   
   
  def _direct_dispatch(data, c):
-     """Direct dispatch: 5 GPU kernel calls with minimal Python overhead."""
+     """Direct dispatch: handles both CK and CKtile paths."""
      hs = data[0]; guw_sh = data[5]; dw_sh = data[6]
      guws_sh = data[7]; dws_sh = data[8]
      tw = data[9]; ti = data[10]
   
      meta = c['meta']
      block_m = c['block_m']
+     ksplit = c['ksplit']
      topk = c['topk']
      E = c['E']
      M = c['M']
  ⋯ 3 unchanged lines
      sid = c['sid']; sw = c['sw']; seid = c['seid']; nvi = c['nvi']
      device = hs.device
   
-     # 1. Sorting (pre-allocated output buffers, fresh moe_buf for atomicAdd)
+     # 1. Sorting
      moe_buf = torch.empty(M, model_dim, dtype=torch.bfloat16, device=device)
      moe_sorting_fwd(ti, tw, sid, sw, seid, nvi, moe_buf,
                      E, block_m, None, None, 0)
   
-     # 2. Activation quantization (fused quant + scale sorting, Triton kernel)
-     a1, a1_scale = fused_dynamic_mxfp4_quant_moe_sort(
-         hs, sorted_ids=sid, num_valid_ids=nvi,
-         token_num=M, topk=1, block_size=block_m)
+     if ksplit > 0:
+         # CKtile path: bf16 activations, no fp4 quantization
+         a1 = hs.to(torch.bfloat16)
+         a2_placeholder = torch.empty(M, topk, inter_dim,
+                                      dtype=torch.bfloat16, device=device)
+         a2 = meta.stage1(
+             a1, guw_sh, dw_sh, sid, seid, nvi, a2_placeholder, topk,
+             block_m=block_m, a1_scale=None,
+             w1_scale=guws_sh.view(dtypes.fp8_e8m0),
+             sorted_weights=None)
+         meta.stage2(
+             a2, guw_sh, dw_sh, sid, seid, nvi, moe_buf, topk,
+             w2_scale=dws_sh.view(dtypes.fp8_e8m0),
+             a2_scale=None, block_m=block_m, sorted_weights=sw)
+     else:
+         # CK path: fp4 quantized activations
+         a1, a1_scale = fused_dynamic_mxfp4_quant_moe_sort(
+             hs, sorted_ids=sid, num_valid_ids=nvi,
+             token_num=M, topk=1, block_size=block_m)
+         a2 = torch.empty(M, topk, inter_dim, dtype=torch.bfloat16, device=device)
+         a2 = meta.stage1(
+             a1, guw_sh, dw_sh, sid, seid, nvi, a2, topk,
+             block_m=block_m, a1_scale=a1_scale,
+             w1_scale=guws_sh.view(dtypes.fp8_e8m0),
+             sorted_weights=None)
+         a2_flat = a2.view(-1, inter_dim)
+         a2_q, a2_scale = fused_dynamic_mxfp4_quant_moe_sort(
+             a2_flat, sorted_ids=sid, num_valid_ids=nvi,
+             token_num=M, topk=topk, block_size=block_m)
+         a2 = a2_q.view(M, topk, -1)
+         meta.stage2(
+             a2, guw_sh, dw_sh, sid, seid, nvi, moe_buf, topk,
+             w2_scale=dws_sh.view(dtypes.fp8_e8m0),
+             a2_scale=a2_scale, block_m=block_m, sorted_weights=sw)
   
-     # 3. Stage 1: gate_up GEMM + SwiGLU (CK or CKtile kernel)
-     a2 = torch.empty(M, topk, inter_dim, dtype=torch.bfloat16, device=device)
-     a2 = meta.stage1(
-         a1, guw_sh, dw_sh, sid, seid, nvi, a2, topk,
-         block_m=block_m,
-         a1_scale=a1_scale,
-         w1_scale=guws_sh.view(dtypes.fp8_e8m0),
-         sorted_weights=None)
-  
-     # 4. Intermediate re-quantization (bf16 → fp4x2, Triton kernel)
-     a2_flat = a2.view(-1, inter_dim)
-     a2_q, a2_scale = fused_dynamic_mxfp4_quant_moe_sort(
-         a2_flat, sorted_ids=sid, num_valid_ids=nvi,
-         token_num=M, topk=topk, block_size=block_m)
-     a2_q = a2_q.view(M, topk, -1)
-  
-     # 5. Stage 2: down GEMM + weighted scatter-reduce (CK/FlyDSL kernel)
-     meta.stage2(
-         a2_q, guw_sh, dw_sh, sid, seid, nvi, moe_buf, topk,
-         w2_scale=dws_sh.view(dtypes.fp8_e8m0),
-         a2_scale=a2_scale,
-         block_m=block_m,
-         sorted_weights=sw)
-  
      return moe_buf
   
   
  ⋯ 24 unchanged lines
      shape_key = (cfg['bs'], cfg['d_expert'],
                   cfg['n_routed_experts'] + cfg['n_shared_experts'])
   
-     # First call per shape: use fused_moe to populate metadata cache
      if shape_key not in _shape_cache:
          return _init_shape(data, shape_key)
   
      c = _shape_cache[shape_key]
   
-     # Fallback if metadata capture failed
      if c is None:
          return _fallback(data)
   
-     # Direct dispatch (bypasses ~8-12µs Python overhead)
      try:
          return _direct_dispatch(data, c)
      except Exception as e:
-         print(f"[V64] Dispatch err {shape_key}: {str(e)[:200]}",
+         print(f"[V72] Dispatch err {shape_key}: {str(e)[:200]}",
                file=sys.stderr, flush=True)
          import traceback
          traceback.print_exc(file=sys.stderr)