File size: 2,798 Bytes
ab48644 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | // 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
);
}
|