gahow/xserv
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

kernels/cuda: paged-attention kernel, dispatch, pinned host memory

Browse Source 
CUDA layer for the paged-KV + swap work:
- csrc: new paged_attention.cu plus updates across attention/gemm/norm/
  activation/embedding/reduce kernels and common.cuh.
- xserv-kernels: new dispatch module and kernel-binding updates.
- xserv-cuda: cudaMallocHost/FreeHost bindings + PinnedBuffer (host swap
  pool backing) and offset-aware D2H/H2D copies used to move KV blocks
  between the GPU pool and pinned host memory.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-05-28 19:58:36 +08:00
parent 3f1c3d429a
commit 4c3f914459
27 changed files with 581 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
csrc/embedding/rope.cu
View File

@@ -1,10 +1,11 @@
#include <cuda_bf16.h>
#include <math.h>
#include "../common.cuh"
// RoPE: Rotary Position Embedding
// For each pair (x[2i], x[2i+1]) at position `pos`:
// y[2i] = x[2i] * cos - x[2i+1] * sin
// y[2i+1] = x[2i] * sin + x[2i+1] * cos
// RoPE: Rotary Position Embedding, using the Qwen/Llama rotate_half layout.
// For each dimension i in the first half at position `pos`:
// y[i] = x[i] * cos - x[i + half_dim] * sin
// y[i + half_dim] = x[i + half_dim] * cos + x[i] * sin
// where cos/sin come from precomputed cos_cache/sin_cache.
//
// cos_cache[pos][i] = cos(pos * freq[i])
@@ -35,11 +36,11 @@ __global__ void rope_f32(
float sin_val = sin_cache[pos * half_dim + pair_idx];
int base = (token_idx * num_heads + head_idx) * head_dim;
float x0 = x[base + 2 * pair_idx];
float x1 = x[base + 2 * pair_idx + 1];
float x0 = x[base + pair_idx];
float x1 = x[base + pair_idx + half_dim];
x[base + 2 * pair_idx] = x0 * cos_val - x1 * sin_val;
x[base + 2 * pair_idx + 1] = x0 * sin_val + x1 * cos_val;
x[base + pair_idx] = x0 * cos_val - x1 * sin_val;
x[base + pair_idx + half_dim] = x1 * cos_val + x0 * sin_val;
}
__global__ void rope_bf16(
@@ -61,11 +62,11 @@ __global__ void rope_bf16(
float sin_val = sin_cache[pos * half_dim + pair_idx];
int base = (token_idx * num_heads + head_idx) * head_dim;
float x0 = __bfloat162float(x[base + 2 * pair_idx]);
float x1 = __bfloat162float(x[base + 2 * pair_idx + 1]);
float x0 = __bfloat162float(x[base + pair_idx]);
float x1 = __bfloat162float(x[base + pair_idx + half_dim]);
x[base + 2 * pair_idx] = __float2bfloat16(x0 * cos_val - x1 * sin_val);
x[base + 2 * pair_idx + 1] = __float2bfloat16(x0 * sin_val + x1 * cos_val);
x[base + pair_idx] = __float2bfloat16(x0 * cos_val - x1 * sin_val);
x[base + pair_idx + half_dim] = __float2bfloat16(x1 * cos_val + x0 * sin_val);
}
// Precompute cos/sin cache on GPU
@@ -94,6 +95,7 @@ void launch_rope_f32(void* x, const void* cos_cache, const void* sin_cache,
rope_f32<<<grid, block, 0, (cudaStream_t)stream>>>(
(float*)x, (const float*)cos_cache, (const float*)sin_cache,
(const int*)positions, num_heads, head_dim);
CUDA_CHECK_LAST_ERROR();
}
void launch_rope_bf16(void* x, const void* cos_cache, const void* sin_cache,
@@ -104,6 +106,7 @@ void launch_rope_bf16(void* x, const void* cos_cache, const void* sin_cache,
rope_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
(__nv_bfloat16*)x, (const float*)cos_cache, (const float*)sin_cache,
(const int*)positions, num_heads, head_dim);
CUDA_CHECK_LAST_ERROR();
}
void launch_compute_rope_cache(void* cos_cache, void* sin_cache,
@@ -111,6 +114,7 @@ void launch_compute_rope_cache(void* cos_cache, void* sin_cache,
void* stream) {
compute_rope_cache<<<max_seq_len, half_dim, 0, (cudaStream_t)stream>>>(
(float*)cos_cache, (float*)sin_cache, max_seq_len, half_dim, theta);
CUDA_CHECK_LAST_ERROR();
}
}