diff --git a/src/llama-context.cpp b/src/llama-context.cpp index eed041eef4..3679931c12 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2472,11 +2472,12 @@ llm_graph_cb llama_context::graph_get_cb() const { ggml_set_name(cur, name); } - // norm may be automatically assigned to the backend of the previous layer, increasing data transfer between backends + // - norm may be automatically assigned to the backend of the previous layer, increasing data transfer between backends + // - force the last op of the layer on the specified backend to avoid running it on the backend of the next layer due to scheduling // FIXME: fix in ggml_backend_sched const bool full_offload = model.n_gpu_layers() > model.hparams.n_layer_all; if (ubatch.n_tokens < 32 || full_offload) { - if (il != -1 && strcmp(name, "norm") == 0) { + if (il != -1 && (strcmp(name, "norm") == 0 || strcmp(name, "l_last") == 0)) { const auto & dev_layer = model.dev_layer(il); for (const auto & backend : backends) { if (ggml_backend_get_device(backend.get()) == dev_layer) { diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index 5ad6473ce2..2d41dace0b 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -1133,6 +1133,10 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p &post, &comb, il); cb(cur, "hc_ffn_pre", il); + ggml_build_forward_expand(gf, residual); + ggml_build_forward_expand(gf, post); + ggml_build_forward_expand(gf, comb); + cur = build_norm(cur, model.layers[il].ffn_norm, nullptr, LLM_NORM_RMS, il); cb(cur, "ffn_norm", il); @@ -1175,7 +1179,7 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p inpL = build_hc_post(cur, residual, post, comb, il); inpL = build_cvec(inpL, il); - cb(inpL, "l_out", il); + cb(inpL, "l_last", il); } if (inp_out_ids) {