#!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 leaderboard /tmp/tmp0gmai1p7
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])
