submission 740330

"""MLA decode Triton v109: exp2 + fast rcp on v94 baseline.

1. tl.exp → tl.math.exp2 with LOG2E prescale. Softmax inputs are ≤0,
   so the range fixup (v_cmp+v_cndmask+v_ldexp) is wasted. exp2 may
   compile to bare v_exp_f32. Saves ~200 VALU in hot loop.
2. IEEE div → v_rcp_f32 via inline asm for acc/esum normalization.
   Saves ~190 instructions in epilogue. FP8 output has ~12% quant
   error anyway.
"""

import torch
import triton
import triton.language as tl
from task import input_t, output_t

NUM_HEADS = 16
KV_LORA_RANK = 512
QK_HEAD_DIM = 576
V_HEAD_DIM = 512
FP8_DTYPE = tl.float8e4nv
LOG2E = tl.constexpr(1.4426950408889634)

_SHAPE_CONFIGS = {
    (4, 1024):   (16, 2),
    (4, 8192):   (32, 3),
    (32, 1024):  (4,  3),
    (32, 8192):  (8,  3),
    (64, 1024):  (4,  3),
    (64, 8192):  (8,  2),
    (256, 1024): (1,  2),
    (256, 8192): (2,  2),
}


@triton.jit
def _remap_xcd(pid, GRID_SIZE, NUM_XCDS: tl.constexpr = 8):
    pids_per_xcd = (GRID_SIZE + NUM_XCDS - 1) // NUM_XCDS
    tall_xcds = GRID_SIZE % NUM_XCDS
    tall_xcds = NUM_XCDS if tall_xcds == 0 else tall_xcds
    xcd = pid % NUM_XCDS
    local_pid = pid // NUM_XCDS
    if xcd < tall_xcds:
        return xcd * pids_per_xcd + local_pid
    else:
        return tall_xcds * pids_per_xcd + (xcd - tall_xcds) * (pids_per_xcd - 1) + local_pid


@triton.jit
def _fast_rcp(x):
    return tl.inline_asm_elementwise(
        "v_rcp_f32_e32 $0, $1", "=v, v", [x],
        dtype=tl.float32, is_pure=True, pack=1,
    )


@triton.jit
def _stage1_kernel(
    Q, KV, Mid_O, Out,
    qo_indptr, kv_indptr,
    sm_scale, kv_scale_ptr,
    stride_qt, stride_qh, stride_kvt,
    stride_mb, stride_mh, stride_ms,
    stride_ot, stride_oh,
    NUM_KV_SPLITS: tl.constexpr, BLOCK_N: tl.constexpr,
    BLOCK_H: tl.constexpr, BLOCK_C: tl.constexpr, BLOCK_R: tl.constexpr,
    DIRECT_OUT: tl.constexpr, GRID_SIZE: tl.constexpr,
):
    raw_pid = tl.program_id(0)
    pid = _remap_xcd(raw_pid, GRID_SIZE)
    cur_batch = pid // NUM_KV_SPLITS
    split_id = pid % NUM_KV_SPLITS
    heads = tl.arange(0, BLOCK_H)
    kv_scale = tl.load(kv_scale_ptr).to(tl.float32)
    score_scale = sm_scale * kv_scale
    q_tok = tl.load(qo_indptr + cur_batch)
    kv_s = tl.load(kv_indptr + cur_batch)
    kv_e = tl.load(kv_indptr + cur_batch + 1)
    kv_len = kv_e - kv_s
    kv_per_split = tl.cdiv(kv_len, NUM_KV_SPLITS)
    sp_start = kv_per_split * split_id
    sp_end = tl.minimum(sp_start + kv_per_split, kv_len)
    offs_c = tl.arange(0, BLOCK_C)
    offs_r = tl.arange(0, BLOCK_R)
    q_base = q_tok * stride_qt
    q_nope = tl.load(Q + q_base + heads[:, None] * stride_qh + offs_c[None, :]).to(FP8_DTYPE)
    q_rope = tl.load(Q + q_base + heads[:, None] * stride_qh + (512 + offs_r[None, :])).to(FP8_DTYPE)
    e_max = tl.full([BLOCK_H], value=float("-inf"), dtype=tl.float32)
    e_sum = tl.zeros([BLOCK_H], dtype=tl.float32)
    acc = tl.zeros([BLOCK_H, BLOCK_C], dtype=tl.float32)
    num_tokens = sp_end - sp_start
    num_full = num_tokens // BLOCK_N
    full_end = sp_start + num_full * BLOCK_N

    for start_n in range(sp_start, full_end, BLOCK_N):
        kv_idx = kv_s + start_n + tl.arange(0, BLOCK_N)
        kv_nope = tl.load(KV + kv_idx[:, None] * stride_kvt + offs_c[None, :],
                          cache_modifier=".cg")
        kv_rope = tl.load(KV + kv_idx[:, None] * stride_kvt + (512 + offs_r[None, :]),
                          cache_modifier=".cg")
        k_nope = tl.trans(kv_nope)
        k_rope = tl.trans(kv_rope)
        qk = tl.dot(q_nope, k_nope) + tl.dot(q_rope, k_rope)
        qk = qk * score_scale
        new_max = tl.maximum(tl.max(qk, 1), e_max)
        rescale = tl.math.exp2((e_max - new_max) * LOG2E)
        p = tl.math.exp2((qk - new_max[:, None]) * LOG2E)
        acc = acc * rescale[:, None]
        acc = acc + tl.dot(p.to(FP8_DTYPE), kv_nope)
        e_sum = e_sum * rescale + tl.sum(p, 1)
        e_max = new_max

    if full_end < sp_end:
        offs_n = full_end + tl.arange(0, BLOCK_N)
        mask_n = offs_n < sp_end
        kv_idx = kv_s + offs_n
        kv_nope = tl.load(KV + kv_idx[:, None] * stride_kvt + offs_c[None, :],
                          mask=mask_n[:, None], other=0.0, cache_modifier=".cg")
        kv_rope = tl.load(KV + kv_idx[:, None] * stride_kvt + (512 + offs_r[None, :]),
                          mask=mask_n[:, None], other=0.0, cache_modifier=".cg")
        k_nope = tl.trans(kv_nope)
        k_rope = tl.trans(kv_rope)
        qk = tl.dot(q_nope, k_nope) + tl.dot(q_rope, k_rope)
        qk = qk * score_scale
        qk = tl.where(mask_n[None, :], qk, float("-inf"))
        new_max = tl.maximum(tl.max(qk, 1), e_max)
        rescale = tl.math.exp2((e_max - new_max) * LOG2E)
        p = tl.math.exp2((qk - new_max[:, None]) * LOG2E)
        acc = acc * rescale[:, None]
        acc = acc + tl.dot(p.to(FP8_DTYPE), kv_nope)
        e_sum = e_sum * rescale + tl.sum(p, 1)
        e_max = new_max

    safe_esum = tl.where(e_sum > 0, e_sum, 1.0)
    inv_esum = _fast_rcp(safe_esum)
    result = acc * inv_esum[:, None]
    if DIRECT_OUT:
        o_base = q_tok * stride_ot
        tl.store(Out + o_base + heads[:, None] * stride_oh + offs_c[None, :],
                 (result * kv_scale).to(tl.bfloat16))
    else:
        mid_base = cur_batch * stride_mb + heads * stride_mh + split_id * stride_ms
        tl.store(Mid_O + mid_base[:, None] + offs_c[None, :], result)
        lse = tl.where(e_sum > 0, e_max + tl.log(e_sum), float("-inf"))
        tl.store(Mid_O + mid_base + 512, lse)


@triton.jit
def _stage2_kernel(
    Mid_O, Out, qo_indptr, kv_scale_ptr,
    stride_mb, stride_mh, stride_ms, stride_ot, stride_oh,
    NUM_KV_SPLITS: tl.constexpr, BLOCK_DV: tl.constexpr,
    batch: tl.constexpr, GRID_SIZE: tl.constexpr,
):
    raw_pid = tl.program_id(0)
    pid = _remap_xcd(raw_pid, GRID_SIZE)
    cur_batch = pid % batch
    cur_head = pid // batch
    kv_scale = tl.load(kv_scale_ptr).to(tl.float32)
    q_tok = tl.load(qo_indptr + cur_batch)
    offs_d = tl.arange(0, BLOCK_DV)
    e_max = float("-inf")
    e_sum = 0.0
    acc = tl.zeros([BLOCK_DV], dtype=tl.float32)
    mid_base = cur_batch * stride_mb + cur_head * stride_mh
    for s in range(NUM_KV_SPLITS):
        tv = tl.load(Mid_O + mid_base + s * stride_ms + offs_d)
        lse = tl.load(Mid_O + mid_base + s * stride_ms + 512)
        new_max = tl.maximum(lse, e_max)
        old_scale = tl.math.exp2((e_max - new_max) * LOG2E)
        exp_lse = tl.math.exp2((lse - new_max) * LOG2E)
        acc = acc * old_scale + exp_lse * tv
        e_sum = e_sum * old_scale + exp_lse
        e_max = new_max
    safe_esum = tl.maximum(e_sum, 1e-12)
    inv_esum = _fast_rcp(safe_esum)
    result = acc * inv_esum * kv_scale
    o_base = q_tok * stride_ot + cur_head * stride_oh
    tl.store(Out + o_base + offs_d, result.to(tl.bfloat16))


def _get_config(batch_size, kv_len):
    key = (batch_size, kv_len)
    if key in _SHAPE_CONFIGS:
        return _SHAPE_CONFIGS[key]
    max_tiles = max(1, kv_len // 64)
    if batch_size >= 256:
        splits = min(2, max_tiles)
    elif batch_size >= 64:
        splits = min(max(1, 512 // batch_size), max_tiles)
    else:
        splits = min(max(1, 768 // batch_size), max_tiles)
    while splits > 1 and batch_size * splits > 912:
        splits -= 1
    grid_size = batch_size * splits
    ns = 3 if grid_size < 512 else 2
    return splits, ns


def custom_kernel(data: input_t) -> output_t:
    q, kv_data, qo_indptr, kv_indptr, config = data
    kv_fp8, kv_scale_t = kv_data["fp8"]
    kv = kv_fp8.view(torch.float8_e4m3fn).view(-1, QK_HEAD_DIM)
    batch_size = config["batch_size"]
    sm_scale = config["sm_scale"]
    total_q = q.shape[0]
    kv_len = kv.shape[0] // batch_size if batch_size > 0 else 0

    num_splits, ns = _get_config(batch_size, kv_len)

    dev = q.device
    direct_out = (num_splits == 1)
    mid_o = torch.empty(
        (batch_size, NUM_HEADS, num_splits, KV_LORA_RANK + 1),
        dtype=torch.float32, device=dev,
    ) if not direct_out else torch.empty(1, dtype=torch.float32, device=dev)

    out = torch.empty((total_q, NUM_HEADS, V_HEAD_DIM), dtype=torch.bfloat16, device=dev)

    grid_size1 = batch_size * num_splits
    _stage1_kernel[(grid_size1,)](
        q, kv, mid_o, out, qo_indptr, kv_indptr,
        sm_scale, kv_scale_t,
        q.stride(0), q.stride(1), kv.stride(0),
        mid_o.stride(0) if not direct_out else 0,
        mid_o.stride(1) if not direct_out else 0,
        mid_o.stride(2) if not direct_out else 0,
        out.stride(0), out.stride(1),
        NUM_KV_SPLITS=num_splits, BLOCK_N=64, BLOCK_H=16,
        BLOCK_C=512, BLOCK_R=64, DIRECT_OUT=direct_out,
        GRID_SIZE=grid_size1,
        num_warps=4, num_stages=ns, waves_per_eu=2,
    )

    if not direct_out:
        grid_size2 = NUM_HEADS * batch_size
        _stage2_kernel[(grid_size2,)](
            mid_o, out, qo_indptr, kv_scale_t,
            mid_o.stride(0), mid_o.stride(1), mid_o.stride(2),
            out.stride(0), out.stride(1),
            NUM_KV_SPLITS=num_splits, BLOCK_DV=512,
            batch=batch_size, GRID_SIZE=grid_size2,
            num_warps=4, num_stages=2,
        )

    return out