submission 639194

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

# DeepSeek R1 MLA constants (forward_absorb path)
NUM_HEADS = 16
NUM_KV_HEADS = 1
QK_HEAD_DIM = 576
V_HEAD_DIM = 512
SM_SCALE = QK_HEAD_DIM ** -0.5
PAGE_SIZE = 1
NUM_KV_SPLITS = 32

FP8 = aiter_dtypes.fp8
_finfo = torch.finfo(FP8)
_FP8_MAX = _finfo.max
_FP8_MIN = _finfo.min

# Per-shape cache: metadata + work buffers + auxiliary tensors
_cache: dict[tuple, dict] = {}


def _build_cache(bs: int, qs: int, kvs: int) -> dict:
    """Pre-compute and cache everything that depends only on (batch_size, q_seq_len, kv_seq_len)."""
    total_q = bs * qs
    total_kv = bs * kvs

    qo_indptr = torch.arange(bs + 1, dtype=torch.int32, device="cuda") * qs
    kv_indptr = torch.arange(bs + 1, dtype=torch.int32, device="cuda") * kvs
    kv_last_page_len = torch.full((bs,), kvs, dtype=torch.int32, device="cuda")
    kv_indices = torch.arange(total_kv, dtype=torch.int32, device="cuda")
    output = torch.empty((total_q, NUM_HEADS, V_HEAD_DIM), dtype=torch.bfloat16, device="cuda")

    # Allocate persistent-mode work buffers
    info = get_mla_metadata_info_v1(
        bs, qs, NUM_HEADS, FP8, FP8,
        is_sparse=False, fast_mode=False,
        num_kv_splits=NUM_KV_SPLITS, intra_batch_mode=True,
    )
    work_metadata, work_indptr, work_info_set, \
        reduce_indptr, reduce_final_map, reduce_partial_map = \
        [torch.empty(s, dtype=t, device="cuda") for s, t in info]

    # Populate metadata once (deterministic for fixed indptrs)
    get_mla_metadata_v1(
        qo_indptr, kv_indptr, kv_last_page_len,
        NUM_HEADS // NUM_KV_HEADS, NUM_KV_HEADS, 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=qs, uni_seqlen_qo=qs,
        fast_mode=False, max_split_per_batch=NUM_KV_SPLITS,
        intra_batch_mode=True,
        dtype_q=FP8, dtype_kv=FP8,
    )

    return {
        "output": output,
        "kv_indices": kv_indices,
        "kv_last_page_len": kv_last_page_len,
        "work_metadata": 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,
    }


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

    key = (bs, qs, kvs)
    c = _cache.get(key)
    if c is None:
        c = _build_cache(bs, qs, kvs)
        _cache[key] = c

    # Dynamic per-tensor FP8 quantization of Q
    amax = q.abs().amax().clamp(min=1e-12)
    scale = amax / _FP8_MAX
    q_fp8 = (q / scale).clamp(min=_FP8_MIN, max=_FP8_MAX).to(FP8)
    q_scale = scale.float().view(1)

    kv_fp8, kv_scale = kv_data["fp8"]

    mla_decode_fwd(
        q_fp8,
        kv_fp8.view(-1, 1, 1, QK_HEAD_DIM),
        c["output"],
        qo_indptr, kv_indptr,
        c["kv_indices"], c["kv_last_page_len"],
        qs,
        page_size=PAGE_SIZE,
        nhead_kv=NUM_KV_HEADS,
        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,
        work_meta_data=c["work_metadata"],
        work_indptr=c["work_indptr"],
        work_info_set=c["work_info_set"],
        reduce_indptr=c["reduce_indptr"],
        reduce_final_map=c["reduce_final_map"],
        reduce_partial_map=c["reduce_partial_map"],
    )

    return c["output"]