import math
from dataclasses import dataclass
import torch
from torch import nn  # KVCache inherits from nn.Module
import torch.nn.functional as F
from task import input_t, output_t  # Assuming these are correctly defined for the environment

import triton
import triton.language as tl


# KVCache class (consistent with problem context)
class KVCache(nn.Module):
    def __init__(self, kv_cache_shape: tuple) -> None:
        super().__init__()
        # buffer 'data' will store KV pairs: (batch_size, max_seq_len, kv_lora_rank + qk_rope_head_dim)
        self.register_buffer(
            "data", torch.zeros(kv_cache_shape, dtype=torch.bfloat16, device="cuda")
        )
        self.seq_len = 0  # current number of tokens in the cache
        self.zero()  # Initialize buffer to zeros

    def zero(self) -> None:
        self.data.zero_()

    def get_data(self) -> torch.Tensor:
        # Returns the full buffer.
        return self.data

    def forward(self, c_kv: torch.Tensor) -> tuple[torch.Tensor, int]:
        # c_kv: new KV data for the current token(s)
        # shape: (batch_size, num_new_tokens, kv_lora_rank + qk_rope_head_dim)
        # For decode, num_new_tokens is 1.

        num_new_tokens = c_kv.size(1)

        if self.seq_len + num_new_tokens > self.data.size(1):  # self.data.size(1) is max_seq_len
            raise ValueError(
                f"KV Cache Exceeded: current_len={self.seq_len}, new_tokens={num_new_tokens}, max_len={self.data.size(1)}"
            )

        if self.data.dtype != c_kv.dtype:
            self.data = self.data.to(c_kv.dtype)

        # Update the cache with new data
        self.data[:, self.seq_len : self.seq_len + num_new_tokens, :] = c_kv
        self.seq_len += num_new_tokens

        # Return the relevant part of the cache (up to new length) and the new total sequence length
        return self.data[:, : self.seq_len, :], self.seq_len


# Config dataclass (as defined in the problem, holding shapes and weights)
@dataclass
class Config:
    batch_size: int
    dim: int
    n_heads: int
    q_lora_rank: int
    kv_lora_rank: int
    qk_nope_head_dim: int
    qk_rope_head_dim: int
    v_head_dim: int
    seq_len: int  # Sequence length of input x, typically 1 for decode
    max_seq_len: int  # Max capacity of KV cache
    kv_cache_shape: tuple  # Full shape of the KV cache buffer
    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


# --- RoPE Triton Kernel and utilities ---

_rope_cos_sin_cache = {}


def get_rope_cos_sin_matrices(d_model: int, max_seq_len: int, device="cuda", dtype=torch.bfloat16):
    """
    Cache and return RoPE cosine and sine matrices.
    cos_cached, sin_cached shapes: (max_seq_len, d_model)
    """
    cache_key = (d_model, max_seq_len, str(device), dtype)  # Device needs to be string for dict key
    if cache_key not in _rope_cos_sin_cache:
        theta = 10000.0 ** (
            -torch.arange(0, d_model // 2, dtype=dtype, device=device) / (d_model // 2)
        )
        positions = torch.arange(max_seq_len, dtype=dtype, device=device)
        # angles shape: (max_seq_len, d_model/2)
        angles = torch.einsum("i,j->ij", positions, theta)
        # angles_full shape: (max_seq_len, d_model) by duplicating for both halves
        angles_full = torch.cat([angles, angles], dim=-1)

        cos_cached = angles_full.cos()
        sin_cached = angles_full.sin()
        _rope_cos_sin_cache[cache_key] = (cos_cached, sin_cached)
    return _rope_cos_sin_cache[cache_key]


@triton.jit
def rope_triton_kernel(
    X_ptr,
    Cos_ptr,
    Sin_ptr,
    Out_ptr,
    stride_x_batch,
    stride_x_head,
    stride_x_seq,
    stride_x_dim,
    stride_cos_seq,
    stride_cos_dim,
    stride_sin_seq,
    stride_sin_dim,
    stride_out_batch,
    stride_out_head,
    stride_out_seq,
    stride_out_dim,
    D_ROPE: tl.constexpr,
    D_ROPE_HALF: tl.constexpr,
    BLOCK_SIZE_D: tl.constexpr,
):
    pid_batch = tl.program_id(0)
    pid_head = tl.program_id(1)
    pid_seq = tl.program_id(
        2
    )  # This corresponds to the sequence dimension of the input tensor X to RoPE

    # Base pointers for the current element in X and Out
    x_base_ptr = (
        X_ptr + pid_batch * stride_x_batch + pid_head * stride_x_head + pid_seq * stride_x_seq
    )
    out_base_ptr = (
        Out_ptr
        + pid_batch * stride_out_batch
        + pid_head * stride_out_head
        + pid_seq * stride_out_seq
    )

    # Base pointers for Cos and Sin. Note that Cos_ptr and Sin_ptr are already sliced by start_pos.
    # So pid_seq here directly indexes into the (already correctly offset) cos/sin matrices for the current token.
    cos_row_ptr = Cos_ptr + pid_seq * stride_cos_seq
    sin_row_ptr = Sin_ptr + pid_seq * stride_sin_seq

    offs_d = tl.arange(0, BLOCK_SIZE_D)
    mask_d = offs_d < D_ROPE

    # Load current x, cos, sin values for the current (batch, head, seq_idx_in_x)
    x_vals = tl.load(x_base_ptr + offs_d * stride_x_dim, mask=mask_d, other=0.0)
    cos_vals = tl.load(
        cos_row_ptr + offs_d * stride_cos_dim, mask=mask_d, other=0.0
    )  # cos for current position
    sin_vals = tl.load(
        sin_row_ptr + offs_d * stride_sin_dim, mask=mask_d, other=0.0
    )  # sin for current position

    # Calculate indices for the "other half" for rotation
    # If offs_d is in the first half, other_half_offs is in the second half.
    # If offs_d is in the second half, other_half_offs is in the first half.
    other_half_offs = tl.where(offs_d < D_ROPE_HALF, offs_d + D_ROPE_HALF, offs_d - D_ROPE_HALF)

    # Load the "other half" of x
    x_other_half_vals = tl.load(x_base_ptr + other_half_offs * stride_x_dim, mask=mask_d, other=0.0)

    # Perform the rotation: (-x2, x1) part
    # If offs_d < D_ROPE_HALF (first half, x1), rotated part is -x2.
    # If offs_d >= D_ROPE_HALF (second half, x2), rotated part is x1.
    x_rotated_vals = tl.where(offs_d < D_ROPE_HALF, -x_other_half_vals, x_other_half_vals)

    # Apply RoPE: x_out = x * cos + x_rotated * sin
    out_vals = x_vals * cos_vals + x_rotated_vals * sin_vals

    tl.store(out_base_ptr + offs_d * stride_out_dim, out_vals, mask=mask_d)


def apply_rope_triton(
    x: torch.Tensor, cos_cached: torch.Tensor, sin_cached: torch.Tensor, start_pos: int
):
    output = torch.empty_like(x)  # Output tensor with the same shape as input x

    batch_size, num_heads, seq_len_x, d_rope = (
        x.shape
    )  # x is (bs, n_heads, seq_len_of_this_tensor, rope_dim)
    assert d_rope % 2 == 0, "D_ROPE must be even"
    d_rope_half = d_rope // 2

    # Slice the cached cos/sin matrices based on start_pos and the sequence length of x
    # current_cos/sin will have shape (seq_len_x, d_rope)
    current_cos = cos_cached[start_pos : start_pos + seq_len_x, :]
    current_sin = sin_cached[start_pos : start_pos + seq_len_x, :]

    # Triton grid: one program instance per (batch, head, token_in_x_sequence)
    grid = (batch_size, num_heads, seq_len_x)

    # Determine BLOCK_SIZE_D for Triton kernel based on d_rope
    # This ensures that the entire d_rope dimension is processed in one go by each thread block instance along that dim.
    if d_rope == 64:
        BLOCK_SIZE_D = 64
    elif d_rope == 128:  # For future flexibility, although problem uses 64
        BLOCK_SIZE_D = 128
    else:  # General case, ensure it's a power of 2 and at least 16 for Triton.
        # Using triton.next_power_of_2 might lead to BLOCK_SIZE_D > d_rope if d_rope is not power of 2.
        # It's better to ensure BLOCK_SIZE_D is exactly d_rope if d_rope is a common power of 2,
        # or make it cover d_rope. For this problem d_rope=64 is fixed.
        BLOCK_SIZE_D = (
            d_rope if d_rope in [16, 32, 64, 128, 256] else triton.next_power_of_2(max(16, d_rope))
        )

    rope_triton_kernel[grid](
        x,
        current_cos,
        current_sin,
        output,
        x.stride(0),
        x.stride(1),
        x.stride(2),
        x.stride(3),  # Strides for X_ptr
        current_cos.stride(0),
        current_cos.stride(1),  # Strides for Cos_ptr
        current_sin.stride(0),
        current_sin.stride(1),  # Strides for Sin_ptr
        output.stride(0),
        output.stride(1),
        output.stride(2),
        output.stride(3),  # Strides for Out_ptr
        D_ROPE=d_rope,
        D_ROPE_HALF=d_rope_half,
        BLOCK_SIZE_D=BLOCK_SIZE_D,  # Triton processes this many elements along the d_rope dimension
    )
    return output


# --- Main MLA forward pass ---


def mla_forward(
    x: torch.Tensor, kv_cache_module: KVCache, config: Config
) -> tuple[torch.Tensor, torch.Tensor]:
    batch_size, seq_len_q, _ = x.size()  # x: (bs, 1, dim)

    # Get precomputed RoPE cosine and sine matrices from cache
    cos_cached, sin_cached = get_rope_cos_sin_matrices(
        config.qk_rope_head_dim, config.max_seq_len, device=x.device, dtype=x.dtype
    )

    # Step 1: Down-projection & KV cache update
    q_lora = F.linear(x, config.Q_proj_down_weight)  # (bs, 1, q_lora_rank)
    kv_lora_current_token = F.linear(
        x, config.KV_proj_down_weight
    )  # (bs, 1, kv_lora_rank + rope_dim)

    # kv_cache_module.forward updates its internal state (self.data, self.seq_len)
    # and returns the slice of kv_lora up to current length, and current total length kv_len.
    kv_lora_full_history, kv_len = kv_cache_module(kv_lora_current_token)
    # kv_lora_full_history: (bs, kv_len, kv_lora_rank + rope_dim)
    # kv_len: new total sequence length in cache (prefill + 1 for decode)

    query_pos = kv_len - 1  # Current token's position for RoPE (0-indexed)

    # Step 2: Up-projection
    # Queries Q
    q_nope_and_rope = F.linear(q_lora, config.Q_proj_up_weight).view(
        batch_size, seq_len_q, config.n_heads, config.qk_nope_head_dim + config.qk_rope_head_dim
    )  # (bs, 1, n_heads, nope_dim + rope_dim)
    q_nope, q_rope_tensor = q_nope_and_rope.split(
        [config.qk_nope_head_dim, config.qk_rope_head_dim], dim=-1
    )

    # Keys K and Values V from full history
    kv_lora_rank_part_history, k_rope_tensor_history = kv_lora_full_history.split(
        [config.kv_lora_rank, config.qk_rope_head_dim], dim=-1
    )  # k_rope_tensor_history: (bs, kv_len, rope_dim)

    kv_nope_and_v_parts = F.linear(kv_lora_rank_part_history, config.KV_proj_up_weight).view(
        batch_size, kv_len, config.n_heads, config.qk_nope_head_dim + config.v_head_dim
    )  # (bs, kv_len, n_heads, nope_dim + v_dim)
    k_nope_history, v_history = kv_nope_and_v_parts.split(
        [config.qk_nope_head_dim, config.v_head_dim], dim=-1
    )

    # Step 3: Apply RoPE
    # Queries Q: RoPE applied only to the current token
    q_rope_tensor_permuted = q_rope_tensor.permute(
        0, 2, 1, 3
    ).contiguous()  # (bs, n_heads, 1, rope_dim)
    q_rope_final = apply_rope_triton(q_rope_tensor_permuted, cos_cached, sin_cached, query_pos)
    q_nope_permuted = q_nope.permute(0, 2, 1, 3)  # (bs, n_heads, 1, nope_dim)
    q_final = torch.cat(
        [q_nope_permuted, q_rope_final], dim=-1
    )  # (bs, n_heads, 1, nope_dim + rope_dim)

    # Keys K: RoPE applied to the entire history
    # k_rope_tensor_history: (bs, kv_len, rope_dim)
    # Unsqueeze for n_heads dimension (it's broadcasted or applied per head if RoPE was head-specific, here it's shared)
    # This is (bs, 1, kv_len, rope_dim) to match Triton kernel's expectation for "num_heads" dim.
    # The RoPE matrices (cos_cached, sin_cached) are not head-specific.
    k_rope_tensor_for_rope = k_rope_tensor_history.unsqueeze(
        1
    ).contiguous()  # (bs, 1, kv_len, rope_dim)
    k_rope_rotated = apply_rope_triton(
        k_rope_tensor_for_rope, cos_cached, sin_cached, 0
    )  # start_pos=0 for keys
    k_rope_final = k_rope_rotated.expand(
        -1, config.n_heads, -1, -1
    )  # (bs, n_heads, kv_len, rope_dim)

    k_nope_permuted = k_nope_history.permute(0, 2, 1, 3)  # (bs, n_heads, kv_len, nope_dim)
    k_final = torch.cat(
        [k_nope_permuted, k_rope_final], dim=-1
    )  # (bs, n_heads, kv_len, nope_dim + rope_dim)

    # Values V (already up-projected, just permute)
    v_final = v_history.permute(0, 2, 1, 3)  # (bs, n_heads, kv_len, v_dim)

    # Step 4: Multi-head Attention
    head_dim_qk = config.qk_nope_head_dim + config.qk_rope_head_dim
    scale = 1.0 / math.sqrt(head_dim_qk)

    # q_final: (bs, n_heads, 1, qk_dim)
    # k_final: (bs, n_heads, kv_len, qk_dim) -> k_final.transpose: (bs, n_heads, qk_dim, kv_len)
    scores = torch.matmul(q_final, k_final.transpose(-1, -2)) * scale  # (bs, n_heads, 1, kv_len)

    # Softmax is applied over the kv_len dimension
    attn_weights = F.softmax(scores, dim=-1, dtype=torch.float32).to(
        x.dtype
    )  # (bs, n_heads, 1, kv_len)

    # attn_weights: (bs, n_heads, 1, kv_len)
    # v_final:      (bs, n_heads, kv_len, v_dim)
    attn_output = torch.matmul(attn_weights, v_final)  # (bs, n_heads, 1, v_dim)

    # Reshape and final output projection
    attn_output_reshaped = attn_output.transpose(1, 2).reshape(
        batch_size, seq_len_q, config.n_heads * config.v_head_dim
    )
    final_output = F.linear(attn_output_reshaped, config.wo_weight)  # (bs, 1, dim)

    # Return final output and the FULL KV cache data tensor (as expected by reference)
    # The kv_cache_module's internal 'data' buffer has been updated by its forward call.
    return final_output, kv_cache_module.get_data()


# --- Entry point for evaluation ---
def custom_kernel(data: input_t) -> output_t:
    # data tuple contains: (Config object, x_tensor, KVCache_module_instance)
    config_obj, x_tensor, kv_cache_module_instance = data

    # mla_forward will use and update kv_cache_module_instance in-place
    final_output, updated_kv_cache_full_data = mla_forward(
        x_tensor, kv_cache_module_instance, config_obj
    )

    return final_output, updated_kv_cache_full_data
python eval.py test /tmp/tmpf3twphvq
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])
