import os

# 1. Control CUDA caching allocator fragmentation
# expandable_segments:True is generally good.
# max_split_size_mb:512 can help reduce fragmentation for very large allocations.
# Set this early, before PyTorch CUDA context is initialized if possible.
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:512,expandable_segments:True"

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

# 2. Enable memory‐efficient and FlashAttention kernels
if torch.cuda.is_available():
    # These settings are generally for CUDA. For HIP (MI300), behavior might differ or options might not be recognized.
    # PyTorch typically enables optimized SDPA paths on ROCm by default where available.
    # Assuming these are beneficial or at least harmless based on common practice.
    torch.backends.cuda.enable_mem_efficient_sdp(True)
    torch.backends.cuda.enable_flash_sdp(True)

# Global cache for RoPE matrices
_rope_cache = {}


def get_rope_cache_params(d_model: int, max_seq_len: int, device: torch.device, dtype: torch.dtype):
    # Using a more robust cache key
    idx = device.index if device.type != "cpu" and device.index is not None else -1
    cache_key = (d_model, max_seq_len, device.type, idx, dtype)

    if cache_key not in _rope_cache:
        # Using bfloat16 for RoPE intermediate calculations as in reference for consistency
        theta_values = 10000.0 ** (
            -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_values)
        angles_full_dim = torch.cat([angles, angles], dim=-1)

        cos_cached = angles_full_dim.cos().to(dtype)  # Ensure final cache is in target dtype
        sin_cached = angles_full_dim.sin().to(dtype)
        _rope_cache[cache_key] = (cos_cached, sin_cached)
    return _rope_cache[cache_key]


def apply_rotary_embeddings(
    x: torch.Tensor, cos_full_params: torch.Tensor, sin_full_params: torch.Tensor, start_pos: int
):
    # x: (..., seq_len, dim)
    # cos_full_params, sin_full_params: (max_seq_len, dim)
    seq_len = x.size(-2)

    # Slicing the cached RoPE matrices for the current positions
    # cos_selected/sin_selected will be (seq_len, dim)
    cos_selected = cos_full_params[start_pos : start_pos + seq_len, :]
    sin_selected = sin_full_params[start_pos : start_pos + seq_len, :]

    # Unsqueeze to allow broadcasting with x's batch/head dimensions
    # If x is (bs, n_heads, seq_len, dim), cos/sin should be (1, 1, seq_len, dim)
    # If x is (bs, seq_len, dim), cos/sin should be (1, seq_len, dim)
    if x.dim() > cos_selected.dim():  # x has more leading dimensions
        # Add leading dimensions to cos_selected and sin_selected to match x
        # E.g., if x is (B,H,S,D) and cos_selected is (S,D), make cos_selected (1,1,S,D)
        num_leading_dims_to_add = x.dim() - cos_selected.dim()
        for _ in range(num_leading_dims_to_add):
            cos_selected = cos_selected.unsqueeze(0)
            sin_selected = sin_selected.unsqueeze(0)

    x1, x2 = x.chunk(2, dim=-1)
    x_rotated = torch.cat([-x2, x1], dim=-1)
    return x * cos_selected + x_rotated * sin_selected


def mla_decode_batched_loop_forward(x_input: torch.Tensor, kv_cache_object, config_params):
    batch_size, current_seq_len, _ = x_input.size()  # current_seq_len is 1 for decode

    device = x_input.device
    dtype = x_input.dtype

    cos_rope_full, sin_rope_full = get_rope_cache_params(
        config_params.qk_rope_head_dim, config_params.max_seq_len, device=device, dtype=dtype
    )

    # Step 1: Down-projection for current token & update KV cache (full batch)
    q_lora_out = F.linear(x_input, config_params.Q_proj_down_weight)
    kv_latent_current_token = F.linear(x_input, config_params.KV_proj_down_weight)

    kv_history_combined_view, current_total_kv_len = kv_cache_object(kv_latent_current_token)
    # kv_history_combined_view: (bs, total_kv_len, kv_lora_rank + qk_rope_head_dim)

    query_start_pos = current_total_kv_len - current_seq_len

    # Step 2: Up-project Q (current token only, small, full batch)
    q_up_proj = F.linear(q_lora_out, config_params.Q_proj_up_weight).view(
        batch_size,
        current_seq_len,
        config_params.n_heads,
        config_params.qk_nope_head_dim + config_params.qk_rope_head_dim,
    )
    q_nope_component, q_rope_latent_component = q_up_proj.split(
        [config_params.qk_nope_head_dim, config_params.qk_rope_head_dim], dim=-1
    )

    # q_rope_latent_component: (bs, 1, n_heads, rope_dim)
    q_rope_latent_permuted = q_rope_latent_component.permute(
        0, 2, 1, 3
    )  # (bs, n_heads, 1, rope_dim)
    q_rope_final_component = apply_rotary_embeddings(
        q_rope_latent_permuted, cos_rope_full, sin_rope_full, query_start_pos
    )

    # q_nope_component: (bs, 1, n_heads, nope_dim)
    q_final_full_batch = torch.cat(
        [q_nope_component.permute(0, 2, 1, 3), q_rope_final_component], dim=-1
    )
    # q_final_full_batch: (bs, n_heads, 1, full_qk_dim)

    # Latent K/V parts from the (updated) KV cache history (full batch views, no large copy yet)
    kv_nope_latent_from_history, k_rope_latent_from_history = kv_history_combined_view.split(
        [config_params.kv_lora_rank, config_params.qk_rope_head_dim], dim=-1
    )
    # kv_nope_latent_from_history: (bs, total_kv_len, kv_lora_rank)
    # k_rope_latent_from_history: (bs, total_kv_len, qk_rope_head_dim)

    attention_outputs_list = []

    # Step 3: Loop over batch items for memory-intensive K/V processing and attention
    for item_idx in range(batch_size):
        q_current_item = q_final_full_batch[item_idx : item_idx + 1]  # (1, n_heads, 1, full_qk_dim)

        kv_nope_latent_current_item = kv_nope_latent_from_history[
            item_idx : item_idx + 1
        ]  # (1, total_kv_len, kv_lora_rank)
        k_rope_latent_current_item = k_rope_latent_from_history[
            item_idx : item_idx + 1
        ]  # (1, total_kv_len, qk_rope_head_dim)

        # Up-project K_nope and V parts for this single batch item (MEMORY BOTTLENECK MITIGATED)
        kv_up_proj_item = F.linear(kv_nope_latent_current_item, config_params.KV_proj_up_weight)
        kv_up_proj_item_reshaped = kv_up_proj_item.view(
            1,
            current_total_kv_len,
            config_params.n_heads,
            config_params.qk_nope_head_dim + config_params.v_head_dim,
        )
        k_nope_item, v_all_history_item = kv_up_proj_item_reshaped.split(
            [config_params.qk_nope_head_dim, config_params.v_head_dim], dim=-1
        )
        # k_nope_item: (1, total_kv_len, n_heads, nope_dim)
        # v_all_history_item: (1, total_kv_len, n_heads, v_dim)

        # Apply RoPE to K_rope part for this single batch item
        # k_rope_latent_current_item: (1, total_kv_len, qk_rope_dim) -> apply_rotary_embeddings handles this 3D tensor
        k_rope_item_applied = apply_rotary_embeddings(
            k_rope_latent_current_item, cos_rope_full, sin_rope_full, 0
        )  # start_pos=0 for cached keys
        # k_rope_item_applied: (1, total_kv_len, qk_rope_dim)
        # Expand to match n_heads for concatenation
        k_rope_final_item_component = k_rope_item_applied.unsqueeze(1).expand(
            -1, config_params.n_heads, -1, -1
        )
        # k_rope_final_item_component: (1, n_heads, total_kv_len, qk_rope_dim)

        # Combine K parts for this item
        k_final_current_item = torch.cat(
            [k_nope_item.permute(0, 2, 1, 3), k_rope_final_item_component], dim=-1
        )
        # k_final_current_item: (1, n_heads, total_kv_len, full_qk_dim)

        # Prepare V for this item
        v_final_current_item = v_all_history_item.permute(
            0, 2, 1, 3
        )  # (1, n_heads, total_kv_len, v_dim)

        # Scaled dot-product attention for this single batch item
        attn_output_current_item = F.scaled_dot_product_attention(
            q_current_item,
            k_final_current_item,
            v_final_current_item,
            attn_mask=None,
            dropout_p=0.0,
            is_causal=False,
        )  # attn_output_current_item: (1, n_heads, 1, v_dim)
        attention_outputs_list.append(attn_output_current_item)

    # Step 4: Concatenate results from all batch items
    attn_output_aggregated = torch.cat(attention_outputs_list, dim=0)  # (bs, n_heads, 1, v_dim)

    # Reshape for output projection
    y_pre_output_proj = (
        attn_output_aggregated.permute(0, 2, 1, 3)
        .contiguous()
        .view(batch_size, current_seq_len, config_params.n_heads * config_params.v_head_dim)
    )

    # Final output projection
    final_model_output = F.linear(y_pre_output_proj, config_params.wo_weight)  # (bs, 1, dim)

    return final_model_output, kv_cache_object.get_data()


def custom_kernel(data: input_t) -> output_t:
    config_obj, x_input_tensor, kv_cache_inst = data

    with torch.no_grad():
        final_output_tensor, updated_kv_data = mla_decode_batched_loop_forward(
            x_input_tensor, kv_cache_inst, config_obj
        )

    return final_output_tensor, updated_kv_data
python eval.py test /tmp/tmp8o_timi7
/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])
