// Hilbert — Rotary Position Embedding (RoPE) Kernel
// Applies rotary encoding in-place on Q and K tensors
// Supports partial rotation (rope_dim <= head_dim)

#include <cuda_fp16.h>
#include <cuda_bf16.h>
#include <math.h>

__global__ void hilbert_rope_kernel(
    half* __restrict__ q,  // [batch, seq_len, n_heads, d_head]
    half* __restrict__ k,  // [batch, seq_len, n_kv_heads, d_head]
    const int seq_len,
    const int n_heads,
    const int n_kv_heads,
    const int d_head,
    const int rope_dim,    // Number of dimensions to rotate (<= d_head)
    const float theta_base,
    const int start_pos    // For KV-cache continuation
) {
    const int tid = blockIdx.x * blockDim.x + threadIdx.x;
    const int batch_idx = blockIdx.z;
    const int seq_idx = blockIdx.y;
    const int pos = start_pos + seq_idx;

    // Each thread handles one (head, dim_pair) for Q or K
    const int total_pairs = n_heads * (rope_dim / 2);
    if (tid >= total_pairs) return;

    const int head_idx = tid / (rope_dim / 2);
    const int pair_idx = tid % (rope_dim / 2);
    const int d0 = pair_idx * 2;
    const int d1 = d0 + 1;

    // Compute rotation angle: theta_i = pos * theta_base^(-2i/rope_dim)
    float freq = 1.0f / powf(theta_base, (float)(2 * pair_idx) / (float)rope_dim);
    float angle = (float)pos * freq;
    float cos_val = cosf(angle);
    float sin_val = sinf(angle);

    // Rotate Q
    {
        int offset = ((batch_idx * seq_len + seq_idx) * n_heads + head_idx) * d_head;
        float q0 = __half2float(q[offset + d0]);
        float q1 = __half2float(q[offset + d1]);

        q[offset + d0] = __float2half_rn(q0 * cos_val - q1 * sin_val);
        q[offset + d1] = __float2half_rn(q0 * sin_val + q1 * cos_val);
    }

    // Rotate K (only if this head maps to a KV head)
    if (head_idx < n_kv_heads) {
        int kv_pair_idx = pair_idx;
        int offset = ((batch_idx * seq_len + seq_idx) * n_kv_heads + head_idx) * d_head;
        float k0 = __half2float(k[offset + d0]);
        float k1 = __half2float(k[offset + d1]);

        k[offset + d0] = __float2half_rn(k0 * cos_val - k1 * sin_val);
        k[offset + d1] = __float2half_rn(k0 * sin_val + k1 * cos_val);
    }
}

void hilbert_rope_launch(
    half* q,
    half* k,
    int batch_size,
    int seq_len,
    int n_heads,
    int n_kv_heads,
    int d_head,
    int rope_dim,
    float theta_base,
    int start_pos,
    cudaStream_t stream
) {
    int pairs_per_seq = n_heads * (rope_dim / 2);
    int threads = 256;
    int blocks_x = (pairs_per_seq + threads - 1) / threads;

    dim3 grid(blocks_x, seq_len, batch_size);
    dim3 block(threads);

    hilbert_rope_kernel<<<grid, block, 0, stream>>>(
        q, k, seq_len, n_heads, n_kv_heads, d_head,
        rope_dim, theta_base, start_pos
    );
}