submission 720045

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

"""
Optimized MLA decode kernel for MI355X — Hybrid a16w8 / a8w8.

Key optimizations:
1. Cache metadata buffers, kv_indices, kv_last_page_len, output per config
2. Hybrid kernel selection:
   - Small/medium workloads: bf16 Q + fp8 KV (a16w8) — avoids Q quantization
   - Large workloads (bs*kv >= 524288): fp8 Q + fp8 KV (a8w8) — MFMA throughput
3. Tuned num_kv_splits per workload
4. Pre-allocated, cached everything possible
"""

import torch
from task import input_t, output_t
from utils import make_match_reference

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

# ---------------------------------------------------------------------------
# DeepSeek R1 latent MQA constants
# ---------------------------------------------------------------------------
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   # 576
V_HEAD_DIM = KV_LORA_RANK                        # 512
SM_SCALE = 1.0 / (QK_HEAD_DIM ** 0.5)

PAGE_SIZE = 1
FP8_DTYPE = aiter_dtypes.fp8
FP8_MAX = torch.finfo(FP8_DTYPE).max
FP8_MIN = torch.finfo(FP8_DTYPE).min

# a8w8 wins only for large workloads (bs>=64 with kv=8192)
A8W8_THRESHOLD = 64 * 8192  # 524288

# Try to import sglang/vllm fused fp8 quantization
_sglang_fp8_quant = None
try:
    from sglang.srt.layers.quantization.fp8_kernel import scaled_fp8_quant
    _sglang_fp8_quant = scaled_fp8_quant
except ImportError:
    pass
if _sglang_fp8_quant is None:
    try:
        from vllm._custom_ops import scaled_fp8_quant
        _sglang_fp8_quant = scaled_fp8_quant
    except ImportError:
        pass


def quantize_fp8_fast(tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Lean per-tensor FP8 quantization."""
    if _sglang_fp8_quant is not None:
        return _sglang_fp8_quant(tensor)
    amax = tensor.abs().amax().clamp(min=1e-12)
    scale = amax / FP8_MAX
    fp8_tensor = (tensor / scale).clamp(min=FP8_MIN, max=FP8_MAX).to(FP8_DTYPE)
    return fp8_tensor, scale.to(torch.float32).reshape(1)


# ---------------------------------------------------------------------------
# Cache for metadata and reusable tensors
# ---------------------------------------------------------------------------
_cache = {}


def _get_num_kv_splits(batch_size, kv_seq_len):
    """Tune splits based on workload."""
    total_kv = batch_size * kv_seq_len
    if total_kv <= 4096:
        return 16
    elif total_kv <= 65536:
        return 32
    else:
        return 32


def _build_cache(batch_size, kv_seq_len, total_q, q_dtype, kv_dtype, qo_indptr, kv_indptr):
    """Build and cache metadata + reusable tensors."""
    num_kv_splits = _get_num_kv_splits(batch_size, kv_seq_len)
    total_kv_len = batch_size * kv_seq_len
    max_q_len = 1
    
    kv_indices = torch.arange(total_kv_len, dtype=torch.int32, device="cuda")
    kv_last_page_len = (kv_indptr[1:] - kv_indptr[:-1]).to(torch.int32)
    output = 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=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,
        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=False,
        max_split_per_batch=num_kv_splits,
        intra_batch_mode=True,
        dtype_q=q_dtype,
        dtype_kv=kv_dtype,
    )

    return (
        work_metadata, work_indptr, work_info_set,
        reduce_indptr, reduce_final_map, reduce_partial_map,
        kv_indices, kv_last_page_len, output, num_kv_splits,
    )


def custom_kernel(data: input_t) -> output_t:
    """Hybrid MLA decode: a16w8 for small workloads, a8w8 for large."""
    q, kv_data, qo_indptr, kv_indptr, config = data

    batch_size = config["batch_size"]
    kv_seq_len = config["kv_seq_len"]
    total_kv = batch_size * kv_seq_len
    total_q = q.shape[0]

    # Use pre-quantized fp8 KV always
    kv_buffer_fp8, kv_scale = kv_data["fp8"]

    # Select kernel variant
    use_a8w8 = total_kv >= A8W8_THRESHOLD
    if use_a8w8:
        q_input, q_scale = quantize_fp8_fast(q)
    else:
        q_input = q
        q_scale = None

    # Cache key includes q dtype to differentiate a8w8 vs a16w8
    cache_key = (batch_size, kv_seq_len, total_q, q_input.dtype, kv_buffer_fp8.dtype)
    cached = _cache.get(cache_key)
    if cached is None:
        cached = _build_cache(
            batch_size, kv_seq_len, total_q,
            q_input.dtype, kv_buffer_fp8.dtype,
            qo_indptr, kv_indptr,
        )
        _cache[cache_key] = cached

    (
        work_metadata, work_indptr, work_info_set,
        reduce_indptr, reduce_final_map, reduce_partial_map,
        kv_indices, kv_last_page_len, output, num_kv_splits,
    ) = cached

    mla_decode_fwd(
        q_input.view(-1, NUM_HEADS, QK_HEAD_DIM),
        kv_buffer_fp8.view(kv_buffer_fp8.shape[0], PAGE_SIZE, NUM_KV_HEADS, QK_HEAD_DIM),
        output,
        qo_indptr,
        kv_indptr,
        kv_indices,
        kv_last_page_len,
        1,
        page_size=PAGE_SIZE,
        nhead_kv=NUM_KV_HEADS,
        sm_scale=SM_SCALE,
        logit_cap=0.0,
        num_kv_splits=num_kv_splits,
        q_scale=q_scale,
        kv_scale=kv_scale,
        intra_batch_mode=True,
        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,
    )
    return output