350 lines
10 KiB
Plaintext
350 lines
10 KiB
Plaintext
#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();
|
|
}
|
|
|
|
}
|