import os
import math
import torch
import torch.nn.functional as F
from task import input_t, output_t

# Keep original settings
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:512,expandable_segments:True"

if torch.cuda.is_available():
    torch.backends.cuda.enable_mem_efficient_sdp(True)
    torch.backends.cuda.enable_flash_sdp(True)

# Keep your original RoPE implementation exactly
_rope_cache = {}


def get_rope_cache(d_model: int, max_seq_len: int, device="cuda"):
    cache_key = (d_model, max_seq_len, device)
    if cache_key not in _rope_cache:
        theta = 10000 ** (
            -torch.arange(0, d_model // 2, dtype=torch.bfloat16, device=device) / (d_model // 2)
        )
        positions = torch.arange(max_seq_len, dtype=torch.bfloat16, device=device)
        angles = torch.einsum("i,j->ij", positions, theta)
        angles = torch.cat([angles, angles], dim=-1)
        cos_cached = angles.cos()
        sin_cached = angles.sin()
        _rope_cache[cache_key] = (cos_cached, sin_cached)
    return _rope_cache[cache_key]


def apply_rope(x: torch.Tensor, cos_cached: torch.Tensor, sin_cached: torch.Tensor, start_pos: int):
    seq_len = x.size(-2)
    cos = cos_cached[start_pos : start_pos + seq_len]
    sin = sin_cached[start_pos : start_pos + seq_len]
    x1, x2 = x.chunk(2, dim=-1)
    rotated = torch.cat((-x2, x1), dim=-1)
    return x * cos + rotated * sin


def mla_forward(x: torch.Tensor, kv_cache, config) -> torch.Tensor:
    batch_size, seq_len, model_dim = x.size()

    # Get cached RoPE matrices
    cos_cached, sin_cached = get_rope_cache(config.qk_rope_head_dim, config.max_seq_len)

    # Step 1: Handle down-projection + KV cache (unchanged)
    q_lora = F.linear(x, config.Q_proj_down_weight)
    kv_lora = F.linear(x, config.KV_proj_down_weight)
    kv_lora, kv_len = kv_cache(kv_lora)
    query_pos = kv_len - 1

    # Step 2: Up-project and prepare NoPE + RoPE (unchanged)
    q_nope_and_rope = F.linear(q_lora, config.Q_proj_up_weight).view(
        batch_size, seq_len, config.n_heads, config.qk_nope_head_dim + config.qk_rope_head_dim
    )
    q_nope, q_rope_tensor = q_nope_and_rope.split(
        [config.qk_nope_head_dim, config.qk_rope_head_dim], dim=-1
    )

    # Handle keys and values K/V. V does not need RoPE (unchanged)
    kv_nope, k_rope_tensor = kv_lora.split([config.kv_lora_rank, config.qk_rope_head_dim], dim=-1)

    # Step 3: Handle RoPE Stream (unchanged)
    q_rope_tensor = q_rope_tensor.permute(0, 2, 1, 3)
    q_rope_tensor = apply_rope(q_rope_tensor, cos_cached, sin_cached, query_pos)
    q_nope = q_nope.permute(0, 2, 1, 3)
    q = torch.cat([q_nope, q_rope_tensor], dim=-1)

    k_rope_tensor = k_rope_tensor.unsqueeze(1)
    k_rope_tensor = apply_rope(k_rope_tensor, cos_cached, sin_cached, 0)
    k_rope_tensor = k_rope_tensor.expand(-1, config.n_heads, -1, -1)

    # Step 4: Memory-efficient KV processing
    # Process in microbatches but with much smaller size and simpler logic
    micro_batch_size = 8  # Much smaller to reduce memory
    num_micro_batches = math.ceil(batch_size / micro_batch_size)

    all_outputs = []

    for mb_idx in range(num_micro_batches):
        start_idx = mb_idx * micro_batch_size
        end_idx = min(batch_size, (mb_idx + 1) * micro_batch_size)

        # Slice everything for this microbatch
        q_mb = q[start_idx:end_idx]
        kv_nope_mb = kv_nope[start_idx:end_idx]
        k_rope_mb = k_rope_tensor[start_idx:end_idx]

        # Do the KV up-projection for this small microbatch
        kv_up_mb = F.linear(kv_nope_mb, config.KV_proj_up_weight)
        kv_up_reshaped = kv_up_mb.view(
            end_idx - start_idx, kv_len, config.n_heads, config.qk_nope_head_dim + config.v_head_dim
        )
        k_nope_mb, v_mb = kv_up_reshaped.split([config.qk_nope_head_dim, config.v_head_dim], dim=-1)

        k_nope_mb = k_nope_mb.permute(0, 2, 1, 3)
        k_mb = torch.cat([k_nope_mb, k_rope_mb], dim=-1)
        v_mb = v_mb.permute(0, 2, 1, 3)

        # Compute attention for this microbatch using your original method
        scale = math.sqrt(config.qk_rope_head_dim + config.qk_nope_head_dim)
        scores = torch.matmul(q_mb, k_mb.transpose(-1, -2)) / scale
        attn = F.softmax(scores, dim=-1).to(torch.bfloat16)
        y_mb = torch.matmul(attn, v_mb).view(end_idx - start_idx, 1, -1)

        all_outputs.append(y_mb)

    # Concatenate all microbatch outputs
    y = torch.cat(all_outputs, dim=0)
    y = F.linear(y, config.wo_weight)

    return y, kv_cache.get_data()


def custom_kernel(data: input_t) -> output_t:
    config, x, kv_cache = data
    output, kv_cache_data = mla_forward(x, kv_cache, config)
    return output, kv_cache_data

python eval.py leaderboard /tmp/tmpwe731koh
/home/runner/_work/discord-cluster-manager/discord-cluster-manager/reference.py:162: UserWarning: expandable_segments not supported on this platform (Triggered internally at /pytorch/c10/hip/HIPAllocatorConfig.h:30.)
  Q_proj_down_weight = torch.randn((dq, dim), dtype=torch.bfloat16, generator=gen, device='cuda') / math.sqrt(dim)

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])
