submission 666334

#!POPCORN leaderboard amd-mixed-mla
#!POPCORN gpu MI355X

"""
MLA v20: Hybrid bf16/fp8 — skip FP8 quant when bf16 a16w8 kernel is faster.
- Small shapes (low total KV tokens): bf16 path (saves 40µs quant overhead)
- Large shapes (high total KV tokens): fp8 path (faster compute on large data)
- Crossover point: ~32K total KV tokens (bs*kv_seq_len)
  bs=4,kv=1024(4K): bf16 30µs vs fp8 50µs → bf16
  bs=32,kv=1024(32K): bf16 33µs vs fp8 60µs → bf16
  bs=32,kv=8192(256K): bf16 93µs vs fp8 103µs → bf16 still wins
  bs=64,kv=8192(512K): bf16 166µs vs fp8 151µs → fp8
  bs=256,kv=8192(2M): bf16 602µs vs fp8 300µs → fp8
"""
import os
os.environ.setdefault('PYTORCH_ROCM_ARCH', 'gfx950')

import torch
import aiter
from task import input_t, output_t
from aiter import dtypes as aiter_dtypes, get_mla_metadata_info_v1, get_mla_metadata_v1

NUM_HEADS = 16
NUM_KV_HEADS = 1
KV_LORA_RANK = 512
QK_ROPE_HEAD_DIM = 64
QK_HEAD_DIM = KV_LORA_RANK + QK_ROPE_HEAD_DIM
V_HEAD_DIM = KV_LORA_RANK
SM_SCALE = 1.0 / (QK_HEAD_DIM ** 0.5)
PAGE_SIZE = 1
FP8_DTYPE = aiter_dtypes.fp8
_FP8_FINFO = torch.finfo(FP8_DTYPE)

# Threshold: use bf16 when total_kv_tokens <= this
# bs=32,kv=8192 = 262144 → bf16 wins (93 vs 103)
# bs=64,kv=8192 = 524288 → fp8 wins (151 vs 166)
BF16_THRESHOLD = 300000

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

def _get_num_kv_splits(batch_size, kv_seq_len):
    if batch_size <= 4:
        return 8 if kv_seq_len <= 2048 else 16
    if batch_size <= 32:
        return 16 if kv_seq_len <= 2048 else 32
    return 32

_cache = {}

def custom_kernel(data: input_t) -> output_t:
    q, kv_data, qo_indptr, kv_indptr, config = data
    batch_size = config["batch_size"]
    nq = config["num_heads"]
    nkv = config["num_kv_heads"]
    dq = config["qk_head_dim"]
    dv = config["v_head_dim"]
    q_seq_len = config["q_seq_len"]
    kv_seq_len = config["kv_seq_len"]

    total_kv_len = batch_size * kv_seq_len
    total_q = q.shape[0]
    max_q_len = q_seq_len
    num_kv_splits = _get_num_kv_splits(batch_size, kv_seq_len)

    # Choose path based on total KV tokens
    use_bf16 = total_kv_len <= BF16_THRESHOLD

    if use_bf16:
        q_input = q
        q_scale = None
        q_dtype = torch.bfloat16
    else:
        q_fp8, q_scale = quantize_fp8(q)
        q_input = q_fp8
        q_dtype = FP8_DTYPE

    kv_buffer_fp8, kv_scale = kv_data["fp8"]
    kv_buffer_4d = kv_buffer_fp8.view(kv_buffer_fp8.shape[0], PAGE_SIZE, nkv, kv_buffer_fp8.shape[-1])

    key = (batch_size, kv_seq_len, num_kv_splits, use_bf16)

    if key not in _cache:
        kv_last_page_len = torch.full(
            (batch_size,), kv_seq_len, dtype=torch.int32, device="cuda"
        )

        info = get_mla_metadata_info_v1(
            batch_size, max_q_len, nq, q_dtype, kv_buffer_fp8.dtype,
            is_sparse=False, fast_mode=False,
            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, kv_last_page_len,
            nq // nkv, nkv, 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=False, max_split_per_batch=num_kv_splits,
            intra_batch_mode=True, dtype_q=q_dtype, dtype_kv=kv_buffer_fp8.dtype,
        )

        P = reduce_partial_map.size(0)
        _cache[key] = {
            'work': work,
            'kv_indices': torch.arange(total_kv_len, dtype=torch.int32, device="cuda"),
            'output': torch.empty((total_q, nq, dv), dtype=torch.bfloat16, device="cuda"),
            'logits': torch.empty((P * max_q_len, 1, nq, dv), dtype=torch.float32, device="cuda"),
            'attn_lse': torch.empty((P * max_q_len, 1, nq, 1), dtype=torch.float32, device="cuda"),
            'kv_last_page_len': kv_last_page_len,
        }

    c = _cache[key]
    (work_metadata, work_indptr, work_info_set,
     reduce_indptr, reduce_final_map, reduce_partial_map) = c['work']
    o = c['output']

    aiter.mla_decode_stage1_asm_fwd(
        q_input.view(-1, nq, dq),
        kv_buffer_4d,
        qo_indptr,
        kv_indptr,
        c['kv_indices'],
        c['kv_last_page_len'],
        None,
        work_metadata,
        work_indptr,
        work_info_set,
        max_q_len,
        PAGE_SIZE,
        nkv,
        SM_SCALE,
        c['logits'],
        c['attn_lse'],
        o,
        q_scale,
        kv_scale,
    )

    aiter.mla_reduce_v1(
        c['logits'],
        c['attn_lse'],
        reduce_indptr,
        reduce_final_map,
        reduce_partial_map,
        max_q_len,
        o,
        None,
    )

    return o