Add Qwen3.6 MoE inference support

This commit is contained in:
2026-07-13 20:24:41 +08:00
parent 588bfd9df3
commit a2de146fb6
27 changed files with 3153 additions and 149 deletions

View File

@@ -36,6 +36,35 @@ __global__ void silu_bf16(const __nv_bfloat16* x, __nv_bfloat16* out, int n) {
if (idx < n) out[idx] = __float2bfloat16(silu_f(__bfloat162float(x[idx])));
}
__global__ void sigmoid_f32(const float* x, float* out, int n) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n) out[idx] = 1.0f / (1.0f + expf(-x[idx]));
}
__global__ void sigmoid_bf16(const __nv_bfloat16* x, __nv_bfloat16* out, int n) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n) {
float v = __bfloat162float(x[idx]);
out[idx] = __float2bfloat16(1.0f / (1.0f + expf(-v)));
}
}
__global__ void softplus_f32(const float* x, float* out, int n) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n) {
float v = x[idx];
out[idx] = log1pf(expf(-fabsf(v))) + fmaxf(v, 0.0f);
}
}
__global__ void softplus_bf16(const __nv_bfloat16* x, __nv_bfloat16* out, int n) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n) {
float v = __bfloat162float(x[idx]);
out[idx] = __float2bfloat16(log1pf(expf(-fabsf(v))) + fmaxf(v, 0.0f));
}
}
__global__ void scale_f32_kernel(const float* x, float* out, float scale, int n) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n) out[idx] = x[idx] * scale;
@@ -108,6 +137,21 @@ __global__ void mul_bf16_kernel(const __nv_bfloat16* a, const __nv_bfloat16* b,
if (idx < n) out[idx] = __float2bfloat16(__bfloat162float(a[idx]) * __bfloat162float(b[idx]));
}
__global__ void row_scale_bf16_kernel(
const __nv_bfloat16* __restrict__ x,
const __nv_bfloat16* __restrict__ scale,
__nv_bfloat16* __restrict__ out,
int rows,
int cols
) {
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = rows * cols;
if (idx >= total) return;
int row = idx / cols;
float v = __bfloat162float(x[idx]) * __bfloat162float(scale[row]);
out[idx] = __float2bfloat16(v);
}
extern "C" {
void launch_gelu_f32(const void* x, void* out, int n, void* stream) {
@@ -140,6 +184,36 @@ void launch_silu_bf16(const void* x, void* out, int n, void* stream) {
CUDA_CHECK_LAST_ERROR();
}
void launch_sigmoid_f32(const void* x, void* out, int n, void* stream) {
int block = 256;
int grid = (n + block - 1) / block;
sigmoid_f32<<<grid, block, 0, (cudaStream_t)stream>>>((const float*)x, (float*)out, n);
CUDA_CHECK_LAST_ERROR();
}
void launch_sigmoid_bf16(const void* x, void* out, int n, void* stream) {
int block = 256;
int grid = (n + block - 1) / block;
sigmoid_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
(const __nv_bfloat16*)x, (__nv_bfloat16*)out, n);
CUDA_CHECK_LAST_ERROR();
}
void launch_softplus_f32(const void* x, void* out, int n, void* stream) {
int block = 256;
int grid = (n + block - 1) / block;
softplus_f32<<<grid, block, 0, (cudaStream_t)stream>>>((const float*)x, (float*)out, n);
CUDA_CHECK_LAST_ERROR();
}
void launch_softplus_bf16(const void* x, void* out, int n, void* stream) {
int block = 256;
int grid = (n + block - 1) / block;
softplus_bf16<<<grid, block, 0, (cudaStream_t)stream>>>(
(const __nv_bfloat16*)x, (__nv_bfloat16*)out, n);
CUDA_CHECK_LAST_ERROR();
}
void launch_scale_f32(const void* x, void* out, float scale, int n, void* stream) {
int block = 256;
int grid = (n + block - 1) / block;
@@ -193,6 +267,15 @@ void launch_mul_bf16(const void* a, const void* b, void* out, int n, void* strea
CUDA_CHECK_LAST_ERROR();
}
void launch_row_scale_bf16(const void* x, const void* scale, void* out, int rows, int cols, void* stream) {
int n = rows * cols;
int block = 256;
int grid = (n + block - 1) / block;
row_scale_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
(const __nv_bfloat16*)x, (const __nv_bfloat16*)scale, (__nv_bfloat16*)out, rows, cols);
CUDA_CHECK_LAST_ERROR();
}
void launch_silu_mul_bf16(const void* gate, const void* up, void* out, int n, void* stream) {
int block = 256;
int grid = (n + block - 1) / block;