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:
@@ -8,10 +8,10 @@ extern "C" {
|
||||
|
||||
#define RPC_PROTO_MAJOR_VERSION 4
|
||||
#define RPC_PROTO_MINOR_VERSION 0
|
||||
#define RPC_PROTO_PATCH_VERSION 2
|
||||
#define RPC_PROTO_PATCH_VERSION 3
|
||||
|
||||
#ifdef __cplusplus
|
||||
static_assert(GGML_OP_COUNT == 98, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
|
||||
static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION");
|
||||
#endif
|
||||
|
||||
#define GGML_RPC_MAX_SERVERS 16
|
||||
|
||||
@@ -571,6 +571,9 @@ extern "C" {
|
||||
GGML_OP_SOLVE_TRI,
|
||||
GGML_OP_GATED_DELTA_NET,
|
||||
GGML_OP_LIGHTNING_INDEXER,
|
||||
GGML_OP_DSV4_HC_COMB,
|
||||
GGML_OP_DSV4_HC_PRE,
|
||||
GGML_OP_DSV4_HC_POST,
|
||||
|
||||
GGML_OP_UNARY,
|
||||
|
||||
@@ -2598,6 +2601,45 @@ extern "C" {
|
||||
struct ggml_tensor * weights,
|
||||
struct ggml_tensor * mask);
|
||||
|
||||
// DeepSeek V4 hyper-connections (ref. https://arxiv.org/pdf/2512.24880)
|
||||
// In short these operations are replacements for the original residual connection (x = transformer(x) + x)
|
||||
// using a richer representation through streams.
|
||||
//
|
||||
// hc_comb: mixes [(2 + hc)*hc, n_tokens], scale [3], base [(2 + hc)*hc]
|
||||
// -> [dst_hc, src_hc, n_tokens]
|
||||
// logits[dst, src, t] = mixes[2*hc + dst + hc*src, t]*scale[2]
|
||||
// + base[2*hc + dst + hc*src]
|
||||
// Softmax over dst, add eps, normalize over src, then repeat normalization
|
||||
// over dst followed by src for iterations 1 through n_iter - 1.
|
||||
GGML_API struct ggml_tensor * ggml_dsv4_hc_comb(
|
||||
struct ggml_context * ctx,
|
||||
struct ggml_tensor * mixes,
|
||||
struct ggml_tensor * scale,
|
||||
struct ggml_tensor * base,
|
||||
float eps,
|
||||
int32_t n_iter);
|
||||
|
||||
// hc_pre: x [n_embd, hc, n_tokens], weights [hc, n_tokens] -> [n_embd, n_tokens]
|
||||
// result[i, t] = sum_h x[i, h, t]*weights[h, t]
|
||||
//
|
||||
GGML_API struct ggml_tensor * ggml_dsv4_hc_pre(
|
||||
struct ggml_context * ctx,
|
||||
struct ggml_tensor * x,
|
||||
struct ggml_tensor * weights);
|
||||
|
||||
// hc_post: x [n_embd, n_tokens], residual [n_embd, hc, n_tokens],
|
||||
// post [hc, n_tokens], comb [dst_hc, src_hc, n_tokens]
|
||||
// -> [n_embd, hc, n_tokens]
|
||||
// result[i, dst, t] = x[i, t]*post[dst, t]
|
||||
// + sum_src residual[i, src, t]*comb[dst, src, t]
|
||||
//
|
||||
GGML_API struct ggml_tensor * ggml_dsv4_hc_post(
|
||||
struct ggml_context * ctx,
|
||||
struct ggml_tensor * x,
|
||||
struct ggml_tensor * residual,
|
||||
struct ggml_tensor * post,
|
||||
struct ggml_tensor * comb);
|
||||
|
||||
// custom operators
|
||||
|
||||
typedef void (*ggml_custom1_op_t)(struct ggml_tensor * dst , const struct ggml_tensor * a, int ith, int nth, void * userdata);
|
||||
|
||||
@@ -984,6 +984,11 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
|
||||
case GGML_OP_GATED_DELTA_NET: {
|
||||
split_state = handle_gated_delta_net(src_ss);
|
||||
} break;
|
||||
case GGML_OP_DSV4_HC_COMB:
|
||||
case GGML_OP_DSV4_HC_PRE:
|
||||
case GGML_OP_DSV4_HC_POST: {
|
||||
split_state = handle_generic(src_ss, /*scalar_only =*/ true);
|
||||
} break;
|
||||
case GGML_OP_UNARY: {
|
||||
split_state = handle_generic(src_ss, /*scalar_only =*/ false);
|
||||
} break;
|
||||
|
||||
@@ -2064,6 +2064,18 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
|
||||
{
|
||||
ggml_compute_forward_lightning_indexer(params, tensor);
|
||||
} break;
|
||||
case GGML_OP_DSV4_HC_COMB:
|
||||
{
|
||||
ggml_compute_forward_dsv4_hc_comb(params, tensor);
|
||||
} break;
|
||||
case GGML_OP_DSV4_HC_PRE:
|
||||
{
|
||||
ggml_compute_forward_dsv4_hc_pre(params, tensor);
|
||||
} break;
|
||||
case GGML_OP_DSV4_HC_POST:
|
||||
{
|
||||
ggml_compute_forward_dsv4_hc_post(params, tensor);
|
||||
} break;
|
||||
case GGML_OP_MAP_CUSTOM1:
|
||||
{
|
||||
ggml_compute_forward_map_custom1(params, tensor);
|
||||
@@ -2244,6 +2256,9 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
|
||||
case GGML_OP_COUNT_EQUAL:
|
||||
case GGML_OP_SOLVE_TRI:
|
||||
case GGML_OP_GATED_DELTA_NET:
|
||||
case GGML_OP_DSV4_HC_COMB:
|
||||
case GGML_OP_DSV4_HC_PRE:
|
||||
case GGML_OP_DSV4_HC_POST:
|
||||
{
|
||||
n_tasks = n_threads;
|
||||
} break;
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -106,6 +106,9 @@ void ggml_compute_forward_solve_tri(const struct ggml_compute_params * params, s
|
||||
void ggml_compute_forward_gla(const struct ggml_compute_params * params, struct ggml_tensor * dst);
|
||||
void ggml_compute_forward_gated_delta_net(const struct ggml_compute_params * params, struct ggml_tensor * dst);
|
||||
void ggml_compute_forward_lightning_indexer(const struct ggml_compute_params * params, struct ggml_tensor * dst);
|
||||
void ggml_compute_forward_dsv4_hc_comb(const struct ggml_compute_params * params, struct ggml_tensor * dst);
|
||||
void ggml_compute_forward_dsv4_hc_pre(const struct ggml_compute_params * params, struct ggml_tensor * dst);
|
||||
void ggml_compute_forward_dsv4_hc_post(const struct ggml_compute_params * params, struct ggml_tensor * dst);
|
||||
void ggml_compute_forward_map_custom1(const struct ggml_compute_params * params, struct ggml_tensor * dst);
|
||||
void ggml_compute_forward_map_custom2(const struct ggml_compute_params * params, struct ggml_tensor * dst);
|
||||
void ggml_compute_forward_map_custom3(const struct ggml_compute_params * params, struct ggml_tensor * dst);
|
||||
|
||||
@@ -0,0 +1,294 @@
|
||||
#include "common.cuh"
|
||||
#include "dsv4-hc.cuh"
|
||||
|
||||
|
||||
static constexpr int DSV4_HC = 4;
|
||||
|
||||
|
||||
static __device__ void dsv4_hc_comb_norm_cols(float * comb, float eps) {
|
||||
for (int idst = 0; idst < DSV4_HC; ++idst) {
|
||||
float sum = eps;
|
||||
for (int isrc = 0; isrc < DSV4_HC; ++isrc) {
|
||||
sum += comb[idst + DSV4_HC*isrc];
|
||||
}
|
||||
|
||||
const float inv_sum = 1.0f / sum;
|
||||
for (int isrc = 0; isrc < DSV4_HC; ++isrc) {
|
||||
comb[idst + DSV4_HC*isrc] *= inv_sum;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static __device__ void dsv4_hc_comb_norm_rows(float * comb, float eps) {
|
||||
for (int isrc = 0; isrc < DSV4_HC; ++isrc) {
|
||||
float sum = eps;
|
||||
for (int idst = 0; idst < DSV4_HC; ++idst) {
|
||||
sum += comb[idst + DSV4_HC*isrc];
|
||||
}
|
||||
|
||||
const float inv_sum = 1.0f / sum;
|
||||
for (int idst = 0; idst < DSV4_HC; ++idst) {
|
||||
comb[idst + DSV4_HC*isrc] *= inv_sum;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static __global__ void dsv4_hc_comb_f32(
|
||||
const float * mixes,
|
||||
const float * scale,
|
||||
const float * base,
|
||||
float * dst,
|
||||
int64_t n_tokens,
|
||||
int64_t sm0,
|
||||
int64_t sm1,
|
||||
int64_t ss0,
|
||||
int64_t sb0,
|
||||
int64_t sd0,
|
||||
int64_t sd1,
|
||||
int64_t sd2,
|
||||
float eps,
|
||||
int32_t n_iter) {
|
||||
constexpr int comb_offset = 2*DSV4_HC;
|
||||
|
||||
ggml_cuda_pdl_lc();
|
||||
const int64_t it = (int64_t) blockIdx.x * blockDim.x + threadIdx.x;
|
||||
|
||||
if (it >= n_tokens) {
|
||||
return;
|
||||
}
|
||||
|
||||
ggml_cuda_pdl_sync();
|
||||
|
||||
const float scale_comb = scale[2*ss0];
|
||||
float comb[DSV4_HC*DSV4_HC];
|
||||
|
||||
for (int isrc = 0; isrc < DSV4_HC; ++isrc) {
|
||||
float max = -INFINITY;
|
||||
for (int idst = 0; idst < DSV4_HC; ++idst) {
|
||||
const int idx = idst + DSV4_HC*isrc;
|
||||
const float v = mixes[(comb_offset + idx)*sm0 + it*sm1] * scale_comb + base[(comb_offset + idx)*sb0];
|
||||
comb[idx] = v;
|
||||
max = fmaxf(max, v);
|
||||
}
|
||||
|
||||
float sum = 0.0f;
|
||||
for (int idst = 0; idst < DSV4_HC; ++idst) {
|
||||
const int idx = idst + DSV4_HC*isrc;
|
||||
const float v = expf(comb[idx] - max);
|
||||
comb[idx] = v;
|
||||
sum += v;
|
||||
}
|
||||
|
||||
const float inv_sum = 1.0f / sum;
|
||||
for (int idst = 0; idst < DSV4_HC; ++idst) {
|
||||
const int idx = idst + DSV4_HC*isrc;
|
||||
comb[idx] = comb[idx] * inv_sum + eps;
|
||||
}
|
||||
}
|
||||
|
||||
dsv4_hc_comb_norm_cols(comb, eps);
|
||||
for (int32_t i = 1; i < n_iter; ++i) {
|
||||
dsv4_hc_comb_norm_rows(comb, eps);
|
||||
dsv4_hc_comb_norm_cols(comb, eps);
|
||||
}
|
||||
|
||||
for (int isrc = 0; isrc < DSV4_HC; ++isrc) {
|
||||
for (int idst = 0; idst < DSV4_HC; ++idst) {
|
||||
const int idx = idst + DSV4_HC*isrc;
|
||||
dst[idst*sd0 + isrc*sd1 + it*sd2] = comb[idx];
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static __global__ void dsv4_hc_pre_f32(
|
||||
const float * x,
|
||||
const float * weights,
|
||||
float * dst,
|
||||
int64_t n_embd,
|
||||
int64_t hc,
|
||||
int64_t n_tokens,
|
||||
int64_t sx0,
|
||||
int64_t sx1,
|
||||
int64_t sx2,
|
||||
int64_t sw0,
|
||||
int64_t sw1,
|
||||
int64_t sd0,
|
||||
int64_t sd1) {
|
||||
ggml_cuda_pdl_lc();
|
||||
const int64_t ir = (int64_t) blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int64_t nr = n_embd * n_tokens;
|
||||
|
||||
if (ir >= nr) {
|
||||
return;
|
||||
}
|
||||
|
||||
ggml_cuda_pdl_sync();
|
||||
|
||||
const int64_t i0 = ir % n_embd;
|
||||
const int64_t it = ir / n_embd;
|
||||
|
||||
float sum = x[i0*sx0 + it*sx2] * weights[it*sw1];
|
||||
for (int64_t ih = 1; ih < hc; ++ih) {
|
||||
const float xv = x[i0*sx0 + ih*sx1 + it*sx2];
|
||||
const float wv = weights[ih*sw0 + it*sw1];
|
||||
sum += xv * wv;
|
||||
}
|
||||
|
||||
dst[i0*sd0 + it*sd1] = sum;
|
||||
}
|
||||
|
||||
static __global__ void dsv4_hc_post_f32(
|
||||
const float * x,
|
||||
const float * residual,
|
||||
const float * post,
|
||||
const float * comb,
|
||||
float * dst,
|
||||
int64_t n_embd,
|
||||
int64_t hc,
|
||||
int64_t n_tokens,
|
||||
int64_t sx0,
|
||||
int64_t sx1,
|
||||
int64_t sr0,
|
||||
int64_t sr1,
|
||||
int64_t sr2,
|
||||
int64_t sp0,
|
||||
int64_t sp1,
|
||||
int64_t sc0,
|
||||
int64_t sc1,
|
||||
int64_t sc2,
|
||||
int64_t sd0,
|
||||
int64_t sd1,
|
||||
int64_t sd2) {
|
||||
ggml_cuda_pdl_lc();
|
||||
const int64_t ir = (int64_t) blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int64_t nr = n_embd * hc * n_tokens;
|
||||
|
||||
if (ir >= nr) {
|
||||
return;
|
||||
}
|
||||
|
||||
ggml_cuda_pdl_sync();
|
||||
|
||||
const int64_t i0 = ir % n_embd;
|
||||
const int64_t idst = (ir / n_embd) % hc;
|
||||
const int64_t it = ir / (n_embd * hc);
|
||||
|
||||
float sum = x[i0*sx0 + it*sx1] * post[idst*sp0 + it*sp1];
|
||||
for (int64_t isrc = 0; isrc < hc; ++isrc) {
|
||||
sum += residual[i0*sr0 + isrc*sr1 + it*sr2] * comb[idst*sc0 + isrc*sc1 + it*sc2];
|
||||
}
|
||||
|
||||
dst[i0*sd0 + idst*sd1 + it*sd2] = sum;
|
||||
}
|
||||
|
||||
void ggml_cuda_op_dsv4_hc_comb(ggml_backend_cuda_context & ctx, 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_mix_dim = (2 + DSV4_HC)*DSV4_HC;
|
||||
|
||||
GGML_ASSERT(mixes->ne[0] == hc_mix_dim);
|
||||
GGML_ASSERT(dst->ne[0] == DSV4_HC);
|
||||
GGML_ASSERT(dst->ne[1] == DSV4_HC);
|
||||
GGML_ASSERT(dst->ne[2] == mixes->ne[1]);
|
||||
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 int64_t n_tokens = mixes->ne[1];
|
||||
const float eps = ggml_get_op_params_f32(dst, 0);
|
||||
const int32_t n_iter = ggml_get_op_params_i32(dst, 1);
|
||||
|
||||
const int block_size = 256;
|
||||
const dim3 block_dims(block_size, 1, 1);
|
||||
const dim3 grid_dims((n_tokens + block_size - 1) / block_size, 1, 1);
|
||||
const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, ctx.stream());
|
||||
|
||||
ggml_cuda_kernel_launch(dsv4_hc_comb_f32, launch_params,
|
||||
(const float *) mixes->data, (const float *) scale->data, (const float *) base->data, (float *) dst->data,
|
||||
n_tokens,
|
||||
nbm0 / sizeof(float), nbm1 / sizeof(float),
|
||||
nbs0 / sizeof(float),
|
||||
nbb0 / sizeof(float),
|
||||
nbd0 / sizeof(float), nbd1 / sizeof(float), nbd2 / sizeof(float),
|
||||
eps, n_iter);
|
||||
}
|
||||
|
||||
void ggml_cuda_op_dsv4_hc_pre(ggml_backend_cuda_context & ctx, 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);
|
||||
|
||||
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 int64_t n_embd = x->ne[0];
|
||||
const int64_t hc = x->ne[1];
|
||||
const int64_t n_tokens = x->ne[2];
|
||||
|
||||
const int block_size = 256;
|
||||
const int64_t nr = n_embd * n_tokens;
|
||||
const dim3 block_dims(block_size, 1, 1);
|
||||
const dim3 grid_dims((nr + block_size - 1) / block_size, 1, 1);
|
||||
const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, ctx.stream());
|
||||
|
||||
ggml_cuda_kernel_launch(dsv4_hc_pre_f32, launch_params,
|
||||
(const float *) x->data, (const float *) weights->data, (float *) dst->data,
|
||||
n_embd, hc, n_tokens,
|
||||
nbx0 / sizeof(float), nbx1 / sizeof(float), nbx2 / sizeof(float),
|
||||
nbw0 / sizeof(float), nbw1 / sizeof(float),
|
||||
nbd0 / sizeof(float), nbd1 / sizeof(float));
|
||||
}
|
||||
|
||||
void ggml_cuda_op_dsv4_hc_post(ggml_backend_cuda_context & ctx, 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);
|
||||
|
||||
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 int64_t n_embd = x->ne[0];
|
||||
const int64_t n_tokens = x->ne[1];
|
||||
const int64_t hc = residual->ne[1];
|
||||
|
||||
const int block_size = 256;
|
||||
const int64_t nr = n_embd * hc * n_tokens;
|
||||
const dim3 block_dims(block_size, 1, 1);
|
||||
const dim3 grid_dims((nr + block_size - 1) / block_size, 1, 1);
|
||||
const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, ctx.stream());
|
||||
|
||||
ggml_cuda_kernel_launch(dsv4_hc_post_f32, launch_params,
|
||||
(const float *) x->data, (const float *) residual->data,
|
||||
(const float *) post->data, (const float *) comb->data, (float *) dst->data,
|
||||
n_embd, hc, n_tokens,
|
||||
nbx0 / sizeof(float), nbx1 / sizeof(float),
|
||||
nbr0 / sizeof(float), nbr1 / sizeof(float), nbr2 / sizeof(float),
|
||||
nbp0 / sizeof(float), nbp1 / sizeof(float),
|
||||
nbc0 / sizeof(float), nbc1 / sizeof(float), nbc2 / sizeof(float),
|
||||
nbd0 / sizeof(float), nbd1 / sizeof(float), nbd2 / sizeof(float));
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
#include "common.cuh"
|
||||
#include "ggml.h"
|
||||
|
||||
void ggml_cuda_op_dsv4_hc_comb(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
|
||||
void ggml_cuda_op_dsv4_hc_pre(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
|
||||
void ggml_cuda_op_dsv4_hc_post(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
|
||||
@@ -58,6 +58,7 @@
|
||||
#include "ggml-cuda/wkv.cuh"
|
||||
#include "ggml-cuda/gla.cuh"
|
||||
#include "ggml-cuda/gated_delta_net.cuh"
|
||||
#include "ggml-cuda/dsv4-hc.cuh"
|
||||
#include "ggml-cuda/set.cuh"
|
||||
#include "ggml-cuda/set-rows.cuh"
|
||||
#include "ggml-cuda/pad_reflect_1d.cuh"
|
||||
@@ -2319,6 +2320,15 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
|
||||
case GGML_OP_GATED_DELTA_NET:
|
||||
ggml_cuda_op_gated_delta_net(ctx, dst);
|
||||
break;
|
||||
case GGML_OP_DSV4_HC_COMB:
|
||||
ggml_cuda_op_dsv4_hc_comb(ctx, dst);
|
||||
break;
|
||||
case GGML_OP_DSV4_HC_PRE:
|
||||
ggml_cuda_op_dsv4_hc_pre(ctx, dst);
|
||||
break;
|
||||
case GGML_OP_DSV4_HC_POST:
|
||||
ggml_cuda_op_dsv4_hc_post(ctx, dst);
|
||||
break;
|
||||
case GGML_OP_RWKV_WKV7:
|
||||
ggml_cuda_op_rwkv_wkv7(ctx, dst);
|
||||
break;
|
||||
@@ -5088,6 +5098,16 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
|
||||
#else
|
||||
return true;
|
||||
#endif // GGML_USE_MUSA
|
||||
case GGML_OP_DSV4_HC_COMB:
|
||||
return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 &&
|
||||
op->src[2]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32;
|
||||
case GGML_OP_DSV4_HC_PRE:
|
||||
return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 &&
|
||||
op->type == GGML_TYPE_F32;
|
||||
case GGML_OP_DSV4_HC_POST:
|
||||
return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 &&
|
||||
op->src[2]->type == GGML_TYPE_F32 && op->src[3]->type == GGML_TYPE_F32 &&
|
||||
op->type == GGML_TYPE_F32;
|
||||
case GGML_OP_FLASH_ATTN_EXT:
|
||||
return ggml_cuda_flash_attn_ext_supported(dev_ctx->device, op);
|
||||
case GGML_OP_CROSS_ENTROPY_LOSS:
|
||||
|
||||
+135
-2
@@ -1080,6 +1080,9 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
|
||||
"SOLVE_TRI",
|
||||
"GATED_DELTA_NET",
|
||||
"LIGHTNING_INDEXER",
|
||||
"DSV4_HC_COMB",
|
||||
"DSV4_HC_PRE",
|
||||
"DSV4_HC_POST",
|
||||
|
||||
"UNARY",
|
||||
|
||||
@@ -1097,7 +1100,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
|
||||
"GLU",
|
||||
};
|
||||
|
||||
static_assert(GGML_OP_COUNT == 98, "GGML_OP_COUNT != 98");
|
||||
static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101");
|
||||
|
||||
static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
|
||||
"none",
|
||||
@@ -1192,6 +1195,9 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
|
||||
"A X = B, A triangular, solve X",
|
||||
"gated_delta_net(q, k, v, g, beta, s)",
|
||||
"lightning_indexer(q, k, weights, mask)",
|
||||
"dsv4_hc_comb(mixes, scale, base)",
|
||||
"dsv4_hc_pre(x, weights)",
|
||||
"dsv4_hc_post(x, residual, post, comb)",
|
||||
|
||||
"unary(x)",
|
||||
|
||||
@@ -1209,7 +1215,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
|
||||
"glu(x)",
|
||||
};
|
||||
|
||||
static_assert(GGML_OP_COUNT == 98, "GGML_OP_COUNT != 98");
|
||||
static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101");
|
||||
|
||||
static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2");
|
||||
|
||||
@@ -5437,6 +5443,7 @@ struct ggml_tensor * ggml_flash_attn_ext(
|
||||
return result;
|
||||
}
|
||||
|
||||
|
||||
void ggml_flash_attn_ext_set_prec(
|
||||
struct ggml_tensor * a,
|
||||
enum ggml_prec prec) {
|
||||
@@ -6337,6 +6344,132 @@ struct ggml_tensor * ggml_lightning_indexer(
|
||||
return result;
|
||||
}
|
||||
|
||||
// ggml_dsv4_hc_comb
|
||||
|
||||
struct ggml_tensor * ggml_dsv4_hc_comb(
|
||||
struct ggml_context * ctx,
|
||||
struct ggml_tensor * mixes,
|
||||
struct ggml_tensor * scale,
|
||||
struct ggml_tensor * base,
|
||||
float eps,
|
||||
int32_t n_iter) {
|
||||
GGML_ASSERT(mixes->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(scale->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(base->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(n_iter > 0);
|
||||
|
||||
const int64_t hc_mix_dim = mixes->ne[0];
|
||||
const int64_t n_tokens = mixes->ne[1];
|
||||
|
||||
int64_t hc = 0;
|
||||
for (int64_t i = 1; i*i + 2*i <= hc_mix_dim; ++i) {
|
||||
if ((2 + i)*i == hc_mix_dim) {
|
||||
hc = i;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
GGML_ASSERT(hc > 0);
|
||||
GGML_ASSERT(hc == 4);
|
||||
GGML_ASSERT(mixes->ne[2] == 1);
|
||||
GGML_ASSERT(mixes->ne[3] == 1);
|
||||
GGML_ASSERT(scale->ne[0] >= 3);
|
||||
GGML_ASSERT(scale->ne[1] == 1);
|
||||
GGML_ASSERT(scale->ne[2] == 1);
|
||||
GGML_ASSERT(scale->ne[3] == 1);
|
||||
GGML_ASSERT(base->ne[0] == hc_mix_dim);
|
||||
GGML_ASSERT(base->ne[1] == 1);
|
||||
GGML_ASSERT(base->ne[2] == 1);
|
||||
GGML_ASSERT(base->ne[3] == 1);
|
||||
|
||||
struct ggml_tensor * result = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, hc, hc, n_tokens);
|
||||
|
||||
ggml_set_op_params_f32(result, 0, eps);
|
||||
ggml_set_op_params_i32(result, 1, n_iter);
|
||||
|
||||
result->op = GGML_OP_DSV4_HC_COMB;
|
||||
result->src[0] = mixes;
|
||||
result->src[1] = scale;
|
||||
result->src[2] = base;
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
// ggml_dsv4_hc_pre
|
||||
|
||||
struct ggml_tensor * ggml_dsv4_hc_pre(
|
||||
struct ggml_context * ctx,
|
||||
struct ggml_tensor * x,
|
||||
struct ggml_tensor * weights) {
|
||||
GGML_ASSERT(x->type == GGML_TYPE_F32);
|
||||
GGML_ASSERT(weights->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(hc > 0);
|
||||
GGML_ASSERT(x->ne[3] == 1);
|
||||
GGML_ASSERT(weights->ne[0] == hc);
|
||||
GGML_ASSERT(weights->ne[1] == n_tokens);
|
||||
GGML_ASSERT(weights->ne[2] == 1);
|
||||
GGML_ASSERT(weights->ne[3] == 1);
|
||||
|
||||
struct ggml_tensor * result = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_tokens);
|
||||
|
||||
result->op = GGML_OP_DSV4_HC_PRE;
|
||||
result->src[0] = x;
|
||||
result->src[1] = weights;
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
// ggml_dsv4_hc_post
|
||||
|
||||
struct ggml_tensor * ggml_dsv4_hc_post(
|
||||
struct ggml_context * ctx,
|
||||
struct ggml_tensor * x,
|
||||
struct ggml_tensor * residual,
|
||||
struct ggml_tensor * post,
|
||||
struct ggml_tensor * comb) {
|
||||
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);
|
||||
|
||||
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(hc > 0);
|
||||
GGML_ASSERT(x->ne[2] == 1);
|
||||
GGML_ASSERT(x->ne[3] == 1);
|
||||
|
||||
GGML_ASSERT(residual->ne[0] == n_embd);
|
||||
GGML_ASSERT(residual->ne[2] == n_tokens);
|
||||
GGML_ASSERT(residual->ne[3] == 1);
|
||||
|
||||
GGML_ASSERT(post->ne[0] == hc);
|
||||
GGML_ASSERT(post->ne[1] == n_tokens);
|
||||
GGML_ASSERT(post->ne[2] == 1);
|
||||
GGML_ASSERT(post->ne[3] == 1);
|
||||
|
||||
GGML_ASSERT(comb->ne[0] == hc);
|
||||
GGML_ASSERT(comb->ne[1] == hc);
|
||||
GGML_ASSERT(comb->ne[2] == n_tokens);
|
||||
GGML_ASSERT(comb->ne[3] == 1);
|
||||
|
||||
struct ggml_tensor * result = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, hc, n_tokens);
|
||||
|
||||
result->op = GGML_OP_DSV4_HC_POST;
|
||||
result->src[0] = x;
|
||||
result->src[1] = residual;
|
||||
result->src[2] = post;
|
||||
result->src[3] = comb;
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
struct ggml_hash_set ggml_hash_set_new(size_t size) {
|
||||
|
||||
Reference in New Issue
Block a user