import math
from dataclasses import dataclass
import torch
from torch import nn
import torch.nn.functional as F
from typing import Tuple

# Import types and utilities from local files
# Assuming task.py and utils.py are available and contain necessary definitions.
# If they are not, they would need to be provided or their content inlined.
from task import input_t, output_t
from utils import make_match_reference

# Dummy RoPE class (its logic will be absorbed into the custom kernel for true optimization)
class RoPE(nn.Module):
    def __init__(self, d_model: int):
        super().__init__()
        self.d_model = d_model
        # Pre-compute theta for RoPE
        theta = 10000 ** (-torch.arange(0, d_model//2,dtype=torch.bfloat16) / (d_model//2))
        self.register_buffer("theta", theta)

    def rotate_half(self, x: torch.Tensor) -> torch.Tensor:
        x1, x2 = x.chunk(2, dim=-1)
        return torch.cat((-x2, x1), dim=-1)

    def forward(self, x: torch.Tensor, start_pos: int = 0) -> torch.Tensor:
        seq_len = x.size(-2)
        d_model = x.size(-1)
        assert d_model == self.d_model, f"Dimension mismatch: expected {self.d_model}, got {d_model}"
        
        # Calculate rotation frequencies for the current sequence segment
        seq_idx = torch.arange(start_pos, start_pos + seq_len, device=x.device)
        idx_theta = torch.einsum('s,d->sd', seq_idx, self.theta)
        idx_theta2 = torch.cat([idx_theta, idx_theta], dim=-1) # Duplicate for both halves
        
        # Compute cosine and sine components
        cos = idx_theta2.cos().to(torch.bfloat16)
        sin = idx_theta2.sin().to(torch.bfloat16)
        
        # Apply RoPE: x * cos + rotate_half(x) * sin
        return x * cos + self.rotate_half(x) * sin

# KVCache class from the user's selection, with device specified.
# Note: In a real optimized kernel, the cache update would be fused directly.
class KVCache(nn.Module):
    def __init__(self, kv_cache_shape: tuple) -> None:
        super().__init__()
        # Initialize data on 'cuda' as specified in the user's selection
        self.register_buffer('data', torch.zeros(kv_cache_shape, dtype=torch.bfloat16, device='cuda'))
        self.seq_len = 0
        self.zero()

    def zero(self) -> None:
        self.data.zero_()
        self.seq_len = 0
    
    def get_data(self) -> torch.Tensor:
        return self.data

    def forward(self, c_kv: torch.Tensor) -> Tuple[torch.Tensor, int]:
        assert self.seq_len + c_kv.size(1) <= self.data.size(1), "KV Cache Exceeded"

        self.data = self.data.to(c_kv.dtype)
        self.data[
            :, self.seq_len : self.seq_len + c_kv.size(1), :
        ] = c_kv
        self.seq_len += c_kv.size(1)

        return self.data[:, :self.seq_len], self.seq_len
    
# --- Custom MLA Decode Kernel (Conceptual Placeholder) ---
# This function represents the highly optimized, fused kernel
# that would be implemented in a low-level language like Triton or HIP/CUDA.
# It would leverage MI300's Matrix Cores and high memory bandwidth.
@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
    max_seq_len: int
    kv_cache_shape: tuple
    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

def _custom_mla_decode_kernel(
    x: torch.Tensor,
    kv_cache_data: torch.Tensor,
    kv_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,
    n_heads: int,
    q_lora_rank: int,
    kv_lora_rank: int,
    nope_head_dim: int,
    rope_head_dim: int,
    v_head_dim: int,
    eps: float
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Conceptual placeholder for a highly optimized MLA decode kernel.

    This single kernel would perform the following steps in a fused manner:
    1.  **Down-projection:** Compute `q_lora = x @ Q_proj_down_weight.T` and
        `kv_lora = x @ KV_proj_down_weight.T`.
    2.  **KV Cache Update (in-place):** Write the new `kv_lora` directly into
        `kv_cache_data` at the `kv_len` position. This avoids a separate copy.
    3.  **Up-projection & Split:**
        * Compute `q_proj_up = q_lora @ Q_proj_up_weight.T`.
        * Split `q_proj_up` into `q_nope` and `q_rope`.
        * Split the *entire* `kv_cache_data` (up to `kv_len`) into its `kv_nope` and `k_rope` components.
        * Compute `kv_proj_up = kv_nope @ KV_proj_up_weight.T`.
        * Split `kv_proj_up` into `k_nope` and `v`.
    4.  **RoPE Application (Fused):**
        * Apply RoPE to `q_rope` for the current token (`query_pos = kv_len - 1`).
        * **Crucially, apply RoPE to the *entire* `k_rope` history.** This is the most computationally intensive part for keys in the reference, as it recomputes RoPE for all `kv_len` tokens at each decode step. An optimized kernel would handle this by:
            * Pre-computing `cos` and `sin` rotation matrices.
            * Applying the rotation efficiently across all `kv_len` elements using vectorized operations and tiling.
    5.  **Concatenation & Permutation:** Combine `q_nope` with RoPE-applied `q_rope` to form `Q`. Combine `k_nope` with RoPE-applied `k_rope` to form `K`. Permute `V` for attention. These should be done without explicit memory copies.
    6.  **Fused Attention (QKV Matmul, Softmax, Attn-V Matmul):**
        * **QK^T Matmul:** Perform `scores = Q @ K.T`. This would be tiled to keep data in fast on-chip memory.
        * **Softmax:** Compute `softmax(scores)` row-wise. This should be numerically stable (e.g., by subtracting max for exponentiation).
        * **Attention-V Matmul:** Perform `attn @ V`.
        * These steps are typically fused to avoid writing intermediate attention scores and attention weights to global memory.
    7.  **Output Projection:** Perform `y = attn_output @ wo_weight.T`.

    **Key Optimization Principles within this kernel:**
    * **Memory Coalescing:** Ensure memory accesses are contiguous to maximize memory bandwidth utilization.
    * **Register/SRAM Usage:** Maximize the use of fast on-chip memory (registers, shared memory/L1 cache) by tiling the computations.
    * **Matrix Core Utilization:** Structure matrix multiplications to leverage MI300's Matrix Cores for BF16 acceleration.
    * **Parallelism:** Exploit parallelism across batches, heads, and within each matrix operation.
    * **Numerical Stability:** Maintain `bfloat16` precision and numerical stability, especially for softmax.

    Args:
        x (torch.Tensor): Input tensor for the current token (batch_size, 1, dim).
        kv_cache_data (torch.Tensor): The pre-allocated KV cache data (batch_size, max_seq_len, kv_dim).
        kv_len (int): The current valid sequence length in the KV cache.
        ... (other model weights and config parameters)

    Returns:
        Tuple[torch.Tensor, torch.Tensor]:
            - attn_output (torch.Tensor): The output of the MLA layer (batch_size, 1, dim).
            - updated_kv_cache_data (torch.Tensor): The KV cache data after update.
    """
    batch_size, seq_len, model_dim = x.size() # seq_len will be 1 for decode

    # --- Step 1: Simulate Down-projection and KV cache update ---
    # In the actual kernel, these would be fused.
    q_lora = F.linear(x, Q_proj_down_weight) # (bs, 1, q_lora_rank)
    kv_lora_new = F.linear(x, KV_proj_down_weight) # (bs, 1, kv_lora_rank + rope_head_dim)

    # In-place update of the KV cache with the new token's data
    # This write would be part of the fused kernel for efficiency.
    kv_cache_data[:, kv_len : kv_len + seq_len, :] = kv_lora_new

    # The full kv_lora sequence now includes the new token
    kv_lora_full = kv_cache_data[:, :kv_len + seq_len, :]
    current_kv_len = kv_len + seq_len
    query_pos = current_kv_len - 1 # Position of the current query

    # --- Step 2 & 3: Up-projection, Split, and RoPE application ---
    # These operations would be fused and optimized within the custom kernel.

    # Q projections and RoPE
    q_nope_and_rope_flat = F.linear(q_lora, Q_proj_up_weight) # (bs, 1, (nope_dim + rope_dim) * n_heads)
    q_nope_and_rope = q_nope_and_rope_flat.view(batch_size, seq_len, n_heads, nope_head_dim + rope_head_dim)
    q_nope, q_rope = torch.split(q_nope_and_rope, [nope_head_dim, rope_head_dim], dim=-1)
    
    # Apply RoPE to Q (current token only)
    q_rope = q_rope.permute(0, 2, 1, 3) # (bs, n_heads, seq_len, rope_head_dim)
    # The RoPE calculation for Q would be fused.
    # For simplicity here, we simulate it with the original RoPE class.
    # In the kernel, this would be computed directly.
    # Dummy RoPE application for Q (conceptual)
    q_rope_rotated = RoPE(rope_head_dim).to(x.device)(q_rope, start_pos=query_pos)
    
    q_nope = q_nope.permute(0, 2, 1, 3) # (bs, n_heads, seq_len, nope_head_dim)
    Q_final = torch.concat([q_nope, q_rope_rotated], dim=-1) # (bs, n_heads, 1, nope_dim + rope_dim)

    # K/V projections and RoPE
    # Split kv_lora_full into kv_lora_rank part and rope_head_dim part
    kv_nope_full, k_rope_full = torch.split(kv_lora_full, [kv_lora_rank, rope_head_dim], dim=-1)
    
    # Apply KV_proj_up to kv_nope_full
    kv_nope_up_flat = F.linear(kv_nope_full, KV_proj_up_weight) # (bs, current_kv_len, (nope_dim + v_dim) * n_heads)
    kv_nope_up = kv_nope_up_flat.view(batch_size, current_kv_len, n_heads, nope_head_dim + v_head_dim)
    k_nope_full, V_full = torch.split(kv_nope_up, [nope_head_dim, v_head_dim], dim=-1)
    
    # Apply RoPE to K (entire history)
    k_rope_full_expanded = k_rope_full[:, None, :, :].expand(-1, n_heads, -1, -1) # (bs, n_heads, current_kv_len, rope_head_dim)
    # The RoPE calculation for K would be fused and highly optimized.
    # For simplicity here, we simulate it with the original RoPE class.
    # In the kernel, this would be computed directly.
    k_rope_rotated_full = RoPE(rope_head_dim).to(x.device)(k_rope_full_expanded)

    k_nope_full = k_nope_full.permute(0, 2, 1, 3) # (bs, n_heads, current_kv_len, nope_head_dim)
    K_final = torch.concat([k_nope_full, k_rope_rotated_full], dim=-1) # (bs, n_heads, current_kv_len, nope_dim + rope_dim)

    V_final = V_full.permute(0, 2, 1, 3) # (bs, n_heads, current_kv_len, v_head_dim)

    # --- Step 4: Fused Multi-head Attention Computation ---
    # This entire block would be a single, highly optimized kernel.
    head_dim_total = nope_head_dim + rope_head_dim
    
    # QK^T Matmul
    # Fix: Cast head_dim_total to bfloat16 before taking the square root
    scores = torch.matmul(Q_final, K_final.transpose(-1, -2)) / torch.sqrt(torch.tensor(head_dim_total, dtype=torch.bfloat16))
    
    # Softmax
    attn_weights = F.softmax(scores, dim=-1).to(torch.bfloat16) # (bs, n_heads, 1, current_kv_len)
    
    # Attention @ V Matmul
    attn_output_per_head = torch.matmul(attn_weights, V_final) # (bs, n_heads, 1, v_head_dim)
    
    # Concatenate heads and apply final output projection
    attn_output_flat = attn_output_per_head.view(batch_size, seq_len, -1) # (bs, 1, n_heads * v_head_dim)
    final_output = F.linear(attn_output_flat, wo_weight) # (bs, 1, dim)

    return final_output, kv_cache_data # Return updated cache data directly

# The optimized MLA class, which calls the conceptual custom kernel
class MLA(nn.Module):
    def __init__(self, config: Config):
        super().__init__()
        self.dim = config.dim
        self.n_heads = config.n_heads
        self.q_lora_rank = config.q_lora_rank
        self.kv_lora_rank = config.kv_lora_rank
        self.nope_head_dim = config.qk_nope_head_dim
        self.rope_head_dim = config.qk_rope_head_dim
        self.v_head_dim = config.v_head_dim
        self.eps = 1e-6 # Epsilon for numerical stability

        # Weights for the linear layers (will be passed to the custom kernel)
        # These are not nn.Linear modules anymore, as their operations are fused.
        self.Q_proj_down_weight = None # Set in forward or when model is loaded
        self.Q_proj_up_weight = None
        self.KV_proj_down_weight = None
        self.KV_proj_up_weight = None
        self.wo_weight = None

    def forward(self, x: torch.Tensor, kv_cache: KVCache) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Forward pass for the MLA module, leveraging a conceptual custom decode kernel.

        Args:
            x (torch.Tensor): Input tensor for the current token (batch_size, 1, dim).
            kv_cache (KVCache): The KV cache object containing the pre-allocated data.

        Returns:
            Tuple[torch.Tensor, torch.Tensor]:
                - attn_output (torch.Tensor): The output of the MLA layer (batch_size, 1, dim).
                - updated_kv_cache_data (torch.Tensor): The KV cache data after update.
        """
        # Retrieve the underlying kv_cache data and its current valid length
        kv_cache_data = kv_cache.get_data()
        kv_len = kv_cache.seq_len # Current valid length of the cache

        # Call the conceptual custom MLA decode kernel
        # This function call replaces all the sequential PyTorch operations
        # from the original MLA forward pass.
        output, updated_kv_data = _custom_mla_decode_kernel(
            x=x,
            kv_cache_data=kv_cache_data,
            kv_len=kv_len,
            Q_proj_down_weight=self.Q_proj_down_weight,
            Q_proj_up_weight=self.Q_proj_up_weight,
            KV_proj_down_weight=self.KV_proj_down_weight,
            KV_proj_up_weight=self.KV_proj_up_weight,
            wo_weight=self.wo_weight,
            n_heads=self.n_heads,
            q_lora_rank=self.q_lora_rank,
            kv_lora_rank=self.kv_lora_rank,
            nope_head_dim=self.nope_head_dim,
            rope_head_dim=self.rope_head_dim,
            v_head_dim=self.v_head_dim,
            eps=self.eps
        )
        
        # Update the KVCache object's internal sequence length based on the kernel's operation.
        # This is crucial for the next decode step.
        # Since the kernel processes one token (seq_len=1), we increment by 1.
        kv_cache.seq_len = kv_len + x.size(1)

        return output, updated_kv_data



def custom_kernel(data: input_t) -> output_t:
    """
    The entry point for the optimized MLA decode kernel.
    It initializes the optimized MLA model and runs its forward pass.
    """
    config, x, kv_cache = data

    # Initialize the optimized MLA model
    model = MLA(config).to('cuda')
    
    # Assign weights to the optimized model (these are not nn.Parameters in MLA)
    model.Q_proj_down_weight = config.Q_proj_down_weight
    model.Q_proj_up_weight = config.Q_proj_up_weight
    model.KV_proj_down_weight = config.KV_proj_down_weight
    model.KV_proj_up_weight = config.KV_proj_up_weight
    model.wo_weight = config.wo_weight

    # Run the forward pass which internally calls the conceptual custom kernel
    output, updated_kv_data = model(x, kv_cache)
    return output, updated_kv_data

python eval.py test /tmp/tmp0qxoivdk
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])
