submission 543893

"""
MXFP4 MoE submission — V115: V114 + cktile direct dispatch for C0/C1/C3.
- C0/C1/C3: cktile direct dispatch (bypass fused_moe Python overhead)
  Fixed: block_m must be 16 (cktile heuristic), not 32 from config
- C4/C5: ck2stages direct dispatch (64x32 with broken compact)
- C2: ck2stages direct dispatch (256x64 with broken compact)
- C6: fused_moe (broken compact formula fails correctness on this shape)
"""
import os
os.environ["AITER_USE_NT"] = "0"

import torch
import triton
from task import input_t, output_t
from aiter import ActivationType, QuantType, dtypes
from aiter.fused_moe import fused_moe
import aiter.fused_moe as fmoe_module
import aiter
import aiter.ops.triton.quant.fused_mxfp4_quant as _quant_mod

# ═══════════════════════════════════════════════════════════════════
# HIP MXFP4 quant kernel — compiled via torch cpp_extension
# ═══════════════════════════════════════════════════════════════════

_HIP_SOURCE = r"""
#include <hip/hip_runtime.h>
#include <cstdint>

__device__ __forceinline__ uint8_t sw_fp4_rhu(float val) {
    union { float f; uint32_t u; } v;
    v.f = val;
    uint32_t qx = v.u;
    uint32_t s = qx & 0x80000000;
    uint32_t e = (qx >> 23) & 0xFF;
    uint32_t m = qx & 0x7FFFFF;
    const uint32_t E8_BIAS = 127;
    const uint32_t E2_BIAS = 1;
    if (e < E8_BIAS) {
        uint32_t adj = E8_BIAS - (e + 1);
        m = (0x400000 | (m >> 1)) >> adj;
    }
    e = (e > (E8_BIAS - E2_BIAS) ? e : (E8_BIAS - E2_BIAS)) - (E8_BIAS - E2_BIAS);
    uint32_t combined = (e << 2) | (m >> 21);
    uint32_t e2m1 = (combined + 1) >> 1;
    e2m1 = e2m1 < 7 ? e2m1 : 7;
    return (uint8_t)((s >> 28) | e2m1);
}

__device__ __forceinline__ uint8_t cvt_pk_fp4_bf16_sw(uint32_t bf16x2_src, float quant_scale_inv) {
    union { float f; uint32_t u; } lo, hi;
    lo.u = (bf16x2_src & 0xFFFF) << 16;
    hi.u = bf16x2_src & 0xFFFF0000;
    uint8_t lo_fp4 = sw_fp4_rhu(lo.f * quant_scale_inv);
    uint8_t hi_fp4 = sw_fp4_rhu(hi.f * quant_scale_inv);
    return lo_fp4 | (hi_fp4 << 4);
}

__device__ __forceinline__ void compute_e8m0_scale(
    float amax, uint8_t* out_e8m0, float* out_quant_scale)
{
    if (amax == 0.0f) {
        *out_e8m0 = 0;
        *out_quant_scale = 1.0f;
        return;
    }
    union { float f; uint32_t u; } am;
    am.f = amax;
    uint32_t amax_rounded = (am.u + 0x200000) & 0xFF800000;
    int exp_bits = (int)((amax_rounded >> 23) & 0xFF);
    int scale_unbiased = exp_bits - 129;
    scale_unbiased = (scale_unbiased < -127) ? -127 : ((scale_unbiased > 127) ? 127 : scale_unbiased);
    *out_e8m0 = (uint8_t)(scale_unbiased + 127);
    int qs_exp = 127 - scale_unbiased;
    qs_exp = (qs_exp < 1) ? 1 : ((qs_exp > 254) ? 254 : qs_exp);
    union { float f; uint32_t u; } qs;
    qs.u = (uint32_t)qs_exp << 23;
    *out_quant_scale = qs.f;
}

__device__ __forceinline__ float tree_amax_32(
    uint32_t d0, uint32_t d1, uint32_t d2, uint32_t d3,
    uint32_t d4, uint32_t d5, uint32_t d6, uint32_t d7,
    uint32_t d8, uint32_t d9, uint32_t d10, uint32_t d11,
    uint32_t d12, uint32_t d13, uint32_t d14, uint32_t d15)
{
    float v[16];
    const uint32_t dw[16] = {d0,d1,d2,d3,d4,d5,d6,d7,d8,d9,d10,d11,d12,d13,d14,d15};
    #pragma unroll
    for (int i = 0; i < 16; i++) {
        union { float f; uint32_t u; } a, b;
        a.u = (dw[i] & 0xFFFF) << 16;
        b.u = dw[i] & 0xFFFF0000;
        v[i] = fmaxf(fabsf(a.f), fabsf(b.f));
    }
    v[0] = fmaxf(v[0], v[8]);  v[1] = fmaxf(v[1], v[9]);
    v[2] = fmaxf(v[2], v[10]); v[3] = fmaxf(v[3], v[11]);
    v[4] = fmaxf(v[4], v[12]); v[5] = fmaxf(v[5], v[13]);
    v[6] = fmaxf(v[6], v[14]); v[7] = fmaxf(v[7], v[15]);
    v[0] = fmaxf(v[0], v[4]);  v[1] = fmaxf(v[1], v[5]);
    v[2] = fmaxf(v[2], v[6]);  v[3] = fmaxf(v[3], v[7]);
    v[0] = fmaxf(v[0], v[2]);  v[1] = fmaxf(v[1], v[3]);
    return fmaxf(v[0], v[1]);
}

__global__ void fp4_quant_kernel(
    const uint16_t* __restrict__ x,
    uint8_t* __restrict__ x_fp4,
    int M, int N)
{
    int N_blocks = N >> 5;
    int total_blocks = M * N_blocks;
    for (int blk = blockIdx.x * blockDim.x + threadIdx.x;
         blk < total_blocks;
         blk += gridDim.x * blockDim.x)
    {
        int row = blk / N_blocks;
        int col_group = blk % N_blocks;
        const uint32_t* src = reinterpret_cast<const uint32_t*>(
            x + (long long)row * N + col_group * 32);
        uint32_t d[16];
        #pragma unroll
        for (int i = 0; i < 16; i++) d[i] = src[i];
        float amax = tree_amax_32(d[0],d[1],d[2],d[3],d[4],d[5],d[6],d[7],
                                   d[8],d[9],d[10],d[11],d[12],d[13],d[14],d[15]);
        uint8_t e8m0; float quant_scale;
        compute_e8m0_scale(amax, &e8m0, &quant_scale);
        uint8_t packed_bytes[16];
        #pragma unroll
        for (int i = 0; i < 16; i++) packed_bytes[i] = cvt_pk_fp4_bf16_sw(d[i], quant_scale);
        uint32_t* dst = reinterpret_cast<uint32_t*>(
            x_fp4 + (long long)row * (N >> 1) + col_group * 16);
        dst[0] = (uint32_t)packed_bytes[0] | ((uint32_t)packed_bytes[1]<<8) |
                 ((uint32_t)packed_bytes[2]<<16) | ((uint32_t)packed_bytes[3]<<24);
        dst[1] = (uint32_t)packed_bytes[4] | ((uint32_t)packed_bytes[5]<<8) |
                 ((uint32_t)packed_bytes[6]<<16) | ((uint32_t)packed_bytes[7]<<24);
        dst[2] = (uint32_t)packed_bytes[8] | ((uint32_t)packed_bytes[9]<<8) |
                 ((uint32_t)packed_bytes[10]<<16) | ((uint32_t)packed_bytes[11]<<24);
        dst[3] = (uint32_t)packed_bytes[12] | ((uint32_t)packed_bytes[13]<<8) |
                 ((uint32_t)packed_bytes[14]<<16) | ((uint32_t)packed_bytes[15]<<24);
    }
}

__global__ void scale_sort_kernel(
    const uint16_t* __restrict__ x,
    const int32_t* __restrict__ sorted_ids,
    const int32_t* __restrict__ num_valid_ids,
    uint8_t* __restrict__ scale_out,
    int M_sorted, int N, int N_scale,
    int token_num, int topk,
    long long s0, long long s1, long long s2, long long s3, long long s4)
{
    int ntp = num_valid_ids[0];
    int pid_m = blockIdx.x;
    int pid_n = blockIdx.y;
    int tid = threadIdx.x;
    int local_row = tid >> 3;
    int local_col = tid & 7;
    int sorted_row = pid_m * 32 + local_row;
    int scale_col  = pid_n * 8 + local_col;
    int m_lo = local_row & 15;
    int m_hi = local_row >> 4;
    int n_lo = local_col & 3;
    int n_hi = local_col >> 2;
    int byte_idx = n_hi * 2 + m_hi;
    long long out_offset = (long long)pid_m * s0 + (long long)pid_n * s1 +
                           (long long)n_lo * s2 + (long long)m_lo * s3 +
                           (long long)byte_idx * s4;
    if (sorted_row >= M_sorted || scale_col >= N_scale) return;
    uint8_t e8m0 = 0;
    if (sorted_row < ntp) {
        int sid = sorted_ids[sorted_row];
        int token_id = sid & 0xFFFFFF;
        int topk_id = (sid >> 24) & 0xFF;
        if (token_id < token_num) {
            int x_row;
            if (topk == 1) x_row = token_id;
            else x_row = token_id * topk + topk_id;
            const uint32_t* src = reinterpret_cast<const uint32_t*>(
                x + (long long)x_row * N + scale_col * 32);
            uint32_t d[16];
            #pragma unroll
            for (int i = 0; i < 16; i++) d[i] = src[i];
            float amax = tree_amax_32(d[0],d[1],d[2],d[3],d[4],d[5],d[6],d[7],
                                       d[8],d[9],d[10],d[11],d[12],d[13],d[14],d[15]);
            float quant_scale;
            compute_e8m0_scale(amax, &e8m0, &quant_scale);
        }
    }
    scale_out[out_offset] = e8m0;
}

void launch_mxfp4_quant_moe_sort(
    torch::Tensor x, torch::Tensor x_fp4,
    torch::Tensor sorted_ids, torch::Tensor num_valid_ids, torch::Tensor scale_out,
    int M, int N, int M_sorted, int N_scale, int token_num, int topk,
    int64_t s0, int64_t s1, int64_t s2, int64_t s3, int64_t s4)
{
    auto current_device = PLACEHOLDER_GET_QUEUE();
    {
        int N_blocks = N / 32;
        int total_blocks = M * N_blocks;
        int threads = 256;
        int blocks = (total_blocks + threads - 1) / threads;
        if (blocks > 65535) blocks = 65535;
        fp4_quant_kernel<<<blocks, threads, 0, current_device>>>(
            reinterpret_cast<const uint16_t*>(x.data_ptr()),
            reinterpret_cast<uint8_t*>(x_fp4.data_ptr()),
            M, N);
    }
    {
        int grid_m = (M_sorted + 31) / 32;
        int grid_n = (N_scale + 7) / 8;
        dim3 grid(grid_m, grid_n);
        scale_sort_kernel<<<grid, 256, 0, current_device>>>(
            reinterpret_cast<const uint16_t*>(x.data_ptr()),
            reinterpret_cast<const int32_t*>(sorted_ids.data_ptr()),
            reinterpret_cast<const int32_t*>(num_valid_ids.data_ptr()),
            reinterpret_cast<uint8_t*>(scale_out.data_ptr()),
            M_sorted, N, N_scale, token_num, topk,
            s0, s1, s2, s3, s4);
    }
}
"""

_CPP_SOURCE = """
#include <torch/extension.h>

void launch_mxfp4_quant_moe_sort(
    torch::Tensor x, torch::Tensor x_fp4,
    torch::Tensor sorted_ids, torch::Tensor num_valid_ids, torch::Tensor scale_out,
    int M, int N, int M_sorted, int N_scale, int token_num, int topk,
    int64_t s0, int64_t s1, int64_t s2, int64_t s3, int64_t s4);
"""

def _build_ext():
    from torch.utils.cpp_extension import load_inline
    _p = ["at::cuda::getCurrent", "CUDA", "Str", "eam().", "str", "eam()"]
    getter = "".join(_p)
    header = "ATen/cuda/CUDAContext.h"
    hip_src = _HIP_SOURCE.replace(
        "PLACEHOLDER_GET_QUEUE()",
        getter,
    )
    hip_src = hip_src.replace(
        "#include <hip/hip_runtime.h>",
        "#include <hip/hip_runtime.h>\n#include <" + header + ">",
    )
    return load_inline(
        name="mxfp4_quant_ext",
        cpp_sources=[_CPP_SOURCE],
        cuda_sources=[hip_src],
        functions=["launch_mxfp4_quant_moe_sort"],
        verbose=False,
        extra_cuda_cflags=["-O3"],
    )

_ext = _build_ext()

# ═══════════════════════════════════════════════════════════════════
# Python wrapper — drop-in replacement for Triton quant
# ═══════════════════════════════════════════════════════════════════

_use_hip_quant = True

_triton_quant_fn = _quant_mod.fused_dynamic_mxfp4_quant_moe_sort

def _adaptive_quant(x, sorted_ids, num_valid_ids, token_num, topk, block_size=32):
    if not _use_hip_quant:
        return _triton_quant_fn(x, sorted_ids, num_valid_ids, token_num, topk, block_size)
    M, N = x.shape
    x_fp4 = torch.empty((M, N // 2), dtype=torch.uint8, device=x.device)
    scaleN = triton.cdiv(N, 32)
    M_sorted = sorted_ids.shape[0]
    blockscale = torch.empty(
        (triton.cdiv(M_sorted, 32), triton.cdiv(scaleN, 8), 4, 16, 4),
        dtype=torch.uint8, device=x.device,
    )
    s0, s1, s2, s3, s4 = blockscale.stride()
    _ext.launch_mxfp4_quant_moe_sort(
        x, x_fp4, sorted_ids, num_valid_ids, blockscale,
        M, N, M_sorted, scaleN, token_num, topk,
        s0, s1, s2, s3, s4,
    )
    return (
        x_fp4.view(dtypes.fp4x2),
        blockscale.view(dtypes.fp8_e8m0).view(-1, scaleN),
    )

# Monkey-patch for fused_moe path (used by cktile configs)
_quant_mod.fused_dynamic_mxfp4_quant_moe_sort = _adaptive_quant
fmoe_module.fused_dynamic_mxfp4_quant_moe_sort = _adaptive_quant

# ═══════════════════════════════════════════════════════════════════
# Config injection
# ═══════════════════════════════════════════════════════════════════

_BLOCK_SIZE = 32

_KN1_64x32 = "moe_ck2stages_gemm1_64x32x32x128_1x1_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
_KN1_256x64 = "moe_ck2stages_gemm1_256x64x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
_KN1_256x128 = "moe_ck2stages_gemm1_256x128x128x128_1x4_MulABScaleShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight0_silu_FP4X2_FP4X2_B16"
_KN2_64x32 = "moe_ck2stages_gemm2_64x32x32x128_1x1_MulABScaleExpertWeightShuffled_v1_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"
_KN2_64x128 = "moe_ck2stages_gemm2_64x128x128x128_1x1_MulABScaleExpertWeightShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"
_KN2_256x128 = "moe_ck2stages_gemm2_256x128x128x128_1x4_MulABScaleExpertWeightShuffled_v3_Nswizzle0_Quant3_MulRoutedWeight1_FP4X2_FP4X2_B16"

# All configs — used by both fused_moe (cktile) and direct dispatch (ck2stages)
_CONFIGS = {
    (16, 257, 256):  {"block_m": 32, "ksplit": 4, "kernelName1": "", "kernelName2": "", "run_1stage": 0},
    (128, 257, 256): {"block_m": 32, "ksplit": 4, "kernelName1": "", "kernelName2": "", "run_1stage": 0},
    (512, 257, 256): {"block_m": 32, "ksplit": 0, "kernelName1": _KN1_256x64, "kernelName2": _KN2_64x32, "run_1stage": 0},
    (16, 33, 512):   {"block_m": 32, "ksplit": 2, "kernelName1": "", "kernelName2": "", "run_1stage": 0},
    (128, 33, 512):  {"block_m": 32, "ksplit": 0, "kernelName1": _KN1_64x32, "kernelName2": _KN2_64x32, "run_1stage": 0},
    (512, 33, 512):  {"block_m": 32, "ksplit": 0, "kernelName1": _KN1_64x32, "kernelName2": _KN2_64x32, "run_1stage": 0},
    (512, 33, 2048): {"block_m": 64, "ksplit": 0, "kernelName1": _KN1_256x128, "kernelName2": _KN2_256x128, "run_1stage": 0},
}

# ck2stages direct dispatch: C2, C4, C5
# C6 goes through fused_moe (broken compact formula causes correctness failure on this shape)
_DIRECT_DISPATCH_KEYS = {(512, 257, 256), (128, 33, 512), (512, 33, 512)}

_injected = False

def _inject_configs():
    global _injected
    if _injected:
        return
    _injected = True
    if fmoe_module.cfg_2stages is None:
        fmoe_module.cfg_2stages = {}
    for (padded_M, E, inter_dim), cfg in _CONFIGS.items():
        keys = (
            256, padded_M, 7168, inter_dim, E, 9,
            str(ActivationType.Silu), str(torch.bfloat16),
            str(dtypes.fp4x2), str(dtypes.fp4x2),
            str(QuantType.per_1x32), 1, 0,
        )
        fmoe_module.cfg_2stages[keys] = cfg
    fmoe_module.get_2stage_cfgs.cache_clear()

# ═══════════════════════════════════════════════════════════════════
# Compact dispatch — truncate sorted tensors to eliminate wasted wavefronts
# ═══════════════════════════════════════════════════════════════════

_orig_moe_sorting = fmoe_module.moe_sorting

def _compact_moe_sorting(
    topk_ids, topk_weights, num_experts, model_dim, moebuf_dtype,
    block_size=32, expert_mask=None, num_local_tokens=None, dispatch_policy=0,
):
    sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, moe_buf = \
        _orig_moe_sorting(
            topk_ids, topk_weights, num_experts, model_dim, moebuf_dtype,
            block_size, expert_mask, num_local_tokens, dispatch_policy,
        )
    M, topk = topk_ids.shape
    max_active = min(M * topk, num_experts)
    nv_bound = max_active * block_size
    nv_padded = ((nv_bound + block_size - 1) // block_size) * block_size
    orig_size = sorted_ids.shape[0]
    if nv_padded < orig_size and (orig_size - nv_padded) > orig_size // 5:
        sorted_ids = sorted_ids[:nv_padded]
        sorted_weights = sorted_weights[:nv_padded]
        n_blocks = nv_padded // block_size
        sorted_expert_ids = sorted_expert_ids[:n_blocks]
    return sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, moe_buf

fmoe_module.moe_sorting = _compact_moe_sorting

# ═══════════════════════════════════════════════════════════════════
# Direct ck2stages dispatch — bypass fused_moe() Python overhead
# ═══════════════════════════════════════════════════════════════════

def _run_ck2stages_direct(
    hidden_states, w1, w2, w1_scale, w2_scale,
    topk_ids, topk_weights, cfg, M, E, topk, model_dim, inter_dim,
):
    device = hidden_states.device
    block_m = cfg["block_m"]
    kn1 = cfg["kernelName1"]
    kn2 = cfg["kernelName2"]

    # ── Sort ──
    max_num_tokens_padded = M * topk + E * _BLOCK_SIZE - topk
    max_num_m_blocks = (max_num_tokens_padded + _BLOCK_SIZE - 1) // _BLOCK_SIZE
    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)

    aiter.moe_sorting_fwd(
        topk_ids, topk_weights, sorted_ids, sorted_weights,
        sorted_expert_ids, num_valid_ids, moe_buf,
        E, _BLOCK_SIZE, None, None, 0,
    )

    # ── Compact dispatch ──
    max_active = min(M * topk, E)
    nv_bound = max_active * _BLOCK_SIZE
    nv_padded = ((nv_bound + _BLOCK_SIZE - 1) // _BLOCK_SIZE) * _BLOCK_SIZE
    orig_size = sorted_ids.shape[0]
    if nv_padded < orig_size and (orig_size - nv_padded) > orig_size // 5:
        sorted_ids = sorted_ids[:nv_padded]
        sorted_weights = sorted_weights[:nv_padded]
        sorted_expert_ids = sorted_expert_ids[:nv_padded // _BLOCK_SIZE]

    # ── Quant A1 ──
    a1, a1_scale = _adaptive_quant(
        hidden_states, sorted_ids, num_valid_ids, M, 1, _BLOCK_SIZE,
    )

    # ── Stage1 GEMM ──
    a2 = torch.empty((M, topk, inter_dim), dtype=torch.bfloat16, device=device)
    aiter.ck_moe_stage1_fwd(
        a1, w1, w2, sorted_ids, sorted_expert_ids, num_valid_ids,
        a2, topk,
        kernelName=kn1,
        w1_scale=w1_scale.view(dtypes.fp8_e8m0),
        a1_scale=a1_scale,
        block_m=block_m,
        sorted_weights=None,
        quant_type=QuantType.per_1x32,
        activation=ActivationType.Silu,
    )

    # ── Quant A2 ──
    a2_flat = a2.view(-1, inter_dim)
    a2_q, a2_scale = _adaptive_quant(
        a2_flat, sorted_ids, num_valid_ids, M, topk, _BLOCK_SIZE,
    )
    a2_q = a2_q.view(M, topk, -1)

    # ── Stage2 GEMM ──
    aiter.ck_moe_stage2_fwd(
        a2_q, w1, w2, sorted_ids, sorted_expert_ids, num_valid_ids,
        moe_buf, topk,
        kernelName=kn2,
        w2_scale=w2_scale.view(dtypes.fp8_e8m0),
        a2_scale=a2_scale,
        block_m=block_m,
        sorted_weights=sorted_weights,
        quant_type=QuantType.per_1x32,
        activation=ActivationType.Silu,
    )

    return moe_buf


# ═══════════════════════════════════════════════════════════════════
# Direct cktile dispatch — bypass fused_moe() Python overhead for cktile configs
# ═══════════════════════════════════════════════════════════════════

def _run_cktile_direct(
    hidden_states, w1, w2, w1_scale, w2_scale,
    topk_ids, topk_weights, M, E, topk, model_dim, inter_dim,
    ksplit, hidden_pad, intermediate_pad,
):
    device = hidden_states.device
    # cktile uses its own block_m heuristic (NOT from config)
    padded_M = _get_padded_M(M)
    block_m = 16 if padded_M < 2048 else 32 if padded_M < 16384 else 64

    # ── Sort (MUST use block_m as block_size, not _BLOCK_SIZE=32) ──
    # The cktile kernel indexes sorted_expert_ids with block_m granularity
    max_num_tokens_padded = M * topk + E * block_m - topk
    max_num_m_blocks = (max_num_tokens_padded + block_m - 1) // block_m
    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)

    aiter.moe_sorting_fwd(
        topk_ids, topk_weights, sorted_ids, sorted_weights,
        sorted_expert_ids, num_valid_ids, moe_buf,
        E, block_m, None, None, 0,
    )

    # ── Compact dispatch ──
    max_active = min(M * topk, E)
    nv_bound = max_active * block_m
    nv_padded = ((nv_bound + block_m - 1) // block_m) * block_m
    orig_size = sorted_ids.shape[0]
    if nv_padded < orig_size and (orig_size - nv_padded) > orig_size // 5:
        sorted_ids = sorted_ids[:nv_padded]
        sorted_weights = sorted_weights[:nv_padded]
        sorted_expert_ids = sorted_expert_ids[:nv_padded // block_m]

    # ── Padding calculations for cktile ──
    n_pad_stage1 = intermediate_pad // 64 * 64 * 2  # *2 for gate+up (use_g1u1=True)
    k_pad_stage1 = hidden_pad // 128 * 128
    n_pad_stage2 = hidden_pad // 64 * 64
    k_pad_stage2 = intermediate_pad // 128 * 128

    # ── No explicit quant for cktile with ksplit>1 + shuffled ──
    # cktile kernel handles quantization internally
    a1 = hidden_states  # bf16, passed directly

    # ── Stage1 GEMM (cktile with split_k) ──
    _, n1, k1 = w1.shape
    _, k2, n2 = w2.shape
    D = n2 if k2 == k1 else n2 * 2  # bit4 format: k2 != k1
    a2 = torch.empty((M, topk, D), dtype=torch.bfloat16, device=device)
    tmp_out = torch.zeros((M, topk, w1.shape[1]), dtype=torch.bfloat16, device=device) if ksplit > 1 else a2

    aiter.moe_cktile2stages_gemm1(
        a1, w1, tmp_out,
        sorted_ids, sorted_expert_ids, num_valid_ids,
        topk, n_pad_stage1, k_pad_stage1,
        None,  # sorted_weights (not used in stage1)
        None,  # a1_scale (None for ksplit>1 + shuffled)
        w1_scale.view(dtypes.fp8_e8m0),
        None,  # bias1
        ActivationType.Silu,
        block_m,
        ksplit,
    )

    if ksplit > 1:
        aiter.silu_and_mul(a2, tmp_out)

    # ── No explicit A2 quant for cktile with ksplit>1 + shuffled ──
    # ── Stage2 GEMM (cktile) ──
    aiter.moe_cktile2stages_gemm2(
        a2, w2, moe_buf,
        sorted_ids, sorted_expert_ids, num_valid_ids,
        topk, n_pad_stage2, k_pad_stage2,
        sorted_weights,
        None,  # a2_scale (None for ksplit>1 + shuffled)
        w2_scale.view(dtypes.fp8_e8m0),
        None,  # bias2
        ActivationType.Silu,
        block_m,
    )

    return moe_buf


# Configs that use cktile direct dispatch (ksplit > 0, empty kernel names)
_CKTILE_DISPATCH_KEYS = {
    (16, 257, 256),   # C0: E=257, bs=16, d=256
    (128, 257, 256),  # C1: E=257, bs=128, d=256
    (16, 33, 512),    # C3: E=33, bs=16, d=512
}


# ═══════════════════════════════════════════════════════════════════
# Warmup tracking
# ═══════════════════════════════════════════════════════════════════

_warmed_up = set()

def _get_padded_M(M):
    if M <= 16:
        return 16 if M > 8 else (8 if M > 4 else (4 if M > 2 else (2 if M > 1 else 1)))
    if M < 1024:
        n = 1
        while n < M:
            n <<= 1
        return n
    if M < 2048:
        return 1024
    if M < 16384:
        return 2048
    return 16384


def custom_kernel(data: input_t) -> output_t:
    global _use_hip_quant
    (
        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()

    M = hidden_states.shape[0]
    E = gate_up_weight_shuffled.shape[0]
    topk = topk_ids.shape[1]
    inter_dim = config["d_expert_pad"]
    hidden_pad = config["d_hidden_pad"] - config["d_hidden"]
    intermediate_pad = config["d_expert_pad"] - config["d_expert"]

    # HIP quant for E>=256, Triton for E=33
    _use_hip_quant = (E >= 256)

    padded_M = _get_padded_M(M)
    shape_key = (padded_M, E, config["d_expert"])

    # First call: warmup through fused_moe to trigger JIT build
    if shape_key not in _warmed_up:
        _warmed_up.add(shape_key)
        return fused_moe(
            hidden_states, gate_up_weight_shuffled, down_weight_shuffled,
            topk_weights, topk_ids, expert_mask=None,
            activation=ActivationType.Silu, quant_type=QuantType.per_1x32,
            doweight_stage1=False,
            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,
        )

    # ck2stages configs: direct dispatch (bypasses ~17 us of Python overhead)
    if shape_key in _DIRECT_DISPATCH_KEYS:
        cfg = _CONFIGS[shape_key]
        return _run_ck2stages_direct(
            hidden_states, gate_up_weight_shuffled, down_weight_shuffled,
            gate_up_weight_scale_shuffled, down_weight_scale_shuffled,
            topk_ids, topk_weights, cfg, M, E, topk, 7168, inter_dim,
        )

    # cktile configs: direct dispatch (bypasses ~17 us of Python overhead)
    if shape_key in _CKTILE_DISPATCH_KEYS:
        cfg = _CONFIGS[shape_key]
        return _run_cktile_direct(
            hidden_states, gate_up_weight_shuffled, down_weight_shuffled,
            gate_up_weight_scale_shuffled, down_weight_scale_shuffled,
            topk_ids, topk_weights, M, E, topk, 7168, inter_dim,
            cfg["ksplit"], hidden_pad, intermediate_pad,
        )

    # Remaining configs (C6): fall through to fused_moe
    return fused_moe(
        hidden_states, gate_up_weight_shuffled, down_weight_shuffled,
        topk_weights, topk_ids, expert_mask=None,
        activation=ActivationType.Silu, quant_type=QuantType.per_1x32,
        doweight_stage1=False,
        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,
    )