submission 696095

"""
MLA decode kernel for AMD MI355X (gfx950, CDNA4).

Strategy: aiter fp8 a8w8 persistent-mode decode with aggressive caching.
- Cache metadata work-buffers per (batch, q_len, nhead, kv_total) config
  → eliminates 6+ CUDA allocations on every call
- Cache kv_indices (torch.arange) per total_kv
  → eliminates tensor creation for large index arrays
- Cache kv_last_page_len per (batch, kv_seq)
  → eliminates tensor subtraction each call
- All caches persist across repeated calls in the same worker process
"""

import torch
from task import input_t, output_t

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

# ── Constants ─────────────────────────────────────────────────────────────────
FP8_DTYPE     = aiter_dtypes.fp8
PAGE_SIZE     = 1
NUM_KV_SPLITS = 32
NUM_KV_HEADS  = 1
QK_HEAD_DIM   = 576    # kv_lora_rank + qk_rope_head_dim
V_HEAD_DIM    = 512    # = kv_lora_rank
SM_SCALE      = 1.0 / (QK_HEAD_DIM ** 0.5)

# ── FP8 quantization (per-tensor dynamic) ─────────────────────────────────────

def quantize_fp8(tensor: torch.Tensor):
    finfo = torch.finfo(FP8_DTYPE)
    amax  = tensor.abs().amax().clamp(min=1e-12)
    scale = amax / finfo.max
    fp8   = (tensor / scale).clamp(finfo.min, finfo.max).to(FP8_DTYPE)
    return fp8, scale.to(torch.float32).reshape(1)


# ── Persistent caches (keyed by config; shared across repeated kernel calls) ──

_meta_cache     = {}   # key → dict of 6 work-buffer tensors
_kv_idx_cache   = {}   # total_kv int → kv_indices tensor
_kv_last_cache  = {}   # (batch_size, kv_seq_len) → kv_last_page_len tensor


def _get_kv_indices(total_kv: int) -> torch.Tensor:
    if total_kv not in _kv_idx_cache:
        _kv_idx_cache[total_kv] = torch.arange(
            total_kv, dtype=torch.int32, device="cuda"
        )
    return _kv_idx_cache[total_kv]


def _get_kv_last_page(batch_size: int, kv_seq_len: int) -> torch.Tensor:
    key = (batch_size, kv_seq_len)
    if key not in _kv_last_cache:
        _kv_last_cache[key] = torch.full(
            (batch_size,), kv_seq_len, dtype=torch.int32, device="cuda"
        )
    return _kv_last_cache[key]


def _get_metadata(
    batch_size: int,
    q_seq_len:  int,
    nhead:      int,
    q_dtype:    torch.dtype,
    kv_dtype:   torch.dtype,
    total_kv:   int,
    qo_indptr:  torch.Tensor,
    kv_indptr:  torch.Tensor,
    kv_last_page_len: torch.Tensor,
) -> dict:
    """Build (or return cached) aiter persistent-mode metadata."""
    key = (batch_size, q_seq_len, nhead, q_dtype, kv_dtype, total_kv, NUM_KV_SPLITS)
    if key not in _meta_cache:
        info = get_mla_metadata_info_v1(
            batch_size, q_seq_len, nhead, q_dtype, kv_dtype,
            is_sparse=False, fast_mode=False,
            num_kv_splits=NUM_KV_SPLITS, intra_batch_mode=True,
        )
        bufs = [torch.empty(s, dtype=t, device="cuda") for s, t in info]
        wm, wi, ws, ri, rfm, rpm = bufs

        get_mla_metadata_v1(
            qo_indptr, kv_indptr, kv_last_page_len,
            nhead // NUM_KV_HEADS, NUM_KV_HEADS, True,
            wm, ws, wi, ri, rfm, rpm,
            page_size=PAGE_SIZE,
            kv_granularity=max(PAGE_SIZE, 16),
            max_seqlen_qo=q_seq_len,
            uni_seqlen_qo=q_seq_len,
            fast_mode=False,
            max_split_per_batch=NUM_KV_SPLITS,
            intra_batch_mode=True,
            dtype_q=q_dtype,
            dtype_kv=kv_dtype,
        )
        _meta_cache[key] = dict(
            work_meta_data=wm,
            work_indptr=wi,
            work_info_set=ws,
            reduce_indptr=ri,
            reduce_final_map=rfm,
            reduce_partial_map=rpm,
        )
    return _meta_cache[key]


# ── Main kernel ───────────────────────────────────────────────────────────────

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

    batch_size  = config["batch_size"]
    nhead       = config["num_heads"]
    q_seq_len   = config["q_seq_len"]
    kv_seq_len  = config["kv_seq_len"]
    dq          = config["qk_head_dim"]   # 576
    dv          = config["v_head_dim"]    # 512
    total_q     = batch_size * q_seq_len
    total_kv    = batch_size * kv_seq_len

    # Quantize Q to fp8 (per-tensor dynamic scale)
    q_fp8, q_scale = quantize_fp8(q)

    # fp8 KV buffer and per-tensor scale
    kv_fp8, kv_scale = kv_data["fp8"]

    # Reuse or create cached auxiliary tensors
    kv_indices  = _get_kv_indices(total_kv)
    kv_last_pg  = _get_kv_last_page(batch_size, kv_seq_len)

    # 4-D KV view: (total_kv, page_size, nkv_heads, kv_dim)
    kv_4d = kv_fp8.view(total_kv, PAGE_SIZE, NUM_KV_HEADS, dq)

    # Retrieve (or build) persistent scheduling metadata
    meta = _get_metadata(
        batch_size, q_seq_len, nhead,
        q_fp8.dtype, kv_fp8.dtype, total_kv,
        qo_indptr, kv_indptr, kv_last_pg,
    )

    o = torch.empty((total_q, nhead, dv), dtype=torch.bfloat16, device="cuda")

    mla_decode_fwd(
        q_fp8.view(total_q, nhead, dq),
        kv_4d, o,
        qo_indptr, kv_indptr,
        kv_indices, kv_last_pg,
        q_seq_len,
        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,
        **meta,
    )
    return o