submission 745692

"""MLA decode Triton v135: constexpr strides.

All tensor strides are known constants. Making them constexprs:
- Enables shift+add instead of multiply for address calculations
  (stride_kvt=576=2^9+2^6, used 128x per tile)
- Eliminates 8 SGPR loads at kernel entry
- Reduces kernel arguments from 16→8 (stage1), 9→4 (stage2)
- Compiler can fully resolve Q/output address patterns at compile time

Based on v128 (GM=32.5us).
"""

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)
LN2 = tl.constexpr(0.6931471805599453)

_SHAPE_CONFIGS = {
    (4, 1024):   16,
    (4, 8192):   32,
    (32, 1024):  4,
    (32, 8192):  8,
    (64, 1024):  4,
    (64, 8192):  8,
    (256, 1024): 1,
    (256, 8192): 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,
    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,
    STRIDE_QT: tl.constexpr, STRIDE_QH: tl.constexpr,
    STRIDE_KVT: tl.constexpr,
    STRIDE_OT: tl.constexpr, STRIDE_OH: tl.constexpr,
    STRIDE_MS: 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:
        stride_mh = NUM_KV_SPLITS * STRIDE_MS
        stride_mb = BLOCK_H * stride_mh
        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.math.log2(e_sum) * LN2, float("-inf"))
        tl.store(Mid_O + mid_base + 512, lse)


@triton.jit
def _stage2_kernel(
    Mid_O, Out, qo_indptr, kv_scale_ptr,
    NUM_KV_SPLITS: tl.constexpr, BLOCK_DV: tl.constexpr,
    batch: tl.constexpr, GRID_SIZE: tl.constexpr,
    STRIDE_OT: tl.constexpr, STRIDE_OH: tl.constexpr,
    STRIDE_MS: tl.constexpr, BLOCK_H: 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)
    stride_mh = NUM_KV_SPLITS * STRIDE_MS
    stride_mb = BLOCK_H * stride_mh
    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_splits(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:
        return min(2, max_tiles)
    elif batch_size >= 64:
        return 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
        return splits


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 = _get_splits(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,
        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,
        STRIDE_QT=NUM_HEADS * QK_HEAD_DIM,
        STRIDE_QH=QK_HEAD_DIM,
        STRIDE_KVT=QK_HEAD_DIM,
        STRIDE_OT=NUM_HEADS * V_HEAD_DIM,
        STRIDE_OH=V_HEAD_DIM,
        STRIDE_MS=KV_LORA_RANK + 1,
        num_warps=4, num_stages=2, 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,
            NUM_KV_SPLITS=num_splits, BLOCK_DV=512,
            batch=batch_size, GRID_SIZE=grid_size2,
            STRIDE_OT=NUM_HEADS * V_HEAD_DIM,
            STRIDE_OH=V_HEAD_DIM,
            STRIDE_MS=KV_LORA_RANK + 1,
            BLOCK_H=16,
            num_warps=4, num_stages=2,
        )

    return out