import math

import torch
import torch.nn as nn
import torch.nn.functional as F

# Default ModelNew is PyTorch-only. run_and_check: use check_kernel=False with backend=cuda
# until the candidate file contains a real CUDA kernel (see kernel_static_checker).

class Model(nn.Module):
    """
    Single-layer multi-head cross-attention with RoPE on Q and K, attention
    dropout, and output projection—intended as a fusion target for robotics
    stacks (e.g. proprio / action queries attending to vision or memory).

    Pipeline (conceptually one fused kernel): linear Q from query stream,
    linear K/V from key–value stream → reshape heads → apply RoPE separately
    along query positions (T_q) and key positions (T_k) → scaled dot-product,
    softmax, dropout on weights → aggregate values → output projection.
    """

    def __init__(
        self,
        dim: int,
        num_heads: int,
        dropout: float = 0.0,
        rope_base: float = 10000.0,
    ):
        super().__init__()
        assert dim % num_heads == 0
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        assert self.head_dim % 2 == 0, "RoPE requires an even head dimension"

        self.q_proj = nn.Linear(dim, dim, bias=False)
        self.k_proj = nn.Linear(dim, dim, bias=False)
        self.v_proj = nn.Linear(dim, dim, bias=False)
        self.out_proj = nn.Linear(dim, dim, bias=False)
        self.dropout = nn.Dropout(dropout)
        self.scale = 1.0 / math.sqrt(self.head_dim)

        inv_freq = 1.0 / (
            rope_base ** (torch.arange(0, self.head_dim, 2).float() / self.head_dim)
        )
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    @staticmethod
    def _rotate_half(x: torch.Tensor) -> torch.Tensor:
        x1, x2 = x.chunk(2, dim=-1)
        return torch.cat([-x2, x1], dim=-1)

    def _rope(self, x: torch.Tensor, seq_len: int) -> torch.Tensor:
        """Apply RoPE along the sequence dimension (last-but-one after head). x: (B, H, L, d)."""
        t = torch.arange(seq_len, device=x.device, dtype=self.inv_freq.dtype)
        freqs = torch.outer(t, self.inv_freq)
        emb = torch.cat([freqs, freqs], dim=-1)
        cos = emb.cos().view(1, 1, seq_len, self.head_dim)
        sin = emb.sin().view(1, 1, seq_len, self.head_dim)
        return x * cos + self._rotate_half(x) * sin

    def forward(self, x_query: torch.Tensor, x_kv: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x_query: (B, T_q, dim) — e.g. robot or decoder tokens
            x_kv:    (B, T_k, dim) — e.g. encoder / memory tokens
        Returns:
            (B, T_q, dim)
        """
        B, T_q, _ = x_query.shape
        _, T_k, _ = x_kv.shape
        h, d = self.num_heads, self.head_dim

        q = self.q_proj(x_query).view(B, T_q, h, d).transpose(1, 2)
        k = self.k_proj(x_kv).view(B, T_k, h, d).transpose(1, 2)
        v = self.v_proj(x_kv).view(B, T_k, h, d).transpose(1, 2)

        q = self._rope(q, T_q)
        k = self._rope(k, T_k)

        attn = torch.matmul(q, k.transpose(-2, -1)) * self.scale
        attn = F.softmax(attn, dim=-1)
        attn = self.dropout(attn)
        out = torch.matmul(attn, v)
        out = out.transpose(1, 2).contiguous().view(B, T_q, self.dim)
        return self.out_proj(out)


class ModelNew(nn.Module):
    """
    KernelBench candidate entry point (`eval_kernel_against_ref` / `run_and_check`).
    Replace this class with a fused CUDA/Triton (etc.) implementation; the default
    delegates to `Model` so you can smoke-test the harness with the same file for
    both `ref_arch_src_path` and `kernel_src_path`.
    """

    def __init__(
        self,
        dim: int,
        num_heads: int,
        dropout: float = 0.0,
        rope_base: float = 10000.0,
    ):
        super().__init__()
        self._impl = Model(dim, num_heads, dropout, rope_base)

    def forward(self, x_query: torch.Tensor, x_kv: torch.Tensor) -> torch.Tensor:
        return self._impl(x_query, x_kv)


batch_size = 4
seq_len_q = 64
seq_len_kv = 128
dim = 256
num_heads = 8
# Use 0.0 so `run_and_check` / multi-trial correctness matches: the harness does not
# re-seed RNG between the reference forward and `ModelNew` forward, so attention
# dropout > 0 would make outputs differ even for identical math.
dropout_p = 0.0


def get_inputs():
    p = popcorn_pri
    bs = p.jitter_int(batch_size)
    tq = p.jitter_int(seq_len_q)
    tk = p.jitter_int(seq_len_kv)
    x_query = torch.randn(bs, tq, dim)
    x_kv = torch.randn(bs, tk, dim)
    return [x_query, x_kv]


def get_init_inputs():
    return [dim, num_heads, dropout_p]

import math
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

# Compile fused CUDA kernels
fused_kernel = load_inline(
    name='fused_attention_rope',
    cpp_sources='',
    cuda_sources=r'''
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <c10/cuda/CUDAException.h>

// Fused RoPE + Attention kernel
// Computes: apply RoPE to Q and K, then scaled dot-product attention with softmax
template <typename scalar_t>
__global__ void fused_rope_attention_kernel(
    const scalar_t* __restrict__ q_in,      // (B, T_q, H, head_dim) - after reshape
    const scalar_t* __restrict__ k_in,      // (B, T_k, H, head_dim) - after reshape  
    const scalar_t* __restrict__ v_in,      // (B, T_k, H, head_dim) - after reshape
    scalar_t* __restrict__ out,             // (B, T_q, H, head_dim)
    const float* __restrict__ cos_q,        // (T_q, head_dim)
    const float* __restrict__ sin_q,        // (T_q, head_dim)
    const float* __restrict__ cos_k,        // (T_k, head_dim)
    const float* __restrict__ sin_k,        // (T_k, head_dim)
    const float scale,
    const int B, const int H, const int T_q, const int T_k, const int head_dim
) {
    // Each block handles one (batch, head, query_pos) tuple
    const int batch = blockIdx.x;
    const int head = blockIdx.y;
    const int q_pos = blockIdx.z;
    
    const int half_dim = head_dim / 2;
    const int thread_idx = threadIdx.x;
    
    // Shared memory for Q values (after RoPE)
    extern __shared__ float shared_mem[];
    float* q_rope = shared_mem;                           // head_dim elements
    float* scores = &shared_mem[head_dim];               // T_k elements (partial per thread)
    
    // Load and apply RoPE to Q for this position
    // Q layout: (B, T_q, H, head_dim)
    for (int d = thread_idx; d < head_dim; d += blockDim.x) {
        const int q_idx = batch * T_q * H * head_dim + q_pos * H * head_dim + head * head_dim + d;
        const float q_val = static_cast<float>(q_in[q_idx]);
        
        const float c = cos_q[q_pos * head_dim + d];
        const float s = sin_q[q_pos * head_dim + d];
        
        // Apply RoPE rotation
        int d_pair = (d < half_dim) ? (d + half_dim) : (d - half_dim);
        const int q_pair_idx = batch * T_q * H * head_dim + q_pos * H * head_dim + head * head_dim + d_pair;
        const float q_pair = static_cast<float>(q_in[q_pair_idx]);
        
        if (d < half_dim) {
            q_rope[d] = q_val * c - q_pair * s;
        } else {
            q_rope[d] = q_val * c + q_pair * s;
        }
    }
    __syncthreads();
    
    // Compute attention scores for this Q position against all K positions
    // Using online softmax for numerical stability
    
    float max_score = -INFINITY;
    float sum_exp = 0.0f;
    
    // Process K positions in chunks to fit in shared memory
    const int CHUNK_SIZE = 32;  // Process 32 K positions at a time
    
    for (int k_chunk_start = 0; k_chunk_start < T_k; k_chunk_start += CHUNK_SIZE) {
        const int k_chunk_end = min(k_chunk_start + CHUNK_SIZE, T_k);
        
        // First pass: compute scores and find max
        for (int k_pos = k_chunk_start + thread_idx; k_pos < k_chunk_end; k_pos += blockDim.x) {
            float dot = 0.0f;
            
            // Compute Q @ K^T with RoPE applied to K
            for (int d = 0; d < head_dim; d++) {
                const int k_idx = batch * T_k * H * head_dim + k_pos * H * head_dim + head * head_dim + d;
                float k_val = static_cast<float>(k_in[k_idx]);
                
                // Apply RoPE to K on-the-fly
                const float c = cos_k[k_pos * head_dim + d];
                const float s = sin_k[k_pos * head_dim + d];
                
                int d_pair = (d < half_dim) ? (d + half_dim) : (d - half_dim);
                const int k_pair_idx = batch * T_k * H * head_dim + k_pos * H * head_dim + head * head_dim + d_pair;
                float k_pair = static_cast<float>(k_in[k_pair_idx]);
                
                float k_rope;
                if (d < half_dim) {
                    k_rope = k_val * c - k_pair * s;
                } else {
                    k_rope = k_val * c + k_pair * s;
                }
                
                dot += q_rope[d] * k_rope;
            }
            
            float score = dot * scale;
            scores[k_pos - k_chunk_start] = score;
            
            if (score > max_score) {
                max_score = score;
            }
        }
        __syncthreads();
        
        // Block-level reduction for max
        __shared__ float block_max;
        if (thread_idx == 0) {
            block_max = -INFINITY;
            for (int i = 0; i < k_chunk_end - k_chunk_start; i++) {
                if (scores[i] > block_max) block_max = scores[i];
            }
        }
        __syncthreads();
        
        if (block_max > max_score) max_score = block_max;
    }
    
    // Second pass: compute exp(score - max) and sum
    float local_sum = 0.0f;
    
    for (int k_chunk_start = 0; k_chunk_start < T_k; k_chunk_start += CHUNK_SIZE) {
        const int k_chunk_end = min(k_chunk_start + CHUNK_SIZE, T_k);
        
        for (int k_pos = k_chunk_start + thread_idx; k_pos < k_chunk_end; k_pos += blockDim.x) {
            // Recompute score (or store in shared memory if space allows)
            float dot = 0.0f;
            for (int d = 0; d < head_dim; d++) {
                const int k_idx = batch * T_k * H * head_dim + k_pos * H * head_dim + head * head_dim + d;
                float k_val = static_cast<float>(k_in[k_idx]);
                
                const float c = cos_k[k_pos * head_dim + d];
                const float s = sin_k[k_pos * head_dim + d];
                
                int d_pair = (d < half_dim) ? (d + half_dim) : (d - half_dim);
                const int k_pair_idx = batch * T_k * H * head_dim + k_pos * H * head_dim + head * head_dim + d_pair;
                float k_pair = static_cast<float>(k_in[k_pair_idx]);
                
                float k_rope;
                if (d < half_dim) {
                    k_rope = k_val * c - k_pair * s;
                } else {
                    k_rope = k_val * c + k_pair * s;
                }
                
                dot += q_rope[d] * k_rope;
            }
            float score = dot * scale;
            float exp_score = expf(score - max_score);
            scores[k_pos - k_chunk_start] = exp_score;
            local_sum += exp_score;
        }
        __syncthreads();
    }
    
    // Reduce sum across threads
    __shared__ float total_sum;
    if (thread_idx == 0) {
        total_sum = 0.0f;
    }
    __syncthreads();
    
    // Use atomic add for simplicity
    atomicAdd(&total_sum, local_sum);
    __syncthreads();
    
    // Third pass: compute weighted sum of V
    // Each thread computes partial output for different output dimensions
    for (int d_out = thread_idx; d_out < head_dim; d_out += blockDim.x) {
        float out_val = 0.0f;
        
        for (int k_chunk_start = 0; k_chunk_start < T_k; k_chunk_start += CHUNK_SIZE) {
            const int k_chunk_end = min(k_chunk_start + CHUNK_SIZE, T_k);
            
            for (int k_pos = k_chunk_start + thread_idx; k_pos < k_chunk_end; k_pos += blockDim.x) {
                // Recompute exp_score or retrieve from shared memory
                float dot = 0.0f;
                for (int d = 0; d < head_dim; d++) {
                    const int k_idx = batch * T_k * H * head_dim + k_pos * H * head_dim + head * head_dim + d;
                    float k_val = static_cast<float>(k_in[k_idx]);
                    
                    const float c = cos_k[k_pos * head_dim + d];
                    const float s = sin_k[k_pos * head_dim + d];
                    
                    int d_pair = (d < half_dim) ? (d + half_dim) : (d - half_dim);
                    const int k_pair_idx = batch * T_k * H * head_dim + k_pos * H * head_dim + head * head_dim + d_pair;
                    float k_pair = static_cast<float>(k_in[k_pair_idx]);
                    
                    float k_rope;
                    if (d < half_dim) {
                        k_rope = k_val * c - k_pair * s;
                    } else {
                        k_rope = k_val * c + k_pair * s;
                    }
                    
                    dot += q_rope[d] * k_rope;
                }
                float score = dot * scale;
                float exp_score = expf(score - max_score);
                float attn_weight = exp_score / total_sum;
                
                const int v_idx = batch * T_k * H * head_dim + k_pos * H * head_dim + head * head_dim + d_out;
                out_val += attn_weight * static_cast<float>(v_in[v_idx]);
            }
        }
        
        // Reduce across threads for this output dimension
        const int out_idx = batch * T_q * H * head_dim + q_pos * H * head_dim + head * head_dim + d_out;
        out[out_idx] = static_cast<scalar_t>(out_val);
    }
}

// Simpler, more efficient fused kernel - processes one attention head per block
template <typename scalar_t>
__global__ void fused_attention_kernel(
    const scalar_t* __restrict__ q,      // (B, H, T_q, head_dim) - already with RoPE
    const scalar_t* __restrict__ k,      // (B, H, T_k, head_dim) - already with RoPE
    const scalar_t* __restrict__ v,      // (B, H, T_k, head_dim)
    scalar_t* __restrict__ out,          // (B, H, T_q, head_dim)
    const float scale,
    const int B, const int H, const int T_q, const int T_k, const int head_dim
) {
    // Each block handles one (batch, head, query_pos)
    const int batch = blockIdx.x / H;
    const int head = blockIdx.x % H;
    const int q_pos = blockIdx.y;
    const int tid = threadIdx.x;
    const int block_size = blockDim.x;
    
    // Shared memory for Q values and attention weights
    extern __shared__ char smem[];
    float* q_local = reinterpret_cast<float*>(smem);
    float* attn_weights = q_local + head_dim;
    
    // Load Q values for this position (coalesced read)
    const int q_base = batch * H * T_q * head_dim + head * T_q * head_dim + q_pos * head_dim;
    for (int d = tid; d < head_dim; d += block_size) {
        q_local[d] = static_cast<float>(q[q_base + d]);
    }
    __syncthreads();
    
    // Compute attention scores: Q @ K^T
    // Each thread handles multiple K positions
    float max_score = -INFINITY;
    
    for (int k_pos = tid; k_pos < T_k; k_pos += block_size) {
        float dot = 0.0f;
        const int k_base = batch * H * T_k * head_dim + head * T_k * head_dim + k_pos * head_dim;
        
        for (int d = 0; d < head_dim; d++) {
            dot += q_local[d] * static_cast<float>(k[k_base + d]);
        }
        
        float score = dot * scale;
        attn_weights[k_pos] = score;
        if (score > max_score) max_score = score;
    }
    
    // Warp-level reduce for max
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        float other = __shfl_down_sync(0xffffffff, max_score, offset);
        if (other > max_score) max_score = other;
    }
    
    // Broadcast max from lane 0
    max_score = __shfl_sync(0xffffffff, max_score, 0);
    
    // Handle multiple warps
    __shared__ float global_max;
    if (tid == 0) global_max = -INFINITY;
    __syncthreads();
    
    if ((tid & 31) == 0) {  // One thread per warp
        atomicMax(reinterpret_cast<int*>(&global_max), __float_as_int(max_score));
    }
    __syncthreads();
    
    // Compute exp(score - max) and sum
    float sum_exp = 0.0f;
    for (int k_pos = tid; k_pos < T_k; k_pos += block_size) {
        float exp_val = expf(attn_weights[k_pos] - global_max);
        attn_weights[k_pos] = exp_val;
        sum_exp += exp_val;
    }
    
    // Warp-level reduce for sum
    #pragma unroll
    for (int offset = 16; offset > 0; offset >>= 1) {
        sum_exp += __shfl_down_sync(0xffffffff, sum_exp, offset);
    }
    
    __shared__ float global_sum;
    if (tid == 0) global_sum = 0.0f;
    __syncthreads();
    
    if ((tid & 31) == 0) {
        atomicAdd(&global_sum, sum_exp);
    }
    __syncthreads();
    
    // Compute weighted sum of V
    for (int d = tid; d < head_dim; d += block_size) {
        float out_val = 0.0f;
        
        for (int k_pos = 0; k_pos < T_k; k_pos++) {
            float attn = attn_weights[k_pos] / global_sum;
            const int v_base = batch * H * T_k * head_dim + head * T_k * head_dim + k_pos * head_dim;
            out_val += attn * static_cast<float>(v[v_base + d]);
        }
        
        const int out_idx = batch * H * T_q * head_dim + head * T_q * head_dim + q_pos * head_dim + d;
        out[out_idx] = static_cast<scalar_t>(out_val);
    }
}

// RoPE kernel - applies rotary positional embeddings
template <typename scalar_t>
__global__ void rope_kernel(
    const scalar_t* __restrict__ x,
    scalar_t* __restrict__ x_out,
    const float* __restrict__ cos,
    const float* __restrict__ sin,
    const int total_elements,
    const int T,
    const int H,
    const int head_dim
) {
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= total_elements) return;
    
    const int d = idx % head_dim;
    const int t = (idx / head_dim) % T;
    const int head = (idx / (head_dim * T)) % H;
    const int batch = idx / (head_dim * T * H);
    
    const int half_dim = head_dim / 2;
    const float c = cos[t * head_dim + d];
    const float s = sin[t * head_dim + d];
    
    int d_pair = (d < half_dim) ? (d + half_dim) : (d - half_dim);
    const int idx_pair = batch * head_dim * T * H + head * head_dim * T + t * head_dim + d_pair;
    
    const scalar_t x_val = x[idx];
    const scalar_t x_pair = x[idx_pair];
    
    if (d < half_dim) {
        x_out[idx] = static_cast<scalar_t>(static_cast<float>(x_val) * c - static_cast<float>(x_pair) * s);
    } else {
        x_out[idx] = static_cast<scalar_t>(static_cast<float>(x_val) * c + static_cast<float>(x_pair) * s);
    }
}

// Host function to compute RoPE frequencies and apply
std::tuple<at::Tensor, at::Tensor> compute_rope_freqs(int seq_len, int head_dim, float rope_base, at::Device device) {
    auto inv_freq = 1.0f / torch::pow(rope_base, torch::arange(0, head_dim, 2, torch::TensorOptions().dtype(torch::kFloat32).device(device)) / head_dim);
    auto t = torch::arange(seq_len, torch::TensorOptions().dtype(torch::kFloat32).device(device));
    auto freqs = torch::outer(t, inv_freq);
    auto emb = torch::cat({freqs, freqs}, -1);
    return {emb.cos(), emb.sin()};
}

// Combined RoPE + Attention forward
at::Tensor fused_rope_attention_forward(
    const at::Tensor& q_proj,     // (B, T_q, H, head_dim)
    const at::Tensor& k_proj,     // (B, T_k, H, head_dim)
    const at::Tensor& v_proj,     // (B, T_k, H, head_dim)
    const at::Tensor& cos_q,
    const at::Tensor& sin_q,
    const at::Tensor& cos_k,
    const at::Tensor& sin_k,
    float scale
) {
    const int B = q_proj.size(0);
    const int T_q = q_proj.size(1);
    const int H = q_proj.size(2);
    const int head_dim = q_proj.size(3);
    const int T_k = k_proj.size(1);
    
    // Output tensor
    auto out = at::empty({B, T_q, H, head_dim}, q_proj.options());
    
    // Apply RoPE to Q
    auto q_rope = at::empty_like(q_proj);
    const int total_q = B * T_q * H * head_dim;
    const int threads = 256;
    const int blocks_q = (total_q + threads - 1) / threads;
    
    AT_DISPATCH_FLOATING_TYPES_AND_HALF(q_proj.scalar_type(), "rope_q", ([&] {
        rope_kernel<scalar_t><<<blocks_q, threads>>>(
            q_proj.data_ptr<scalar_t>(),
            q_rope.data_ptr<scalar_t>(),
            cos_q.data_ptr<float>(),
            sin_q.data_ptr<float>(),
            total_q, T_q, H, head_dim
        );
    }));
    
    // Apply RoPE to K
    auto k_rope = at::empty_like(k_proj);
    const int total_k = B * T_k * H * head_dim;
    const int blocks_k = (total_k + threads - 1) / threads;
    
    AT_DISPATCH_FLOATING_TYPES_AND_HALF(k_proj.scalar_type(), "rope_k", ([&] {
        rope_kernel<scalar_t><<<blocks_k, threads>>>(
            k_proj.data_ptr<scalar_t>(),
            k_rope.data_ptr<scalar_t>(),
            cos_k.data_ptr<float>(),
            sin_k.data_ptr<float>(),
            total_k, T_k, H, head_dim
        );
    }));
    
    // Fused attention kernel
    // Grid: (B * H) blocks in x, T_q blocks in y
    dim3 grid(B * H, T_q);
    dim3 block(min(256, max(32, T_k)));
    const int shared_mem = (head_dim + T_k) * sizeof(float);
    
    AT_DISPATCH_FLOATING_TYPES_AND_HALF(q_proj.scalar_type(), "fused_attention", ([&] {
        fused_attention_kernel<scalar_t><<<grid, block, shared_mem>>>(
            q_rope.data_ptr<scalar_t>(),
            k_rope.data_ptr<scalar_t>(),
            v_proj.data_ptr<scalar_t>(),
            out.data_ptr<scalar_t>(),
            scale,
            B, H, T_q, T_k, head_dim
        );
    }));
    
    return out;
}
''',
    functions=['fused_rope_attention_forward'],
    extra_cuda_cflags=['-O3', '--use_fast_math'],
)

class ModelNew(nn.Module):
    def __init__(
        self,
        dim: int,
        num_heads: int,
        dropout: float = 0.0,
        rope_base: float = 10000.0,
    ):
        super().__init__()
        assert dim % num_heads == 0
        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        assert self.head_dim % 2 == 0, "RoPE requires an even head dimension"

        # Keep projections as standard PyTorch layers (well-optimized)
        self.q_proj = nn.Linear(dim, dim, bias=False)
        self.k_proj = nn.Linear(dim, dim, bias=False)
        self.v_proj = nn.Linear(dim, dim, bias=False)
        self.out_proj = nn.Linear(dim, dim, bias=False)
        self.dropout = nn.Dropout(dropout)
        self.scale = 1.0 / math.sqrt(self.head_dim)
        self.rope_base = rope_base

        # Precompute RoPE frequencies
        inv_freq = 1.0 / (
            rope_base ** (torch.arange(0, self.head_dim, 2).float() / self.head_dim)
        )
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    def _compute_rope_embeddings(self, seq_len: int, device: torch.device) -> tuple:
        """Compute cos and sin for RoPE."""
        t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
        freqs = torch.outer(t, self.inv_freq.to(device))
        emb = torch.cat([freqs, freqs], dim=-1)
        return emb.cos(), emb.sin()

    def forward(self, x_query: torch.Tensor, x_kv: torch.Tensor) -> torch.Tensor:
        B, T_q, _ = x_query.shape
        _, T_k, _ = x_kv.shape
        H, d = self.num_heads, self.head_dim

        # Linear projections
        q = self.q_proj(x_query)
        k = self.k_proj(x_kv)
        v = self.v_proj(x_kv)

        # Reshape to (B, T, H, head_dim) for kernel
        q = q.view(B, T_q, H, d)
        k = k.view(B, T_k, H, d)
        v = v.view(B, T_k, H, d)

        # Compute RoPE embeddings
        cos_q, sin_q = self._compute_rope_embeddings(T_q, x_query.device)
        cos_k, sin_k = self._compute_rope_embeddings(T_k, x_kv.device)

        # Fused RoPE + Attention
        out = fused_kernel.fused_rope_attention_forward(
            q, k, v, cos_q, sin_q, cos_k, sin_k, self.scale
        )

        # Reshape back: (B, T_q, H, head_dim) -> (B, T_q, dim)
        out = out.reshape(B, T_q, self.dim)

        # Output projection
        return self.out_proj(out)