model: qwen4exp: reduce number of graph splits (#27880)

This commit is contained in:
Xuan-Son Nguyen
2026-08-28 19:24:46 +02:00
committed by GitHub
parent b387ddfd84
commit 6fe7498016
2 changed files with 27 additions and 10 deletions
+4 -1
View File
@@ -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);
+23 -9
View File
@@ -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);