wehub-resource-sync
78 lines
2.3 KiB
Plaintext
78 lines
2.3 KiB
Plaintext
// DS4 ROCm output/pointwise kernels.
|
|
//
|
|
// Included from ds4_cuda.cu in the same translation unit; API launch glue stays
|
|
// in ds4_cuda.cu for now.
|
|
|
|
__global__ static void output_hc_weights_kernel(
|
|
float *out,
|
|
const float *pre,
|
|
const float *scale,
|
|
const float *base,
|
|
uint32_t n_hc,
|
|
uint32_t n_tokens,
|
|
float epsv) {
|
|
uint32_t gid = blockIdx.x * blockDim.x + threadIdx.x;
|
|
uint32_t n = n_tokens * n_hc;
|
|
if (gid >= n) return;
|
|
uint32_t h = gid % n_hc;
|
|
float z = pre[gid] * scale[0] + base[h];
|
|
out[gid] = 1.0f / (1.0f + expf(-z)) + epsv;
|
|
}
|
|
|
|
|
|
__global__ static void swiglu_kernel(float *out, const float *gate, const float *up, uint32_t n, float clamp, float weight) {
|
|
uint32_t i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (i >= n) return;
|
|
float g = gate[i];
|
|
float u = up[i];
|
|
if (clamp > 1.0e-6f) {
|
|
g = fminf(g, clamp);
|
|
u = fminf(fmaxf(u, -clamp), clamp);
|
|
}
|
|
float s = g / (1.0f + expf(-g));
|
|
out[i] = s * u * weight;
|
|
}
|
|
|
|
__global__ static void add_kernel(float *out, const float *a, const float *b, uint32_t n) {
|
|
uint32_t i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (i >= n) return;
|
|
out[i] = a[i] + b[i];
|
|
}
|
|
|
|
__global__ static void directional_steering_project_kernel(
|
|
float *x,
|
|
const float *directions,
|
|
uint32_t layer,
|
|
uint32_t width,
|
|
uint32_t rows,
|
|
float scale) {
|
|
const uint32_t row = blockIdx.x;
|
|
if (row >= rows || width == 0) return;
|
|
|
|
float *xr = x + (uint64_t)row * width;
|
|
const float *dir = directions + (uint64_t)layer * width;
|
|
float sum = 0.0f;
|
|
for (uint32_t i = threadIdx.x; i < width; i += blockDim.x) {
|
|
sum += xr[i] * dir[i];
|
|
}
|
|
|
|
__shared__ float partial[256];
|
|
partial[threadIdx.x] = sum;
|
|
__syncthreads();
|
|
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
|
|
if (threadIdx.x < stride) partial[threadIdx.x] += partial[threadIdx.x + stride];
|
|
__syncthreads();
|
|
}
|
|
|
|
const float coeff = scale * partial[0];
|
|
for (uint32_t i = threadIdx.x; i < width; i += blockDim.x) {
|
|
xr[i] -= coeff * dir[i];
|
|
}
|
|
}
|
|
|
|
__global__ static void zero_kernel(float *out, uint64_t n) {
|
|
uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
|
if (i < n) out[i] = 0.0f;
|
|
}
|
|
|