submission 663325

import torch
from task import input_t, output_t

from aiter.mla import mla_decode_fwd
from aiter import dtypes as aiter_dtypes

NUM_HEADS = 16
NUM_KV_HEADS = 1
QK_HEAD_DIM = 576
V_HEAD_DIM = 512
SM_SCALE = 1.0 / (QK_HEAD_DIM ** 0.5)
PAGE_SIZE = 1


class _CachedState:
    __slots__ = [
        'kv_indices', 'kv_last_page_len', 'max_q_len', 'o',
        'total_kv',
    ]

_shape_cache = {}


def _build_shape_cache(batch_size, q_seq_len, kv_seq_len, total_q, total_kv, qo_indptr, kv_indptr):
    s = _CachedState()
    s.total_kv = total_kv
    s.max_q_len = q_seq_len
    s.kv_indices = torch.arange(total_kv, dtype=torch.int32, device="cuda")
    s.kv_last_page_len = (kv_indptr[1:] - kv_indptr[:-1]).to(torch.int32)
    s.o = torch.empty((total_q, NUM_HEADS, V_HEAD_DIM), dtype=torch.bfloat16, device="cuda")
    return s


def custom_kernel(data: input_t) -> output_t:
    q, kv_data, qo_indptr, kv_indptr, config = data
    bs = config["batch_size"]
    kvlen = config["kv_seq_len"]
    shape_key = (bs, kvlen)

    s = _shape_cache.get(shape_key)
    if s is None:
        total_q = q.shape[0]
        total_kv = int(kv_indptr[-1].item())
        s = _build_shape_cache(bs, config["q_seq_len"], kvlen, total_q, total_kv, qo_indptr, kv_indptr)
        _shape_cache[shape_key] = s

    # All bf16: no Q quantization needed, no correctness issues
    kv_bf16 = kv_data["bf16"]
    kv_4d = kv_bf16.view(s.total_kv, PAGE_SIZE, NUM_KV_HEADS, QK_HEAD_DIM)

    mla_decode_fwd(
        q, kv_4d, s.o,
        qo_indptr, kv_indptr, s.kv_indices, s.kv_last_page_len,
        s.max_q_len,
        page_size=PAGE_SIZE, nhead_kv=NUM_KV_HEADS,
        sm_scale=SM_SCALE, logit_cap=0.0,
    )

    return s.o