submission 638261

"""
test_v24_hw_quant_occ: HW quant with higher occupancy.
v23 had 64 threads/block = 64 blocks = 3% occupancy.
Use 256 threads/block = 16 blocks for better GPU utilization.
"""
import functools
import hashlib
import os

os.environ.setdefault("PYTORCH_ROCM_ARCH", "gfx950")
os.environ.setdefault("CXX", "clang++")

import torch
import triton
import triton.language as tl
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t

SCALE_GROUP_SIZE = 32

# ============================================================================
# HIP C++ quantization kernel (for M>16, proven fast)
# ============================================================================

CPP_SRC = r"""
#include <torch/extension.h>
void quantize_a_hip(torch::Tensor a, torch::Tensor a_q, torch::Tensor a_scale);
torch::Tensor gemm_fused_16x16_hip(torch::Tensor A, torch::Tensor Bq,
    torch::Tensor Bsc, int64_t N, torch::Tensor out);
"""

HIP_SRC = r"""
#include <torch/extension.h>
#include <hip/hip_runtime.h>
#include <hip/hip_ext_ocp.h>
#include <hip/amd_detail/amd_hip_bf16.h>
#include <stdint.h>

#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA/ROCm tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")

namespace {

constexpr int QUANT_GROUP = 32;
constexpr int GROUPS_PER_BLOCK = 4;
constexpr int BLOCK_THREADS = QUANT_GROUP * GROUPS_PER_BLOCK;  // 128

__device__ inline uint32_t float_to_bits(float x) {
    union { float f; uint32_t u; } v;
    v.f = x;
    return v.u;
}

__device__ inline float bits_to_float(uint32_t x) {
    union { uint32_t u; float f; } v;
    v.u = x;
    return v.f;
}

__device__ inline uint8_t encode_mxfp4(float x) {
    constexpr uint32_t FP32_EXP_BIAS = 127;
    constexpr uint32_t FP4_EXP_BIAS = 1;
    constexpr uint32_t FP32_MBITS = 23;
    constexpr uint32_t FP4_MBITS = 1;
    constexpr uint32_t FP4_EBITS = 2;
    constexpr uint8_t SIGN_MASK = 1u << (FP4_EBITS + FP4_MBITS);
    constexpr float MAX_NORMAL = 6.0f;
    constexpr float MIN_NORMAL = 1.0f;
    constexpr int32_t DENORM_EXP =
        (FP32_EXP_BIAS - FP4_EXP_BIAS) + (FP32_MBITS - FP4_MBITS) + 1;
    constexpr int32_t DENORM_MASK_INT = DENORM_EXP << FP32_MBITS;
    constexpr uint32_t NORMAL_ROUND_BIAS_BITS = 0xC11FFFFFu;

    const uint32_t bits = float_to_bits(x);
    const uint32_t sign = bits & 0x80000000u;
    const uint32_t abs_bits = bits ^ sign;
    const float abs_x = bits_to_float(abs_bits);

    bool saturate = abs_x >= MAX_NORMAL;
    bool denormal = (!saturate) && (abs_x < MIN_NORMAL);

    uint8_t value = 0x7;
    if (denormal) {
        const float denorm_float =
            abs_x + bits_to_float(static_cast<uint32_t>(DENORM_MASK_INT));
        const uint32_t denorm_bits = float_to_bits(denorm_float);
        value = static_cast<uint8_t>(
            denorm_bits - static_cast<uint32_t>(DENORM_MASK_INT)
        );
    } else if (!saturate) {
        const uint32_t mant_odd = (abs_bits >> (FP32_MBITS - FP4_MBITS)) & 1u;
        const int32_t rounded =
            static_cast<int32_t>(abs_bits)
            + static_cast<int32_t>(NORMAL_ROUND_BIAS_BITS)
            + static_cast<int32_t>(mant_odd);
        value = static_cast<uint8_t>(
            static_cast<uint32_t>(rounded) >> (FP32_MBITS - FP4_MBITS)
        );
    }

    return value | (sign ? SIGN_MASK : 0u);
}

// Warp-parallel quantize with software encode_mxfp4.
// 32 threads per group: each loads 1 bf16, warp-reduce for amax.
// Much higher occupancy than single-thread-per-group (1024 vs 16 blocks for s5).

__global__ void __launch_bounds__(BLOCK_THREADS)
quantize_a_kernel(
    const __hip_bfloat16* __restrict__ a,
    uint8_t* __restrict__ a_q,
    uint8_t* __restrict__ a_scale,
    int m,
    int k
) {
    const int row = blockIdx.y;
    const int group_in_block = threadIdx.x / QUANT_GROUP;
    const int lane = threadIdx.x % QUANT_GROUP;
    const int k_block = blockIdx.x * GROUPS_PER_BLOCK + group_in_block;

    if (k_block >= k / QUANT_GROUP) return;

    const int col = k_block * QUANT_GROUP + lane;

    float value = 0.0f;
    if (row < m && col < k) {
        value = static_cast<float>(a[row * k + col]);
    }

    float abs_val = fabsf(value);
    for (int offset = 16; offset > 0; offset >>= 1) {
        float other = __shfl_xor(abs_val, offset);
        abs_val = fmaxf(abs_val, other);
    }

    const float amax = abs_val;
    uint32_t amax_bits = float_to_bits(amax);
    amax_bits = (amax_bits + 0x00200000u) & 0xFF800000u;
    int ef = (amax_bits >> 23u) & 0xFFu;
    int scale_exp = (amax_bits == 0u) ? -127 : max(-127, min(127, ef - 127 - 2));
    float quant_scale = (scale_exp <= 126) ? bits_to_float(static_cast<uint32_t>(127 - scale_exp) << 23) : 0.0f;
    uint8_t scale_byte = static_cast<uint8_t>(scale_exp + 127);

    if (lane == 0) {
        a_scale[row * (k / QUANT_GROUP) + k_block] = scale_byte;
    }

    uint8_t nibble = encode_mxfp4(value * quant_scale);
    uint8_t partner_nibble = __shfl_xor(nibble, 1);
    if ((lane & 1) == 0 && row < m) {
        uint8_t packed = nibble | (partner_nibble << 4);
        const int out_col = k_block * 16 + (lane / 2);
        a_q[row * (k / 2) + out_col] = packed;
    }
}

// ══════ HIP 16x16 MFMA kernel types and helpers ══════
typedef int __attribute__((ext_vector_type(8))) i32x8;
typedef float __attribute__((ext_vector_type(4))) f32x4;
typedef __bf16 bf16v2_t __attribute__((ext_vector_type(2)));

__device__ __forceinline__ void quant_32_hw(const uint16_t* ap, int out[4], int32_t& spk) {
    uint16_t mx = 0;
    #pragma unroll
    for (int j = 0; j < 32; j++) mx = max(mx, (uint16_t)(ap[j] & 0x7FFF));
    uint32_t au = (((uint32_t)mx << 16) + 0x200000u) & 0xFF800000u;
    int ef = (au >> 23u) & 0xFFu;
    int su = (au == 0u) ? -127 : max(-127, min(127, ef - 127 - 2));
    float hs = (su >= -126) ? __uint_as_float((uint32_t)(su + 127) << 23) : 0.0f;
    #pragma unroll
    for (int j = 0; j < 4; j++) {
        out[j] = 0;
        out[j] = __builtin_amdgcn_cvt_scalef32_pk_fp4_bf16(out[j], *reinterpret_cast<const bf16v2_t*>(&ap[j*8+0]), hs, 0);
        out[j] = __builtin_amdgcn_cvt_scalef32_pk_fp4_bf16(out[j], *reinterpret_cast<const bf16v2_t*>(&ap[j*8+2]), hs, 1);
        out[j] = __builtin_amdgcn_cvt_scalef32_pk_fp4_bf16(out[j], *reinterpret_cast<const bf16v2_t*>(&ap[j*8+4]), hs, 2);
        out[j] = __builtin_amdgcn_cvt_scalef32_pk_fp4_bf16(out[j], *reinterpret_cast<const bf16v2_t*>(&ap[j*8+6]), hs, 3);
    }
    spk = (int32_t)(uint8_t)(su + 127);
}

__device__ __forceinline__ int32_t load_b_scale_16(const uint8_t* Bsc,
    int ng, int bb, int SNG, int N) {
    if (ng >= N) return 127;
    int ifl = (ng/32)*(SNG*256) + (bb/8)*256 + (bb%4)*64
            + (ng%16)*4 + ((bb%8)/4)*2 + (ng%32)/16;
    return (int32_t)Bsc[ifl];
}

__global__ void __launch_bounds__(256, 3)
gemm_fused_16x16_kern(
    const uint16_t* __restrict__ A, const uint8_t* __restrict__ Bq,
    const uint8_t* __restrict__ Bsc, uint16_t* __restrict__ C,
    int M, int N, int K, int strA, int strBq, int BscSN)
{
    const int mt = blockIdx.x, nt = blockIdx.y;
    const int wid = threadIdx.x >> 6;
    const int lid = threadIdx.x & 63;
    const int t_row = lid & 15;
    const int kpart = lid >> 4;
    const int SNG = BscSN >> 3;
    const int m_row = mt * 16 + t_row;
    const int a_row_off = m_row * strA;
    const int b_ng = nt * 16 + t_row;

    int total_steps = K >> 7;
    int steps_per_warp = (total_steps + 3) >> 2;
    int k_start = wid * steps_per_warp * 128;
    int k_end = min(k_start + steps_per_warp * 128, K);

    f32x4 c_acc = {0.0f, 0.0f, 0.0f, 0.0f};

    for (int kb = k_start; kb < k_end; kb += 128) {
        int k_off = kb + kpart * 32;
        int bk = (kb >> 1) + kpart * 16;
        uint4 a0, a1, a2, a3;
        if (m_row < M && k_off + 31 < K) {
            const uint4* src = reinterpret_cast<const uint4*>(&A[a_row_off + k_off]);
            a0 = src[0]; a1 = src[1]; a2 = src[2]; a3 = src[3];
        } else {
            a0 = {0,0,0,0}; a1 = {0,0,0,0}; a2 = {0,0,0,0}; a3 = {0,0,0,0};
        }
        int b_i32[4];
        if (b_ng < N && bk + 15 < (K >> 1)) {
            uint4 bd = *reinterpret_cast<const uint4*>(&Bq[b_ng * strBq + bk]);
            b_i32[0]=((int*)&bd)[0]; b_i32[1]=((int*)&bd)[1];
            b_i32[2]=((int*)&bd)[2]; b_i32[3]=((int*)&bd)[3];
        } else { b_i32[0]=0; b_i32[1]=0; b_i32[2]=0; b_i32[3]=0; }
        int32_t b_spk = load_b_scale_16(Bsc, b_ng, (kb >> 5) + kpart, SNG, N);
        uint16_t a_local[32];
        uint4* dst = reinterpret_cast<uint4*>(a_local);
        dst[0] = a0; dst[1] = a1; dst[2] = a2; dst[3] = a3;
        int a_i32[4]; int32_t a_spk;
        quant_32_hw(a_local, a_i32, a_spk);
        i32x8 am = {a_i32[0], a_i32[1], a_i32[2], a_i32[3], 0, 0, 0, 0};
        i32x8 bm = {b_i32[0], b_i32[1], b_i32[2], b_i32[3], 0, 0, 0, 0};
        c_acc = __builtin_amdgcn_mfma_scale_f32_16x16x128_f8f6f4(
            am, bm, c_acc, 4, 4, 0, a_spk, 0, b_spk);
    }

    __shared__ float reduce16[4][64 * 4 + 1];
    #pragma unroll
    for (int i = 0; i < 4; i++)
        reduce16[wid][lid * 4 + i] = c_acc[i];
    __syncthreads();
    if (wid == 0) {
        for (int j = 0; j < 4; j++) {
            float sum = reduce16[0][lid*4+j] + reduce16[1][lid*4+j]
                      + reduce16[2][lid*4+j] + reduce16[3][lid*4+j];
            int mo = mt * 16 + kpart * 4 + j;
            int no = nt * 16 + t_row;
            if (mo < M && no < N) {
                uint32_t fp = __float_as_uint(sum);
                fp += 0x7FFFu + ((fp >> 16) & 1u);
                C[mo * N + no] = (uint16_t)(fp >> 16u);
            }
        }
    }
}

}  // namespace

void quantize_a_hip(torch::Tensor a, torch::Tensor a_q, torch::Tensor a_scale) {
    CHECK_CUDA(a);
    CHECK_CONTIGUOUS(a);
    TORCH_CHECK(a.scalar_type() == torch::kBFloat16, "a must be bfloat16");
    TORCH_CHECK(a.dim() == 2, "a must be 2D");

    const int m = static_cast<int>(a.size(0));
    const int k = static_cast<int>(a.size(1));

    const int num_k_groups = k / QUANT_GROUP;
    const dim3 q_block(BLOCK_THREADS);
    const dim3 q_grid((num_k_groups + GROUPS_PER_BLOCK - 1) / GROUPS_PER_BLOCK, m);
    quantize_a_kernel<<<q_grid, q_block>>>(
        reinterpret_cast<const __hip_bfloat16*>(a.data_ptr<at::BFloat16>()),
        a_q.data_ptr<uint8_t>(),
        a_scale.data_ptr<uint8_t>(),
        m,
        k
    );
    TORCH_CHECK(hipGetLastError() == hipSuccess, "quantize_a_kernel launch failed");
}

torch::Tensor gemm_fused_16x16_hip(torch::Tensor A, torch::Tensor Bq,
    torch::Tensor Bsc, int64_t N, torch::Tensor out) {
    CHECK_CUDA(A);
    int M = (int)A.size(0), K = (int)A.size(1);
    dim3 grid((M + 15) / 16, ((int)N + 15) / 16);
    gemm_fused_16x16_kern<<<grid, 256>>>(
        reinterpret_cast<const uint16_t*>(A.data_ptr()),
        reinterpret_cast<const uint8_t*>(Bq.data_ptr()),
        reinterpret_cast<const uint8_t*>(Bsc.data_ptr()),
        reinterpret_cast<uint16_t*>(out.data_ptr()),
        M, (int)N, K, (int)A.stride(0),
        (int)(Bq.stride(0) * Bq.element_size()),
        (int)Bsc.size(1));
    return out;
}
"""


@functools.lru_cache(maxsize=1)
def _hip_module():
    digest = hashlib.sha1(HIP_SRC.encode("utf-8")).hexdigest()[:16]
    return load_inline(
        name=f"hipquant_{digest}",
        cpp_sources=[CPP_SRC],
        cuda_sources=[HIP_SRC],
        functions=["quantize_a_hip", "gemm_fused_16x16_hip"],
        extra_cflags=["-O3"],
        extra_cuda_cflags=["-O3", "--offload-arch=gfx950", "-std=c++17"],
        with_cuda=True,
        verbose=False,
    )


# ============================================================================
# Inline MXFP4 quantization op (for fused kernel path)
# ============================================================================

@triton.jit
def _mxfp4_quant_inline(
    x,
    BLOCK_SIZE_K: tl.constexpr,
    BLOCK_SIZE_M: tl.constexpr,
    MXFP4_QUANT_BLOCK_SIZE: tl.constexpr,
):
    EXP_BIAS_FP32: tl.constexpr = 127
    EXP_BIAS_FP4: tl.constexpr = 1
    MBITS_F32: tl.constexpr = 23
    MBITS_FP4: tl.constexpr = 1
    EBITS_F32: tl.constexpr = 8
    EBITS_FP4: tl.constexpr = 2

    max_normal: tl.constexpr = 6
    min_normal: tl.constexpr = 1

    NUM_QUANT_BLOCKS: tl.constexpr = BLOCK_SIZE_K // MXFP4_QUANT_BLOCK_SIZE
    x = x.reshape(BLOCK_SIZE_M, NUM_QUANT_BLOCKS, MXFP4_QUANT_BLOCK_SIZE)

    amax = tl.max(tl.abs(x), axis=-1, keep_dims=True)
    amax = amax.to(tl.int32, bitcast=True)
    amax = (amax + 0x200000).to(tl.uint32, bitcast=True) & 0xFF800000
    amax = amax.to(tl.float32, bitcast=True)
    scale_e8m0_unbiased = tl.log2(amax).floor() - 2
    scale_e8m0_unbiased = tl.clamp(scale_e8m0_unbiased, min=-127, max=127)

    bs_e8m0 = scale_e8m0_unbiased.to(tl.uint8) + 127
    quant_scale = tl.exp2(-scale_e8m0_unbiased)

    qx = x * quant_scale
    qx = qx.to(tl.uint32, bitcast=True)

    s = qx & 0x80000000
    qx = qx ^ s

    qx_fp32 = qx.to(tl.float32, bitcast=True)
    saturate_mask = qx_fp32 >= max_normal
    denormal_mask = (not saturate_mask) & (qx_fp32 < min_normal)
    normal_mask = not (saturate_mask | denormal_mask)

    denorm_exp: tl.constexpr = (EXP_BIAS_FP32 - EXP_BIAS_FP4) + (MBITS_F32 - MBITS_FP4) + 1
    denorm_mask_int: tl.constexpr = denorm_exp << MBITS_F32
    denorm_mask_float: tl.constexpr = tl.cast(denorm_mask_int, tl.float32, bitcast=True)

    denormal_x = qx_fp32 + denorm_mask_float
    denormal_x = denormal_x.to(tl.uint32, bitcast=True)
    denormal_x -= denorm_mask_int
    denormal_x = denormal_x.to(tl.uint8)

    normal_x = qx
    mant_odd = (normal_x >> (MBITS_F32 - MBITS_FP4)) & 1
    val_to_add = ((EXP_BIAS_FP4 - EXP_BIAS_FP32) << MBITS_F32) + (1 << 21) - 1
    normal_x += val_to_add
    normal_x += mant_odd
    normal_x = normal_x >> (MBITS_F32 - MBITS_FP4)
    normal_x = normal_x.to(tl.uint8)

    e2m1_value = tl.full(qx.type.get_block_shapes(), 0x7, dtype=tl.uint8)
    e2m1_value = tl.where(normal_mask, normal_x, e2m1_value)
    e2m1_value = tl.where(denormal_mask, denormal_x, e2m1_value)

    sign_lp = s >> (MBITS_F32 + EBITS_F32 - MBITS_FP4 - EBITS_FP4)
    sign_lp = sign_lp.to(tl.uint8)
    e2m1_value = e2m1_value | sign_lp

    e2m1_value = tl.reshape(
        e2m1_value, [BLOCK_SIZE_M, NUM_QUANT_BLOCKS, MXFP4_QUANT_BLOCK_SIZE // 2, 2]
    )
    evens, odds = tl.split(e2m1_value)
    x_fp4 = evens | (odds << 4)
    x_fp4 = x_fp4.reshape(BLOCK_SIZE_M, BLOCK_SIZE_K // 2)

    return x_fp4, bs_e8m0.reshape(BLOCK_SIZE_M, NUM_QUANT_BLOCKS)


# ============================================================================
# Triton helpers
# ============================================================================

@triton.jit
def remap_xcd(pid, GRID_MN, NUM_XCDS: tl.constexpr = 8):
    pids_per_xcd = (GRID_MN + NUM_XCDS - 1) // NUM_XCDS
    tall_xcds = GRID_MN % NUM_XCDS
    tall_xcds = NUM_XCDS if tall_xcds == 0 else tall_xcds
    xcd = pid % NUM_XCDS
    local_pid = pid // NUM_XCDS
    if xcd < tall_xcds:
        pid = xcd * pids_per_xcd + local_pid
    else:
        pid = (
            tall_xcds * pids_per_xcd
            + (xcd - tall_xcds) * (pids_per_xcd - 1)
            + local_pid
        )
    return pid


@triton.jit
def pid_grid(pid: int, num_pid_m: int, num_pid_n: int, GROUP_SIZE_M: tl.constexpr = 1):
    if GROUP_SIZE_M == 1:
        pid_m = pid // num_pid_n
        pid_n = pid % num_pid_n
    else:
        num_pid_in_group = GROUP_SIZE_M * num_pid_n
        group_id = pid // num_pid_in_group
        first_pid_m = group_id * GROUP_SIZE_M
        group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
        tl.assume(group_size_m >= 0)
        pid_m = first_pid_m + (pid % group_size_m)
        pid_n = (pid % num_pid_in_group) // group_size_m
    return pid_m, pid_n


@triton.jit
def _shuffled_b_scale_offset(row, col, Ks_stride):
    """Compute flat byte offset into shuffled E8M0 B_scale_sh for logical (row, col).

    The e8m0_shuffle layout views [N, K_scale] as [N//32, K_scale//8, 4, 16, 2, 2]
    with permutation (0,2,4,1,3,5) applied (shuffle) / (0,5,3,1,4,2) to undo.

    Given logical unshuffled (row, col), the shuffled flat offset is:
      s0=row//32, s1=col//8, s2=col%4, s3=row%16, s4=(col//4)%2, s5=(row//16)%2
      flat = s0*(Ks*32) + s1*256 + s2*64 + s3*4 + s4*2 + s5

    Ks_stride = B_scale_sh.shape[1] (= K_scale, must be multiple of 8).
    """
    return (
        (row // 32) * (Ks_stride * 32)
        + (col // 8) * 256
        + (col % 4) * 64
        + (row % 16) * 4
        + ((col // 4) % 2) * 2
        + ((row // 16) % 2)
    )


# ============================================================================
# Fused Quant+GEMM kernel (for small M, loads B_scale from shuffled layout)
# ============================================================================

@triton.heuristics(
    {
        "EVEN_K": lambda args: (args["K"] % (args["BLOCK_SIZE_K"] // 2) == 0)
        and (args["SPLITK_BLOCK_SIZE"] % args["BLOCK_SIZE_K"] == 0)
        and (args["K"] % (args["SPLITK_BLOCK_SIZE"] // 2) == 0),
    }
)
@triton.jit
def _fused_gemm_fp4_kernel(
    a_bf16_ptr, b_ptr, c_ptr,
    b_scales_ptr,
    M, N, K,
    actual_K,
    Ks_stride,
    stride_am, stride_ak,
    stride_bk, stride_bn,
    stride_ck, stride_cm, stride_cn,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
    NUM_KSPLIT: tl.constexpr,
    SPLITK_BLOCK_SIZE: tl.constexpr,
    EVEN_K: tl.constexpr,
    num_warps: tl.constexpr,
    num_stages: tl.constexpr,
    waves_per_eu: tl.constexpr,
    matrix_instr_nonkdim: tl.constexpr,
):
    tl.assume(stride_am > 0)
    tl.assume(stride_ak > 0)
    tl.assume(stride_bk > 0)
    tl.assume(stride_bn > 0)
    tl.assume(stride_cm > 0)
    tl.assume(stride_cn > 0)

    GRID_MN = tl.cdiv(M, BLOCK_SIZE_M) * tl.cdiv(N, BLOCK_SIZE_N)
    pid_unified = tl.program_id(axis=0)
    pid_unified = remap_xcd(pid_unified, GRID_MN * NUM_KSPLIT, NUM_XCDS=8)

    pid_k = pid_unified % NUM_KSPLIT
    pid = pid_unified // NUM_KSPLIT
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)

    if NUM_KSPLIT == 1:
        pid_m, pid_n = pid_grid(pid, num_pid_m, num_pid_n, GROUP_SIZE_M=GROUP_SIZE_M)
    else:
        pid_m = pid // num_pid_n
        pid_n = pid % num_pid_n

    tl.assume(pid_m >= 0)
    tl.assume(pid_n >= 0)

    SCALE_GROUP_SIZE: tl.constexpr = 32
    SCALES_PER_KBLOCK: tl.constexpr = BLOCK_SIZE_K // SCALE_GROUP_SIZE

    if (pid_k * SPLITK_BLOCK_SIZE // 2) < K:
        num_k_iter = tl.cdiv(SPLITK_BLOCK_SIZE // 2, BLOCK_SIZE_K // 2)

        offs_k_packed = tl.arange(0, BLOCK_SIZE_K // 2)
        offs_k_packed_split = pid_k * (SPLITK_BLOCK_SIZE // 2) + offs_k_packed

        offs_k_actual = tl.arange(0, BLOCK_SIZE_K)
        offs_k_actual_split = pid_k * SPLITK_BLOCK_SIZE + offs_k_actual

        offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
        offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N

        a_bf16_ptrs = a_bf16_ptr + (offs_am[:, None] * stride_am + offs_k_actual_split[None, :] * stride_ak)
        b_ptrs = b_ptr + (offs_k_packed_split[:, None] * stride_bk + offs_bn[None, :] * stride_bn)

        # Shuffled B scale loading
        ks_base = pid_k * (SPLITK_BLOCK_SIZE // SCALE_GROUP_SIZE)
        offs_ks_local = tl.arange(0, SCALES_PER_KBLOCK)

        accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)

        for k_iter in range(pid_k * num_k_iter, (pid_k + 1) * num_k_iter):
            if EVEN_K:
                a_bf16 = tl.load(a_bf16_ptrs)
            else:
                a_bf16 = tl.load(a_bf16_ptrs, mask=offs_k_actual[None, :] < actual_K - k_iter * BLOCK_SIZE_K, other=0.0)

            a_f32 = a_bf16.to(tl.float32)
            a_fp4, a_scales = _mxfp4_quant_inline(a_f32, BLOCK_SIZE_K, BLOCK_SIZE_M, SCALE_GROUP_SIZE)

            cur_offs_ks = ks_base + offs_ks_local
            b_scale_offsets = _shuffled_b_scale_offset(offs_bn[:, None], cur_offs_ks[None, :], Ks_stride)
            b_scales = tl.load(b_scales_ptr + b_scale_offsets, cache_modifier=".cg")

            if EVEN_K:
                b = tl.load(b_ptrs, cache_modifier=".cg")
            else:
                b = tl.load(b_ptrs, mask=offs_k_packed[:, None] < K - k_iter * (BLOCK_SIZE_K // 2), other=0, cache_modifier=".cg")

            accumulator = tl.dot_scaled(a_fp4, a_scales, "e2m1", b, b_scales, "e2m1", accumulator)

            a_bf16_ptrs += BLOCK_SIZE_K * stride_ak
            b_ptrs += (BLOCK_SIZE_K // 2) * stride_bk
            ks_base += SCALES_PER_KBLOCK

        c = accumulator.to(c_ptr.type.element_ty)

        offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
        offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)
        c_ptrs = (
            c_ptr
            + stride_cm * offs_cm[:, None]
            + stride_cn * offs_cn[None, :]
            + pid_k * stride_ck
        )
        c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
        tl.store(c_ptrs, c, mask=c_mask)


# ============================================================================
# Standard GEMM kernel (for larger M, loads B_scale from shuffled layout)
# ============================================================================

@triton.heuristics(
    {
        "EVEN_K": lambda args: (args["K"] % (args["BLOCK_SIZE_K"] // 2) == 0)
        and (args["SPLITK_BLOCK_SIZE"] % args["BLOCK_SIZE_K"] == 0)
        and (args["K"] % (args["SPLITK_BLOCK_SIZE"] // 2) == 0),
    }
)
@triton.jit
def _gemm_fp4_kernel(
    a_ptr, b_ptr, c_ptr,
    a_scales_ptr, b_scales_ptr,
    M, N, K,
    Ks_stride,
    stride_am, stride_ak,
    stride_bk, stride_bn,
    stride_ck, stride_cm, stride_cn,
    stride_asm, stride_ask,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
    NUM_KSPLIT: tl.constexpr,
    SPLITK_BLOCK_SIZE: tl.constexpr,
    EVEN_K: tl.constexpr,
    num_warps: tl.constexpr,
    num_stages: tl.constexpr,
    waves_per_eu: tl.constexpr,
    matrix_instr_nonkdim: tl.constexpr,
):
    tl.assume(stride_am > 0)
    tl.assume(stride_ak > 0)
    tl.assume(stride_bk > 0)
    tl.assume(stride_bn > 0)
    tl.assume(stride_cm > 0)
    tl.assume(stride_cn > 0)
    tl.assume(stride_asm > 0)
    tl.assume(stride_ask > 0)

    GRID_MN = tl.cdiv(M, BLOCK_SIZE_M) * tl.cdiv(N, BLOCK_SIZE_N)
    pid_unified = tl.program_id(axis=0)
    pid_unified = remap_xcd(pid_unified, GRID_MN * NUM_KSPLIT, NUM_XCDS=8)

    pid_k = pid_unified % NUM_KSPLIT
    pid = pid_unified // NUM_KSPLIT
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)

    if NUM_KSPLIT == 1:
        pid_m, pid_n = pid_grid(pid, num_pid_m, num_pid_n, GROUP_SIZE_M=GROUP_SIZE_M)
    else:
        pid_m = pid // num_pid_n
        pid_n = pid % num_pid_n

    tl.assume(pid_m >= 0)
    tl.assume(pid_n >= 0)

    SCALE_GROUP_SIZE: tl.constexpr = 32
    SCALES_PER_KBLOCK: tl.constexpr = BLOCK_SIZE_K // SCALE_GROUP_SIZE

    if (pid_k * SPLITK_BLOCK_SIZE // 2) < K:
        num_k_iter = tl.cdiv(SPLITK_BLOCK_SIZE // 2, BLOCK_SIZE_K // 2)

        offs_k = tl.arange(0, BLOCK_SIZE_K // 2)
        offs_k_split = pid_k * (SPLITK_BLOCK_SIZE // 2) + offs_k
        offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
        offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N

        a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k_split[None, :] * stride_ak)
        b_ptrs = b_ptr + (offs_k_split[:, None] * stride_bk + offs_bn[None, :] * stride_bn)

        # A scale: standard stride-based loading (A_scale is unshuffled)
        offs_a_ks = (pid_k * (SPLITK_BLOCK_SIZE // SCALE_GROUP_SIZE)) + tl.arange(0, SCALES_PER_KBLOCK)
        a_scale_ptrs = a_scales_ptr + offs_am[:, None] * stride_asm + offs_a_ks[None, :] * stride_ask

        # B scale: shuffled offset loading (no unshuffle needed)
        ks_base = pid_k * (SPLITK_BLOCK_SIZE // SCALE_GROUP_SIZE)
        offs_ks_local = tl.arange(0, SCALES_PER_KBLOCK)

        accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)

        for k in range(pid_k * num_k_iter, (pid_k + 1) * num_k_iter):
            a_scales = tl.load(a_scale_ptrs)

            cur_offs_ks = ks_base + offs_ks_local
            b_scale_offsets = _shuffled_b_scale_offset(offs_bn[:, None], cur_offs_ks[None, :], Ks_stride)
            b_scales = tl.load(b_scales_ptr + b_scale_offsets, cache_modifier=".cg")

            if EVEN_K:
                a = tl.load(a_ptrs)
                b = tl.load(b_ptrs, cache_modifier=".cg")
            else:
                a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * (BLOCK_SIZE_K // 2), other=0)
                b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * (BLOCK_SIZE_K // 2), other=0, cache_modifier=".cg")

            accumulator = tl.dot_scaled(a, a_scales, "e2m1", b, b_scales, "e2m1", accumulator)

            a_ptrs += (BLOCK_SIZE_K // 2) * stride_ak
            b_ptrs += (BLOCK_SIZE_K // 2) * stride_bk
            a_scale_ptrs += SCALES_PER_KBLOCK * stride_ask
            ks_base += SCALES_PER_KBLOCK

        c = accumulator.to(c_ptr.type.element_ty)

        offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
        offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)
        c_ptrs = (
            c_ptr
            + stride_cm * offs_cm[:, None]
            + stride_cn * offs_cn[None, :]
            + pid_k * stride_ck
        )
        c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
        tl.store(c_ptrs, c, mask=c_mask)


@triton.jit
def _reduce_kernel(
    c_in_ptr, c_out_ptr,
    M, N,
    stride_c_in_k, stride_c_in_m, stride_c_in_n,
    stride_c_out_m, stride_c_out_n,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    ACTUAL_KSPLIT: tl.constexpr,
    MAX_KSPLIT: tl.constexpr,
):
    pid_m = tl.program_id(axis=0)
    pid_n = tl.program_id(axis=1)

    offs_m = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
    offs_n = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
    offs_k = tl.arange(0, MAX_KSPLIT)

    c_in_ptrs = (
        c_in_ptr
        + (offs_k[:, None, None] * stride_c_in_k)
        + (offs_m[None, :, None] * stride_c_in_m)
        + (offs_n[None, None, :] * stride_c_in_n)
    )

    if ACTUAL_KSPLIT == MAX_KSPLIT:
        c = tl.load(c_in_ptrs)
    else:
        c = tl.load(c_in_ptrs, mask=offs_k[:, None, None] < ACTUAL_KSPLIT)
    c = tl.sum(c, axis=0)
    c = c.to(c_out_ptr.type.element_ty)

    c_out_ptrs = (
        c_out_ptr
        + (offs_m[:, None] * stride_c_out_m)
        + (offs_n[None, :] * stride_c_out_n)
    )
    tl.store(c_out_ptrs, c)


# ============================================================================
# Config selection (from tuned v14)
# ============================================================================

SHAPE_CONFIGS = {
    (4, 2880, 512): dict(BLOCK_SIZE_M=16, BLOCK_SIZE_N=64, BLOCK_SIZE_K=512, GROUP_SIZE_M=1,
                         num_warps=4, num_stages=1, waves_per_eu=4, matrix_instr_nonkdim=16, NUM_KSPLIT=1),
    (16, 2112, 7168): dict(BLOCK_SIZE_M=16, BLOCK_SIZE_N=64, BLOCK_SIZE_K=512, GROUP_SIZE_M=1,
                           num_warps=4, num_stages=1, waves_per_eu=2, matrix_instr_nonkdim=16, NUM_KSPLIT=8),
    (32, 4096, 512): dict(BLOCK_SIZE_M=32, BLOCK_SIZE_N=64, BLOCK_SIZE_K=512, GROUP_SIZE_M=1,
                          num_warps=4, num_stages=1, waves_per_eu=4, matrix_instr_nonkdim=16, NUM_KSPLIT=1),
    (32, 2880, 512): dict(BLOCK_SIZE_M=32, BLOCK_SIZE_N=64, BLOCK_SIZE_K=512, GROUP_SIZE_M=1,
                          num_warps=4, num_stages=1, waves_per_eu=4, matrix_instr_nonkdim=16, NUM_KSPLIT=1),
    (64, 7168, 2048): dict(BLOCK_SIZE_M=32, BLOCK_SIZE_N=64, BLOCK_SIZE_K=512, GROUP_SIZE_M=4,
                           num_warps=4, num_stages=5, waves_per_eu=1, matrix_instr_nonkdim=16, NUM_KSPLIT=1),
    (256, 3072, 1536): dict(BLOCK_SIZE_M=64, BLOCK_SIZE_N=64, BLOCK_SIZE_K=512, GROUP_SIZE_M=4,
                            num_warps=8, num_stages=4, waves_per_eu=2, matrix_instr_nonkdim=16, NUM_KSPLIT=1),
}

CONFIGS = {
    4: dict(BLOCK_SIZE_M=16, BLOCK_SIZE_N=128, BLOCK_SIZE_K=256, GROUP_SIZE_M=1,
            num_warps=4, num_stages=2, waves_per_eu=3, matrix_instr_nonkdim=16, NUM_KSPLIT=16),
    16: dict(BLOCK_SIZE_M=16, BLOCK_SIZE_N=128, BLOCK_SIZE_K=256, GROUP_SIZE_M=1,
             num_warps=4, num_stages=2, waves_per_eu=3, matrix_instr_nonkdim=16, NUM_KSPLIT=16),
    32: dict(BLOCK_SIZE_M=32, BLOCK_SIZE_N=128, BLOCK_SIZE_K=256, GROUP_SIZE_M=1,
             num_warps=4, num_stages=2, waves_per_eu=3, matrix_instr_nonkdim=16, NUM_KSPLIT=1),
    64: dict(BLOCK_SIZE_M=64, BLOCK_SIZE_N=256, BLOCK_SIZE_K=256, GROUP_SIZE_M=1,
             num_warps=4, num_stages=3, waves_per_eu=2, matrix_instr_nonkdim=32, NUM_KSPLIT=1),
    256: dict(BLOCK_SIZE_M=128, BLOCK_SIZE_N=256, BLOCK_SIZE_K=256, GROUP_SIZE_M=2,
              num_warps=4, num_stages=3, waves_per_eu=2, matrix_instr_nonkdim=32, NUM_KSPLIT=1),
}

def get_config(M, N=None, K=None):
    if N is not None and K is not None:
        key = (M, N, K)
        if key in SHAPE_CONFIGS:
            return SHAPE_CONFIGS[key].copy()
    for threshold in sorted(CONFIGS.keys()):
        if M <= threshold:
            return CONFIGS[threshold].copy()
    return CONFIGS[256].copy()


def get_splitk(K, BLOCK_SIZE_K, NUM_KSPLIT):
    NUM_KSPLIT_STEP = 2
    BLOCK_SIZE_K_STEP = 2
    SPLITK_BLOCK_SIZE = (
        triton.cdiv((2 * triton.cdiv(K, NUM_KSPLIT)), BLOCK_SIZE_K) * BLOCK_SIZE_K
    )
    while NUM_KSPLIT > 1 and BLOCK_SIZE_K > 16:
        if (
            K % (SPLITK_BLOCK_SIZE // 2) == 0
            and SPLITK_BLOCK_SIZE % BLOCK_SIZE_K == 0
            and K % (BLOCK_SIZE_K // 2) == 0
        ):
            break
        elif K % (SPLITK_BLOCK_SIZE // 2) != 0 and NUM_KSPLIT > 1:
            NUM_KSPLIT = NUM_KSPLIT // NUM_KSPLIT_STEP
        elif SPLITK_BLOCK_SIZE % BLOCK_SIZE_K != 0:
            if NUM_KSPLIT > 1:
                NUM_KSPLIT = NUM_KSPLIT // NUM_KSPLIT_STEP
            elif BLOCK_SIZE_K > 16:
                BLOCK_SIZE_K = BLOCK_SIZE_K // BLOCK_SIZE_K_STEP
        elif K % (BLOCK_SIZE_K // 2) != 0 and BLOCK_SIZE_K > 16:
            BLOCK_SIZE_K = BLOCK_SIZE_K // BLOCK_SIZE_K_STEP
        else:
            break
        SPLITK_BLOCK_SIZE = (
            triton.cdiv((2 * triton.cdiv(K, NUM_KSPLIT)), BLOCK_SIZE_K) * BLOCK_SIZE_K
        )
    NUM_KSPLIT = triton.cdiv(K, (SPLITK_BLOCK_SIZE // 2))
    return SPLITK_BLOCK_SIZE, BLOCK_SIZE_K, NUM_KSPLIT


# ============================================================================
# Safe caches (shape-keyed only, no input-derived data)
# ============================================================================
_OUT_CACHE: dict = {}
_AQ_CACHE: dict = {}
_HIP16_OUT_CACHE: dict = {}

FUSED_M_THRESHOLD = 32
HIP16_THRESHOLD_K = 1024  # Use HIP 16x16 for M<=32 and K<=this


def _get_aq_buffers(m, k, device):
    key = (m, k, device)
    cached = _AQ_CACHE.get(key)
    if cached is not None:
        return cached
    a_q = torch.empty((m, k // 2), dtype=torch.uint8, device=device)
    a_scale = torch.empty((m, k // SCALE_GROUP_SIZE), dtype=torch.uint8, device=device)
    _AQ_CACHE[key] = (a_q, a_scale)
    return a_q, a_scale


def custom_kernel(data: input_t) -> output_t:
    A, B, B_q, B_shuffle, B_scale_sh = data
    A = A.contiguous()
    m, k = A.shape
    n = B_q.shape[0]
    K_packed = k // 2
    k_scale = k // SCALE_GROUP_SIZE

    # HIP 16x16 MFMA path for small M with small K (fused quant+GEMM, single kernel)
    if m <= 32 and k <= HIP16_THRESHOLD_K:
        key16 = (m, n, A.device)
        if key16 not in _HIP16_OUT_CACHE:
            _HIP16_OUT_CACHE[key16] = torch.empty((m, n), dtype=torch.bfloat16, device=A.device)
        out16 = _HIP16_OUT_CACHE[key16]
        B_q_u8 = B_q.view(torch.uint8) if B_q.dtype != torch.uint8 else B_q
        B_scale_sh_u8 = B_scale_sh.view(torch.uint8) if B_scale_sh.dtype != torch.uint8 else B_scale_sh
        return _hip_module().gemm_fused_16x16_hip(A, B_q_u8, B_scale_sh_u8, n, out16)

    # B_q transpose: free view (no copy, no allocation)
    B_q_u8 = B_q.view(torch.uint8) if B_q.dtype != torch.uint8 else B_q
    B_t = B_q_u8.T

    # B_scale_sh: pass directly to kernel (no unshuffle!)
    B_scale_sh_u8 = B_scale_sh.view(torch.uint8) if B_scale_sh.dtype != torch.uint8 else B_scale_sh
    Ks_stride = B_scale_sh_u8.shape[1]

    use_fused = (m <= FUSED_M_THRESHOLD)

    if not use_fused:
        A_q, A_scale = _get_aq_buffers(m, k, A.device)
        _hip_module().quantize_a_hip(A, A_q, A_scale)

    config = get_config(m, n, k)

    if config["NUM_KSPLIT"] > 1:
        SPLITK_BLOCK_SIZE, BLOCK_SIZE_K, NUM_KSPLIT = get_splitk(
            K_packed, config["BLOCK_SIZE_K"], config["NUM_KSPLIT"]
        )
        config["SPLITK_BLOCK_SIZE"] = SPLITK_BLOCK_SIZE
        config["BLOCK_SIZE_K"] = BLOCK_SIZE_K
        config["NUM_KSPLIT"] = NUM_KSPLIT
    else:
        config["SPLITK_BLOCK_SIZE"] = 2 * K_packed

    if config["BLOCK_SIZE_K"] >= 2 * K_packed:
        config["BLOCK_SIZE_K"] = triton.next_power_of_2(2 * K_packed)
        config["SPLITK_BLOCK_SIZE"] = 2 * K_packed
        config["NUM_KSPLIT"] = 1

    config["BLOCK_SIZE_K"] = max(config["BLOCK_SIZE_K"], 128)

    NUM_KSPLIT = config["NUM_KSPLIT"]

    # Allocate output (shape-keyed cache, safe)
    out_key = (m, n, NUM_KSPLIT, A.device)
    cached_out = _OUT_CACHE.get(out_key)
    if cached_out is not None:
        if NUM_KSPLIT > 1:
            y_pp, y = cached_out
        else:
            y = cached_out
            y_pp = None
    elif NUM_KSPLIT > 1:
        y_pp = torch.empty((NUM_KSPLIT, m, n), dtype=torch.float32, device=A.device)
        y = torch.empty((m, n), dtype=torch.bfloat16, device=A.device)
        _OUT_CACHE[out_key] = (y_pp, y)
    else:
        y = torch.empty((m, n), dtype=torch.bfloat16, device=A.device)
        _OUT_CACHE[out_key] = y
        y_pp = None

    out_tensor = y if NUM_KSPLIT == 1 else y_pp

    grid = lambda META: (
        META["NUM_KSPLIT"]
        * triton.cdiv(m, META["BLOCK_SIZE_M"])
        * triton.cdiv(n, META["BLOCK_SIZE_N"]),
    )

    if use_fused:
        _fused_gemm_fp4_kernel[grid](
            A, B_t, out_tensor,
            B_scale_sh_u8,
            m, n, K_packed,
            k,
            Ks_stride,
            A.stride(0), A.stride(1),
            B_t.stride(0), B_t.stride(1),
            0 if NUM_KSPLIT == 1 else y_pp.stride(0),
            out_tensor.stride(-2), out_tensor.stride(-1),
            SPLITK_BLOCK_SIZE=config["SPLITK_BLOCK_SIZE"],
            BLOCK_SIZE_M=config["BLOCK_SIZE_M"],
            BLOCK_SIZE_N=config["BLOCK_SIZE_N"],
            BLOCK_SIZE_K=config["BLOCK_SIZE_K"],
            GROUP_SIZE_M=config["GROUP_SIZE_M"],
            NUM_KSPLIT=NUM_KSPLIT,
            num_warps=config["num_warps"],
            num_stages=config["num_stages"],
            waves_per_eu=config["waves_per_eu"],
            matrix_instr_nonkdim=config["matrix_instr_nonkdim"],
        )
    else:
        _gemm_fp4_kernel[grid](
            A_q, B_t, out_tensor,
            A_scale, B_scale_sh_u8,
            m, n, K_packed,
            Ks_stride,
            A_q.stride(0), A_q.stride(1),
            B_t.stride(0), B_t.stride(1),
            0 if NUM_KSPLIT == 1 else y_pp.stride(0),
            out_tensor.stride(-2), out_tensor.stride(-1),
            A_scale.stride(0), A_scale.stride(1),
            SPLITK_BLOCK_SIZE=config["SPLITK_BLOCK_SIZE"],
            BLOCK_SIZE_M=config["BLOCK_SIZE_M"],
            BLOCK_SIZE_N=config["BLOCK_SIZE_N"],
            BLOCK_SIZE_K=config["BLOCK_SIZE_K"],
            GROUP_SIZE_M=config["GROUP_SIZE_M"],
            NUM_KSPLIT=NUM_KSPLIT,
            num_warps=config["num_warps"],
            num_stages=config["num_stages"],
            waves_per_eu=config["waves_per_eu"],
            matrix_instr_nonkdim=config["matrix_instr_nonkdim"],
        )

    if NUM_KSPLIT > 1:
        REDUCE_BLOCK_SIZE_M = 16
        REDUCE_BLOCK_SIZE_N = 32
        ACTUAL_KSPLIT = triton.cdiv(K_packed, (config["SPLITK_BLOCK_SIZE"] // 2))

        grid_reduce = (
            triton.cdiv(m, REDUCE_BLOCK_SIZE_M),
            triton.cdiv(n, REDUCE_BLOCK_SIZE_N),
        )
        _reduce_kernel[grid_reduce](
            y_pp, y,
            m, n,
            y_pp.stride(0), y_pp.stride(1), y_pp.stride(2),
            y.stride(0), y.stride(1),
            BLOCK_SIZE_M=REDUCE_BLOCK_SIZE_M,
            BLOCK_SIZE_N=REDUCE_BLOCK_SIZE_N,
            ACTUAL_KSPLIT=ACTUAL_KSPLIT,
            MAX_KSPLIT=triton.next_power_of_2(NUM_KSPLIT),
        )

    return y