Refactor mhc helpers, graph base

Author: timkhronos

src/llama-graph.cpp

@@ -3410,181 +3410,6 @@ llm_graph_input_dsv4 * llm_graph_context::build_inp_dsv4() const {
     return (llm_graph_input_dsv4 *) res->add_input(std::move(inp));
 }
 
-// manifold-constrained hyper-connections (mHC), deepseek4 and glm5-next
-
-static ggml_tensor * hc_view_1d(ggml_context * ctx, ggml_tensor * t, int64_t ne0, int64_t i0) {
-    return ggml_view_1d(ctx, t, ne0, ggml_row_size(t->type, i0));
-}
-
-static ggml_tensor * hc_view_2d(ggml_context * ctx, ggml_tensor * t, int64_t ne0, int64_t ne1, int64_t i0) {
-    return ggml_view_2d(ctx, t, ne0, ne1, t->nb[1], ggml_row_size(t->type, i0));
-}
-
-static ggml_tensor * hc_affine(ggml_context * ctx, ggml_tensor * x, ggml_tensor * scale, ggml_tensor * base) {
-    x = ggml_mul(ctx, x, scale);
-    x = ggml_add(ctx, x, base);
-    return x;
-}
-
-ggml_tensor * llm_graph_context::build_hc_pre(
-        ggml_tensor * x,
-        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];
-
-    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);
-        result = result ? ggml_add(ctx0, result, cur) : cur;
-    }
-
-    return result;
-}
-
-ggml_tensor * llm_graph_context::build_hc_sinkhorn(
-        ggml_tensor * comb,
-        int           il) const {
-    GGML_UNUSED(il);
-
-    // comb is [dst_hc, src_hc, n_tokens]. Sinkhorn follows the reference:
-    // row softmax over dst, one column normalization, then repeated row/column normalization.
-    comb = ggml_soft_max(ctx0, comb);
-
-    ggml_tensor * eps = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, 1);
-    eps = ggml_fill(ctx0, eps, hparams.dsv4_hc_eps);
-
-    comb = ggml_add(ctx0, comb, eps);
-
-    auto norm_cols = [&]() {
-        ggml_tensor * comb_src_dst = ggml_cont(ctx0, ggml_permute(ctx0, comb, 1, 0, 2, 3));
-        ggml_tensor * col_sum = ggml_sum_rows(ctx0, comb_src_dst);
-        col_sum = ggml_add(ctx0, col_sum, eps);
-        col_sum = ggml_permute(ctx0, col_sum, 1, 0, 2, 3);
-        comb = ggml_div(ctx0, comb, col_sum);
-    };
-
-    auto norm_rows = [&]() {
-        ggml_tensor * row_sum = ggml_sum_rows(ctx0, comb);
-        row_sum = ggml_add(ctx0, row_sum, eps);
-        comb = ggml_div(ctx0, comb, row_sum);
-    };
-
-    norm_cols();
-    for (uint32_t i = 1; i < hparams.dsv4_hc_sinkhorn_iters; ++i) {
-        norm_rows();
-        norm_cols();
-    }
-
-    return comb;
-}
-
-ggml_tensor * llm_graph_context::build_hc_pre(
-        ggml_tensor * x,
-        ggml_tensor * hc_fn,
-        ggml_tensor * hc_scale,
-        ggml_tensor * hc_base,
-        ggml_tensor ** post,
-        ggml_tensor ** comb,
-        int il) const {
-    const int64_t hc         = hparams.dsv4_hc_mult;
-    const int64_t hc_dim     = hc*n_embd;
-    const int64_t hc_mix_dim = (2 + hc)*hc;
-    const int64_t nt         = x->ne[2];
-
-    GGML_ASSERT(hc == 4);
-    GGML_ASSERT(hc_fn->ne[1] == hc_mix_dim);
-
-    ggml_tensor * flat = ggml_reshape_2d(ctx0, x, hc_dim, nt);
-    ggml_tensor * flat_norm = ggml_rms_norm(ctx0, flat, norm_rms_eps);
-    ggml_tensor * mixes = ggml_mul_mat(ctx0, hc_fn, flat_norm);
-    cb(mixes, "hc_mixes", il);
-
-    ggml_tensor * scale_pre  = hc_view_1d(ctx0, hc_scale, 1, 0);
-    ggml_tensor * scale_post = hc_view_1d(ctx0, hc_scale, 1, 1);
-
-    ggml_tensor * base_pre  = hc_view_1d(ctx0, hc_base, hc, 0);
-    ggml_tensor * base_post = hc_view_1d(ctx0, hc_base, hc, hc);
-
-    ggml_tensor * pre = hc_view_2d(ctx0, mixes, hc, nt, 0);
-    pre = hc_affine(ctx0, pre, scale_pre, base_pre);
-    pre = ggml_sigmoid(ctx0, pre);
-    pre = ggml_scale_bias(ctx0, pre, 1.0f, hparams.dsv4_hc_eps);
-    cb(pre, "hc_pre", il);
-
-    *post = hc_view_2d(ctx0, mixes, hc, nt, hc);
-    *post = hc_affine(ctx0, *post, scale_post, base_post);
-    *post = ggml_sigmoid(ctx0, *post);
-    *post = ggml_scale(ctx0, *post, 2.0f);
-    cb(*post, "hc_post", 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 = hc_view_1d(ctx0, hc_scale, 1, 2);
-        ggml_tensor * base_comb  = hc_view_1d(ctx0, hc_base, hc*hc, 2*hc);
-
-        *comb = hc_view_2d(ctx0, mixes, hc*hc, nt, 2*hc);
-        *comb = 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);
-
-    ggml_tensor * result = build_hc_pre(x, pre, il);
-    return result;
-}
-
-ggml_tensor * llm_graph_context::build_hc_post(
-        ggml_tensor * x,
-        ggml_tensor * residual,
-        ggml_tensor * post,
-        ggml_tensor * comb,
-        int il) const {
-    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];
-
-    ggml_tensor * out = nullptr;
-    for (int64_t dst = 0; dst < hc; ++dst) {
-        ggml_tensor * post_dst = ggml_view_2d(ctx0, post, 1, nt, post->nb[1], dst*post->nb[0]);
-        ggml_tensor * cur = ggml_mul(ctx0, x, post_dst);
-
-        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]);
-            cur = ggml_add(ctx0, cur, ggml_mul(ctx0, res_src, comb_src_dst));
-        }
-
-        cur = ggml_reshape_3d(ctx0, cur, n_embd, 1, nt);
-        out = out ? ggml_concat(ctx0, out, cur, 1) : cur;
-    }
-
-    return out;
-}
-
 ggml_tensor * llm_graph_context::build_rs(
         ggml_tensor * s,
         ggml_tensor * state_copy_main,

src/llama-graph.h

@@ -1339,36 +1339,6 @@ struct llm_graph_context {
     // hybrid
     //
 
-    // hyper-connections (mHC)
-
-
-    // collapse the hc streams with per-stream weights
-    ggml_tensor * build_hc_pre(
-            ggml_tensor * x,
-            ggml_tensor * weights,
-            int il) const;
-
-    // returns the collapsed input and fills the post / comb weights
-    ggml_tensor * build_hc_pre(
-            ggml_tensor * x,
-            ggml_tensor * hc_fn,
-            ggml_tensor * hc_scale,
-            ggml_tensor * hc_base,
-            ggml_tensor ** post,
-            ggml_tensor ** comb,
-            int il) const;
-
-    ggml_tensor * build_hc_post(
-            ggml_tensor * x,
-            ggml_tensor * residual,
-            ggml_tensor * post,
-            ggml_tensor * comb,
-            int il) const;
-
-    ggml_tensor * build_hc_sinkhorn(
-            ggml_tensor * comb,
-            int il) const;
-
     llm_graph_input_mem_hybrid * build_inp_mem_hybrid() const;
     llm_graph_input_mem_hybrid_k * build_inp_mem_hybrid_k() const;