chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,377 @@
|
||||
// DS4 ROCm common embedding/dense-matmul kernels and device helpers.
|
||||
//
|
||||
// Included from ds4_cuda.cu before more specialized modules; these helpers are
|
||||
// intentionally kept static in the single translation unit.
|
||||
|
||||
__global__ static void fill_f32_kernel(float *x, uint64_t n, float v) {
|
||||
uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i < n) x[i] = v;
|
||||
}
|
||||
|
||||
__global__ static void embed_token_hc_kernel(float *out, const unsigned short *w, uint32_t token, uint32_t n_embd, uint32_t n_hc) {
|
||||
uint32_t i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint32_t n = n_embd * n_hc;
|
||||
if (i >= n) return;
|
||||
uint32_t e = i % n_embd;
|
||||
out[i] = __half2float(reinterpret_cast<const __half *>(w)[(uint64_t)token * n_embd + e]);
|
||||
}
|
||||
|
||||
__device__ static float embed_q8_0_scale(const unsigned char *blk) {
|
||||
const uint16_t bits = (uint16_t)blk[0] | ((uint16_t)blk[1] << 8);
|
||||
return __half2float(__ushort_as_half((unsigned short)bits));
|
||||
}
|
||||
|
||||
__global__ static void embed_token_hc_q8_0_kernel(
|
||||
float *out,
|
||||
const unsigned char *w,
|
||||
uint32_t token,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc) {
|
||||
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint64_t n = (uint64_t)n_embd * n_hc;
|
||||
if (gid >= n) return;
|
||||
const uint32_t d = gid % n_embd;
|
||||
const uint32_t blocks = (n_embd + 31u) / 32u;
|
||||
const uint32_t b = d >> 5u;
|
||||
const uint32_t j = d & 31u;
|
||||
const unsigned char *blk = w + ((uint64_t)token * blocks + b) * 34u;
|
||||
out[gid] = embed_q8_0_scale(blk) * (float)((const int8_t *)(blk + 2u))[j];
|
||||
}
|
||||
|
||||
__global__ static void embed_tokens_hc_q8_0_kernel(
|
||||
float *out,
|
||||
const int32_t *tokens,
|
||||
const unsigned char *w,
|
||||
uint32_t n_vocab,
|
||||
uint32_t n_tokens,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc) {
|
||||
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint64_t n = (uint64_t)n_tokens * n_hc * n_embd;
|
||||
if (gid >= n) return;
|
||||
const uint32_t d = gid % n_embd;
|
||||
uint64_t tmp = gid / n_embd;
|
||||
const uint32_t t = tmp / n_hc;
|
||||
int32_t tok_i = tokens[t];
|
||||
uint32_t tok = tok_i < 0 ? 0u : (uint32_t)tok_i;
|
||||
if (tok >= n_vocab) tok = 0;
|
||||
const uint32_t blocks = (n_embd + 31u) / 32u;
|
||||
const uint32_t b = d >> 5u;
|
||||
const uint32_t j = d & 31u;
|
||||
const unsigned char *blk = w + ((uint64_t)tok * blocks + b) * 34u;
|
||||
out[gid] = embed_q8_0_scale(blk) * (float)((const int8_t *)(blk + 2u))[j];
|
||||
}
|
||||
|
||||
__global__ static void embed_tokens_hc_kernel(
|
||||
float *out,
|
||||
const int32_t *tokens,
|
||||
const __half *w,
|
||||
uint32_t n_vocab,
|
||||
uint32_t n_tokens,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc) {
|
||||
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint64_t n = (uint64_t)n_tokens * n_hc * n_embd;
|
||||
if (gid >= n) return;
|
||||
uint32_t d = gid % n_embd;
|
||||
uint64_t tmp = gid / n_embd;
|
||||
uint32_t t = tmp / n_hc;
|
||||
int32_t tok_i = tokens[t];
|
||||
uint32_t tok = tok_i < 0 ? 0u : (uint32_t)tok_i;
|
||||
if (tok >= n_vocab) tok = 0;
|
||||
out[gid] = __half2float(w[(uint64_t)tok * n_embd + d]);
|
||||
}
|
||||
|
||||
__device__ static float warp_sum_f32(float v);
|
||||
|
||||
__global__ static void matmul_f16_kernel(
|
||||
float *out,
|
||||
const __half *w,
|
||||
const float *x,
|
||||
uint64_t in_dim,
|
||||
uint64_t out_dim,
|
||||
uint64_t n_tok) {
|
||||
uint64_t row = (uint64_t)blockIdx.x;
|
||||
uint64_t tok = (uint64_t)blockIdx.y;
|
||||
if (row >= out_dim || tok >= n_tok) return;
|
||||
|
||||
float sum = 0.0f;
|
||||
const __half *wr = w + row * in_dim;
|
||||
const float *xr = x + tok * in_dim;
|
||||
for (uint64_t i = threadIdx.x; i < in_dim; i += blockDim.x) {
|
||||
sum += __half2float(wr[i]) * xr[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();
|
||||
}
|
||||
if (threadIdx.x == 0) out[tok * out_dim + row] = partial[0];
|
||||
}
|
||||
|
||||
__global__ static void matmul_f16_ordered_chunks_kernel(
|
||||
float *out,
|
||||
const __half *w,
|
||||
const float *x,
|
||||
uint64_t in_dim,
|
||||
uint64_t out_dim,
|
||||
uint64_t n_tok) {
|
||||
uint64_t row = (uint64_t)blockIdx.x;
|
||||
uint64_t tok = (uint64_t)blockIdx.y;
|
||||
if (row >= out_dim || tok >= n_tok) return;
|
||||
|
||||
__shared__ float partial[32];
|
||||
const uint32_t tid = threadIdx.x;
|
||||
float sum = 0.0f;
|
||||
const uint64_t chunk = (in_dim + 31u) / 32u;
|
||||
const uint64_t k0 = (uint64_t)tid * chunk;
|
||||
uint64_t k1 = k0 + chunk;
|
||||
if (k1 > in_dim) k1 = in_dim;
|
||||
const __half *wr = w + row * in_dim;
|
||||
const float *xr = x + tok * in_dim;
|
||||
for (uint64_t i = k0; i < k1; i++) {
|
||||
sum += __half2float(wr[i]) * xr[i];
|
||||
}
|
||||
partial[tid] = sum;
|
||||
__syncthreads();
|
||||
if (tid == 0) {
|
||||
float total = 0.0f;
|
||||
for (uint32_t i = 0; i < 32u; i++) total += partial[i];
|
||||
out[tok * out_dim + row] = total;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ static void matmul_f16_f32_sharedx_warp_rows_w32_kernel(
|
||||
float *out,
|
||||
const __half *w,
|
||||
const float *x,
|
||||
uint32_t in_dim,
|
||||
uint64_t out_dim) {
|
||||
extern __shared__ float shx[];
|
||||
const uint32_t tid = threadIdx.x;
|
||||
const uint32_t lane = tid & 31u;
|
||||
const uint32_t wave = tid >> 5u;
|
||||
const uint32_t rows_per_block = blockDim.x >> 5u;
|
||||
for (uint32_t i = tid; i < in_dim; i += blockDim.x) shx[i] = x[i];
|
||||
__syncthreads();
|
||||
|
||||
const uint64_t row = (uint64_t)blockIdx.x * rows_per_block + wave;
|
||||
if (row >= out_dim) return;
|
||||
const __half *wr = w + row * (uint64_t)in_dim;
|
||||
float acc = 0.0f;
|
||||
uint32_t i = lane;
|
||||
for (; i + 224u < in_dim; i += 256u) {
|
||||
acc += __half2float(wr[i]) * shx[i];
|
||||
acc += __half2float(wr[i + 32u]) * shx[i + 32u];
|
||||
acc += __half2float(wr[i + 64u]) * shx[i + 64u];
|
||||
acc += __half2float(wr[i + 96u]) * shx[i + 96u];
|
||||
acc += __half2float(wr[i + 128u]) * shx[i + 128u];
|
||||
acc += __half2float(wr[i + 160u]) * shx[i + 160u];
|
||||
acc += __half2float(wr[i + 192u]) * shx[i + 192u];
|
||||
acc += __half2float(wr[i + 224u]) * shx[i + 224u];
|
||||
}
|
||||
for (; i < in_dim; i += 32u) {
|
||||
acc += __half2float(wr[i]) * shx[i];
|
||||
}
|
||||
acc = warp_sum_f32(acc);
|
||||
if (lane == 0u) out[row] = acc;
|
||||
}
|
||||
|
||||
__global__ static void matmul_f16_pair_f32_sharedx_warp_rows_w32_kernel(
|
||||
float *out0,
|
||||
float *out1,
|
||||
const __half *w0,
|
||||
const __half *w1,
|
||||
const float *x,
|
||||
uint32_t in_dim,
|
||||
uint64_t out_dim) {
|
||||
extern __shared__ float shx[];
|
||||
const uint32_t tid = threadIdx.x;
|
||||
const uint32_t lane = tid & 31u;
|
||||
const uint32_t wave = tid >> 5u;
|
||||
const uint32_t rows_per_block = blockDim.x >> 5u;
|
||||
for (uint32_t i = tid; i < in_dim; i += blockDim.x) shx[i] = x[i];
|
||||
__syncthreads();
|
||||
|
||||
const uint64_t row = (uint64_t)blockIdx.x * rows_per_block + wave;
|
||||
if (row >= out_dim) return;
|
||||
const __half *wr0 = w0 + row * (uint64_t)in_dim;
|
||||
const __half *wr1 = w1 + row * (uint64_t)in_dim;
|
||||
float acc0 = 0.0f;
|
||||
float acc1 = 0.0f;
|
||||
uint32_t i = lane;
|
||||
for (; i + 224u < in_dim; i += 256u) {
|
||||
float xv = shx[i];
|
||||
acc0 += __half2float(wr0[i]) * xv;
|
||||
acc1 += __half2float(wr1[i]) * xv;
|
||||
xv = shx[i + 32u];
|
||||
acc0 += __half2float(wr0[i + 32u]) * xv;
|
||||
acc1 += __half2float(wr1[i + 32u]) * xv;
|
||||
xv = shx[i + 64u];
|
||||
acc0 += __half2float(wr0[i + 64u]) * xv;
|
||||
acc1 += __half2float(wr1[i + 64u]) * xv;
|
||||
xv = shx[i + 96u];
|
||||
acc0 += __half2float(wr0[i + 96u]) * xv;
|
||||
acc1 += __half2float(wr1[i + 96u]) * xv;
|
||||
xv = shx[i + 128u];
|
||||
acc0 += __half2float(wr0[i + 128u]) * xv;
|
||||
acc1 += __half2float(wr1[i + 128u]) * xv;
|
||||
xv = shx[i + 160u];
|
||||
acc0 += __half2float(wr0[i + 160u]) * xv;
|
||||
acc1 += __half2float(wr1[i + 160u]) * xv;
|
||||
xv = shx[i + 192u];
|
||||
acc0 += __half2float(wr0[i + 192u]) * xv;
|
||||
acc1 += __half2float(wr1[i + 192u]) * xv;
|
||||
xv = shx[i + 224u];
|
||||
acc0 += __half2float(wr0[i + 224u]) * xv;
|
||||
acc1 += __half2float(wr1[i + 224u]) * xv;
|
||||
}
|
||||
for (; i < in_dim; i += 32u) {
|
||||
const float xv = shx[i];
|
||||
acc0 += __half2float(wr0[i]) * xv;
|
||||
acc1 += __half2float(wr1[i]) * xv;
|
||||
}
|
||||
acc0 = warp_sum_f32(acc0);
|
||||
acc1 = warp_sum_f32(acc1);
|
||||
if (lane == 0u) {
|
||||
out0[row] = acc0;
|
||||
out1[row] = acc1;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ static void matmul_f16_pair_ordered_chunks_kernel(
|
||||
float *out0,
|
||||
float *out1,
|
||||
const __half *w0,
|
||||
const __half *w1,
|
||||
const float *x,
|
||||
uint64_t in_dim,
|
||||
uint64_t out0_dim,
|
||||
uint64_t out1_dim) {
|
||||
uint64_t row = (uint64_t)blockIdx.x;
|
||||
if (row >= out0_dim && row >= out1_dim) return;
|
||||
|
||||
__shared__ float partial0[32];
|
||||
__shared__ float partial1[32];
|
||||
const uint32_t tid = threadIdx.x;
|
||||
float sum0 = 0.0f;
|
||||
float sum1 = 0.0f;
|
||||
const uint64_t chunk = (in_dim + 31u) / 32u;
|
||||
const uint64_t k0 = (uint64_t)tid * chunk;
|
||||
uint64_t k1 = k0 + chunk;
|
||||
if (k1 > in_dim) k1 = in_dim;
|
||||
const __half *wr0 = row < out0_dim ? w0 + row * in_dim : w0;
|
||||
const __half *wr1 = row < out1_dim ? w1 + row * in_dim : w1;
|
||||
for (uint64_t i = k0; i < k1; i++) {
|
||||
const float xv = x[i];
|
||||
if (row < out0_dim) sum0 += __half2float(wr0[i]) * xv;
|
||||
if (row < out1_dim) sum1 += __half2float(wr1[i]) * xv;
|
||||
}
|
||||
partial0[tid] = sum0;
|
||||
partial1[tid] = sum1;
|
||||
__syncthreads();
|
||||
if (tid == 0) {
|
||||
float total0 = 0.0f;
|
||||
float total1 = 0.0f;
|
||||
for (uint32_t i = 0; i < 32u; i++) {
|
||||
total0 += partial0[i];
|
||||
total1 += partial1[i];
|
||||
}
|
||||
if (row < out0_dim) out0[row] = total0;
|
||||
if (row < out1_dim) out1[row] = total1;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ static void matmul_f32_kernel(
|
||||
float *out,
|
||||
const float *w,
|
||||
const float *x,
|
||||
uint64_t in_dim,
|
||||
uint64_t out_dim,
|
||||
uint64_t n_tok) {
|
||||
uint64_t row = (uint64_t)blockIdx.x;
|
||||
uint64_t tok = (uint64_t)blockIdx.y;
|
||||
if (row >= out_dim || tok >= n_tok) return;
|
||||
|
||||
float sum = 0.0f;
|
||||
const float *wr = w + row * in_dim;
|
||||
const float *xr = x + tok * in_dim;
|
||||
for (uint64_t i = threadIdx.x; i < in_dim; i += blockDim.x) {
|
||||
sum += wr[i] * xr[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();
|
||||
}
|
||||
if (threadIdx.x == 0) out[tok * out_dim + row] = partial[0];
|
||||
}
|
||||
|
||||
__global__ static void repeat_hc_kernel(float *out, const float *row, uint32_t n_embd, uint32_t n_hc) {
|
||||
uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint64_t n = (uint64_t)n_embd * n_hc;
|
||||
if (i >= n) return;
|
||||
out[i] = row[i % n_embd];
|
||||
}
|
||||
|
||||
__global__ static void f32_to_f16_kernel(__half *out, const float *x, uint64_t n) {
|
||||
uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i < n) out[i] = __float2half(x[i]);
|
||||
}
|
||||
|
||||
__device__ static float warp_sum_f32(float v) {
|
||||
for (int offset = 16; offset > 0; offset >>= 1) {
|
||||
#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
|
||||
v += __shfl_down(v, offset, 32);
|
||||
#else
|
||||
v += __shfl_down_sync(FULL_WARP_MASK, v, offset, 32);
|
||||
#endif
|
||||
}
|
||||
return v;
|
||||
}
|
||||
|
||||
__device__ static float warp_max_f32(float v) {
|
||||
for (int offset = 16; offset > 0; offset >>= 1) {
|
||||
#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
|
||||
v = fmaxf(v, __shfl_down(v, offset, 32));
|
||||
#else
|
||||
v = fmaxf(v, __shfl_down_sync(FULL_WARP_MASK, v, offset, 32));
|
||||
#endif
|
||||
}
|
||||
return v;
|
||||
}
|
||||
|
||||
__device__ static uint16_t f32_to_f16_bits_hip_round(float f) {
|
||||
union { float f; uint32_t u; } v;
|
||||
v.f = f;
|
||||
uint32_t sign = (v.u >> 16) & 0x8000u;
|
||||
int32_t exp = (int32_t)((v.u >> 23) & 0xffu) - 127 + 15;
|
||||
uint32_t mant = v.u & 0x7fffffu;
|
||||
if (exp <= 0) {
|
||||
if (exp < -10) return (uint16_t)sign;
|
||||
mant |= 0x800000u;
|
||||
uint32_t shift = (uint32_t)(14 - exp);
|
||||
uint32_t half_mant = mant >> shift;
|
||||
if ((mant >> (shift - 1)) & 1u) half_mant++;
|
||||
return (uint16_t)(sign | half_mant);
|
||||
}
|
||||
if (exp >= 31) return (uint16_t)(sign | 0x7c00u);
|
||||
uint32_t half = sign | ((uint32_t)exp << 10) | (mant >> 13);
|
||||
if (mant & 0x1000u) half++;
|
||||
return (uint16_t)half;
|
||||
}
|
||||
|
||||
__device__ static float f16_bits_to_f32(uint16_t bits) {
|
||||
return __half2float(__ushort_as_half((unsigned short)bits));
|
||||
}
|
||||
|
||||
__device__ static float dot4_f32(float4 a, float4 b) {
|
||||
return a.x * b.x + a.y * b.y + a.z * b.z + a.w * b.w;
|
||||
}
|
||||
|
||||
@@ -0,0 +1,557 @@
|
||||
__global__ static void compressor_store_kernel(
|
||||
const float *kv,
|
||||
const float *sc,
|
||||
float *state_kv,
|
||||
float *state_score,
|
||||
const void *model_map,
|
||||
uint64_t ape_offset,
|
||||
uint32_t ape_type,
|
||||
uint32_t head_dim,
|
||||
uint32_t ratio,
|
||||
uint32_t pos0,
|
||||
uint32_t n_tokens) {
|
||||
uint32_t coff = ratio == 4u ? 2u : 1u;
|
||||
uint32_t width = coff * head_dim;
|
||||
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint64_t n = (uint64_t)n_tokens * width;
|
||||
if (gid >= n) return;
|
||||
uint32_t t = gid / width;
|
||||
uint32_t j = gid - (uint64_t)t * width;
|
||||
uint32_t pos_mod = (pos0 + t) % ratio;
|
||||
uint32_t dst_row = ratio == 4u ? ratio + pos_mod : pos_mod;
|
||||
state_kv[(uint64_t)dst_row * width + j] = kv[(uint64_t)t * width + j];
|
||||
state_score[(uint64_t)dst_row * width + j] =
|
||||
sc[(uint64_t)t * width + j] + model_ape_value_dev(model_map, ape_offset, ape_type, width, pos_mod, j);
|
||||
}
|
||||
|
||||
__global__ static void compressor_set_rows_kernel(
|
||||
float *state_kv,
|
||||
float *state_score,
|
||||
const float *kv,
|
||||
const float *sc,
|
||||
const void *model_map,
|
||||
uint64_t ape_offset,
|
||||
uint32_t ape_type,
|
||||
uint32_t width,
|
||||
uint32_t ratio,
|
||||
uint32_t pos0,
|
||||
uint32_t src0,
|
||||
uint32_t dst0,
|
||||
uint32_t rows) {
|
||||
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint64_t n = (uint64_t)rows * width;
|
||||
if (gid >= n) return;
|
||||
uint32_t r = gid / width;
|
||||
uint32_t j = gid - (uint64_t)r * width;
|
||||
uint32_t src = src0 + r;
|
||||
uint32_t dst = dst0 + r;
|
||||
uint32_t phase = (pos0 + src) % ratio;
|
||||
state_kv[(uint64_t)dst * width + j] = kv[(uint64_t)src * width + j];
|
||||
state_score[(uint64_t)dst * width + j] =
|
||||
sc[(uint64_t)src * width + j] + model_ape_value_dev(model_map, ape_offset, ape_type, width, phase, j);
|
||||
}
|
||||
|
||||
__global__ static void compressor_prefill_pool_kernel(
|
||||
float *comp,
|
||||
const float *kv,
|
||||
const float *sc,
|
||||
const float *state_kv,
|
||||
const float *state_score,
|
||||
const void *model_map,
|
||||
uint64_t ape_offset,
|
||||
uint32_t ape_type,
|
||||
uint32_t head_dim,
|
||||
uint32_t ratio,
|
||||
uint32_t pos0,
|
||||
uint32_t n_comp,
|
||||
uint32_t replay) {
|
||||
uint32_t d = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint32_t c = blockIdx.y;
|
||||
if (d >= head_dim || c >= n_comp) return;
|
||||
uint32_t coff = ratio == 4u ? 2u : 1u;
|
||||
uint32_t width = coff * head_dim;
|
||||
float vals[DS4_ROCM_COMPRESSOR_MAX_RATIO];
|
||||
float scores[DS4_ROCM_COMPRESSOR_MAX_RATIO];
|
||||
float max_s = -INFINITY;
|
||||
uint32_t n_cand = 0;
|
||||
if (ratio == 4u) {
|
||||
if (replay && c == 0) {
|
||||
for (uint32_t r = 0; r < 4; r++) {
|
||||
vals[n_cand] = state_kv[(uint64_t)r * width + d];
|
||||
scores[n_cand] = state_score[(uint64_t)r * width + d];
|
||||
max_s = fmaxf(max_s, scores[n_cand++]);
|
||||
}
|
||||
} else if (c > 0) {
|
||||
uint32_t base = (c - 1u) * ratio;
|
||||
for (uint32_t r = 0; r < 4; r++) {
|
||||
uint32_t t = base + r;
|
||||
float ape = model_ape_value_dev(model_map, ape_offset, ape_type, width, (pos0 + t) % ratio, d);
|
||||
vals[n_cand] = kv[(uint64_t)t * width + d];
|
||||
scores[n_cand] = sc[(uint64_t)t * width + d] + ape;
|
||||
max_s = fmaxf(max_s, scores[n_cand++]);
|
||||
}
|
||||
}
|
||||
uint32_t base = c * ratio;
|
||||
for (uint32_t r = 0; r < 4; r++) {
|
||||
uint32_t t = base + r;
|
||||
float ape = model_ape_value_dev(model_map, ape_offset, ape_type, width, (pos0 + t) % ratio, head_dim + d);
|
||||
vals[n_cand] = kv[(uint64_t)t * width + head_dim + d];
|
||||
scores[n_cand] = sc[(uint64_t)t * width + head_dim + d] + ape;
|
||||
max_s = fmaxf(max_s, scores[n_cand++]);
|
||||
}
|
||||
} else {
|
||||
uint32_t base = c * ratio;
|
||||
for (uint32_t r = 0; r < ratio; r++) {
|
||||
uint32_t t = base + r;
|
||||
float ape = model_ape_value_dev(model_map, ape_offset, ape_type, width, (pos0 + t) % ratio, d);
|
||||
vals[n_cand] = kv[(uint64_t)t * width + d];
|
||||
scores[n_cand] = sc[(uint64_t)t * width + d] + ape;
|
||||
max_s = fmaxf(max_s, scores[n_cand++]);
|
||||
}
|
||||
}
|
||||
float den = 0.0f, acc = 0.0f;
|
||||
for (uint32_t i = 0; i < n_cand; i++) {
|
||||
float w = expf(scores[i] - max_s);
|
||||
den += w;
|
||||
acc += vals[i] * w;
|
||||
}
|
||||
comp[(uint64_t)c * head_dim + d] = den != 0.0f ? acc / den : 0.0f;
|
||||
}
|
||||
|
||||
__global__ static void compressor_update_pool_kernel(
|
||||
float *row,
|
||||
const float *state_kv,
|
||||
const float *state_score,
|
||||
uint32_t head_dim,
|
||||
uint32_t ratio) {
|
||||
uint32_t d = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (d >= head_dim) return;
|
||||
uint32_t coff = ratio == 4u ? 2u : 1u;
|
||||
uint32_t width = coff * head_dim;
|
||||
float vals[DS4_ROCM_COMPRESSOR_MAX_RATIO];
|
||||
float scores[DS4_ROCM_COMPRESSOR_MAX_RATIO];
|
||||
float max_s = -INFINITY;
|
||||
uint32_t n_cand = 0;
|
||||
if (ratio == 4u) {
|
||||
for (uint32_t r = 0; r < 4; r++) {
|
||||
vals[n_cand] = state_kv[(uint64_t)r * width + d];
|
||||
scores[n_cand] = state_score[(uint64_t)r * width + d];
|
||||
max_s = fmaxf(max_s, scores[n_cand++]);
|
||||
}
|
||||
for (uint32_t r = 0; r < 4; r++) {
|
||||
vals[n_cand] = state_kv[(uint64_t)(ratio + r) * width + head_dim + d];
|
||||
scores[n_cand] = state_score[(uint64_t)(ratio + r) * width + head_dim + d];
|
||||
max_s = fmaxf(max_s, scores[n_cand++]);
|
||||
}
|
||||
} else {
|
||||
for (uint32_t r = 0; r < ratio; r++) {
|
||||
vals[n_cand] = state_kv[(uint64_t)r * width + d];
|
||||
scores[n_cand] = state_score[(uint64_t)r * width + d];
|
||||
max_s = fmaxf(max_s, scores[n_cand++]);
|
||||
}
|
||||
}
|
||||
float den = 0.0f, acc = 0.0f;
|
||||
for (uint32_t i = 0; i < n_cand; i++) {
|
||||
float w = expf(scores[i] - max_s);
|
||||
den += w;
|
||||
acc += vals[i] * w;
|
||||
}
|
||||
row[d] = den != 0.0f ? acc / den : 0.0f;
|
||||
}
|
||||
|
||||
__global__ static void compressor_shift_ratio4_kernel(float *state_kv, float *state_score, uint32_t width) {
|
||||
uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint64_t half = 4ull * width;
|
||||
if (i >= half) return;
|
||||
float v = state_kv[half + i];
|
||||
float s = state_score[half + i];
|
||||
state_kv[i] = v;
|
||||
state_score[i] = s;
|
||||
state_kv[half + i] = v;
|
||||
state_score[half + i] = s;
|
||||
}
|
||||
|
||||
static uint64_t cuda_tensor_2d_bytes(uint32_t type, uint64_t width, uint64_t rows) {
|
||||
if (type == 0u) return width * rows * sizeof(float);
|
||||
if (type == 1u) return width * rows * sizeof(uint16_t);
|
||||
if (type == 8u) return rows * (((width + 31u) / 32u) * 34u);
|
||||
return 0;
|
||||
}
|
||||
|
||||
static int cuda_ape_type_supported(uint32_t type) {
|
||||
return type == 0u || type == 1u || type == 8u;
|
||||
}
|
||||
|
||||
static int cuda_compressor_shape_supported(uint32_t head_dim, uint32_t ratio) {
|
||||
if (head_dim == 0u || ratio == 0u || ratio > DS4_ROCM_COMPRESSOR_MAX_RATIO) return 0;
|
||||
const uint32_t coff = ratio == 4u ? 2u : 1u;
|
||||
return head_dim <= UINT32_MAX / coff;
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_compressor_store_batch_tensor(
|
||||
const ds4_gpu_tensor *kv,
|
||||
const ds4_gpu_tensor *sc,
|
||||
ds4_gpu_tensor *state_kv,
|
||||
ds4_gpu_tensor *state_score,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t ape_offset,
|
||||
uint32_t ape_type,
|
||||
uint32_t head_dim,
|
||||
uint32_t ratio,
|
||||
uint32_t pos0,
|
||||
uint32_t n_tokens) {
|
||||
if (!kv || !sc || !state_kv || !state_score || !model_map ||
|
||||
!cuda_compressor_shape_supported(head_dim, ratio) || n_tokens == 0 ||
|
||||
!cuda_ape_type_supported(ape_type)) {
|
||||
return 0;
|
||||
}
|
||||
const uint32_t coff = ratio == 4u ? 2u : 1u;
|
||||
const uint32_t width = coff * head_dim;
|
||||
const uint32_t state_rows = coff * ratio;
|
||||
uint64_t kv_bytes = 0, state_bytes = 0;
|
||||
const uint64_t ape_bytes = cuda_tensor_2d_bytes(ape_type, width, ratio);
|
||||
if (!cuda_u64_mul3_checked(n_tokens, width, sizeof(float), &kv_bytes) ||
|
||||
!cuda_u64_mul3_checked(state_rows, width, sizeof(float), &state_bytes) ||
|
||||
!cuda_model_range_fits(model_size, ape_offset, ape_bytes) ||
|
||||
!cuda_tensor_has_bytes(kv, kv_bytes) || !cuda_tensor_has_bytes(sc, kv_bytes) ||
|
||||
!cuda_tensor_has_bytes(state_kv, state_bytes) || !cuda_tensor_has_bytes(state_score, state_bytes)) {
|
||||
return 0;
|
||||
}
|
||||
const char *ape = cuda_model_range_ptr(model_map, ape_offset, ape_bytes, "compressor_ape");
|
||||
if (!ape) return 0;
|
||||
uint64_t n = (uint64_t)n_tokens * width;
|
||||
compressor_store_kernel<<<(n + 255) / 256, 256>>>(
|
||||
(const float *)kv->ptr,
|
||||
(const float *)sc->ptr,
|
||||
(float *)state_kv->ptr,
|
||||
(float *)state_score->ptr,
|
||||
ape,
|
||||
0,
|
||||
ape_type,
|
||||
head_dim,
|
||||
ratio,
|
||||
pos0,
|
||||
n_tokens);
|
||||
return cuda_ok(cudaGetLastError(), "compressor store launch");
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_compressor_update_tensor(
|
||||
const ds4_gpu_tensor *kv_cur,
|
||||
const ds4_gpu_tensor *sc_cur,
|
||||
ds4_gpu_tensor *state_kv,
|
||||
ds4_gpu_tensor *state_score,
|
||||
ds4_gpu_tensor *comp_cache,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t ape_offset,
|
||||
uint32_t ape_type,
|
||||
uint64_t norm_offset,
|
||||
uint32_t norm_type,
|
||||
uint32_t head_dim,
|
||||
uint32_t ratio,
|
||||
uint32_t pos,
|
||||
uint32_t comp_row,
|
||||
uint32_t n_rot,
|
||||
uint32_t n_ctx_orig,
|
||||
float freq_base,
|
||||
float freq_scale,
|
||||
float ext_factor,
|
||||
float attn_factor,
|
||||
float beta_fast,
|
||||
float beta_slow,
|
||||
float rms_eps) {
|
||||
if (!kv_cur || !sc_cur || !state_kv || !state_score || !comp_cache ||
|
||||
!model_map || !cuda_compressor_shape_supported(head_dim, ratio) ||
|
||||
n_rot > head_dim || (n_rot & 1u) != 0 ||
|
||||
!cuda_ape_type_supported(ape_type) || norm_type != 0u) {
|
||||
return 0;
|
||||
}
|
||||
const uint32_t coff = ratio == 4u ? 2u : 1u;
|
||||
const uint32_t width = coff * head_dim;
|
||||
const uint32_t state_rows = coff * ratio;
|
||||
const uint32_t emit = ((pos + 1u) % ratio) == 0u ? 1u : 0u;
|
||||
uint64_t kv_bytes = 0, state_bytes = 0, comp_bytes = 0, norm_bytes = 0;
|
||||
const uint64_t ape_bytes = cuda_tensor_2d_bytes(ape_type, width, ratio);
|
||||
if (!cuda_u64_mul_checked(width, sizeof(float), &kv_bytes) ||
|
||||
!cuda_u64_mul3_checked(state_rows, width, sizeof(float), &state_bytes) ||
|
||||
!cuda_u64_mul3_checked((uint64_t)comp_row + (emit ? 1u : 0u), head_dim, sizeof(float), &comp_bytes) ||
|
||||
!cuda_u64_mul_checked(head_dim, sizeof(float), &norm_bytes) ||
|
||||
!cuda_model_range_fits(model_size, ape_offset, ape_bytes) ||
|
||||
!cuda_model_range_fits(model_size, norm_offset, norm_bytes) ||
|
||||
!cuda_tensor_has_bytes(kv_cur, kv_bytes) || !cuda_tensor_has_bytes(sc_cur, kv_bytes) ||
|
||||
!cuda_tensor_has_bytes(state_kv, state_bytes) || !cuda_tensor_has_bytes(state_score, state_bytes) ||
|
||||
(emit && !cuda_tensor_has_bytes(comp_cache, comp_bytes))) {
|
||||
return 0;
|
||||
}
|
||||
if (!ds4_gpu_compressor_store_batch_tensor(kv_cur, sc_cur, state_kv, state_score,
|
||||
model_map, model_size, ape_offset, ape_type,
|
||||
head_dim, ratio, pos, 1)) {
|
||||
return 0;
|
||||
}
|
||||
if (!emit) return 1;
|
||||
ds4_gpu_tensor *comp_row_view = ds4_gpu_tensor_view(
|
||||
comp_cache,
|
||||
(uint64_t)comp_row * head_dim * sizeof(float),
|
||||
(uint64_t)head_dim * sizeof(float));
|
||||
if (!comp_row_view) return 0;
|
||||
compressor_update_pool_kernel<<<(head_dim + 255) / 256, 256>>>(
|
||||
(float *)comp_row_view->ptr,
|
||||
(const float *)state_kv->ptr,
|
||||
(const float *)state_score->ptr,
|
||||
head_dim,
|
||||
ratio);
|
||||
int ok = cuda_ok(cudaGetLastError(), "compressor update pool launch");
|
||||
if (ok) ok = ds4_gpu_rms_norm_weight_rows_tensor(comp_row_view, comp_row_view,
|
||||
model_map, model_size, norm_offset,
|
||||
head_dim, 1, rms_eps);
|
||||
if (ok) ok = ds4_gpu_rope_tail_tensor(comp_row_view, 1, 1, head_dim, n_rot,
|
||||
pos + 1u - ratio, n_ctx_orig, false,
|
||||
freq_base, freq_scale, ext_factor, attn_factor,
|
||||
beta_fast, beta_slow);
|
||||
ds4_gpu_tensor_free(comp_row_view);
|
||||
if (ok && ratio == 4u) {
|
||||
uint64_t half = 4ull * width;
|
||||
compressor_shift_ratio4_kernel<<<(half + 255) / 256, 256>>>(
|
||||
(float *)state_kv->ptr, (float *)state_score->ptr, width);
|
||||
ok = cuda_ok(cudaGetLastError(), "compressor ratio4 shift launch");
|
||||
}
|
||||
return ok;
|
||||
}
|
||||
extern "C" int ds4_gpu_compressor_prefill_tensor(
|
||||
ds4_gpu_tensor *comp_cache,
|
||||
ds4_gpu_tensor *state_kv,
|
||||
ds4_gpu_tensor *state_score,
|
||||
const ds4_gpu_tensor *kv,
|
||||
const ds4_gpu_tensor *sc,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t ape_offset,
|
||||
uint32_t ape_type,
|
||||
uint64_t norm_offset,
|
||||
uint32_t norm_type,
|
||||
uint32_t head_dim,
|
||||
uint32_t ratio,
|
||||
uint32_t pos0,
|
||||
uint32_t n_tokens,
|
||||
uint32_t n_rot,
|
||||
uint32_t n_ctx_orig,
|
||||
bool quantize_fp8,
|
||||
float freq_base,
|
||||
float freq_scale,
|
||||
float ext_factor,
|
||||
float attn_factor,
|
||||
float beta_fast,
|
||||
float beta_slow,
|
||||
float rms_eps) {
|
||||
if (!comp_cache || !state_kv || !state_score || !kv || !sc || !model_map ||
|
||||
!cuda_compressor_shape_supported(head_dim, ratio) || n_tokens == 0 ||
|
||||
n_rot > head_dim || (n_rot & 1u) != 0 ||
|
||||
!cuda_ape_type_supported(ape_type) || norm_type != 0u) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
const uint32_t coff = ratio == 4u ? 2u : 1u;
|
||||
const uint32_t width = coff * head_dim;
|
||||
const uint32_t state_rows = coff * ratio;
|
||||
const uint32_t n_comp = n_tokens / ratio;
|
||||
const uint32_t cutoff = n_comp * ratio;
|
||||
const uint32_t rem = n_tokens - cutoff;
|
||||
uint64_t kv_bytes = 0, state_bytes = 0, comp_bytes = 0, norm_bytes = 0;
|
||||
const uint64_t ape_bytes = cuda_tensor_2d_bytes(ape_type, width, ratio);
|
||||
|
||||
if (!cuda_u64_mul3_checked(n_tokens, width, sizeof(float), &kv_bytes) ||
|
||||
!cuda_u64_mul3_checked(state_rows, width, sizeof(float), &state_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_comp, head_dim, sizeof(float), &comp_bytes) ||
|
||||
!cuda_u64_mul_checked(head_dim, sizeof(float), &norm_bytes) ||
|
||||
!cuda_model_range_fits(model_size, ape_offset, ape_bytes) ||
|
||||
!cuda_model_range_fits(model_size, norm_offset, norm_bytes) ||
|
||||
!cuda_tensor_has_bytes(kv, kv_bytes) || !cuda_tensor_has_bytes(sc, kv_bytes) ||
|
||||
!cuda_tensor_has_bytes(state_kv, state_bytes) || !cuda_tensor_has_bytes(state_score, state_bytes) ||
|
||||
(n_comp && !cuda_tensor_has_bytes(comp_cache, comp_bytes))) {
|
||||
return 0;
|
||||
}
|
||||
const char *ape = cuda_model_range_ptr(model_map, ape_offset, ape_bytes, "compressor_ape");
|
||||
if (!ape) return 0;
|
||||
|
||||
uint64_t state_n = (uint64_t)state_rows * width;
|
||||
if (!cuda_ok(cudaMemsetAsync(state_kv->ptr, 0, (size_t)(state_n * sizeof(float))),
|
||||
"compressor state kv zero")) return 0;
|
||||
fill_f32_kernel<<<(state_n + 255) / 256, 256>>>((float *)state_score->ptr, state_n, -INFINITY);
|
||||
if (!cuda_ok(cudaGetLastError(), "compressor state score fill launch")) return 0;
|
||||
|
||||
if (ratio == 4u) {
|
||||
if (cutoff >= ratio) {
|
||||
uint32_t prev_start = cutoff - ratio;
|
||||
uint64_t n = (uint64_t)ratio * width;
|
||||
compressor_set_rows_kernel<<<(n + 255) / 256, 256>>>(
|
||||
(float *)state_kv->ptr, (float *)state_score->ptr,
|
||||
(const float *)kv->ptr, (const float *)sc->ptr,
|
||||
ape, 0, ape_type, width, ratio, pos0,
|
||||
prev_start, 0, ratio);
|
||||
if (!cuda_ok(cudaGetLastError(), "compressor prefill prev state launch")) return 0;
|
||||
}
|
||||
if (rem != 0) {
|
||||
uint64_t n = (uint64_t)rem * width;
|
||||
compressor_set_rows_kernel<<<(n + 255) / 256, 256>>>(
|
||||
(float *)state_kv->ptr, (float *)state_score->ptr,
|
||||
(const float *)kv->ptr, (const float *)sc->ptr,
|
||||
ape, 0, ape_type, width, ratio, pos0,
|
||||
cutoff, ratio, rem);
|
||||
if (!cuda_ok(cudaGetLastError(), "compressor prefill rem state launch")) return 0;
|
||||
}
|
||||
} else if (rem != 0) {
|
||||
uint64_t n = (uint64_t)rem * width;
|
||||
compressor_set_rows_kernel<<<(n + 255) / 256, 256>>>(
|
||||
(float *)state_kv->ptr, (float *)state_score->ptr,
|
||||
(const float *)kv->ptr, (const float *)sc->ptr,
|
||||
ape, 0, ape_type, width, ratio, pos0,
|
||||
cutoff, 0, rem);
|
||||
if (!cuda_ok(cudaGetLastError(), "compressor prefill rem state launch")) return 0;
|
||||
}
|
||||
if (n_comp != 0) {
|
||||
dim3 grid((head_dim + 255) / 256, n_comp, 1);
|
||||
compressor_prefill_pool_kernel<<<grid, 256>>>(
|
||||
(float *)comp_cache->ptr,
|
||||
(const float *)kv->ptr,
|
||||
(const float *)sc->ptr,
|
||||
(const float *)state_kv->ptr,
|
||||
(const float *)state_score->ptr,
|
||||
ape, 0, ape_type, head_dim, ratio, pos0, n_comp, 0);
|
||||
if (!cuda_ok(cudaGetLastError(), "compressor prefill pool launch")) return 0;
|
||||
if (!ds4_gpu_rms_norm_weight_rows_tensor(comp_cache, comp_cache,
|
||||
model_map, model_size, norm_offset,
|
||||
head_dim, n_comp, rms_eps)) return 0;
|
||||
if (n_rot != 0 && !cuda_rope_tail_stride_tensor(comp_cache, n_comp, 1, head_dim,
|
||||
n_rot, pos0, ratio, n_ctx_orig, false,
|
||||
freq_base, freq_scale, ext_factor,
|
||||
attn_factor, beta_fast, beta_slow)) return 0;
|
||||
if (quantize_fp8 && !ds4_gpu_dsv4_fp8_kv_quantize_tensor(comp_cache, n_comp, head_dim, n_rot)) return 0;
|
||||
}
|
||||
return 1;
|
||||
}
|
||||
extern "C" int ds4_gpu_compressor_prefill_ratio4_replay_tensor(
|
||||
ds4_gpu_tensor *comp_cache,
|
||||
ds4_gpu_tensor *state_kv,
|
||||
ds4_gpu_tensor *state_score,
|
||||
const ds4_gpu_tensor *kv,
|
||||
const ds4_gpu_tensor *sc,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t ape_offset,
|
||||
uint32_t ape_type,
|
||||
uint64_t norm_offset,
|
||||
uint32_t norm_type,
|
||||
uint32_t head_dim,
|
||||
uint32_t pos0,
|
||||
uint32_t n_tokens,
|
||||
uint32_t n_rot,
|
||||
uint32_t n_ctx_orig,
|
||||
bool quantize_fp8,
|
||||
float freq_base,
|
||||
float freq_scale,
|
||||
float ext_factor,
|
||||
float attn_factor,
|
||||
float beta_fast,
|
||||
float beta_slow,
|
||||
float rms_eps) {
|
||||
if (!comp_cache || !state_kv || !state_score || !kv || !sc || !model_map ||
|
||||
head_dim == 0 || n_tokens == 0 || (n_tokens & 3u) != 0 || (pos0 & 3u) != 0 ||
|
||||
n_rot > head_dim || (n_rot & 1u) != 0 ||
|
||||
!cuda_ape_type_supported(ape_type) || norm_type != 0u) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
const uint32_t ratio = 4u;
|
||||
const uint32_t width = 2u * head_dim;
|
||||
const uint32_t state_rows = 8u;
|
||||
const uint32_t n_comp = n_tokens / ratio;
|
||||
uint64_t kv_bytes = 0, state_bytes = 0, comp_bytes = 0, norm_bytes = 0;
|
||||
const uint64_t ape_bytes = cuda_tensor_2d_bytes(ape_type, width, ratio);
|
||||
if (!cuda_u64_mul3_checked(n_tokens, width, sizeof(float), &kv_bytes) ||
|
||||
!cuda_u64_mul3_checked(state_rows, width, sizeof(float), &state_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_comp, head_dim, sizeof(float), &comp_bytes) ||
|
||||
!cuda_u64_mul_checked(head_dim, sizeof(float), &norm_bytes) ||
|
||||
!cuda_model_range_fits(model_size, ape_offset, ape_bytes) ||
|
||||
!cuda_model_range_fits(model_size, norm_offset, norm_bytes) ||
|
||||
!cuda_tensor_has_bytes(kv, kv_bytes) || !cuda_tensor_has_bytes(sc, kv_bytes) ||
|
||||
!cuda_tensor_has_bytes(state_kv, state_bytes) || !cuda_tensor_has_bytes(state_score, state_bytes) ||
|
||||
!cuda_tensor_has_bytes(comp_cache, comp_bytes)) {
|
||||
return 0;
|
||||
}
|
||||
const char *ape = cuda_model_range_ptr(model_map, ape_offset, ape_bytes, "compressor_ape");
|
||||
if (!ape) return 0;
|
||||
dim3 grid((head_dim + 255) / 256, n_comp, 1);
|
||||
compressor_prefill_pool_kernel<<<grid, 256>>>(
|
||||
(float *)comp_cache->ptr,
|
||||
(const float *)kv->ptr,
|
||||
(const float *)sc->ptr,
|
||||
(const float *)state_kv->ptr,
|
||||
(const float *)state_score->ptr,
|
||||
ape, 0, ape_type, head_dim, ratio, pos0, n_comp, 1);
|
||||
if (!cuda_ok(cudaGetLastError(), "compressor replay pool launch")) return 0;
|
||||
if (!ds4_gpu_rms_norm_weight_rows_tensor(comp_cache, comp_cache,
|
||||
model_map, model_size, norm_offset,
|
||||
head_dim, n_comp, rms_eps)) return 0;
|
||||
if (n_rot != 0 && !cuda_rope_tail_stride_tensor(comp_cache, n_comp, 1, head_dim,
|
||||
n_rot, pos0, ratio, n_ctx_orig, false,
|
||||
freq_base, freq_scale, ext_factor,
|
||||
attn_factor, beta_fast, beta_slow)) return 0;
|
||||
if (quantize_fp8 && !ds4_gpu_dsv4_fp8_kv_quantize_tensor(comp_cache, n_comp, head_dim, n_rot)) return 0;
|
||||
|
||||
uint64_t state_n = (uint64_t)state_rows * width;
|
||||
if (!cuda_ok(cudaMemsetAsync(state_kv->ptr, 0, (size_t)(state_n * sizeof(float))),
|
||||
"compressor replay state kv zero")) return 0;
|
||||
fill_f32_kernel<<<(state_n + 255) / 256, 256>>>((float *)state_score->ptr, state_n, -INFINITY);
|
||||
if (!cuda_ok(cudaGetLastError(), "compressor replay state score fill launch")) return 0;
|
||||
uint32_t prev_start = n_tokens - ratio;
|
||||
uint64_t n = (uint64_t)ratio * width;
|
||||
compressor_set_rows_kernel<<<(n + 255) / 256, 256>>>(
|
||||
(float *)state_kv->ptr, (float *)state_score->ptr,
|
||||
(const float *)kv->ptr, (const float *)sc->ptr,
|
||||
ape, 0, ape_type, width, ratio, pos0,
|
||||
prev_start, 0, ratio);
|
||||
return cuda_ok(cudaGetLastError(), "compressor replay state launch");
|
||||
}
|
||||
extern "C" int ds4_gpu_compressor_prefill_state_ratio4_tensor(
|
||||
ds4_gpu_tensor *state_kv,
|
||||
ds4_gpu_tensor *state_score,
|
||||
const ds4_gpu_tensor *kv_tail,
|
||||
const ds4_gpu_tensor *sc_tail,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t ape_offset,
|
||||
uint32_t ape_type,
|
||||
uint32_t head_dim,
|
||||
uint32_t pos0) {
|
||||
if (!state_kv || !state_score || !kv_tail || !sc_tail || !model_map ||
|
||||
head_dim == 0 || !cuda_ape_type_supported(ape_type)) {
|
||||
return 0;
|
||||
}
|
||||
const uint32_t ratio = 4u;
|
||||
const uint32_t width = 2u * head_dim;
|
||||
const uint32_t state_rows = 8u;
|
||||
uint64_t tail_bytes = 0, state_bytes = 0;
|
||||
const uint64_t ape_bytes = cuda_tensor_2d_bytes(ape_type, width, ratio);
|
||||
if (!cuda_u64_mul3_checked(ratio, width, sizeof(float), &tail_bytes) ||
|
||||
!cuda_u64_mul3_checked(state_rows, width, sizeof(float), &state_bytes) ||
|
||||
!cuda_model_range_fits(model_size, ape_offset, ape_bytes) ||
|
||||
!cuda_tensor_has_bytes(kv_tail, tail_bytes) || !cuda_tensor_has_bytes(sc_tail, tail_bytes) ||
|
||||
!cuda_tensor_has_bytes(state_kv, state_bytes) || !cuda_tensor_has_bytes(state_score, state_bytes)) {
|
||||
return 0;
|
||||
}
|
||||
const char *ape = cuda_model_range_ptr(model_map, ape_offset, ape_bytes, "compressor_ape");
|
||||
if (!ape) return 0;
|
||||
uint64_t state_n = (uint64_t)state_rows * width;
|
||||
if (!cuda_ok(cudaMemsetAsync(state_kv->ptr, 0, (size_t)(state_n * sizeof(float))),
|
||||
"compressor state kv zero")) return 0;
|
||||
fill_f32_kernel<<<(state_n + 255) / 256, 256>>>((float *)state_score->ptr, state_n, -INFINITY);
|
||||
if (!cuda_ok(cudaGetLastError(), "compressor state score fill launch")) return 0;
|
||||
uint64_t n = (uint64_t)ratio * width;
|
||||
compressor_set_rows_kernel<<<(n + 255) / 256, 256>>>(
|
||||
(float *)state_kv->ptr, (float *)state_score->ptr,
|
||||
(const float *)kv_tail->ptr, (const float *)sc_tail->ptr,
|
||||
ape, 0, ape_type, width, ratio, pos0,
|
||||
0, 0, ratio);
|
||||
return cuda_ok(cudaGetLastError(), "compressor state set launch");
|
||||
}
|
||||
@@ -0,0 +1,346 @@
|
||||
extern "C" int ds4_gpu_signal_selected_readback_ready(uint64_t *event_value) {
|
||||
if (!event_value) return 0;
|
||||
*event_value = 0;
|
||||
if (!g_selected_readback_event) {
|
||||
cudaError_t err =
|
||||
cudaEventCreateWithFlags(&g_selected_readback_event,
|
||||
cudaEventDisableTiming);
|
||||
if (err != cudaSuccess) {
|
||||
fprintf(stderr,
|
||||
DS4_GPU_LOG_PREFIX "selected readback event creation failed: %s\n",
|
||||
cudaGetErrorString(err));
|
||||
(void)cudaGetLastError();
|
||||
return 0;
|
||||
}
|
||||
}
|
||||
cudaError_t err = cudaEventRecord(g_selected_readback_event, 0);
|
||||
if (err != cudaSuccess) {
|
||||
fprintf(stderr,
|
||||
DS4_GPU_LOG_PREFIX "selected readback event record failed: %s\n",
|
||||
cudaGetErrorString(err));
|
||||
(void)cudaGetLastError();
|
||||
return 0;
|
||||
}
|
||||
*event_value = ++g_selected_readback_event_value;
|
||||
return 1;
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_commit_and_wait_selected_readback(uint64_t event_value, const char *label) {
|
||||
if (event_value == 0 || !g_selected_readback_event) return 0;
|
||||
cudaError_t err = cudaEventSynchronize(g_selected_readback_event);
|
||||
if (err != cudaSuccess) {
|
||||
fprintf(stderr,
|
||||
DS4_GPU_LOG_PREFIX "selected readback wait failed for %s: %s\n",
|
||||
label ? label : "selected-id readback",
|
||||
cudaGetErrorString(err));
|
||||
(void)cudaGetLastError();
|
||||
return 0;
|
||||
}
|
||||
return 1;
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_wait_selected_readback_ready(uint64_t event_value, const char *label) {
|
||||
return ds4_gpu_commit_and_wait_selected_readback(event_value, label);
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_tensor_read_after_selected_event(
|
||||
const ds4_gpu_tensor *tensor,
|
||||
uint64_t offset,
|
||||
void *data,
|
||||
uint64_t bytes,
|
||||
uint64_t event_value,
|
||||
const char *label) {
|
||||
if (!tensor || !data || offset > tensor->bytes ||
|
||||
bytes > tensor->bytes - offset ||
|
||||
event_value == 0 ||
|
||||
!g_selected_readback_event) {
|
||||
return 0;
|
||||
}
|
||||
if (!g_selected_readback_stream) {
|
||||
cudaError_t err =
|
||||
cudaStreamCreateWithFlags(&g_selected_readback_stream,
|
||||
cudaStreamNonBlocking);
|
||||
if (err != cudaSuccess) {
|
||||
fprintf(stderr,
|
||||
DS4_GPU_LOG_PREFIX "selected readback stream creation failed: %s\n",
|
||||
cudaGetErrorString(err));
|
||||
(void)cudaGetLastError();
|
||||
return 0;
|
||||
}
|
||||
}
|
||||
#ifdef __HIP_PLATFORM_AMD__
|
||||
cudaError_t err = hipStreamWaitEvent(g_selected_readback_stream,
|
||||
g_selected_readback_event,
|
||||
0);
|
||||
#else
|
||||
cudaError_t err = cudaStreamWaitEvent(g_selected_readback_stream,
|
||||
g_selected_readback_event,
|
||||
0);
|
||||
#endif
|
||||
if (err != cudaSuccess) {
|
||||
fprintf(stderr,
|
||||
DS4_GPU_LOG_PREFIX "selected readback stream wait failed for %s: %s\n",
|
||||
label ? label : "selected-id readback",
|
||||
cudaGetErrorString(err));
|
||||
(void)cudaGetLastError();
|
||||
return 0;
|
||||
}
|
||||
err = cudaMemcpyAsync(data,
|
||||
(const char *)tensor->ptr + offset,
|
||||
(size_t)bytes,
|
||||
cudaMemcpyDeviceToHost,
|
||||
g_selected_readback_stream);
|
||||
if (err != cudaSuccess) {
|
||||
fprintf(stderr,
|
||||
DS4_GPU_LOG_PREFIX "selected readback copy failed for %s: %s\n",
|
||||
label ? label : "selected-id readback",
|
||||
cudaGetErrorString(err));
|
||||
(void)cudaGetLastError();
|
||||
return 0;
|
||||
}
|
||||
err = cudaStreamSynchronize(g_selected_readback_stream);
|
||||
if (err != cudaSuccess) {
|
||||
fprintf(stderr,
|
||||
DS4_GPU_LOG_PREFIX "selected readback sync failed for %s: %s\n",
|
||||
label ? label : "selected-id readback",
|
||||
cudaGetErrorString(err));
|
||||
(void)cudaGetLastError();
|
||||
return 0;
|
||||
}
|
||||
return 1;
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_set_model_fd_for_map(int fd, const void *model_map) {
|
||||
int ok = ds4_gpu_set_model_fd(fd);
|
||||
g_model_fd_host_base = model_map ? model_map : g_model_host_base;
|
||||
return ok;
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_tensor_copy_f32_to_f16(
|
||||
ds4_gpu_tensor *dst,
|
||||
uint64_t dst_offset,
|
||||
const ds4_gpu_tensor *src,
|
||||
uint64_t src_offset,
|
||||
uint64_t count) {
|
||||
if (!dst || !src || !dst->ptr || !src->ptr) return 0;
|
||||
if ((dst_offset % sizeof(__half)) != 0 || (src_offset % sizeof(float)) != 0) return 0;
|
||||
if (dst_offset > dst->bytes || src_offset > src->bytes) return 0;
|
||||
if (count > (UINT64_MAX / sizeof(__half)) || count > (UINT64_MAX / sizeof(float))) return 0;
|
||||
uint64_t dst_bytes = count * sizeof(__half);
|
||||
uint64_t src_bytes = count * sizeof(float);
|
||||
if (dst_bytes > dst->bytes - dst_offset || src_bytes > src->bytes - src_offset) return 0;
|
||||
if (count == 0) return 1;
|
||||
f32_to_f16_kernel<<<(count + 255u) / 256u, 256>>>(
|
||||
(__half *)((char *)dst->ptr + dst_offset),
|
||||
(const float *)((const char *)src->ptr + src_offset),
|
||||
count);
|
||||
return cuda_ok(cudaGetLastError(), "tensor copy f32 to f16 launch");
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_pro_q4_expert_table_auto_available(void) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_preload_q4_expert_tables(
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t gate_offset,
|
||||
uint64_t up_offset,
|
||||
uint64_t down_offset,
|
||||
uint64_t gate_expert_bytes,
|
||||
uint64_t down_expert_bytes,
|
||||
uint32_t n_total_expert) {
|
||||
(void)model_map;
|
||||
(void)model_size;
|
||||
(void)gate_offset;
|
||||
(void)up_offset;
|
||||
(void)down_offset;
|
||||
(void)gate_expert_bytes;
|
||||
(void)down_expert_bytes;
|
||||
(void)n_total_expert;
|
||||
return 0;
|
||||
}
|
||||
|
||||
extern "C" void ds4_gpu_set_ssd_streaming(bool enabled) {
|
||||
g_ssd_streaming_mode = enabled ? 1 : 0;
|
||||
cuda_model_range_release_all();
|
||||
cuda_q8_f16_cache_release_all();
|
||||
g_routed_moe_selected_override_n = 0;
|
||||
g_stream_selected_cache.loaded = 0;
|
||||
g_stream_batch_selected_cache.loaded = 0;
|
||||
}
|
||||
|
||||
extern "C" void ds4_gpu_set_streaming_expert_cache_budget(uint32_t experts) {
|
||||
g_stream_expert_cache_budget = experts;
|
||||
}
|
||||
|
||||
extern "C" void ds4_gpu_set_streaming_expert_cache_expert_bytes(uint64_t bytes) {
|
||||
(void)bytes;
|
||||
}
|
||||
|
||||
extern "C" uint64_t ds4_gpu_recommended_working_set_size(void) {
|
||||
size_t free_b = 0;
|
||||
size_t total_b = 0;
|
||||
if (cudaMemGetInfo(&free_b, &total_b) != cudaSuccess) {
|
||||
(void)cudaGetLastError();
|
||||
return 0;
|
||||
}
|
||||
(void)free_b;
|
||||
return (uint64_t)total_b;
|
||||
}
|
||||
|
||||
extern "C" uint32_t ds4_gpu_stream_expert_cache_configured_count(void) {
|
||||
return g_ssd_streaming_mode ? g_stream_expert_cache_budget : 0;
|
||||
}
|
||||
|
||||
extern "C" uint32_t ds4_gpu_stream_expert_cache_current_count(void) {
|
||||
return (uint32_t)g_stream_resident_experts.size();
|
||||
}
|
||||
|
||||
extern "C" void ds4_gpu_stream_expert_cache_reset_route_hotness(void) {
|
||||
}
|
||||
|
||||
extern "C" void ds4_gpu_stream_expert_cache_release_resident(void) {
|
||||
cuda_stream_resident_cache_release();
|
||||
}
|
||||
|
||||
extern "C" uint32_t ds4_gpu_stream_expert_cache_budget_for_expert_size(
|
||||
uint64_t gate_expert_bytes,
|
||||
uint64_t down_expert_bytes) {
|
||||
(void)gate_expert_bytes;
|
||||
(void)down_expert_bytes;
|
||||
return ds4_gpu_stream_expert_cache_configured_count();
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_stream_expert_cache_seed_selected(
|
||||
const ds4_gpu_stream_expert_table *table,
|
||||
const int32_t *selected_ids,
|
||||
uint32_t n_selected) {
|
||||
if (!table) return 0;
|
||||
if (!cuda_stream_selected_load(table->model_map,
|
||||
table->model_size,
|
||||
table->layer,
|
||||
selected_ids,
|
||||
table->n_total_expert,
|
||||
n_selected,
|
||||
table->gate_offset,
|
||||
table->up_offset,
|
||||
table->down_offset,
|
||||
table->gate_expert_bytes,
|
||||
table->down_expert_bytes)) {
|
||||
return 0;
|
||||
}
|
||||
return cuda_stream_selected_finish_pending_missing(0);
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_stream_expert_cache_begin_selected_load(
|
||||
const ds4_gpu_stream_expert_table *table,
|
||||
const int32_t *selected_ids,
|
||||
uint32_t n_selected) {
|
||||
if (!table) return 0;
|
||||
return cuda_stream_selected_load(table->model_map,
|
||||
table->model_size,
|
||||
table->layer,
|
||||
selected_ids,
|
||||
table->n_total_expert,
|
||||
n_selected,
|
||||
table->gate_offset,
|
||||
table->up_offset,
|
||||
table->down_offset,
|
||||
table->gate_expert_bytes,
|
||||
table->down_expert_bytes);
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_stream_expert_cache_prepare_selected_batch(
|
||||
const ds4_gpu_stream_expert_table *table,
|
||||
const int32_t *selected_ids,
|
||||
uint32_t n_tokens,
|
||||
uint32_t n_selected) {
|
||||
if (!table) return 0;
|
||||
const ds4_gpu_tensor *selected_exec = NULL;
|
||||
const char **gate_ptrs = NULL;
|
||||
const char **up_ptrs = NULL;
|
||||
const char **down_ptrs = NULL;
|
||||
uint32_t unique = 0;
|
||||
return cuda_stream_batch_selected_prepare_from_host(table->model_map,
|
||||
table->model_size,
|
||||
table->layer,
|
||||
selected_ids,
|
||||
n_tokens,
|
||||
table->n_total_expert,
|
||||
n_selected,
|
||||
table->gate_offset,
|
||||
table->up_offset,
|
||||
table->down_offset,
|
||||
table->gate_expert_bytes,
|
||||
table->down_expert_bytes,
|
||||
&selected_exec,
|
||||
&gate_ptrs,
|
||||
&up_ptrs,
|
||||
&down_ptrs,
|
||||
&unique,
|
||||
1);
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_stream_expert_cache_load_layer(
|
||||
const ds4_gpu_stream_expert_table *table) {
|
||||
if (!table) return 0;
|
||||
return cuda_stream_layer_expert_cache_load(table->model_map,
|
||||
table->model_size,
|
||||
table->layer,
|
||||
table->n_total_expert,
|
||||
table->gate_offset,
|
||||
table->up_offset,
|
||||
table->down_offset,
|
||||
table->gate_expert_bytes,
|
||||
table->down_expert_bytes);
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_stream_expert_cache_seed_from_layer_selected(
|
||||
const ds4_gpu_stream_expert_table *table,
|
||||
const ds4_gpu_tensor *selected,
|
||||
uint32_t n_tokens,
|
||||
uint32_t n_seed_tokens,
|
||||
uint32_t n_selected) {
|
||||
if (!table) return 0;
|
||||
return cuda_stream_layer_expert_cache_seed_selected(table->model_map,
|
||||
table->layer,
|
||||
selected,
|
||||
n_tokens,
|
||||
n_seed_tokens,
|
||||
table->n_total_expert,
|
||||
n_selected,
|
||||
table->gate_offset,
|
||||
table->up_offset,
|
||||
table->down_offset,
|
||||
table->gate_expert_bytes,
|
||||
table->down_expert_bytes);
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_stream_expert_cache_release_layer_cache(void) {
|
||||
cuda_stream_layer_expert_cache_release();
|
||||
return 1;
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_stream_expert_cache_seed_experts(
|
||||
const ds4_gpu_stream_expert_table *table,
|
||||
const int32_t *expert_ids,
|
||||
const uint32_t *expert_priorities,
|
||||
uint32_t n_experts) {
|
||||
(void)table;
|
||||
(void)expert_ids;
|
||||
(void)expert_priorities;
|
||||
(void)n_experts;
|
||||
return 1;
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_routed_moe_set_selected_override(
|
||||
const int32_t *selected,
|
||||
uint32_t n_selected) {
|
||||
if (n_selected > DS4_ROCM_N_EXPERT_USED || (!selected && n_selected != 0)) return 0;
|
||||
for (uint32_t i = 0; i < n_selected; i++) {
|
||||
g_routed_moe_selected_override[i] = selected[i];
|
||||
}
|
||||
g_routed_moe_selected_override_n = n_selected;
|
||||
return 1;
|
||||
}
|
||||
@@ -0,0 +1,113 @@
|
||||
extern "C" int ds4_gpu_embed_token_hc_tensor(ds4_gpu_tensor *out_hc, const void *model_map, uint64_t model_size, uint64_t weight_offset, uint32_t n_vocab, uint32_t token, uint32_t n_embd, uint32_t n_hc) {
|
||||
if (!out_hc || !model_map || n_vocab == 0u || token >= n_vocab || n_embd == 0u || n_hc == 0u ||
|
||||
(uint64_t)n_embd * n_hc > UINT32_MAX) return 0;
|
||||
uint64_t weight_bytes = 0;
|
||||
uint64_t out_bytes = 0;
|
||||
if (!cuda_u64_mul3_checked(n_vocab, n_embd, sizeof(uint16_t), &weight_bytes) ||
|
||||
!cuda_model_range_fits(model_size, weight_offset, weight_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_hc, n_embd, sizeof(float), &out_bytes) ||
|
||||
!cuda_tensor_has_bytes(out_hc, out_bytes)) return 0;
|
||||
const char *wptr = cuda_model_range_ptr(model_map, weight_offset, weight_bytes, "token_embd");
|
||||
if (!wptr) return 0;
|
||||
uint64_t n = (uint64_t)n_embd * n_hc;
|
||||
embed_token_hc_kernel<<<(n + 255u) / 256u, 256>>>((float *)out_hc->ptr, (const unsigned short *)wptr, token, n_embd, n_hc);
|
||||
return cuda_ok(cudaGetLastError(), "embed token launch");
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_embed_tokens_hc_tensor(
|
||||
ds4_gpu_tensor *out_hc,
|
||||
const ds4_gpu_tensor *tokens_t,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t weight_offset,
|
||||
uint32_t n_vocab,
|
||||
uint32_t n_tokens,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc) {
|
||||
if (!out_hc || !tokens_t || !model_map || n_vocab == 0u || n_tokens == 0u || n_embd == 0u || n_hc == 0u) {
|
||||
return 0;
|
||||
}
|
||||
uint64_t weight_bytes = 0, n = 0;
|
||||
if (!cuda_u64_mul3_checked(n_vocab, n_embd, sizeof(uint16_t), &weight_bytes) ||
|
||||
!cuda_model_range_fits(model_size, weight_offset, weight_bytes) ||
|
||||
!cuda_tensor_has_i32(tokens_t, n_tokens) ||
|
||||
!cuda_u64_mul_checked(n_tokens, n_hc, &n) ||
|
||||
!cuda_u64_mul_checked(n, n_embd, &n) ||
|
||||
!cuda_tensor_has_f32(out_hc, n)) {
|
||||
return 0;
|
||||
}
|
||||
const char *wptr = cuda_model_range_ptr(model_map, weight_offset, weight_bytes, "token_embd");
|
||||
if (!wptr) return 0;
|
||||
embed_tokens_hc_kernel<<<(n + 255) / 256, 256>>>(
|
||||
(float *)out_hc->ptr,
|
||||
(const int32_t *)tokens_t->ptr,
|
||||
(const __half *)wptr,
|
||||
n_vocab, n_tokens, n_embd, n_hc);
|
||||
return cuda_ok(cudaGetLastError(), "embed tokens launch");
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_embed_token_hc_q8_0_tensor(
|
||||
ds4_gpu_tensor *out_hc,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t weight_offset,
|
||||
uint32_t n_vocab,
|
||||
uint32_t token,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc) {
|
||||
if (!out_hc || !model_map || n_vocab == 0u || token >= n_vocab || n_embd == 0 || n_hc == 0) return 0;
|
||||
const uint64_t blocks = (n_embd + 31u) / 32u;
|
||||
uint64_t row_bytes = 0, weight_bytes = 0, out_bytes = 0;
|
||||
if (!cuda_u64_mul_checked(blocks, 34u, &row_bytes) ||
|
||||
!cuda_u64_mul_checked(n_vocab, row_bytes, &weight_bytes) ||
|
||||
!cuda_model_range_fits(model_size, weight_offset, weight_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_embd, n_hc, sizeof(float), &out_bytes) ||
|
||||
!cuda_tensor_has_bytes(out_hc, out_bytes)) {
|
||||
return 0;
|
||||
}
|
||||
const char *wptr = cuda_model_range_ptr(model_map, weight_offset, weight_bytes, "token_embd_q8_0");
|
||||
if (!wptr) return 0;
|
||||
const uint64_t n = (uint64_t)n_embd * n_hc;
|
||||
embed_token_hc_q8_0_kernel<<<(n + 255u) / 256u, 256>>>(
|
||||
(float *)out_hc->ptr,
|
||||
(const unsigned char *)wptr,
|
||||
token,
|
||||
n_embd,
|
||||
n_hc);
|
||||
return cuda_ok(cudaGetLastError(), "embed token q8_0 launch");
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_embed_tokens_hc_q8_0_tensor(
|
||||
ds4_gpu_tensor *out_hc,
|
||||
const ds4_gpu_tensor *tokens_t,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t weight_offset,
|
||||
uint32_t n_vocab,
|
||||
uint32_t n_tokens,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc) {
|
||||
if (!out_hc || !tokens_t || !model_map || n_vocab == 0u || n_tokens == 0 || n_embd == 0 || n_hc == 0) return 0;
|
||||
const uint64_t blocks = (n_embd + 31u) / 32u;
|
||||
uint64_t row_bytes = 0, weight_bytes = 0, n = 0;
|
||||
if (!cuda_u64_mul_checked(blocks, 34u, &row_bytes) ||
|
||||
!cuda_u64_mul_checked(n_vocab, row_bytes, &weight_bytes) ||
|
||||
!cuda_model_range_fits(model_size, weight_offset, weight_bytes) ||
|
||||
!cuda_tensor_has_i32(tokens_t, n_tokens) ||
|
||||
!cuda_u64_mul_checked(n_tokens, n_hc, &n) ||
|
||||
!cuda_u64_mul_checked(n, n_embd, &n) ||
|
||||
!cuda_tensor_has_f32(out_hc, n)) {
|
||||
return 0;
|
||||
}
|
||||
const char *wptr = cuda_model_range_ptr(model_map, weight_offset, weight_bytes, "token_embd_q8_0");
|
||||
if (!wptr) return 0;
|
||||
embed_tokens_hc_q8_0_kernel<<<(n + 255u) / 256u, 256>>>(
|
||||
(float *)out_hc->ptr,
|
||||
(const int32_t *)tokens_t->ptr,
|
||||
(const unsigned char *)wptr,
|
||||
n_vocab,
|
||||
n_tokens,
|
||||
n_embd,
|
||||
n_hc);
|
||||
return cuda_ok(cudaGetLastError(), "embed tokens q8_0 launch");
|
||||
}
|
||||
@@ -0,0 +1,40 @@
|
||||
// DS4 ROCm FP8 KV quantization and raw-cache store kernels.
|
||||
//
|
||||
// This file is included from ds4_cuda.cu (same translation unit) so these
|
||||
// kernels can reuse the backend's existing device helpers without HIP device
|
||||
// linking or behavior changes.
|
||||
|
||||
__global__ static void fp8_kv_quantize_kernel(float *x, uint32_t n_tok, uint32_t head_dim, uint32_t n_rot) {
|
||||
const uint32_t row = blockIdx.x;
|
||||
const uint32_t grp = blockIdx.y;
|
||||
const uint32_t tid = threadIdx.x;
|
||||
const uint32_t n_nope = head_dim - n_rot;
|
||||
const uint32_t off = grp * 64u;
|
||||
if (row >= n_tok || off >= n_nope) return;
|
||||
float *xr = x + (uint64_t)row * head_dim;
|
||||
__shared__ float scratch[64];
|
||||
float v = 0.0f;
|
||||
if (tid < 64u && off + tid < n_nope) v = xr[off + tid];
|
||||
scratch[tid] = (tid < 64u && off + tid < n_nope) ? fabsf(v) : 0.0f;
|
||||
__syncthreads();
|
||||
for (uint32_t stride = 32; stride > 0; stride >>= 1) {
|
||||
if (tid < stride) scratch[tid] = fmaxf(scratch[tid], scratch[tid + stride]);
|
||||
__syncthreads();
|
||||
}
|
||||
const float scale = exp2f(ceilf(log2f(fmaxf(scratch[0], 1.0e-4f) / 448.0f)));
|
||||
if (tid < 64u && off + tid < n_nope) {
|
||||
const float q = dsv4_e4m3fn_dequant_dev(fminf(448.0f, fmaxf(-448.0f, v / scale))) * scale;
|
||||
xr[off + tid] = q;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ static void store_raw_kv_batch_kernel(float *raw, const float *kv, uint32_t raw_cap, uint32_t pos0, uint32_t n_tokens, uint32_t head_dim) {
|
||||
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint64_t n = (uint64_t)n_tokens * head_dim;
|
||||
if (gid >= n) return;
|
||||
uint32_t d = gid % head_dim;
|
||||
uint32_t t = gid / head_dim;
|
||||
uint32_t row = (pos0 + t) % raw_cap;
|
||||
const uint16_t hb = f32_to_f16_bits_hip_round(kv[(uint64_t)t * head_dim + d]);
|
||||
raw[(uint64_t)row * head_dim + d] = f16_bits_to_f32(hb);
|
||||
}
|
||||
@@ -0,0 +1,9 @@
|
||||
extern "C" int ds4_gpu_dsv4_fp8_kv_quantize_tensor(ds4_gpu_tensor *x, uint32_t n_tok, uint32_t head_dim, uint32_t n_rot) {
|
||||
if (n_rot > head_dim || !cuda_tensor_has_elems2(x, n_tok, head_dim, sizeof(float))) return 0;
|
||||
if (n_tok == 0u || head_dim == 0u) return 1;
|
||||
const uint32_t n_nope = head_dim - n_rot;
|
||||
if (n_nope == 0) return 1;
|
||||
const uint32_t groups = (n_nope + 63u) / 64u;
|
||||
fp8_kv_quantize_kernel<<<dim3(n_tok, groups), 64>>>((float *)x->ptr, n_tok, head_dim, n_rot);
|
||||
return cuda_ok(cudaGetLastError(), "fp8_kv_quantize launch");
|
||||
}
|
||||
@@ -0,0 +1,373 @@
|
||||
// DS4 ROCm hierarchical-composition (HC) split/sum/expand kernels.
|
||||
//
|
||||
// Included from ds4_cuda.cu in the same translation unit to preserve current
|
||||
// static helper visibility and launch behavior.
|
||||
|
||||
__device__ static void hc4_split_one(float *out, const float *mix, const float *scale, const float *base, uint32_t sinkhorn_iters, float epsv) {
|
||||
const float pre_scale = scale[0];
|
||||
const float post_scale = scale[1];
|
||||
const float comb_scale = scale[2];
|
||||
for (int i = 0; i < 4; i++) {
|
||||
float z = mix[i] * pre_scale + base[i];
|
||||
out[i] = 1.0f / (1.0f + expf(-z)) + epsv;
|
||||
}
|
||||
for (int i = 0; i < 4; i++) {
|
||||
float z = mix[4 + i] * post_scale + base[4 + i];
|
||||
out[4 + i] = 2.0f / (1.0f + expf(-z));
|
||||
}
|
||||
float c[16];
|
||||
for (int r = 0; r < 4; r++) {
|
||||
float m = -INFINITY;
|
||||
for (int col = 0; col < 4; col++) {
|
||||
float v = mix[8 + r * 4 + col] * comb_scale + base[8 + r * 4 + col];
|
||||
c[r * 4 + col] = v;
|
||||
m = fmaxf(m, v);
|
||||
}
|
||||
float s = 0.0f;
|
||||
for (int col = 0; col < 4; col++) {
|
||||
float v = expf(c[r * 4 + col] - m);
|
||||
c[r * 4 + col] = v;
|
||||
s += v;
|
||||
}
|
||||
for (int col = 0; col < 4; col++) c[r * 4 + col] = c[r * 4 + col] / s + epsv;
|
||||
}
|
||||
for (int col = 0; col < 4; col++) {
|
||||
float s = epsv;
|
||||
for (int r = 0; r < 4; r++) s += c[r * 4 + col];
|
||||
for (int r = 0; r < 4; r++) c[r * 4 + col] /= s;
|
||||
}
|
||||
for (uint32_t iter = 1; iter < sinkhorn_iters; iter++) {
|
||||
for (int r = 0; r < 4; r++) {
|
||||
float s = epsv;
|
||||
for (int col = 0; col < 4; col++) s += c[r * 4 + col];
|
||||
for (int col = 0; col < 4; col++) c[r * 4 + col] /= s;
|
||||
}
|
||||
for (int col = 0; col < 4; col++) {
|
||||
float s = epsv;
|
||||
for (int r = 0; r < 4; r++) s += c[r * 4 + col];
|
||||
for (int r = 0; r < 4; r++) c[r * 4 + col] /= s;
|
||||
}
|
||||
}
|
||||
for (int i = 0; i < 16; i++) out[8 + i] = c[i];
|
||||
}
|
||||
|
||||
__global__ static void hc_split_sinkhorn_kernel(float *out, const float *mix, const float *scale, const float *base, uint32_t n_rows, uint32_t sinkhorn_iters, float epsv) {
|
||||
uint32_t row = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (row >= n_rows) return;
|
||||
hc4_split_one(out + (uint64_t)row * 24, mix + (uint64_t)row * 24, scale, base, sinkhorn_iters, epsv);
|
||||
}
|
||||
|
||||
__global__ static void hc_weighted_sum_kernel(float *out, const float *x, const float *w, uint32_t n_embd, uint32_t n_hc, uint32_t n_tokens, uint32_t weight_stride_f32) {
|
||||
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint64_t n = (uint64_t)n_embd * n_tokens;
|
||||
if (gid >= n) return;
|
||||
uint32_t d = gid % n_embd;
|
||||
uint32_t t = gid / n_embd;
|
||||
float acc = 0.0f;
|
||||
for (uint32_t h = 0; h < n_hc; h++) {
|
||||
acc += x[(uint64_t)t * n_hc * n_embd + (uint64_t)h * n_embd + d] *
|
||||
w[(uint64_t)t * weight_stride_f32 + h];
|
||||
}
|
||||
out[(uint64_t)t * n_embd + d] = acc;
|
||||
}
|
||||
|
||||
__global__ static void hc_expand_kernel(
|
||||
float *out_hc,
|
||||
const float *block_out,
|
||||
const float *block_add,
|
||||
const float *residual_hc,
|
||||
const float *post,
|
||||
const float *comb,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc,
|
||||
uint32_t n_tokens,
|
||||
uint32_t post_stride,
|
||||
uint32_t comb_stride,
|
||||
int has_add) {
|
||||
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd;
|
||||
if (gid >= n_elem) return;
|
||||
uint32_t d = gid % n_embd;
|
||||
uint64_t tmp = gid / n_embd;
|
||||
uint32_t dst_hc = tmp % n_hc;
|
||||
uint32_t t = tmp / n_hc;
|
||||
|
||||
float block_v = block_out[(uint64_t)t * n_embd + d];
|
||||
if (has_add) block_v += block_add[(uint64_t)t * n_embd + d];
|
||||
float acc = block_v * post[(uint64_t)t * post_stride + dst_hc];
|
||||
for (uint32_t src_hc = 0; src_hc < n_hc; src_hc++) {
|
||||
float comb_v = comb[(uint64_t)t * comb_stride + dst_hc + (uint64_t)src_hc * n_hc];
|
||||
float res_v = residual_hc[(uint64_t)t * n_hc * n_embd + (uint64_t)src_hc * n_embd + d];
|
||||
acc += comb_v * res_v;
|
||||
}
|
||||
out_hc[(uint64_t)t * n_hc * n_embd + (uint64_t)dst_hc * n_embd + d] = acc;
|
||||
}
|
||||
|
||||
__global__ static void hc_expand_half_kernel(
|
||||
float *out_hc,
|
||||
const __half *block_out,
|
||||
const float *residual_hc,
|
||||
const float *post,
|
||||
const float *comb,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc,
|
||||
uint32_t n_tokens,
|
||||
uint32_t post_stride,
|
||||
uint32_t comb_stride) {
|
||||
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd;
|
||||
if (gid >= n_elem) return;
|
||||
uint32_t d = gid % n_embd;
|
||||
uint64_t tmp = gid / n_embd;
|
||||
uint32_t dst_hc = tmp % n_hc;
|
||||
uint32_t t = tmp / n_hc;
|
||||
|
||||
const float block_v = __half2float(block_out[(uint64_t)t * n_embd + d]);
|
||||
float acc = block_v * post[(uint64_t)t * post_stride + dst_hc];
|
||||
for (uint32_t src_hc = 0; src_hc < n_hc; src_hc++) {
|
||||
float comb_v = comb[(uint64_t)t * comb_stride + dst_hc + (uint64_t)src_hc * n_hc];
|
||||
float res_v = residual_hc[(uint64_t)t * n_hc * n_embd + (uint64_t)src_hc * n_embd + d];
|
||||
acc += comb_v * res_v;
|
||||
}
|
||||
out_hc[(uint64_t)t * n_hc * n_embd + (uint64_t)dst_hc * n_embd + d] = acc;
|
||||
}
|
||||
|
||||
__global__ static void hc_expand_add_half_kernel(
|
||||
float *out_hc,
|
||||
const float *block_out,
|
||||
const __half *block_add,
|
||||
const float *residual_hc,
|
||||
const float *post,
|
||||
const float *comb,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc,
|
||||
uint32_t n_tokens,
|
||||
uint32_t post_stride,
|
||||
uint32_t comb_stride) {
|
||||
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd;
|
||||
if (gid >= n_elem) return;
|
||||
uint32_t d = gid % n_embd;
|
||||
uint64_t tmp = gid / n_embd;
|
||||
uint32_t dst_hc = tmp % n_hc;
|
||||
uint32_t t = tmp / n_hc;
|
||||
|
||||
const float block_v = block_out[(uint64_t)t * n_embd + d] +
|
||||
__half2float(block_add[(uint64_t)t * n_embd + d]);
|
||||
float acc = block_v * post[(uint64_t)t * post_stride + dst_hc];
|
||||
for (uint32_t src_hc = 0; src_hc < n_hc; src_hc++) {
|
||||
float comb_v = comb[(uint64_t)t * comb_stride + dst_hc + (uint64_t)src_hc * n_hc];
|
||||
float res_v = residual_hc[(uint64_t)t * n_hc * n_embd + (uint64_t)src_hc * n_embd + d];
|
||||
acc += comb_v * res_v;
|
||||
}
|
||||
out_hc[(uint64_t)t * n_hc * n_embd + (uint64_t)dst_hc * n_embd + d] = acc;
|
||||
}
|
||||
|
||||
__global__ static void hc_expand4_kernel(
|
||||
float *out_hc,
|
||||
const float *block_out,
|
||||
const float *residual_hc,
|
||||
const float *split,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_tokens) {
|
||||
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const uint64_t n = (uint64_t)n_tokens * n_embd;
|
||||
if (gid >= n) return;
|
||||
const uint32_t d = gid % n_embd;
|
||||
const uint32_t t = gid / n_embd;
|
||||
const uint64_t td = (uint64_t)t * n_embd + d;
|
||||
const uint64_t hc_base = (uint64_t)t * 4u * n_embd + d;
|
||||
const float bv = block_out[td];
|
||||
const float r0 = residual_hc[hc_base + 0u * (uint64_t)n_embd];
|
||||
const float r1 = residual_hc[hc_base + 1u * (uint64_t)n_embd];
|
||||
const float r2 = residual_hc[hc_base + 2u * (uint64_t)n_embd];
|
||||
const float r3 = residual_hc[hc_base + 3u * (uint64_t)n_embd];
|
||||
const float *sp = split + (uint64_t)t * 24u;
|
||||
const float *post = sp + 4u;
|
||||
const float *comb = sp + 8u;
|
||||
#pragma unroll
|
||||
for (uint32_t dst = 0; dst < 4u; dst++) {
|
||||
float acc = bv * post[dst];
|
||||
acc += comb[0u * 4u + dst] * r0;
|
||||
acc += comb[1u * 4u + dst] * r1;
|
||||
acc += comb[2u * 4u + dst] * r2;
|
||||
acc += comb[3u * 4u + dst] * r3;
|
||||
out_hc[hc_base + (uint64_t)dst * n_embd] = acc;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ static void hc_expand4_add_kernel(
|
||||
float *out_hc,
|
||||
const float *block_out,
|
||||
const float *block_add,
|
||||
const float *residual_hc,
|
||||
const float *split,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_tokens) {
|
||||
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const uint64_t n = (uint64_t)n_tokens * n_embd;
|
||||
if (gid >= n) return;
|
||||
const uint32_t d = gid % n_embd;
|
||||
const uint32_t t = gid / n_embd;
|
||||
const uint64_t td = (uint64_t)t * n_embd + d;
|
||||
const uint64_t hc_base = (uint64_t)t * 4u * n_embd + d;
|
||||
const float bv = block_out[td] + block_add[td];
|
||||
const float r0 = residual_hc[hc_base + 0u * (uint64_t)n_embd];
|
||||
const float r1 = residual_hc[hc_base + 1u * (uint64_t)n_embd];
|
||||
const float r2 = residual_hc[hc_base + 2u * (uint64_t)n_embd];
|
||||
const float r3 = residual_hc[hc_base + 3u * (uint64_t)n_embd];
|
||||
const float *sp = split + (uint64_t)t * 24u;
|
||||
const float *post = sp + 4u;
|
||||
const float *comb = sp + 8u;
|
||||
#pragma unroll
|
||||
for (uint32_t dst = 0; dst < 4u; dst++) {
|
||||
float acc = bv * post[dst];
|
||||
acc += comb[0u * 4u + dst] * r0;
|
||||
acc += comb[1u * 4u + dst] * r1;
|
||||
acc += comb[2u * 4u + dst] * r2;
|
||||
acc += comb[3u * 4u + dst] * r3;
|
||||
out_hc[hc_base + (uint64_t)dst * n_embd] = acc;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ static void hc_expand4_half_kernel(
|
||||
float *out_hc,
|
||||
const __half *block_out,
|
||||
const float *residual_hc,
|
||||
const float *split,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_tokens) {
|
||||
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const uint64_t n = (uint64_t)n_tokens * n_embd;
|
||||
if (gid >= n) return;
|
||||
const uint32_t d = gid % n_embd;
|
||||
const uint32_t t = gid / n_embd;
|
||||
const uint64_t td = (uint64_t)t * n_embd + d;
|
||||
const uint64_t hc_base = (uint64_t)t * 4u * n_embd + d;
|
||||
const float bv = __half2float(block_out[td]);
|
||||
const float r0 = residual_hc[hc_base + 0u * (uint64_t)n_embd];
|
||||
const float r1 = residual_hc[hc_base + 1u * (uint64_t)n_embd];
|
||||
const float r2 = residual_hc[hc_base + 2u * (uint64_t)n_embd];
|
||||
const float r3 = residual_hc[hc_base + 3u * (uint64_t)n_embd];
|
||||
const float *sp = split + (uint64_t)t * 24u;
|
||||
const float *post = sp + 4u;
|
||||
const float *comb = sp + 8u;
|
||||
#pragma unroll
|
||||
for (uint32_t dst = 0; dst < 4u; dst++) {
|
||||
float acc = bv * post[dst];
|
||||
acc += comb[0u * 4u + dst] * r0;
|
||||
acc += comb[1u * 4u + dst] * r1;
|
||||
acc += comb[2u * 4u + dst] * r2;
|
||||
acc += comb[3u * 4u + dst] * r3;
|
||||
out_hc[hc_base + (uint64_t)dst * n_embd] = acc;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ static void hc_expand4_add_half_kernel(
|
||||
float *out_hc,
|
||||
const float *block_out,
|
||||
const __half *block_add,
|
||||
const float *residual_hc,
|
||||
const float *split,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_tokens) {
|
||||
uint64_t gid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const uint64_t n = (uint64_t)n_tokens * n_embd;
|
||||
if (gid >= n) return;
|
||||
const uint32_t d = gid % n_embd;
|
||||
const uint32_t t = gid / n_embd;
|
||||
const uint64_t td = (uint64_t)t * n_embd + d;
|
||||
const uint64_t hc_base = (uint64_t)t * 4u * n_embd + d;
|
||||
const float bv = block_out[td] + __half2float(block_add[td]);
|
||||
const float r0 = residual_hc[hc_base + 0u * (uint64_t)n_embd];
|
||||
const float r1 = residual_hc[hc_base + 1u * (uint64_t)n_embd];
|
||||
const float r2 = residual_hc[hc_base + 2u * (uint64_t)n_embd];
|
||||
const float r3 = residual_hc[hc_base + 3u * (uint64_t)n_embd];
|
||||
const float *sp = split + (uint64_t)t * 24u;
|
||||
const float *post = sp + 4u;
|
||||
const float *comb = sp + 8u;
|
||||
#pragma unroll
|
||||
for (uint32_t dst = 0; dst < 4u; dst++) {
|
||||
float acc = bv * post[dst];
|
||||
acc += comb[0u * 4u + dst] * r0;
|
||||
acc += comb[1u * 4u + dst] * r1;
|
||||
acc += comb[2u * 4u + dst] * r2;
|
||||
acc += comb[3u * 4u + dst] * r3;
|
||||
out_hc[hc_base + (uint64_t)dst * n_embd] = acc;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ static void hc_split_weighted_sum_fused_kernel(
|
||||
float *out,
|
||||
float *split,
|
||||
const float *mix,
|
||||
const float *residual_hc,
|
||||
const float *scale,
|
||||
const float *base,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc,
|
||||
uint32_t n_rows,
|
||||
uint32_t sinkhorn_iters,
|
||||
float epsv) {
|
||||
uint32_t t = blockIdx.x;
|
||||
uint32_t d = threadIdx.x;
|
||||
if (t >= n_rows || n_hc != 4) return;
|
||||
const uint32_t mix_hc = 24;
|
||||
float *sp = split + (uint64_t)t * mix_hc;
|
||||
if (d == 0) hc4_split_one(sp, mix + (uint64_t)t * mix_hc, scale, base, sinkhorn_iters, epsv);
|
||||
__syncthreads();
|
||||
for (uint32_t col = d; col < n_embd; col += blockDim.x) {
|
||||
float acc = 0.0f;
|
||||
for (uint32_t h = 0; h < 4; h++) {
|
||||
acc += residual_hc[(uint64_t)t * 4u * n_embd + (uint64_t)h * n_embd + col] * sp[h];
|
||||
}
|
||||
out[(uint64_t)t * n_embd + col] = acc;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ static void hc_split_weighted_sum_norm_fused_kernel(
|
||||
float *out,
|
||||
float *norm_out,
|
||||
float *split,
|
||||
const float *mix,
|
||||
const float *residual_hc,
|
||||
const float *scale,
|
||||
const float *base,
|
||||
const float *norm_w,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc,
|
||||
uint32_t n_rows,
|
||||
uint32_t sinkhorn_iters,
|
||||
float epsv,
|
||||
float norm_eps) {
|
||||
const uint32_t t = blockIdx.x;
|
||||
const uint32_t d = threadIdx.x;
|
||||
if (t >= n_rows || n_hc != 4) return;
|
||||
const uint32_t mix_hc = 24;
|
||||
float *sp = split + (uint64_t)t * mix_hc;
|
||||
if (d == 0) hc4_split_one(sp, mix + (uint64_t)t * mix_hc, scale, base, sinkhorn_iters, epsv);
|
||||
__syncthreads();
|
||||
|
||||
float sum = 0.0f;
|
||||
for (uint32_t col = d; col < n_embd; col += blockDim.x) {
|
||||
float acc = 0.0f;
|
||||
for (uint32_t h = 0; h < 4; h++) {
|
||||
acc += residual_hc[(uint64_t)t * 4u * n_embd + (uint64_t)h * n_embd + col] * sp[h];
|
||||
}
|
||||
out[(uint64_t)t * n_embd + col] = acc;
|
||||
sum += acc * acc;
|
||||
}
|
||||
|
||||
__shared__ float partial[256];
|
||||
partial[d] = sum;
|
||||
__syncthreads();
|
||||
for (uint32_t stride = blockDim.x >> 1; stride > 0; stride >>= 1) {
|
||||
if (d < stride) partial[d] += partial[d + stride];
|
||||
__syncthreads();
|
||||
}
|
||||
const float norm_scale = rsqrtf(partial[0] / (float)n_embd + norm_eps);
|
||||
for (uint32_t col = d; col < n_embd; col += blockDim.x) {
|
||||
const float v = out[(uint64_t)t * n_embd + col];
|
||||
norm_out[(uint64_t)t * n_embd + col] = v * norm_scale * norm_w[col];
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,437 @@
|
||||
extern "C" int ds4_gpu_repeat_hc_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *row, uint32_t n_embd, uint32_t n_hc) {
|
||||
uint64_t n = 0;
|
||||
if (n_embd == 0 || n_hc == 0 ||
|
||||
!cuda_u64_mul_checked(n_embd, n_hc, &n) ||
|
||||
!cuda_tensor_has_f32(row, n_embd) || !cuda_tensor_has_f32(out, n)) {
|
||||
return 0;
|
||||
}
|
||||
repeat_hc_kernel<<<(n + 255) / 256, 256>>>((float *)out->ptr, (const float *)row->ptr, n_embd, n_hc);
|
||||
return cuda_ok(cudaGetLastError(), "repeat_hc launch");
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_hc_split_sinkhorn_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *mix, const void *model_map, uint64_t model_size, uint64_t scale_offset, uint64_t base_offset, uint32_t n_hc, uint32_t sinkhorn_iters, float eps) {
|
||||
if (!out || !mix || !model_map || n_hc != 4) return 0;
|
||||
const uint64_t mix_bytes = 24ull * sizeof(float);
|
||||
if (!cuda_model_range_fits(model_size, scale_offset, 3ull * sizeof(float)) ||
|
||||
!cuda_model_range_fits(model_size, base_offset, mix_bytes) ||
|
||||
!cuda_tensor_has_bytes(mix, mix_bytes) || !cuda_tensor_has_bytes(out, mix_bytes)) return 0;
|
||||
const float *scale = (const float *)cuda_model_range_ptr(model_map, scale_offset, 3ull * sizeof(float), "hc_scale");
|
||||
const float *base = (const float *)cuda_model_range_ptr(model_map, base_offset, mix_bytes, "hc_base");
|
||||
if (!scale || !base) return 0;
|
||||
uint32_t n_rows = (uint32_t)(mix->bytes / mix_bytes);
|
||||
if (out->bytes / mix_bytes < n_rows) n_rows = (uint32_t)(out->bytes / mix_bytes);
|
||||
hc_split_sinkhorn_kernel<<<(n_rows + 255) / 256, 256>>>(
|
||||
(float *)out->ptr, (const float *)mix->ptr,
|
||||
scale,
|
||||
base,
|
||||
n_rows, sinkhorn_iters, eps);
|
||||
return cuda_ok(cudaGetLastError(), "hc_split_sinkhorn launch");
|
||||
}
|
||||
static int cuda_hc_flat_token_count(const ds4_gpu_tensor *out, uint32_t n_embd, uint64_t *n_tokens) {
|
||||
if (!out || n_embd == 0u || !n_tokens) return 0;
|
||||
uint64_t row_bytes = 0;
|
||||
if (!cuda_u64_mul3_checked(n_embd, 1u, sizeof(float), &row_bytes) ||
|
||||
row_bytes == 0u || out->bytes < row_bytes || (out->bytes % row_bytes) != 0u) return 0;
|
||||
*n_tokens = out->bytes / row_bytes;
|
||||
return *n_tokens != 0u && *n_tokens <= UINT32_MAX;
|
||||
}
|
||||
|
||||
static int cuda_hc_hc_token_count(const ds4_gpu_tensor *out_hc, uint32_t n_embd, uint32_t n_hc, uint64_t *n_tokens) {
|
||||
if (!out_hc || n_embd == 0u || n_hc == 0u || !n_tokens) return 0;
|
||||
uint64_t row_elems = 0, row_bytes = 0;
|
||||
if (!cuda_u64_mul_checked(n_hc, n_embd, &row_elems) ||
|
||||
!cuda_u64_mul_checked(row_elems, sizeof(float), &row_bytes) ||
|
||||
row_bytes == 0u || out_hc->bytes < row_bytes || (out_hc->bytes % row_bytes) != 0u) return 0;
|
||||
*n_tokens = out_hc->bytes / row_bytes;
|
||||
return *n_tokens != 0u && *n_tokens <= UINT32_MAX;
|
||||
}
|
||||
|
||||
static int cuda_hc_mix_width(uint32_t n_hc, uint64_t *mix_hc) {
|
||||
if (n_hc == 0u || !mix_hc) return 0;
|
||||
const uint64_t h = (uint64_t)n_hc;
|
||||
const uint64_t mix = 2ull * h + h * h;
|
||||
if (mix > UINT32_MAX) return 0;
|
||||
*mix_hc = mix;
|
||||
return 1;
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_hc_weighted_sum_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *weights, uint32_t n_embd, uint32_t n_hc) {
|
||||
uint64_t n_tokens64 = 0, residual_bytes = 0, weights_bytes = 0;
|
||||
if (!out || !residual_hc || !weights || n_hc == 0u ||
|
||||
!cuda_hc_flat_token_count(out, n_embd, &n_tokens64) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, (uint64_t)n_hc * n_embd, sizeof(float), &residual_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, n_hc, sizeof(float), &weights_bytes) ||
|
||||
residual_hc->bytes < residual_bytes || weights->bytes < weights_bytes) return 0;
|
||||
uint32_t n_tokens = (uint32_t)n_tokens64;
|
||||
hc_weighted_sum_kernel<<<((uint64_t)n_embd * n_tokens + 255) / 256, 256>>>(
|
||||
(float *)out->ptr, (const float *)residual_hc->ptr, (const float *)weights->ptr,
|
||||
n_embd, n_hc, n_tokens, n_hc);
|
||||
return cuda_ok(cudaGetLastError(), "hc_weighted_sum launch");
|
||||
}
|
||||
extern "C" int ds4_gpu_hc_weighted_sum_split_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *split, uint32_t n_embd, uint32_t n_hc) {
|
||||
uint64_t n_tokens64 = 0, residual_bytes = 0, split_bytes = 0, mix_hc = 0;
|
||||
if (!out || !residual_hc || !split ||
|
||||
!cuda_hc_flat_token_count(out, n_embd, &n_tokens64) ||
|
||||
!cuda_hc_mix_width(n_hc, &mix_hc) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, (uint64_t)n_hc * n_embd, sizeof(float), &residual_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, mix_hc, sizeof(float), &split_bytes) ||
|
||||
residual_hc->bytes < residual_bytes || split->bytes < split_bytes) return 0;
|
||||
uint32_t n_tokens = (uint32_t)n_tokens64;
|
||||
uint32_t stride = (uint32_t)mix_hc;
|
||||
hc_weighted_sum_kernel<<<((uint64_t)n_embd * n_tokens + 255) / 256, 256>>>(
|
||||
(float *)out->ptr, (const float *)residual_hc->ptr, (const float *)split->ptr,
|
||||
n_embd, n_hc, n_tokens, stride);
|
||||
return cuda_ok(cudaGetLastError(), "hc_weighted_sum_split launch");
|
||||
}
|
||||
extern "C" int ds4_gpu_hc_split_weighted_sum_tensor(
|
||||
ds4_gpu_tensor *out,
|
||||
ds4_gpu_tensor *split,
|
||||
const ds4_gpu_tensor *mix,
|
||||
const ds4_gpu_tensor *residual_hc,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t scale_offset,
|
||||
uint64_t base_offset,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc,
|
||||
uint32_t sinkhorn_iters,
|
||||
float eps) {
|
||||
if (!out || !split || !mix || !residual_hc || !model_map ||
|
||||
n_embd == 0 || n_hc != 4) {
|
||||
return 0;
|
||||
}
|
||||
const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc;
|
||||
const uint64_t mix_bytes = mix_hc * sizeof(float);
|
||||
const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float);
|
||||
const uint64_t residual_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float);
|
||||
if (out->bytes < out_row_bytes || out->bytes % out_row_bytes != 0 ||
|
||||
scale_offset > model_size || 3ull * sizeof(float) > model_size - scale_offset ||
|
||||
base_offset > model_size || mix_bytes > model_size - base_offset) {
|
||||
return 0;
|
||||
}
|
||||
uint64_t n_rows = out->bytes / out_row_bytes;
|
||||
if (mix->bytes < n_rows * mix_bytes ||
|
||||
split->bytes < n_rows * mix_bytes ||
|
||||
residual_hc->bytes < n_rows * residual_row_bytes) {
|
||||
return 0;
|
||||
}
|
||||
const float *scale = (const float *)cuda_model_range_ptr(model_map, scale_offset, 3ull * sizeof(float), "hc_scale");
|
||||
const float *base = (const float *)cuda_model_range_ptr(model_map, base_offset, mix_bytes, "hc_base");
|
||||
if (!scale || !base) return 0;
|
||||
hc_split_weighted_sum_fused_kernel<<<(uint32_t)n_rows, 256>>>(
|
||||
(float *)out->ptr,
|
||||
(float *)split->ptr,
|
||||
(const float *)mix->ptr,
|
||||
(const float *)residual_hc->ptr,
|
||||
scale,
|
||||
base,
|
||||
n_embd, n_hc, (uint32_t)n_rows, sinkhorn_iters, eps);
|
||||
return cuda_ok(cudaGetLastError(), "hc split weighted sum launch");
|
||||
}
|
||||
extern "C" int ds4_gpu_hc_split_weighted_sum_norm_tensor(
|
||||
ds4_gpu_tensor *out,
|
||||
ds4_gpu_tensor *norm_out,
|
||||
ds4_gpu_tensor *split,
|
||||
const ds4_gpu_tensor *mix,
|
||||
const ds4_gpu_tensor *residual_hc,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t scale_offset,
|
||||
uint64_t base_offset,
|
||||
uint64_t norm_weight_offset,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc,
|
||||
uint32_t sinkhorn_iters,
|
||||
float eps,
|
||||
float norm_eps) {
|
||||
if (!out || !norm_out || !split || !mix || !residual_hc || !model_map ||
|
||||
n_embd == 0 || n_hc != 4) {
|
||||
return 0;
|
||||
}
|
||||
const uint64_t mix_hc = 2ull * n_hc + (uint64_t)n_hc * n_hc;
|
||||
const uint64_t mix_bytes = mix_hc * sizeof(float);
|
||||
const uint64_t out_row_bytes = (uint64_t)n_embd * sizeof(float);
|
||||
const uint64_t residual_row_bytes = (uint64_t)n_hc * n_embd * sizeof(float);
|
||||
if (out->bytes < out_row_bytes || out->bytes % out_row_bytes != 0 ||
|
||||
norm_out->bytes < out->bytes ||
|
||||
scale_offset > model_size || 3ull * sizeof(float) > model_size - scale_offset ||
|
||||
base_offset > model_size || mix_bytes > model_size - base_offset ||
|
||||
norm_weight_offset > model_size ||
|
||||
(uint64_t)n_embd * sizeof(float) > model_size - norm_weight_offset) {
|
||||
return 0;
|
||||
}
|
||||
uint64_t n_rows = out->bytes / out_row_bytes;
|
||||
if (n_rows == 1) {
|
||||
if (mix->bytes < n_rows * mix_bytes ||
|
||||
split->bytes < n_rows * mix_bytes ||
|
||||
residual_hc->bytes < n_rows * residual_row_bytes) {
|
||||
return 0;
|
||||
}
|
||||
const float *scale = (const float *)cuda_model_range_ptr(model_map, scale_offset,
|
||||
3ull * sizeof(float), "hc_scale");
|
||||
const float *base = (const float *)cuda_model_range_ptr(model_map, base_offset,
|
||||
mix_bytes, "hc_base");
|
||||
const float *norm_w = (const float *)cuda_model_range_ptr(model_map, norm_weight_offset,
|
||||
(uint64_t)n_embd * sizeof(float), "hc_norm_weight");
|
||||
if (!scale || !base || !norm_w) return 0;
|
||||
hc_split_weighted_sum_norm_fused_kernel<<<(uint32_t)n_rows, 256>>>(
|
||||
(float *)out->ptr,
|
||||
(float *)norm_out->ptr,
|
||||
(float *)split->ptr,
|
||||
(const float *)mix->ptr,
|
||||
(const float *)residual_hc->ptr,
|
||||
scale,
|
||||
base,
|
||||
norm_w,
|
||||
n_embd, n_hc, (uint32_t)n_rows, sinkhorn_iters, eps, norm_eps);
|
||||
return cuda_ok(cudaGetLastError(), "hc split weighted sum norm launch");
|
||||
}
|
||||
return ds4_gpu_hc_split_weighted_sum_tensor(out, split, mix, residual_hc,
|
||||
model_map, model_size,
|
||||
scale_offset, base_offset,
|
||||
n_embd, n_hc,
|
||||
sinkhorn_iters, eps) &&
|
||||
ds4_gpu_rms_norm_weight_tensor(norm_out, out, model_map, model_size,
|
||||
norm_weight_offset, n_embd, norm_eps);
|
||||
}
|
||||
extern "C" int ds4_gpu_output_hc_weights_tensor(
|
||||
ds4_gpu_tensor *out,
|
||||
const ds4_gpu_tensor *pre,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t scale_offset,
|
||||
uint64_t base_offset,
|
||||
uint32_t n_hc,
|
||||
float eps) {
|
||||
if (!out || !pre || !model_map || n_hc == 0) return 0;
|
||||
const uint64_t row_bytes = (uint64_t)n_hc * sizeof(float);
|
||||
if (row_bytes == 0 || out->bytes < row_bytes || out->bytes % row_bytes != 0 ||
|
||||
pre->bytes < out->bytes ||
|
||||
scale_offset > model_size || sizeof(float) > model_size - scale_offset ||
|
||||
base_offset > model_size || row_bytes > model_size - base_offset) {
|
||||
return 0;
|
||||
}
|
||||
const uint64_t n_tokens = out->bytes / row_bytes;
|
||||
const float *scale = (const float *)cuda_model_range_ptr(model_map, scale_offset, sizeof(float), "output_hc_scale");
|
||||
const float *base = (const float *)cuda_model_range_ptr(model_map, base_offset, row_bytes, "output_hc_base");
|
||||
if (!scale || !base) return 0;
|
||||
uint64_t n = n_tokens * n_hc;
|
||||
output_hc_weights_kernel<<<(n + 255) / 256, 256>>>(
|
||||
(float *)out->ptr,
|
||||
(const float *)pre->ptr,
|
||||
scale,
|
||||
base,
|
||||
n_hc,
|
||||
(uint32_t)n_tokens,
|
||||
eps);
|
||||
return cuda_ok(cudaGetLastError(), "output hc weights launch");
|
||||
}
|
||||
extern "C" int ds4_gpu_hc_expand_tensor(ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *post, const ds4_gpu_tensor *comb, uint32_t n_embd, uint32_t n_hc) {
|
||||
uint64_t n_tokens64 = 0, flat_bytes = 0, hc_bytes = 0, post_bytes = 0, comb_bytes = 0, comb_stride = 0;
|
||||
if (!out_hc || !block_out || !residual_hc || !post || !comb ||
|
||||
!cuda_hc_hc_token_count(out_hc, n_embd, n_hc, &n_tokens64) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, n_embd, sizeof(float), &flat_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, (uint64_t)n_hc * n_embd, sizeof(float), &hc_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, n_hc, sizeof(float), &post_bytes) ||
|
||||
!cuda_u64_mul_checked(n_hc, n_hc, &comb_stride) || comb_stride > UINT32_MAX ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, comb_stride, sizeof(float), &comb_bytes) ||
|
||||
block_out->bytes < flat_bytes || residual_hc->bytes < hc_bytes ||
|
||||
post->bytes < post_bytes || comb->bytes < comb_bytes) return 0;
|
||||
uint32_t n_tokens = (uint32_t)n_tokens64;
|
||||
uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd;
|
||||
hc_expand_kernel<<<(n_elem + 255) / 256, 256>>>((float *)out_hc->ptr,
|
||||
(const float *)block_out->ptr,
|
||||
(const float *)block_out->ptr,
|
||||
(const float *)residual_hc->ptr,
|
||||
(const float *)post->ptr,
|
||||
(const float *)comb->ptr,
|
||||
n_embd, n_hc, n_tokens,
|
||||
n_hc, (uint32_t)comb_stride, 0);
|
||||
return cuda_ok(cudaGetLastError(), "hc_expand launch");
|
||||
}
|
||||
extern "C" int ds4_gpu_hc_expand_split_tensor(ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *split, uint32_t n_embd, uint32_t n_hc) {
|
||||
uint64_t n_tokens64 = 0, flat_bytes = 0, hc_bytes = 0, split_bytes = 0, mix_hc64 = 0;
|
||||
if (!out_hc || !block_out || !residual_hc || !split ||
|
||||
!cuda_hc_hc_token_count(out_hc, n_embd, n_hc, &n_tokens64) ||
|
||||
!cuda_hc_mix_width(n_hc, &mix_hc64) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, n_embd, sizeof(float), &flat_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, (uint64_t)n_hc * n_embd, sizeof(float), &hc_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, mix_hc64, sizeof(float), &split_bytes) ||
|
||||
block_out->bytes < flat_bytes || residual_hc->bytes < hc_bytes || split->bytes < split_bytes) return 0;
|
||||
uint32_t n_tokens = (uint32_t)n_tokens64;
|
||||
if (n_hc == 4u) {
|
||||
const uint64_t n = (uint64_t)n_tokens * n_embd;
|
||||
hc_expand4_kernel<<<(n + 255) / 256, 256>>>((float *)out_hc->ptr,
|
||||
(const float *)block_out->ptr,
|
||||
(const float *)residual_hc->ptr,
|
||||
(const float *)split->ptr,
|
||||
n_embd,
|
||||
n_tokens);
|
||||
return cuda_ok(cudaGetLastError(), "hc_expand_split4 launch");
|
||||
}
|
||||
uint32_t mix_hc = (uint32_t)mix_hc64;
|
||||
uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd;
|
||||
const float *base = (const float *)split->ptr;
|
||||
hc_expand_kernel<<<(n_elem + 255) / 256, 256>>>((float *)out_hc->ptr,
|
||||
(const float *)block_out->ptr,
|
||||
(const float *)block_out->ptr,
|
||||
(const float *)residual_hc->ptr,
|
||||
base + n_hc,
|
||||
base + 2u * n_hc,
|
||||
n_embd, n_hc, n_tokens,
|
||||
mix_hc, mix_hc, 0);
|
||||
return cuda_ok(cudaGetLastError(), "hc_expand_split launch");
|
||||
}
|
||||
extern "C" int ds4_gpu_hc_expand_split_half_tensor(ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out_h, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *split, uint32_t n_embd, uint32_t n_hc) {
|
||||
uint64_t n_tokens64 = 0, flat_half_bytes = 0, hc_bytes = 0, split_bytes = 0, mix_hc64 = 0;
|
||||
if (!out_hc || !block_out_h || !residual_hc || !split ||
|
||||
!cuda_hc_hc_token_count(out_hc, n_embd, n_hc, &n_tokens64) ||
|
||||
!cuda_hc_mix_width(n_hc, &mix_hc64) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, n_embd, sizeof(__half), &flat_half_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, (uint64_t)n_hc * n_embd, sizeof(float), &hc_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, mix_hc64, sizeof(float), &split_bytes) ||
|
||||
block_out_h->bytes < flat_half_bytes || residual_hc->bytes < hc_bytes || split->bytes < split_bytes) return 0;
|
||||
uint32_t n_tokens = (uint32_t)n_tokens64;
|
||||
if (n_hc == 4u) {
|
||||
const uint64_t n = (uint64_t)n_tokens * n_embd;
|
||||
hc_expand4_half_kernel<<<(n + 255) / 256, 256>>>((float *)out_hc->ptr,
|
||||
(const __half *)block_out_h->ptr,
|
||||
(const float *)residual_hc->ptr,
|
||||
(const float *)split->ptr,
|
||||
n_embd,
|
||||
n_tokens);
|
||||
return cuda_ok(cudaGetLastError(), "hc_expand_split_half4 launch");
|
||||
}
|
||||
uint32_t mix_hc = (uint32_t)mix_hc64;
|
||||
uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd;
|
||||
const float *base = (const float *)split->ptr;
|
||||
hc_expand_half_kernel<<<(n_elem + 255) / 256, 256>>>((float *)out_hc->ptr,
|
||||
(const __half *)block_out_h->ptr,
|
||||
(const float *)residual_hc->ptr,
|
||||
base + n_hc,
|
||||
base + 2u * n_hc,
|
||||
n_embd, n_hc, n_tokens,
|
||||
mix_hc, mix_hc);
|
||||
return cuda_ok(cudaGetLastError(), "hc_expand_split_half launch");
|
||||
}
|
||||
extern "C" int ds4_gpu_hc_expand_add_split_tensor(ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out, const ds4_gpu_tensor *block_add, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *split, uint32_t n_embd, uint32_t n_hc) {
|
||||
uint64_t n_tokens64 = 0, flat_bytes = 0, hc_bytes = 0, split_bytes = 0, mix_hc64 = 0;
|
||||
if (!out_hc || !block_out || !block_add || !residual_hc || !split ||
|
||||
!cuda_hc_hc_token_count(out_hc, n_embd, n_hc, &n_tokens64) ||
|
||||
!cuda_hc_mix_width(n_hc, &mix_hc64) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, n_embd, sizeof(float), &flat_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, (uint64_t)n_hc * n_embd, sizeof(float), &hc_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, mix_hc64, sizeof(float), &split_bytes) ||
|
||||
block_out->bytes < flat_bytes || block_add->bytes < flat_bytes ||
|
||||
residual_hc->bytes < hc_bytes || split->bytes < split_bytes) return 0;
|
||||
uint32_t n_tokens = (uint32_t)n_tokens64;
|
||||
if (n_hc == 4u) {
|
||||
const uint64_t n = (uint64_t)n_tokens * n_embd;
|
||||
hc_expand4_add_kernel<<<(n + 255) / 256, 256>>>((float *)out_hc->ptr,
|
||||
(const float *)block_out->ptr,
|
||||
(const float *)block_add->ptr,
|
||||
(const float *)residual_hc->ptr,
|
||||
(const float *)split->ptr,
|
||||
n_embd,
|
||||
n_tokens);
|
||||
return cuda_ok(cudaGetLastError(), "hc_expand_add_split4 launch");
|
||||
}
|
||||
uint32_t mix_hc = (uint32_t)mix_hc64;
|
||||
uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd;
|
||||
const float *base = (const float *)split->ptr;
|
||||
hc_expand_kernel<<<(n_elem + 255) / 256, 256>>>((float *)out_hc->ptr,
|
||||
(const float *)block_out->ptr,
|
||||
(const float *)block_add->ptr,
|
||||
(const float *)residual_hc->ptr,
|
||||
base + n_hc,
|
||||
base + 2u * n_hc,
|
||||
n_embd, n_hc, n_tokens,
|
||||
mix_hc, mix_hc, 1);
|
||||
return cuda_ok(cudaGetLastError(), "hc_expand_add_split launch");
|
||||
}
|
||||
extern "C" int ds4_gpu_hc_expand_add_split_half_add_tensor(ds4_gpu_tensor *out_hc, const ds4_gpu_tensor *block_out, const ds4_gpu_tensor *block_add_h, const ds4_gpu_tensor *residual_hc, const ds4_gpu_tensor *split, uint32_t n_embd, uint32_t n_hc) {
|
||||
uint64_t n_tokens64 = 0, flat_bytes = 0, flat_half_bytes = 0, hc_bytes = 0, split_bytes = 0, mix_hc64 = 0;
|
||||
if (!out_hc || !block_out || !block_add_h || !residual_hc || !split ||
|
||||
!cuda_hc_hc_token_count(out_hc, n_embd, n_hc, &n_tokens64) ||
|
||||
!cuda_hc_mix_width(n_hc, &mix_hc64) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, n_embd, sizeof(float), &flat_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, n_embd, sizeof(__half), &flat_half_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, (uint64_t)n_hc * n_embd, sizeof(float), &hc_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tokens64, mix_hc64, sizeof(float), &split_bytes) ||
|
||||
block_out->bytes < flat_bytes || block_add_h->bytes < flat_half_bytes ||
|
||||
residual_hc->bytes < hc_bytes || split->bytes < split_bytes) return 0;
|
||||
uint32_t n_tokens = (uint32_t)n_tokens64;
|
||||
if (n_hc == 4u) {
|
||||
const uint64_t n = (uint64_t)n_tokens * n_embd;
|
||||
hc_expand4_add_half_kernel<<<(n + 255) / 256, 256>>>((float *)out_hc->ptr,
|
||||
(const float *)block_out->ptr,
|
||||
(const __half *)block_add_h->ptr,
|
||||
(const float *)residual_hc->ptr,
|
||||
(const float *)split->ptr,
|
||||
n_embd,
|
||||
n_tokens);
|
||||
return cuda_ok(cudaGetLastError(), "hc_expand_add_split_half4 launch");
|
||||
}
|
||||
uint32_t mix_hc = (uint32_t)mix_hc64;
|
||||
uint64_t n_elem = (uint64_t)n_tokens * n_hc * n_embd;
|
||||
const float *base = (const float *)split->ptr;
|
||||
hc_expand_add_half_kernel<<<(n_elem + 255) / 256, 256>>>((float *)out_hc->ptr,
|
||||
(const float *)block_out->ptr,
|
||||
(const __half *)block_add_h->ptr,
|
||||
(const float *)residual_hc->ptr,
|
||||
base + n_hc,
|
||||
base + 2u * n_hc,
|
||||
n_embd, n_hc, n_tokens,
|
||||
mix_hc, mix_hc);
|
||||
return cuda_ok(cudaGetLastError(), "hc_expand_add_split_half_add launch");
|
||||
}
|
||||
extern "C" int ds4_gpu_shared_down_hc_expand_q8_0_tensor(
|
||||
ds4_gpu_tensor *out_hc,
|
||||
ds4_gpu_tensor *shared_out,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t weight_offset,
|
||||
uint64_t in_dim,
|
||||
uint64_t out_dim,
|
||||
const ds4_gpu_tensor *shared_mid,
|
||||
const ds4_gpu_tensor *routed_out,
|
||||
const ds4_gpu_tensor *residual_hc,
|
||||
const ds4_gpu_tensor *split,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc) {
|
||||
return cuda_matmul_q8_0_hc_expand_tensor_labeled(out_hc, shared_out,
|
||||
model_map, model_size,
|
||||
weight_offset,
|
||||
in_dim, out_dim,
|
||||
shared_mid,
|
||||
routed_out,
|
||||
residual_hc,
|
||||
split,
|
||||
n_embd, n_hc,
|
||||
"shared_down_hc_expand");
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_matmul_q8_0_hc_expand_tensor(
|
||||
ds4_gpu_tensor *out_hc,
|
||||
ds4_gpu_tensor *block_out,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t weight_offset,
|
||||
uint64_t in_dim,
|
||||
uint64_t out_dim,
|
||||
const ds4_gpu_tensor *x,
|
||||
const ds4_gpu_tensor *residual_hc,
|
||||
const ds4_gpu_tensor *split,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc) {
|
||||
return cuda_matmul_q8_0_hc_expand_tensor_labeled(out_hc, block_out,
|
||||
model_map, model_size,
|
||||
weight_offset,
|
||||
in_dim, out_dim,
|
||||
x,
|
||||
NULL,
|
||||
residual_hc,
|
||||
split,
|
||||
n_embd, n_hc,
|
||||
"q8_hc_expand");
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
/* HIP-only hipBLASLt state and helpers.
|
||||
* Included from ds4_cuda.cu under __HIP_PLATFORM_AMD__ to keep ROCm
|
||||
* planning/cache code out of the CUDA host runtime body. */
|
||||
|
||||
static hipblasLtHandle_t g_hipblaslt;
|
||||
static int g_hipblaslt_ready;
|
||||
struct cuda_hipblaslt_gemm_plan {
|
||||
uint32_t out_dim;
|
||||
uint32_t n_tok;
|
||||
uint32_t in_dim;
|
||||
hipblasLtMatmulDesc_t desc;
|
||||
hipblasLtMatrixLayout_t a_desc;
|
||||
hipblasLtMatrixLayout_t b_desc;
|
||||
hipblasLtMatrixLayout_t c_desc;
|
||||
hipblasLtMatrixLayout_t d_desc;
|
||||
hipblasLtMatmulAlgo_t algo;
|
||||
};
|
||||
static std::vector<cuda_hipblaslt_gemm_plan> g_hipblaslt_gemm_plans;
|
||||
|
||||
static void hipblaslt_gemm_plan_clear(void) {
|
||||
for (size_t i = 0; i < g_hipblaslt_gemm_plans.size(); i++) {
|
||||
cuda_hipblaslt_gemm_plan &p = g_hipblaslt_gemm_plans[i];
|
||||
if (p.d_desc) (void)hipblasLtMatrixLayoutDestroy(p.d_desc);
|
||||
if (p.c_desc) (void)hipblasLtMatrixLayoutDestroy(p.c_desc);
|
||||
if (p.b_desc) (void)hipblasLtMatrixLayoutDestroy(p.b_desc);
|
||||
if (p.a_desc) (void)hipblasLtMatrixLayoutDestroy(p.a_desc);
|
||||
if (p.desc) (void)hipblasLtMatmulDescDestroy(p.desc);
|
||||
}
|
||||
g_hipblaslt_gemm_plans.clear();
|
||||
}
|
||||
|
||||
static int hipblaslt_ok(hipblasStatus_t st, const char *what) {
|
||||
if (st == HIPBLAS_STATUS_SUCCESS) return 1;
|
||||
fprintf(stderr, "ds4: hipBLASLt %s failed: status %d\n", what, (int)st);
|
||||
return 0;
|
||||
}
|
||||
|
||||
static cuda_hipblaslt_gemm_plan *hipblaslt_gemm_plan_get(
|
||||
uint32_t out_dim,
|
||||
uint32_t n_tok,
|
||||
uint32_t in_dim,
|
||||
const char *label) {
|
||||
for (size_t i = 0; i < g_hipblaslt_gemm_plans.size(); i++) {
|
||||
cuda_hipblaslt_gemm_plan &p = g_hipblaslt_gemm_plans[i];
|
||||
if (p.out_dim == out_dim && p.n_tok == n_tok && p.in_dim == in_dim) return &p;
|
||||
}
|
||||
|
||||
hipblasLtMatmulDesc_t desc = NULL;
|
||||
hipblasLtMatrixLayout_t a_desc = NULL, b_desc = NULL, c_desc = NULL, d_desc = NULL;
|
||||
hipblasLtMatmulPreference_t pref = NULL;
|
||||
hipblasLtMatmulHeuristicResult_t heur[8];
|
||||
int returned = 0;
|
||||
int ok = 0;
|
||||
do {
|
||||
if (!hipblaslt_ok(hipblasLtMatmulDescCreate(&desc, HIPBLAS_COMPUTE_32F, HIP_R_32F),
|
||||
"matmul desc create")) break;
|
||||
hipblasOperation_t op_a = HIPBLAS_OP_T;
|
||||
hipblasOperation_t op_b = HIPBLAS_OP_N;
|
||||
if (!hipblaslt_ok(hipblasLtMatmulDescSetAttribute(desc, HIPBLASLT_MATMUL_DESC_TRANSA,
|
||||
&op_a, sizeof(op_a)),
|
||||
"set transA")) break;
|
||||
if (!hipblaslt_ok(hipblasLtMatmulDescSetAttribute(desc, HIPBLASLT_MATMUL_DESC_TRANSB,
|
||||
&op_b, sizeof(op_b)),
|
||||
"set transB")) break;
|
||||
if (!hipblaslt_ok(hipblasLtMatrixLayoutCreate(&a_desc, HIP_R_16F, in_dim, out_dim, in_dim),
|
||||
"A layout create")) break;
|
||||
if (!hipblaslt_ok(hipblasLtMatrixLayoutCreate(&b_desc, HIP_R_16F, in_dim, n_tok, in_dim),
|
||||
"B layout create")) break;
|
||||
if (!hipblaslt_ok(hipblasLtMatrixLayoutCreate(&c_desc, HIP_R_16F, out_dim, n_tok, out_dim),
|
||||
"C layout create")) break;
|
||||
if (!hipblaslt_ok(hipblasLtMatrixLayoutCreate(&d_desc, HIP_R_16F, out_dim, n_tok, out_dim),
|
||||
"D layout create")) break;
|
||||
if (!hipblaslt_ok(hipblasLtMatmulPreferenceCreate(&pref), "preference create")) break;
|
||||
const size_t max_workspace = 0;
|
||||
if (!hipblaslt_ok(hipblasLtMatmulPreferenceSetAttribute(
|
||||
pref, HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
|
||||
&max_workspace, sizeof(max_workspace)),
|
||||
"set max workspace")) break;
|
||||
if (!hipblaslt_ok(hipblasLtMatmulAlgoGetHeuristic(g_hipblaslt, desc,
|
||||
a_desc, b_desc, c_desc, d_desc,
|
||||
pref, 8, heur, &returned),
|
||||
"algo heuristic")) break;
|
||||
if (returned <= 0 || heur[0].state != HIPBLAS_STATUS_SUCCESS) {
|
||||
fprintf(stderr, "ds4: hipBLASLt no algo for %s m=%u n=%u k=%u\n",
|
||||
label ? label : "gemm", out_dim, n_tok, in_dim);
|
||||
break;
|
||||
}
|
||||
ok = 1;
|
||||
} while (0);
|
||||
if (pref) (void)hipblasLtMatmulPreferenceDestroy(pref);
|
||||
if (!ok) {
|
||||
if (d_desc) (void)hipblasLtMatrixLayoutDestroy(d_desc);
|
||||
if (c_desc) (void)hipblasLtMatrixLayoutDestroy(c_desc);
|
||||
if (b_desc) (void)hipblasLtMatrixLayoutDestroy(b_desc);
|
||||
if (a_desc) (void)hipblasLtMatrixLayoutDestroy(a_desc);
|
||||
if (desc) (void)hipblasLtMatmulDescDestroy(desc);
|
||||
return NULL;
|
||||
}
|
||||
|
||||
cuda_hipblaslt_gemm_plan p;
|
||||
p.out_dim = out_dim;
|
||||
p.n_tok = n_tok;
|
||||
p.in_dim = in_dim;
|
||||
p.desc = desc;
|
||||
p.a_desc = a_desc;
|
||||
p.b_desc = b_desc;
|
||||
p.c_desc = c_desc;
|
||||
p.d_desc = d_desc;
|
||||
p.algo = heur[0].algo;
|
||||
g_hipblaslt_gemm_plans.push_back(p);
|
||||
return &g_hipblaslt_gemm_plans.back();
|
||||
}
|
||||
|
||||
static int hipblaslt_gemm_tn_f16_out_f16(
|
||||
__half *out,
|
||||
const __half *w_rowmajor_out_in,
|
||||
const __half *x_rowmajor_tok_in,
|
||||
uint32_t out_dim,
|
||||
uint32_t n_tok,
|
||||
uint32_t in_dim,
|
||||
const char *label) {
|
||||
if (!g_hipblaslt_ready || !out || !w_rowmajor_out_in || !x_rowmajor_tok_in ||
|
||||
out_dim == 0 || n_tok == 0 || in_dim == 0) return 0;
|
||||
cuda_hipblaslt_gemm_plan *p = hipblaslt_gemm_plan_get(out_dim, n_tok, in_dim, label);
|
||||
if (!p) return 0;
|
||||
const float alpha = 1.0f;
|
||||
const float beta = 0.0f;
|
||||
return hipblaslt_ok(hipblasLtMatmul(g_hipblaslt, p->desc, &alpha,
|
||||
w_rowmajor_out_in, p->a_desc,
|
||||
x_rowmajor_tok_in, p->b_desc,
|
||||
&beta,
|
||||
out, p->c_desc,
|
||||
out, p->d_desc,
|
||||
&p->algo,
|
||||
NULL, 0, 0),
|
||||
label ? label : "gemm");
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,791 @@
|
||||
template <uint32_t BT>
|
||||
static void cuda_launch_q8_batch_sharedx_bt(
|
||||
float *out,
|
||||
const unsigned char *w,
|
||||
const float *x,
|
||||
uint32_t n_blocks,
|
||||
uint32_t out_dim,
|
||||
uint32_t n_tok,
|
||||
uint64_t row_bytes,
|
||||
dim3 grid,
|
||||
uint32_t rows_per_block,
|
||||
uint32_t tile) {
|
||||
const size_t shmem = (size_t)tile * BT * 32u * sizeof(float);
|
||||
if (tile == 2u) {
|
||||
matmul_q8_0_f32_batch_sharedx_warp_rows_w32_toktile_kernel<2u, BT><<<grid, rows_per_block * 32u, shmem>>>(out, w, x, n_blocks, out_dim, n_tok, row_bytes);
|
||||
} else if (tile == 4u) {
|
||||
matmul_q8_0_f32_batch_sharedx_warp_rows_w32_toktile_kernel<4u, BT><<<grid, rows_per_block * 32u, shmem>>>(out, w, x, n_blocks, out_dim, n_tok, row_bytes);
|
||||
} else if (tile == 8u) {
|
||||
matmul_q8_0_f32_batch_sharedx_warp_rows_w32_toktile_kernel<8u, BT><<<grid, rows_per_block * 32u, shmem>>>(out, w, x, n_blocks, out_dim, n_tok, row_bytes);
|
||||
} else if (tile == 16u) {
|
||||
matmul_q8_0_f32_batch_sharedx_warp_rows_w32_toktile_kernel<16u, BT><<<grid, rows_per_block * 32u, shmem>>>(out, w, x, n_blocks, out_dim, n_tok, row_bytes);
|
||||
} else {
|
||||
matmul_q8_0_f32_batch_sharedx_warp_rows_w32_toktile_kernel<32u, BT><<<grid, rows_per_block * 32u, shmem>>>(out, w, x, n_blocks, out_dim, n_tok, row_bytes);
|
||||
}
|
||||
}
|
||||
|
||||
static void cuda_launch_q8_batch_sharedx(
|
||||
float *out,
|
||||
const unsigned char *w,
|
||||
const float *x,
|
||||
uint32_t n_blocks,
|
||||
uint32_t out_dim,
|
||||
uint32_t n_tok,
|
||||
uint64_t row_bytes,
|
||||
uint32_t rows_per_block,
|
||||
uint32_t tile,
|
||||
uint32_t block_tile) {
|
||||
const dim3 grid((out_dim + rows_per_block - 1u) / rows_per_block,
|
||||
(n_tok + tile - 1u) / tile,
|
||||
1u);
|
||||
if (block_tile == 8u) {
|
||||
cuda_launch_q8_batch_sharedx_bt<8u>(out, w, x, n_blocks, out_dim, n_tok, row_bytes, grid, rows_per_block, tile);
|
||||
} else if (block_tile == 32u) {
|
||||
cuda_launch_q8_batch_sharedx_bt<32u>(out, w, x, n_blocks, out_dim, n_tok, row_bytes, grid, rows_per_block, tile);
|
||||
} else {
|
||||
cuda_launch_q8_batch_sharedx_bt<16u>(out, w, x, n_blocks, out_dim, n_tok, row_bytes, grid, rows_per_block, tile);
|
||||
}
|
||||
}
|
||||
|
||||
template <uint32_t BT>
|
||||
static void cuda_launch_grouped_q8_a_sharedx_bt(
|
||||
float *low,
|
||||
const unsigned char *w,
|
||||
const float *heads,
|
||||
uint32_t n_tokens,
|
||||
uint32_t n_groups,
|
||||
uint32_t n_blocks,
|
||||
uint32_t rank,
|
||||
uint64_t row_bytes,
|
||||
dim3 grid,
|
||||
uint32_t rows_per_block,
|
||||
uint32_t tile) {
|
||||
const size_t shmem = (size_t)tile * BT * 32u * sizeof(float);
|
||||
if (tile == 2u) {
|
||||
grouped_q8_0_a_f32_batch_sharedx_chunked_w32_kernel<2u, BT><<<grid, rows_per_block * 32u, shmem>>>(low, w, heads, n_tokens, n_groups, n_blocks, rank, row_bytes);
|
||||
} else if (tile == 4u) {
|
||||
grouped_q8_0_a_f32_batch_sharedx_chunked_w32_kernel<4u, BT><<<grid, rows_per_block * 32u, shmem>>>(low, w, heads, n_tokens, n_groups, n_blocks, rank, row_bytes);
|
||||
} else if (tile == 8u) {
|
||||
grouped_q8_0_a_f32_batch_sharedx_chunked_w32_kernel<8u, BT><<<grid, rows_per_block * 32u, shmem>>>(low, w, heads, n_tokens, n_groups, n_blocks, rank, row_bytes);
|
||||
} else if (tile == 16u) {
|
||||
grouped_q8_0_a_f32_batch_sharedx_chunked_w32_kernel<16u, BT><<<grid, rows_per_block * 32u, shmem>>>(low, w, heads, n_tokens, n_groups, n_blocks, rank, row_bytes);
|
||||
} else {
|
||||
grouped_q8_0_a_f32_batch_sharedx_chunked_w32_kernel<32u, BT><<<grid, rows_per_block * 32u, shmem>>>(low, w, heads, n_tokens, n_groups, n_blocks, rank, row_bytes);
|
||||
}
|
||||
}
|
||||
|
||||
static void cuda_launch_grouped_q8_a_sharedx(
|
||||
float *low,
|
||||
const unsigned char *w,
|
||||
const float *heads,
|
||||
uint32_t n_tokens,
|
||||
uint32_t n_groups,
|
||||
uint32_t n_blocks,
|
||||
uint32_t rank,
|
||||
uint64_t row_bytes,
|
||||
uint32_t rows_per_block,
|
||||
uint32_t tile,
|
||||
uint32_t block_tile) {
|
||||
const uint32_t row_blocks = (rank + rows_per_block - 1u) / rows_per_block;
|
||||
const dim3 grid(n_groups * row_blocks,
|
||||
(n_tokens + tile - 1u) / tile,
|
||||
1u);
|
||||
if (block_tile == 8u) {
|
||||
cuda_launch_grouped_q8_a_sharedx_bt<8u>(low, w, heads, n_tokens, n_groups, n_blocks, rank, row_bytes, grid, rows_per_block, tile);
|
||||
} else if (block_tile == 32u) {
|
||||
cuda_launch_grouped_q8_a_sharedx_bt<32u>(low, w, heads, n_tokens, n_groups, n_blocks, rank, row_bytes, grid, rows_per_block, tile);
|
||||
} else {
|
||||
cuda_launch_grouped_q8_a_sharedx_bt<16u>(low, w, heads, n_tokens, n_groups, n_blocks, rank, row_bytes, grid, rows_per_block, tile);
|
||||
}
|
||||
}
|
||||
|
||||
static int cuda_matmul_q8_0_tensor_f16_gemm(
|
||||
ds4_gpu_tensor *out,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t weight_offset,
|
||||
uint64_t in_dim,
|
||||
uint64_t out_dim,
|
||||
const ds4_gpu_tensor *x,
|
||||
uint64_t n_tok,
|
||||
const char *label) {
|
||||
if (!g_cublas_ready || !out || !x || !model_map ||
|
||||
in_dim == 0u || out_dim == 0u || n_tok == 0u ||
|
||||
in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) return 0;
|
||||
const uint64_t blocks = (in_dim + 31u) / 32u;
|
||||
uint64_t row_bytes = 0, weight_bytes = 0, x_bytes = 0, out_bytes = 0;
|
||||
if (weight_offset > model_size ||
|
||||
!cuda_u64_mul_checked(blocks, 34u, &row_bytes) ||
|
||||
!cuda_u64_mul_checked(out_dim, row_bytes, &weight_bytes) ||
|
||||
weight_bytes > model_size - weight_offset ||
|
||||
!cuda_u64_mul3_checked(n_tok, in_dim, sizeof(float), &x_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tok, out_dim, sizeof(float), &out_bytes) ||
|
||||
x->bytes < x_bytes || out->bytes < out_bytes) return 0;
|
||||
const __half *w_f16 = cuda_q8_f16_ptr(model_map, weight_offset, weight_bytes, in_dim, out_dim, label);
|
||||
if (!w_f16) return 0;
|
||||
const uint64_t xh_count = n_tok * in_dim;
|
||||
__half *xh = (__half *)cuda_tmp_alloc(xh_count * sizeof(__half), "q8 f16 gemm activations");
|
||||
if (!xh) return 0;
|
||||
f32_to_f16_kernel<<<(xh_count + 255u) / 256u, 256>>>(xh, (const float *)x->ptr, xh_count);
|
||||
if (!cuda_ok(cudaGetLastError(), "q8 f16 activation convert launch")) return 0;
|
||||
const float alpha = 1.0f;
|
||||
const float beta = 0.0f;
|
||||
cublasStatus_t st = cublasGemmEx(g_cublas,
|
||||
CUBLAS_OP_T,
|
||||
CUBLAS_OP_N,
|
||||
(int)out_dim,
|
||||
(int)n_tok,
|
||||
(int)in_dim,
|
||||
&alpha,
|
||||
w_f16,
|
||||
CUDA_R_16F,
|
||||
(int)in_dim,
|
||||
xh,
|
||||
CUDA_R_16F,
|
||||
(int)in_dim,
|
||||
&beta,
|
||||
out->ptr,
|
||||
CUDA_R_32F,
|
||||
(int)out_dim,
|
||||
CUBLAS_COMPUTE_32F,
|
||||
CUBLAS_GEMM_DEFAULT);
|
||||
if (st == CUBLAS_STATUS_SUCCESS) return 1;
|
||||
fprintf(stderr, "ds4: " DS4_GPU_BLAS_NAME " q8 f16 matmul failed: status %d\n", (int)st);
|
||||
cuda_q8_f16_cache_disable_after_failure(DS4_GPU_BLAS_NAME " f16 matmul failure",
|
||||
in_dim * out_dim * sizeof(__half));
|
||||
return 0;
|
||||
}
|
||||
|
||||
static int cuda_matmul_q8_0_tensor_f16_gemm_out_half(
|
||||
ds4_gpu_tensor *out_h,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t weight_offset,
|
||||
uint64_t in_dim,
|
||||
uint64_t out_dim,
|
||||
const ds4_gpu_tensor *x,
|
||||
uint64_t n_tok,
|
||||
const char *label) {
|
||||
if (!g_cublas_ready || !out_h || !x || !model_map ||
|
||||
in_dim == 0u || out_dim == 0u || n_tok == 0u ||
|
||||
in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) return 0;
|
||||
const uint64_t blocks = (in_dim + 31u) / 32u;
|
||||
uint64_t row_bytes = 0, weight_bytes = 0, x_bytes = 0, out_bytes = 0;
|
||||
if (weight_offset > model_size ||
|
||||
!cuda_u64_mul_checked(blocks, 34u, &row_bytes) ||
|
||||
!cuda_u64_mul_checked(out_dim, row_bytes, &weight_bytes) ||
|
||||
weight_bytes > model_size - weight_offset ||
|
||||
!cuda_u64_mul3_checked(n_tok, in_dim, sizeof(float), &x_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tok, out_dim, sizeof(__half), &out_bytes) ||
|
||||
x->bytes < x_bytes || out_h->bytes < out_bytes) return 0;
|
||||
const __half *w_f16 = cuda_q8_f16_ptr(model_map, weight_offset, weight_bytes, in_dim, out_dim, label);
|
||||
if (!w_f16) return 0;
|
||||
const uint64_t xh_count = n_tok * in_dim;
|
||||
__half *xh = (__half *)cuda_tmp_alloc(xh_count * sizeof(__half), "q8 f16-out gemm activations");
|
||||
if (!xh) return 0;
|
||||
f32_to_f16_kernel<<<(xh_count + 255u) / 256u, 256>>>(xh, (const float *)x->ptr, xh_count);
|
||||
if (!cuda_ok(cudaGetLastError(), "q8 f16-out activation convert launch")) return 0;
|
||||
const float alpha = 1.0f;
|
||||
const float beta = 0.0f;
|
||||
cublasStatus_t st = cublasGemmEx(g_cublas,
|
||||
CUBLAS_OP_T,
|
||||
CUBLAS_OP_N,
|
||||
(int)out_dim,
|
||||
(int)n_tok,
|
||||
(int)in_dim,
|
||||
&alpha,
|
||||
w_f16,
|
||||
CUDA_R_16F,
|
||||
(int)in_dim,
|
||||
xh,
|
||||
CUDA_R_16F,
|
||||
(int)in_dim,
|
||||
&beta,
|
||||
out_h->ptr,
|
||||
CUDA_R_16F,
|
||||
(int)out_dim,
|
||||
CUBLAS_COMPUTE_32F,
|
||||
CUBLAS_GEMM_DEFAULT);
|
||||
if (st == CUBLAS_STATUS_SUCCESS) return 1;
|
||||
fprintf(stderr, "ds4: " DS4_GPU_BLAS_NAME " q8 f16-out matmul failed: status %d\n", (int)st);
|
||||
cuda_q8_f16_cache_disable_after_failure(DS4_GPU_BLAS_NAME " f16-out matmul failure",
|
||||
in_dim * out_dim * sizeof(__half));
|
||||
return 0;
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_matmul_q8_0_f16_out_tensor(
|
||||
ds4_gpu_tensor *out_h,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t weight_offset,
|
||||
uint64_t in_dim,
|
||||
uint64_t out_dim,
|
||||
const ds4_gpu_tensor *x,
|
||||
uint64_t n_tok) {
|
||||
return cuda_matmul_q8_0_tensor_f16_gemm_out_half(out_h, model_map, model_size,
|
||||
weight_offset, in_dim, out_dim,
|
||||
x, n_tok, "q8_f16_out");
|
||||
}
|
||||
|
||||
static int cuda_matmul_q8_0_tensor_labeled(ds4_gpu_tensor *out, const void *model_map, uint64_t model_size, uint64_t weight_offset, uint64_t in_dim, uint64_t out_dim, const ds4_gpu_tensor *x, uint64_t n_tok, const char *label) {
|
||||
if (!out || !x || !model_map ||
|
||||
in_dim == 0u || out_dim == 0u || n_tok == 0u ||
|
||||
in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) return 0;
|
||||
uint64_t blocks = (in_dim + 31u) / 32u;
|
||||
uint64_t row_bytes = 0, weight_bytes = 0, x_bytes = 0, out_bytes = 0;
|
||||
if (weight_offset > model_size ||
|
||||
!cuda_u64_mul_checked(blocks, 34u, &row_bytes) ||
|
||||
!cuda_u64_mul_checked(out_dim, row_bytes, &weight_bytes) ||
|
||||
weight_bytes > model_size - weight_offset ||
|
||||
!cuda_u64_mul3_checked(n_tok, in_dim, sizeof(float), &x_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tok, out_dim, sizeof(float), &out_bytes) ||
|
||||
x->bytes < x_bytes || out->bytes < out_bytes) return 0;
|
||||
if (n_tok > 1 && !g_quality_mode &&
|
||||
cuda_runtime_config()->shared_down_cublas && in_dim == 2048u && out_dim == 4096u &&
|
||||
cuda_matmul_q8_0_tensor_f16_gemm(out, model_map, model_size, weight_offset,
|
||||
in_dim, out_dim, x, n_tok, label ? label : "shared_expert")) {
|
||||
return 1;
|
||||
}
|
||||
const char *wptr = cuda_model_range_ptr(model_map, weight_offset, weight_bytes, "q8_0");
|
||||
if (!wptr) return 0;
|
||||
if (n_tok == 1 && !cuda_runtime_config()->q8_prequant_decode) {
|
||||
if ((in_dim & 31u) == 0u && in_dim <= 8192u) {
|
||||
const unsigned rows_per_block = 32u;
|
||||
const unsigned threads = rows_per_block * 32u;
|
||||
matmul_q8_0_f32_sharedx_warp_rows_w32_kernel<<<
|
||||
(unsigned)((out_dim + rows_per_block - 1u) / rows_per_block),
|
||||
threads,
|
||||
(size_t)in_dim * sizeof(float)>>>(
|
||||
(float *)out->ptr,
|
||||
reinterpret_cast<const unsigned char *>(wptr),
|
||||
(const float *)x->ptr,
|
||||
(uint32_t)blocks,
|
||||
out_dim,
|
||||
blocks * 34u);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_q8_0 f32 sharedx launch");
|
||||
}
|
||||
matmul_q8_0_f32_warp8_kernel<<<((unsigned)out_dim + 7u) / 8u, 256>>>(
|
||||
(float *)out->ptr,
|
||||
reinterpret_cast<const unsigned char *>(wptr),
|
||||
(const float *)x->ptr,
|
||||
in_dim,
|
||||
out_dim,
|
||||
blocks);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_q8_0 f32 warp launch");
|
||||
}
|
||||
if (n_tok > 1) {
|
||||
#if defined(__HIP_PLATFORM_AMD__) || defined(__HIPCC__)
|
||||
if (!g_quality_mode && (in_dim % 32u) == 0u &&
|
||||
out_dim >= 1024u &&
|
||||
n_tok >= 256u &&
|
||||
in_dim <= UINT32_MAX && out_dim <= UINT32_MAX && n_tok <= UINT32_MAX) {
|
||||
const dim3 grid((uint32_t)((out_dim + 63u) / 64u),
|
||||
(uint32_t)((n_tok + 63u) / 64u),
|
||||
1u);
|
||||
matmul_q8_0_f32_batch_wmma_4w_kernel<<<grid, 128u>>>(
|
||||
(float *)out->ptr,
|
||||
reinterpret_cast<const unsigned char *>(wptr),
|
||||
(const float *)x->ptr,
|
||||
(uint32_t)n_tok,
|
||||
(uint32_t)in_dim,
|
||||
(uint32_t)out_dim,
|
||||
blocks * 34u);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_q8_0 f32 batch wmma 4w launch");
|
||||
}
|
||||
#endif
|
||||
if ((in_dim & 31u) == 0u && out_dim <= UINT32_MAX && n_tok <= UINT32_MAX) {
|
||||
const uint32_t rows_per_block = 32u;
|
||||
const uint32_t tile = 32u;
|
||||
const uint32_t block_tile = 16u;
|
||||
cuda_launch_q8_batch_sharedx((float *)out->ptr,
|
||||
reinterpret_cast<const unsigned char *>(wptr),
|
||||
(const float *)x->ptr,
|
||||
(uint32_t)blocks,
|
||||
(uint32_t)out_dim,
|
||||
(uint32_t)n_tok,
|
||||
blocks * 34u,
|
||||
rows_per_block,
|
||||
tile,
|
||||
block_tile);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_q8_0 f32 batch sharedx launch");
|
||||
}
|
||||
dim3 bgrid(((unsigned)out_dim + 7u) / 8u, (unsigned)n_tok, 1);
|
||||
matmul_q8_0_f32_batch_warp8_kernel<<<bgrid, 256>>>(
|
||||
(float *)out->ptr,
|
||||
reinterpret_cast<const unsigned char *>(wptr),
|
||||
(const float *)x->ptr,
|
||||
in_dim,
|
||||
out_dim,
|
||||
n_tok,
|
||||
blocks);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_q8_0 f32 batch warp launch");
|
||||
}
|
||||
if (g_cublas_ready && n_tok > 1) {
|
||||
const __half *w_f16 = cuda_q8_f16_ptr(model_map, weight_offset, weight_bytes, in_dim, out_dim, label);
|
||||
if (w_f16) {
|
||||
const uint64_t xh_count = n_tok * in_dim;
|
||||
__half *xh = (__half *)cuda_tmp_alloc(xh_count * sizeof(__half), "q8 f16 gemm activations");
|
||||
if (!xh) return 0;
|
||||
f32_to_f16_kernel<<<(xh_count + 255) / 256, 256>>>(xh, (const float *)x->ptr, xh_count);
|
||||
if (!cuda_ok(cudaGetLastError(), "q8 f16 activation convert launch")) return 0;
|
||||
const float alpha = 1.0f;
|
||||
const float beta = 0.0f;
|
||||
cublasStatus_t st = cublasGemmEx(g_cublas,
|
||||
CUBLAS_OP_T,
|
||||
CUBLAS_OP_N,
|
||||
(int)out_dim,
|
||||
(int)n_tok,
|
||||
(int)in_dim,
|
||||
&alpha,
|
||||
w_f16,
|
||||
CUDA_R_16F,
|
||||
(int)in_dim,
|
||||
xh,
|
||||
CUDA_R_16F,
|
||||
(int)in_dim,
|
||||
&beta,
|
||||
out->ptr,
|
||||
CUDA_R_32F,
|
||||
(int)out_dim,
|
||||
CUBLAS_COMPUTE_32F,
|
||||
CUBLAS_GEMM_DEFAULT);
|
||||
if (st == CUBLAS_STATUS_SUCCESS) return 1;
|
||||
fprintf(stderr, "ds4: " DS4_GPU_BLAS_NAME " q8 f16 matmul failed: status %d\n", (int)st);
|
||||
cuda_q8_f16_cache_disable_after_failure(DS4_GPU_BLAS_NAME " f16 matmul failure",
|
||||
in_dim * out_dim * sizeof(__half));
|
||||
/* The F16 expansion cache is only an optimization. If cuBLAS
|
||||
* rejects the cached path under memory pressure, retry the same
|
||||
* operation through the native Q8 kernels below. */
|
||||
}
|
||||
}
|
||||
const uint64_t xq_bytes = n_tok * blocks * 32u;
|
||||
const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull;
|
||||
const uint64_t tmp_bytes = scale_offset + n_tok * blocks * sizeof(float);
|
||||
void *tmp = cuda_tmp_alloc(tmp_bytes, "q8_0 prequant");
|
||||
if (!tmp) return 0;
|
||||
int8_t *xq = (int8_t *)tmp;
|
||||
float *xscale = (float *)((char *)tmp + scale_offset);
|
||||
const ds4_rocm_runtime_config *cfg = cuda_runtime_config();
|
||||
const int use_dp4a = 1;
|
||||
dim3 qgrid((unsigned)blocks, (unsigned)n_tok, 1);
|
||||
quantize_q8_0_f32_kernel<<<qgrid, 32>>>(xq, xscale, (const float *)x->ptr, in_dim, blocks);
|
||||
if (!cuda_ok(cudaGetLastError(), "matmul_q8_0 quantize launch")) return 0;
|
||||
if (n_tok == 1) {
|
||||
uint32_t rows_per_block = cfg->q8_decode_rpb;
|
||||
matmul_q8_0_preq_rows_w32_kernel<<<((unsigned)out_dim + rows_per_block - 1u) / rows_per_block,
|
||||
rows_per_block * 32u>>>(
|
||||
(float *)out->ptr,
|
||||
reinterpret_cast<const unsigned char *>(wptr),
|
||||
xq,
|
||||
xscale,
|
||||
in_dim,
|
||||
out_dim,
|
||||
blocks,
|
||||
rows_per_block,
|
||||
use_dp4a);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_q8_0 rows launch");
|
||||
}
|
||||
if (blocks <= 32u) {
|
||||
dim3 bgrid(((unsigned)out_dim + 7u) / 8u, (unsigned)n_tok, 1);
|
||||
matmul_q8_0_preq_batch_warp8_kernel<<<bgrid, 256>>>(
|
||||
(float *)out->ptr,
|
||||
reinterpret_cast<const unsigned char *>(wptr),
|
||||
xq,
|
||||
xscale,
|
||||
in_dim,
|
||||
out_dim,
|
||||
n_tok,
|
||||
blocks,
|
||||
use_dp4a);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_q8_0 batch warp launch");
|
||||
}
|
||||
dim3 grid((unsigned)out_dim, (unsigned)n_tok, 1);
|
||||
matmul_q8_0_preq_kernel<<<grid, 256>>>((float *)out->ptr,
|
||||
reinterpret_cast<const unsigned char *>(wptr),
|
||||
xq,
|
||||
xscale,
|
||||
in_dim, out_dim, n_tok, blocks,
|
||||
use_dp4a);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_q8_0 launch");
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_matmul_q8_0_tensor(ds4_gpu_tensor *out, const void *model_map, uint64_t model_size, uint64_t weight_offset, uint64_t in_dim, uint64_t out_dim, const ds4_gpu_tensor *x, uint64_t n_tok) {
|
||||
return cuda_matmul_q8_0_tensor_labeled(out, model_map, model_size, weight_offset,
|
||||
in_dim, out_dim, x, n_tok, "q8_0");
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_matmul_q8_0_pair_tensor(
|
||||
ds4_gpu_tensor *out0,
|
||||
ds4_gpu_tensor *out1,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t weight0_offset,
|
||||
uint64_t weight1_offset,
|
||||
uint64_t in_dim,
|
||||
uint64_t out0_dim,
|
||||
uint64_t out1_dim,
|
||||
const ds4_gpu_tensor *x,
|
||||
uint64_t n_tok) {
|
||||
if (!out0 || !out1 || !x || !model_map ||
|
||||
in_dim == 0 || out0_dim == 0 || out1_dim == 0 || n_tok == 0 ||
|
||||
in_dim > UINT32_MAX || out0_dim > UINT32_MAX || out1_dim > UINT32_MAX || n_tok > UINT32_MAX) {
|
||||
return 0;
|
||||
}
|
||||
if (n_tok != 1) {
|
||||
return cuda_matmul_q8_0_tensor_labeled(out0, model_map, model_size, weight0_offset,
|
||||
in_dim, out0_dim, x, n_tok, "q8_0_pair0") &&
|
||||
cuda_matmul_q8_0_tensor_labeled(out1, model_map, model_size, weight1_offset,
|
||||
in_dim, out1_dim, x, n_tok, "q8_0_pair1");
|
||||
}
|
||||
const uint64_t blocks = (in_dim + 31u) / 32u;
|
||||
uint64_t row_bytes = 0, weight0_bytes = 0, weight1_bytes = 0;
|
||||
if (weight0_offset > model_size || weight1_offset > model_size ||
|
||||
!cuda_u64_mul_checked(blocks, 34u, &row_bytes) ||
|
||||
!cuda_u64_mul_checked(out0_dim, row_bytes, &weight0_bytes) ||
|
||||
!cuda_u64_mul_checked(out1_dim, row_bytes, &weight1_bytes)) {
|
||||
return 0;
|
||||
}
|
||||
if (weight0_bytes > model_size - weight0_offset ||
|
||||
weight1_bytes > model_size - weight1_offset ||
|
||||
x->bytes < in_dim * sizeof(float) ||
|
||||
out0->bytes < out0_dim * sizeof(float) ||
|
||||
out1->bytes < out1_dim * sizeof(float)) {
|
||||
return 0;
|
||||
}
|
||||
const char *w0 = cuda_model_range_ptr(model_map, weight0_offset, weight0_bytes, "q8_0_pair0");
|
||||
const char *w1 = cuda_model_range_ptr(model_map, weight1_offset, weight1_bytes, "q8_0_pair1");
|
||||
if (!w0 || !w1) return 0;
|
||||
if (!cuda_runtime_config()->q8_prequant_decode) {
|
||||
const uint64_t max_out = out0_dim > out1_dim ? out0_dim : out1_dim;
|
||||
if ((in_dim & 31u) == 0u && in_dim <= 8192u) {
|
||||
const unsigned rows_per_block = 32u;
|
||||
const unsigned threads = rows_per_block * 32u;
|
||||
matmul_q8_0_pair_f32_sharedx_warp_rows_w32_kernel<<<
|
||||
(unsigned)((max_out + rows_per_block - 1u) / rows_per_block),
|
||||
threads,
|
||||
(size_t)in_dim * sizeof(float)>>>(
|
||||
(float *)out0->ptr,
|
||||
(float *)out1->ptr,
|
||||
reinterpret_cast<const unsigned char *>(w0),
|
||||
reinterpret_cast<const unsigned char *>(w1),
|
||||
(const float *)x->ptr,
|
||||
(uint32_t)blocks,
|
||||
out0_dim,
|
||||
out1_dim,
|
||||
blocks * 34u);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair f32 sharedx launch");
|
||||
}
|
||||
matmul_q8_0_pair_f32_warp8_kernel<<<((unsigned)max_out + 7u) / 8u, 256>>>(
|
||||
(float *)out0->ptr,
|
||||
(float *)out1->ptr,
|
||||
reinterpret_cast<const unsigned char *>(w0),
|
||||
reinterpret_cast<const unsigned char *>(w1),
|
||||
(const float *)x->ptr,
|
||||
in_dim,
|
||||
out0_dim,
|
||||
out1_dim,
|
||||
blocks);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair f32 warp launch");
|
||||
}
|
||||
|
||||
const uint64_t xq_bytes = blocks * 32u;
|
||||
const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull;
|
||||
const uint64_t tmp_bytes = scale_offset + blocks * sizeof(float);
|
||||
void *tmp = cuda_tmp_alloc(tmp_bytes, "q8_0 pair prequant");
|
||||
if (!tmp) return 0;
|
||||
int8_t *xq = (int8_t *)tmp;
|
||||
float *xscale = (float *)((char *)tmp + scale_offset);
|
||||
const int use_dp4a = 1;
|
||||
dim3 qgrid((unsigned)blocks, 1, 1);
|
||||
quantize_q8_0_f32_kernel<<<qgrid, 32>>>(xq, xscale, (const float *)x->ptr, in_dim, blocks);
|
||||
if (!cuda_ok(cudaGetLastError(), "matmul_q8_0 pair quantize launch")) return 0;
|
||||
const uint64_t max_out = out0_dim > out1_dim ? out0_dim : out1_dim;
|
||||
matmul_q8_0_pair_preq_warp8_kernel<<<((unsigned)max_out + 7u) / 8u, 256>>>(
|
||||
(float *)out0->ptr,
|
||||
(float *)out1->ptr,
|
||||
reinterpret_cast<const unsigned char *>(w0),
|
||||
reinterpret_cast<const unsigned char *>(w1),
|
||||
xq,
|
||||
xscale,
|
||||
in_dim,
|
||||
out0_dim,
|
||||
out1_dim,
|
||||
blocks,
|
||||
use_dp4a);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_q8_0 pair warp launch");
|
||||
}
|
||||
|
||||
static int cuda_matmul_q8_0_hc_expand_tensor_labeled(
|
||||
ds4_gpu_tensor *out_hc,
|
||||
ds4_gpu_tensor *block_out,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t weight_offset,
|
||||
uint64_t in_dim,
|
||||
uint64_t out_dim,
|
||||
const ds4_gpu_tensor *x,
|
||||
const ds4_gpu_tensor *block_add,
|
||||
const ds4_gpu_tensor *residual_hc,
|
||||
const ds4_gpu_tensor *split,
|
||||
uint32_t n_embd,
|
||||
uint32_t n_hc,
|
||||
const char *label) {
|
||||
if (!out_hc || !block_out || !x || !residual_hc || !split || !model_map ||
|
||||
in_dim == 0 || out_dim == 0 || n_embd == 0 || n_hc == 0 ||
|
||||
out_dim != (uint64_t)n_embd) {
|
||||
return 0;
|
||||
}
|
||||
const uint64_t blocks = (in_dim + 31) / 32;
|
||||
if (weight_offset > model_size || out_dim > UINT64_MAX / (blocks * 34)) return 0;
|
||||
const uint64_t weight_bytes = out_dim * blocks * 34;
|
||||
const uint64_t hc_bytes = (uint64_t)n_hc * n_embd * sizeof(float);
|
||||
const uint64_t split_bytes = (uint64_t)(2u * n_hc + n_hc * n_hc) * sizeof(float);
|
||||
if (weight_bytes > model_size - weight_offset ||
|
||||
x->bytes < in_dim * sizeof(float) ||
|
||||
block_out->bytes < out_dim * sizeof(float) ||
|
||||
residual_hc->bytes < hc_bytes ||
|
||||
split->bytes < split_bytes ||
|
||||
out_hc->bytes < hc_bytes ||
|
||||
(block_add && block_add->bytes < out_dim * sizeof(float))) {
|
||||
return 0;
|
||||
}
|
||||
const char *wptr = cuda_model_range_ptr(model_map, weight_offset, weight_bytes, label ? label : "q8_0_hc_expand");
|
||||
if (!wptr) return 0;
|
||||
if (!cuda_runtime_config()->q8_prequant_decode) {
|
||||
if ((in_dim & 31u) == 0u && in_dim <= 8192u) {
|
||||
const unsigned rows_per_block = 32u;
|
||||
const unsigned threads = rows_per_block * 32u;
|
||||
matmul_q8_0_hc_expand_f32_sharedx_warp_rows_w32_kernel<<<
|
||||
(unsigned)((out_dim + rows_per_block - 1u) / rows_per_block),
|
||||
threads,
|
||||
(size_t)in_dim * sizeof(float)>>>(
|
||||
(float *)out_hc->ptr,
|
||||
(float *)block_out->ptr,
|
||||
block_add ? (const float *)block_add->ptr : (const float *)block_out->ptr,
|
||||
(const float *)residual_hc->ptr,
|
||||
(const float *)split->ptr,
|
||||
reinterpret_cast<const unsigned char *>(wptr),
|
||||
(const float *)x->ptr,
|
||||
(uint32_t)blocks,
|
||||
out_dim,
|
||||
blocks * 34u,
|
||||
n_embd,
|
||||
n_hc,
|
||||
block_add ? 1 : 0);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_q8_0_hc_expand f32 sharedx launch");
|
||||
}
|
||||
matmul_q8_0_hc_expand_f32_warp8_kernel<<<((unsigned)out_dim + 7u) / 8u, 256>>>(
|
||||
(float *)out_hc->ptr,
|
||||
(float *)block_out->ptr,
|
||||
block_add ? (const float *)block_add->ptr : (const float *)block_out->ptr,
|
||||
(const float *)residual_hc->ptr,
|
||||
(const float *)split->ptr,
|
||||
reinterpret_cast<const unsigned char *>(wptr),
|
||||
(const float *)x->ptr,
|
||||
in_dim,
|
||||
out_dim,
|
||||
n_embd,
|
||||
n_hc,
|
||||
blocks,
|
||||
block_add ? 1 : 0);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_q8_0_hc_expand f32 launch");
|
||||
}
|
||||
|
||||
const uint64_t xq_bytes = blocks * 32u;
|
||||
const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull;
|
||||
const uint64_t tmp_bytes = scale_offset + blocks * sizeof(float);
|
||||
void *tmp = cuda_tmp_alloc(tmp_bytes, "q8_0 hc expand prequant");
|
||||
if (!tmp) return 0;
|
||||
int8_t *xq = (int8_t *)tmp;
|
||||
float *xscale = (float *)((char *)tmp + scale_offset);
|
||||
const ds4_rocm_runtime_config *cfg = cuda_runtime_config();
|
||||
const int use_dp4a = 1;
|
||||
quantize_q8_0_f32_kernel<<<(unsigned)blocks, 32>>>(xq, xscale, (const float *)x->ptr, in_dim, blocks);
|
||||
if (!cuda_ok(cudaGetLastError(), "matmul_q8_0_hc_expand quantize launch")) return 0;
|
||||
uint32_t rows_per_block = cfg->q8_hc_decode_rpb;
|
||||
matmul_q8_0_hc_expand_preq_rows_w32_kernel<<<((unsigned)out_dim + rows_per_block - 1u) / rows_per_block,
|
||||
rows_per_block * 32u>>>(
|
||||
(float *)out_hc->ptr,
|
||||
(float *)block_out->ptr,
|
||||
block_add ? (const float *)block_add->ptr : (const float *)block_out->ptr,
|
||||
(const float *)residual_hc->ptr,
|
||||
(const float *)split->ptr,
|
||||
reinterpret_cast<const unsigned char *>(wptr),
|
||||
xq,
|
||||
xscale,
|
||||
in_dim,
|
||||
out_dim,
|
||||
n_embd,
|
||||
n_hc,
|
||||
blocks,
|
||||
rows_per_block,
|
||||
block_add ? 1 : 0,
|
||||
use_dp4a);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_q8_0_hc_expand rows launch");
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_matmul_f16_tensor(ds4_gpu_tensor *out, const void *model_map, uint64_t model_size, uint64_t weight_offset, uint64_t in_dim, uint64_t out_dim, const ds4_gpu_tensor *x, uint64_t n_tok) {
|
||||
if (!out || !x || !model_map ||
|
||||
in_dim == 0u || out_dim == 0u || n_tok == 0u ||
|
||||
in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) return 0;
|
||||
uint64_t weight_bytes = 0, x_bytes = 0, out_bytes = 0;
|
||||
if (weight_offset > model_size ||
|
||||
!cuda_u64_mul3_checked(out_dim, in_dim, sizeof(uint16_t), &weight_bytes) ||
|
||||
weight_bytes > model_size - weight_offset ||
|
||||
!cuda_u64_mul3_checked(n_tok, in_dim, sizeof(float), &x_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tok, out_dim, sizeof(float), &out_bytes) ||
|
||||
x->bytes < x_bytes || out->bytes < out_bytes) return 0;
|
||||
const char *wptr = cuda_model_range_ptr(model_map, weight_offset, weight_bytes, "f16");
|
||||
if (!wptr) return 0;
|
||||
const __half *w = (const __half *)wptr;
|
||||
const int ordered_decode = n_tok == 1u;
|
||||
if (g_cublas_ready && n_tok > 1) {
|
||||
const uint64_t xh_count = n_tok * in_dim;
|
||||
__half *xh = (__half *)cuda_tmp_alloc(xh_count * sizeof(__half), "f16 gemm activations");
|
||||
if (!xh) return 0;
|
||||
f32_to_f16_kernel<<<(xh_count + 255) / 256, 256>>>(xh, (const float *)x->ptr, xh_count);
|
||||
if (!cuda_ok(cudaGetLastError(), "f16 activation convert launch")) return 0;
|
||||
const float alpha = 1.0f;
|
||||
const float beta = 0.0f;
|
||||
cublasStatus_t st = cublasGemmEx(g_cublas,
|
||||
CUBLAS_OP_T,
|
||||
CUBLAS_OP_N,
|
||||
(int)out_dim,
|
||||
(int)n_tok,
|
||||
(int)in_dim,
|
||||
&alpha,
|
||||
w,
|
||||
CUDA_R_16F,
|
||||
(int)in_dim,
|
||||
xh,
|
||||
CUDA_R_16F,
|
||||
(int)in_dim,
|
||||
&beta,
|
||||
out->ptr,
|
||||
CUDA_R_32F,
|
||||
(int)out_dim,
|
||||
CUBLAS_COMPUTE_32F,
|
||||
CUBLAS_GEMM_DEFAULT);
|
||||
return cublas_ok(st, "f16 matmul");
|
||||
}
|
||||
/* The 4096x256 F16 router projection is latency-bound and the ordered
|
||||
* 32-thread row kernel is at least as fast on gfx1151; keep shared-X for
|
||||
* compressor/indexer F16 decode where reusing x across rows is the win. */
|
||||
const bool f16_decode_router_shape = (in_dim == 4096u && out_dim == 256u);
|
||||
if (n_tok == 1u && !g_quality_mode && !cuda_runtime_config()->graph_dump &&
|
||||
!f16_decode_router_shape) {
|
||||
if (in_dim <= 8192u && in_dim * sizeof(float) <= 65536u) {
|
||||
const uint32_t rows_per_block = 32u;
|
||||
matmul_f16_f32_sharedx_warp_rows_w32_kernel<<<
|
||||
((unsigned)out_dim + rows_per_block - 1u) / rows_per_block,
|
||||
rows_per_block * 32u,
|
||||
(size_t)in_dim * sizeof(float)>>>(
|
||||
(float *)out->ptr, w, (const float *)x->ptr, (uint32_t)in_dim, out_dim);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_f16 sharedx launch");
|
||||
}
|
||||
}
|
||||
dim3 grid((unsigned)out_dim, (unsigned)n_tok, 1);
|
||||
if (ordered_decode) {
|
||||
matmul_f16_ordered_chunks_kernel<<<grid, 32>>>((float *)out->ptr, w, (const float *)x->ptr, in_dim, out_dim, n_tok);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_f16_ordered_chunks launch");
|
||||
}
|
||||
matmul_f16_kernel<<<grid, 256>>>((float *)out->ptr, w, (const float *)x->ptr, in_dim, out_dim, n_tok);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_f16 launch");
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_matmul_f16_pair_tensor(
|
||||
ds4_gpu_tensor *out0,
|
||||
ds4_gpu_tensor *out1,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t weight0_offset,
|
||||
uint64_t weight1_offset,
|
||||
uint64_t in_dim,
|
||||
uint64_t out_dim,
|
||||
const ds4_gpu_tensor *x,
|
||||
uint64_t n_tok) {
|
||||
if (!out0 || !out1 || !x || !model_map || in_dim == 0 || out_dim == 0 || n_tok == 0 ||
|
||||
in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) {
|
||||
return 0;
|
||||
}
|
||||
if (n_tok != 1) {
|
||||
return ds4_gpu_matmul_f16_tensor(out0, model_map, model_size, weight0_offset,
|
||||
in_dim, out_dim, x, n_tok) &&
|
||||
ds4_gpu_matmul_f16_tensor(out1, model_map, model_size, weight1_offset,
|
||||
in_dim, out_dim, x, n_tok);
|
||||
}
|
||||
uint64_t weight_bytes = 0;
|
||||
if (weight0_offset > model_size || weight1_offset > model_size ||
|
||||
!cuda_u64_mul3_checked(out_dim, in_dim, sizeof(uint16_t), &weight_bytes)) {
|
||||
return 0;
|
||||
}
|
||||
if (weight_bytes > model_size - weight0_offset ||
|
||||
weight_bytes > model_size - weight1_offset ||
|
||||
x->bytes < in_dim * sizeof(float) ||
|
||||
out0->bytes < out_dim * sizeof(float) ||
|
||||
out1->bytes < out_dim * sizeof(float)) {
|
||||
return 0;
|
||||
}
|
||||
const __half *w0 = (const __half *)cuda_model_range_ptr(model_map, weight0_offset, weight_bytes, "f16_pair0");
|
||||
const __half *w1 = (const __half *)cuda_model_range_ptr(model_map, weight1_offset, weight_bytes, "f16_pair1");
|
||||
if (!w0 || !w1) return 0;
|
||||
if (!g_quality_mode && !cuda_runtime_config()->graph_dump) {
|
||||
if (in_dim <= 8192u && in_dim * sizeof(float) <= 65536u) {
|
||||
const uint32_t rows_per_block = 32u;
|
||||
matmul_f16_pair_f32_sharedx_warp_rows_w32_kernel<<<
|
||||
((unsigned)out_dim + rows_per_block - 1u) / rows_per_block,
|
||||
rows_per_block * 32u,
|
||||
(size_t)in_dim * sizeof(float)>>>(
|
||||
(float *)out0->ptr, (float *)out1->ptr, w0, w1,
|
||||
(const float *)x->ptr, (uint32_t)in_dim, out_dim);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_f16_pair sharedx launch");
|
||||
}
|
||||
}
|
||||
matmul_f16_pair_ordered_chunks_kernel<<<(unsigned)out_dim, 32>>>(
|
||||
(float *)out0->ptr,
|
||||
(float *)out1->ptr,
|
||||
w0,
|
||||
w1,
|
||||
(const float *)x->ptr,
|
||||
in_dim,
|
||||
out_dim,
|
||||
out_dim);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_f16_pair_ordered_chunks launch");
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_matmul_f32_tensor(ds4_gpu_tensor *out, const void *model_map, uint64_t model_size, uint64_t weight_offset, uint64_t in_dim, uint64_t out_dim, const ds4_gpu_tensor *x, uint64_t n_tok) {
|
||||
if (!out || !x || !model_map || in_dim == 0 || out_dim == 0 || n_tok == 0 ||
|
||||
in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX) return 0;
|
||||
uint64_t weight_bytes = 0, x_bytes = 0, out_bytes = 0;
|
||||
if (weight_offset > model_size ||
|
||||
!cuda_u64_mul3_checked(out_dim, in_dim, sizeof(float), &weight_bytes) ||
|
||||
weight_bytes > model_size - weight_offset ||
|
||||
!cuda_u64_mul3_checked(n_tok, in_dim, sizeof(float), &x_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tok, out_dim, sizeof(float), &out_bytes) ||
|
||||
x->bytes < x_bytes || out->bytes < out_bytes) return 0;
|
||||
const char *wptr = cuda_model_range_ptr(model_map, weight_offset, weight_bytes, "f32");
|
||||
if (!wptr) return 0;
|
||||
const float *w = (const float *)wptr;
|
||||
if (g_cublas_ready && n_tok > 1) {
|
||||
const float alpha = 1.0f;
|
||||
const float beta = 0.0f;
|
||||
cublasStatus_t st = cublasSgemm(g_cublas,
|
||||
CUBLAS_OP_T,
|
||||
CUBLAS_OP_N,
|
||||
(int)out_dim,
|
||||
(int)n_tok,
|
||||
(int)in_dim,
|
||||
&alpha,
|
||||
w,
|
||||
(int)in_dim,
|
||||
(const float *)x->ptr,
|
||||
(int)in_dim,
|
||||
&beta,
|
||||
(float *)out->ptr,
|
||||
(int)out_dim);
|
||||
return cublas_ok(st, "f32 matmul");
|
||||
}
|
||||
dim3 grid((unsigned)out_dim, (unsigned)n_tok, 1);
|
||||
matmul_f32_kernel<<<grid, 256>>>((float *)out->ptr, w, (const float *)x->ptr, in_dim, out_dim, n_tok);
|
||||
return cuda_ok(cudaGetLastError(), "matmul_f32 launch");
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
extern "C" int ds4_gpu_add_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *a, const ds4_gpu_tensor *b, uint32_t n) {
|
||||
if (!cuda_tensor_has_f32(out, n) || !cuda_tensor_has_f32(a, n) || !cuda_tensor_has_f32(b, n)) return 0;
|
||||
if (n == 0u) return 1;
|
||||
add_kernel<<<(n + 255) / 256, 256>>>((float *)out->ptr, (const float *)a->ptr, (const float *)b->ptr, n);
|
||||
return cuda_ok(cudaGetLastError(), "add launch");
|
||||
}
|
||||
extern "C" int ds4_gpu_directional_steering_project_tensor(
|
||||
ds4_gpu_tensor *x,
|
||||
const ds4_gpu_tensor *directions,
|
||||
uint32_t layer,
|
||||
uint32_t width,
|
||||
uint32_t rows,
|
||||
float scale) {
|
||||
if (!x || !directions || width == 0 || rows == 0) return 0;
|
||||
uint64_t x_bytes = 0, dir_bytes = 0;
|
||||
if (!cuda_u64_mul3_checked(width, rows, sizeof(float), &x_bytes) ||
|
||||
!cuda_u64_mul3_checked((uint64_t)layer + 1u, width, sizeof(float), &dir_bytes) ||
|
||||
x->bytes < x_bytes || directions->bytes < dir_bytes) return 0;
|
||||
if (scale == 0.0f) return 1;
|
||||
|
||||
uint32_t nth = 256u;
|
||||
while (nth > width && nth > 1u) nth >>= 1;
|
||||
directional_steering_project_kernel<<<rows, nth>>>(
|
||||
(float *)x->ptr,
|
||||
(const float *)directions->ptr,
|
||||
layer,
|
||||
width,
|
||||
rows,
|
||||
scale);
|
||||
return cuda_ok(cudaGetLastError(), "directional steering launch");
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,591 @@
|
||||
__global__ static void rms_norm_plain_kernel(float *out, const float *x, uint32_t n, uint32_t rows, float eps) {
|
||||
uint32_t row = blockIdx.x;
|
||||
if (row >= rows) return;
|
||||
const float *xr = x + (uint64_t)row * n;
|
||||
float *orow = out + (uint64_t)row * n;
|
||||
float sum = 0.0f;
|
||||
for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) {
|
||||
float v = xr[i];
|
||||
sum += v * v;
|
||||
}
|
||||
__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();
|
||||
}
|
||||
float scale = rsqrtf(partial[0] / (float)n + eps);
|
||||
for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) {
|
||||
orow[i] = xr[i] * scale;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ static void rms_norm_weight_kernel(float *out, const float *x, const float *w, uint32_t n, uint32_t rows, float eps) {
|
||||
uint32_t row = blockIdx.x;
|
||||
if (row >= rows) return;
|
||||
const float *xr = x + (uint64_t)row * n;
|
||||
float *orow = out + (uint64_t)row * n;
|
||||
float sum = 0.0f;
|
||||
for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) {
|
||||
float v = xr[i];
|
||||
sum += v * v;
|
||||
}
|
||||
__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();
|
||||
}
|
||||
float scale = rsqrtf(partial[0] / (float)n + eps);
|
||||
for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) {
|
||||
orow[i] = xr[i] * scale * w[i];
|
||||
}
|
||||
}
|
||||
|
||||
__global__ static void dsv4_qkv_rms_norm_rows_kernel(
|
||||
float *q_out,
|
||||
const float *q,
|
||||
const float *q_w,
|
||||
uint32_t q_n,
|
||||
float *kv_out,
|
||||
const float *kv,
|
||||
const float *kv_w,
|
||||
uint32_t kv_n,
|
||||
uint32_t rows,
|
||||
float eps) {
|
||||
const uint32_t row = blockIdx.x;
|
||||
const uint32_t which = blockIdx.y;
|
||||
if (row >= rows || which > 1u) return;
|
||||
const uint32_t n = which == 0u ? q_n : kv_n;
|
||||
const float *xr = (which == 0u ? q : kv) + (uint64_t)row * n;
|
||||
float *orow = (which == 0u ? q_out : kv_out) + (uint64_t)row * n;
|
||||
const float *w = which == 0u ? q_w : kv_w;
|
||||
float sum = 0.0f;
|
||||
for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) {
|
||||
const float v = xr[i];
|
||||
sum += v * v;
|
||||
}
|
||||
__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 scale = rsqrtf(partial[0] / (float)n + eps);
|
||||
for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) {
|
||||
orow[i] = xr[i] * scale * w[i];
|
||||
}
|
||||
}
|
||||
|
||||
__global__ static void head_rms_norm_kernel(float *x, uint32_t n_tok, uint32_t n_head, uint32_t head_dim, float eps) {
|
||||
uint32_t row = blockIdx.x;
|
||||
if (row >= n_tok * n_head) return;
|
||||
float *xr = x + (uint64_t)row * head_dim;
|
||||
float sum = 0.0f;
|
||||
for (uint32_t i = threadIdx.x; i < head_dim; i += blockDim.x) {
|
||||
float v = xr[i];
|
||||
sum += v * v;
|
||||
}
|
||||
__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();
|
||||
}
|
||||
float scale = rsqrtf(partial[0] / (float)head_dim + eps);
|
||||
for (uint32_t i = threadIdx.x; i < head_dim; i += blockDim.x) xr[i] *= scale;
|
||||
}
|
||||
|
||||
__device__ static float rope_yarn_ramp_dev(float low, float high, int i0);
|
||||
|
||||
__global__ static void head_rms_norm_rope_tail_kernel(
|
||||
float *x,
|
||||
uint32_t n_tok,
|
||||
uint32_t n_head,
|
||||
uint32_t head_dim,
|
||||
uint32_t n_rot,
|
||||
uint32_t pos0,
|
||||
uint32_t n_ctx_orig,
|
||||
int inverse,
|
||||
float freq_base,
|
||||
float freq_scale,
|
||||
float ext_factor,
|
||||
float attn_factor,
|
||||
float beta_fast,
|
||||
float beta_slow,
|
||||
float eps) {
|
||||
uint32_t row = blockIdx.x;
|
||||
if (row >= n_tok * n_head) return;
|
||||
uint32_t t = row / n_head;
|
||||
float *xr = x + (uint64_t)row * head_dim;
|
||||
float sum = 0.0f;
|
||||
for (uint32_t i = threadIdx.x; i < head_dim; i += blockDim.x) {
|
||||
float v = xr[i];
|
||||
sum += v * v;
|
||||
}
|
||||
__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 scale = rsqrtf(partial[0] / (float)head_dim + eps);
|
||||
const uint32_t n_nope = head_dim - n_rot;
|
||||
for (uint32_t i = threadIdx.x; i < n_nope; i += blockDim.x) {
|
||||
xr[i] *= scale;
|
||||
}
|
||||
|
||||
float corr0 = 0.0f, corr1 = 0.0f;
|
||||
if (ext_factor != 0.0f) {
|
||||
float denom = 2.0f * logf(freq_base);
|
||||
corr0 = floorf((float)n_rot * logf((float)n_ctx_orig / (beta_fast * 2.0f * (float)M_PI)) / denom);
|
||||
corr1 = ceilf((float)n_rot * logf((float)n_ctx_orig / (beta_slow * 2.0f * (float)M_PI)) / denom);
|
||||
corr0 = fmaxf(0.0f, corr0);
|
||||
corr1 = fminf((float)(n_rot - 1), corr1);
|
||||
}
|
||||
const float theta_scale = powf(freq_base, -2.0f / (float)n_rot);
|
||||
for (uint32_t pair = threadIdx.x; pair < n_rot / 2; pair += blockDim.x) {
|
||||
uint32_t i = pair * 2u;
|
||||
float theta_extrap = (float)(pos0 + t) * powf(theta_scale, (float)pair);
|
||||
float theta_interp = freq_scale * theta_extrap;
|
||||
float theta = theta_interp;
|
||||
float mscale = attn_factor;
|
||||
if (ext_factor != 0.0f) {
|
||||
float ramp_mix = rope_yarn_ramp_dev(corr0, corr1, (int)i) * ext_factor;
|
||||
theta = theta_interp * (1.0f - ramp_mix) + theta_extrap * ramp_mix;
|
||||
mscale *= 1.0f + 0.1f * logf(1.0f / freq_scale);
|
||||
}
|
||||
float c = cosf(theta) * mscale;
|
||||
float s = sinf(theta) * mscale;
|
||||
if (inverse) s = -s;
|
||||
float *tail = xr + n_nope;
|
||||
float x0 = tail[i] * scale;
|
||||
float x1 = tail[i + 1] * scale;
|
||||
tail[i] = x0 * c - x1 * s;
|
||||
tail[i + 1] = x0 * s + x1 * c;
|
||||
}
|
||||
}
|
||||
|
||||
__global__ static void head_rms_norm_rope_tail_from_half_kernel(
|
||||
float *out,
|
||||
const __half *x,
|
||||
uint32_t n_tok,
|
||||
uint32_t n_head,
|
||||
uint32_t head_dim,
|
||||
uint32_t n_rot,
|
||||
uint32_t pos0,
|
||||
uint32_t n_ctx_orig,
|
||||
int inverse,
|
||||
float freq_base,
|
||||
float freq_scale,
|
||||
float ext_factor,
|
||||
float attn_factor,
|
||||
float beta_fast,
|
||||
float beta_slow,
|
||||
float eps) {
|
||||
uint32_t row = blockIdx.x;
|
||||
if (row >= n_tok * n_head) return;
|
||||
uint32_t t = row / n_head;
|
||||
const __half *xr = x + (uint64_t)row * head_dim;
|
||||
float *orow = out + (uint64_t)row * head_dim;
|
||||
float sum = 0.0f;
|
||||
for (uint32_t i = threadIdx.x; i < head_dim; i += blockDim.x) {
|
||||
float v = __half2float(xr[i]);
|
||||
sum += v * v;
|
||||
}
|
||||
__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 scale = rsqrtf(partial[0] / (float)head_dim + eps);
|
||||
const uint32_t n_nope = head_dim - n_rot;
|
||||
for (uint32_t i = threadIdx.x; i < n_nope; i += blockDim.x) {
|
||||
orow[i] = __half2float(xr[i]) * scale;
|
||||
}
|
||||
|
||||
float corr0 = 0.0f, corr1 = 0.0f;
|
||||
if (ext_factor != 0.0f) {
|
||||
float denom = 2.0f * logf(freq_base);
|
||||
corr0 = floorf((float)n_rot * logf((float)n_ctx_orig / (beta_fast * 2.0f * (float)M_PI)) / denom);
|
||||
corr1 = ceilf((float)n_rot * logf((float)n_ctx_orig / (beta_slow * 2.0f * (float)M_PI)) / denom);
|
||||
corr0 = fmaxf(0.0f, corr0);
|
||||
corr1 = fminf((float)(n_rot - 1), corr1);
|
||||
}
|
||||
const float theta_scale = powf(freq_base, -2.0f / (float)n_rot);
|
||||
const __half *tail = xr + n_nope;
|
||||
float *otail = orow + n_nope;
|
||||
for (uint32_t pair = threadIdx.x; pair < n_rot / 2; pair += blockDim.x) {
|
||||
uint32_t i = pair * 2u;
|
||||
float theta_extrap = (float)(pos0 + t) * powf(theta_scale, (float)pair);
|
||||
float theta_interp = freq_scale * theta_extrap;
|
||||
float theta = theta_interp;
|
||||
float mscale = attn_factor;
|
||||
if (ext_factor != 0.0f) {
|
||||
float ramp_mix = rope_yarn_ramp_dev(corr0, corr1, (int)i) * ext_factor;
|
||||
theta = theta_interp * (1.0f - ramp_mix) + theta_extrap * ramp_mix;
|
||||
mscale *= 1.0f + 0.1f * logf(1.0f / freq_scale);
|
||||
}
|
||||
float c = cosf(theta) * mscale;
|
||||
float s = sinf(theta) * mscale;
|
||||
if (inverse) s = -s;
|
||||
float x0 = __half2float(tail[i]) * scale;
|
||||
float x1 = __half2float(tail[i + 1]) * scale;
|
||||
otail[i] = x0 * c - x1 * s;
|
||||
otail[i + 1] = x0 * s + x1 * c;
|
||||
}
|
||||
}
|
||||
|
||||
__device__ static float rope_yarn_ramp_dev(float low, float high, int i0) {
|
||||
float y = ((float)(i0 / 2) - low) / fmaxf(0.001f, high - low);
|
||||
return 1.0f - fminf(1.0f, fmaxf(0.0f, y));
|
||||
}
|
||||
|
||||
__global__ static void rope_tail_kernel(
|
||||
float *x,
|
||||
uint32_t n_tok,
|
||||
uint32_t n_head,
|
||||
uint32_t head_dim,
|
||||
uint32_t n_rot,
|
||||
uint32_t pos0,
|
||||
uint32_t pos_stride,
|
||||
uint32_t n_ctx_orig,
|
||||
int inverse,
|
||||
float freq_base,
|
||||
float freq_scale,
|
||||
float ext_factor,
|
||||
float attn_factor,
|
||||
float beta_fast,
|
||||
float beta_slow) {
|
||||
uint32_t gid = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
uint32_t pairs = n_tok * n_head * (n_rot / 2);
|
||||
if (gid >= pairs) return;
|
||||
uint32_t pair = gid % (n_rot / 2);
|
||||
uint32_t tmp = gid / (n_rot / 2);
|
||||
uint32_t h = tmp % n_head;
|
||||
uint32_t t = tmp / n_head;
|
||||
uint32_t n_nope = head_dim - n_rot;
|
||||
uint32_t i = pair * 2;
|
||||
|
||||
float corr0 = 0.0f, corr1 = 0.0f;
|
||||
if (ext_factor != 0.0f) {
|
||||
float denom = 2.0f * logf(freq_base);
|
||||
corr0 = floorf((float)n_rot * logf((float)n_ctx_orig / (beta_fast * 2.0f * (float)M_PI)) / denom);
|
||||
corr1 = ceilf((float)n_rot * logf((float)n_ctx_orig / (beta_slow * 2.0f * (float)M_PI)) / denom);
|
||||
corr0 = fmaxf(0.0f, corr0);
|
||||
corr1 = fminf((float)(n_rot - 1), corr1);
|
||||
}
|
||||
|
||||
const float theta_scale = powf(freq_base, -2.0f / (float)n_rot);
|
||||
float theta_extrap = (float)(pos0 + t * pos_stride) * powf(theta_scale, (float)pair);
|
||||
float theta_interp = freq_scale * theta_extrap;
|
||||
float theta = theta_interp;
|
||||
float mscale = attn_factor;
|
||||
if (ext_factor != 0.0f) {
|
||||
float ramp_mix = rope_yarn_ramp_dev(corr0, corr1, (int)i) * ext_factor;
|
||||
theta = theta_interp * (1.0f - ramp_mix) + theta_extrap * ramp_mix;
|
||||
mscale *= 1.0f + 0.1f * logf(1.0f / freq_scale);
|
||||
}
|
||||
float c = cosf(theta) * mscale;
|
||||
float s = sinf(theta) * mscale;
|
||||
if (inverse) s = -s;
|
||||
|
||||
float *tail = x + ((uint64_t)t * n_head + h) * head_dim + n_nope;
|
||||
float x0 = tail[i];
|
||||
float x1 = tail[i + 1];
|
||||
tail[i] = x0 * c - x1 * s;
|
||||
tail[i + 1] = x0 * s + x1 * c;
|
||||
}
|
||||
|
||||
__device__ static float dsv4_e4m3fn_value_dev(int i) {
|
||||
int exp = (i >> 3) & 15;
|
||||
int mant = i & 7;
|
||||
if (exp == 0) return (float)mant * 0.001953125f;
|
||||
return (1.0f + (float)mant * 0.125f) * exp2f((float)exp - 7.0f);
|
||||
}
|
||||
|
||||
__device__ static float dsv4_e4m3fn_dequant_dev(float x) {
|
||||
float sign = x < 0.0f ? -1.0f : 1.0f;
|
||||
float ax = fminf(fabsf(x), 448.0f);
|
||||
int lo = 0, hi = 126;
|
||||
while (lo < hi) {
|
||||
int mid = (lo + hi + 1) >> 1;
|
||||
if (dsv4_e4m3fn_value_dev(mid) <= ax) lo = mid;
|
||||
else hi = mid - 1;
|
||||
}
|
||||
int best = lo;
|
||||
if (best < 126) {
|
||||
float bd = fabsf(ax - dsv4_e4m3fn_value_dev(best));
|
||||
float nd = fabsf(ax - dsv4_e4m3fn_value_dev(best + 1));
|
||||
if (nd < bd || (nd == bd && (((best + 1) & 1) == 0) && ((best & 1) != 0))) best++;
|
||||
}
|
||||
return sign * dsv4_e4m3fn_value_dev(best);
|
||||
}
|
||||
|
||||
__device__ static float dsv4_e2m1fn_value_dev(int i) {
|
||||
switch (i & 7) {
|
||||
case 0: return 0.0f;
|
||||
case 1: return 0.5f;
|
||||
case 2: return 1.0f;
|
||||
case 3: return 1.5f;
|
||||
case 4: return 2.0f;
|
||||
case 5: return 3.0f;
|
||||
case 6: return 4.0f;
|
||||
default: return 6.0f;
|
||||
}
|
||||
}
|
||||
|
||||
__device__ static float dsv4_e2m1fn_dequant_dev(float x) {
|
||||
float sign = x < 0.0f ? -1.0f : 1.0f;
|
||||
float ax = fminf(fabsf(x), 6.0f);
|
||||
int best = 0;
|
||||
float best_diff = fabsf(ax - dsv4_e2m1fn_value_dev(0));
|
||||
for (int i = 1; i < 8; i++) {
|
||||
float diff = fabsf(ax - dsv4_e2m1fn_value_dev(i));
|
||||
if (diff < best_diff || (diff == best_diff && ((i & 1) == 0) && ((best & 1) != 0))) {
|
||||
best = i;
|
||||
best_diff = diff;
|
||||
}
|
||||
}
|
||||
return sign * dsv4_e2m1fn_value_dev(best);
|
||||
}
|
||||
|
||||
__device__ static float model_scalar_dev(const void *base, uint64_t offset, uint32_t type, uint64_t idx) {
|
||||
const char *p = (const char *)base + offset;
|
||||
if (type == 1u) return __half2float(((const __half *)p)[idx]);
|
||||
return ((const float *)p)[idx];
|
||||
}
|
||||
|
||||
__device__ static float model_ape_value_dev(const void *base, uint64_t offset, uint32_t type,
|
||||
uint32_t width, uint32_t row, uint32_t col) {
|
||||
const char *p = (const char *)base + offset;
|
||||
if (type == 1u) return __half2float(((const __half *)p)[(uint64_t)row * width + col]);
|
||||
if (type == 8u) {
|
||||
const uint64_t row_bytes = ((uint64_t)width + 31u) / 32u * 34u;
|
||||
const unsigned char *blk = (const unsigned char *)p + (uint64_t)row * row_bytes + (uint64_t)(col >> 5) * 34u;
|
||||
const float d = q8_0_scale_scalar(blk);
|
||||
const int8_t q = ((const int8_t *)(blk + 2u))[col & 31u];
|
||||
return d * (float)q;
|
||||
}
|
||||
return ((const float *)p)[(uint64_t)row * width + col];
|
||||
}
|
||||
|
||||
__device__ static float rope_yarn_ramp_cpu_equiv_dev(float low, float high, int i0) {
|
||||
float y = ((float)(i0 / 2) - low) / fmaxf(0.001f, high - low);
|
||||
return 1.0f - fminf(1.0f, fmaxf(0.0f, y));
|
||||
}
|
||||
|
||||
__device__ static DS4_ROCM_UNUSED void rope_tail_one_dev(float *x, uint32_t head_dim, uint32_t n_rot, uint32_t pos, uint32_t n_ctx_orig, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow) {
|
||||
uint32_t n_nope = head_dim - n_rot;
|
||||
float corr0 = 0.0f, corr1 = 0.0f;
|
||||
if (ext_factor != 0.0f) {
|
||||
float denom = 2.0f * logf(freq_base);
|
||||
corr0 = fmaxf(0.0f, floorf((float)n_rot * logf((float)n_ctx_orig / (beta_fast * 2.0f * (float)M_PI)) / denom));
|
||||
corr1 = fminf((float)(n_rot - 1), ceilf((float)n_rot * logf((float)n_ctx_orig / (beta_slow * 2.0f * (float)M_PI)) / denom));
|
||||
}
|
||||
for (uint32_t i = 0; i < n_rot; i += 2) {
|
||||
float theta_extrap = (float)pos * powf(freq_base, -((float)i) / (float)n_rot);
|
||||
float theta_interp = freq_scale * theta_extrap;
|
||||
float theta = theta_interp;
|
||||
float mscale = attn_factor;
|
||||
if (ext_factor != 0.0f) {
|
||||
float mix = rope_yarn_ramp_cpu_equiv_dev(corr0, corr1, (int)i) * ext_factor;
|
||||
theta = theta_interp * (1.0f - mix) + theta_extrap * mix;
|
||||
mscale *= 1.0f + 0.1f * logf(1.0f / freq_scale);
|
||||
}
|
||||
float c = cosf(theta) * mscale;
|
||||
float s = sinf(theta) * mscale;
|
||||
float x0 = x[n_nope + i];
|
||||
float x1 = x[n_nope + i + 1];
|
||||
x[n_nope + i] = x0 * c - x1 * s;
|
||||
x[n_nope + i + 1] = x0 * s + x1 * c;
|
||||
}
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_rms_norm_plain_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *x, uint32_t n, float eps) {
|
||||
if (!cuda_tensor_has_f32(out, n) || !cuda_tensor_has_f32(x, n)) return 0;
|
||||
if (n == 0u) return 1;
|
||||
rms_norm_plain_kernel<<<1, 256>>>((float *)out->ptr, (const float *)x->ptr, n, 1, eps);
|
||||
return cuda_ok(cudaGetLastError(), "rms_norm_plain launch");
|
||||
}
|
||||
extern "C" int ds4_gpu_rms_norm_plain_rows_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *x, uint32_t n, uint32_t rows, float eps) {
|
||||
if (!cuda_tensor_has_elems2(out, n, rows, sizeof(float)) ||
|
||||
!cuda_tensor_has_elems2(x, n, rows, sizeof(float))) return 0;
|
||||
if (n == 0u || rows == 0u) return 1;
|
||||
rms_norm_plain_kernel<<<rows, 256>>>((float *)out->ptr, (const float *)x->ptr, n, rows, eps);
|
||||
return cuda_ok(cudaGetLastError(), "rms_norm_plain launch");
|
||||
}
|
||||
extern "C" int ds4_gpu_rms_norm_weight_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *x, const void *model_map, uint64_t model_size, uint64_t weight_offset, uint32_t n, float eps) {
|
||||
uint64_t weight_bytes = 0;
|
||||
if (!model_map || !cuda_u64_mul_checked(n, sizeof(float), &weight_bytes) ||
|
||||
!cuda_model_range_fits(model_size, weight_offset, weight_bytes) ||
|
||||
!cuda_tensor_has_f32(out, n) || !cuda_tensor_has_f32(x, n)) return 0;
|
||||
if (n == 0u) return 1;
|
||||
const char *wptr = cuda_model_range_ptr(model_map, weight_offset, weight_bytes, "rms_weight");
|
||||
if (!wptr) return 0;
|
||||
const float *w = (const float *)wptr;
|
||||
rms_norm_weight_kernel<<<1, 256>>>((float *)out->ptr, (const float *)x->ptr, w, n, 1, eps);
|
||||
return cuda_ok(cudaGetLastError(), "rms_norm_weight launch");
|
||||
}
|
||||
extern "C" int ds4_gpu_rms_norm_weight_rows_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *x, const void *model_map, uint64_t model_size, uint64_t weight_offset, uint32_t n, uint32_t rows, float eps) {
|
||||
uint64_t weight_bytes = 0;
|
||||
if (!model_map || !cuda_u64_mul_checked(n, sizeof(float), &weight_bytes) ||
|
||||
!cuda_model_range_fits(model_size, weight_offset, weight_bytes) ||
|
||||
!cuda_tensor_has_elems2(out, n, rows, sizeof(float)) ||
|
||||
!cuda_tensor_has_elems2(x, n, rows, sizeof(float))) return 0;
|
||||
if (n == 0u || rows == 0u) return 1;
|
||||
const char *wptr = cuda_model_range_ptr(model_map, weight_offset, weight_bytes, "rms_weight");
|
||||
if (!wptr) return 0;
|
||||
const float *w = (const float *)wptr;
|
||||
rms_norm_weight_kernel<<<rows, 256>>>((float *)out->ptr, (const float *)x->ptr, w, n, rows, eps);
|
||||
return cuda_ok(cudaGetLastError(), "rms_norm_weight launch");
|
||||
}
|
||||
extern "C" int ds4_gpu_dsv4_qkv_rms_norm_rows_tensor(
|
||||
ds4_gpu_tensor *q_out,
|
||||
const ds4_gpu_tensor *q,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t q_weight_offset,
|
||||
uint32_t q_n,
|
||||
ds4_gpu_tensor *kv_out,
|
||||
const ds4_gpu_tensor *kv,
|
||||
uint64_t kv_weight_offset,
|
||||
uint32_t kv_n,
|
||||
uint32_t rows,
|
||||
float eps) {
|
||||
uint64_t q_weight_bytes = 0, kv_weight_bytes = 0;
|
||||
if (!model_map || !cuda_u64_mul_checked(q_n, sizeof(float), &q_weight_bytes) ||
|
||||
!cuda_u64_mul_checked(kv_n, sizeof(float), &kv_weight_bytes) ||
|
||||
!cuda_model_range_fits(model_size, q_weight_offset, q_weight_bytes) ||
|
||||
!cuda_model_range_fits(model_size, kv_weight_offset, kv_weight_bytes) ||
|
||||
!cuda_tensor_has_elems2(q_out, q_n, rows, sizeof(float)) ||
|
||||
!cuda_tensor_has_elems2(q, q_n, rows, sizeof(float)) ||
|
||||
!cuda_tensor_has_elems2(kv_out, kv_n, rows, sizeof(float)) ||
|
||||
!cuda_tensor_has_elems2(kv, kv_n, rows, sizeof(float))) {
|
||||
return 0;
|
||||
}
|
||||
if ((q_n == 0u && kv_n == 0u) || rows == 0u) return 1;
|
||||
const float *q_w = (const float *)cuda_model_range_ptr(model_map,
|
||||
q_weight_offset, q_weight_bytes, "q_rms_weight");
|
||||
const float *kv_w = (const float *)cuda_model_range_ptr(model_map,
|
||||
kv_weight_offset, kv_weight_bytes, "kv_rms_weight");
|
||||
if (!q_w || !kv_w) return 0;
|
||||
dim3 grid(rows, 2u, 1u);
|
||||
dsv4_qkv_rms_norm_rows_kernel<<<grid, 256>>>(
|
||||
(float *)q_out->ptr,
|
||||
(const float *)q->ptr,
|
||||
q_w,
|
||||
q_n,
|
||||
(float *)kv_out->ptr,
|
||||
(const float *)kv->ptr,
|
||||
kv_w,
|
||||
kv_n,
|
||||
rows,
|
||||
eps);
|
||||
return cuda_ok(cudaGetLastError(), "dsv4 qkv rms norm rows launch");
|
||||
}
|
||||
extern "C" int ds4_gpu_head_rms_norm_tensor(ds4_gpu_tensor *x, uint32_t n_tok, uint32_t n_head, uint32_t head_dim, float eps) {
|
||||
uint64_t rows64 = 0;
|
||||
if (!cuda_u64_mul_checked(n_tok, n_head, &rows64) || rows64 > UINT32_MAX ||
|
||||
!cuda_tensor_has_elems3(x, n_tok, n_head, head_dim, sizeof(float))) return 0;
|
||||
if (rows64 == 0u || head_dim == 0u) return 1;
|
||||
head_rms_norm_kernel<<<(uint32_t)rows64, 256>>>((float *)x->ptr, n_tok, n_head, head_dim, eps);
|
||||
return cuda_ok(cudaGetLastError(), "head_rms_norm launch");
|
||||
}
|
||||
extern "C" int ds4_gpu_head_rms_norm_rope_tail_tensor(ds4_gpu_tensor *x, uint32_t n_tok, uint32_t n_head, uint32_t head_dim, uint32_t n_rot, uint32_t pos0, uint32_t n_ctx_orig, bool inverse, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow, float eps) {
|
||||
uint64_t rows64 = 0;
|
||||
if (n_rot > head_dim || (n_rot & 1u) ||
|
||||
!cuda_u64_mul_checked(n_tok, n_head, &rows64) || rows64 > UINT32_MAX ||
|
||||
!cuda_tensor_has_elems3(x, n_tok, n_head, head_dim, sizeof(float))) return 0;
|
||||
if (rows64 == 0u || head_dim == 0u) return 1;
|
||||
head_rms_norm_rope_tail_kernel<<<(uint32_t)rows64, 256>>>((float *)x->ptr, n_tok, n_head, head_dim, n_rot, pos0, n_ctx_orig, inverse ? 1 : 0, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow, eps);
|
||||
return cuda_ok(cudaGetLastError(), "head_rms_norm_rope_tail launch");
|
||||
}
|
||||
extern "C" int ds4_gpu_attn_q_b_f16_head_rms_rope_tail_tensor(
|
||||
ds4_gpu_tensor *out,
|
||||
ds4_gpu_tensor *q_half,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t weight_offset,
|
||||
uint64_t in_dim,
|
||||
uint64_t out_dim,
|
||||
const ds4_gpu_tensor *x,
|
||||
uint32_t n_tok,
|
||||
uint32_t n_head,
|
||||
uint32_t head_dim,
|
||||
uint32_t n_rot,
|
||||
uint32_t pos0,
|
||||
uint32_t n_ctx_orig,
|
||||
bool inverse,
|
||||
float freq_base,
|
||||
float freq_scale,
|
||||
float ext_factor,
|
||||
float attn_factor,
|
||||
float beta_fast,
|
||||
float beta_slow,
|
||||
float eps) {
|
||||
if (!g_cublas_ready || !out || !q_half || !x || !model_map || n_tok == 0 ||
|
||||
n_rot > head_dim || (n_rot & 1u) || out_dim != (uint64_t)n_head * head_dim ||
|
||||
x->bytes < (uint64_t)n_tok * in_dim * sizeof(float) ||
|
||||
out->bytes < (uint64_t)n_tok * out_dim * sizeof(float) ||
|
||||
q_half->bytes < (uint64_t)n_tok * out_dim * sizeof(__half)) return 0;
|
||||
const uint64_t blocks = (in_dim + 31u) / 32u;
|
||||
if (weight_offset > model_size || out_dim > UINT64_MAX / (blocks * 34u)) return 0;
|
||||
const uint64_t weight_bytes = out_dim * blocks * 34u;
|
||||
if (weight_bytes > model_size - weight_offset) return 0;
|
||||
const __half *w_f16 = cuda_q8_f16_ptr(model_map, weight_offset, weight_bytes, in_dim, out_dim, "attn_q_b");
|
||||
if (!w_f16) return 0;
|
||||
const uint64_t xh_count = (uint64_t)n_tok * in_dim;
|
||||
__half *xh = (__half *)cuda_tmp_alloc(xh_count * sizeof(__half), "attn q_b f16 activations");
|
||||
if (!xh) return 0;
|
||||
f32_to_f16_kernel<<<(xh_count + 255u) / 256u, 256>>>(xh, (const float *)x->ptr, xh_count);
|
||||
if (!cuda_ok(cudaGetLastError(), "attn q_b f16 activation convert launch")) return 0;
|
||||
const float alpha = 1.0f;
|
||||
const float beta = 0.0f;
|
||||
cublasStatus_t st = cublasGemmEx(g_cublas,
|
||||
CUBLAS_OP_T,
|
||||
CUBLAS_OP_N,
|
||||
(int)out_dim,
|
||||
(int)n_tok,
|
||||
(int)in_dim,
|
||||
&alpha,
|
||||
w_f16,
|
||||
CUDA_R_16F,
|
||||
(int)in_dim,
|
||||
xh,
|
||||
CUDA_R_16F,
|
||||
(int)in_dim,
|
||||
&beta,
|
||||
q_half->ptr,
|
||||
CUDA_R_16F,
|
||||
(int)out_dim,
|
||||
CUBLAS_COMPUTE_32F,
|
||||
CUBLAS_GEMM_DEFAULT);
|
||||
if (st != CUBLAS_STATUS_SUCCESS) {
|
||||
fprintf(stderr, "ds4: " DS4_GPU_BLAS_NAME " attn q_b f16-out matmul failed: status %d\n", (int)st);
|
||||
return 0;
|
||||
}
|
||||
head_rms_norm_rope_tail_from_half_kernel<<<n_tok * n_head, 256>>>(
|
||||
(float *)out->ptr, (const __half *)q_half->ptr, n_tok, n_head, head_dim, n_rot,
|
||||
pos0, n_ctx_orig, inverse ? 1 : 0, freq_base, freq_scale, ext_factor, attn_factor,
|
||||
beta_fast, beta_slow, eps);
|
||||
return cuda_ok(cudaGetLastError(), "attn q_b f16-out head_rms_norm_rope launch");
|
||||
}
|
||||
|
||||
static int cuda_rope_tail_stride_tensor(ds4_gpu_tensor *x, uint32_t n_tok, uint32_t n_head, uint32_t head_dim, uint32_t n_rot, uint32_t pos0, uint32_t pos_stride, uint32_t n_ctx_orig, bool inverse, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow) {
|
||||
if (!x || n_rot > head_dim || (n_rot & 1) || x->bytes < (uint64_t)n_tok * n_head * head_dim * sizeof(float)) return 0;
|
||||
uint32_t pairs = n_tok * n_head * (n_rot / 2);
|
||||
rope_tail_kernel<<<(pairs + 255) / 256, 256>>>((float *)x->ptr, n_tok, n_head, head_dim, n_rot, pos0, pos_stride, n_ctx_orig, inverse ? 1 : 0, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);
|
||||
return cuda_ok(cudaGetLastError(), "rope_tail launch");
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_rope_tail_tensor(ds4_gpu_tensor *x, uint32_t n_tok, uint32_t n_head, uint32_t head_dim, uint32_t n_rot, uint32_t pos0, uint32_t n_ctx_orig, bool inverse, float freq_base, float freq_scale, float ext_factor, float attn_factor, float beta_fast, float beta_slow) {
|
||||
return cuda_rope_tail_stride_tensor(x, n_tok, n_head, head_dim, n_rot, pos0, 1u, n_ctx_orig, inverse, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow);
|
||||
}
|
||||
@@ -0,0 +1,77 @@
|
||||
// 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;
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,232 @@
|
||||
__device__ static float softplus_dev(float x) {
|
||||
if (x > 20.0f) return x;
|
||||
if (x < -10.0f) return ds4_precise_expf(x);
|
||||
return ds4_precise_log1pf(ds4_precise_expf(x));
|
||||
}
|
||||
|
||||
__device__ __forceinline__ static bool router_score_better(float av, uint32_t ai, float bv, uint32_t bi) {
|
||||
return av > bv || (av == bv && ai < bi);
|
||||
}
|
||||
|
||||
template <uint32_t N_EXPERT>
|
||||
__global__ static void router_select_warp_topk_kernel(
|
||||
int32_t *selected,
|
||||
float *weights,
|
||||
float *probs,
|
||||
const float *bias,
|
||||
const int32_t *hash,
|
||||
const float *logits,
|
||||
const int32_t *tokens,
|
||||
int32_t token_scalar,
|
||||
uint32_t hash_rows,
|
||||
uint32_t n_tokens,
|
||||
float expert_weight_scale,
|
||||
int has_bias,
|
||||
int hash_mode) {
|
||||
const uint32_t lane = threadIdx.x;
|
||||
const uint32_t row_in_block = threadIdx.y;
|
||||
const uint32_t t = blockIdx.x * blockDim.y + row_in_block;
|
||||
if (t >= n_tokens || lane >= 32u) return;
|
||||
|
||||
const float *log = logits + (uint64_t)t * N_EXPERT;
|
||||
float *prob = probs + (uint64_t)t * N_EXPERT;
|
||||
int32_t *sel = selected + (uint64_t)t * DS4_ROCM_N_EXPERT_USED;
|
||||
float *w = weights + (uint64_t)t * DS4_ROCM_N_EXPERT_USED;
|
||||
__shared__ float sprob[4][N_EXPERT];
|
||||
float local_prob[N_EXPERT / 32u];
|
||||
float local_score[N_EXPERT / 32u];
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < N_EXPERT / 32u; j++) {
|
||||
const uint32_t e = lane + j * 32u;
|
||||
const float p = ds4_precise_sqrtf(softplus_dev(log[e]));
|
||||
local_prob[j] = p;
|
||||
local_score[j] = p + (has_bias ? bias[e] : 0.0f);
|
||||
sprob[row_in_block][e] = p;
|
||||
prob[e] = p;
|
||||
}
|
||||
__syncwarp();
|
||||
|
||||
if (hash_mode) {
|
||||
if (lane == 0) {
|
||||
int32_t tok = tokens ? tokens[t] : token_scalar;
|
||||
if (tok < 0 || (uint32_t)tok >= hash_rows) tok = 0;
|
||||
const int32_t *row = hash + (uint64_t)tok * DS4_ROCM_N_EXPERT_USED;
|
||||
float sum = 0.0f;
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < DS4_ROCM_N_EXPERT_USED; j++) {
|
||||
const int32_t e = row[j];
|
||||
sel[j] = e;
|
||||
const float v = (e >= 0 && e < N_EXPERT) ? sprob[row_in_block][(uint32_t)e] : 0.0f;
|
||||
w[j] = v;
|
||||
sum += v;
|
||||
}
|
||||
sum = fmaxf(sum, 6.103515625e-5f);
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < DS4_ROCM_N_EXPERT_USED; j++) w[j] = w[j] / sum * expert_weight_scale;
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
float out_prob[DS4_ROCM_N_EXPERT_USED] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
|
||||
uint32_t out_idx[DS4_ROCM_N_EXPERT_USED] = {0, 0, 0, 0, 0, 0};
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < DS4_ROCM_N_EXPERT_USED; k++) {
|
||||
float best_score = -INFINITY;
|
||||
float best_prob = 0.0f;
|
||||
uint32_t best_idx = UINT32_MAX;
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < N_EXPERT / 32u; j++) {
|
||||
const uint32_t e = lane + j * 32u;
|
||||
const float s = local_score[j];
|
||||
if (router_score_better(s, e, best_score, best_idx)) {
|
||||
best_score = s;
|
||||
best_prob = local_prob[j];
|
||||
best_idx = e;
|
||||
}
|
||||
}
|
||||
#pragma unroll
|
||||
for (uint32_t mask = 16u; mask > 0u; mask >>= 1u) {
|
||||
const float other_score = __shfl_xor_sync(FULL_WARP_MASK, best_score, mask);
|
||||
const float other_prob = __shfl_xor_sync(FULL_WARP_MASK, best_prob, mask);
|
||||
const uint32_t other_idx = __shfl_xor_sync(FULL_WARP_MASK, best_idx, mask);
|
||||
if (router_score_better(other_score, other_idx, best_score, best_idx)) {
|
||||
best_score = other_score;
|
||||
best_prob = other_prob;
|
||||
best_idx = other_idx;
|
||||
}
|
||||
}
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < N_EXPERT / 32u; j++) {
|
||||
const uint32_t e = lane + j * 32u;
|
||||
if (e == best_idx) local_score[j] = -INFINITY;
|
||||
}
|
||||
if (lane == 0) {
|
||||
out_idx[k] = best_idx;
|
||||
out_prob[k] = best_prob;
|
||||
}
|
||||
}
|
||||
|
||||
if (lane == 0) {
|
||||
float sum = 0.0f;
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < DS4_ROCM_N_EXPERT_USED; j++) {
|
||||
sel[j] = (int32_t)out_idx[j];
|
||||
w[j] = out_prob[j];
|
||||
sum += out_prob[j];
|
||||
}
|
||||
sum = fmaxf(sum, 6.103515625e-5f);
|
||||
#pragma unroll
|
||||
for (uint32_t j = 0; j < DS4_ROCM_N_EXPERT_USED; j++) w[j] = w[j] / sum * expert_weight_scale;
|
||||
}
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_router_select_tensor(ds4_gpu_tensor *selected, ds4_gpu_tensor *weights, ds4_gpu_tensor *probs, const void *model_map, uint64_t model_size, uint64_t bias_offset, uint64_t hash_offset, uint32_t hash_rows, uint32_t token, uint32_t n_expert, uint32_t n_expert_used, float expert_weight_scale, uint32_t n_expert_groups, uint32_t n_group_used, bool has_bias, bool hash_mode, const ds4_gpu_tensor *logits) {
|
||||
const uint32_t active_n_expert = n_expert != 0u ? n_expert : DS4_ROCM_N_EXPERT;
|
||||
const float active_scale = expert_weight_scale != 0.0f ? expert_weight_scale : DS4_ROCM_EXPERT_WEIGHT_SCALE;
|
||||
if (!selected || !weights || !probs || !logits || !model_map || n_expert_groups > 1u || n_group_used > 0u ||
|
||||
(active_n_expert != DS4_ROCM_N_EXPERT && active_n_expert != DS4_ROCM_MAX_N_EXPERT) ||
|
||||
(n_expert_used != 0u && n_expert_used != DS4_ROCM_N_EXPERT_USED) ||
|
||||
!(active_scale > 0.0f) ||
|
||||
!cuda_tensor_has_f32(logits, active_n_expert) ||
|
||||
!cuda_tensor_has_f32(probs, active_n_expert) ||
|
||||
!cuda_tensor_has_i32(selected, DS4_ROCM_N_EXPERT_USED) ||
|
||||
!cuda_tensor_has_f32(weights, DS4_ROCM_N_EXPERT_USED)) return 0;
|
||||
int32_t tok = (int32_t)token;
|
||||
int ok = 1;
|
||||
const float *bias = NULL;
|
||||
const int32_t *hash = NULL;
|
||||
if (ok && has_bias && !hash_mode) {
|
||||
if (!cuda_model_range_fits(model_size, bias_offset, active_n_expert * sizeof(float))) ok = 0;
|
||||
else bias = (const float *)cuda_model_range_ptr(model_map, bias_offset, active_n_expert * sizeof(float), "router_bias");
|
||||
if (!bias) ok = 0;
|
||||
}
|
||||
if (ok && hash_mode) {
|
||||
if (hash_rows == 0u) ok = 0;
|
||||
else {
|
||||
const uint64_t hash_bytes = (uint64_t)hash_rows * DS4_ROCM_N_EXPERT_USED * sizeof(int32_t);
|
||||
if (!cuda_model_range_fits(model_size, hash_offset, hash_bytes)) ok = 0;
|
||||
else hash = (const int32_t *)cuda_model_range_ptr(model_map, hash_offset, hash_bytes, "router_hash");
|
||||
if (!hash) ok = 0;
|
||||
}
|
||||
}
|
||||
if (ok) {
|
||||
dim3 block(32, 4, 1);
|
||||
if (active_n_expert == DS4_ROCM_MAX_N_EXPERT) {
|
||||
router_select_warp_topk_kernel<DS4_ROCM_MAX_N_EXPERT><<<1, block>>>(
|
||||
(int32_t *)selected->ptr, (float *)weights->ptr, (float *)probs->ptr,
|
||||
bias, hash, (const float *)logits->ptr, NULL, tok, hash_rows, 1,
|
||||
active_scale, has_bias && !hash_mode, hash_mode);
|
||||
} else {
|
||||
router_select_warp_topk_kernel<DS4_ROCM_N_EXPERT><<<1, block>>>(
|
||||
(int32_t *)selected->ptr, (float *)weights->ptr, (float *)probs->ptr,
|
||||
bias, hash, (const float *)logits->ptr, NULL, tok, hash_rows, 1,
|
||||
active_scale, has_bias && !hash_mode, hash_mode);
|
||||
}
|
||||
ok = cuda_ok(cudaGetLastError(), "router_select launch");
|
||||
}
|
||||
return ok;
|
||||
}
|
||||
extern "C" int ds4_gpu_router_select_batch_tensor(ds4_gpu_tensor *selected, ds4_gpu_tensor *weights, ds4_gpu_tensor *probs, const void *model_map, uint64_t model_size, uint64_t bias_offset, uint64_t hash_offset, uint32_t hash_rows, uint32_t n_expert_groups, uint32_t n_group_used, bool has_bias, bool hash_mode, const ds4_gpu_tensor *logits, const ds4_gpu_tensor *tokens, uint32_t n_expert, uint32_t n_expert_used, float expert_weight_scale, uint32_t n_tokens) {
|
||||
const uint32_t active_n_expert = n_expert != 0u ? n_expert : DS4_ROCM_N_EXPERT;
|
||||
const float active_scale = expert_weight_scale != 0.0f ? expert_weight_scale : DS4_ROCM_EXPERT_WEIGHT_SCALE;
|
||||
if (!selected || !weights || !probs || !logits || !tokens || !model_map || n_tokens == 0 ||
|
||||
n_expert_groups > 1u || n_group_used > 0u ||
|
||||
(active_n_expert != DS4_ROCM_N_EXPERT && active_n_expert != DS4_ROCM_MAX_N_EXPERT) ||
|
||||
(n_expert_used != 0u && n_expert_used != DS4_ROCM_N_EXPERT_USED) ||
|
||||
!(active_scale > 0.0f) ||
|
||||
!cuda_tensor_has_i32(tokens, n_tokens) ||
|
||||
!cuda_tensor_has_elems2(logits, n_tokens, active_n_expert, sizeof(float)) ||
|
||||
!cuda_tensor_has_elems2(probs, n_tokens, active_n_expert, sizeof(float)) ||
|
||||
!cuda_tensor_has_elems2(selected, n_tokens, DS4_ROCM_N_EXPERT_USED, sizeof(int32_t)) ||
|
||||
!cuda_tensor_has_elems2(weights, n_tokens, DS4_ROCM_N_EXPERT_USED, sizeof(float))) {
|
||||
return 0;
|
||||
}
|
||||
const float *bias = NULL;
|
||||
const int32_t *hash = NULL;
|
||||
if (has_bias && !hash_mode) {
|
||||
if (!cuda_model_range_fits(model_size, bias_offset, active_n_expert * sizeof(float))) return 0;
|
||||
bias = (const float *)cuda_model_range_ptr(model_map, bias_offset, active_n_expert * sizeof(float), "router_bias");
|
||||
if (!bias) return 0;
|
||||
}
|
||||
if (hash_mode) {
|
||||
if (hash_rows == 0u) return 0;
|
||||
const uint64_t hash_bytes = (uint64_t)hash_rows * DS4_ROCM_N_EXPERT_USED * sizeof(int32_t);
|
||||
if (!cuda_model_range_fits(model_size, hash_offset, hash_bytes)) return 0;
|
||||
hash = (const int32_t *)cuda_model_range_ptr(model_map, hash_offset, hash_bytes, "router_hash");
|
||||
if (!hash) return 0;
|
||||
}
|
||||
dim3 block(32, 4, 1);
|
||||
if (active_n_expert == DS4_ROCM_MAX_N_EXPERT) {
|
||||
router_select_warp_topk_kernel<DS4_ROCM_MAX_N_EXPERT><<<(n_tokens + 3u) / 4u, block>>>(
|
||||
(int32_t *)selected->ptr,
|
||||
(float *)weights->ptr,
|
||||
(float *)probs->ptr,
|
||||
bias,
|
||||
hash,
|
||||
(const float *)logits->ptr,
|
||||
(const int32_t *)tokens->ptr,
|
||||
0,
|
||||
hash_rows,
|
||||
n_tokens,
|
||||
active_scale,
|
||||
has_bias && !hash_mode,
|
||||
hash_mode);
|
||||
} else {
|
||||
router_select_warp_topk_kernel<DS4_ROCM_N_EXPERT><<<(n_tokens + 3u) / 4u, block>>>(
|
||||
(int32_t *)selected->ptr,
|
||||
(float *)weights->ptr,
|
||||
(float *)probs->ptr,
|
||||
bias,
|
||||
hash,
|
||||
(const float *)logits->ptr,
|
||||
(const int32_t *)tokens->ptr,
|
||||
0,
|
||||
hash_rows,
|
||||
n_tokens,
|
||||
active_scale,
|
||||
has_bias && !hash_mode,
|
||||
hash_mode);
|
||||
}
|
||||
return cuda_ok(cudaGetLastError(), "router_select launch");
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,310 @@
|
||||
extern "C" int ds4_gpu_swiglu_tensor(ds4_gpu_tensor *out, const ds4_gpu_tensor *gate, const ds4_gpu_tensor *up, uint32_t n, float clamp, float weight) {
|
||||
if (!cuda_tensor_has_f32(out, n) || !cuda_tensor_has_f32(gate, n) || !cuda_tensor_has_f32(up, n)) return 0;
|
||||
if (n == 0u) return 1;
|
||||
swiglu_kernel<<<(n + 255) / 256, 256>>>((float *)out->ptr, (const float *)gate->ptr, (const float *)up->ptr, n, clamp, weight);
|
||||
return cuda_ok(cudaGetLastError(), "swiglu launch");
|
||||
}
|
||||
extern "C" int ds4_gpu_shared_gate_up_swiglu_q8_0_tensor(
|
||||
ds4_gpu_tensor *gate,
|
||||
ds4_gpu_tensor *up,
|
||||
ds4_gpu_tensor *mid,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t gate_offset,
|
||||
uint64_t up_offset,
|
||||
uint64_t in_dim,
|
||||
uint64_t out_dim,
|
||||
const ds4_gpu_tensor *x,
|
||||
float clamp) {
|
||||
if (!gate || !up || !mid || !model_map || !x ||
|
||||
in_dim == 0u || out_dim == 0u || in_dim > UINT32_MAX || out_dim > UINT32_MAX) {
|
||||
return 0;
|
||||
}
|
||||
const uint64_t blocks = (in_dim + 31u) / 32u;
|
||||
uint64_t row_bytes = 0;
|
||||
uint64_t weight_bytes = 0;
|
||||
uint64_t x_bytes = 0;
|
||||
uint64_t out_bytes = 0;
|
||||
if (!cuda_u64_mul_checked(blocks, 34u, &row_bytes) ||
|
||||
!cuda_u64_mul_checked(out_dim, row_bytes, &weight_bytes) ||
|
||||
!cuda_u64_mul3_checked(in_dim, 1u, sizeof(float), &x_bytes) ||
|
||||
!cuda_u64_mul3_checked(out_dim, 1u, sizeof(float), &out_bytes) ||
|
||||
!cuda_tensor_has_bytes(x, x_bytes) || !cuda_tensor_has_bytes(gate, out_bytes) ||
|
||||
!cuda_tensor_has_bytes(up, out_bytes) || !cuda_tensor_has_bytes(mid, out_bytes)) {
|
||||
return 0;
|
||||
}
|
||||
if (in_dim == 4096u && (in_dim & 31u) == 0u &&
|
||||
cuda_model_range_fits(model_size, gate_offset, weight_bytes) &&
|
||||
cuda_model_range_fits(model_size, up_offset, weight_bytes) &&
|
||||
!cuda_runtime_config()->disable_shared_gate_up_fused_w32) {
|
||||
const char *wg = cuda_model_range_ptr(model_map, gate_offset, weight_bytes, "shared_gate_q8");
|
||||
const char *wu = cuda_model_range_ptr(model_map, up_offset, weight_bytes, "shared_up_q8");
|
||||
if (!wg || !wu) return 0;
|
||||
const int store_gate_up = (g_quality_mode || cuda_runtime_config()->graph_dump) ? 1 : 0;
|
||||
const unsigned rows_per_block = 32u;
|
||||
shared_gate_up_swiglu_q8_0_rows_w32_kernel<<<
|
||||
(unsigned)((out_dim + rows_per_block - 1u) / rows_per_block),
|
||||
rows_per_block * 32u>>>(
|
||||
(float *)gate->ptr,
|
||||
(float *)up->ptr,
|
||||
(float *)mid->ptr,
|
||||
reinterpret_cast<const unsigned char *>(wg),
|
||||
reinterpret_cast<const unsigned char *>(wu),
|
||||
(const float *)x->ptr,
|
||||
(uint32_t)blocks,
|
||||
out_dim,
|
||||
row_bytes,
|
||||
store_gate_up,
|
||||
clamp);
|
||||
return cuda_ok(cudaGetLastError(), "shared gate/up fused q8 launch");
|
||||
}
|
||||
return ds4_gpu_matmul_q8_0_pair_tensor(gate, up,
|
||||
model_map, model_size,
|
||||
gate_offset, up_offset,
|
||||
in_dim, out_dim, out_dim,
|
||||
x, 1) &&
|
||||
ds4_gpu_swiglu_tensor(mid, gate, up, (uint32_t)out_dim, clamp, 1.0f);
|
||||
}
|
||||
|
||||
static cudaStream_t g_shared_gate_up_stream = NULL;
|
||||
static cudaEvent_t g_shared_gate_up_ready_event = NULL;
|
||||
static void *g_shared_gate_up_tmp = NULL;
|
||||
static uint64_t g_shared_gate_up_tmp_bytes = 0;
|
||||
static int g_shared_gate_up_pending = 0;
|
||||
|
||||
static int cuda_shared_gate_up_async_wait_internal(void) {
|
||||
if (!g_shared_gate_up_pending) return 1;
|
||||
cudaError_t err = cudaStreamSynchronize(g_shared_gate_up_stream);
|
||||
g_shared_gate_up_pending = 0;
|
||||
if (err != cudaSuccess) {
|
||||
fprintf(stderr, DS4_GPU_LOG_PREFIX "shared gate/up async wait failed: %s\n", cudaGetErrorString(err));
|
||||
(void)cudaGetLastError();
|
||||
return 0;
|
||||
}
|
||||
return 1;
|
||||
}
|
||||
|
||||
static void *cuda_shared_gate_up_async_tmp_alloc(uint64_t bytes) {
|
||||
if (bytes == 0) return NULL;
|
||||
if (g_shared_gate_up_tmp_bytes >= bytes) return g_shared_gate_up_tmp;
|
||||
if (g_shared_gate_up_tmp) {
|
||||
(void)cuda_shared_gate_up_async_wait_internal();
|
||||
(void)cudaFree(g_shared_gate_up_tmp);
|
||||
g_shared_gate_up_tmp = NULL;
|
||||
g_shared_gate_up_tmp_bytes = 0;
|
||||
}
|
||||
void *ptr = NULL;
|
||||
cudaError_t err = cudaMalloc(&ptr, (size_t)bytes);
|
||||
if (err != cudaSuccess) {
|
||||
fprintf(stderr, DS4_GPU_LOG_PREFIX "shared gate/up async temp alloc failed (%.2f MiB): %s\n",
|
||||
(double)bytes / 1048576.0, cudaGetErrorString(err));
|
||||
(void)cudaGetLastError();
|
||||
return NULL;
|
||||
}
|
||||
g_shared_gate_up_tmp = ptr;
|
||||
g_shared_gate_up_tmp_bytes = bytes;
|
||||
return g_shared_gate_up_tmp;
|
||||
}
|
||||
|
||||
static void cuda_shared_gate_up_async_cleanup(void) {
|
||||
if (g_shared_gate_up_stream) {
|
||||
(void)cuda_shared_gate_up_async_wait_internal();
|
||||
}
|
||||
if (g_shared_gate_up_tmp) {
|
||||
(void)cudaFree(g_shared_gate_up_tmp);
|
||||
g_shared_gate_up_tmp = NULL;
|
||||
g_shared_gate_up_tmp_bytes = 0;
|
||||
}
|
||||
if (g_shared_gate_up_ready_event) {
|
||||
(void)cudaEventDestroy(g_shared_gate_up_ready_event);
|
||||
g_shared_gate_up_ready_event = NULL;
|
||||
}
|
||||
if (g_shared_gate_up_stream) {
|
||||
(void)cudaStreamDestroy(g_shared_gate_up_stream);
|
||||
g_shared_gate_up_stream = NULL;
|
||||
}
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_shared_gate_up_swiglu_q8_0_async_tensor(
|
||||
ds4_gpu_tensor *gate,
|
||||
ds4_gpu_tensor *up,
|
||||
ds4_gpu_tensor *mid,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t gate_offset,
|
||||
uint64_t up_offset,
|
||||
uint64_t in_dim,
|
||||
uint64_t out_dim,
|
||||
const ds4_gpu_tensor *x,
|
||||
float clamp) {
|
||||
if (g_quality_mode || cuda_runtime_config()->graph_dump) return 0;
|
||||
if (g_shared_gate_up_pending && !cuda_shared_gate_up_async_wait_internal()) return 0;
|
||||
if (!gate || !up || !mid || !model_map || !x ||
|
||||
in_dim == 0u || out_dim == 0u || in_dim > UINT32_MAX || out_dim > UINT32_MAX) {
|
||||
return 0;
|
||||
}
|
||||
const uint64_t blocks = (in_dim + 31u) / 32u;
|
||||
uint64_t row_bytes = 0;
|
||||
uint64_t weight_bytes = 0;
|
||||
if (!cuda_u64_mul_checked(blocks, 34u, &row_bytes) ||
|
||||
!cuda_u64_mul_checked(out_dim, row_bytes, &weight_bytes)) {
|
||||
return 0;
|
||||
}
|
||||
if (g_quality_mode ||
|
||||
!gate || !up || !mid || !model_map || !x ||
|
||||
in_dim == 0u || out_dim == 0u || in_dim > UINT32_MAX || out_dim > UINT32_MAX ||
|
||||
gate_offset > model_size || up_offset > model_size ||
|
||||
weight_bytes > model_size - gate_offset ||
|
||||
weight_bytes > model_size - up_offset ||
|
||||
x->bytes < in_dim * sizeof(float) ||
|
||||
gate->bytes < out_dim * sizeof(float) ||
|
||||
up->bytes < out_dim * sizeof(float) ||
|
||||
mid->bytes < out_dim * sizeof(float)) {
|
||||
return 0;
|
||||
}
|
||||
const char *wg = cuda_model_range_ptr(model_map, gate_offset, weight_bytes, "shared_gate_q8_pair_async");
|
||||
const char *wu = cuda_model_range_ptr(model_map, up_offset, weight_bytes, "shared_up_q8_pair_async");
|
||||
if (!wg || !wu) return 0;
|
||||
if (!g_shared_gate_up_stream) {
|
||||
int least_priority = 0;
|
||||
int greatest_priority = 0;
|
||||
#ifdef __HIP_PLATFORM_AMD__
|
||||
hipError_t err = hipDeviceGetStreamPriorityRange(&least_priority, &greatest_priority);
|
||||
if (err == hipSuccess) {
|
||||
err = hipStreamCreateWithPriority(&g_shared_gate_up_stream, cudaStreamNonBlocking, least_priority);
|
||||
} else {
|
||||
(void)cudaGetLastError();
|
||||
err = hipStreamCreateWithFlags(&g_shared_gate_up_stream, cudaStreamNonBlocking);
|
||||
}
|
||||
if (err != hipSuccess) return 0;
|
||||
#else
|
||||
cudaError_t err = cudaDeviceGetStreamPriorityRange(&least_priority, &greatest_priority);
|
||||
if (err == cudaSuccess) {
|
||||
err = cudaStreamCreateWithPriority(&g_shared_gate_up_stream, cudaStreamNonBlocking, least_priority);
|
||||
} else {
|
||||
(void)cudaGetLastError();
|
||||
err = cudaStreamCreateWithFlags(&g_shared_gate_up_stream, cudaStreamNonBlocking);
|
||||
}
|
||||
if (err != cudaSuccess) return 0;
|
||||
#endif
|
||||
}
|
||||
if (!g_shared_gate_up_ready_event) {
|
||||
cudaError_t err = cudaEventCreateWithFlags(&g_shared_gate_up_ready_event, cudaEventDisableTiming);
|
||||
if (err != cudaSuccess) {
|
||||
fprintf(stderr, DS4_GPU_LOG_PREFIX "shared gate/up async event create failed: %s\n", cudaGetErrorString(err));
|
||||
(void)cudaGetLastError();
|
||||
return 0;
|
||||
}
|
||||
}
|
||||
/*
|
||||
* This stream is intentionally non-blocking so it can overlap routed MoE.
|
||||
* Non-blocking streams do not inherit default-stream ordering, so explicitly
|
||||
* wait until the default-stream producer of x (ffn_norm) has completed before
|
||||
* quantizing it here.
|
||||
*/
|
||||
cudaError_t dep_err = cudaEventRecord(g_shared_gate_up_ready_event, 0);
|
||||
if (dep_err != cudaSuccess) {
|
||||
fprintf(stderr, DS4_GPU_LOG_PREFIX "shared gate/up async dependency record failed: %s\n", cudaGetErrorString(dep_err));
|
||||
(void)cudaGetLastError();
|
||||
return 0;
|
||||
}
|
||||
#ifdef __HIP_PLATFORM_AMD__
|
||||
dep_err = hipStreamWaitEvent(g_shared_gate_up_stream, g_shared_gate_up_ready_event, 0);
|
||||
#else
|
||||
dep_err = cudaStreamWaitEvent(g_shared_gate_up_stream, g_shared_gate_up_ready_event, 0);
|
||||
#endif
|
||||
if (dep_err != cudaSuccess) {
|
||||
fprintf(stderr, DS4_GPU_LOG_PREFIX "shared gate/up async dependency wait failed: %s\n", cudaGetErrorString(dep_err));
|
||||
(void)cudaGetLastError();
|
||||
return 0;
|
||||
}
|
||||
const uint64_t xq_bytes = blocks * 32u;
|
||||
const uint64_t scale_offset = (xq_bytes + 15u) & ~15ull;
|
||||
const uint64_t tmp_bytes = scale_offset + blocks * sizeof(float);
|
||||
void *tmp = cuda_shared_gate_up_async_tmp_alloc(tmp_bytes);
|
||||
if (!tmp) return 0;
|
||||
int8_t *xq = (int8_t *)tmp;
|
||||
float *xscale = (float *)((char *)tmp + scale_offset);
|
||||
const int use_dp4a = 1;
|
||||
dim3 qgrid((unsigned)blocks, 1, 1);
|
||||
quantize_q8_0_f32_kernel<<<qgrid, 32, 0, g_shared_gate_up_stream>>>(xq, xscale, (const float *)x->ptr, in_dim, blocks);
|
||||
if (!cuda_ok(cudaGetLastError(), "shared gate/up async quantize launch")) return 0;
|
||||
matmul_q8_0_pair_preq_warp8_kernel<<<((unsigned)out_dim + 7u) / 8u, 256, 0, g_shared_gate_up_stream>>>(
|
||||
(float *)gate->ptr,
|
||||
(float *)up->ptr,
|
||||
reinterpret_cast<const unsigned char *>(wg),
|
||||
reinterpret_cast<const unsigned char *>(wu),
|
||||
xq,
|
||||
xscale,
|
||||
in_dim,
|
||||
out_dim,
|
||||
out_dim,
|
||||
blocks,
|
||||
use_dp4a);
|
||||
if (!cuda_ok(cudaGetLastError(), "shared gate/up async pair launch")) return 0;
|
||||
swiglu_kernel<<<((unsigned)out_dim + 255u) / 256u, 256, 0, g_shared_gate_up_stream>>>(
|
||||
(float *)mid->ptr,
|
||||
(const float *)gate->ptr,
|
||||
(const float *)up->ptr,
|
||||
(uint32_t)out_dim,
|
||||
clamp,
|
||||
1.0f);
|
||||
if (!cuda_ok(cudaGetLastError(), "shared gate/up async swiglu launch")) return 0;
|
||||
g_shared_gate_up_pending = 1;
|
||||
return 1;
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_shared_gate_up_async_wait(void) {
|
||||
return cuda_shared_gate_up_async_wait_internal();
|
||||
}
|
||||
|
||||
extern "C" int ds4_gpu_shared_gate_up_swiglu_q8_0_batch_tensor(
|
||||
ds4_gpu_tensor *gate,
|
||||
ds4_gpu_tensor *up,
|
||||
ds4_gpu_tensor *mid,
|
||||
const void *model_map,
|
||||
uint64_t model_size,
|
||||
uint64_t gate_offset,
|
||||
uint64_t up_offset,
|
||||
uint64_t in_dim,
|
||||
uint64_t out_dim,
|
||||
const ds4_gpu_tensor *x,
|
||||
uint64_t n_tok) {
|
||||
uint64_t x_bytes = 0, out_bytes = 0;
|
||||
if (!gate || !up || !mid || !model_map || !x || n_tok == 0 ||
|
||||
(in_dim & 31u) != 0u || in_dim == 0u || out_dim == 0u ||
|
||||
in_dim > UINT32_MAX || out_dim > UINT32_MAX || n_tok > UINT32_MAX ||
|
||||
!cuda_u64_mul3_checked(n_tok, in_dim, sizeof(float), &x_bytes) ||
|
||||
!cuda_u64_mul3_checked(n_tok, out_dim, sizeof(float), &out_bytes) ||
|
||||
x->bytes < x_bytes || gate->bytes < out_bytes || up->bytes < out_bytes || mid->bytes < out_bytes) {
|
||||
return 0;
|
||||
}
|
||||
const uint64_t blocks = (in_dim + 31u) / 32u;
|
||||
uint64_t row_bytes = 0, weight_bytes = 0;
|
||||
if (!cuda_u64_mul_checked(blocks, 34u, &row_bytes) ||
|
||||
!cuda_u64_mul_checked(out_dim, row_bytes, &weight_bytes)) return 0;
|
||||
if (gate_offset > model_size || up_offset > model_size ||
|
||||
weight_bytes > model_size - gate_offset || weight_bytes > model_size - up_offset) {
|
||||
return 0;
|
||||
}
|
||||
const char *wg = cuda_model_range_ptr(model_map, gate_offset, weight_bytes, "shared_gate_q8_batch");
|
||||
const char *wu = cuda_model_range_ptr(model_map, up_offset, weight_bytes, "shared_up_q8_batch");
|
||||
if (!wg || !wu) return 0;
|
||||
|
||||
const uint32_t rows_per_block = 32u;
|
||||
const uint32_t tile = 16u;
|
||||
const uint32_t block_tile = 16u;
|
||||
const dim3 grid((uint32_t)((out_dim + rows_per_block - 1u) / rows_per_block),
|
||||
(uint32_t)((n_tok + tile - 1u) / tile),
|
||||
1u);
|
||||
const size_t shmem = (size_t)tile * block_tile * 32u * sizeof(float);
|
||||
const int store_gate_up = (g_quality_mode || cuda_runtime_config()->graph_dump) ? 1 : 0;
|
||||
#define DS4_LAUNCH_SHARED_GU_BATCH(TT, BT) \
|
||||
shared_gate_up_swiglu_q8_0_batch_sharedx_w32_kernel<TT, BT><<<grid, rows_per_block * 32u, shmem>>>( \
|
||||
(float *)gate->ptr, (float *)up->ptr, (float *)mid->ptr, \
|
||||
reinterpret_cast<const unsigned char *>(wg), reinterpret_cast<const unsigned char *>(wu), \
|
||||
(const float *)x->ptr, (uint32_t)blocks, (uint32_t)out_dim, (uint32_t)n_tok, row_bytes, store_gate_up)
|
||||
DS4_LAUNCH_SHARED_GU_BATCH(16u, 16u);
|
||||
#undef DS4_LAUNCH_SHARED_GU_BATCH
|
||||
return cuda_ok(cudaGetLastError(), "shared gate/up fused q8 batch launch");
|
||||
}
|
||||
Reference in New Issue
Block a user