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:
Aman Gupta
2026-07-17 00:33:33 +08:00
committed by GitHub
parent b2dd28a3b6
commit 0dc74e332e
16 changed files with 1062 additions and 22 deletions
+31
View File
@@ -61,6 +61,24 @@ static const llm_fused_op_probe llm_fused_op_lid_probe = {
/*.n_tokens_per_seq =*/ 1,
};
static const llm_fused_op_probe llm_fused_op_dsv4_hc_pre_probe = {
/*.op =*/ LLM_FUSED_OP_DSV4_HC_PRE,
/*.name =*/ "fused DeepSeek V4 HC pre",
/*.n_tokens_per_seq =*/ 1,
};
static const llm_fused_op_probe llm_fused_op_dsv4_hc_comb_probe = {
/*.op =*/ LLM_FUSED_OP_DSV4_HC_COMB,
/*.name =*/ "fused DeepSeek V4 HC comb",
/*.n_tokens_per_seq =*/ 1,
};
static const llm_fused_op_probe llm_fused_op_dsv4_hc_post_probe = {
/*.op =*/ LLM_FUSED_OP_DSV4_HC_POST,
/*.name =*/ "fused DeepSeek V4 HC post",
/*.n_tokens_per_seq =*/ 1,
};
llama_context::llama_context(
const llama_model & model,
llama_context_params params) :
@@ -235,6 +253,11 @@ llama_context::llama_context(
cparams.fused_lid = true;
cparams.auto_flid = true;
cparams.fused_dsv4_hc_pre = true;
cparams.fused_dsv4_hc_comb = true;
cparams.fused_dsv4_hc_post = true;
cparams.auto_fhc = true;
// with causal attention, the batch size is limited by the context size
cparams.n_batch = cparams.causal_attn ? std::min(cparams.n_ctx, params.n_batch) : params.n_batch;
@@ -537,6 +560,14 @@ void llama_context::resolve_fused_ops(const llama_memory_context_i * mctx, uint3
resolve(llm_fused_op_lid_probe, cparams.fused_lid);
cparams.auto_flid = false;
}
if (cparams.auto_fhc) {
LLAMA_LOG_INFO("%s: resolving fused DeepSeek V4 HC support:\n", func);
resolve(llm_fused_op_dsv4_hc_pre_probe, cparams.fused_dsv4_hc_pre);
resolve(llm_fused_op_dsv4_hc_comb_probe, cparams.fused_dsv4_hc_comb);
resolve(llm_fused_op_dsv4_hc_post_probe, cparams.fused_dsv4_hc_post);
cparams.auto_fhc = false;
}
}
void llama_context::sched_reserve() {
+4
View File
@@ -43,6 +43,10 @@ struct llama_cparams {
bool auto_fgdn;
bool fused_lid; // use fused lightning indexer
bool auto_flid;
bool fused_dsv4_hc_pre;
bool fused_dsv4_hc_comb;
bool fused_dsv4_hc_post;
bool auto_fhc;
bool no_perf;
bool warmup; // TODO: remove [TAG_LLAMA_GRAPH_NO_WARMUP]
bool op_offload;
+3
View File
@@ -43,6 +43,9 @@ enum llm_fused_op {
LLM_FUSED_OP_GDN_AR,
LLM_FUSED_OP_GDN_CH,
LLM_FUSED_OP_LIGHTNING_INDEXER,
LLM_FUSED_OP_DSV4_HC_PRE,
LLM_FUSED_OP_DSV4_HC_COMB,
LLM_FUSED_OP_DSV4_HC_POST,
};
enum llm_ffn_op_type : int {
+41 -16
View File
@@ -197,22 +197,31 @@ static ggml_tensor * dsv4_hc_affine(
return x;
}
ggml_tensor * llama_model_deepseek4::graph::build_hc_weighted_sum(
ggml_tensor * llama_model_deepseek4::graph::build_hc_pre(
ggml_tensor * x,
ggml_tensor * weights) const {
ggml_tensor * weights,
int il) const {
GGML_ASSERT(x->ne[0] == n_embd);
GGML_ASSERT(x->ne[1] == hparams.dsv4_hc_mult);
const int64_t hc = hparams.dsv4_hc_mult;
const int64_t nt = x->ne[2];
ggml_tensor * acc = nullptr;
if (cparams.fused_dsv4_hc_pre && il >= 0) {
ggml_tensor * result = ggml_dsv4_hc_pre(ctx0, x, weights);
res->add_fused_node({LLM_FUSED_OP_DSV4_HC_PRE, result, il});
return result;
}
ggml_tensor * result = nullptr;
for (int64_t ih = 0; ih < hc; ++ih) {
ggml_tensor * xh = ggml_view_2d(ctx0, x, n_embd, nt, x->nb[2], ih*x->nb[1]);
ggml_tensor * wh = ggml_view_2d(ctx0, weights, 1, nt, weights->nb[1], ih*weights->nb[0]);
ggml_tensor * cur = ggml_mul(ctx0, xh, wh);
acc = acc ? ggml_add(ctx0, acc, cur) : cur;
result = result ? ggml_add(ctx0, result, cur) : cur;
}
return acc;
return result;
}
ggml_tensor * llama_model_deepseek4::graph::build_hc_sinkhorn(
@@ -275,11 +284,9 @@ ggml_tensor * llama_model_deepseek4::graph::build_hc_pre(
ggml_tensor * scale_pre = dsv4_view_1d(ctx0, hc_scale, 1, 0);
ggml_tensor * scale_post = dsv4_view_1d(ctx0, hc_scale, 1, 1);
ggml_tensor * scale_comb = dsv4_view_1d(ctx0, hc_scale, 1, 2);
ggml_tensor * base_pre = dsv4_view_1d(ctx0, hc_base, hc, 0);
ggml_tensor * base_post = dsv4_view_1d(ctx0, hc_base, hc, hc);
ggml_tensor * base_comb = dsv4_view_1d(ctx0, hc_base, hc*hc, 2*hc);
ggml_tensor * pre = dsv4_view_2d(ctx0, mixes, hc, nt, 0);
pre = dsv4_hc_affine(ctx0, pre, scale_pre, base_pre);
@@ -293,13 +300,23 @@ ggml_tensor * llama_model_deepseek4::graph::build_hc_pre(
*post = ggml_scale(ctx0, *post, 2.0f);
cb(*post, "hc_post", il);
*comb = dsv4_view_2d(ctx0, mixes, hc*hc, nt, 2*hc);
*comb = dsv4_hc_affine(ctx0, *comb, scale_comb, base_comb);
*comb = ggml_reshape_3d(ctx0, *comb, hc, hc, nt);
*comb = build_hc_sinkhorn(*comb, il);
if (cparams.fused_dsv4_hc_comb) {
*comb = ggml_dsv4_hc_comb(ctx0, mixes, hc_scale, hc_base, hparams.dsv4_hc_eps,
(int32_t) hparams.dsv4_hc_sinkhorn_iters);
res->add_fused_node({LLM_FUSED_OP_DSV4_HC_COMB, *comb, il});
} else {
ggml_tensor * scale_comb = dsv4_view_1d(ctx0, hc_scale, 1, 2);
ggml_tensor * base_comb = dsv4_view_1d(ctx0, hc_base, hc*hc, 2*hc);
*comb = dsv4_view_2d(ctx0, mixes, hc*hc, nt, 2*hc);
*comb = dsv4_hc_affine(ctx0, *comb, scale_comb, base_comb);
*comb = ggml_reshape_3d(ctx0, *comb, hc, hc, nt);
*comb = build_hc_sinkhorn(*comb, il);
}
cb(*comb, "hc_comb", il);
return build_hc_weighted_sum(x, pre);
ggml_tensor * result = build_hc_pre(x, pre, il);
return result;
}
ggml_tensor * llama_model_deepseek4::graph::build_hc_post(
@@ -308,7 +325,14 @@ ggml_tensor * llama_model_deepseek4::graph::build_hc_post(
ggml_tensor * post,
ggml_tensor * comb,
int il) const {
GGML_UNUSED(il);
GGML_ASSERT(x->ne[0] == n_embd);
GGML_ASSERT(residual->ne[1] == hparams.dsv4_hc_mult);
if (cparams.fused_dsv4_hc_post) {
ggml_tensor * result = ggml_dsv4_hc_post(ctx0, x, residual, post, comb);
res->add_fused_node({LLM_FUSED_OP_DSV4_HC_POST, result, il});
return result;
}
const int64_t hc = hparams.dsv4_hc_mult;
const int64_t nt = x->ne[1];
@@ -320,7 +344,8 @@ ggml_tensor * llama_model_deepseek4::graph::build_hc_post(
for (int64_t src = 0; src < hc; ++src) {
ggml_tensor * res_src = ggml_view_2d(ctx0, residual, n_embd, nt, residual->nb[2], src*residual->nb[1]);
ggml_tensor * comb_src_dst = ggml_view_2d(ctx0, comb, 1, nt, comb->nb[2], dst*comb->nb[0] + src*comb->nb[1]);
ggml_tensor * comb_src_dst = ggml_view_2d(ctx0, comb, 1, nt, comb->nb[2],
dst*comb->nb[0] + src*comb->nb[1]);
cur = ggml_add(ctx0, cur, ggml_mul(ctx0, res_src, comb_src_dst));
}
@@ -350,7 +375,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_hc_head(
pre = ggml_scale_bias(ctx0, pre, 1.0f, hparams.dsv4_hc_eps);
cb(pre, "hc_head_pre", -1);
return build_hc_weighted_sum(x, pre);
return build_hc_pre(x, pre, -1);
}
ggml_tensor * llama_model_deepseek4::graph::build_hca_compressed_kv_from_state(
+3 -2
View File
@@ -1187,9 +1187,10 @@ struct llama_model_deepseek4 : public llama_model_base {
float kq_scale,
int il) const;
ggml_tensor * build_hc_weighted_sum(
ggml_tensor * build_hc_pre(
ggml_tensor * x,
ggml_tensor * weights) const;
ggml_tensor * weights,
int il) const;
ggml_tensor * build_hc_sinkhorn(
ggml_tensor * comb,