submission 649394

"""MLA decode: adaptive BF16/FP8 Q + zero overhead. Best of both worlds.
BF16 Q for small KV (skip quant saves ~20µs), FP8 Q for large KV (2x bandwidth)."""
import torch
from task import input_t, output_t
from aiter.mla import mla_decode_fwd
from aiter import dtypes as ad
from aiter import get_mla_metadata_info_v1, get_mla_metadata_v1

_FP8 = ad.fp8
_fi = torch.finfo(_FP8)
_mx = _fi.max
_SM = 1.0 / (576 ** 0.5)
_c = {}


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

    # Adaptive: FP8 Q for large KV (bandwidth bound), BF16 Q for small (overhead bound)
    use_fp8_q = tkv > 200000
    if use_fp8_q:
        am = q.abs().amax().clamp(min=1e-12)
        qs = (am / _mx).to(torch.float32).reshape(1)
        qi = (q / qs).to(_FP8)
        qd = _FP8
    else:
        qi = q
        qs = None
        qd = torch.bfloat16

    kv_fp8, kvs = kv_data["fp8"]

    key = (bs, kvsl, qd)
    if key not in _c:
        ki = torch.arange(tkv, dtype=torch.int32, device="cuda")
        kl = torch.full([bs], kvsl, dtype=torch.int32, device="cuda")
        o = torch.empty((bs, 16, 512), dtype=torch.bfloat16, device="cuda")
        info = get_mla_metadata_info_v1(
            bs, 1, 16, qd, _FP8,
            is_sparse=False, fast_mode=True,
            num_kv_splits=32, intra_batch_mode=True,
        )
        w = [torch.empty(s, dtype=t, device="cuda") for s, t in info]
        get_mla_metadata_v1(
            qo_indptr, kv_indptr, kl, 16, 1, True,
            w[0], w[2], w[1], w[3], w[4], w[5],
            page_size=1, kv_granularity=16,
            max_seqlen_qo=1, uni_seqlen_qo=1,
            fast_mode=True, max_split_per_batch=32,
            intra_batch_mode=True, dtype_q=qd, dtype_kv=_FP8,
        )
        _c[key] = (ki, kl, o, w[0], w[1], w[2], w[3], w[4], w[5])

    ki, kl, o, wm, wi, wis, ri, rfm, rpm = _c[key]

    mla_decode_fwd(
        qi, kv_fp8.view(tkv, 1, 1, 576), o,
        qo_indptr, kv_indptr, ki, kl, 1,
        page_size=1, nhead_kv=1,
        sm_scale=_SM, logit_cap=0.0, num_kv_splits=32,
        q_scale=qs, kv_scale=kvs,
        intra_batch_mode=True,
        work_meta_data=wm, work_indptr=wi, work_info_set=wis,
        reduce_indptr=ri, reduce_final_map=rfm, reduce_partial_map=rpm,
    )
    return o