submission 592885

#!POPCORN leaderboard amd-mixed-mla
#!POPCORN gpu MI355X
"""MLA Decode v1 — use aiter's ASM-optimized MLA decode."""
from task import input_t, output_t
import torch
from aiter.mla import mla_decode_fwd

def custom_kernel(data: input_t) -> output_t:
    q, kv_data, qo_indptr, kv_indptr, config = data
    kv_bf16 = kv_data["bf16"]

    batch_size = config['batch_size']
    num_heads = config['num_heads']
    qk_head_dim = config['qk_head_dim']
    v_head_dim = config['v_head_dim']
    sm_scale = config['sm_scale']

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

    # Treat contiguous KV as paged with page_size=1
    # kv_bf16 shape: [total_kv, nhead_kv, head_dim] → reshape to [total_kv, 1, nhead_kv, head_dim]
    if kv_bf16.dim() == 3:
        kv_buffer = kv_bf16.unsqueeze(1)  # [total_kv, 1, nhead_kv, head_dim]
    elif kv_bf16.dim() == 2:
        kv_buffer = kv_bf16.unsqueeze(1).unsqueeze(2)  # [total_kv, 1, 1, head_dim]
    else:
        kv_buffer = kv_bf16

    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=sm_scale,
    )

    return output