#!POPCORN leaderboard amd-mla-decode
#!POPCORN gpus MI300

import math
import torch
import torch.nn.functional as F


def precompute_basis(
    seq_len: int, head_dim: int, theta0: float = 10000
) -> tuple[torch.Tensor, torch.Tensor]:
    freq = theta0 ** (
        -torch.arange(0, head_dim, 2, dtype=torch.bfloat16, device="cuda") / head_dim
    )
    pos = torch.arange(seq_len, dtype=torch.bfloat16, device="cuda")
    theta = torch.outer(pos, freq)

    cos = torch.cos(theta).repeat((1, 2))
    sin = torch.sin(theta).repeat((1, 2))

    return cos, sin


def rope(
    x: torch.Tensor, basis: tuple[torch.Tensor, torch.Tensor], pos: int = 0
) -> torch.Tensor:
    cos, sin = basis
    x1, x2 = x.chunk(2, dim=-1)
    return (
        x * cos[pos : pos + x.shape[-2]]
        + torch.cat((-x2, x1), dim=-1) * sin[pos : pos + x.shape[-2]]
    )


rope = torch.compile(rope, mode="max-autotune")


@torch.no_grad
def compute(
    x: torch.Tensor,
    kv_cache: torch.Tensor,
    prev_len: int,
    Q_proj_down_weight: torch.Tensor,
    Q_proj_up_weight: torch.Tensor,
    KV_proj_down_weight: torch.Tensor,
    KV_proj_up_weight: torch.Tensor,
    wo_weight: torch.Tensor,
):
    # Constants
    nh = 128
    dq_nope = 128
    dq_rope = 64
    dkv = 512
    dk_nope = 128
    dk_rope = 64
    dv = 128

    B, C, d = x.shape
    total_len = prev_len + C
    inv_sqrt_dk = 1.0 / math.sqrt(dk_nope + dq_rope)

    # Fused Q projection and reshape - avoid intermediate tensor
    q_down = torch.mm(x.view(-1, d), Q_proj_down_weight.T)  # (B*C, Q_down_dim)
    q_up = torch.mm(q_down, Q_proj_up_weight.T).view(B, C, nh, dq_nope + dq_rope)
    q_up = q_up.transpose(1, 2)  # B, nh, C, dq_nope + dq_rope

    # Split Q without creating new tensors
    q_nope = q_up[..., :dq_nope]
    q_rope = q_up[..., dq_nope:]

    # Apply RoPE to q_rope in-place if possible
    q_rope = rope(q_rope, basis, prev_len)

    # Concatenate q without intermediate allocation
    q = torch.cat([q_nope, q_rope], dim=-1)

    # KV projection - fused operations
    kv_down = torch.mm(x.view(-1, d), KV_proj_down_weight.T).view(B, C, -1)

    # Update cache in-place
    kv_cache[:, prev_len:total_len] = kv_down

    # Get full context from cache
    kv_full = kv_cache[:, :total_len]  # B, total_len, dkv + dk_rope

    # Split KV cache efficiently
    kv_down_part = kv_full[..., :dkv]
    k_rope_full = kv_full[..., dkv:]

    # Fused KV up projection
    kv_up = torch.mm(kv_down_part.reshape(-1, dkv), KV_proj_up_weight.T).view(
        B, total_len, nh, dk_nope + dv
    )
    kv_up = kv_up.transpose(1, 2)  # B, nh, total_len, dk_nope + dv

    # Split K and V
    k_nope = kv_up[..., :dk_nope]
    v = kv_up[..., dk_nope:]

    # Apply RoPE to k_rope for full sequence
    k_rope_rotated = rope(k_rope_full, basis, 0)  # Apply RoPE from position 0
    k_rope_expanded = k_rope_rotated.unsqueeze(1).expand(-1, nh, -1, -1)

    # Concatenate K
    k = torch.cat([k_nope, k_rope_expanded], dim=-1)

    # Fused attention computation
    # Use torch.baddbmm for better performance when possible
    B_nh = B * nh
    q_flat = q.reshape(B_nh, C, -1)
    k_flat = k.reshape(B_nh, total_len, -1)
    v_flat = v.reshape(B_nh, total_len, dv)

    # Compute attention scores
    scores = torch.bmm(q_flat, k_flat.transpose(-1, -2)) * inv_sqrt_dk

    # Apply softmax
    attn_weights = F.softmax(scores, dim=-1)

    # Apply attention to values
    y_flat = torch.bmm(attn_weights, v_flat)  # B_nh, C, dv

    # Reshape and apply output projection
    y = y_flat.reshape(B, nh, C, dv).transpose(1, 2).reshape(B, C, -1)

    # Final output projection
    out = torch.mm(y.view(-1, nh * dv), wo_weight.T).view(B, C, -1)

    return out


compute = torch.compile(compute, mode="max-autotune")

basis = precompute_basis(6145, 64)


def custom_kernel(data: tuple) -> tuple:
    config, x, kv_cache = data
    prev_len = kv_cache.seq_len
    out = compute(
        x,
        kv_cache.data,
        prev_len,
        config.Q_proj_down_weight,
        config.Q_proj_up_weight,
        config.KV_proj_down_weight,
        config.KV_proj_up_weight,
        config.wo_weight,
    )
    return out, kv_cache.data

python eval.py test /tmp/tmpjufyux1m
AUTOTUNE mm(16512x512, 512x32768)
strides: [512, 1], [1, 512]
dtypes: torch.bfloat16, torch.bfloat16
  triton_mm_137 1.1974 ms 100.0% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=256, BLOCK_N=256, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_128 1.2624 ms 94.9% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=128, BLOCK_N=128, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_134 1.3001 ms 92.1% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=256, BLOCK_N=128, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=2, num_stages=2, num_warps=4
  mm 1.3291 ms 90.1% 
  triton_mm_135 1.3368 ms 89.6% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=256, BLOCK_N=128, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_131 1.3627 ms 87.9% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=128, BLOCK_N=256, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=2, num_stages=2, num_warps=4
  triton_mm_125 1.3795 ms 86.8% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=128, BLOCK_N=128, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=2, num_stages=2, num_warps=4
  triton_mm_136 1.3922 ms 86.0% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=256, BLOCK_N=128, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_129 1.4349 ms 83.5% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=128, BLOCK_N=128, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_130 1.4814 ms 80.8% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=128, BLOCK_N=128, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
SingleProcess AUTOTUNE benchmarking takes 7.3810 seconds and 0.0005 seconds precompiling for 37 choices
AUTOTUNE mm(128x7168, 7168x1536)
strides: [7168, 1], [1, 7168]
dtypes: torch.bfloat16, torch.bfloat16
  triton_mm_1 0.0293 ms 100.0% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=256, BLOCK_M=32, BLOCK_N=16, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=2
  triton_mm_3 0.0359 ms 81.7% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=32, BLOCK_N=32, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  mm 0.0391 ms 74.9% 
  triton_mm_5 0.0407 ms 72.0% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=64, BLOCK_N=16, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=2, num_stages=2, num_warps=4
  triton_mm_12 0.0442 ms 66.3% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=64, BLOCK_N=64, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_9 0.0449 ms 65.3% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=64, BLOCK_N=32, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_13 0.0473 ms 61.9% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=256, BLOCK_M=64, BLOCK_N=64, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_8 0.0498 ms 58.9% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=64, BLOCK_N=32, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_17 0.0694 ms 42.2% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=64, BLOCK_N=128, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_22 0.0696 ms 42.1% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=128, BLOCK_N=64, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
SingleProcess AUTOTUNE benchmarking takes 4.1384 seconds and 0.0003 seconds precompiling for 35 choices
AUTOTUNE mm(128x1536, 1536x24576)
strides: [1536, 1], [1, 1536]
dtypes: torch.bfloat16, torch.bfloat16
  triton_mm_62 0.0344 ms 100.0% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=128, BLOCK_N=128, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_55 0.0345 ms 99.6% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=128, BLOCK_N=64, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_67 0.0367 ms 93.7% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=128, BLOCK_N=128, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_61 0.0380 ms 90.6% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=128, BLOCK_N=128, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_60 0.0381 ms 90.3% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=128, BLOCK_N=128, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_45 0.0385 ms 89.3% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=64, BLOCK_N=64, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_53 0.0416 ms 82.7% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=128, BLOCK_N=32, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_65 0.0425 ms 80.9% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=128, BLOCK_N=64, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_56 0.0434 ms 79.3% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=128, BLOCK_N=64, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_59 0.0436 ms 78.9% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=128, BLOCK_N=128, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=1, num_stages=2, num_warps=8
SingleProcess AUTOTUNE benchmarking takes 4.0769 seconds and 0.0003 seconds precompiling for 35 choices
AUTOTUNE mm(128x7168, 7168x576)
strides: [7168, 1], [1, 7168]
dtypes: torch.bfloat16, torch.bfloat16
  triton_mm_69 0.0232 ms 100.0% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=256, BLOCK_M=32, BLOCK_N=16, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=2
  triton_mm_73 0.0280 ms 83.0% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=64, BLOCK_N=16, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=2, num_stages=2, num_warps=4
  triton_mm_77 0.0336 ms 69.0% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=64, BLOCK_N=32, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_71 0.0339 ms 68.5% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=32, BLOCK_N=32, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_75 0.0348 ms 66.8% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=64, BLOCK_N=32, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_76 0.0350 ms 66.4% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=64, BLOCK_N=32, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  mm 0.0435 ms 53.4% 
  triton_mm_80 0.0440 ms 52.8% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=64, BLOCK_N=64, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_81 0.0471 ms 49.3% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=256, BLOCK_M=64, BLOCK_N=64, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_87 0.0500 ms 46.5% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=128, BLOCK_N=32, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
SingleProcess AUTOTUNE benchmarking takes 4.1177 seconds and 0.0003 seconds precompiling for 35 choices
AUTOTUNE mm(128x16384, 16384x7168)
strides: [16384, 1], [1, 16384]
dtypes: torch.bfloat16, torch.bfloat16
  mm 0.0814 ms 100.0% 
  triton_mm_151 0.1535 ms 53.0% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=256, BLOCK_M=64, BLOCK_N=64, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_155 0.1544 ms 52.7% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=64, BLOCK_N=128, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_154 0.1625 ms 50.1% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=64, BLOCK_N=128, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_169 0.1633 ms 49.8% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=128, BLOCK_N=64, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_143 0.1650 ms 49.3% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=64, BLOCK_N=16, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=2, num_stages=2, num_warps=4
  triton_mm_150 0.1653 ms 49.2% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=64, BLOCK_N=64, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_147 0.1655 ms 49.2% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=64, BLOCK_N=32, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_160 0.1701 ms 47.9% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=128, BLOCK_N=64, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_145 0.1756 ms 46.3% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=64, BLOCK_N=32, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
SingleProcess AUTOTUNE benchmarking takes 4.4528 seconds and 0.0004 seconds precompiling for 35 choices
skipping cudagraphs due to mutated inputs (1 instances)
skipping cudagraphs due to mutated inputs (1 instances)
skipping cudagraphs due to mutated inputs (1 instances)
skipping cudagraphs due to mutated inputs (1 instances)

Buff name: data, shape: torch.Size([128, 8192, 576])
Buff name: data, shape: torch.Size([128, 8192, 576])
Buff name: data, shape: torch.Size([128, 8192, 576])
Buff name: data, shape: torch.Size([128, 8192, 576])
