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;
|
||||
|
||||
349
csrc/attention/deltanet.cu
Normal file
349
csrc/attention/deltanet.cu
Normal file
@@ -0,0 +1,349 @@
|
||||
#include <cuda_bf16.h>
|
||||
#include <math.h>
|
||||
#include "../common.cuh"
|
||||
|
||||
__global__ void depthwise_causal_conv1d_silu_bf16_kernel(
|
||||
const __nv_bfloat16* __restrict__ x,
|
||||
const __nv_bfloat16* __restrict__ w,
|
||||
__nv_bfloat16* __restrict__ out,
|
||||
int seq_len,
|
||||
int channels,
|
||||
int kernel
|
||||
) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int total = seq_len * channels;
|
||||
if (idx >= total) return;
|
||||
|
||||
int c = idx % channels;
|
||||
int t = idx / channels;
|
||||
float acc = 0.0f;
|
||||
for (int k = 0; k < kernel; ++k) {
|
||||
int src_t = t - (kernel - 1 - k);
|
||||
if (src_t >= 0) {
|
||||
float xv = __bfloat162float(x[src_t * channels + c]);
|
||||
float wv = __bfloat162float(w[c * kernel + k]);
|
||||
acc += xv * wv;
|
||||
}
|
||||
}
|
||||
float y = acc / (1.0f + expf(-acc));
|
||||
out[idx] = __float2bfloat16(y);
|
||||
}
|
||||
|
||||
__global__ void deltanet_ar_bf16_kernel(
|
||||
const __nv_bfloat16* __restrict__ q,
|
||||
const __nv_bfloat16* __restrict__ k,
|
||||
const __nv_bfloat16* __restrict__ v,
|
||||
const __nv_bfloat16* __restrict__ beta,
|
||||
const __nv_bfloat16* __restrict__ alpha_softplus,
|
||||
const __nv_bfloat16* __restrict__ a_log,
|
||||
__nv_bfloat16* __restrict__ state,
|
||||
__nv_bfloat16* __restrict__ out,
|
||||
int key_heads,
|
||||
int value_heads,
|
||||
int head_dim
|
||||
) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int total = value_heads * head_dim;
|
||||
if (idx >= total) return;
|
||||
|
||||
int h = idx / head_dim;
|
||||
int j = idx % head_dim;
|
||||
int key_h = h * key_heads / value_heads;
|
||||
|
||||
float beta_h = __bfloat162float(beta[h]);
|
||||
float gate = expf(__bfloat162float(alpha_softplus[h]) * __bfloat162float(a_log[h]));
|
||||
|
||||
float sk = 0.0f;
|
||||
for (int i = 0; i < head_dim; ++i) {
|
||||
float s = __bfloat162float(state[(h * head_dim + i) * head_dim + j]);
|
||||
float ki = __bfloat162float(k[key_h * head_dim + i]);
|
||||
s *= gate;
|
||||
state[(h * head_dim + i) * head_dim + j] = __float2bfloat16(s);
|
||||
sk += s * ki;
|
||||
}
|
||||
|
||||
float vj = __bfloat162float(v[h * head_dim + j]);
|
||||
float d = (vj - sk) * beta_h;
|
||||
|
||||
float out_j = 0.0f;
|
||||
for (int i = 0; i < head_dim; ++i) {
|
||||
float ki = __bfloat162float(k[key_h * head_dim + i]);
|
||||
float qi = __bfloat162float(q[key_h * head_dim + i]) / sqrtf((float)head_dim);
|
||||
int sidx = (h * head_dim + i) * head_dim + j;
|
||||
float s = __bfloat162float(state[sidx]) + ki * d;
|
||||
state[sidx] = __float2bfloat16(s);
|
||||
out_j += s * qi;
|
||||
}
|
||||
out[h * head_dim + j] = __float2bfloat16(out_j);
|
||||
}
|
||||
|
||||
// Stateful depthwise causal convolution for hybrid recurrent attention.
|
||||
// `state` stores the preceding `kernel - 1` inputs in chronological order.
|
||||
// One thread owns a channel, so the state update is race-free for any seq_len.
|
||||
__global__ void depthwise_causal_conv1d_stateful_silu_bf16_kernel(
|
||||
const __nv_bfloat16* __restrict__ x,
|
||||
const __nv_bfloat16* __restrict__ w,
|
||||
__nv_bfloat16* __restrict__ state,
|
||||
__nv_bfloat16* __restrict__ out,
|
||||
int seq_len,
|
||||
int channels,
|
||||
int kernel
|
||||
) {
|
||||
int c = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (c >= channels) return;
|
||||
|
||||
const int history = kernel - 1;
|
||||
for (int t = 0; t < seq_len; ++t) {
|
||||
float acc = 0.0f;
|
||||
for (int k = 0; k < kernel; ++k) {
|
||||
int src_t = t - history + k;
|
||||
float xv;
|
||||
if (src_t < 0) {
|
||||
xv = __bfloat162float(state[c * history + history + src_t]);
|
||||
} else {
|
||||
xv = __bfloat162float(x[(long long)src_t * channels + c]);
|
||||
}
|
||||
acc += xv * __bfloat162float(w[c * kernel + k]);
|
||||
}
|
||||
out[(long long)t * channels + c] =
|
||||
__float2bfloat16(acc / (1.0f + expf(-acc)));
|
||||
}
|
||||
|
||||
// Keep the last `history` raw projection values for the next call.
|
||||
for (int h = 0; h < history; ++h) {
|
||||
int src_t = seq_len - history + h;
|
||||
__nv_bfloat16 value;
|
||||
if (src_t < 0) {
|
||||
value = state[c * history + history + src_t];
|
||||
} else {
|
||||
value = x[(long long)src_t * channels + c];
|
||||
}
|
||||
state[c * history + h] = value;
|
||||
}
|
||||
}
|
||||
|
||||
// L2-normalize each row using FP32 accumulation. Qwen3.5/3.6 applies this to
|
||||
// convolved Q and K before the delta rule (this is not RMSNorm).
|
||||
__global__ void l2_normalize_rows_bf16_kernel(
|
||||
const __nv_bfloat16* __restrict__ x,
|
||||
__nv_bfloat16* __restrict__ out,
|
||||
int rows,
|
||||
int cols,
|
||||
float eps
|
||||
) {
|
||||
int row = blockIdx.x;
|
||||
if (row >= rows) return;
|
||||
|
||||
float local = 0.0f;
|
||||
for (int i = threadIdx.x; i < cols; i += blockDim.x) {
|
||||
float v = __bfloat162float(x[(long long)row * cols + i]);
|
||||
local += v * v;
|
||||
}
|
||||
__shared__ float sums[256];
|
||||
sums[threadIdx.x] = local;
|
||||
__syncthreads();
|
||||
for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
|
||||
if (threadIdx.x < stride) sums[threadIdx.x] += sums[threadIdx.x + stride];
|
||||
__syncthreads();
|
||||
}
|
||||
// Match torch.nn.functional.normalize / llama.cpp: clamp the norm by eps,
|
||||
// rather than adding eps to every non-zero norm.
|
||||
float inv = rsqrtf(fmaxf(sums[0], eps * eps));
|
||||
for (int i = threadIdx.x; i < cols; i += blockDim.x) {
|
||||
out[(long long)row * cols + i] =
|
||||
__float2bfloat16(__bfloat162float(x[(long long)row * cols + i]) * inv);
|
||||
}
|
||||
}
|
||||
|
||||
// Recurrent Gated DeltaNet over one or more tokens. HF stores V heads grouped
|
||||
// by K head, so V head h uses K/Q head floor(h / (value_heads/key_heads)).
|
||||
// The recurrent matrix is FP32 as required by `mamba_ssm_dtype=float32`.
|
||||
__global__ void deltanet_recurrent_bf16_f32_kernel(
|
||||
const __nv_bfloat16* __restrict__ q,
|
||||
const __nv_bfloat16* __restrict__ k,
|
||||
const __nv_bfloat16* __restrict__ v,
|
||||
const __nv_bfloat16* __restrict__ beta,
|
||||
const __nv_bfloat16* __restrict__ alpha_softplus,
|
||||
const __nv_bfloat16* __restrict__ a_log,
|
||||
float* __restrict__ state,
|
||||
__nv_bfloat16* __restrict__ out,
|
||||
int seq_len,
|
||||
int key_heads,
|
||||
int value_heads,
|
||||
int head_dim
|
||||
) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int total = value_heads * head_dim;
|
||||
if (idx >= total) return;
|
||||
|
||||
int h = idx / head_dim;
|
||||
int j = idx % head_dim;
|
||||
int key_h = h * key_heads / value_heads;
|
||||
float a = __bfloat162float(a_log[h]);
|
||||
float neg_a = -expf(a);
|
||||
float q_scale = rsqrtf((float)head_dim);
|
||||
|
||||
for (int t = 0; t < seq_len; ++t) {
|
||||
const __nv_bfloat16* q_t = q + ((long long)t * key_heads + key_h) * head_dim;
|
||||
const __nv_bfloat16* k_t = k + ((long long)t * key_heads + key_h) * head_dim;
|
||||
const __nv_bfloat16* v_t = v + ((long long)t * value_heads + h) * head_dim;
|
||||
float b = __bfloat162float(beta[(long long)t * value_heads + h]);
|
||||
float alpha = __bfloat162float(alpha_softplus[(long long)t * value_heads + h]);
|
||||
float decay = expf(neg_a * alpha);
|
||||
|
||||
float sk = 0.0f;
|
||||
for (int i = 0; i < head_dim; ++i) {
|
||||
long long sidx = ((long long)h * head_dim + i) * head_dim + j;
|
||||
float s = state[sidx] * decay;
|
||||
state[sidx] = s;
|
||||
sk += s * __bfloat162float(k_t[i]);
|
||||
}
|
||||
|
||||
float d = (__bfloat162float(v_t[j]) - sk) * b;
|
||||
float out_j = 0.0f;
|
||||
for (int i = 0; i < head_dim; ++i) {
|
||||
long long sidx = ((long long)h * head_dim + i) * head_dim + j;
|
||||
float s = state[sidx] + __bfloat162float(k_t[i]) * d;
|
||||
state[sidx] = s;
|
||||
out_j += s * (__bfloat162float(q_t[i]) * q_scale);
|
||||
}
|
||||
out[((long long)t * value_heads + h) * head_dim + j] = __float2bfloat16(out_j);
|
||||
}
|
||||
}
|
||||
|
||||
extern "C" {
|
||||
|
||||
void launch_depthwise_causal_conv1d_silu_bf16(
|
||||
const void* x,
|
||||
const void* w,
|
||||
void* out,
|
||||
int seq_len,
|
||||
int channels,
|
||||
int kernel,
|
||||
void* stream
|
||||
) {
|
||||
int total = seq_len * channels;
|
||||
int block = 256;
|
||||
int grid = (total + block - 1) / block;
|
||||
depthwise_causal_conv1d_silu_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
|
||||
(const __nv_bfloat16*)x,
|
||||
(const __nv_bfloat16*)w,
|
||||
(__nv_bfloat16*)out,
|
||||
seq_len,
|
||||
channels,
|
||||
kernel
|
||||
);
|
||||
CUDA_CHECK_LAST_ERROR();
|
||||
}
|
||||
|
||||
void launch_deltanet_ar_bf16(
|
||||
const void* q,
|
||||
const void* k,
|
||||
const void* v,
|
||||
const void* beta,
|
||||
const void* alpha_softplus,
|
||||
const void* a_log,
|
||||
void* state,
|
||||
void* out,
|
||||
int key_heads,
|
||||
int value_heads,
|
||||
int head_dim,
|
||||
void* stream
|
||||
) {
|
||||
int total = value_heads * head_dim;
|
||||
int block = 128;
|
||||
int grid = (total + block - 1) / block;
|
||||
deltanet_ar_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
|
||||
(const __nv_bfloat16*)q,
|
||||
(const __nv_bfloat16*)k,
|
||||
(const __nv_bfloat16*)v,
|
||||
(const __nv_bfloat16*)beta,
|
||||
(const __nv_bfloat16*)alpha_softplus,
|
||||
(const __nv_bfloat16*)a_log,
|
||||
(__nv_bfloat16*)state,
|
||||
(__nv_bfloat16*)out,
|
||||
key_heads,
|
||||
value_heads,
|
||||
head_dim
|
||||
);
|
||||
CUDA_CHECK_LAST_ERROR();
|
||||
}
|
||||
|
||||
void launch_depthwise_causal_conv1d_stateful_silu_bf16(
|
||||
const void* x,
|
||||
const void* w,
|
||||
void* state,
|
||||
void* out,
|
||||
int seq_len,
|
||||
int channels,
|
||||
int kernel,
|
||||
void* stream
|
||||
) {
|
||||
int block = 256;
|
||||
int grid = (channels + block - 1) / block;
|
||||
depthwise_causal_conv1d_stateful_silu_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
|
||||
(const __nv_bfloat16*)x,
|
||||
(const __nv_bfloat16*)w,
|
||||
(__nv_bfloat16*)state,
|
||||
(__nv_bfloat16*)out,
|
||||
seq_len,
|
||||
channels,
|
||||
kernel
|
||||
);
|
||||
CUDA_CHECK_LAST_ERROR();
|
||||
}
|
||||
|
||||
void launch_l2_normalize_rows_bf16(
|
||||
const void* x,
|
||||
void* out,
|
||||
int rows,
|
||||
int cols,
|
||||
float eps,
|
||||
void* stream
|
||||
) {
|
||||
l2_normalize_rows_bf16_kernel<<<rows, 256, 0, (cudaStream_t)stream>>>(
|
||||
(const __nv_bfloat16*)x,
|
||||
(__nv_bfloat16*)out,
|
||||
rows,
|
||||
cols,
|
||||
eps
|
||||
);
|
||||
CUDA_CHECK_LAST_ERROR();
|
||||
}
|
||||
|
||||
void launch_deltanet_recurrent_bf16_f32(
|
||||
const void* q,
|
||||
const void* k,
|
||||
const void* v,
|
||||
const void* beta,
|
||||
const void* alpha_softplus,
|
||||
const void* a_log,
|
||||
void* state,
|
||||
void* out,
|
||||
int seq_len,
|
||||
int key_heads,
|
||||
int value_heads,
|
||||
int head_dim,
|
||||
void* stream
|
||||
) {
|
||||
int total = value_heads * head_dim;
|
||||
int block = 128;
|
||||
int grid = (total + block - 1) / block;
|
||||
deltanet_recurrent_bf16_f32_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
|
||||
(const __nv_bfloat16*)q,
|
||||
(const __nv_bfloat16*)k,
|
||||
(const __nv_bfloat16*)v,
|
||||
(const __nv_bfloat16*)beta,
|
||||
(const __nv_bfloat16*)alpha_softplus,
|
||||
(const __nv_bfloat16*)a_log,
|
||||
(float*)state,
|
||||
(__nv_bfloat16*)out,
|
||||
seq_len,
|
||||
key_heads,
|
||||
value_heads,
|
||||
head_dim
|
||||
);
|
||||
CUDA_CHECK_LAST_ERROR();
|
||||
}
|
||||
|
||||
}
|
||||
@@ -13,13 +13,13 @@
|
||||
// O [batch, num_q_heads, q_len, head_dim]
|
||||
//
|
||||
// Shared memory (BF16):
|
||||
// smem_q[BR][head_dim] — 64 * 128 * 2 = 16 KB (loaded once per Q tile)
|
||||
// smem_kv[BC][head_dim] — 64 * 128 * 2 = 16 KB (alternates K and V)
|
||||
// Total: 32 KB (fits in default 48 KB shared memory)
|
||||
// Use 32-row tiles so head_dim=256 still fits in the default 48 KB shared
|
||||
// memory limit: 2 * 32 * 256 * 2 = 32 KB.
|
||||
|
||||
#define BR 64
|
||||
#define BC 64
|
||||
#define BR 32
|
||||
#define BC 32
|
||||
#define THREADS_PER_BLOCK 128
|
||||
#define FLASH_HEAD_DIM_MAX 256
|
||||
|
||||
__global__ void flash_attention_bf16_kernel(
|
||||
const __nv_bfloat16* __restrict__ Q,
|
||||
@@ -73,8 +73,8 @@ __global__ void flash_attention_bf16_kernel(
|
||||
// Thread t (0 <= t < q_tile_rows) owns Q row t
|
||||
bool owns_row = (tid < q_tile_rows);
|
||||
|
||||
// Per-thread FP32 accumulators (head_dim up to 128)
|
||||
float O_acc[128];
|
||||
// Per-thread FP32 accumulators (head_dim up to 256)
|
||||
float O_acc[FLASH_HEAD_DIM_MAX];
|
||||
float m_val = -INFINITY;
|
||||
float l_val = 0.0f;
|
||||
if (owns_row) {
|
||||
@@ -250,7 +250,7 @@ __global__ void flash_attention_sinks_bf16_kernel(
|
||||
|
||||
bool owns_row = (tid < q_tile_rows);
|
||||
|
||||
float O_acc[128];
|
||||
float O_acc[FLASH_HEAD_DIM_MAX];
|
||||
float m_val = -INFINITY;
|
||||
float l_val = 0.0f;
|
||||
if (owns_row) {
|
||||
@@ -384,7 +384,7 @@ __global__ void flash_attention_sinks_bf16_kernel(
|
||||
// ============================================================
|
||||
|
||||
#define DECODE_THREADS 256
|
||||
#define HEAD_DIM_MAX 128
|
||||
#define HEAD_DIM_MAX 256
|
||||
|
||||
__global__ void decode_attention_bf16_kernel(
|
||||
const __nv_bfloat16* __restrict__ Q,
|
||||
@@ -413,7 +413,7 @@ __global__ void decode_attention_bf16_kernel(
|
||||
const __nv_bfloat16* V_base = V + ((long long)batch_idx * num_kv_heads + kv_head) * kv_len * head_dim;
|
||||
__nv_bfloat16* O_ptr = O + ((long long)batch_idx * num_q_heads + q_head) * head_dim;
|
||||
|
||||
// Load Q vector into registers (head_dim <= 128)
|
||||
// Load Q vector into registers (head_dim <= 256)
|
||||
float q_reg[HEAD_DIM_MAX];
|
||||
for (int d = 0; d < head_dim; d++) {
|
||||
q_reg[d] = __bfloat162float(Q_ptr[d]);
|
||||
|
||||
@@ -21,11 +21,11 @@
|
||||
// active batch) so we just index by blockIdx.x_seq.
|
||||
// context_lens [batch] int32 — number of valid tokens per sequence.
|
||||
//
|
||||
// One CUDA block: 256 threads, head_dim <= 128.
|
||||
// One CUDA block: 256 threads, head_dim <= 256.
|
||||
|
||||
#define PAGED_BLOCK_SIZE 16
|
||||
#define PAGED_THREADS 256
|
||||
#define PAGED_HEAD_DIM_MAX 128
|
||||
#define PAGED_HEAD_DIM_MAX 256
|
||||
|
||||
__global__ void paged_decode_attention_bf16_kernel(
|
||||
const __nv_bfloat16* __restrict__ Q,
|
||||
|
||||
@@ -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) {
|
||||
|
||||
@@ -34,6 +34,51 @@
|
||||
|
||||
#define SPARSE_TILE_N 8 // output columns per block (= warps per block)
|
||||
|
||||
// Weights FP8 E4M3 [local_experts, N, K], activations BF16 (W8A16).
|
||||
// Decode is memory-bound (~2 FLOP/byte), so dequant-in-registers GEMV
|
||||
// loses nothing to tensor cores and skips activation quantization.
|
||||
__global__ void moe_sparse_gemv_bf16_bf16_kernel(
|
||||
const __nv_bfloat16* __restrict__ x, // [T, K] or [T*top_k, K]
|
||||
const __nv_bfloat16* __restrict__ w, // [local_experts, N, K]
|
||||
const int* __restrict__ topk_ids, // [T, top_k] global expert ids
|
||||
__nv_bfloat16* __restrict__ y, // [T, top_k, N]
|
||||
int N, int K, int top_k,
|
||||
int expert_start, int local_experts,
|
||||
int x_per_slot
|
||||
) {
|
||||
int token = blockIdx.z;
|
||||
int slot = blockIdx.y;
|
||||
int eid = topk_ids[token * top_k + slot];
|
||||
int lid = eid - expert_start;
|
||||
if (lid < 0 || lid >= local_experts) return;
|
||||
|
||||
extern __shared__ float xs[];
|
||||
const __nv_bfloat16* xrow =
|
||||
x + (long long)(x_per_slot ? token * top_k + slot : token) * K;
|
||||
for (int i = threadIdx.x; i < K; i += blockDim.x) {
|
||||
xs[i] = __bfloat162float(xrow[i]);
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
int n = blockIdx.x * SPARSE_TILE_N + (threadIdx.x >> 5);
|
||||
if (n >= N) return;
|
||||
int lane = threadIdx.x & 31;
|
||||
|
||||
const __nv_bfloat16* wrow = w + ((long long)lid * N + n) * K;
|
||||
float acc = 0.0f;
|
||||
for (int i = lane; i < K; i += 32) {
|
||||
acc += xs[i] * __bfloat162float(wrow[i]);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int o = 16; o > 0; o >>= 1) {
|
||||
acc += __shfl_down_sync(0xffffffffu, acc, o);
|
||||
}
|
||||
if (lane == 0) {
|
||||
y[((long long)token * top_k + slot) * N + n] = __float2bfloat16(acc);
|
||||
}
|
||||
}
|
||||
|
||||
// Weights FP8 E4M3 [local_experts, N, K], activations BF16 (W8A16).
|
||||
// Decode is memory-bound (~2 FLOP/byte), so dequant-in-registers GEMV
|
||||
// loses nothing to tensor cores and skips activation quantization.
|
||||
@@ -194,6 +239,23 @@ __global__ void moe_weighted_sum_sparse_bf16_kernel(
|
||||
|
||||
extern "C" {
|
||||
|
||||
void launch_moe_sparse_gemv_bf16_bf16(
|
||||
const void* x, const void* w, const void* topk_ids, void* y,
|
||||
int num_tokens, int N, int K, int top_k,
|
||||
int expert_start, int local_experts, int x_per_slot,
|
||||
void* stream
|
||||
) {
|
||||
dim3 grid((N + SPARSE_TILE_N - 1) / SPARSE_TILE_N, top_k, num_tokens);
|
||||
int block = SPARSE_TILE_N * 32;
|
||||
size_t smem = (size_t)K * sizeof(float);
|
||||
moe_sparse_gemv_bf16_bf16_kernel<<<grid, block, smem, (cudaStream_t)stream>>>(
|
||||
(const __nv_bfloat16*)x, (const __nv_bfloat16*)w, (const int*)topk_ids,
|
||||
(__nv_bfloat16*)y,
|
||||
N, K, top_k, expert_start, local_experts, x_per_slot
|
||||
);
|
||||
CUDA_CHECK_LAST_ERROR();
|
||||
}
|
||||
|
||||
void launch_moe_sparse_gemv_fp8_bf16(
|
||||
const void* x, const void* w, const void* w_scales, const void* bias,
|
||||
const void* topk_ids, void* y,
|
||||
|
||||
Reference in New Issue
Block a user