diff --git a/src/models/models.h b/src/models/models.h index af60764c2f..9b87a40d5a 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -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); diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp index acfdd5b500..abf6a0502f 100644 --- a/src/models/qwen4exp.cpp +++ b/src/models/qwen4exp.cpp @@ -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);