import math

import torch
import torch.nn.functional as F

from task import input_t, output_t


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

def rope(x: torch.Tensor, theta: torch.Tensor ,start_pos: int = 0) -> torch.Tensor:
    seq_len = x.size(-2)
    d_model = x.size(-1)
    seq_idx = torch.arange(start_pos, start_pos + seq_len, device=x.device)
    idx_theta = torch.einsum('s,d->sd', seq_idx, theta)
    idx_theta2 = torch.cat([idx_theta, idx_theta], dim=-1)
    cos = idx_theta2.cos().to(torch.bfloat16)
    sin = idx_theta2.sin().to(torch.bfloat16)
    return x * cos + rotate_half(x) * sin


def custom_kernel(data: input_t) -> output_t:
    config, x, kv_cache = data

    theta = 10000 ** (-torch.arange(0, config.qk_rope_head_dim//2,dtype=torch.bfloat16) / (config.qk_rope_head_dim//2)).cuda()

    # seq_len = 1 always here
    batch_size, seq_len, model_dim = x.size()

    ################################################################################
    #                 Step 1: Handle down-projection + KV cache                    #
    ################################################################################

    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                  #
    ################################################################################

    # Handle queries Q first
    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 = torch.split(q_nope_and_rope, [config.qk_nope_head_dim, config.qk_rope_head_dim], dim=-1)

    # Handle keys and values K/V. V does not need RoPE
    kv_nope, k_rope = torch.split(kv_lora, [config.kv_lora_rank, config.qk_rope_head_dim], dim=-1)
    kv_nope = F.linear(kv_nope, config.KV_proj_up_weight).view(
        batch_size, kv_len, config.n_heads, config.qk_nope_head_dim + config.v_head_dim)
    k_nope, v = torch.split(kv_nope, [config.qk_nope_head_dim, config.v_head_dim], dim=-1)

    ################################################################################
    #                    Step 3: Handle RoPE Stream                                #
    ################################################################################

    # Compute RoPE for queries and combine with no-RoPE part
    q_rope = q_rope.permute(0, 2, 1, 3) # bs x n_heads x seq_len x rope_head_dim
    q_rope = rope(q_rope, theta, start_pos=query_pos)
    q_nope = q_nope.permute(0, 2, 1, 3) # bs x n_heads x seq_len x rope_head_dim
    q = torch.concat([q_nope, q_rope], dim=-1)

    # Compute RoPE for keys and combine with no-RoPE part
    k_rope = k_rope[:, None, :, :]
    k_rope = rope(k_rope, theta).expand(-1,config.n_heads,-1,-1)
    k_nope = k_nope.permute(0, 2, 1, 3) # bs x kv_len x n_heads x rope_head_dim
    k = torch.concat([k_nope, k_rope], dim=-1)

    ################################################################################
    #                        Compute Multi-head Attention                          #
    ################################################################################

    v = v.permute(0, 2, 1, 3) # bs x n_heads x kv_len x v_head_dim
    scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(config.qk_rope_head_dim + config.qk_nope_head_dim)
    attn = F.softmax(scores, dim=-1).to(torch.bfloat16)
    y = torch.matmul(attn, v).view(batch_size, 1, -1)
    y = F.linear(y, config.wo_weight)

    return y, kv_cache.get_data()

python eval.py benchmark /tmp/tmp_v8gb540
Buff name: data, shape: torch.Size([128, 8192, 576])
Buff name: data, shape: torch.Size([128, 8192, 576])
