submission 603814

"""
Attempt 130: Direct stage1 ASM + reduce_v1 calls with pre-allocated buffers.

Key insight from reading mla_decode_fwd source:
- Persistent mode allocates `logits` (fp32) and `attn_lse` (fp32) EVERY CALL
- These allocations are sized by reduce_partial_map.size(0) * max_seqlen_q
- By calling mla_decode_stage1_asm_fwd + mla_reduce_v1 directly with
  pre-allocated buffers, we eliminate per-call allocation overhead.

Based on attempt_094 (current best: 33.5μs ranked).
"""

import torch
from task import input_t, output_t

import aiter
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

NUM_HEADS = 16
NUM_KV_HEADS = 1
QK_HEAD_DIM = 576
V_HEAD_DIM = 512
SM_SCALE = 1.0 / (QK_HEAD_DIM ** 0.5)
FP8_DTYPE = aiter_dtypes.fp8
NUM_KV_SPLITS = 16
Q_SCALE = torch.ones(1, dtype=torch.float32, device="cuda")

_cache = {}


def _get_page_size(batch_size, kv_seq_len):
    if kv_seq_len > 1024:
        return 8
    elif batch_size >= 64:
        return 2
    else:
        return 1


def _build_cache(batch_size, kv_seq_len, total_q, total_kv, qo_indptr, kv_indptr):
    max_q_len = 1
    q_dtype = FP8_DTYPE
    kv_dtype = FP8_DTYPE
    page_size = _get_page_size(batch_size, kv_seq_len)
    fast_mode = kv_seq_len > 1024

    seq_lens = kv_indptr[1:] - kv_indptr[:-1]
    if page_size == 1:
        kv_last_page_len = seq_lens.to(torch.int32)
        kv_indptr_pages = kv_indptr
        num_pages = total_kv
    else:
        pages_per_seq = (seq_lens + page_size - 1) // page_size
        kv_last_page_len = ((seq_lens - 1) % page_size + 1).to(torch.int32)
        kv_indptr_pages = torch.zeros(batch_size + 1, dtype=torch.int32, device="cuda")
        kv_indptr_pages[1:] = torch.cumsum(pages_per_seq, dim=0).to(torch.int32)
        num_pages = int(kv_indptr_pages[-1].item())

    kv_indices = torch.arange(num_pages, dtype=torch.int32, device="cuda")
    o = torch.empty((total_q, NUM_HEADS, V_HEAD_DIM), dtype=torch.bfloat16, device="cuda")

    info = get_mla_metadata_info_v1(
        batch_size, max_q_len, NUM_HEADS, q_dtype, kv_dtype,
        is_sparse=False, fast_mode=fast_mode,
        num_kv_splits=NUM_KV_SPLITS, intra_batch_mode=True,
    )
    work = [torch.empty(s, dtype=t, device="cuda") for s, t in info]
    (work_metadata, work_indptr, work_info_set,
     reduce_indptr, reduce_final_map, reduce_partial_map) = work

    get_mla_metadata_v1(
        qo_indptr, kv_indptr_pages, 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=max_q_len, uni_seqlen_qo=max_q_len,
        fast_mode=fast_mode, max_split_per_batch=NUM_KV_SPLITS,
        intra_batch_mode=True, dtype_q=q_dtype, dtype_kv=kv_dtype,
    )

    # Pre-allocate intermediate buffers (this is what mla_decode_fwd allocates every call)
    num_partials = reduce_partial_map.size(0)
    logits = torch.empty(
        (num_partials * max_q_len, 1, NUM_HEADS, V_HEAD_DIM),
        dtype=aiter_dtypes.fp32, device="cuda",
    )
    attn_lse = torch.empty(
        (num_partials * max_q_len, 1, NUM_HEADS, 1),
        dtype=aiter_dtypes.fp32, device="cuda",
    )

    return {
        "work_meta_data": 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,
        "kv_indices": kv_indices,
        "kv_last_page_len": kv_last_page_len,
        "kv_indptr_pages": kv_indptr_pages,
        "o": o,
        "logits": logits,
        "attn_lse": attn_lse,
        "total_kv": total_kv,
        "page_size": page_size,
    }


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

    batch_size = config["batch_size"]
    kv_seq_len = config["kv_seq_len"]

    kv_buffer_fp8, kv_scale = kv_data["fp8"]
    total_kv = kv_buffer_fp8.shape[0]
    total_q = q.shape[0]

    cache_key = (batch_size, kv_seq_len)
    if cache_key not in _cache or _cache[cache_key]["total_kv"] != total_kv:
        _cache[cache_key] = _build_cache(
            batch_size, kv_seq_len, total_q, total_kv,
            qo_indptr, kv_indptr,
        )
    c = _cache[cache_key]

    ps = c["page_size"]
    q_fp8_view = q.to(FP8_DTYPE).view(-1, NUM_HEADS, QK_HEAD_DIM)
    kv_4d = kv_buffer_fp8.view(total_kv // ps, ps, NUM_KV_HEADS, QK_HEAD_DIM)

    # Direct stage1 ASM call (bypasses mla_decode_fwd Python overhead)
    aiter.mla_decode_stage1_asm_fwd(
        q_fp8_view,
        kv_4d,
        qo_indptr,
        c["kv_indptr_pages"],
        c["kv_indices"],
        c["kv_last_page_len"],
        None,  # num_kv_splits_indptr (None for persistent mode)
        c["work_meta_data"],
        c["work_indptr"],
        c["work_info_set"],
        1,  # max_seqlen_q
        ps,  # page_size
        NUM_KV_HEADS,
        SM_SCALE,
        c["logits"],  # pre-allocated
        c["attn_lse"],  # pre-allocated
        c["o"],
        Q_SCALE,
        kv_scale,
    )

    # Direct reduce call
    aiter.mla_reduce_v1(
        c["logits"],
        c["attn_lse"],
        c["reduce_indptr"],
        c["reduce_final_map"],
        c["reduce_partial_map"],
        1,  # max_seqlen_q
        c["o"],
        None,  # final_lse (not needed)
    )

    return c["o"]