submission 754225

"""
MLA (Multi-head Latent Attention) decode kernel - optimized implementation.

Custom HIP kernel for fused MXFP4 attention (JIT compiled on gfx950).
MXFP4 Q + MXFP4 KV using MFMA FP4 intrinsics for QK^T (4x bandwidth savings).

DeepSeek R1 forward_absorb MLA config:
  total_num_heads  = 128    (query heads before TP split)
  num_heads        = 128 // tp  (query heads per device, tp=4 -> 32, tp=8 -> 16)
  num_kv_heads     = 1      (shared latent KV head)
  kv_lora_rank     = 512    (latent dim)
  qk_rope_head_dim = 64     (RoPE dim)
  qk_head_dim      = 576    (kv_lora_rank + qk_rope_head_dim, absorbed q/k dim)
  v_head_dim       = 512    (= kv_lora_rank, output dim)
  sm_scale         = 1/sqrt(576)

KV buffer format (forward_absorb):
  - Full 576 dims used as keys (for Q@K^T score computation)
  - First 512 dims (kv_lora_rank) used as values (for output computation)
"""

import torch
from task import input_t, output_t

# ---------------------------------------------------------------------------
# Embedded HIP Kernel for MXFP4 MLA Decode (using hip-python)
# ---------------------------------------------------------------------------

# HIP kernel source code - optimized for AMD MI355X (gfx950)
# Key optimizations:
# 1. Vectorized loads (float4/uint4) for coalesced memory access
# 2. LDS-cached FP4 LUT for faster dequantization
# 3. Better work distribution with 2D thread blocks
# 4. KV tiling to improve cache locality
# 5. Fused softmax with online normalization
# 6. Reduced shared memory bank conflicts
# 7. Loop unrolling for MXFP4 block processing
MLA_MXFP4_HIP_SOURCE = b'''
// v4 fused attn kernel
#include <hip/hip_runtime.h>
#include <hip/hip_fp16.h>
#include <hip/hip_bfloat16.h>
#include <cstdint>
#include <cmath>

// Constants
constexpr int QK_HEAD_DIM = 576;
constexpr int NUM_HEADS = 16;
constexpr int FMT_FP4_MFMA = 4;

typedef float __attribute__((ext_vector_type(4))) float4_t;
typedef float __attribute__((ext_vector_type(16))) float16_t;
typedef int __attribute__((ext_vector_type(8))) int8_vec;
typedef uint32_t __attribute__((ext_vector_type(4))) uint128_vec;

// Profiling macros
#define PROFILE_START(phase_var) \
    if constexpr (PROFILE_PHASES) { \
        if (is_timer_block && tid == 0) \
            phase_var = __builtin_amdgcn_s_memrealtime(); \
    }

#define PROFILE_ACCUM(accum_var, phase_var) \
    if constexpr (PROFILE_PHASES) { \
        if (is_timer_block && tid == 0) { \
            uint64_t t_now = __builtin_amdgcn_s_memrealtime(); \
            accum_var += t_now - phase_var; \
            phase_var = t_now; \
        } \
    }

// Fast E8M0 to float using bit reinterpretation
__device__ __forceinline__ float e8m0_to_float_fast(uint8_t e8m0) {
    // E8M0: pure exponent format, value = 2^(e8m0 - 127)
    // Construct IEEE754 float directly: exponent = e8m0, mantissa = 0
    uint32_t bits = (static_cast<uint32_t>(e8m0)) << 23;
    return __uint_as_float(bits);
}

struct QuantBlock {
    uint32_t data[4];  // 16 packed bytes (32 FP4 values)
    uint8_t e8m0;
};

// =============================================================================
// QK^T via MFMA: warp 0 computes query-key dot products for all heads
// =============================================================================
__device__ __forceinline__ int32_t mla_broadcast_scale(uint8_t e8m0);

template <typename Traits, int HEADS, int KV_SUBTILE, int NUM_BLOCKS,
          int A_K_HALF, int K_HALF, int B_SCALE_STRIDE, int WARPS, int KV_PAD = 0>
__device__ __forceinline__ void mla_qkt_mfma(
    const uint8_t q_data[][A_K_HALF],
    const uint8_t q_scale[],
    const uint8_t kv_data[][K_HALF + KV_PAD],
    const float kv_scale[][B_SCALE_STRIDE],
    float scores[][KV_SUBTILE],
    float sm_scale,
    int lane, int warp_id
) {
    using acc_t = typename Traits::acc_t;
    constexpr int IM = Traits::IM;
    constexpr int BPC = Traits::BPC;
    constexpr int K_ITERS = (NUM_BLOCKS + BPC - 1) / BPC;
    constexpr int QKT_BATCH = IM;
    constexpr int TOTAL_QKT_ITERS = KV_SUBTILE / QKT_BATCH;
    constexpr int ITERS_PER_WARP = (TOTAL_QKT_ITERS + WARPS - 1) / WARPS;
    const int my_qkt_start = warp_id * ITERS_PER_WARP;
    const int my_qkt_end = (my_qkt_start + ITERS_PER_WARP < TOTAL_QKT_ITERS) ? (my_qkt_start + ITERS_PER_WARP) : TOTAL_QKT_ITERS;

    for (int qkt_iter = my_qkt_start; qkt_iter < my_qkt_end; qkt_iter++) {
        acc_t mfma_acc = {};
        for (int ki = 0; ki < K_ITERS; ki++) {
            int blk0 = ki * BPC;
            int qa_row = lane % IM;
            int qa_blk = blk0 + lane / IM;
            uint32_t qa_reg[8] = {};
            if (qa_row < HEADS && qa_blk < NUM_BLOCKS) {
                *reinterpret_cast<uint128_vec*>(&qa_reg[0]) =
                    *reinterpret_cast<const uint128_vec*>(&q_data[qa_row][qa_blk * 16]);
            }
            uint8_t qa_e = 127;
            if (qa_row < HEADS && qa_blk < NUM_BLOCKS)
                qa_e = q_scale[qa_row + qa_blk * HEADS];
            else
                qa_reg[0] = qa_reg[1] = qa_reg[2] = qa_reg[3] = 0;
            int32_t qa_sc = mla_broadcast_scale(qa_e);

            int kb_row = lane % IM;
            int kb_abs = qkt_iter * QKT_BATCH + kb_row;
            int kb_blk = blk0 + lane / IM;
            uint32_t kb_reg[8] = {};
            if (kb_abs < KV_SUBTILE && kb_blk < NUM_BLOCKS) {
                // Row-major layout: kv_data[kv_pos][byte_offset]
                // Load 16 contiguous bytes for this KV position via vectorized load
                *reinterpret_cast<uint128_vec*>(&kb_reg[0]) =
                    *reinterpret_cast<const uint128_vec*>(&kv_data[kb_abs][kb_blk * 16]);
            }
            uint8_t kb_e = 127;
            if (kb_abs < KV_SUBTILE && kb_blk < NUM_BLOCKS)
                kb_e = (uint8_t)(__float_as_uint(kv_scale[kb_abs][kb_blk]) >> 23);
            else
                kb_reg[0] = kb_reg[1] = kb_reg[2] = kb_reg[3] = 0;
            int32_t kb_sc = mla_broadcast_scale(kb_e);

            int8_vec qa_vec = {(int)qa_reg[0],(int)qa_reg[1],(int)qa_reg[2],(int)qa_reg[3],
                               (int)qa_reg[4],(int)qa_reg[5],(int)qa_reg[6],(int)qa_reg[7]};
            int8_vec kb_vec = {(int)kb_reg[0],(int)kb_reg[1],(int)kb_reg[2],(int)kb_reg[3],
                               (int)kb_reg[4],(int)kb_reg[5],(int)kb_reg[6],(int)kb_reg[7]};

            mfma_acc = Traits::mfma(qa_vec, qa_sc, kb_vec, kb_sc, mfma_acc);
        }
        Traits::template store_scores<HEADS, KV_SUBTILE>(
            mfma_acc, lane, qkt_iter * QKT_BATCH,
            sm_scale, scores);
    }
}

// =============================================================================
// KV tile load: cooperative load of KV data + scales from HBM to LDS
// LDS layout: kv_buf[KV_SUBTILE][K_HALF] (KV position major, byte offset minor)
// Vectorized: 16-byte uint128_vec copies (K_HALF=288 = 18 * 16)
// =============================================================================
template <int KV_SUBTILE, int K_HALF, int B_SCALE_STRIDE, int BLOCK_SIZE, int KV_PAD = 0>
__device__ __forceinline__ void mla_load_kv_tile(
    uint8_t kv_buf[][K_HALF + KV_PAD],
    float kv_scale_buf[][B_SCALE_STRIDE],
    const uint8_t kv_src[][K_HALF],
    const uint8_t kv_scale_src[][B_SCALE_STRIDE],
    int kv_row_start, int tid
) {
    // Vectorized load: 16 bytes per iteration via uint128_vec
    constexpr int CHUNK_SIZE = 16;
    constexpr int CHUNKS_PER_ROW = K_HALF / CHUNK_SIZE;
    constexpr int TOTAL_CHUNKS = KV_SUBTILE * CHUNKS_PER_ROW;
    for (int i = tid; i < TOTAL_CHUNKS; i += BLOCK_SIZE) {
        const int row = i / CHUNKS_PER_ROW;
        const int chunk = i % CHUNKS_PER_ROW;
        const int col = chunk * CHUNK_SIZE;
        *reinterpret_cast<uint128_vec*>(&kv_buf[row][col]) =
            *reinterpret_cast<const uint128_vec*>(&kv_src[kv_row_start + row][col]);
    }
    // Handle remaining bytes if K_HALF is not divisible by 16
    constexpr int REM_BYTES = K_HALF % CHUNK_SIZE;
    if constexpr (REM_BYTES > 0) {
        constexpr int REM_START = CHUNKS_PER_ROW * CHUNK_SIZE;
        for (int i = tid; i < KV_SUBTILE * REM_BYTES; i += BLOCK_SIZE) {
            const int row = i / REM_BYTES;
            const int col = REM_START + i % REM_BYTES;
            kv_buf[row][col] = kv_src[kv_row_start + row][col];
        }
    }
    // Scale copy - B_SCALE_STRIDE is small (18 or 24), byte copy is fine
    constexpr int TOTAL_SC = KV_SUBTILE * B_SCALE_STRIDE;
    for (int i = tid; i < TOTAL_SC; i += BLOCK_SIZE) {
        const int row = i / B_SCALE_STRIDE;
        const int blk = i % B_SCALE_STRIDE;
        kv_scale_buf[row][blk] = e8m0_to_float_fast(kv_scale_src[kv_row_start + row][blk]);
    }
}

// =============================================================================
// Parallel softmax: all threads participate (16 per head, 8 elements each)
// 4 phases: parallel max -> reduce+rescale -> parallel exp+sum -> reduce sum
// =============================================================================
template <typename Traits, int HEADS, int KV_SUBTILE, int BLOCK_SIZE,
          int WARPS, int SM_ELEMS_PER_THREAD, int CHUNKS_PER_WARP>
__device__ __forceinline__ void mla_parallel_softmax(
    float scores[][KV_SUBTILE],
    float running_max[HEADS],
    float running_sum[HEADS],
    typename Traits::acc_t v_acc[CHUNKS_PER_WARP],
    int tid, int lane, int warp_id
) {
    __shared__ float rescale_lds[HEADS];
    __shared__ float new_max_lds[HEADS];
    __shared__ float softmax_scratch[HEADS][WARPS];

    const int sm_h = tid % HEADS;
    const int sm_group = tid / HEADS;
    const int sm_base = sm_group * SM_ELEMS_PER_THREAD;

    // Phase A: each thread finds max over its elements
    float local_max = -INFINITY;
    for (int j = 0; j < SM_ELEMS_PER_THREAD; j++) {
        const int ki = sm_base + j;
        if (ki < KV_SUBTILE)
            local_max = fmaxf(local_max, scores[sm_h][ki]);
    }

    // Warp-level reduction: 4 threads per head within each 64-lane wavefront
    local_max = fmaxf(local_max, __shfl_xor(local_max, 16));
    local_max = fmaxf(local_max, __shfl_xor(local_max, 32));
    if (lane < HEADS)
        softmax_scratch[sm_h][warp_id] = local_max;
    if constexpr (WARPS > 1) {
        __syncthreads();
    }

    // Phase B: master threads (0..15) reduce across warps, compute rescale
    if (tid < HEADS) {
        const int h = tid;
        float subtile_max = -INFINITY;
        for (int w = 0; w < WARPS; w++)
            subtile_max = fmaxf(subtile_max, softmax_scratch[h][w]);
        const float old_max = running_max[h];
        const float new_max = fmaxf(old_max, subtile_max);
        const float prev_rescale = __expf(old_max - new_max);
        rescale_lds[h] = prev_rescale;
        new_max_lds[h] = new_max;
        running_max[h] = new_max;
        running_sum[h] *= prev_rescale;
    }

    if constexpr (WARPS > 1) {
        __syncthreads();
    }

    // Phase C: all threads rescale v_acc, compute exp weights, partial sums
    Traits::template rescale_v_acc<HEADS, CHUNKS_PER_WARP>(
        v_acc, lane, rescale_lds);

    const float nm = new_max_lds[sm_h];
    float local_sum = 0.0f;
    for (int j = 0; j < SM_ELEMS_PER_THREAD; j++) {
        const int ki = sm_base + j;
        if (ki < KV_SUBTILE) {
            float w = __expf(scores[sm_h][ki] - nm);
            scores[sm_h][ki] = w;
            local_sum += w;
        }
    }
    // Warp-level sum reduction
    local_sum += __shfl_xor(local_sum, 16);
    local_sum += __shfl_xor(local_sum, 32);
    if (lane < HEADS)
        softmax_scratch[sm_h][warp_id] = local_sum;

    if constexpr (WARPS > 1) {
        __syncthreads();
    }

    // Phase D: master threads reduce sums across warps
    if (tid < HEADS) {
        const int h = tid;
        float subtile_sum = 0.0f;
        for (int w = 0; w < WARPS; w++)
            subtile_sum += softmax_scratch[h][w];
        running_sum[h] += subtile_sum;
    }
}

struct E8M0Scale {
    uint8_t e8m0;
    float quant_scale;
};

__device__ constexpr uint8_t E8M0_LUT[256] = {
    0,0,0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23,24,25,26,27,28,29,
    30,31,32,33,34,35,36,37,38,39,40,41,42,43,44,45,46,47,48,49,50,51,52,53,54,55,56,57,58,59,60,61,
    62,63,64,65,66,67,68,69,70,71,72,73,74,75,76,77,78,79,80,81,82,83,84,85,86,87,88,89,90,91,92,93,
    94,95,96,97,98,99,100,101,102,103,104,105,106,107,108,109,110,111,112,113,114,115,116,117,118,119,120,121,122,123,124,125,
    126,127,128,129,130,131,132,133,134,135,136,137,138,139,140,141,142,143,144,145,146,147,148,149,150,151,152,153,154,155,156,157,
    158,159,160,161,162,163,164,165,166,167,168,169,170,171,172,173,174,175,176,177,178,179,180,181,182,183,184,185,186,187,188,189,
    190,191,192,193,194,195,196,197,198,199,200,201,202,203,204,205,206,207,208,209,210,211,212,213,214,215,216,217,218,219,220,221,
    222,223,224,225,226,227,228,229,230,231,232,233,234,235,236,237,238,239,240,241,242,243,244,245,246,247,248,249,250,251,252,253,
};

__device__ constexpr uint32_t QUANT_SCALE_RECIP_LUT[256] = {
    0x00000000u,0x00800000u,0x01000000u,0x01800000u,0x02000000u,0x02800000u,0x03000000u,0x03800000u,
    0x04000000u,0x04800000u,0x05000000u,0x05800000u,0x06000000u,0x06800000u,0x07000000u,0x07800000u,
    0x08000000u,0x08800000u,0x09000000u,0x09800000u,0x0A000000u,0x0A800000u,0x0B000000u,0x0B800000u,
    0x0C000000u,0x0C800000u,0x0D000000u,0x0D800000u,0x0E000000u,0x0E800000u,0x0F000000u,0x0F800000u,
    0x10000000u,0x10800000u,0x11000000u,0x11800000u,0x12000000u,0x12800000u,0x13000000u,0x13800000u,
    0x14000000u,0x14800000u,0x15000000u,0x15800000u,0x16000000u,0x16800000u,0x17000000u,0x17800000u,
    0x18000000u,0x18800000u,0x19000000u,0x19800000u,0x1A000000u,0x1A800000u,0x1B000000u,0x1B800000u,
    0x1C000000u,0x1C800000u,0x1D000000u,0x1D800000u,0x1E000000u,0x1E800000u,0x1F000000u,0x1F800000u,
    0x20000000u,0x20800000u,0x21000000u,0x21800000u,0x22000000u,0x22800000u,0x23000000u,0x23800000u,
    0x24000000u,0x24800000u,0x25000000u,0x25800000u,0x26000000u,0x26800000u,0x27000000u,0x27800000u,
    0x28000000u,0x28800000u,0x29000000u,0x29800000u,0x2A000000u,0x2A800000u,0x2B000000u,0x2B800000u,
    0x2C000000u,0x2C800000u,0x2D000000u,0x2D800000u,0x2E000000u,0x2E800000u,0x2F000000u,0x2F800000u,
    0x30000000u,0x30800000u,0x31000000u,0x31800000u,0x32000000u,0x32800000u,0x33000000u,0x33800000u,
    0x34000000u,0x34800000u,0x35000000u,0x35800000u,0x36000000u,0x36800000u,0x37000000u,0x37800000u,
    0x38000000u,0x38800000u,0x39000000u,0x39800000u,0x3A000000u,0x3A800000u,0x3B000000u,0x3B800000u,
    0x3C000000u,0x3C800000u,0x3D000000u,0x3D800000u,0x3E000000u,0x3E800000u,0x3F000000u,0x3F800000u,
    0x40000000u,0x40800000u,0x41000000u,0x41800000u,0x42000000u,0x42800000u,0x43000000u,0x43800000u,
    0x44000000u,0x44800000u,0x45000000u,0x45800000u,0x46000000u,0x46800000u,0x47000000u,0x47800000u,
    0x48000000u,0x48800000u,0x49000000u,0x49800000u,0x4A000000u,0x4A800000u,0x4B000000u,0x4B800000u,
    0x4C000000u,0x4C800000u,0x4D000000u,0x4D800000u,0x4E000000u,0x4E800000u,0x4F000000u,0x4F800000u,
    0x50000000u,0x50800000u,0x51000000u,0x51800000u,0x52000000u,0x52800000u,0x53000000u,0x53800000u,
    0x54000000u,0x54800000u,0x55000000u,0x55800000u,0x56000000u,0x56800000u,0x57000000u,0x57800000u,
    0x58000000u,0x58800000u,0x59000000u,0x59800000u,0x5A000000u,0x5A800000u,0x5B000000u,0x5B800000u,
    0x5C000000u,0x5C800000u,0x5D000000u,0x5D800000u,0x5E000000u,0x5E800000u,0x5F000000u,0x5F800000u,
    0x60000000u,0x60800000u,0x61000000u,0x61800000u,0x62000000u,0x62800000u,0x63000000u,0x63800000u,
    0x64000000u,0x64800000u,0x65000000u,0x65800000u,0x66000000u,0x66800000u,0x67000000u,0x67800000u,
    0x68000000u,0x68800000u,0x69000000u,0x69800000u,0x6A000000u,0x6A800000u,0x6B000000u,0x6B800000u,
    0x6C000000u,0x6C800000u,0x6D000000u,0x6D800000u,0x6E000000u,0x6E800000u,0x6F000000u,0x6F800000u,
    0x70000000u,0x70800000u,0x71000000u,0x71800000u,0x72000000u,0x72800000u,0x73000000u,0x73800000u,
    0x74000000u,0x74800000u,0x75000000u,0x75800000u,0x76000000u,0x76800000u,0x77000000u,0x77800000u,
    0x78000000u,0x78800000u,0x79000000u,0x79800000u,0x7A000000u,0x7A800000u,0x7B000000u,0x7B800000u,
    0x7C000000u,0x7C800000u,0x7D000000u,0x7D800000u,0x7E000000u,0x7E800000u,0x7F000000u,0x00000000u,
};

__device__ __forceinline__ E8M0Scale compute_e8m0_scale(hip_bfloat16 amax_bf16) {
    E8M0Scale result;
    uint16_t amax_bits = amax_bf16.data;
    if (amax_bits == 0) {
        result.e8m0 = 0;
        result.quant_scale = 0.0f;
    } else {
        uint16_t rounded = (amax_bits + 0x0020u) & 0xFF80u;
        int raw_exp = (int)((rounded >> 7) & 0xFF);
        result.e8m0 = E8M0_LUT[raw_exp];
        result.quant_scale = __uint_as_float(QUANT_SCALE_RECIP_LUT[result.e8m0]);
    }
    return result;
}

// E8M0 scale computation from FP32 amax (avoids BF16 intermediate)
__device__ __forceinline__ E8M0Scale compute_e8m0_scale_f32(float amax) {
    E8M0Scale result;
    uint32_t amax_bits = __float_as_uint(amax);
    if ((amax_bits & 0x7FFFFFFFu) == 0) {
        result.e8m0 = 0;
        result.quant_scale = 0.0f;
    } else {
        // Round mantissa up to nearest power of 2 (same logic as BF16 but for FP32)
        // FP32: [sign(1)][exp(8)][mant(23)], round bit at position 22
        uint32_t rounded = (amax_bits + 0x00400000u) & 0xFF800000u;
        int raw_exp = (int)((rounded >> 23) & 0xFF);
        result.e8m0 = E8M0_LUT[raw_exp];
        result.quant_scale = __uint_as_float(QUANT_SCALE_RECIP_LUT[result.e8m0]);
    }
    return result;
}

// Pack 4 FP4 bytes into a uint32 using hardware byte permutation.
__device__ __forceinline__ uint32_t pack_fp4_to_u32(uint8_t p0, uint8_t p1, uint8_t p2, uint8_t p3) {
    uint32_t srcA = (uint32_t)p0 | ((uint32_t)p1 << 8);
    uint32_t srcB = (uint32_t)p2 | ((uint32_t)p3 << 8);
    return __builtin_amdgcn_perm(srcB, srcA, 0x05040100);
}

// Hardware FP4 conversion from BF16: converts 2 BF16 values to packed FP4 byte.
__device__ __forceinline__ uint8_t quantize_fp4_pair_hw_bf16(
    hip_bfloat16 v0, hip_bfloat16 v1, float quant_scale
) {
    using bf16x2 = uint16_t __attribute__((ext_vector_type(2)));
    bf16x2 pair = {v0.data, v1.data};
    union { uint32_t u32; uint8_t u8[4]; } cvt = {0};
    cvt.u32 = __builtin_amdgcn_cvt_scalef32_pk_fp4_bf16(cvt.u32, pair, quant_scale, 0);
    return cvt.u8[0];
}

// Hardware FP4 conversion from FP32: converts 2 FP32 values to packed FP4 byte.
__device__ __forceinline__ uint8_t quantize_fp4_pair_hw(float v0, float v1, float quant_scale) {
    union { uint32_t u32; uint8_t u8[4]; } cvt = {0};
    cvt.u32 = __builtin_amdgcn_cvt_scalef32_pk_fp4_f32(
        cvt.u32, v0, v1, quant_scale, 0);
    return cvt.u8[0];
}

// Quantize 32 BF16 values to packed FP4 + E8M0 scale using BF16 hw intrinsic.
__device__ __forceinline__ QuantBlock quantize_fp4_block_bf16(const hip_bfloat16* src) {
    uint16_t amax_bits = 0;
    for (int i = 0; i < 32; i++) {
        uint16_t bits = *reinterpret_cast<const uint16_t*>(&src[i]) & 0x7FFF;
        amax_bits = (bits > amax_bits) ? bits : amax_bits;
    }
    hip_bfloat16 amax_bf16 = *reinterpret_cast<const hip_bfloat16*>(&amax_bits);

    E8M0Scale sc = compute_e8m0_scale(amax_bf16);

    uint32_t pack[4];
    for (int j = 0; j < 4; j++) {
        int base = j << 3;
        uint8_t p0 = quantize_fp4_pair_hw_bf16(src[base], src[base+1], sc.quant_scale);
        uint8_t p1 = quantize_fp4_pair_hw_bf16(src[base+2], src[base+3], sc.quant_scale);
        uint8_t p2 = quantize_fp4_pair_hw_bf16(src[base+4], src[base+5], sc.quant_scale);
        uint8_t p3 = quantize_fp4_pair_hw_bf16(src[base+6], src[base+7], sc.quant_scale);
        pack[j] = pack_fp4_to_u32(p0, p1, p2, p3);
    }

    QuantBlock result;
    *reinterpret_cast<uint128_vec*>(&result.data) = *reinterpret_cast<uint128_vec*>(&pack);
    result.e8m0 = sc.e8m0;
    return result;
}

// Quantize 32 float values with pre-computed amax (skips amax search).
__device__ __forceinline__ QuantBlock quantize_fp4_block_with_scale(const float vals[32], E8M0Scale sc) {
    uint32_t pack[4];
    for (int j = 0; j < 4; j++) {
        int base = j << 3;
        uint8_t p0 = quantize_fp4_pair_hw(vals[base], vals[base+1], sc.quant_scale);
        uint8_t p1 = quantize_fp4_pair_hw(vals[base+2], vals[base+3], sc.quant_scale);
        uint8_t p2 = quantize_fp4_pair_hw(vals[base+4], vals[base+5], sc.quant_scale);
        uint8_t p3 = quantize_fp4_pair_hw(vals[base+6], vals[base+7], sc.quant_scale);
        pack[j] = pack_fp4_to_u32(p0, p1, p2, p3);
    }

    QuantBlock result;
    *reinterpret_cast<uint128_vec*>(&result.data) = *reinterpret_cast<uint128_vec*>(&pack);
    result.e8m0 = sc.e8m0;
    return result;
}

// =====================================================================
// MLA QK^T kernel: MFMA FP4, fp32 output, batched, linear B scales
// =====================================================================

__device__ __forceinline__ int32_t mla_broadcast_scale(uint8_t e8m0) {
    return (int32_t)e8m0 * 0x01010101;
}

// =====================================================================
// MLA MFMA traits: abstracts 16x16x128 vs 32x32x64 differences
// =====================================================================
template <bool USE_32x32>
struct MlaMfmaTraits {
    using acc_t = typename std::conditional<USE_32x32, float16_t, float4_t>::type;
    static constexpr int IM = USE_32x32 ? 32 : 16;
    static constexpr int IK = USE_32x32 ? 64 : 128;
    static constexpr int BPC = IK / 32;
    static constexpr int ACC_SIZE = USE_32x32 ? 16 : 4;
    static constexpr int VALID_ACC = USE_32x32 ? 8 : 4;

    static __device__ __forceinline__ acc_t mfma(
        int8_vec a, int32_t a_sc, int8_vec b, int32_t b_sc, acc_t acc
    ) {
        if constexpr (USE_32x32) {
            return __builtin_amdgcn_mfma_scale_f32_32x32x64_f8f6f4(
                a, b, acc, FMT_FP4_MFMA, FMT_FP4_MFMA, 0, a_sc, 0, b_sc);
        } else {
            return __builtin_amdgcn_mfma_scale_f32_16x16x128_f8f6f4(
                a, b, acc, FMT_FP4_MFMA, FMT_FP4_MFMA, 0, a_sc, 0, b_sc);
        }
    }

    // Map accumulator element index + lane to head index
    static __device__ __forceinline__ int acc_to_head(int i, int lane) {
        if constexpr (USE_32x32) { // i : [0, 16)
            int half = lane / 32;
            return (i % 4) + 4 * half + 8 * (i / 4);
        } else {  // i : [0, 4)
            int quad = lane / 16;
            return i + 4 * quad;
        }
    }

    // Map lane to output column (N-dimension)
    static __device__ __forceinline__ int lane_to_col(int lane) {
        return lane % IM;
    }

    // Store accumulator values to scores_lds with sm_scale
    template <int HEADS, int KV_SUBTILE>
    static __device__ __forceinline__ void store_scores(
        const acc_t& acc, int lane, int batch_col_offset,
        float sm_scale, float scores[][KV_SUBTILE]
    ) {
        int col = lane_to_col(lane);
        int eff_col = batch_col_offset + col;
        if (eff_col < KV_SUBTILE) {
            for (int i = 0; i < VALID_ACC; i++) {
                int head = acc_to_head(i, lane);
                scores[head][eff_col] = acc[i] * sm_scale;
            }
        }
    }

    // Rescale V accumulators by per-head factors
    template <int HEADS, int CHUNKS_PER_WARP>
    static __device__ __forceinline__ void rescale_v_acc(
        acc_t v_acc[], int lane, const float rescale[]
    ) {
        for (int ci = 0; ci < CHUNKS_PER_WARP; ci++) {
            for (int i = 0; i < VALID_ACC; i++) {
                int head = acc_to_head(i, lane);
                v_acc[ci][i] *= rescale[head];
            }
        }
    }

    // Store V accumulator to partial output
    template <int HEADS, int V_DIM, int CHUNKS_PER_WARP, int KV_SUBTILE>
    static __device__ __forceinline__ void store_v_output(
        const acc_t v_acc[], int lane, int warp_id,
        int out_row, float partial_out[][16 * 512],
        const float running_sum_lds[][KV_SUBTILE]
    ) {
        int col = lane_to_col(lane);
        for (int ci = 0; ci < CHUNKS_PER_WARP; ci++) {
            int chunk = warp_id * CHUNKS_PER_WARP + ci;
            for (int i = 0; i < VALID_ACC; i++) {
                int head = acc_to_head(i, lane);
                int v_dim = chunk * IM + col;
                if (v_dim < V_DIM) {
                    float rs = running_sum_lds[head][0];
                    float inv_sum = (rs > 0.0f) ? (1.0f / rs) : 0.0f;
                    partial_out[out_row][head * V_DIM + v_dim] = v_acc[ci][i] * inv_sum;
                }
            }
        }
    }
};

// =============================================================================
// MFMA Attn*V: all warps accumulate weighted V using MFMA
// Absorbs KV scales into attention weights before requantizing to FP4
// =============================================================================
template <typename Traits, int HEADS, int KV_SUBTILE, int K_HALF,
          int B_SCALE_STRIDE, int CHUNKS_PER_WARP, int KV_PAD = 0>
__device__ __forceinline__ void mla_mfma_attn_v(
    const float scores[][KV_SUBTILE],
    const uint8_t kv_data[][K_HALF + KV_PAD],
    const float kv_scale[][B_SCALE_STRIDE],
    typename Traits::acc_t v_acc[CHUNKS_PER_WARP],
    int32_t b_sc_one,
    int lane, int warp_id
) {
    using acc_t = typename Traits::acc_t;
    constexpr int IM = Traits::IM;
    constexpr int V_DIM = 512;
    constexpr int V_CHUNK_DIM = IM;
    constexpr int V_CHUNKS = V_DIM / V_CHUNK_DIM;
    constexpr int WARPS = V_CHUNKS / CHUNKS_PER_WARP;
    constexpr int MFMAS_PER_SCALE = 32 / IM;
    constexpr int SCALE_BLOCKS_PER_WARP = CHUNKS_PER_WARP / MFMAS_PER_SCALE;

    const int mfma_a_row = lane % IM;
    const int mfma_a_kgrp = lane / IM;
    const int mfma_b_col = mfma_a_row;
    const int mfma_b_kgrp = mfma_a_kgrp;

    for (int si = 0; si < SCALE_BLOCKS_PER_WARP; si++) {
        const int vscale_blk = warp_id * SCALE_BLOCKS_PER_WARP + si;
        const int k_base_a = mfma_a_kgrp * 32;
        float scaled_attn[32];
        float amax = 0.0f;
        for (int j = 0; j < 32; j++) {
            float aw = (IM <= HEADS || mfma_a_row < HEADS)
                ? scores[mfma_a_row][k_base_a + j] : 0.0f;
            const float vs = kv_scale[k_base_a + j][vscale_blk];
            float val = aw * vs;
            scaled_attn[j] = val;
            float a = fabsf(val);
            amax = (a > amax) ? a : amax;
        }
        auto scale = compute_e8m0_scale_f32(amax);
        QuantBlock aqb = quantize_fp4_block_with_scale(scaled_attn, scale);

        uint32_t a_reg[8] = {};
        *reinterpret_cast<uint128_vec*>(&a_reg[0]) =
            *reinterpret_cast<uint128_vec*>(&aqb.data);
        int32_t a_sc = mla_broadcast_scale(aqb.e8m0);

        int8_vec a_vec = {(int)a_reg[0], (int)a_reg[1], (int)a_reg[2], (int)a_reg[3],
                          0, 0, 0, 0};

        for (int mi = 0; mi < MFMAS_PER_SCALE; mi++) {
            const int acc_idx = si * MFMAS_PER_SCALE + mi;
            const int chunk = warp_id * CHUNKS_PER_WARP + acc_idx;

            const int k_base_b = mfma_b_kgrp * 32;
            const int v_d = chunk * IM + mfma_b_col;
            const int nib_shift = (v_d & 1) * 4;
            const int byte_off = v_d / 2;

            // Row-major layout: kv_data[kv_pos][byte_offset]
            // Load one byte per KV position (stride = K_HALF + KV_PAD)
            uint8_t raw_bytes[32];
            for (int b = 0; b < 32; b++) {
                int k = k_base_b + b;
                raw_bytes[b] = (k < KV_SUBTILE) ? kv_data[k][byte_off] : 0;
            }

            // Extract nibble pairs and pack directly into B registers using pack_fp4_to_u32
            uint32_t b_reg[8] = {};
            for (int r = 0; r < 4; r++) {
                int base = r * 8;  // 4 pairs per uint32
                uint8_t p0 = ((raw_bytes[base]     >> nib_shift) & 0x0F) | (((raw_bytes[base + 1] >> nib_shift) & 0x0F) << 4);
                uint8_t p1 = ((raw_bytes[base + 2] >> nib_shift) & 0x0F) | (((raw_bytes[base + 3] >> nib_shift) & 0x0F) << 4);
                uint8_t p2 = ((raw_bytes[base + 4] >> nib_shift) & 0x0F) | (((raw_bytes[base + 5] >> nib_shift) & 0x0F) << 4);
                uint8_t p3 = ((raw_bytes[base + 6] >> nib_shift) & 0x0F) | (((raw_bytes[base + 7] >> nib_shift) & 0x0F) << 4);
                b_reg[r] = pack_fp4_to_u32(p0, p1, p2, p3);
            }

            int8_vec b_vec = {(int)b_reg[0], (int)b_reg[1], (int)b_reg[2], (int)b_reg[3],
                              0, 0, 0, 0};
            v_acc[acc_idx] = Traits::mfma(a_vec, a_sc, b_vec, b_sc_one, v_acc[acc_idx]);
        }
    }
}

// =============================================================================
// FUSED FlashAttention-style MLA decode kernel (v5 - MFMA V)
//
// Replaces: QK^T GEMM + softmax + attnV with a SINGLE kernel that reads KV
// data ONCE from HBM. Uses online softmax with V accumulator rescaling.
//
// QK^T: single-warp MFMA 16x16x128, 8 iterations over 128-position subtile
// V accumulation: MFMA 16x16x128 with scale-absorption pattern
//   (same as mla_attn_v_mfma_head_merged_kernel)
//
// Grid: (BATCH_SIZE, KV_SPLITS)
// Block: 256 threads (4 warps) - warp 0 does QK^T, all 4 do MFMA V
//
// LDS: Q data+scale (4.9KB) + KV tile (36.9KB) + KV scales (3.1KB)
//      + scores (8.2KB) = ~53 KB
// =============================================================================

template <int N, int BLOCK_SIZE, int B_SCALE_STRIDE, int KV_SPLITS,
          int K_HALF, int NUM_BLOCKS, int A_K_HALF,
          int KV_SUBTILE, bool USE_32x32 = false, bool PROFILE_PHASES = false>
__global__ __launch_bounds__(BLOCK_SIZE)
void mla_fused_attn_kernel(
    const hip_bfloat16* __restrict__ Q_bf16,
    const uint8_t kv_mxfp4[][K_HALF],
    const uint8_t kv_scale[][B_SCALE_STRIDE],
    float sm_scale,
    float partial_out[][16 * 512],
    float partial_lse[][16]
) {
    static_assert(N % KV_SPLITS == 0);
    constexpr int V_DIM = 512;
    constexpr int HEADS = 16;
    constexpr int KV_PER_SPLIT = N / KV_SPLITS;
    static_assert(KV_PER_SPLIT % KV_SUBTILE == 0);

    // MFMA variant traits
    using Traits = MlaMfmaTraits<USE_32x32>;
    using acc_t = typename Traits::acc_t;
    constexpr int IM = Traits::IM;
    constexpr int BPC = Traits::BPC;
    constexpr int ACC_SIZE = Traits::ACC_SIZE;
    constexpr int K_ITERS = (NUM_BLOCKS + BPC - 1) / BPC;
    constexpr int QKT_BATCH = IM;
    constexpr int V_CHUNK_DIM = IM;
    constexpr int V_CHUNKS = V_DIM / V_CHUNK_DIM;

    constexpr int WARPS = BLOCK_SIZE / 64;
    constexpr int CHUNKS_PER_WARP = V_CHUNKS / WARPS;
    constexpr int SM_ELEMS_PER_THREAD = KV_SUBTILE / (BLOCK_SIZE / HEADS);

    const int batch_idx = __builtin_amdgcn_readfirstlane(blockIdx.x);
    const int split_idx = __builtin_amdgcn_readfirstlane(blockIdx.y);
    const int tid = threadIdx.x;
    const int warp_id = __builtin_amdgcn_readfirstlane(tid / 64);
    const int lane = tid % 64;

    const int kv_start_global = __builtin_amdgcn_readfirstlane(split_idx * KV_PER_SPLIT);

    [[maybe_unused]] const bool is_timer_block = PROFILE_PHASES && (blockIdx.x == 0 && blockIdx.y == 0);
    [[maybe_unused]] uint64_t t_start = 0, t_load_total = 0, t_qkt_total = 0, t_store = 0;
    [[maybe_unused]] uint64_t t_softmax_total = 0, t_mfma_v_total = 0;
    [[maybe_unused]] uint64_t t_phase = 0;

    PROFILE_START(t_start)

    // ---- LDS layout ----
    __shared__ uint8_t q_lds_data[HEADS][A_K_HALF];
    __shared__ uint8_t q_lds_scale[NUM_BLOCKS * HEADS];

    // ---- Per-warp MFMA V accumulators + online softmax state ----
    acc_t v_acc[CHUNKS_PER_WARP];
    float running_max[HEADS];
    float running_sum[HEADS];
    for (int c = 0; c < CHUNKS_PER_WARP; c++) {
        v_acc[c] = {};
    }
    for (int h = 0; h < HEADS; h++) {
        running_max[h] = -INFINITY;
        running_sum[h] = 0.0f;
    }

    const int out_row = __builtin_amdgcn_readfirstlane(batch_idx * KV_SPLITS + split_idx);

    // MFMA lane indices (used for both QK^T and V accumulation)
    const int mfma_a_row = lane % IM;
    const int mfma_a_kgrp = lane / IM;
    const int mfma_b_col = mfma_a_row;
    const int mfma_b_kgrp = mfma_a_kgrp;
    const int32_t b_sc_one = mla_broadcast_scale(127);

    // ---- Step 1: Quantize Q from BF16 to MXFP4 directly into LDS ----
    {
        constexpr int K = NUM_BLOCKS * 32;
        const hip_bfloat16* q_batch = reinterpret_cast<const hip_bfloat16*>(Q_bf16) + batch_idx * HEADS * K;
        constexpr int TOTAL_Q_PAIRS = HEADS * NUM_BLOCKS;
        for (int i = tid; i < TOTAL_Q_PAIRS; i += BLOCK_SIZE) {
            const int h = i / NUM_BLOCKS;
            const int blk = i % NUM_BLOCKS;
            QuantBlock qb = quantize_fp4_block_bf16(&q_batch[h * K + blk * 32]);
            *reinterpret_cast<uint128_vec*>(&q_lds_data[h][blk * 16]) =
                *reinterpret_cast<uint128_vec*>(&qb.data);
            q_lds_scale[h + blk * HEADS] = qb.e8m0;
        }
    }
    if (kv_start_global >= N) {
        for (int i = tid; i < HEADS * V_DIM; i += BLOCK_SIZE) {
            partial_out[out_row][i] = 0.0f;
        }
        if (tid < HEADS) partial_lse[out_row][tid] = -INFINITY;
        return;
    }

    if constexpr (WARPS > 1) {
        __syncthreads();
    }

    // ---- Step 2: Iterate over KV in subtiles ----
    const int kv_row_base = __builtin_amdgcn_readfirstlane(batch_idx * N);
    constexpr int NUM_SUBTILES = KV_PER_SPLIT / KV_SUBTILE;
    constexpr int KV_BUFS = (NUM_SUBTILES > 1) ? 2 : 1;
    constexpr int KV_PAD = (KV_SUBTILE >= 64) ? 4 : 0;
    __shared__ uint8_t kv_lds_data[KV_BUFS][KV_SUBTILE][K_HALF + KV_PAD];
    __shared__ float kv_lds_scale_tile[KV_BUFS][KV_SUBTILE][B_SCALE_STRIDE];
    __shared__ float   scores_lds[HEADS][KV_SUBTILE];

    PROFILE_START(t_phase)

    mla_load_kv_tile<KV_SUBTILE, K_HALF, B_SCALE_STRIDE, BLOCK_SIZE, KV_PAD>(
        kv_lds_data[0], kv_lds_scale_tile[0],
        kv_mxfp4, kv_scale, kv_row_base + kv_start_global, tid);

    PROFILE_ACCUM(t_load_total, t_phase)

    const auto kv_end_pos = kv_start_global + KV_PER_SPLIT;
    for (int kv_pos = kv_start_global, cur = 0; kv_pos < kv_end_pos; kv_pos += KV_SUBTILE, cur = (cur + 1) % KV_BUFS) {

        if constexpr (WARPS > 1) { __syncthreads(); }

        PROFILE_START(t_phase)

        // ---- Prefetch next subtile (overlaps with compute on current for double-buffer) ----
        if (KV_BUFS > 1 || kv_pos > kv_start_global) {
              int kv_row = kv_row_base + kv_pos;
              if constexpr (KV_BUFS > 1) {
                // look-ahead to next tile
                kv_row += KV_SUBTILE;
            }

            // Double-buffer: skip if no next subtile to prefetch
            // Single-buffer: always load (iter 0 skip is handled by outer guard)
            if (KV_BUFS == 1 || kv_pos + KV_SUBTILE < kv_end_pos) {
                const int nxt = (cur + 1) % KV_BUFS;
                mla_load_kv_tile<KV_SUBTILE, K_HALF, B_SCALE_STRIDE, BLOCK_SIZE, KV_PAD>(
                    kv_lds_data[nxt], kv_lds_scale_tile[nxt],
                    kv_mxfp4, kv_scale, kv_row, tid);
            }

            if constexpr (KV_BUFS == 1 && WARPS > 1) {
                __syncthreads();
            }
        }

        PROFILE_ACCUM(t_load_total, t_phase)

        // ---- 2b: QK^T via MFMA (all warps) ----
        mla_qkt_mfma<Traits, HEADS, KV_SUBTILE, NUM_BLOCKS,
            A_K_HALF, K_HALF, B_SCALE_STRIDE, WARPS, KV_PAD>(
            q_lds_data, q_lds_scale, kv_lds_data[cur], kv_lds_scale_tile[cur],
            scores_lds, sm_scale, lane, warp_id);

        if constexpr (WARPS > 1) { __syncthreads(); }

        PROFILE_ACCUM(t_qkt_total, t_phase)

        // ---- 2c: Softmax ----
        mla_parallel_softmax<Traits, HEADS, KV_SUBTILE, BLOCK_SIZE,
            WARPS, SM_ELEMS_PER_THREAD, CHUNKS_PER_WARP>(
            scores_lds, running_max, running_sum, v_acc,
            tid, lane, warp_id);

        PROFILE_ACCUM(t_softmax_total, t_phase)

        // ---- 2d: MFMA Attn*V ----
        mla_mfma_attn_v<Traits, HEADS, KV_SUBTILE, K_HALF,
            B_SCALE_STRIDE, CHUNKS_PER_WARP, KV_PAD>(
            scores_lds, kv_lds_data[cur], kv_lds_scale_tile[cur],
            v_acc, b_sc_one, lane, warp_id);

        PROFILE_ACCUM(t_mfma_v_total, t_phase)
    }

    PROFILE_START(t_phase)

    // ---- Step 3: Broadcast running_sum via LDS, then write partial output ----
    // running_sum is only correct on threads 0..15; broadcast to all via scores_lds
    if (tid < HEADS) {
        scores_lds[tid][0] = running_sum[tid];
    }
    if constexpr (WARPS > 1) {
        __syncthreads();
    }

    {
        Traits::template store_v_output<HEADS, V_DIM, CHUNKS_PER_WARP>(
            v_acc, lane, warp_id, out_row, partial_out, scores_lds);
    }

    if (tid < HEADS) {
        float lse = running_max[tid] + logf(fmaxf(running_sum[tid], 1e-20f));
        partial_lse[out_row][tid] = lse;
    }

    if constexpr (PROFILE_PHASES) {
        if (is_timer_block && tid == 0) {
            t_store = __builtin_amdgcn_s_memrealtime() - t_phase;
            uint64_t t_total = __builtin_amdgcn_s_memrealtime() - t_start;
            printf("[Fused Attn] N=%d splits=%d | load=%llu qkt=%llu softmax=%llu mfma_v=%llu store=%llu total=%llu cycles\\n",
                   N, KV_SPLITS,
                   (unsigned long long)t_load_total,
                   (unsigned long long)t_qkt_total,
                   (unsigned long long)t_softmax_total,
                   (unsigned long long)t_mfma_v_total,
                   (unsigned long long)t_store,
                   (unsigned long long)t_total);
        }
    }
}

// =============================================================================
// Fused reduce kernel: combines independently-normalized splits using LSE
// 2-pass: first find max_lse, then accumulate weighted partials
// No large stack arrays - avoids register spill for large KV_SPLITS
// =============================================================================

template <int BLOCK_SIZE, int KV_SPLITS>
__global__ __launch_bounds__(BLOCK_SIZE)
void mla_fused_reduce_kernel(
    const float partial_out[][16 * 512],         // (batch * KV_SPLITS, 16 * 512)
    const float partial_lse[][16],               // (batch * KV_SPLITS, 16)
    hip_bfloat16 output[][512]                   // (batch * 16, 512)
) {
    constexpr int V_DIM = 512;
    constexpr int HEADS = 16;

    const int batch_idx = blockIdx.x;
    const int head_idx = blockIdx.y;
    const int tid = threadIdx.x;

    const int v_dim0 = tid * 2;
    const int v_dim1 = tid * 2 + 1;

    // Pass 1: Find max LSE across all splits (no array needed)
    float max_lse = -INFINITY;
    for (int s = 0; s < KV_SPLITS; s++) {
        float lse = partial_lse[batch_idx * KV_SPLITS + s][head_idx];
        max_lse = fmaxf(max_lse, lse);
    }

    // Pass 2: Accumulate weighted partials using max_lse
    float sum0 = 0.0f, sum1 = 0.0f, denom = 0.0f;
    for (int s = 0; s < KV_SPLITS; s++) {
        float lse = partial_lse[batch_idx * KV_SPLITS + s][head_idx];
        float w = expf(lse - max_lse);
        denom += w;
        sum0 += w * partial_out[batch_idx * KV_SPLITS + s][(int64_t)head_idx * V_DIM + v_dim0];
        sum1 += w * partial_out[batch_idx * KV_SPLITS + s][(int64_t)head_idx * V_DIM + v_dim1];
    }

    float inv_denom = (denom > 0.0f) ? (1.0f / denom) : 0.0f;
    output[batch_idx * HEADS + head_idx][v_dim0] = hip_bfloat16(sum0 * inv_denom);
    output[batch_idx * HEADS + head_idx][v_dim1] = hip_bfloat16(sum1 * inv_denom);
}

'''

# C++ wrapper for PyTorch load_inline compilation - optimized for MI355X
MLA_MXFP4_CPP_SOURCE = r'''
// v4 fused attn kernel rewrite
#include <torch/extension.h>
#include <hip/hip_runtime.h>
#include <hip/hip_bfloat16.h>
#include <ATen/hip/HIPContext.h>
#include <unordered_map>

// =============================================================================
// Fused MLA pipeline: Q quant + fused attn (QK^T + online softmax + V accum) + reduce
// For kv=8192 shapes where reading KV once instead of twice saves ~640MB of HBM traffic.
// =============================================================================

// Split-K for fused pipeline: fewer splits since each block does more work
// and the reduce kernel has per-split overhead
template <int BATCH_SIZE, int N>
constexpr int get_fused_kv_split() {
    if constexpr (BATCH_SIZE == 4 && N == 1024) return 16;
    else if constexpr (BATCH_SIZE == 4 && N == 8192) return 48;
    else if constexpr (BATCH_SIZE == 32 && N == 1024) return 8;
    else if constexpr (BATCH_SIZE == 32 && N == 8192) return 16;
    else if constexpr (BATCH_SIZE == 64 && N == 1024) return 4;
    else if constexpr (BATCH_SIZE == 64 && N == 8192) return 16;
    else if constexpr (BATCH_SIZE == 256 && N == 1024) return 2;
    else if constexpr (BATCH_SIZE == 256 && N == 8192) return 8;
    else return 4;
}

template <int BS, int N>
constexpr int getLen() {
    if constexpr (N <= 1024)
        return 1024;

    return 1536;
}

template <int BATCH_SIZE, int N, int STRIDE, int BS, int KV_SPLITS, int KV_SUBTILE, bool USE_32x32>
torch::Tensor mla_fused_pipeline_impl(
    torch::Tensor Q_bf16,
    torch::Tensor KV_data,
    torch::Tensor KV_scale,
    float sm_scale,
    bool profile
) {
    constexpr int M = 16;
    constexpr int K = 576;
    constexpr int K_HALF = 288;
    constexpr int NUM_BLOCKS = K / 32; // 18
    constexpr int NUM_HEADS = 16;
    constexpr int V_DIM = 512;
    constexpr int A_K_HALF = NUM_BLOCKS * 16;
    constexpr int TOTAL_HEADS = BATCH_SIZE * NUM_HEADS;
    static torch::Tensor partial_v_buf, partial_lse_buf;
    static int last_n = 0, last_splits = 0;

    bool need_realloc = (N != last_n || KV_SPLITS != last_splits);
    if (need_realloc) {
        auto f32opts = torch::TensorOptions().dtype(torch::kFloat32).device(Q_bf16.device());

        partial_v_buf = torch::empty({BATCH_SIZE * KV_SPLITS * NUM_HEADS * V_DIM}, f32opts);
        partial_lse_buf = torch::empty({BATCH_SIZE * KV_SPLITS * NUM_HEADS}, f32opts);

        last_n = N;
        last_splits = KV_SPLITS;
    }

    // ---- Profiling ----
    struct PerfStats {
        float t_fused = 0, t_reduce = 0;
        int count = 0;
    };
    static std::unordered_map<int, std::unordered_map<int, PerfStats>> perf_map;
    constexpr int PROFILE_INTERVAL = 10;
    auto& stats = perf_map[BATCH_SIZE][N];
    stats.count++;
    bool do_profile = profile && (stats.count % PROFILE_INTERVAL == 0);

    hipEvent_t e0, e2, e3;
    if (do_profile) {
        (void)hipEventCreate(&e0);
        (void)hipEventCreate(&e2); (void)hipEventCreate(&e3);
        (void)hipEventRecord(e0);
    }

    // ---- Step 1+2: Fused Q quantization + attention (QK^T MFMA + online softmax + V accumulation) ----
    {
        dim3 grid(BATCH_SIZE, KV_SPLITS);
        dim3 block(BS);
        constexpr bool KERNEL_TIMER = false;
        mla_fused_attn_kernel<N, BS, STRIDE, KV_SPLITS, K_HALF, NUM_BLOCKS, A_K_HALF, KV_SUBTILE, USE_32x32, KERNEL_TIMER>
            <<<grid, block>>>(
            reinterpret_cast<const hip_bfloat16*>(Q_bf16.data_ptr()),
            reinterpret_cast<const uint8_t(*)[K_HALF]>(KV_data.data_ptr()),
            reinterpret_cast<const uint8_t(*)[STRIDE]>(KV_scale.data_ptr()),
            sm_scale,
            reinterpret_cast<float(*)[16 * 512]>(partial_v_buf.data_ptr()),
            reinterpret_cast<float(*)[16]>(partial_lse_buf.data_ptr()));
    }
    if (do_profile) (void)hipEventRecord(e2);

    // ---- Step 3: LSE-corrected reduce across splits ----
    static torch::Tensor output_buf;
    static int last_total_heads = 0;
    if (TOTAL_HEADS != last_total_heads) {
        output_buf = torch::empty({TOTAL_HEADS, V_DIM},
            torch::TensorOptions().dtype(torch::kBFloat16).device(Q_bf16.device()));
        last_total_heads = TOTAL_HEADS;
    }
    auto output = output_buf;
    {
        constexpr int R_BLOCK = 256;
        dim3 r_grid(BATCH_SIZE, NUM_HEADS);
        dim3 r_block(R_BLOCK);
        mla_fused_reduce_kernel<R_BLOCK, KV_SPLITS><<<r_grid, r_block>>>(
            reinterpret_cast<const float(*)[16 * 512]>(partial_v_buf.data_ptr()),
            reinterpret_cast<const float(*)[16]>(partial_lse_buf.data_ptr()),
            reinterpret_cast<hip_bfloat16(*)[512]>(output.data_ptr()));
    }

    if (do_profile) {
        (void)hipEventRecord(e3);
        (void)hipEventSynchronize(e3);
        float d02, d23;
        (void)hipEventElapsedTime(&d02, e0, e2);
        (void)hipEventElapsedTime(&d23, e2, e3);
        stats.t_fused += d02; stats.t_reduce += d23;
        int n = stats.count / PROFILE_INTERVAL;
        if (n < 5) {
            printf("[MLA FUSED] bs=%d kv=%d splits=%d | "
                "fused_attn=%.1fus reduce=%.1fus | "
                "total=%.1fus (avg over %d)\n",
                BATCH_SIZE, N, KV_SPLITS,
                stats.t_fused/n*1000, stats.t_reduce/n*1000,
                (stats.t_fused+stats.t_reduce)/n*1000, n);
        }
        (void)hipEventDestroy(e0);
        (void)hipEventDestroy(e2); (void)hipEventDestroy(e3);
    }

    return output;
}

// ---- Dispatch wrapper ----
torch::Tensor mla_mxfp4_pipeline(
    torch::Tensor Q_bf16,
    torch::Tensor KV_data,
    torch::Tensor KV_scale,
    int batch_size,
    int kv_seq_len,
    float sm_scale,
    bool profile
) {
    int B_SCALE_STRIDE = KV_scale.stride(0);
    assert(B_SCALE_STRIDE == 18 || B_SCALE_STRIDE == 24);
    assert(kv_seq_len == 1024 || kv_seq_len == 8192);

#define MLA_FUSED(BS_VAL, N, STR, BLOCK, KV_SUBTILE, USE32) \
    if (batch_size == BS_VAL && kv_seq_len == N && B_SCALE_STRIDE == STR) { \
        constexpr int EFFECTIVE_LEN = getLen<BS_VAL, N>(); \
        constexpr int KV_SPLITS = get_fused_kv_split<BS_VAL, N>(); \
        return mla_fused_pipeline_impl<BS_VAL, EFFECTIVE_LEN, STR, BLOCK, KV_SPLITS, KV_SUBTILE, USE32>(Q_bf16, KV_data, KV_scale, sm_scale, profile); \
    }

    // Tuned dispatch: MLA_FUSED(batch_size, kv_len, stride, block_size, kv_subtile, use_32x32, double_buffer)
    //
    // === TUNING LOG ===
    // Baseline (pre-optimization): all 16x16x128, KV_SUBTILE=128, single-buffer
    //   bs=4/kv=1024: 41.0µs | bs=4/kv=8192: 47.2µs
    //   bs=32/kv=1024: 40.8µs | bs=32/kv=8192: 176µs
    //   bs=64/kv=1024: 75.6µs | bs=256/kv=1024: 155µs
    //
    // After parallel softmax + pack_fp4_to_u32 (single-buffer):
    //   bs=4/kv=1024: 38.6µs | bs=4/kv=8192: 43.0µs
    //   bs=32/kv=1024: 36.7µs | bs=32/kv=8192: 140µs
    //   bs=64/kv=1024: 65.7µs | bs=256/kv=1024: 140µs
    //
    // Trial 1: DB=true for kv=1024, DB=false for kv=8192
    //   bs=4/kv=1024: 35.9µs ✅ | bs=4/kv=8192: 39.4µs ✅
    //   bs=32/kv=1024: 32.8µs ✅ | bs=32/kv=8192: 189µs 🔴 (was 140)
    //   bs=64/kv=1024: 60.2µs ✅ | bs=64/kv=8192: 306µs 🔴 (was 218 aiter)
    //   bs=256/kv=1024: 188µs 🔴 (was 140) | bs=256/kv=8192: 1038µs 🔴
    //
    // Trial 2: BS=256 for kv=8192, BS=512 for bs=256/kv=1024
    //   bs=4/kv=1024: 35.5µs ✅ | bs=4/kv=8192: 39.7µs
    //   bs=32/kv=1024: 34.2µs | bs=32/kv=8192: 189µs
    //   bs=64/kv=1024: 60.3µs | bs=64/kv=8192: 347µs 🔴 (BS=256 worse)
    //   bs=256/kv=1024: 139µs ✅ (BS=512 no-DB) | bs=256/kv=8192: 1538µs 🔴 (BS=256 worse)
    //
    // Trial 3 (BEST OF): pick winners from Trial 1+2
    //   bs=4/kv=1024: 35.6µs | bs=4/kv=8192: 39.5µs
    //   bs=32/kv=1024: 32.9µs | bs=32/kv=8192: 193µs (MXFP4)
    //   bs=64/kv=1024: 61.6µs | bs=64/kv=8192: 309µs (MXFP4)
    //   bs=256/kv=1024: 143µs | bs=256/kv=8192: 1037µs (MXFP4)
    //
    // Trial 4: Hybrid dispatch (MXFP4 for bs≤4 or kv≤1024, aiter FP8 for rest)
    //   bs=4/kv=1024: 35.7µs MXFP4 | bs=4/kv=8192: 39.3µs MXFP4
    //   bs=32/kv=1024: 32.8µs MXFP4 | bs=32/kv=8192: 184µs aiter
    //   bs=64/kv=1024: 61.6µs MXFP4 | bs=64/kv=8192: 234µs aiter
    //   bs=256/kv=1024: 143µs MXFP4 | bs=256/kv=8192: 385µs aiter
    //
    MLA_FUSED(4, 1024, 18, 256, 64, false);
    MLA_FUSED(4, 1024, 24, 256, 64, false);
    MLA_FUSED(4, 8192, 18, 512, 32, false);
    MLA_FUSED(4, 8192, 24, 512, 32, false);
    MLA_FUSED(32, 1024, 18, 512, 128, false);
    MLA_FUSED(32, 1024, 24, 512, 128, false);
    MLA_FUSED(32, 8192, 18, 256, 96, false);
    MLA_FUSED(32, 8192, 24, 256, 96, false);
    MLA_FUSED(64, 1024, 18, 256, 128, false);
    MLA_FUSED(64, 1024, 24, 256, 128, false);
    MLA_FUSED(64, 8192, 18, 256, 96, false);
    MLA_FUSED(64, 8192, 24, 256, 96, false);
    MLA_FUSED(256, 1024, 18, 512, 128, false);
    MLA_FUSED(256, 1024, 24, 512, 128, false);
    MLA_FUSED(256, 8192, 18, 256, 48, false);
    MLA_FUSED(256, 8192, 24, 256, 48, false);
    TORCH_CHECK(false, "Unsupported batch_size: ", batch_size);
}


PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("mla_mxfp4_pipeline", &mla_mxfp4_pipeline);
}
'''

# Global state for compiled kernel
_torch_hip_module = None  # For PyTorch load_inline


def _try_compile_hip_kernel_torch():
    """Compile HIP kernel using PyTorch load_inline (Method 2)."""
    global _torch_hip_module

    import time

    try:
        print("[PyTorch] Starting compilation...")
        t0 = time.time()

        from torch.utils.cpp_extension import load_inline
        import os

        # Check if we're on ROCm
        if not torch.cuda.is_available():
            print("[PyTorch] CUDA/ROCm not available")
            return False

        # Check for HIP/ROCm
        if not hasattr(torch.version, 'hip') or torch.version.hip is None:
            print("[PyTorch] Not running on HIP/ROCm")
            return False

        rocm_home = os.environ.get('ROCM_HOME', '/opt/rocm')
        t1 = time.time()
        print(f"[PyTorch] Setup took {t1-t0:.2f}s")

        # Combine kernel source with C++ wrapper
        # Convert bytes to string for load_inline
        kernel_source = MLA_MXFP4_HIP_SOURCE.decode('utf-8') + '\n' + MLA_MXFP4_CPP_SOURCE

        print(f"[PyTorch] Starting load_inline...")
        os.environ['PYTORCH_ROCM_ARCH'] = 'gfx950'
        os.environ['MAX_JOBS'] = '4'
        _torch_hip_module = load_inline(
            name='mla_mxfp4_hip_torch_v13',
            cpp_sources='',
            cuda_sources=[kernel_source],
            extra_cflags=['-O3'],
            extra_cuda_cflags=['-O3', '-ffast-math', '-munsafe-fp-atomics', '--offload-arch=gfx950', '-DHIP_ENABLE_EXTRA_WARP_SYNC_TYPES'],
            extra_include_paths=[f'{rocm_home}/include'],
            verbose=True,  # Enable verbose to see what's happening
        )
        t2 = time.time()
        print(f"[PyTorch] load_inline took {t2-t1:.2f}s")

        print(f"[PyTorch] Total compilation time: {t2-t0:.2f}s")
        return True

    except Exception as e:
        print(f"[PyTorch] Compilation failed: {e}")
        import traceback
        traceback.print_exc()
        return False


def _try_compile_hip_kernel():
    """Try both compilation methods, preferring hip-python."""
    if _try_compile_hip_kernel_torch():
        return True
    return False


# Attempt compilation at module load (disabled by default)
# Uncomment the line below to enable JIT compilation of HIP kernel
HAS_HIP_KERNEL = _try_compile_hip_kernel()

def custom_kernel(data: input_t) -> output_t:
    return custom_kernel_mxfp4_qkt(data)


def custom_kernel_mxfp4_qkt(data):
    q, kv_data, _, _, config = data
    batch_size = config["batch_size"]
    kv_seq_len = config["kv_seq_len"]

    kv_buffer_mxfp4, kv_scale_mxfp4 = kv_data["mxfp4"]

    q_flat = q.reshape(batch_size * 16, 576)
    kv_data_flat = kv_buffer_mxfp4.reshape(-1, 288)
    kv_scale_flat = kv_scale_mxfp4.reshape(-1, kv_scale_mxfp4.shape[-1])

    output = _torch_hip_module.mla_mxfp4_pipeline(
        q_flat, kv_data_flat, kv_scale_flat,
        batch_size, kv_seq_len, config["sm_scale"], False)

    return output.reshape(q.shape[0], 16, 512)