submission 610716

#!POPCORN leaderboard amd-mixed-mla
#!POPCORN gpu MI355X
"""
MLA Decode v6 — bf16 Q + fp8 KV with metadata caching + pre-allocated output
"""
import torch
from task import input_t, output_t

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
QK_HEAD_DIM = 576
V_HEAD_DIM = 512
SM_SCALE = 1.0 / (576 ** 0.5)
PAGE_SIZE = 1
NUM_KV_SPLITS = 16
FP8_DTYPE = aiter_dtypes.fp8

# Global cache for metadata and pre-allocated tensors
_cache = {}


def _get_or_build_metadata(batch_size, q_len, nq, nkv, q_dtype, kv_dtype,
                           qo_indptr, kv_indptr, kv_last_page_len, total_kv, total_q):
    key = (batch_size, q_len, total_kv, q_dtype, kv_dtype)
    if key not in _cache:
        info = get_mla_metadata_info_v1(
            batch_size, q_len, nq, q_dtype, kv_dtype,
            is_sparse=False, fast_mode=False,
            num_kv_splits=NUM_KV_SPLITS, intra_batch_mode=True,
        )
        buffers = [torch.empty(s, dtype=t, device="cuda") for s, t in info]
        (meta, indptr, info_set, red_indptr, red_final, red_partial) = buffers

        get_mla_metadata_v1(
            qo_indptr, kv_indptr, kv_last_page_len,
            nq // nkv, nkv, True,
            meta, info_set, indptr,
            red_indptr, red_final, red_partial,
            page_size=PAGE_SIZE,
            kv_granularity=max(PAGE_SIZE, 16),
            max_seqlen_qo=q_len,
            uni_seqlen_qo=q_len,
            fast_mode=False,
            max_split_per_batch=NUM_KV_SPLITS,
            intra_batch_mode=True,
            dtype_q=q_dtype,
            dtype_kv=kv_dtype,
        )

        metadata = {
            "work_meta_data": meta,
            "work_indptr": indptr,
            "work_info_set": info_set,
            "reduce_indptr": red_indptr,
            "reduce_final_map": red_final,
            "reduce_partial_map": red_partial,
        }

        # Pre-allocate output and kv_indices
        output = torch.empty((total_q, nq, V_HEAD_DIM), dtype=torch.bfloat16, device="cuda")
        kv_indices = torch.arange(total_kv, dtype=torch.int32, device="cuda")

        _cache[key] = (metadata, output, kv_indices)

    return _cache[key]


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

    batch_size = config["batch_size"]
    nq = config["num_heads"]
    nkv = config["num_kv_heads"]
    dq = config["qk_head_dim"]
    dv = config["v_head_dim"]
    q_len = config["q_seq_len"]

    q_input = q
    q_scale = None

    kv_fp8, kv_scale = kv_data["fp8"]

    total_kv = int(kv_indptr[-1].item())
    total_q = q.shape[0]
    kv_last_page_len = (kv_indptr[1:] - kv_indptr[:-1]).to(torch.int32)
    kv_4d = kv_fp8.view(kv_fp8.shape[0], PAGE_SIZE, nkv, kv_fp8.shape[-1])

    metadata, output, kv_indices = _get_or_build_metadata(
        batch_size, q_len, nq, nkv,
        q_input.dtype, kv_fp8.dtype,
        qo_indptr, kv_indptr, kv_last_page_len,
        total_kv, total_q,
    )

    mla_decode_fwd(
        q_input.view(-1, nq, dq),
        kv_4d,
        output,
        qo_indptr, kv_indptr, kv_indices, kv_last_page_len,
        q_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,
        **metadata,
    )

    return output