submission 745058

- """
- V112: V98 + Enable non_temporal_load for E=257 bs=512.
- CSV path forces use_non_temporal_load=False, but heuristic says True
- for E=257 M=512 (token*topk//E = 17 < 64). Fix by patching metadata.
- """
  #!POPCORN leaderboard amd-moe-mxfp4
  #!POPCORN gpu MI355X
- import torch
+  
+ """
+ v168: Pre-allocate quantization output buffers (x_fp4, blockscale_e8m0_sorted)
+ for both stage1 and stage2 quant calls. Inline the fused_dynamic_mxfp4_quant_moe_sort
+ Triton kernel launch with cached output tensors to eliminate 4 torch.empty allocations
+ per forward pass on CK 2-stage shapes.
+ """
  import os
+ import functools
+ import torch
+ import triton
+ from task import input_t, output_t
+  
  import aiter
  from aiter import ActivationType, QuantType, dtypes
  from aiter.fused_moe import (
-     fused_moe, get_2stage_cfgs, get_inter_dim, get_padded_M,
-     fused_dynamic_mxfp4_quant_moe_sort,
+     get_2stage_cfgs, get_padded_M, get_inter_dim,
  )
- from task import input_t, output_t
+ import aiter.fused_moe as _fused_moe_module
+ import aiter.ops.flydsl.moe_kernels as _flydsl_moe_kernels
+ from aiter.ops.triton._triton_kernels.quant.fused_mxfp4_quant import (
+     _fused_dynamic_mxfp4_quant_moe_sort_kernel,
+ )
   
- # Stage1 kernels
- CK_S1_256WG_64 = "moe_ck2stages_gemm1_256x64x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
- CK_S1_64WG_32 = "moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
+ def _register_flydsl_kernel(name, tile_m, tile_n, tile_k=128):
+     _flydsl_moe_kernels._KERNEL_PARAMS[name] = {
+         "stage": 2,
+         "a_dtype": "fp4",
+         "b_dtype": "fp4",
+         "out_dtype": "bf16",
+         "tile_m": tile_m,
+         "tile_n": tile_n,
+         "tile_k": tile_k,
+         "mode": "atomic",
+         "MPerBlock": tile_m,
+     }
   
- # Stage2 kernels
- FLY_S2_32x256A = "flydsl_moe2_afp4_wfp4_bf16_t32x256x256_atomic"
- FLY_S2_64x256R = "flydsl_moe2_afp4_wfp4_bf16_t64x256x256_reduce"
- FLY_S2_64x128R = "flydsl_moe2_afp4_wfp4_bf16_t64x128x256_reduce"
- CK_S2_64WG_32 = "moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"
   
- _CSV_HEADER = "cu_num,token,model_dim,inter_dim,expert,topk,act_type,dtype,q_dtype_a,q_dtype_w,q_type,use_g1u1,doweight_stage1,block_m,ksplit,us1,kernelName1,err1,us2,kernelName2,err2,us,run_1stage,tflops,bw"
+ # 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),
+     ("flydsl_moe2_afp4_wfp4_bf16_t16x256x128_atomic", 16, 256),
+     ("flydsl_moe2_afp4_wfp4_bf16_t16x128x128_atomic", 16, 128),
+ ):
+     _register_flydsl_kernel(_name, _tm, _tn)
   
- def _row(cu, tok, mdim, idim, E, topk, bm, k1, k2, ks=0):
-     return f"{cu},{tok},{mdim},{idim},{E},{topk},ActivationType.Silu,torch.bfloat16,torch.float4_e2m1fn_x2,torch.float4_e2m1fn_x2,QuantType.per_1x32,True,False,{bm},{ks},0,{k1},0,0,{k2},0,0,False,0,0"
+ # Shape configs
+ _CUSTOM_CONFIGS = {}
   
- def _build_custom_csv():
-     rows = [_CSV_HEADER]
-     # E=33 shapes (same as V92/V97)
-     rows.append(_row(256, 512, 7168, 512, 33, 9, 64, CK_S1_256WG_64, FLY_S2_64x256R))
-     rows.append(_row(256, 512, 7168, 2048, 33, 9, 64, CK_S1_256WG_64, FLY_S2_64x128R))
-     # E=257 bs=512: Try FlyDSL atomic stage2 with block_m=32
-     rows.append(_row(256, 512, 7168, 256, 257, 9, 32, CK_S1_64WG_32, FLY_S2_32x256A))
-     return "\n".join(rows) + "\n"
   
- _CKTILE_SHAPES = {(16, 257), (128, 257), (16, 33), (128, 33)}
- _CKTILE_KSPLIT = {
-     (16, 257): 7,
-     (128, 257): 4,
-     (16, 33): 2,
-     (128, 33): 2,
- }
+ def _add_cfg(token, inter_dim, expert, block_m, ksplit, kernelName1="", kernelName2="", use_non_temporal_load=None):
+     cfg = {
+         "block_m": block_m,
+         "ksplit": ksplit,
+         "kernelName1": kernelName1,
+         "kernelName2": kernelName2,
+         "run_1stage": False,
+     }
+     if use_non_temporal_load is not None:
+         cfg["use_non_temporal_load"] = use_non_temporal_load
+     _CUSTOM_CONFIGS[_make_key(token, inter_dim, expert)] = cfg
   
- _cache = {}
- _initialized = False
- _sort_fn = None
- _DISPATCH_POLICY = 2
+ 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,
+     )
   
+ _4WG_STAGE1_M128 = "moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
+ _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"
   
- def custom_kernel(data: input_t) -> output_t:
-     global _initialized, _sort_fn
+ # E=33 shapes
+ _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)
   
-     (
-         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
+ # E=257 shapes
+ _add_cfg(16, 256, 257, block_m=16, ksplit=2)
+ _add_cfg(128, 256, 257, block_m=16, ksplit=2)
+ _add_cfg(
+     512,
+     256,
+     257,
+     block_m=32,
+     ksplit=0,
+     kernelName1=_4WG_STAGE1_M32,
+     kernelName2=_FLYDSL_STAGE2_M16_K128,
+     use_non_temporal_load=True,
+ )
   
-     M = hidden_states.shape[0]
-     d_hidden = config["d_hidden"]
-     d_hidden_pad = config["d_hidden_pad"]
-     d_expert = config["d_expert"]
-     d_expert_pad = config["d_expert_pad"]
-     E = config["n_routed_experts"] + config["n_shared_experts"]
-     topk = config["total_top_k"]
-     hidden_pad = d_hidden_pad - d_hidden
-     intermediate_pad = d_expert_pad - d_expert
+ # Pre-allocated buffer cache
+ _buffer_cache = {}
   
-     w1s = gate_up_weight_scale_shuffled.view(dtypes.fp8_e8m0)
-     w2s = down_weight_scale_shuffled.view(dtypes.fp8_e8m0)
+ def _get_or_alloc_sorting_buffers(M, E, topk, model_dim, block_size_M, device):
+     """Pre-allocate moe_sorting output buffers."""
+     key = ("sort", M, E, topk, model_dim, block_size_M)
+     if key in _buffer_cache:
+         return _buffer_cache[key]
   
-     if not _initialized:
-         csv_path = "/tmp/custom_tuned_fmoe_v112.csv"
-         with open(csv_path, 'w') as f:
-             f.write(_build_custom_csv())
-         aiter_root = os.path.dirname(os.path.abspath(aiter.__file__))
-         default_csv = os.path.join(aiter_root, "configs", "tuned_fmoe.csv")
-         dsv3_csv = os.path.join(aiter_root, "configs", "model_configs", "dsv3_fp4_tuned_fmoe.csv")
-         paths = []
-         if os.path.exists(default_csv):
-             paths.append(default_csv)
-         if os.path.exists(dsv3_csv):
-             paths.append(dsv3_csv)
-         paths.append(csv_path)
-         os.environ["AITER_CONFIG_FMOE"] = ":".join(paths)
-         os.environ["AITER_KSPLIT"] = "0"
-         get_2stage_cfgs.cache_clear()
-         _sort_fn = getattr(aiter, 'moe_sorting_opus_fwd', aiter.moe_sorting_fwd)
-         _initialized = True
+     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)
   
-     shape_key = (M, E, d_expert)
-     is_cktile = (M, E) in _CKTILE_SHAPES
+     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),
+         "sorted_expert_ids": torch.empty(max_num_m_blocks, dtype=dtypes.i32, device=device),
+         "num_valid_ids": torch.empty(2, dtype=dtypes.i32, device=device),
+         "moe_buf": torch.empty((M, model_dim), dtype=torch.bfloat16, device=device),
+     }
+     _buffer_cache[key] = bufs
+     return bufs
   
-     if shape_key not in _cache:
-         if is_cktile:
-             ks = _CKTILE_KSPLIT.get((M, E), 2)
-             os.environ["AITER_KSPLIT"] = str(ks)
-             os.environ["AITER_BYPASS_TUNE_CONFIG"] = "1"
+ def _get_or_alloc_a2(M, topk, inter_dim, device):
+     """Pre-allocate a2 intermediate buffer."""
+     key = ("a2", M, topk, inter_dim)
+     if key in _buffer_cache:
+         return _buffer_cache[key]
+     buf = torch.empty((M, topk, inter_dim), dtype=torch.bfloat16, device=device)
+     _buffer_cache[key] = buf
+     return buf
+  
+ def _get_or_alloc_quant_buffers(M, N, sorted_ids_len, topk, device):
+     """Pre-allocate quantization output buffers for fused_dynamic_mxfp4_quant_moe_sort."""
+     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,
+     )
+  
+     bufs = {"x_fp4": x_fp4, "blockscale": blockscale_e8m0_sorted}
+     _buffer_cache[key] = bufs
+     return bufs
+  
+ def _quant_prealloc(x, sorted_ids, num_valid_ids, token_num, topk, block_size, device):
+     """Inline fused_dynamic_mxfp4_quant_moe_sort with pre-allocated output buffers."""
+     M, N = x.shape
+     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]
+  
+     # Get pre-allocated buffers
+     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)
+  
+     _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,
+         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,
+     )
+  
+  
+ _injected = False
+  
+ def _inject_configs():
+     global _injected
+     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",
+             ]
+             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:
-             os.environ["AITER_KSPLIT"] = "0"
-             os.environ["AITER_BYPASS_TUNE_CONFIG"] = "0"
-         get_2stage_cfgs.cache_clear()
+             _fused_moe_module.cfg_2stages = {}
   
-         result = fused_moe(
-             hidden_states, gate_up_weight_shuffled, down_weight_shuffled,
-             topk_weights, topk_ids,
-             activation=ActivationType.Silu, quant_type=QuantType.per_1x32,
-             w1_scale=w1s, w2_scale=w2s,
-             hidden_pad=hidden_pad, intermediate_pad=intermediate_pad,
-         )
+     _fused_moe_module.cfg_2stages.update(_CUSTOM_CONFIGS)
   
-         E2, model_dim, inter_dim = get_inter_dim(
-             gate_up_weight_shuffled.shape, down_weight_shuffled.shape
+     # 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,
          )
-         is_shuffled = getattr(gate_up_weight_shuffled, "is_shuffled", False)
-         metadata = get_2stage_cfgs(
-             get_padded_M(M), model_dim, inter_dim, E, topk,
-             torch.bfloat16, dtypes.fp4x2, dtypes.fp4x2,
-             QuantType.per_1x32, True, ActivationType.Silu,
-             False, 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,
          )
-         block_m = metadata.block_m
+         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,
+                         )
+         return metadata
   
-         device = hidden_states.device
-         max_tok = M * topk + E * block_m
-         max_mb = (max_tok + block_m - 1) // block_m
+     _fused_moe_module.get_2stage_cfgs = _patched_get_2stage_cfgs
   
-         sorted_ids = torch.empty(max_tok, dtype=torch.int32, device=device)
-         sorted_weights = torch.empty(max_tok, dtype=torch.float32, device=device)
-         sorted_expert_ids = torch.empty(max_mb, dtype=torch.int32, device=device)
-         num_valid_ids = torch.empty(2, dtype=torch.int32, device=device)
-         moe_buf = torch.zeros((M, d_hidden_pad), dtype=torch.bfloat16, device=device)
   
-         entry = {
-             'block_m': block_m,
-             'metadata': metadata,
-             'sorted_ids': sorted_ids,
-             'sorted_weights': sorted_weights,
-             'sorted_expert_ids': sorted_expert_ids,
-             'num_valid_ids': num_valid_ids,
-             'moe_buf': moe_buf,
-             'is_cktile': is_cktile,
-             'inter_dim': inter_dim,
-         }
+ 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
   
-         if not is_cktile:
-             entry['a2_buf'] = torch.empty(
-                 (M, topk, inter_dim),
-                 dtype=torch.bfloat16,
-                 device=device,
-             )
+     _inject_configs()
   
-         # Enable non_temporal_load for E=257 CSV shapes (heuristic says True)
-         if E == 257 and not is_cktile:
-             metadata.stage1.keywords['use_non_temporal_load'] = True
-             metadata.stage2.keywords['use_non_temporal_load'] = True
+     hidden_pad = config["d_hidden_pad"] - config["d_hidden"]
+     intermediate_pad = config["d_expert_pad"] - config["d_expert"]
   
-         _cache[shape_key] = entry
-         return result
+     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)
   
-     c = _cache[shape_key]
+     padded_M = get_padded_M(M)
+     metadata = get_2stage_cfgs(
+         padded_M, model_dim, inter_dim, E, topk,
+         torch.bfloat16, dtypes.fp4x2, dtypes.fp4x2,
+         QuantType.per_1x32, True, ActivationType.Silu,
+         False, hidden_pad, intermediate_pad, True,
+     )
   
-     _sort_fn(topk_ids, topk_weights, c['sorted_ids'], c['sorted_weights'],
-              c['sorted_expert_ids'], c['num_valid_ids'], c['moe_buf'],
-              E, c['block_m'], None, None, _DISPATCH_POLICY)
+     block_size_M = int(metadata.block_m)
   
-     if c['is_cktile']:
-         a2 = c['metadata'].stage1(
-             hidden_states,
-             gate_up_weight_shuffled, down_weight_shuffled,
-             c['sorted_ids'], c['sorted_expert_ids'], c['num_valid_ids'],
-             None, topk,
-             block_m=c['block_m'], a1_scale=None, w1_scale=w1s,
+     # === Pre-allocated moe_sorting ===
+     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,
+     )
+  
+     # === Inline 2-stage pipeline ===
+     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)
+     a2_buf = _get_or_alloc_a2(M, topk, inter_dim, device)
+  
+     if metadata.ksplit > 1:
+         # cktile_moe path: bf16 activations, no fp4 quant
+         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,
          )
-         c['metadata'].stage2(
+  
+         # cktile_moe stage2: a2 is bf16, no inter-stage requant
+         _run_stage2(
+             metadata,
              a2,
-             gate_up_weight_shuffled, down_weight_shuffled,
-             c['sorted_ids'], c['sorted_expert_ids'], c['num_valid_ids'],
-             c['moe_buf'], topk,
-             w2_scale=w2s, a2_scale=None, block_m=c['block_m'],
-             sorted_weights=c['sorted_weights'],
+             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_fp4, a1_scale_sorted = fused_dynamic_mxfp4_quant_moe_sort(
-             hidden_states,
-             sorted_ids=c['sorted_ids'],
-             num_valid_ids=c['num_valid_ids'],
-             token_num=M,
-             topk=1,
-             block_size=c['block_m'],
+         # CK 2-stage path: fp4 activation quant with pre-allocated buffers
+         # Stage 1: quant activations + gate_up GEMM + SwiGLU
+         a1, a1_scale = _quant_prealloc(
+             hidden_states, sorted_ids, num_valid_ids,
+             token_num, 1, block_size_M, device,
          )
-         c['metadata'].stage1(
-             a1_fp4,
-             gate_up_weight_shuffled, down_weight_shuffled,
-             c['sorted_ids'], c['sorted_expert_ids'], c['num_valid_ids'],
-             c['a2_buf'], topk,
-             block_m=c['block_m'],
-             a1_scale=a1_scale_sorted,
-             w1_scale=w1s,
+  
+         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,
          )
-         a2_flat = c['a2_buf'].view(-1, c['inter_dim'])
-         a2_fp4, a2_scale_sorted = fused_dynamic_mxfp4_quant_moe_sort(
-             a2_flat,
-             sorted_ids=c['sorted_ids'],
-             num_valid_ids=c['num_valid_ids'],
-             token_num=M,
-             topk=topk,
-             block_size=c['block_m'],
+  
+         # Inter-stage requant: bf16 -> fp4 with pre-allocated buffers
+         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_fp4_3d = a2_fp4.view(M, topk, -1)
-         c['metadata'].stage2(
-             a2_fp4_3d,
-             gate_up_weight_shuffled, down_weight_shuffled,
-             c['sorted_ids'], c['sorted_expert_ids'], c['num_valid_ids'],
-             c['moe_buf'], topk,
-             w2_scale=w2s,
-             a2_scale=a2_scale_sorted,
-             block_m=c['block_m'],
-             sorted_weights=c['sorted_weights'],
+         a2_quant = a2_quant.view(token_num, topk, -1)
+  
+         # Stage 2: down GEMM + weighted reduction
+         _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 c['moe_buf']
+     return moe_out