submission 684392

import torch
import os
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t

# ---------------------------------------------------------------------------
# MLA decode v158: Interleaved SV MFMA pairs for pipeline fill
# - Process SV_CHUNKS in pairs: issue B-operand construction + MFMA alternating
# - Independent vacc[vc] and vacc[vc+1] overlap in MFMA pipeline
# - Based on v154 (interleaved QK + non-volatile MFMA)
# ---------------------------------------------------------------------------

os.environ["PYTORCH_ROCM_ARCH"] = "gfx950"

HIP_SRC = r"""
#include <torch/extension.h>
#include <hip/hip_runtime.h>

static constexpr int WARP_SIZE = 64;
static constexpr int NUM_WARPS = 4;
static constexpr int BLOCK_SIZE = WARP_SIZE * NUM_WARPS;

static constexpr int QK_DIM = 576;
static constexpr int V_DIM  = 512;
static constexpr int NUM_HEADS = 16;
static constexpr int NUM_K_CHUNKS = QK_DIM / 32;  // 18
static constexpr int NUM_K128_CHUNKS = (QK_DIM + 127) / 128;  // 5
static constexpr int SUPER_TILE = 32;
static constexpr int SV_CHUNKS = 8;
static constexpr int KV_TILE_BYTES = SUPER_TILE * QK_DIM;  // 18432

// 18432 bytes / 16 bytes per uint4 / 256 threads = 4.5 -> 5 rounds
static constexpr int PF_UINT4S = KV_TILE_BYTES / 16;  // 1152
static constexpr int PF_ROUNDS = (PF_UINT4S + BLOCK_SIZE - 1) / BLOCK_SIZE;  // 5

typedef float __attribute__((ext_vector_type(4))) v4f32;
typedef unsigned int __attribute__((ext_vector_type(4))) u32x4;
typedef int __attribute__((ext_vector_type(4))) i32x4;
typedef int __attribute__((ext_vector_type(8))) i32x8;
typedef unsigned int __attribute__((address_space(3)))* lds_ptr_t;

extern "C" __device__ void __llvm_amdgcn_raw_buffer_load_lds(
    i32x4 rsrc, lds_ptr_t lds_ptr, int size,
    int voffset, int soffset, int offset, int aux)
    __asm("llvm.amdgcn.raw.buffer.load.lds");

struct buffer_resource { uint64_t ptr; uint32_t range; uint32_t config; };

__device__ __forceinline__ i32x4 make_buffer_rsrc(const void* p, uint32_t bytes) {
    buffer_resource r = {reinterpret_cast<uint64_t>(p), bytes, 0x110000};
    return *reinterpret_cast<i32x4*>(&r);
}

__device__ __forceinline__ float bf16_to_f32(unsigned short v) {
    return __uint_as_float(static_cast<unsigned int>(v) << 16);
}

__device__ __forceinline__ unsigned short f32_to_bf16(float v) {
    unsigned int bits = __float_as_uint(v);
    bits += 0x7FFF + ((bits >> 16) & 1);
    return static_cast<unsigned short>(bits >> 16);
}

__device__ __forceinline__ float fp8_to_f32(unsigned char b) {
    return __builtin_amdgcn_cvt_f32_fp8(static_cast<int>(b), 0);
}

__device__ __forceinline__ v4f32 mfma_f32_16x16x128_fp8(
    i32x8 A, i32x8 B, v4f32 C)
{
    v4f32 D;
    asm(
        "v_mfma_f32_16x16x128_f8f6f4 %0, %1, %2, %3 cbsz:0 blgp:0"
        : "=v"(D) : "v"(A), "v"(B), "v"(C));
    return D;
}

// =========================================================================
// Full MFMA pipeline kernel with double-buffered LDS + K=128 QK MFMA
// =========================================================================

__global__ __launch_bounds__(256, 3)
void mla_mfma_pipeline_kernel(
    const unsigned short* __restrict__ q_ptr,
    const unsigned char*  __restrict__ kv_ptr,
    float*                __restrict__ partial_m,
    float*                __restrict__ partial_l,
    float*                __restrict__ partial_acc,
    unsigned short*       __restrict__ out_ptr,
    const int*            __restrict__ qo_indptr,
    const int*            __restrict__ kv_indptr,
    const float*          __restrict__ kv_scale_ptr,
    const int  num_splits,
    const float sm_scale)
{
    const int split_idx = blockIdx.x;
    const int batch_idx = blockIdx.y;
    const int warp_id = threadIdx.x / WARP_SIZE;
    const int lane_id = threadIdx.x % WARP_SIZE;
    const int tid = threadIdx.x;

    const float score_scale = sm_scale * (*kv_scale_ptr);

    const int q_start = qo_indptr[batch_idx];
    const int kv_start = kv_indptr[batch_idx];
    const int kv_end = kv_indptr[batch_idx + 1];
    const int kv_len = kv_end - kv_start;

    const int tps = (kv_len + num_splits - 1) / num_splits;
    const int split_kv_start = kv_start + split_idx * tps;
    const int split_kv_end = min(split_kv_start + tps, kv_end);

    const int mr = lane_id & 0xF;
    const int kg = lane_id >> 4;

    if (split_kv_start >= kv_end) {
        if (lane_id < 16 && warp_id == 0) {
            int head = lane_id;
            int off = (batch_idx * NUM_HEADS + head) * num_splits + split_idx;
            partial_m[off] = -1e30f;
            partial_l[off] = 0.0f;
        }
        return;
    }

    // ===== LDS: double-buffered KV + per-warp W =====
    __shared__ __align__(16) unsigned char kv_lds[2][KV_TILE_BYTES];
    __shared__ float s_W[NUM_WARPS][16][33];

    // ===== Super-tile iteration setup =====
    const int total_tokens = split_kv_end - split_kv_start;
    const int num_st = (total_tokens + SUPER_TILE - 1) / SUPER_TILE;

    // ===== PROLOGUE: issue DMA FIRST, then Q prep overlaps with DMA =====
    {
        const int first_bytes = min(SUPER_TILE, total_tokens) * QK_DIM;
        const unsigned char* __restrict__ src0 = kv_ptr +
            static_cast<long long>(split_kv_start) * QK_DIM;
        i32x4 srsrc = make_buffer_rsrc(src0, first_bytes);
        #pragma unroll
        for (int r = 0; r < PF_ROUNDS; r++) {
            int off = tid * 16 + r * BLOCK_SIZE * 16;
            lds_ptr_t ldp = (lds_ptr_t)(reinterpret_cast<uintptr_t>(kv_lds[0]) + off);
            __llvm_amdgcn_raw_buffer_load_lds(srsrc, ldp, 16, off, 0, 0, 2);
        }
    }

    // ===== Q preload as FP8 (overlapped with DMA in flight) =====
    const unsigned short* qh = q_ptr +
        (static_cast<long long>(q_start) * NUM_HEADS + mr) * QK_DIM;

    i32x8 q_128[NUM_K128_CHUNKS];
    #pragma unroll
    for (int c = 0; c < NUM_K128_CHUNKS; c++) {
        unsigned int w[8];
        int base1 = c * 128 + 16 * kg;
        #pragma unroll
        for (int i = 0; i < 4; i++) {
            int d = base1 + i * 4;
            float f0 = (d     < QK_DIM) ? bf16_to_f32(qh[d])     : 0.f;
            float f1 = (d + 1 < QK_DIM) ? bf16_to_f32(qh[d + 1]) : 0.f;
            float f2 = (d + 2 < QK_DIM) ? bf16_to_f32(qh[d + 2]) : 0.f;
            float f3 = (d + 3 < QK_DIM) ? bf16_to_f32(qh[d + 3]) : 0.f;
            unsigned int pk = __builtin_amdgcn_cvt_pk_fp8_f32(f0, f1, 0, false);
            pk = __builtin_amdgcn_cvt_pk_fp8_f32(f2, f3, pk, true);
            w[i] = pk;
        }
        int base2 = c * 128 + 64 + 16 * kg;
        #pragma unroll
        for (int i = 0; i < 4; i++) {
            int d = base2 + i * 4;
            float f0 = (d     < QK_DIM) ? bf16_to_f32(qh[d])     : 0.f;
            float f1 = (d + 1 < QK_DIM) ? bf16_to_f32(qh[d + 1]) : 0.f;
            float f2 = (d + 2 < QK_DIM) ? bf16_to_f32(qh[d + 2]) : 0.f;
            float f3 = (d + 3 < QK_DIM) ? bf16_to_f32(qh[d + 3]) : 0.f;
            unsigned int pk = __builtin_amdgcn_cvt_pk_fp8_f32(f0, f1, 0, false);
            pk = __builtin_amdgcn_cvt_pk_fp8_f32(f2, f3, pk, true);
            w[4 + i] = pk;
        }
        q_128[c] = *reinterpret_cast<i32x8*>(w);
    }

    // ===== V accumulators + softmax state =====
    float vacc[SV_CHUNKS][4];
    #pragma unroll
    for (int i = 0; i < SV_CHUNKS; i++)
        vacc[i][0] = vacc[i][1] = vacc[i][2] = vacc[i][3] = 0.0f;

    float mv[4] = {-1e30f, -1e30f, -1e30f, -1e30f};
    float lv[4] = {0.0f, 0.0f, 0.0f, 0.0f};

    // ===== Wait for DMA (Q prep ran while DMA was in flight) =====
    asm volatile("s_waitcnt vmcnt(0)" ::: "memory");
    __syncthreads();

    int cur_buf = 0;

    // ===== MAIN LOOP =====
    for (int st_idx = 0; st_idx < num_st; st_idx++) {
        const int stcnt = min(SUPER_TILE, total_tokens - st_idx * SUPER_TILE);
        const int ta = min(16, stcnt);
        const int tb = max(0, stcnt - 16);
        const unsigned char* kv_cur = kv_lds[cur_buf];

        // ---- PREFETCH: GLOBAL_LOAD_LDS for NEXT tile ----
        __builtin_amdgcn_s_setprio(3);
        const int nxt_buf = cur_buf ^ 1;
        const bool has_next = (st_idx + 1 < num_st);

        if (has_next) {
            const int nxt_start = split_kv_start + (st_idx + 1) * SUPER_TILE;
            const int nxt_bytes = min(SUPER_TILE, split_kv_end - nxt_start) * QK_DIM;
            const unsigned char* __restrict__ nsrc = kv_ptr +
                static_cast<long long>(nxt_start) * QK_DIM;
            i32x4 srsrc = make_buffer_rsrc(nsrc, nxt_bytes);

            #pragma unroll
            for (int r = 0; r < PF_ROUNDS; r++) {
                int off = tid * 16 + r * BLOCK_SIZE * 16;
                lds_ptr_t ldp = (lds_ptr_t)(reinterpret_cast<uintptr_t>(kv_lds[nxt_buf]) + off);
                __llvm_amdgcn_raw_buffer_load_lds(srsrc, ldp, 16, off, 0, 0, 2);
            }
        }

        // ---- QK: K=128 MFMA (FP8xFP8), INTERLEAVED CK layout ----
        // Interleave ca and cb MFMAs for MFMA pipeline fill
        __builtin_amdgcn_s_setprio(0);
        v4f32 ca = {0, 0, 0, 0};
        v4f32 cb = {0, 0, 0, 0};
        #pragma unroll
        for (int c = 0; c < NUM_K128_CHUNKS; c++) {
            i32x8 ba = {};
            if (mr < ta) {
                int base1 = mr * QK_DIM + c * 128 + 16 * kg;
                #pragma unroll
                for (int i = 0; i < 4; i++) {
                    int off = base1 + i * 4;
                    if (c * 128 + 16 * kg + i * 4 + 4 <= QK_DIM)
                        ba[i] = *reinterpret_cast<const int*>(&kv_cur[off]);
                }
                int base2 = mr * QK_DIM + c * 128 + 64 + 16 * kg;
                #pragma unroll
                for (int i = 0; i < 4; i++) {
                    int off = base2 + i * 4;
                    if (c * 128 + 64 + 16 * kg + i * 4 + 4 <= QK_DIM)
                        ba[4 + i] = *reinterpret_cast<const int*>(&kv_cur[off]);
                }
            }
            i32x8 bb = {};
            if (mr < tb) {
                int base1 = (16 + mr) * QK_DIM + c * 128 + 16 * kg;
                #pragma unroll
                for (int i = 0; i < 4; i++) {
                    int off = base1 + i * 4;
                    if (c * 128 + 16 * kg + i * 4 + 4 <= QK_DIM)
                        bb[i] = *reinterpret_cast<const int*>(&kv_cur[off]);
                }
                int base2 = (16 + mr) * QK_DIM + c * 128 + 64 + 16 * kg;
                #pragma unroll
                for (int i = 0; i < 4; i++) {
                    int off = base2 + i * 4;
                    if (c * 128 + 64 + 16 * kg + i * 4 + 4 <= QK_DIM)
                        bb[4 + i] = *reinterpret_cast<const int*>(&kv_cur[off]);
                }
            }
            ca = mfma_f32_16x16x128_fp8(q_128[c], ba, ca);
            cb = mfma_f32_16x16x128_fp8(q_128[c], bb, cb);
        }

        // ---- Online softmax (16-lane reduce only: offsets 8,4,2,1) ----
        float sa0 = ca[0] * score_scale, sa1 = ca[1] * score_scale;
        float sa2 = ca[2] * score_scale, sa3 = ca[3] * score_scale;
        float sb0 = cb[0] * score_scale, sb1 = cb[1] * score_scale;
        float sb2 = cb[2] * score_scale, sb3 = cb[3] * score_scale;

        if (mr >= ta) { sa0 = sa1 = sa2 = sa3 = -1e30f; }
        if (mr >= tb) { sb0 = sb1 = sb2 = sb3 = -1e30f; }

        float tm0 = fmaxf(sa0, sb0), tm1 = fmaxf(sa1, sb1);
        float tm2 = fmaxf(sa2, sb2), tm3 = fmaxf(sa3, sb3);
        #pragma unroll
        for (int off = 8; off >= 1; off >>= 1) {
            tm0 = fmaxf(tm0, __shfl_xor(tm0, off));
            tm1 = fmaxf(tm1, __shfl_xor(tm1, off));
            tm2 = fmaxf(tm2, __shfl_xor(tm2, off));
            tm3 = fmaxf(tm3, __shfl_xor(tm3, off));
        }

        float nm0 = fmaxf(mv[0], tm0), nm1 = fmaxf(mv[1], tm1);
        float nm2 = fmaxf(mv[2], tm2), nm3 = fmaxf(mv[3], tm3);
        float rc0 = __expf(mv[0] - nm0), rc1 = __expf(mv[1] - nm1);
        float rc2 = __expf(mv[2] - nm2), rc3 = __expf(mv[3] - nm3);
        mv[0] = nm0; mv[1] = nm1; mv[2] = nm2; mv[3] = nm3;

        #pragma unroll
        for (int vc = 0; vc < SV_CHUNKS; vc++) {
            vacc[vc][0] *= rc0; vacc[vc][1] *= rc1;
            vacc[vc][2] *= rc2; vacc[vc][3] *= rc3;
        }

        float wa0 = (mr < ta) ? __expf(sa0 - nm0) : 0.f;
        float wa1 = (mr < ta) ? __expf(sa1 - nm1) : 0.f;
        float wa2 = (mr < ta) ? __expf(sa2 - nm2) : 0.f;
        float wa3 = (mr < ta) ? __expf(sa3 - nm3) : 0.f;
        float wb0 = (mr < tb) ? __expf(sb0 - nm0) : 0.f;
        float wb1 = (mr < tb) ? __expf(sb1 - nm1) : 0.f;
        float wb2 = (mr < tb) ? __expf(sb2 - nm2) : 0.f;
        float wb3 = (mr < tb) ? __expf(sb3 - nm3) : 0.f;

        float dl0 = wa0 + wb0, dl1 = wa1 + wb1;
        float dl2 = wa2 + wb2, dl3 = wa3 + wb3;
        #pragma unroll
        for (int off = 8; off >= 1; off >>= 1) {
            dl0 += __shfl_xor(dl0, off); dl1 += __shfl_xor(dl1, off);
            dl2 += __shfl_xor(dl2, off); dl3 += __shfl_xor(dl3, off);
        }
        lv[0] = lv[0] * rc0 + dl0; lv[1] = lv[1] * rc1 + dl1;
        lv[2] = lv[2] * rc2 + dl2; lv[3] = lv[3] * rc3 + dl3;

        // ---- W to per-warp LDS ----
        s_W[warp_id][kg * 4    ][mr]      = wa0;
        s_W[warp_id][kg * 4 + 1][mr]      = wa1;
        s_W[warp_id][kg * 4 + 2][mr]      = wa2;
        s_W[warp_id][kg * 4 + 3][mr]      = wa3;
        s_W[warp_id][kg * 4    ][16 + mr] = wb0;
        s_W[warp_id][kg * 4 + 1][16 + mr] = wb1;
        s_W[warp_id][kg * 4 + 2][16 + mr] = wb2;
        s_W[warp_id][kg * 4 + 3][16 + mr] = wb3;

        asm volatile("s_waitcnt lgkmcnt(0)" ::: "memory");

        float wvals[8];
        #pragma unroll
        for (int i = 0; i < 8; i++)
            wvals[i] = s_W[warp_id][mr][i * 4 + kg];

        unsigned int wlo = __builtin_amdgcn_cvt_pk_fp8_f32(wvals[0], wvals[1], 0, false);
        wlo = __builtin_amdgcn_cvt_pk_fp8_f32(wvals[2], wvals[3], wlo, true);
        unsigned int whi = __builtin_amdgcn_cvt_pk_fp8_f32(wvals[4], wvals[5], 0, false);
        whi = __builtin_amdgcn_cvt_pk_fp8_f32(wvals[6], wvals[7], whi, true);
        long w_a = static_cast<long>(wlo) | (static_cast<long>(whi) << 32);

        // ---- SV MFMA from current LDS (each warp: 128 V dims) ----
        // Interleave pairs of vc chunks for MFMA pipeline fill
        __builtin_amdgcn_s_setprio(1);
        const int v_warp_base = warp_id * SV_CHUNKS * 16;

        #pragma unroll
        for (int vc = 0; vc < SV_CHUNKS; vc += 2) {
            const int vd0 = v_warp_base + vc * 16 + mr;
            const int vd0_align = vd0 & ~3;
            const int vd0_shift = (vd0 & 3) * 8;
            const int vd1 = v_warp_base + (vc + 1) * 16 + mr;
            const int vd1_align = vd1 & ~3;
            const int vd1_shift = (vd1 & 3) * 8;

            unsigned int blo0 = 0, bhi0 = 0;
            unsigned int blo1 = 0, bhi1 = 0;
            if (vd0 < V_DIM) {
                #pragma unroll
                for (int i = 0; i < 4; i++) {
                    int tok = i * 4 + kg;
                    unsigned int dw0 = (tok < stcnt) ?
                        *reinterpret_cast<const unsigned int*>(&kv_cur[tok * QK_DIM + vd0_align]) : 0u;
                    unsigned int dw1 = (tok < stcnt) ?
                        *reinterpret_cast<const unsigned int*>(&kv_cur[tok * QK_DIM + vd1_align]) : 0u;
                    blo0 |= ((dw0 >> vd0_shift) & 0xFF) << (i * 8);
                    blo1 |= ((dw1 >> vd1_shift) & 0xFF) << (i * 8);
                }
                #pragma unroll
                for (int i = 0; i < 4; i++) {
                    int tok = (i + 4) * 4 + kg;
                    unsigned int dw0 = (tok < stcnt) ?
                        *reinterpret_cast<const unsigned int*>(&kv_cur[tok * QK_DIM + vd0_align]) : 0u;
                    unsigned int dw1 = (tok < stcnt) ?
                        *reinterpret_cast<const unsigned int*>(&kv_cur[tok * QK_DIM + vd1_align]) : 0u;
                    bhi0 |= ((dw0 >> vd0_shift) & 0xFF) << (i * 8);
                    bhi1 |= ((dw1 >> vd1_shift) & 0xFF) << (i * 8);
                }
            }
            long v_b0 = static_cast<long>(blo0) | (static_cast<long>(bhi0) << 32);
            long v_b1 = static_cast<long>(blo1) | (static_cast<long>(bhi1) << 32);

            v4f32 sc0 = {vacc[vc][0], vacc[vc][1], vacc[vc][2], vacc[vc][3]};
            v4f32 sc1 = {vacc[vc+1][0], vacc[vc+1][1], vacc[vc+1][2], vacc[vc+1][3]};
            sc0 = __builtin_amdgcn_mfma_f32_16x16x32_fp8_fp8(w_a, v_b0, sc0, 0, 0, 0);
            sc1 = __builtin_amdgcn_mfma_f32_16x16x32_fp8_fp8(w_a, v_b1, sc1, 0, 0, 0);
            vacc[vc][0] = sc0[0]; vacc[vc][1] = sc0[1];
            vacc[vc][2] = sc0[2]; vacc[vc][3] = sc0[3];
            vacc[vc+1][0] = sc1[0]; vacc[vc+1][1] = sc1[1];
            vacc[vc+1][2] = sc1[2]; vacc[vc+1][3] = sc1[3];
        }

        // ---- Wait for GLOBAL_LOAD_LDS and flip ----
        __builtin_amdgcn_s_setprio(3);
        if (has_next) {
            asm volatile("s_waitcnt vmcnt(0)" ::: "memory");
            __syncthreads();
            cur_buf = nxt_buf;
        }
    }

    if (num_splits == 1) {
        const float kv_scale = *kv_scale_ptr;
        const int vwb = warp_id * SV_CHUNKS * 16;
        #pragma unroll
        for (int vc = 0; vc < SV_CHUNKS; vc++) {
            int vd = vwb + vc * 16 + mr;
            if (vd < V_DIM) {
                #pragma unroll
                for (int r = 0; r < 4; r++) {
                    int head = kg * 4 + r;
                    float inv_l = (lv[r] > 0.f) ? (kv_scale / lv[r]) : 0.f;
                    long long idx = (static_cast<long long>(q_start) * NUM_HEADS + head) * V_DIM + vd;
                    out_ptr[idx] = f32_to_bf16(vacc[vc][r] * inv_l);
                }
            }
        }
    } else {
        if (warp_id == 0 && mr == 0) {
            #pragma unroll
            for (int r = 0; r < 4; r++) {
                int head = kg * 4 + r;
                int off = (batch_idx * NUM_HEADS + head) * num_splits + split_idx;
                partial_m[off] = mv[r];
                partial_l[off] = lv[r];
            }
        }
        const int vwb = warp_id * SV_CHUNKS * 16;
        #pragma unroll
        for (int vc = 0; vc < SV_CHUNKS; vc++) {
            int vd = vwb + vc * 16 + mr;
            if (vd < V_DIM) {
                #pragma unroll
                for (int r = 0; r < 4; r++) {
                    int head = kg * 4 + r;
                    int off = (batch_idx * NUM_HEADS + head) * num_splits + split_idx;
                    partial_acc[static_cast<long long>(off) * V_DIM + vd] = vacc[vc][r];
                }
            }
        }
    }
}

// =========================================================================
// Reduce kernel
// =========================================================================

__global__ __launch_bounds__(512)
void mla_reduce_kernel(
    const float*          __restrict__ partial_m,
    const float*          __restrict__ partial_l,
    const float*          __restrict__ partial_acc,
    unsigned short*       __restrict__ out_ptr,
    const float*          __restrict__ kv_scale_ptr,
    const int  num_splits)
{
    const int item_idx = blockIdx.x;
    const int tid = threadIdx.x;
    const float kv_scale = *kv_scale_ptr;

    __shared__ float s_corr[128];
    __shared__ float s_inv_l;

    if (tid == 0) {
        float merged_m = -1e30f;
        for (int s = 0; s < num_splits; ++s)
            merged_m = fmaxf(merged_m, partial_m[item_idx * num_splits + s]);
        float merged_l = 0.0f;
        for (int s = 0; s < num_splits; ++s) {
            float l = partial_l[item_idx * num_splits + s];
            float c = (l > 0.f) ? __expf(partial_m[item_idx * num_splits + s] - merged_m) : 0.f;
            s_corr[s] = c;
            merged_l += l * c;
        }
        s_inv_l = (merged_l > 0.f) ? (kv_scale / merged_l) : 0.f;
    }
    __syncthreads();

    if (tid < V_DIM) {
        float val = 0.0f;
        for (int s = 0; s < num_splits; ++s) {
            val += partial_acc[(static_cast<long long>(item_idx) * num_splits + s) * V_DIM + tid]
                   * s_corr[s];
        }
        out_ptr[static_cast<long long>(item_idx) * V_DIM + tid] = f32_to_bf16(val * s_inv_l);
    }
}

torch::Tensor mla_decode(
    torch::Tensor q, torch::Tensor kv_buffer,
    torch::Tensor qo_indptr, torch::Tensor kv_indptr,
    torch::Tensor kv_scale_tensor,
    int64_t num_heads, int64_t num_splits,
    float sm_scale,
    torch::Tensor partial_m, torch::Tensor partial_l,
    torch::Tensor partial_acc,
    torch::Tensor output)
{
    const int batch_size = qo_indptr.size(0) - 1;
    const int num_items = batch_size * static_cast<int>(num_heads);

    dim3 grid1(static_cast<int>(num_splits), batch_size);
    dim3 block1(BLOCK_SIZE);
    mla_mfma_pipeline_kernel<<<grid1, block1>>>(
        reinterpret_cast<const unsigned short*>(q.data_ptr()),
        reinterpret_cast<const unsigned char*>(kv_buffer.data_ptr()),
        partial_m.data_ptr<float>(), partial_l.data_ptr<float>(),
        partial_acc.data_ptr<float>(),
        reinterpret_cast<unsigned short*>(output.data_ptr()),
        qo_indptr.data_ptr<int>(), kv_indptr.data_ptr<int>(),
        kv_scale_tensor.data_ptr<float>(),
        static_cast<int>(num_splits), sm_scale);
    if (num_splits == 1) return output;

    dim3 grid2(num_items);
    dim3 block2(512);
    mla_reduce_kernel<<<grid2, block2>>>(
        partial_m.data_ptr<float>(), partial_l.data_ptr<float>(),
        partial_acc.data_ptr<float>(),
        reinterpret_cast<unsigned short*>(output.data_ptr()),
        kv_scale_tensor.data_ptr<float>(),
        static_cast<int>(num_splits));

    return output;
}
"""

CPP_DECL = """
torch::Tensor mla_decode(
    torch::Tensor q, torch::Tensor kv_buffer,
    torch::Tensor qo_indptr, torch::Tensor kv_indptr,
    torch::Tensor kv_scale_tensor,
    int64_t num_heads, int64_t num_splits,
    float sm_scale,
    torch::Tensor partial_m, torch::Tensor partial_l,
    torch::Tensor partial_acc,
    torch::Tensor output);
"""

_module = load_inline(
    name="mla_hip_v158_interleaved_sv",
    cpp_sources=CPP_DECL,
    cuda_sources=HIP_SRC,
    functions=["mla_decode"],
    extra_cuda_cflags=[
        "-O3", "-std=c++17",
        "-ffast-math", "-funsafe-math-optimizations", "-ffp-contract=fast",
        "-fno-gpu-rdc",
        "-mllvm", "-amdgpu-early-inline-all=true",
        "-mllvm", "-amdgpu-function-calls=false",
        "-mllvm", "-amdgpu-max-memory-clause=64",
        "-mllvm", "-amdgpu-load-store-vectorizer",
        "-mllvm", "-amdgpu-early-ifcvt",
        "-mllvm", "-amdgpu-internalize-symbols",
        "-mllvm", "-amdgpu-scalarize-global-loads",
        "-mllvm", "-amdgpu-dpp-combine",
        "-mllvm", "-amdgpu-enable-pre-ra-optimizations",
        "-mllvm", "-amdgpu-promote-alloca-to-vector-limit=256",
    ],
    verbose=False,
)

_buf_cache = {}


def _get_bufs(num_items, num_splits, total_q, num_heads, device):
    key = (num_items, num_splits, total_q, device)
    if key not in _buf_cache:
        _buf_cache[key] = (
            torch.empty((num_items * num_splits,), dtype=torch.float32, device=device),
            torch.empty((num_items * num_splits,), dtype=torch.float32, device=device),
            torch.empty((num_items * num_splits, 512), dtype=torch.float32, device=device),
            torch.empty((total_q, num_heads, 512), dtype=torch.bfloat16, device=device),
        )
    return _buf_cache[key]


def _choose_splits(batch_size, kv_len):
    tiles = max(1, kv_len // 32)
    max_useful = max(1, tiles // 2)

    if tiles > 64:
        ideal = max(1, -(-768 // batch_size))
    else:
        target_wgs = max(512, batch_size * 8)
        ideal = max(1, target_wgs // batch_size)

    splits = max(1, min(ideal, max_useful, 64))

    while splits > 1 and batch_size * splits > 912:
        splits -= 1

    return splits


def custom_kernel(data: input_t) -> output_t:
    q, kv_data, qo_indptr, kv_indptr, config = data

    kv_buffer_fp8, kv_scale = kv_data["fp8"]
    kv_buffer = kv_buffer_fp8.view(-1, 576)

    batch_size = config["batch_size"]
    num_heads = config["num_heads"]
    sm_scale = config["sm_scale"]
    total_q = q.size(0)
    num_items = batch_size * num_heads

    total_kv = kv_buffer.shape[0]
    kv_len = total_kv // batch_size

    num_splits = _choose_splits(batch_size, kv_len)

    pm, pl, pa, out = _get_bufs(num_items, num_splits, total_q, num_heads, q.device)

    return _module.mla_decode(
        q, kv_buffer, qo_indptr, kv_indptr,
        kv_scale,
        num_heads, num_splits,
        sm_scale,
        pm, pl, pa, out)