DeepseekV4: Add fused hyper-connection ops (#25585)
* dsv4 hc-ops * add missing files; * add cparams * update rpc version * address review comments * address review comments
This commit is contained in:
@@ -10944,6 +10944,291 @@ void ggml_compute_forward_gated_delta_net(
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// ggml_compute_forward_dsv4_hc_comb
|
||||
|
||||
static void ggml_dsv4_hc_comb_norm_cols(float * comb, float eps) {
|
||||
constexpr int64_t hc = 4;
|
||||
|
||||
for (int64_t idst = 0; idst < hc; ++idst) {
|
||||
float sum = eps;
|
||||
for (int64_t isrc = 0; isrc < hc; ++isrc) {
|
||||
sum += comb[idst + hc*isrc];
|
||||
}
|
||||
|
||||
const float inv_sum = 1.0f / sum;
|
||||
for (int64_t isrc = 0; isrc < hc; ++isrc) {
|
||||
comb[idst + hc*isrc] *= inv_sum;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static void ggml_dsv4_hc_comb_norm_rows(float * comb, float eps) {
|
||||
constexpr int64_t hc = 4;
|
||||
|
||||
for (int64_t isrc = 0; isrc < hc; ++isrc) {
|
||||
float sum = eps;
|
||||
for (int64_t idst = 0; idst < hc; ++idst) {
|
||||
sum += comb[idst + hc*isrc];
|
||||
}
|
||||
|
||||
const float inv_sum = 1.0f / sum;
|
||||
for (int64_t idst = 0; idst < hc; ++idst) {
|
||||
comb[idst + hc*isrc] *= inv_sum;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static void ggml_compute_forward_dsv4_hc_comb_f32(
|
||||
const ggml_compute_params * params,
|
||||
ggml_tensor * dst) {
|
||||
const ggml_tensor * mixes = dst->src[0];
|
||||
const ggml_tensor * scale = dst->src[1];
|
||||
const ggml_tensor * base = dst->src[2];
|
||||
|
||||
GGML_ASSERT(mixes->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(scale->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(base->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(dst->type == GGML_TYPE_F32);
|
||||
|
||||
constexpr int64_t hc = 4;
|
||||
constexpr int64_t comb_offset = 2*hc;
|
||||
constexpr int64_t hc_mix_dim = (2 + hc)*hc;
|
||||
|
||||
const int64_t n_tokens = mixes->ne[1];
|
||||
|
||||
GGML_ASSERT(mixes->ne[0] == hc_mix_dim);
|
||||
GGML_ASSERT(dst->ne[0] == hc);
|
||||
GGML_ASSERT(dst->ne[1] == hc);
|
||||
GGML_ASSERT(dst->ne[2] == n_tokens);
|
||||
GGML_ASSERT(scale->ne[0] >= 3);
|
||||
GGML_ASSERT(base->ne[0] == hc_mix_dim);
|
||||
|
||||
GGML_TENSOR_LOCALS(size_t, nbm, mixes, nb);
|
||||
GGML_TENSOR_LOCALS(size_t, nbs, scale, nb);
|
||||
GGML_TENSOR_LOCALS(size_t, nbb, base, nb);
|
||||
GGML_TENSOR_LOCALS(size_t, nbd, dst, nb);
|
||||
|
||||
const float eps = ggml_get_op_params_f32(dst, 0);
|
||||
const int32_t n_iter = ggml_get_op_params_i32(dst, 1);
|
||||
GGML_ASSERT(n_iter > 0);
|
||||
|
||||
const int ith = params->ith;
|
||||
const int nth = params->nth;
|
||||
|
||||
const int64_t dr = (n_tokens + nth - 1) / nth;
|
||||
const int64_t it0 = dr * ith;
|
||||
const int64_t it1 = MIN(it0 + dr, n_tokens);
|
||||
|
||||
const float scale_comb = *(const float *) ((const char *) scale->data + 2*nbs0);
|
||||
|
||||
for (int64_t it = it0; it < it1; ++it) {
|
||||
float comb[hc*hc];
|
||||
|
||||
for (int64_t isrc = 0; isrc < hc; ++isrc) {
|
||||
float max = -INFINITY;
|
||||
for (int64_t idst = 0; idst < hc; ++idst) {
|
||||
const int64_t idx = idst + hc*isrc;
|
||||
const float xv = *(const float *) ((const char *) mixes->data + (comb_offset + idx)*nbm0 + it*nbm1);
|
||||
const float bv = *(const float *) ((const char *) base->data + (comb_offset + idx)*nbb0);
|
||||
const float v = xv * scale_comb + bv;
|
||||
comb[idx] = v;
|
||||
max = MAX(max, v);
|
||||
}
|
||||
|
||||
float sum = 0.0f;
|
||||
for (int64_t idst = 0; idst < hc; ++idst) {
|
||||
const int64_t idx = idst + hc*isrc;
|
||||
const float v = expf(comb[idx] - max);
|
||||
comb[idx] = v;
|
||||
sum += v;
|
||||
}
|
||||
|
||||
const float inv_sum = 1.0f / sum;
|
||||
for (int64_t idst = 0; idst < hc; ++idst) {
|
||||
const int64_t idx = idst + hc*isrc;
|
||||
comb[idx] = comb[idx] * inv_sum + eps;
|
||||
}
|
||||
}
|
||||
|
||||
ggml_dsv4_hc_comb_norm_cols(comb, eps);
|
||||
for (int32_t i = 1; i < n_iter; ++i) {
|
||||
ggml_dsv4_hc_comb_norm_rows(comb, eps);
|
||||
ggml_dsv4_hc_comb_norm_cols(comb, eps);
|
||||
}
|
||||
|
||||
for (int64_t isrc = 0; isrc < hc; ++isrc) {
|
||||
for (int64_t idst = 0; idst < hc; ++idst) {
|
||||
const int64_t idx = idst + hc*isrc;
|
||||
*(float *) ((char *) dst->data + idst*nbd0 + isrc*nbd1 + it*nbd2) = comb[idx];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void ggml_compute_forward_dsv4_hc_comb(
|
||||
const ggml_compute_params * params,
|
||||
ggml_tensor * dst) {
|
||||
const ggml_tensor * src0 = dst->src[0];
|
||||
|
||||
switch (src0->type) {
|
||||
case GGML_TYPE_F32:
|
||||
{
|
||||
ggml_compute_forward_dsv4_hc_comb_f32(params, dst);
|
||||
} break;
|
||||
default:
|
||||
{
|
||||
GGML_ABORT("fatal error");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ggml_compute_forward_dsv4_hc_pre
|
||||
|
||||
static void ggml_compute_forward_dsv4_hc_pre_f32(
|
||||
const ggml_compute_params * params,
|
||||
ggml_tensor * dst) {
|
||||
const ggml_tensor * x = dst->src[0];
|
||||
const ggml_tensor * weights = dst->src[1];
|
||||
|
||||
GGML_ASSERT(x->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(weights->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(dst->type == GGML_TYPE_F32);
|
||||
|
||||
const int64_t n_embd = x->ne[0];
|
||||
const int64_t hc = x->ne[1];
|
||||
const int64_t n_tokens = x->ne[2];
|
||||
|
||||
GGML_ASSERT(dst->ne[0] == n_embd);
|
||||
GGML_ASSERT(dst->ne[1] == n_tokens);
|
||||
GGML_ASSERT(weights->ne[0] == hc);
|
||||
GGML_ASSERT(weights->ne[1] == n_tokens);
|
||||
|
||||
GGML_TENSOR_LOCALS(size_t, nbx, x, nb);
|
||||
GGML_TENSOR_LOCALS(size_t, nbw, weights, nb);
|
||||
GGML_TENSOR_LOCALS(size_t, nbd, dst, nb);
|
||||
|
||||
const int ith = params->ith;
|
||||
const int nth = params->nth;
|
||||
|
||||
const int64_t nr = n_embd * n_tokens;
|
||||
const int64_t dr = (nr + nth - 1) / nth;
|
||||
const int64_t ir0 = dr * ith;
|
||||
const int64_t ir1 = MIN(ir0 + dr, nr);
|
||||
|
||||
for (int64_t ir = ir0; ir < ir1; ++ir) {
|
||||
const int64_t i0 = ir % n_embd;
|
||||
const int64_t it = ir / n_embd;
|
||||
|
||||
float sum = 0.0f;
|
||||
for (int64_t ih = 0; ih < hc; ++ih) {
|
||||
const float xv = *(const float *) ((const char *) x->data + i0*nbx0 + ih*nbx1 + it*nbx2);
|
||||
const float wv = *(const float *) ((const char *) weights->data + ih*nbw0 + it*nbw1);
|
||||
sum += xv * wv;
|
||||
}
|
||||
|
||||
*(float *) ((char *) dst->data + i0*nbd0 + it*nbd1) = sum;
|
||||
}
|
||||
}
|
||||
|
||||
void ggml_compute_forward_dsv4_hc_pre(
|
||||
const ggml_compute_params * params,
|
||||
ggml_tensor * dst) {
|
||||
const ggml_tensor * src0 = dst->src[0];
|
||||
|
||||
switch (src0->type) {
|
||||
case GGML_TYPE_F32:
|
||||
{
|
||||
ggml_compute_forward_dsv4_hc_pre_f32(params, dst);
|
||||
} break;
|
||||
default:
|
||||
{
|
||||
GGML_ABORT("fatal error");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ggml_compute_forward_dsv4_hc_post
|
||||
|
||||
static void ggml_compute_forward_dsv4_hc_post_f32(
|
||||
const ggml_compute_params * params,
|
||||
ggml_tensor * dst) {
|
||||
const ggml_tensor * x = dst->src[0];
|
||||
const ggml_tensor * residual = dst->src[1];
|
||||
const ggml_tensor * post = dst->src[2];
|
||||
const ggml_tensor * comb = dst->src[3];
|
||||
|
||||
GGML_ASSERT(x->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(residual->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(post->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(comb->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(dst->type == GGML_TYPE_F32);
|
||||
|
||||
const int64_t n_embd = x->ne[0];
|
||||
const int64_t n_tokens = x->ne[1];
|
||||
const int64_t hc = residual->ne[1];
|
||||
|
||||
GGML_ASSERT(dst->ne[0] == n_embd);
|
||||
GGML_ASSERT(dst->ne[1] == hc);
|
||||
GGML_ASSERT(dst->ne[2] == n_tokens);
|
||||
GGML_ASSERT(residual->ne[0] == n_embd);
|
||||
GGML_ASSERT(residual->ne[2] == n_tokens);
|
||||
GGML_ASSERT(post->ne[0] == hc);
|
||||
GGML_ASSERT(post->ne[1] == n_tokens);
|
||||
GGML_ASSERT(comb->ne[0] == hc);
|
||||
GGML_ASSERT(comb->ne[1] == hc);
|
||||
GGML_ASSERT(comb->ne[2] == n_tokens);
|
||||
|
||||
GGML_TENSOR_LOCALS(size_t, nbx, x, nb);
|
||||
GGML_TENSOR_LOCALS(size_t, nbr, residual, nb);
|
||||
GGML_TENSOR_LOCALS(size_t, nbp, post, nb);
|
||||
GGML_TENSOR_LOCALS(size_t, nbc, comb, nb);
|
||||
GGML_TENSOR_LOCALS(size_t, nbd, dst, nb);
|
||||
|
||||
const int ith = params->ith;
|
||||
const int nth = params->nth;
|
||||
|
||||
const int64_t nr = n_embd * hc * n_tokens;
|
||||
const int64_t dr = (nr + nth - 1) / nth;
|
||||
const int64_t ir0 = dr * ith;
|
||||
const int64_t ir1 = MIN(ir0 + dr, nr);
|
||||
|
||||
for (int64_t ir = ir0; ir < ir1; ++ir) {
|
||||
const int64_t i0 = ir % n_embd;
|
||||
const int64_t idst = (ir / n_embd) % hc;
|
||||
const int64_t it = ir / (n_embd * hc);
|
||||
|
||||
const float xv = *(const float *) ((const char *) x->data + i0*nbx0 + it*nbx1);
|
||||
const float pv = *(const float *) ((const char *) post->data + idst*nbp0 + it*nbp1);
|
||||
|
||||
float sum = xv * pv;
|
||||
for (int64_t isrc = 0; isrc < hc; ++isrc) {
|
||||
const float rv = *(const float *) ((const char *) residual->data + i0*nbr0 + isrc*nbr1 + it*nbr2);
|
||||
const float cv = *(const float *) ((const char *) comb->data + idst*nbc0 + isrc*nbc1 + it*nbc2);
|
||||
sum += rv * cv;
|
||||
}
|
||||
|
||||
*(float *) ((char *) dst->data + i0*nbd0 + idst*nbd1 + it*nbd2) = sum;
|
||||
}
|
||||
}
|
||||
|
||||
void ggml_compute_forward_dsv4_hc_post(
|
||||
const ggml_compute_params * params,
|
||||
ggml_tensor * dst) {
|
||||
const ggml_tensor * src0 = dst->src[0];
|
||||
|
||||
switch (src0->type) {
|
||||
case GGML_TYPE_F32:
|
||||
{
|
||||
ggml_compute_forward_dsv4_hc_post_f32(params, dst);
|
||||
} break;
|
||||
default:
|
||||
{
|
||||
GGML_ABORT("fatal error");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ggml_compute_forward_rwkv_wkv7
|
||||
|
||||
static void ggml_compute_forward_rwkv_wkv7_f32(
|
||||
|
||||
Reference in New Issue
Block a user