Add Qwen3.6 MoE inference support
This commit is contained in:
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user