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

Add Qwen3.6 MoE inference support

Browse Source
This commit is contained in:
Gahow Wang
2026-07-13 20:24:41 +08:00
parent 588bfd9df3
commit a2de146fb6
27 changed files with 3153 additions and 149 deletions
Download Patch File Download Diff File Expand all files Collapse all files

54
csrc/embedding/rope.cu
View File

@@ -69,6 +69,42 @@ __global__ void rope_bf16(
x[base + pair_idx + half_dim] = __float2bfloat16(x1 * cos_val + x0 * sin_val);
}
__global__ void partial_rope_bf16(
const __nv_bfloat16* __restrict__ x,
__nv_bfloat16* __restrict__ out,
const int* __restrict__ positions,
int num_heads, int head_dim, int n_rot, float theta
) {
int token_idx = blockIdx.x;
int head_idx = blockIdx.y;
int pair_idx = threadIdx.x;
int half_rot = n_rot / 2;
if (pair_idx >= half_rot) return;
int pos = positions[token_idx];
float freq = 1.0f / powf(theta, (float)(2 * pair_idx) / (float)n_rot);
float angle = (float)pos * freq;
float sin_val, cos_val;
sincosf(angle, &sin_val, &cos_val);
int base = (token_idx * num_heads + head_idx) * head_dim;
float x0 = __bfloat162float(x[base + pair_idx]);
float x1 = __bfloat162float(x[base + pair_idx + half_rot]);
out[base + pair_idx] = __float2bfloat16(x0 * cos_val - x1 * sin_val);
out[base + pair_idx + half_rot] = __float2bfloat16(x1 * cos_val + x0 * sin_val);
}
__global__ void copy_partial_rope_tail_bf16(
const __nv_bfloat16* __restrict__ x,
__nv_bfloat16* __restrict__ out,
int total
) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx >= total) return;
out[idx] = x[idx];
}
// Precompute cos/sin cache on GPU
__global__ void compute_rope_cache(
float* __restrict__ cos_cache, // [max_seq_len, half_dim]
@@ -109,6 +145,24 @@ void launch_rope_bf16(void* x, const void* cos_cache, const void* sin_cache,
CUDA_CHECK_LAST_ERROR();
}
void launch_partial_rope_bf16(const void* x, void* out, const void* positions,
int num_tokens, int num_heads, int head_dim,
int n_rot, float theta, void* stream) {
int total = num_tokens * num_heads * head_dim;
int block_copy = 256;
int grid_copy = (total + block_copy - 1) / block_copy;
copy_partial_rope_tail_bf16<<<grid_copy, block_copy, 0, (cudaStream_t)stream>>>(
(const __nv_bfloat16*)x, (__nv_bfloat16*)out, total);
CUDA_CHECK_LAST_ERROR();
dim3 grid(num_tokens, num_heads);
int block = n_rot / 2;
partial_rope_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
(const __nv_bfloat16*)x, (__nv_bfloat16*)out, (const int*)positions,
num_heads, head_dim, n_rot, theta);
CUDA_CHECK_LAST_ERROR();
}
void launch_compute_rope_cache(void* cos_cache, void* sin_cache,
int max_seq_len, int half_dim, float theta,
void* stream) {