submission 739435

import os 
os.environ['PYTORCH_ROCM_ARCH'] = 'gfx950'


from torch.utils.cpp_extension import load_inline
from task import input_t, output_t
CPP_WRAPPER = """
void mla_ksplit(torch::Tensor q, torch::Tensor kv, torch::Tensor output, torch::Tensor sum);
"""

CUDA_SRC = """
#include <hip/amd_detail/amd_hip_fp8.h>
#include <hip/amd_detail/amd_hip_bf16.h>

using i16x4 = __attribute__((__vector_size__(4 * sizeof(short)))) short;
using f32x4 = __attribute__((__vector_size__(4 * sizeof(float)))) float;

typedef i16x4 (*qkv_t)[16][36][4];

#define SM_SCALE_LOG2 0.06011229337f

__device__ static i16x4 convert4_f32_to_i16(f32x4 x) {
    __hip_bfloat16 tmp[4] = { 
        (__hip_bfloat16)x[0], 
        (__hip_bfloat16)x[1], 
        (__hip_bfloat16)x[2], 
        (__hip_bfloat16)x[3], 
    };
    return *(i16x4*)tmp;
}

__launch_bounds__(256, 4)
__global__ void mla_ksplit_kernel(
    const qkv_t __restrict__ q,
    const qkv_t __restrict__ kv, 
    i16x4 (* __restrict__ output)[16][32][4],
    __hip_bfloat16 * __restrict__ out_sum,       
    const size_t numkgroup
) {
    __shared__ i16x4 q_buf[32][64];
    __shared__ i16x4 kv_buf[2][32][64];
    __shared__ f32x4 psum_buf[4][64];
    __shared__ f32x4 result[32][64];
    const size_t bID = blockIdx.x + blockIdx.y * gridDim.x;
    const size_t trID = (threadIdx.x & 3) << 4 | threadIdx.x >> 2;
    const size_t offset = bID * numkgroup;

    #pragma unroll
    for (size_t i = 0; i < 8; ++i) {
        size_t j = i + threadIdx.y * 8;
        q_buf[j][threadIdx.x] = q[blockIdx.y][threadIdx.x & 15][j][threadIdx.x >> 4];
        result[j][threadIdx.x] = { 0.0f };
    }

    i16x4 q_rope = q[blockIdx.y][threadIdx.x & 15][32 + threadIdx.y][threadIdx.x >> 4];

    for (size_t i = 0; i < 8; ++i) {
        size_t j = i + threadIdx.y * 8;
        kv_buf[0][j][threadIdx.x] = kv[offset][threadIdx.x & 15][j][threadIdx.x >> 4];
    }

    float sum = 0.0f;
    for (size_t n = 0; n < numkgroup; ++n) {
        size_t m = n & 1;
        f32x4 psum = __builtin_amdgcn_mfma_f32_16x16x16bf16_1k(kv[offset + n][threadIdx.x & 15][32 + threadIdx.y][threadIdx.x >> 4], q_rope, { }, 0, 0, 0);

        if (n < numkgroup - 1)
            for (size_t i = 0; i < 8; ++i) {
                size_t j = i + threadIdx.y * 8;
                kv_buf[m ^ 1][j][threadIdx.x] = kv[offset + n + 1][threadIdx.x & 15][j][threadIdx.x >> 4];
            }

        for (size_t i = 0; i < 8; ++i) {
            size_t j = i + threadIdx.y * 8;
            psum = __builtin_amdgcn_mfma_f32_16x16x16bf16_1k(kv_buf[m][j][threadIdx.x], q_buf[j][threadIdx.x], psum, 0, 0, 0);
        }

        psum_buf[threadIdx.y][threadIdx.x] = psum;
        __syncthreads();
        psum_buf[0][threadIdx.x][threadIdx.y] = __builtin_amdgcn_exp2f((psum_buf[0][threadIdx.x][threadIdx.y] + psum_buf[1][threadIdx.x][threadIdx.y] +
            psum_buf[2][threadIdx.x][threadIdx.y] + psum_buf[3][threadIdx.x][threadIdx.y]) * SM_SCALE_LOG2);
        __syncthreads();
        psum = psum_buf[0][threadIdx.x];
        __syncthreads();
        if (threadIdx.y == 0)
            sum += psum[0] + psum[1] + psum[2] + psum[3];

        i16x4 qkv = convert4_f32_to_i16(psum);
        for (size_t i = 0; i < 8; ++i) {
            size_t j = i + threadIdx.y * 8;
            i16x4 kvt = __builtin_amdgcn_ds_read_tr16_b64_v4i16((__attribute__((address_space(3))) i16x4*)&kv_buf[m][j][trID]);
            result[j][threadIdx.x] = __builtin_amdgcn_mfma_f32_16x16x16bf16_1k(kvt, qkv, result[j][threadIdx.x], 0, 0, 0);
        }
    }

    for (size_t i = 0; i < 8; ++i) {
        size_t j = i + threadIdx.y * 8;
        output[bID][threadIdx.x & 15][j][threadIdx.x >> 4] = convert4_f32_to_i16(result[j][threadIdx.x]);
    }

    if (threadIdx.y == 0) {
        sum += __shfl_xor_sync(0xFFFFFFFFFFFFFFFF, sum, 16);
        sum += __shfl_xor_sync(0xFFFFFFFFFFFFFFFF, sum, 32);
        if (threadIdx.x < 16) {
            out_sum[bID * 16 + threadIdx.x] = (__hip_bfloat16)sum;
        }
    }
}

void mla_ksplit(torch::Tensor q, torch::Tensor kv, torch::Tensor output, torch::Tensor sum) {
    size_t by = kv.size(0);
    size_t bx = kv.size(1);
    size_t numkgroup = kv.size(2) / 16;
    mla_ksplit_kernel<<<dim3(bx, by), dim3(64, 4), 0, 0>>> ((qkv_t)q.data_ptr(), (qkv_t)kv.data_ptr(), (i16x4(*)[16][32][4])output.data_ptr(), (__hip_bfloat16*)sum.data_ptr(), numkgroup);
}
"""

import os
os.environ["CXX"] = "clang++"

module = load_inline(
    name='mla_ksplit',
    cpp_sources=[CPP_WRAPPER],
    cuda_sources=[CUDA_SRC],
    functions=['mla_ksplit'],
    verbose=True,
    extra_cuda_cflags=["--offload-arch=gfx950", "-std=c++20"],
)

import torch
import torch.nn.functional as F
from task import input_t, output_t
from utils import make_match_reference


NUM_HEADS = 16
NUM_KV_HEADS = 1
KV_LORA_RANK = 512
QK_ROPE_HEAD_DIM = 64
QK_HEAD_DIM = KV_LORA_RANK + QK_ROPE_HEAD_DIM   # 576
V_HEAD_DIM = KV_LORA_RANK                        # 512
PAGE_SIZE = 1
NUM_KV_SPLITS = 32


SM_SCALE = QK_HEAD_DIM ** -0.5
SM_SCALE_LOG2 = 0.06011229337

from torch.utils.cpp_extension import load_inline

def mla_step1(q: torch.tensor, kv_buffer: torch.tensor) -> tuple[torch.tensor, torch.tensor]:
    output = torch.empty(kv_buffer.shape[0], kv_buffer.shape[1], q.shape[1], 512, dtype=q.dtype, device=q.device)
    sum = torch.empty(kv_buffer.shape[0], kv_buffer.shape[1], q.shape[1], 1, dtype=q.dtype, device=q.device)
    module.mla_ksplit(q, kv_buffer, output, sum)
    return (output, sum)


def mla_ksplit(q: torch.tensor, kv_buffer: torch.tensor, batch_size: int) -> torch.tensor:
    kv_buffer = kv_buffer.view(batch_size, 256 // batch_size, -1, 576)
    # dispatch (batch_size, 256 / batch_size)
    result, qkv_sum = mla_step1(q, kv_buffer) 
    # qkv = (torch.einsum('bij, bsjk -> bsik', q, kv_buffer.transpose(2, 3)) * SM_SCALE_LOG2).exp2_()
    # qkv_sum = qkv.sum(dim=-1, keepdim=True)
    # result = torch.matmul(qkv, kv_buffer[:, :, :, :512])
    # reduce
    return (result.sum(dim=1) * qkv_sum.sum(dim=1).reciprocal_()).squeeze()

# def mla(q: torch.tensor, kv_buffer: torch.tensor, batch_size: int) -> torch.tensor:
#     kv_buffer = kv_buffer.view(batch_size, -1, 576) # (batch_size, kvseqlen, 576)
#     qkv = torch.bmm(q, kv_buffer.transpose(1, 2)) * SM_SCALE # (batch_size, num_heads, kvseqlen)
#     qkv = torch.softmax(qkv, dim=-1)    
#     return torch.bmm(qkv, kv_buffer[:, :, :512]) # (batch_size, num_heads, 512)

def custom_kernel(data: input_t) -> output_t:
    q, kv_data, qo_indptr, kv_indptr, config = data
    kv_buffer = kv_data["bf16"]
    batch_size = config["batch_size"]
    return mla_ksplit(q, kv_buffer, batch_size)