#!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/tmpb2b5w25o
AUTOTUNE mm(16512x512, 512x32768)
strides: [512, 1], [1, 512]
dtypes: torch.bfloat16, torch.bfloat16
  triton_mm_137 1.1950 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.2711 ms 94.0% 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
  mm 1.3184 ms 90.6% 
  triton_mm_136 1.3348 ms 89.5% 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_135 1.3512 ms 88.4% 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_134 1.3582 ms 88.0% 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
  triton_mm_131 1.3796 ms 86.6% 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.4330 ms 83.4% 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_129 1.4774 ms 80.9% 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.4922 ms 80.1% 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.4288 seconds and 0.0004 seconds precompiling for 37 choices
AUTOTUNE mm(128x7168, 7168x1536)
strides: [7168, 1], [1, 7168]
dtypes: torch.bfloat16, torch.bfloat16
  triton_mm_1 0.0295 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.0360 ms 82.0% 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 75.5% 
  triton_mm_5 0.0405 ms 73.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.0443 ms 66.7% 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.0444 ms 66.5% 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.0474 ms 62.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_8 0.0501 ms 59.0% 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.0692 ms 42.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_22 0.0696 ms 42.5% 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.1288 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_55 0.0338 ms 100.0% 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_62 0.0343 ms 98.5% 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_61 0.0371 ms 91.0% 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_67 0.0378 ms 89.3% 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_60 0.0382 ms 88.4% 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 87.7% 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.0423 ms 79.9% 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.0427 ms 79.2% 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.0432 ms 78.2% 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_54 0.0435 ms 77.6% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=128, BLOCK_N=64, 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
SingleProcess AUTOTUNE benchmarking takes 4.0731 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.0282 ms 82.4% 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_71 0.0335 ms 69.3% 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_77 0.0338 ms 68.6% 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_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
  triton_mm_75 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=4
  mm 0.0407 ms 57.0% 
  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.2% 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.0504 ms 46.1% 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.0783 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.0836 ms 100.0% 
  triton_mm_150 0.1465 ms 57.1% 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_151 0.1532 ms 54.6% 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.1551 ms 53.9% 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_145 0.1587 ms 52.7% 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_147 0.1597 ms 52.4% 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_143 0.1609 ms 52.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_154 0.1628 ms 51.4% 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.1645 ms 50.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_157 0.1697 ms 49.3% 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.4646 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])
