submission 714276

#!POPCORN leaderboard amd-mixed-mla
#!POPCORN gpu MI355X

import torch
from task import input_t, output_t
from utils import make_match_reference

from aiter.mla import mla_decode_fwd
from aiter import dtypes as aiter_dtypes
from aiter import get_mla_metadata_info_v1, get_mla_metadata_v1

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
V_HEAD_DIM = KV_LORA_RANK
SM_SCALE = 1.0 / (QK_HEAD_DIM**0.5)

FP8_DTYPE = aiter_dtypes.fp8

_META_CACHE: dict[tuple, dict[str, torch.Tensor]] = {}
_KV_PAGE_INDICES_CACHE: dict[int, torch.Tensor] = {}
_KV_STRUCT_CACHE: dict[tuple[int, int, int], tuple[torch.Tensor, torch.Tensor]] = {}
_KVLAST_TOK_CACHE: dict[tuple[int, int], torch.Tensor] = {}
_Q_SCALE_ONE = torch.ones((1,), device="cuda", dtype=torch.float32)

_SPLITS_TABLE = {
    (4, 1024): 16,
    (4, 8192): 32,
    (32, 1024): 8,
    (32, 8192): 16,
    (64, 1024): 6,
    (64, 8192): 6,
    (256, 1024): 6,
    (256, 8192): 6,
}


def quantize_fp8_noscale(tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    return tensor.to(FP8_DTYPE), _Q_SCALE_ONE


def choose_page_size(kv_seq_len: int) -> int:
    if kv_seq_len >= 8192:
        return 8
    if kv_seq_len >= 1024:
        return 2
    return 1


def choose_num_kv_splits(batch_size: int, kv_seq_len: int) -> int:
    v = _SPLITS_TABLE.get((batch_size, kv_seq_len))
    if v is not None:
        return v
    return 8


def get_kv_page_indices(total_pages: int) -> torch.Tensor:
    cached = _KV_PAGE_INDICES_CACHE.get(total_pages)
    if cached is None:
        cached = torch.arange(total_pages, dtype=torch.int32, device="cuda")
        _KV_PAGE_INDICES_CACHE[total_pages] = cached
    return cached


def get_kv_struct(
    batch_size: int, kv_seq_len: int, page_size: int
) -> tuple[torch.Tensor, torch.Tensor]:
    key = (batch_size, kv_seq_len, page_size)
    cached = _KV_STRUCT_CACHE.get(key)
    if cached is not None:
        return cached

    pages_per_batch = (kv_seq_len + page_size - 1) // page_size
    kv_indptr_pages = (
        torch.arange(batch_size + 1, device="cuda", dtype=torch.int32) * pages_per_batch
    )
    kv_last_page_len = torch.full(
        (batch_size,),
        kv_seq_len - (pages_per_batch - 1) * page_size,
        device="cuda",
        dtype=torch.int32,
    )
    _KV_STRUCT_CACHE[key] = (kv_indptr_pages, kv_last_page_len)
    return kv_indptr_pages, kv_last_page_len


def get_kv_last_len_tok(batch_size: int, kv_seq_len: int) -> torch.Tensor:
    key = (batch_size, kv_seq_len)
    cached = _KVLAST_TOK_CACHE.get(key)
    if cached is None:
        cached = torch.full((batch_size,), kv_seq_len, device="cuda", dtype=torch.int32)
        _KVLAST_TOK_CACHE[key] = cached
    return cached


def get_cached_meta(
    batch_size: int,
    q_seq_len: int,
    kv_seq_len: int,
    nhead: int,
    nhead_kv: int,
    q_dtype: torch.dtype,
    kv_dtype: torch.dtype,
    page_size: int,
    num_kv_splits: int,
    qo_indptr: torch.Tensor,
    kv_indptr: torch.Tensor,
    kv_last_page_len: torch.Tensor,
):
    key = (
        batch_size,
        q_seq_len,
        kv_seq_len,
        nhead,
        nhead_kv,
        str(q_dtype),
        str(kv_dtype),
        page_size,
        num_kv_splits,
    )
    cached = _META_CACHE.get(key)
    if cached is not None:
        return cached

    info = get_mla_metadata_info_v1(
        batch_size,
        q_seq_len,
        nhead,
        q_dtype,
        kv_dtype,
        is_sparse=False,
        fast_mode=True,
        num_kv_splits=num_kv_splits,
        intra_batch_mode=True,
    )
    work = [torch.empty(s, dtype=t, device="cuda") for s, t in info]
    (
        work_metadata,
        work_indptr,
        work_info_set,
        reduce_indptr,
        reduce_final_map,
        reduce_partial_map,
    ) = work

    get_mla_metadata_v1(
        qo_indptr,
        kv_indptr,
        kv_last_page_len,
        nhead // nhead_kv,
        nhead_kv,
        True,
        work_metadata,
        work_info_set,
        work_indptr,
        reduce_indptr,
        reduce_final_map,
        reduce_partial_map,
        page_size=page_size,
        kv_granularity=max(page_size, 16),
        max_seqlen_qo=q_seq_len,
        uni_seqlen_qo=q_seq_len,
        fast_mode=True,
        max_split_per_batch=num_kv_splits,
        intra_batch_mode=True,
        dtype_q=q_dtype,
        dtype_kv=kv_dtype,
    )

    payload = {
        "work_meta_data": work_metadata,
        "work_indptr": work_indptr,
        "work_info_set": work_info_set,
        "reduce_indptr": reduce_indptr,
        "reduce_final_map": reduce_final_map,
        "reduce_partial_map": reduce_partial_map,
    }
    _META_CACHE[key] = payload
    return payload


def custom_kernel(data: input_t) -> output_t:
    q, kv_data, qo_indptr, kv_indptr_tok, config = data
    q_input, q_scale = quantize_fp8_noscale(q)
    kv_buffer_fp8, kv_scale = kv_data["fp8"]

    batch_size = int(config["batch_size"])
    nq = int(config["num_heads"])
    nkv = int(config["num_kv_heads"])
    dq = int(config["qk_head_dim"])
    dv = int(config["v_head_dim"])
    q_seq_len = int(config["q_seq_len"])
    kv_seq_len = int(config["kv_seq_len"])

    page_size = choose_page_size(kv_seq_len)
    if page_size > 1 and (kv_seq_len % page_size) != 0:
        page_size = 1

    num_kv_splits = choose_num_kv_splits(batch_size, kv_seq_len)

    if page_size == 1:
        kv_buffer_4d = kv_buffer_fp8.view(
            kv_buffer_fp8.shape[0], 1, nkv, kv_buffer_fp8.shape[-1]
        )
        kv_indptr = kv_indptr_tok
        total_kv = batch_size * kv_seq_len
        kv_page_indices = get_kv_page_indices(total_kv)
        kv_last_page_len = get_kv_last_len_tok(batch_size, kv_seq_len)
    else:
        pages_per_batch = kv_seq_len // page_size
        total_pages = batch_size * pages_per_batch
        kv_buffer_4d = kv_buffer_fp8.view(
            total_pages, page_size, nkv, kv_buffer_fp8.shape[-1]
        )
        kv_indptr, kv_last_page_len = get_kv_struct(batch_size, kv_seq_len, page_size)
        kv_page_indices = get_kv_page_indices(total_pages)

    meta = get_cached_meta(
        batch_size,
        q_seq_len,
        kv_seq_len,
        nq,
        nkv,
        q_input.dtype,
        kv_buffer_fp8.dtype,
        page_size,
        num_kv_splits,
        qo_indptr,
        kv_indptr,
        kv_last_page_len,
    )

    o = torch.empty((q.shape[0], nq, dv), dtype=torch.bfloat16, device="cuda")
    mla_decode_fwd(
        q_input.view(-1, nq, dq),
        kv_buffer_4d,
        o,
        qo_indptr,
        kv_indptr,
        kv_page_indices,
        kv_last_page_len,
        q_seq_len,
        page_size=page_size,
        nhead_kv=nkv,
        sm_scale=SM_SCALE,
        logit_cap=0.0,
        num_kv_splits=num_kv_splits,
        q_scale=q_scale,
        kv_scale=kv_scale,
        intra_batch_mode=True,
        **meta,
    )
    return o


check_implementation = make_match_reference(custom_kernel, rtol=2e-02, atol=8e-03)