submission 618575

#!POPCORN leaderboard amd-mixed-mla
#!POPCORN gpu MI355X
"""MLA v46 — v44 + optimal kv=8192 splits from sweep.

Key finding: splits=16 for bs=4/kv=8192 gives 19% speedup (38→31µs).
Per-shape optimal splits based on v45a/v45b benchmark data.

Dispatch:
  bs≤4, kv≤1024   → BF16 non-persistent (26µs)
  bs=32, kv=1024   → a16w8 ps=2 splits=16 (35µs)
  bs=64, kv=1024   → BF16 non-persistent (48µs, safe fallback)
  bs=256, kv=1024  → a16w8 ps=2 splits=8 (61µs)
  bs≤4, kv=8192    → a16w8 ps=8 splits=16 (31µs — was 38µs!)
  bs=32, kv=8192   → a16w8 ps=8 splits=8 (42µs)
  bs≥64, kv=8192   → a16w8 ps=8 splits=4 (51-98µs)
"""
from task import input_t, output_t
import torch
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

_meta_cache = {}


def _get_or_make_metadata(batch_size, total_kv, num_heads, nhead_kv, num_splits, page_size,
                          q_dtype, kv_dtype, qo_indptr, kv_indptr, kv_last_page_lens, device):
    key = (batch_size, total_kv, num_heads, num_splits, page_size, str(q_dtype), str(kv_dtype))
    cached = _meta_cache.get(key)
    if cached is not None:
        return cached

    info = get_mla_metadata_info_v1(
        batch_size, 1, num_heads, q_dtype, kv_dtype,
        is_sparse=False, fast_mode=True,
        num_kv_splits=num_splits, intra_batch_mode=False,
    )
    work = [torch.empty(s, dtype=t, device=device) for s, t in info]
    (wmd, wi, wis, ri, rfm, rpm) = work

    get_mla_metadata_v1(
        qo_indptr, kv_indptr, kv_last_page_lens,
        num_heads // nhead_kv, nhead_kv, False,
        wmd, wis, wi, ri, rfm, rpm,
        page_size=page_size, kv_granularity=max(page_size, 16),
        max_seqlen_qo=1, uni_seqlen_qo=1, fast_mode=True,
        max_split_per_batch=num_splits, intra_batch_mode=False,
        dtype_q=q_dtype, dtype_kv=kv_dtype,
    )

    result = dict(
        work_meta_data=wmd, work_indptr=wi, work_info_set=wis,
        reduce_indptr=ri, reduce_final_map=rfm, reduce_partial_map=rpm,
    )
    _meta_cache[key] = result
    return result


def _run_bf16(q, kv_bf16, output, qo_indptr, kv_indptr, config):
    """BF16 non-persistent — exact correctness."""
    batch_size = config['batch_size']
    total_kv = kv_bf16.shape[0]
    kv_buffer = kv_bf16.unsqueeze(1)
    kv_indices = torch.arange(total_kv, device=q.device, dtype=torch.int32)
    kv_last_page_lens = torch.ones(batch_size, device=q.device, dtype=torch.int32)

    mla_decode_fwd(
        q=q, kv_buffer=kv_buffer, o=output,
        qo_indptr=qo_indptr, kv_indptr=kv_indptr,
        kv_indices=kv_indices, kv_last_page_lens=kv_last_page_lens,
        max_seqlen_q=1, page_size=1, nhead_kv=1, sm_scale=config['sm_scale'],
    )


def _run_a16w8(q, kv_fp8_data, kv_fp8_scale, output, qo_indptr, kv_indptr, config, page_size, num_splits=4):
    """a16w8: BF16 Q + FP8 KV persistent."""
    batch_size = config['batch_size']
    num_heads = config['num_heads']
    total_kv = kv_fp8_data.shape[0]

    num_pages = total_kv // page_size
    kv_buffer = kv_fp8_data.view(num_pages, page_size, 1, 576)
    kv_indices = torch.arange(num_pages, device=q.device, dtype=torch.int32)
    kv_indptr_pages = kv_indptr // page_size
    kv_last_page_lens = torch.full((batch_size,), page_size, device=q.device, dtype=torch.int32)

    meta = _get_or_make_metadata(
        batch_size, total_kv, num_heads, 1, num_splits, page_size,
        torch.bfloat16, aiter_dtypes.fp8,
        qo_indptr, kv_indptr_pages, kv_last_page_lens, q.device,
    )

    mla_decode_fwd(
        q, kv_buffer, output,
        qo_indptr, kv_indptr_pages, kv_indices, kv_last_page_lens,
        1, page_size=page_size, nhead_kv=1, sm_scale=config['sm_scale'],
        logit_cap=0.0, num_kv_splits=num_splits,
        q_scale=None, kv_scale=kv_fp8_scale,
        intra_batch_mode=False, **meta,
    )


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

    batch_size = config['batch_size']
    num_heads = config['num_heads']
    v_head_dim = config['v_head_dim']
    kv_seq_len = config['kv_seq_len']

    output = torch.empty((q.shape[0], num_heads, v_head_dim), dtype=q.dtype, device=q.device)

    if kv_seq_len <= 1024 and batch_size <= 4:
        _run_bf16(q, kv_data["bf16"], output, qo_indptr, kv_indptr, config)
    elif kv_seq_len <= 1024 and batch_size == 64:
        _run_bf16(q, kv_data["bf16"], output, qo_indptr, kv_indptr, config)
    elif kv_seq_len <= 1024:
        kv_fp8_data, kv_fp8_scale = kv_data["fp8"]
        splits = 16 if batch_size <= 32 else 8
        _run_a16w8(q, kv_fp8_data, kv_fp8_scale, output, qo_indptr, kv_indptr, config, page_size=2, num_splits=splits)
    else:
        kv_fp8_data, kv_fp8_scale = kv_data["fp8"]
        # Optimal splits from v45a/v45b sweep:
        # bs=4: s=16 (30.9µs, 19% faster than s=8)
        # bs=32: s=8 (42µs, unchanged)
        # bs=64: s=4 (51µs, s=8 is neutral)
        # bs=256: s=4 (98-100µs, s=2/8 marginal)
        if batch_size <= 4:
            splits = 16
        elif batch_size <= 32:
            splits = 8
        else:
            splits = 4
        _run_a16w8(q, kv_fp8_data, kv_fp8_scale, output, qo_indptr, kv_indptr, config, page_size=8, num_splits=splits)

    return output