ggml-org/llama.cpp · #27880

model: qwen4exp: reduce number of graph splits

ngxson · merged Aug 28, 20262 files · 27 + / 10
src/models/models.h4 + / 1
@@ -2360,9 +2360,12 @@ struct llama_model_qwen4exp : public llama_model_base {                         int64_t   channels,                             int   il); +        ggml_tensor * build_inp_ple(+  const llama_memory_hybrid_idx_context * mctx_hyb);+         ggml_tensor * build_ple(              llm_graph_input_rs * inp,-  const llama_memory_hybrid_idx_context * mctx_hyb,+                    ggml_tensor * emb,                     ggml_tensor * hidden,                             int   il); 
src/models/qwen4exp.cpp23 + / 9
@@ -296,6 +296,7 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa      ggml_tensor * inpL = build_inp_embd(model.tok_embd);     cb(inpL, "model.input_embed", -1);+    ggml_build_forward_expand(gf, inpL);      auto * inp = build_inp_mem_hybrid(); @@ -312,6 +313,13 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa     ggml_tensor * inp_pos     = build_inp_pos();     ggml_tensor * inp_out_ids = build_inp_out_ids(); +    ggml_tensor * ple_emb = nullptr;+    if (hparams.ple_n_heads > 0) {+        ple_emb = build_inp_ple(mctx_hyb);+        // make sure ple_emb and build_inp_embd are in the same graph split+        ggml_build_forward_expand(gf, ple_emb);+    }+     // the wide residual starts as hc identical copies of the embedding     ggml_tensor * res_hc = ggml_repeat_4d(ctx0,             ggml_reshape_3d(ctx0, inpL, n_embd, 1, n_tokens),@@ -322,7 +330,7 @@ llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_pa         res->t_layer_inp[il] = res_hc;          if (hparams.is_ple(il)) {-            res_hc = build_ple(inp->get_recr(), mctx_hyb, res_hc, il);+            res_hc = build_ple(inp->get_recr(), ple_emb, res_hc, il);         }          ggml_tensor * inject = nullptr;@@ -1090,13 +1098,8 @@ ggml_tensor * llama_model_qwen4exp::graph::build_conv_state_at(     return conv_input; } -ggml_tensor * llama_model_qwen4exp::graph::build_ple(-        llm_graph_input_rs * inp,-        const llama_memory_hybrid_idx_context * mctx_hyb,-        ggml_tensor *        hidden,-        int                  il) {-    const int64_t hc      = hparams.dsv4_hc_mult;-    const int64_t hc_dim  = hc * n_embd;+ggml_tensor * llama_model_qwen4exp::graph::build_inp_ple(+        const llama_memory_hybrid_idx_context * mctx_hyb) {     const int64_t n_heads = hparams.ple_n_heads;      // the attention cells see every ubatch regardless of the layer types@@ -1111,7 +1114,18 @@ ggml_tensor * llama_model_qwen4exp::graph::build_ple(     // gather then flatten the heads: get_rows lays the head dimension out slowest, as the reference does     ggml_tensor * emb = ggml_get_rows(ctx0, model.per_layer_tok_embd, rows);     emb = ggml_reshape_2d(ctx0, emb, hparams.ple_head_dim * n_heads, n_tokens);-    cb(emb, "ple_embd", il);+    cb(emb, "ple_embd", -1);++    return emb;+}++ggml_tensor * llama_model_qwen4exp::graph::build_ple(+        llm_graph_input_rs * inp,+        ggml_tensor *        emb,+        ggml_tensor *        hidden,+        int                  il) {+    const int64_t hc      = hparams.dsv4_hc_mult;+    const int64_t hc_dim  = hc * n_embd;      ggml_tensor * key   = build_lora_mm(model.layers[il].ple_key,   emb);     ggml_tensor * value = build_lora_mm(model.layers[il].ple_value, emb);