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

19
csrc/embedding/embedding.cu
View File

@@ -1,4 +1,5 @@
#include <cuda_bf16.h>
#include "../common.cuh"
// Embedding lookup: out[seq_idx] = table[token_ids[seq_idx]]
// Grid: num_tokens, Block: handles hidden_size elements per token.
@@ -7,10 +8,12 @@ __global__ void embedding_f32(
const float* __restrict__ table, // [vocab_size, hidden_size]
const int* __restrict__ token_ids, // [num_tokens]
float* __restrict__ out, // [num_tokens, hidden_size]
int hidden_size
int hidden_size,
int vocab_size
) {
int token_idx = blockIdx.x;
int tid = token_ids[token_idx];
if (tid < 0 || tid >= vocab_size) return;
const float* row = table + tid * hidden_size;
float* dst = out + token_idx * hidden_size;
@@ -23,10 +26,12 @@ __global__ void embedding_bf16(
const __nv_bfloat16* __restrict__ table,
const int* __restrict__ token_ids,
__nv_bfloat16* __restrict__ out,
int hidden_size
int hidden_size,
int vocab_size
) {
int token_idx = blockIdx.x;
int tid = token_ids[token_idx];
if (tid < 0 || tid >= vocab_size) return;
const __nv_bfloat16* row = table + tid * hidden_size;
__nv_bfloat16* dst = out + token_idx * hidden_size;
@@ -38,18 +43,20 @@ __global__ void embedding_bf16(
extern "C" {
void launch_embedding_f32(const void* table, const void* token_ids, void* out,
int num_tokens, int hidden_size, void* stream) {
int num_tokens, int hidden_size, int vocab_size, void* stream) {
int block = (hidden_size < 256) ? hidden_size : 256;
embedding_f32<<<num_tokens, block, 0, (cudaStream_t)stream>>>(
(const float*)table, (const int*)token_ids, (float*)out, hidden_size);
(const float*)table, (const int*)token_ids, (float*)out, hidden_size, vocab_size);
CUDA_CHECK_LAST_ERROR();
}
void launch_embedding_bf16(const void* table, const void* token_ids, void* out,
int num_tokens, int hidden_size, void* stream) {
int num_tokens, int hidden_size, int vocab_size, void* stream) {
int block = (hidden_size < 256) ? hidden_size : 256;
embedding_bf16<<<num_tokens, block, 0, (cudaStream_t)stream>>>(
(const __nv_bfloat16*)table, (const int*)token_ids,
(__nv_bfloat16*)out, hidden_size);
(__nv_bfloat16*)out, hidden_size, vocab_size);
CUDA_CHECK_LAST_ERROR();
}
}

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();
}
}

8
csrc/embedding/transpose.cu
View File

@@ -1,4 +1,5 @@
#include <cuda_bf16.h>
#include "../common.cuh"
// Transpose between [S, H, D] and [H, S, D] layouts (used for RoPE and attention).
// Also handles [S, H*D] → [H, S, D] (reshape_heads) and reverse (merge_heads).
@@ -169,6 +170,7 @@ void launch_reshape_heads_bf16(const void* in, void* out,
int grid = (total + block - 1) / block;
reshape_heads_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
(const __nv_bfloat16*)in, (__nv_bfloat16*)out, seq_len, num_heads, head_dim);
CUDA_CHECK_LAST_ERROR();
}
void launch_merge_heads_bf16(const void* in, void* out,
@@ -178,6 +180,7 @@ void launch_merge_heads_bf16(const void* in, void* out,
int grid = (total + block - 1) / block;
merge_heads_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
(const __nv_bfloat16*)in, (__nv_bfloat16*)out, seq_len, num_heads, head_dim);
CUDA_CHECK_LAST_ERROR();
}
void launch_transpose_hsd_to_shd_bf16(const void* in, void* out,
@@ -187,6 +190,7 @@ void launch_transpose_hsd_to_shd_bf16(const void* in, void* out,
int grid = (total + block - 1) / block;
transpose_hsd_to_shd_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
(const __nv_bfloat16*)in, (__nv_bfloat16*)out, seq_len, num_heads, head_dim);
CUDA_CHECK_LAST_ERROR();
}
void launch_transpose_shd_to_hsd_bf16(const void* in, void* out,
@@ -196,6 +200,7 @@ void launch_transpose_shd_to_hsd_bf16(const void* in, void* out,
int grid = (total + block - 1) / block;
transpose_shd_to_hsd_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
(const __nv_bfloat16*)in, (__nv_bfloat16*)out, seq_len, num_heads, head_dim);
CUDA_CHECK_LAST_ERROR();
}
void launch_repeat_kv_bf16(const void* in, void* out,
@@ -205,6 +210,7 @@ void launch_repeat_kv_bf16(const void* in, void* out,
int grid = (total + block - 1) / block;
repeat_kv_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
(const __nv_bfloat16*)in, (__nv_bfloat16*)out, kv_heads, n_rep, seq_len, head_dim);
CUDA_CHECK_LAST_ERROR();
}
void launch_strided_copy_bf16(const void* in, void* out, int numel, int ndim,
@@ -217,6 +223,7 @@ void launch_strided_copy_bf16(const void* in, void* out, int numel, int ndim,
(const __nv_bfloat16*)in, (__nv_bfloat16*)out, numel, ndim,
shape0, shape1, shape2, shape3,
in_stride0, in_stride1, in_stride2, in_stride3, in_offset);
CUDA_CHECK_LAST_ERROR();
}
void launch_strided_copy_f32(const void* in, void* out, int numel, int ndim,
@@ -229,6 +236,7 @@ void launch_strided_copy_f32(const void* in, void* out, int numel, int ndim,
(const float*)in, (float*)out, numel, ndim,
shape0, shape1, shape2, shape3,
in_stride0, in_stride1, in_stride2, in_stride3, in_offset);
CUDA_CHECK_LAST_ERROR();
}
}