submission 742151

"""
mixed custom×aiter
4x1024:   bf16 splitk (triton_flash_v1)
4x8192:   fp8 splitk (mk19g_c)
32x1024:  fp8 splitk (mk19_FINAL_c)
32x8192:  fp8 splitk mk19f (champ, 74.4µs)
64x1024:  fp8 splitk (mk1a_c)
~~64x8192:  fp8 splitk mk1a (champ, 124.3µs)~~ aiter
256x1024: fp8 stage1 splitk (hip_v17_mk1g)
~~256x8192: fp8 splitk mk4j (champ, 319.4µs)~~ aiter
"""

import torch
import triton
import triton.language as tl
from task import input_t, output_t
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

QK_HEAD_DIM = 576
V_HEAD_DIM = 512
ROPE_DIM = 64
NOPE_DIM = 512
NUM_HEADS = 16
SM_SCALE = 1.0 / (QK_HEAD_DIM ** 0.5)
FP8_DTYPE = aiter_dtypes.fp8
FP8_MAX = 448.0


# ═══════════════════════════════════════════════════════════════════════
# Shared: Reduce kernel (used by all triton splitk paths)
# ═══════════════════════════════════════════════════════════════════════

@triton.jit
def _reduce(
    output_ptr, segm_output_ptr, segm_lse_ptr,
    BLOCK_M: tl.constexpr, V_DIM: tl.constexpr, NUM_SEGMENTS: tl.constexpr,
):
    batch_idx = tl.program_id(0)
    head_idx = tl.program_id(1)
    offs_v = tl.arange(0, V_DIM)
    seg_offs = tl.arange(0, NUM_SEGMENTS)
    stat_base = batch_idx * NUM_SEGMENTS * BLOCK_M + head_idx
    segm_lse = tl.load(segm_lse_ptr + stat_base + seg_offs * BLOCK_M)
    overall_max = tl.max(segm_lse)
    weights = tl.exp(segm_lse - overall_max)
    overall_sum = tl.sum(weights)
    acc = tl.zeros([V_DIM], dtype=tl.float32)
    for s in range(NUM_SEGMENTS):
        o_off = (segm_output_ptr + batch_idx * NUM_SEGMENTS * BLOCK_M * V_DIM
                 + s * BLOCK_M * V_DIM + head_idx * V_DIM)
        seg_out = tl.load(o_off + offs_v)
        seg_lse = tl.load(segm_lse_ptr + batch_idx * NUM_SEGMENTS * BLOCK_M + s * BLOCK_M + head_idx)
        acc += seg_out * tl.exp(seg_lse - overall_max)
    result = tl.where(overall_sum == 0.0, 0.0, acc / overall_sum)
    out_off = output_ptr + batch_idx * BLOCK_M * V_DIM + head_idx * V_DIM
    tl.store(out_off + offs_v, result.to(tl.bfloat16))

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

@triton.jit
def _mla_splitk_pure_fp8(
    segm_output_ptr, segm_lse_ptr,
    q_bf16_ptr, kv_fp8_ptr, kv_indptr, kv_scale_ptr,
    BLOCK_M: tl.constexpr, TILE_SIZE: tl.constexpr,
    NOPE_DIM: tl.constexpr, ROPE_DIM: tl.constexpr,
    V_DIM: tl.constexpr, KV_STRIDE: tl.constexpr,
    NUM_SEGMENTS: tl.constexpr, SM_SCALE: tl.constexpr,
):
    batch_idx = tl.program_id(0)
    segm_idx = tl.program_id(1)

    kv_start = tl.load(kv_indptr + batch_idx)
    kv_end = tl.load(kv_indptr + batch_idx + 1)
    kv_len = kv_end - kv_start

    kv_s = tl.load(kv_scale_ptr)
    qk_s = kv_s * SM_SCALE

    tiles_per_segment = tl.cdiv(kv_len, NUM_SEGMENTS * TILE_SIZE)
    seg_tile_start = segm_idx * tiles_per_segment
    seg_tile_end = tl.minimum((segm_idx + 1) * tiles_per_segment, tl.cdiv(kv_len, TILE_SIZE))

    offs_m = tl.arange(0, BLOCK_M)
    offs_nope = tl.arange(0, NOPE_DIM)
    offs_rope = tl.arange(0, ROPE_DIM)
    offs_t = tl.arange(0, TILE_SIZE)

    q_base = q_bf16_ptr + batch_idx * BLOCK_M * KV_STRIDE
    Q_nope = tl.load(q_base + offs_m[:, None] * KV_STRIDE + offs_nope[None, :]).to(tl.float8e4nv)
    Q_rope = tl.load(q_base + offs_m[:, None] * KV_STRIDE + (NOPE_DIM + offs_rope)[None, :]).to(tl.float8e4nv)

    M_val = tl.full([BLOCK_M], -float("inf"), dtype=tl.float32)
    L = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
    acc = tl.zeros([BLOCK_M, V_DIM], dtype=tl.float32)

    for j in range(seg_tile_start, seg_tile_end):
        seq_offset = j * TILE_SIZE + offs_t
        tile_mask = seq_offset < kv_len
        kv_idx = kv_start + seq_offset

        K_nope = tl.load(kv_fp8_ptr + kv_idx[None, :] * KV_STRIDE + offs_nope[:, None],
                         mask=tile_mask[None, :], other=0.0)
        K_rope = tl.load(kv_fp8_ptr + kv_idx[None, :] * KV_STRIDE + (NOPE_DIM + offs_rope)[:, None],
                         mask=tile_mask[None, :], other=0.0)

        S = qk_s * (tl.dot(Q_nope, K_nope) + tl.dot(Q_rope, K_rope))
        S = tl.where(tile_mask[None, :], S, -float("inf"))

        m_j = tl.maximum(M_val, tl.max(S, axis=1))
        m_j = tl.where(m_j > -float("inf"), m_j, 0.0)
        P = tl.exp(S - m_j[:, None])
        l_j = tl.sum(P, axis=1)
        alpha = tl.exp(M_val - m_j)
        acc = acc * alpha[:, None]
        L = L * alpha + l_j
        M_val = m_j

        V = tl.trans(K_nope)
        acc += tl.dot(P.to(tl.float8e4nv), V)

    acc = acc * kv_s
    o_base = (segm_output_ptr + batch_idx * NUM_SEGMENTS * BLOCK_M * V_DIM
              + segm_idx * BLOCK_M * V_DIM)
    offs_v = tl.arange(0, V_DIM)
    tl.store(o_base + offs_m[:, None] * V_DIM + offs_v[None, :], acc / L[:, None])
    lse_base = segm_lse_ptr + batch_idx * NUM_SEGMENTS * BLOCK_M + segm_idx * BLOCK_M
    tl.store(lse_base + offs_m, M_val + tl.log(L))


def _run_fp8_splitk(q, kv_fp8, kv_scale, kv_indptr, config):
    bs, nh, vd = config["batch_size"], config["num_heads"], config["v_head_dim"]
    kvsl = config["kv_seq_len"]
    ns, tile, warps, stages = FP8_SPLITK_CONFIGS[(bs, kvsl)]
    kv_flat = kv_fp8.view(-1, QK_HEAD_DIM)
    o = torch.empty((bs, nh, vd), dtype=torch.bfloat16, device=q.device)
    mid_o = torch.empty((bs, ns, nh, vd), dtype=torch.float32, device=q.device)
    mid_lse = torch.empty((bs, ns, nh), dtype=torch.float32, device=q.device)
    _mla_splitk_pure_fp8[(bs, ns)](
        mid_o, mid_lse,
        q, kv_flat, kv_indptr, kv_scale,
        BLOCK_M=nh, TILE_SIZE=tile,
        NOPE_DIM=NOPE_DIM, ROPE_DIM=ROPE_DIM,
        V_DIM=vd, KV_STRIDE=QK_HEAD_DIM,
        NUM_SEGMENTS=ns, SM_SCALE=SM_SCALE,
        num_warps=warps, num_stages=stages,
        allow_flush_denorm=True)
    _reduce[(bs, nh)](
        o, mid_o, mid_lse,
        BLOCK_M=nh, V_DIM=vd, NUM_SEGMENTS=ns,
        num_warps=4, num_stages=1,
        allow_flush_denorm=True)
    return o

@triton.jit
def _mla_mk1g_fused(
    output_ptr,
    q_bf16_ptr, kv_fp8_ptr, kv_indptr, kv_scale_ptr,
    BLOCK_M: tl.constexpr, TILE_SIZE: tl.constexpr,
    NOPE_DIM: tl.constexpr, ROPE_DIM: tl.constexpr,
    V_DIM: tl.constexpr, KV_STRIDE: tl.constexpr,
    SM_SCALE: tl.constexpr, FP8_MAX: tl.constexpr,
):
    batch_idx = tl.program_id(0)

    kv_start = tl.load(kv_indptr + batch_idx)
    kv_end = tl.load(kv_indptr + batch_idx + 1)
    kv_len = kv_end - kv_start
    kv_s = tl.load(kv_scale_ptr)

    offs_m = tl.arange(0, BLOCK_M)
    offs_nope = tl.arange(0, NOPE_DIM)
    offs_rope = tl.arange(0, ROPE_DIM)
    offs_t = tl.arange(0, TILE_SIZE)

    q_base = q_bf16_ptr + batch_idx * BLOCK_M * KV_STRIDE
    Q_nope_bf16 = tl.load(q_base + offs_m[:, None] * KV_STRIDE + offs_nope[None, :])
    Q_rope_bf16 = tl.load(q_base + offs_m[:, None] * KV_STRIDE + (NOPE_DIM + offs_rope)[None, :])

    amax = tl.maximum(tl.max(tl.abs(Q_nope_bf16)), tl.max(tl.abs(Q_rope_bf16)))
    amax = tl.maximum(amax, 1e-12)
    q_scale = amax / FP8_MAX
    inv_q_scale = FP8_MAX / amax

    Q_nope = (Q_nope_bf16 * inv_q_scale).to(tl.float8e4nv)
    Q_rope = (Q_rope_bf16 * inv_q_scale).to(tl.float8e4nv)
    combined_scale = q_scale * kv_s * SM_SCALE

    M_val = tl.full([BLOCK_M], -float("inf"), dtype=tl.float32)
    L = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
    acc = tl.zeros([BLOCK_M, V_DIM], dtype=tl.float32)

    num_tiles = tl.cdiv(kv_len, TILE_SIZE)
    for j in range(0, num_tiles):
        seq_offset = j * TILE_SIZE + offs_t
        tile_mask = seq_offset < kv_len
        kv_idx = kv_start + seq_offset

        K_nope = tl.load(kv_fp8_ptr + kv_idx[None, :] * KV_STRIDE + offs_nope[:, None],
                         mask=tile_mask[None, :], other=0.0)
        K_rope = tl.load(kv_fp8_ptr + kv_idx[None, :] * KV_STRIDE + (NOPE_DIM + offs_rope)[:, None],
                         mask=tile_mask[None, :], other=0.0)

        S = combined_scale * (tl.dot(Q_nope, K_nope) + tl.dot(Q_rope, K_rope))
        S = tl.where(tile_mask[None, :], S, -float("inf"))

        V = tl.trans(K_nope)

        m_j = tl.maximum(M_val, tl.max(S, axis=1))
        m_j = tl.where(m_j > -float("inf"), m_j, 0.0)
        P = tl.exp(S - m_j[:, None])
        l_j = tl.sum(P, axis=1)
        alpha = tl.exp(M_val - m_j)
        acc = acc * alpha[:, None]
        L = L * alpha + l_j
        M_val = m_j

        acc += tl.dot(P.to(tl.float8e4nv), V)

    acc = (acc * kv_s) / L[:, None]
    out_off = output_ptr + batch_idx * BLOCK_M * V_DIM + offs_m[:, None] * V_DIM + tl.arange(0, V_DIM)[None, :]
    tl.store(out_off, acc.to(tl.bfloat16))


# ═══════════════════════════════════════════════════════════════════════
# Champ: mk19 pure fp8 splitk (bs32/8192, bs64/8192)
# Separate qk_scale/v_scale passed from CPU, no in-kernel Q scaling
# ═══════════════════════════════════════════════════════════════════════

@triton.jit
def _mla_splitk_mk19(
    segm_output_ptr, segm_lse_ptr,
    q_bf16_ptr, kv_fp8_ptr, kv_indptr,
    kv_scale_ptr,
    BLOCK_M: tl.constexpr, TILE_SIZE: tl.constexpr,
    NOPE_DIM: tl.constexpr, ROPE_DIM: tl.constexpr,
    V_DIM: tl.constexpr, KV_STRIDE: tl.constexpr,
    NUM_SEGMENTS: tl.constexpr, SM_SCALE: tl.constexpr,
):
    batch_idx = tl.program_id(0)
    segm_idx = tl.program_id(1)

    kv_start = tl.load(kv_indptr + batch_idx)
    kv_end = tl.load(kv_indptr + batch_idx + 1)
    kv_len = kv_end - kv_start

    kv_s = tl.load(kv_scale_ptr)
    qk_s = kv_s * SM_SCALE
    v_s = kv_s

    tiles_per_segment = tl.cdiv(kv_len, NUM_SEGMENTS * TILE_SIZE)
    seg_tile_start = segm_idx * tiles_per_segment
    seg_tile_end = tl.minimum((segm_idx + 1) * tiles_per_segment, tl.cdiv(kv_len, TILE_SIZE))

    offs_m = tl.arange(0, BLOCK_M)
    offs_nope = tl.arange(0, NOPE_DIM)
    offs_rope = tl.arange(0, ROPE_DIM)
    offs_t = tl.arange(0, TILE_SIZE)

    q_base = q_bf16_ptr + batch_idx * BLOCK_M * KV_STRIDE
    Q_nope = tl.load(q_base + offs_m[:, None] * KV_STRIDE + offs_nope[None, :]).to(tl.float8e4nv)
    Q_rope = tl.load(q_base + offs_m[:, None] * KV_STRIDE + (NOPE_DIM + offs_rope)[None, :]).to(tl.float8e4nv)

    M_val = tl.full([BLOCK_M], -float("inf"), dtype=tl.float32)
    L = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
    acc = tl.zeros([BLOCK_M, V_DIM], dtype=tl.float32)

    for j in range(seg_tile_start, seg_tile_end):
        seq_offset = j * TILE_SIZE + offs_t
        tile_mask = seq_offset < kv_len
        kv_idx = kv_start + seq_offset

        K_nope = tl.load(kv_fp8_ptr + kv_idx[None, :] * KV_STRIDE + offs_nope[:, None],
                         mask=tile_mask[None, :], other=0.0)
        K_rope = tl.load(kv_fp8_ptr + kv_idx[None, :] * KV_STRIDE + (NOPE_DIM + offs_rope)[:, None],
                         mask=tile_mask[None, :], other=0.0)

        S = qk_s * (tl.dot(Q_nope, K_nope) + tl.dot(Q_rope, K_rope))
        S = tl.where(tile_mask[None, :], S, -float("inf"))

        m_j = tl.maximum(M_val, tl.max(S, axis=1))
        m_j = tl.where(m_j > -float("inf"), m_j, 0.0)
        P = tl.exp(S - m_j[:, None])
        l_j = tl.sum(P, axis=1)
        alpha = tl.exp(M_val - m_j)
        acc = acc * alpha[:, None]
        L = L * alpha + l_j
        M_val = m_j

        V = tl.trans(K_nope)
        acc += tl.dot(P.to(tl.float8e4nv), V)

    acc = acc * v_s
    o_base = (segm_output_ptr + batch_idx * NUM_SEGMENTS * BLOCK_M * V_DIM
              + segm_idx * BLOCK_M * V_DIM)
    offs_v = tl.arange(0, V_DIM)
    tl.store(o_base + offs_m[:, None] * V_DIM + offs_v[None, :], acc / L[:, None])
    lse_base = segm_lse_ptr + batch_idx * NUM_SEGMENTS * BLOCK_M + segm_idx * BLOCK_M
    tl.store(lse_base + offs_m, M_val + tl.log(L))


# (ns, tile, warps, stages)
MK19_SPLITK_CONFIGS = {
    (32, 8192): (8, 64, 4, 2),
    (64, 8192): (4, 64, 8, 2),
}

def _run_mk19_splitk(q, kv_fp8, kv_scale, kv_indptr, config):
    bs, nh, vd = config["batch_size"], config["num_heads"], config["v_head_dim"]
    kvsl = config["kv_seq_len"]
    ns, tile, warps, stages = MK19_SPLITK_CONFIGS[(bs, kvsl)]
    kv_flat = kv_fp8.view(-1, QK_HEAD_DIM)
    o = torch.empty((bs, nh, vd), dtype=torch.bfloat16, device=q.device)
    mid_o = torch.empty((bs, ns, nh, vd), dtype=torch.float32, device=q.device)
    mid_lse = torch.empty((bs, ns, nh), dtype=torch.float32, device=q.device)
    _mla_splitk_mk19[(bs, ns)](
        mid_o, mid_lse,
        q, kv_flat, kv_indptr,
        kv_scale,
        BLOCK_M=nh, TILE_SIZE=tile,
        NOPE_DIM=NOPE_DIM, ROPE_DIM=ROPE_DIM,
        V_DIM=vd, KV_STRIDE=QK_HEAD_DIM,
        NUM_SEGMENTS=ns, SM_SCALE=SM_SCALE,
        num_warps=warps, num_stages=stages,
        allow_flush_denorm=True)
    _reduce[(bs, nh)](
        o, mid_o, mid_lse,
        BLOCK_M=nh, V_DIM=vd, NUM_SEGMENTS=ns,
        num_warps=4, num_stages=1,
        allow_flush_denorm=True)
    return o


# ═══════════════════════════════════════════════════════════════════════
# Champ: hip_v16_mk4j split-K fp8pv (bs256/8192)
# In-kernel Q quantization + split-K=2 for 2 waves/SIMD latency hiding
# ═══════════════════════════════════════════════════════════════════════

@triton.jit
def _mla_splitk_fp8pv(
    segm_output_ptr, segm_lse_ptr,
    q_bf16_ptr, kv_fp8_ptr, kv_indptr,
    kv_scale_ptr,
    BLOCK_M: tl.constexpr, TILE_SIZE: tl.constexpr,
    NOPE_DIM: tl.constexpr, ROPE_DIM: tl.constexpr,
    V_DIM: tl.constexpr, KV_STRIDE: tl.constexpr,
    NUM_SEGMENTS: tl.constexpr,
    SM_SCALE: tl.constexpr,
    FP8_MAX: tl.constexpr,
):
    batch_idx = tl.program_id(0)
    segm_idx = tl.program_id(1)

    kv_start = tl.load(kv_indptr + batch_idx)
    kv_end = tl.load(kv_indptr + batch_idx + 1)
    kv_len = kv_end - kv_start

    kv_s = tl.load(kv_scale_ptr)

    tiles_per_segment = tl.cdiv(kv_len, NUM_SEGMENTS * TILE_SIZE)
    seg_tile_start = segm_idx * tiles_per_segment
    seg_tile_end = tl.minimum((segm_idx + 1) * tiles_per_segment, tl.cdiv(kv_len, TILE_SIZE))

    offs_m = tl.arange(0, BLOCK_M)
    offs_nope = tl.arange(0, NOPE_DIM)
    offs_rope = tl.arange(0, ROPE_DIM)
    offs_t = tl.arange(0, TILE_SIZE)

    q_base = q_bf16_ptr + batch_idx * BLOCK_M * KV_STRIDE
    Q_nope_bf16 = tl.load(q_base + offs_m[:, None] * KV_STRIDE + offs_nope[None, :])
    Q_rope_bf16 = tl.load(q_base + offs_m[:, None] * KV_STRIDE + (NOPE_DIM + offs_rope)[None, :])

    amax_nope = tl.max(tl.abs(Q_nope_bf16))
    amax_rope = tl.max(tl.abs(Q_rope_bf16))
    amax = tl.maximum(amax_nope, amax_rope)
    amax = tl.maximum(amax, 1e-12)
    q_scale = amax / FP8_MAX
    inv_q_scale = FP8_MAX / amax

    Q_nope = (Q_nope_bf16 * inv_q_scale).to(tl.float8e4nv)
    Q_rope = (Q_rope_bf16 * inv_q_scale).to(tl.float8e4nv)

    combined_scale = q_scale * kv_s * SM_SCALE

    M_val = tl.full([BLOCK_M], -float("inf"), dtype=tl.float32)
    L = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
    acc = tl.zeros([BLOCK_M, V_DIM], dtype=tl.float32)

    for j in range(seg_tile_start, seg_tile_end):
        seq_offset = j * TILE_SIZE + offs_t
        tile_mask = seq_offset < kv_len
        kv_idx = kv_start + seq_offset

        K_nope = tl.load(kv_fp8_ptr + kv_idx[None, :] * KV_STRIDE + offs_nope[:, None],
                         mask=tile_mask[None, :], other=0.0)
        K_rope = tl.load(kv_fp8_ptr + kv_idx[None, :] * KV_STRIDE + (NOPE_DIM + offs_rope)[:, None],
                         mask=tile_mask[None, :], other=0.0)

        S = combined_scale * (tl.dot(Q_nope, K_nope) + tl.dot(Q_rope, K_rope))
        S = tl.where(tile_mask[None, :], S, -float("inf"))

        m_j = tl.maximum(M_val, tl.max(S, axis=1))
        m_j = tl.where(m_j > -float("inf"), m_j, 0.0)
        P = tl.exp(S - m_j[:, None])
        l_j = tl.sum(P, axis=1)
        alpha = tl.exp(M_val - m_j)
        acc = acc * alpha[:, None]
        L = L * alpha + l_j
        M_val = m_j

        V = tl.trans(K_nope)
        acc += tl.dot(P.to(tl.float8e4nv), V)

    acc = acc * kv_s
    o_base = (segm_output_ptr + batch_idx * NUM_SEGMENTS * BLOCK_M * V_DIM
              + segm_idx * BLOCK_M * V_DIM)
    offs_v = tl.arange(0, V_DIM)
    tl.store(o_base + offs_m[:, None] * V_DIM + offs_v[None, :], acc / L[:, None])
    lse_base = segm_lse_ptr + batch_idx * NUM_SEGMENTS * BLOCK_M + segm_idx * BLOCK_M
    tl.store(lse_base + offs_m, M_val + tl.log(L))


def _run_mk4j_splitk(q, kv_fp8, kv_scale, kv_indptr, config):
    bs, nh, vd = config["batch_size"], config["num_heads"], config["v_head_dim"]
    kv_flat = kv_fp8.view(-1, QK_HEAD_DIM)
    NS, TILE = 2, 64
    o = torch.empty((bs, nh, vd), dtype=torch.bfloat16, device=q.device)
    mid_o = torch.empty((bs, NS, nh, vd), dtype=torch.float32, device=q.device)
    mid_lse = torch.empty((bs, NS, nh), dtype=torch.float32, device=q.device)
    _mla_splitk_fp8pv[(bs, NS)](
        mid_o, mid_lse,
        q.view(bs, nh, QK_HEAD_DIM), kv_flat, kv_indptr,
        kv_scale.reshape(1),
        BLOCK_M=nh, TILE_SIZE=TILE,
        NOPE_DIM=NOPE_DIM, ROPE_DIM=ROPE_DIM,
        V_DIM=vd, KV_STRIDE=QK_HEAD_DIM,
        NUM_SEGMENTS=NS,
        SM_SCALE=SM_SCALE, FP8_MAX=FP8_MAX,
        num_warps=4, num_stages=2,
        allow_flush_denorm=True)
    _reduce[(bs, nh)](
        o, mid_o, mid_lse,
        BLOCK_M=nh, V_DIM=vd, NUM_SEGMENTS=NS,
        num_warps=4, num_stages=1,
        allow_flush_denorm=True)
    return o


def _run_mk1g_fused(q, kv_fp8, kv_scale, kv_indptr, config):
    bs, nh, vd = config["batch_size"], config["num_heads"], config["v_head_dim"]
    kv_flat = kv_fp8.view(-1, QK_HEAD_DIM)
    o = torch.empty((bs, nh, vd), dtype=torch.bfloat16, device=q.device)
    _mla_mk1g_fused[(bs,)](
        o, q, kv_flat, kv_indptr, kv_scale.reshape(1),
        BLOCK_M=nh, TILE_SIZE=64,
        NOPE_DIM=NOPE_DIM, ROPE_DIM=ROPE_DIM,
        V_DIM=vd, KV_STRIDE=QK_HEAD_DIM,
        SM_SCALE=SM_SCALE, FP8_MAX=FP8_MAX,
        num_warps=8, num_stages=2,
        allow_flush_denorm=True, 
        matrix_instr_nonkdim=16,)
    return o



@triton.jit
def _fp8_static_pack_kernel(
    x_ptr, y_ptr, kv_scale_ptr, q_scale_out_ptr,
    n_rows, n_cols, stride_x_row, stride_y_row,
    FP8_MAX_CST: tl.constexpr,
    BLOCK_ROWS: tl.constexpr, BLOCK_COLS: tl.constexpr,
):
    """Single-pass FP8 quant: prescale = FP8_MAX / (2 * kv_scale * FP8_MAX) = 1/(2*kv_scale)."""
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    kv_scale = tl.load(kv_scale_ptr)
    q_scale = kv_scale
    prescale = FP8_MAX_CST / (q_scale * FP8_MAX_CST)  # = 1 / (2 * kv_scale)
    if pid_m == 0 and pid_n == 0:
        tl.store(q_scale_out_ptr, q_scale)
    rows = pid_m * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)
    cols = pid_n * BLOCK_COLS + tl.arange(0, BLOCK_COLS)
    mask = (rows[:, None] < n_rows) & (cols[None, :] < n_cols)
    x = tl.load(x_ptr + rows[:, None] * stride_x_row + cols[None, :],
                mask=mask, other=0.0).to(tl.float32) * prescale
    x = tl.clamp(x, -FP8_MAX_CST, FP8_MAX_CST)
    HALF: tl.constexpr = BLOCK_COLS // 2
    x_even, x_odd = tl.split(x.reshape(BLOCK_ROWS, HALF, 2))
    packed = tl.inline_asm_elementwise(
        "v_cvt_pk_fp8_f32 $0, $1, $2;",
        constraints="=&v,v,v", args=[x_even, x_odd],
        dtype=tl.int16, is_pure=True, pack=1,
    )
    y_cols = pid_n * BLOCK_COLS // 2 + tl.arange(0, HALF)
    y_mask = (rows[:, None] < n_rows) & (y_cols[None, :] < n_cols // 2)
    tl.store(y_ptr + rows[:, None] * stride_y_row + y_cols[None, :],
             packed, mask=y_mask)


def _quantize_fp8_static(t, kv_scale):
    """Single-kernel FP8 quant using 2*kv_scale as static scale estimate."""
    orig_shape = t.shape
    t2d = t.reshape(-1, orig_shape[-1])
    M, N = t2d.shape
    BR, BC = 32, 32
    grid = (triton.cdiv(M, BR), triton.cdiv(N, BC))
    q_scale = torch.empty(1, dtype=torch.float32, device=t.device)
    out_i16 = torch.empty(M, N // 2, dtype=torch.int16, device=t.device)
    _fp8_static_pack_kernel[grid](
        t2d, out_i16, kv_scale, q_scale, M, N,
        t2d.stride(0), out_i16.stride(0),
        FP8_MAX_CST=FP8_MAX,
        BLOCK_ROWS=BR, BLOCK_COLS=BC, num_warps=1,
    )
    return out_i16.view(torch.uint8).view(FP8_DTYPE).reshape(orig_shape), q_scale

PAGE_SIZE = 1
NUM_KV_SPLITS_AITER = 1024

def _make_mla_decode_metadata(batch_size, max_q_len, nhead, nhead_kv,
                               q_dtype, kv_dtype, qo_indptr, kv_indptr,
                               kv_last_page_len, num_kv_splits=NUM_KV_SPLITS_AITER):
    info = get_mla_metadata_info_v1(batch_size, max_q_len, nhead, 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]
    (wm, wi, wis, ri, rfm, rpm) = work
    get_mla_metadata_v1(qo_indptr, kv_indptr, kv_last_page_len,
        nhead // nhead_kv, nhead_kv, True, wm, wis, wi, ri, rfm, rpm,
        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_meta_data": wm, "work_indptr": wi, "work_info_set": wis,
            "reduce_indptr": ri, "reduce_final_map": rfm, "reduce_partial_map": rpm}


def reference_fallback(data):
    q, kv_data, qo_indptr, kv_indptr, config = data
    kv_fp8, kv_scale = kv_data["fp8"]
    q_fp8, q_scale = _quantize_fp8_static(q, kv_scale)
    bs = config["batch_size"]
    nq, nkv, dq, dv = config["num_heads"], config["num_kv_heads"], config["qk_head_dim"], config["v_head_dim"]
    total_kv = int(kv_indptr[-1].item())
    kv_indices = torch.arange(total_kv, dtype=torch.int32, device=q.device)
    kv_4d = kv_fp8.view(kv_fp8.shape[0], PAGE_SIZE, nkv, kv_fp8.shape[-1])
    kv_last = (kv_indptr[1:] - kv_indptr[:-1]).to(torch.int32)
    meta = _make_mla_decode_metadata(bs, config["q_seq_len"], nq, nkv,
        q_fp8.dtype, kv_fp8.dtype, qo_indptr, kv_indptr, kv_last)
    o = torch.empty((q.shape[0], nq, dv), dtype=torch.bfloat16, device=q.device)
    mla_decode_fwd(q_fp8.view(-1, nq, dq), kv_4d, o, qo_indptr, kv_indptr, kv_indices,
        kv_last, config["q_seq_len"], page_size=PAGE_SIZE, nhead_kv=nkv,
        sm_scale=config["sm_scale"], logit_cap=0.0,
        num_kv_splits=NUM_KV_SPLITS_AITER, q_scale=q_scale, kv_scale=kv_scale,
        intra_batch_mode=True, **meta)
    return o

def custom_kernel(data: input_t) -> output_t:
    q, kv_data, qo_indptr, kv_indptr, config = data
    bs = config["batch_size"]
    kvsl = config["kv_seq_len"]
    key = (bs, kvsl)

    # 4x1024, 4x8192, 32x1024, 64x1024: fp8 splitk
    if key in FP8_SPLITK_CONFIGS:
        kv_fp8, kv_scale = kv_data["fp8"]
        return _run_fp8_splitk(q, kv_fp8, kv_scale, kv_indptr, config)

    # 256x1024: mk1g fused (in-kernel Q quant, tile=64, warps=8)
    if key == (256, 1024):
        kv_fp8, kv_scale = kv_data["fp8"]
        return _run_mk1g_fused(q, kv_fp8, kv_scale, kv_indptr, config)

    # kv8192 bs>32: aiter path
    if kvsl == 8192 and bs > 32:
        kv_fp8, kv_scale = kv_data["fp8"]
        return _run_aiter(q, kv_fp8, kv_scale, qo_indptr, kv_indptr, config)

    # 32x8192: mk19f pure fp8 splitk (champ, 74.4µs vs 94.7µs aiter)
    # 64x8192: mk1a pure fp8 splitk (champ, 124.3µs vs 139.1µs aiter)
    if key in MK19_SPLITK_CONFIGS:
        kv_fp8, kv_scale = kv_data["fp8"]
        return _run_mk19_splitk(q, kv_fp8, kv_scale, kv_indptr, config)

    # 256x8192: hip_v16_mk4j split-K=2 fp8pv (champ, 319.4µs vs 320.5µs aiter)
    if key == (256, 8192):
        kv_fp8, kv_scale = kv_data["fp8"]
        return _run_mk4j_splitk(q, kv_fp8, kv_scale, kv_indptr, config)

    return reference_fallback(data)

# ═══════════════════════════════════════════════════════════════════════
# Aiter MLA decode path (kv8192, bs>4)
# ═══════════════════════════════════════════════════════════════════════

def _build_aiter_meta(bs, kv_len, q_len, nq, nkv, qo_indptr, kv_indptr, ps, ns):
    num_pages = (bs * kv_len) // ps
    kv_indptr_paged = kv_indptr // ps
    seq_lens = kv_indptr[1:] - kv_indptr[:-1]
    last_page_len = (seq_lens % ps).to(torch.int32)
    last_page_len = torch.where(last_page_len == 0, ps, last_page_len)
    kv_gran = max(1, 16 // ps)
    info = get_mla_metadata_info_v1(
        bs, q_len, nq, FP8_DTYPE, FP8_DTYPE,
        is_sparse=False, fast_mode=False,
        num_kv_splits=ns, intra_batch_mode=True,
    )
    work_meta, work_indptr, work_info, reduce_indptr, reduce_final, reduce_partial = \
        [torch.empty(s, dtype=t, device="cuda") for s, t in info]
    get_mla_metadata_v1(
        qo_indptr, kv_indptr_paged, last_page_len,
        nq // nkv, nkv, True,
        work_meta, work_info, work_indptr,
        reduce_indptr, reduce_final, reduce_partial,
        page_size=ps, kv_granularity=kv_gran,
        max_seqlen_qo=q_len, uni_seqlen_qo=q_len,
        fast_mode=False, max_split_per_batch=ns,
        intra_batch_mode=True, dtype_q=FP8_DTYPE, dtype_kv=FP8_DTYPE,
    )
    kv_indices = torch.arange(num_pages, dtype=torch.int32, device="cuda")
    return dict(
        work_meta_data=work_meta, work_indptr=work_indptr, work_info_set=work_info,
        reduce_indptr=reduce_indptr, reduce_final_map=reduce_final, reduce_partial_map=reduce_partial,
        kv_indices=kv_indices, last_page_len=last_page_len, kv_indptr_paged=kv_indptr_paged,
    )


def _run_aiter(q, kv_fp8, kv_scale, qo_indptr, kv_indptr, config):
    bs = config["batch_size"]
    nq, nkv = config["num_heads"], config["num_kv_heads"]
    dq, dv = config["qk_head_dim"], config["v_head_dim"]
    q_len = config["q_seq_len"]
    kv_len = config["kv_seq_len"]
    ps = 8
    ns = 32

    meta = _build_aiter_meta(bs, kv_len, q_len, nq, nkv, qo_indptr, kv_indptr, ps, ns)

    q_fp8, q_scale = _quantize_fp8_static(q, kv_scale)
    kv_4d = kv_fp8.view(-1, ps, nkv, kv_fp8.shape[-1])
    output = torch.empty((q.shape[0], nq, dv), dtype=torch.bfloat16, device=q.device)

    mla_decode_fwd(
        q_fp8.view(q.shape[0], nq, dq), kv_4d, output,
        qo_indptr, meta["kv_indptr_paged"], meta["kv_indices"], meta["last_page_len"],
        q_len, page_size=ps, nhead_kv=nkv,
        sm_scale=config["sm_scale"], logit_cap=0.0, num_kv_splits=ns,
        q_scale=q_scale, kv_scale=kv_scale,
        intra_batch_mode=True,
        work_meta_data=meta["work_meta_data"], work_indptr=meta["work_indptr"],
        work_info_set=meta["work_info_set"], reduce_indptr=meta["reduce_indptr"],
        reduce_final_map=meta["reduce_final_map"], reduce_partial_map=meta["reduce_partial_map"],
    )
    return output