Skip to content

Commit c0a925a

Browse files
Alex7MVAlex7MV
committed
feat: deep-copy Phase 2 graph for persistent replay (#huanglawsded)
- deep_copy_phase2_graph() in moe-hijacker duplicates all 4790 nodes and 1163 leafs into stable heap context, preserving names/ops/data - REPLAY path (Token 2+) skips model.build_graph() entirely, uses persistent_gf directly — zero name parsing, zero node creation - BUILD path (Token 1) builds graph, cascade-forces, deep-copies, caches persistent copy for all future tokens - galloc fast-path reuses sched_phase2 layout on Token 2+ - n_reused increments on every REPLAY token
1 parent 1daed02 commit c0a925a

3 files changed

Lines changed: 214 additions & 66 deletions

File tree

‎src/llama-context.cpp‎

Lines changed: 109 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -1817,76 +1817,121 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll
18171817
const bool do_cuda = (ubatch.n_tokens == 1);
18181818
ggml_cgraph * phase2_gf;
18191819

1820-
// Unified native compute via isolated sched_phase2.
1821-
// galloc fast-path (backend_ids_changed=false) skips reserve_n on Token 2+.
1822-
if (!sched_phase2) {
1823-
const size_t phase2_n = std::max((size_t)graph_max_nodes(ubatch.n_tokens), (size_t)10000);
1824-
sched_phase2.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), phase2_n, false, cparams.op_offload));
1825-
}
1826-
res->reset();
1827-
ggml_backend_sched_reset(sched_phase2.get());
1828-
moe_weight_cache.build_layer_only = -1;
1829-
moe_weight_cache.build_phase = 2;
1830-
phase2_gf = model.build_graph(gparams, nullptr, nullptr, &moe_weight_cache);
1831-
if (!phase2_gf) { ret = GGML_STATUS_FAILED; return nullptr; }
1832-
ggml_backend_sched_reset(sched_phase2.get());
1833-
for (int i = 0; i < ggml_graph_n_nodes(phase2_gf); i++) {
1834-
ggml_tensor * t = ggml_graph_node(phase2_gf, i);
1835-
auto [tid, il] = moe::match_hijack_name(h2_hijack, t->name);
1836-
if (tid >= 0)
1837-
ggml_backend_sched_set_tensor_backend(sched_phase2.get(), t, gpu);
1838-
}
1839-
for (int i = 0; i < ggml_graph_n_leafs(phase2_gf); i++) {
1840-
ggml_tensor * t = ggml_graph_leaf(phase2_gf, i);
1841-
auto [tid, il] = moe::match_hijack_name(h2_hijack, t->name);
1842-
if (tid >= 0)
1843-
ggml_backend_sched_set_tensor_backend(sched_phase2.get(), t, gpu);
1844-
}
1845-
moe::cascade_force_moe_consumers(h2_hijack, phase2_gf, sched_phase2.get(), gpu);
1846-
if (backend_cpu) {
1847-
auto force = [&](ggml_tensor * t) {
1848-
if (t) ggml_backend_sched_set_tensor_backend(sched_phase2.get(), t, backend_cpu);
1849-
};
1850-
for (auto & inp : res->inputs) {
1851-
auto * base = inp.get();
1852-
if (auto * akv = dynamic_cast<llm_graph_input_attn_kv *>(base)) {
1853-
force(akv->self_k_idxs); force(akv->self_v_idxs); force(akv->self_kq_mask);
1854-
} else if (auto * ak = dynamic_cast<llm_graph_input_attn_k *>(base)) {
1855-
force(ak->self_k_idxs); force(ak->self_kq_mask);
1856-
} else if (auto * dsa = dynamic_cast<llm_graph_input_attn_k_dsa *>(base)) {
1857-
force(dsa->self_k_idxs_mla); force(dsa->self_k_idxs_lid);
1858-
force(dsa->self_kq_mask_mla); force(dsa->self_kq_mask_lid);
1859-
} else if (auto * iswa = dynamic_cast<llm_graph_input_attn_kv_iswa *>(base)) {
1860-
force(iswa->self_k_idxs); force(iswa->self_v_idxs);
1861-
force(iswa->self_k_idxs_swa); force(iswa->self_v_idxs_swa);
1862-
force(iswa->self_kq_mask); force(iswa->self_kq_mask_swa);
1863-
} else if (auto * hyb = dynamic_cast<llm_graph_input_mem_hybrid *>(base)) {
1864-
force(hyb->inp_attn->self_k_idxs); force(hyb->inp_attn->self_v_idxs);
1865-
force(hyb->inp_attn->self_kq_mask);
1866-
} else if (auto * hybk = dynamic_cast<llm_graph_input_mem_hybrid_k *>(base)) {
1867-
force(hybk->inp_attn->self_k_idxs); force(hybk->inp_attn->self_kq_mask);
1868-
} else if (auto * hybiswa = dynamic_cast<llm_graph_input_mem_hybrid_iswa *>(base)) {
1869-
force(hybiswa->inp_attn->self_k_idxs); force(hybiswa->inp_attn->self_v_idxs);
1870-
force(hybiswa->inp_attn->self_k_idxs_swa); force(hybiswa->inp_attn->self_v_idxs_swa);
1871-
force(hybiswa->inp_attn->self_kq_mask); force(hybiswa->inp_attn->self_kq_mask_swa);
1872-
} else if (auto * oid = dynamic_cast<llm_graph_input_out_ids *>(base)) {
1873-
force(oid->out_ids);
1820+
if (do_cuda && phase2_cache.valid) {
1821+
// REPLAY (Token 2+) — use persistent copy, skip build+cascade
1822+
res->reset();
1823+
ggml_backend_sched_reset(sched_phase2.get());
1824+
phase2_gf = phase2_cache.persistent_gf;
1825+
for (int i = 0; i < ggml_graph_n_nodes(phase2_gf); i++) {
1826+
ggml_tensor * t = ggml_graph_node(phase2_gf, i);
1827+
auto [tid, il] = moe::match_hijack_name(h2_hijack, t->name);
1828+
if (tid >= 0)
1829+
ggml_backend_sched_set_tensor_backend(sched_phase2.get(), t, gpu);
1830+
}
1831+
for (int i = 0; i < ggml_graph_n_leafs(phase2_gf); i++) {
1832+
ggml_tensor * t = ggml_graph_leaf(phase2_gf, i);
1833+
auto [tid, il] = moe::match_hijack_name(h2_hijack, t->name);
1834+
if (tid >= 0)
1835+
ggml_backend_sched_set_tensor_backend(sched_phase2.get(), t, gpu);
1836+
}
1837+
moe::cascade_force_moe_consumers(h2_hijack, phase2_gf, sched_phase2.get(), gpu);
1838+
if (backend_cpu) {
1839+
auto force = [&](ggml_tensor * t) {
1840+
if (t) ggml_backend_sched_set_tensor_backend(sched_phase2.get(), t, backend_cpu);
1841+
};
1842+
for (auto & inp : res->inputs) {
1843+
auto * base = inp.get();
1844+
if (auto * akv = dynamic_cast<llm_graph_input_attn_kv *>(base)) {
1845+
force(akv->self_k_idxs); force(akv->self_v_idxs); force(akv->self_kq_mask);
1846+
} else if (auto * ak = dynamic_cast<llm_graph_input_attn_k *>(base)) {
1847+
force(ak->self_k_idxs); force(ak->self_kq_mask);
1848+
} else if (auto * dsa = dynamic_cast<llm_graph_input_attn_k_dsa *>(base)) {
1849+
force(dsa->self_k_idxs_mla); force(dsa->self_k_idxs_lid);
1850+
force(dsa->self_kq_mask_mla); force(dsa->self_kq_mask_lid);
1851+
} else if (auto * iswa = dynamic_cast<llm_graph_input_attn_kv_iswa *>(base)) {
1852+
force(iswa->self_k_idxs); force(iswa->self_v_idxs);
1853+
force(iswa->self_k_idxs_swa); force(iswa->self_v_idxs_swa);
1854+
force(iswa->self_kq_mask); force(iswa->self_kq_mask_swa);
1855+
} else if (auto * hyb = dynamic_cast<llm_graph_input_mem_hybrid *>(base)) {
1856+
force(hyb->inp_attn->self_k_idxs); force(hyb->inp_attn->self_v_idxs);
1857+
force(hyb->inp_attn->self_kq_mask);
1858+
} else if (auto * hybk = dynamic_cast<llm_graph_input_mem_hybrid_k *>(base)) {
1859+
force(hybk->inp_attn->self_k_idxs); force(hybk->inp_attn->self_kq_mask);
1860+
} else if (auto * hybiswa = dynamic_cast<llm_graph_input_mem_hybrid_iswa *>(base)) {
1861+
force(hybiswa->inp_attn->self_k_idxs); force(hybiswa->inp_attn->self_v_idxs);
1862+
force(hybiswa->inp_attn->self_k_idxs_swa); force(hybiswa->inp_attn->self_v_idxs_swa);
1863+
force(hybiswa->inp_attn->self_kq_mask); force(hybiswa->inp_attn->self_kq_mask_swa);
1864+
} else if (auto * oid = dynamic_cast<llm_graph_input_out_ids *>(base)) {
1865+
force(oid->out_ids);
1866+
}
18741867
}
18751868
}
1876-
}
1877-
1878-
// alloc_graph: Token 1 = full reserve+alloc; Token 2+ = galloc fast-path (no reserve_n)
1879-
if (!ggml_backend_sched_alloc_graph(sched_phase2.get(), phase2_gf)) { ret = GGML_STATUS_ALLOC_FAILED; return nullptr; }
1880-
1881-
// Track cache hit for server metrics
1882-
if (do_cuda && phase2_cache.valid) {
18831869
n_reused++;
1884-
}
1885-
if (!phase2_cache.valid) {
1886-
phase2_cache.capture();
1870+
} else {
1871+
// BUILD (Token 1) — build, cascade, deep-copy
1872+
if (!sched_phase2) {
1873+
const size_t phase2_n = std::max((size_t)graph_max_nodes(ubatch.n_tokens), (size_t)10000);
1874+
sched_phase2.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), phase2_n, false, cparams.op_offload));
1875+
}
1876+
res->reset();
1877+
ggml_backend_sched_reset(sched_phase2.get());
1878+
moe_weight_cache.build_layer_only = -1;
1879+
moe_weight_cache.build_phase = 2;
1880+
phase2_gf = model.build_graph(gparams, nullptr, nullptr, &moe_weight_cache);
1881+
if (!phase2_gf) { ret = GGML_STATUS_FAILED; return nullptr; }
1882+
ggml_backend_sched_reset(sched_phase2.get());
1883+
for (int i = 0; i < ggml_graph_n_nodes(phase2_gf); i++) {
1884+
ggml_tensor * t = ggml_graph_node(phase2_gf, i);
1885+
auto [tid, il] = moe::match_hijack_name(h2_hijack, t->name);
1886+
if (tid >= 0)
1887+
ggml_backend_sched_set_tensor_backend(sched_phase2.get(), t, gpu);
1888+
}
1889+
for (int i = 0; i < ggml_graph_n_leafs(phase2_gf); i++) {
1890+
ggml_tensor * t = ggml_graph_leaf(phase2_gf, i);
1891+
auto [tid, il] = moe::match_hijack_name(h2_hijack, t->name);
1892+
if (tid >= 0)
1893+
ggml_backend_sched_set_tensor_backend(sched_phase2.get(), t, gpu);
1894+
}
1895+
moe::cascade_force_moe_consumers(h2_hijack, phase2_gf, sched_phase2.get(), gpu);
1896+
if (backend_cpu) {
1897+
auto force = [&](ggml_tensor * t) {
1898+
if (t) ggml_backend_sched_set_tensor_backend(sched_phase2.get(), t, backend_cpu);
1899+
};
1900+
for (auto & inp : res->inputs) {
1901+
auto * base = inp.get();
1902+
if (auto * akv = dynamic_cast<llm_graph_input_attn_kv *>(base)) {
1903+
force(akv->self_k_idxs); force(akv->self_v_idxs); force(akv->self_kq_mask);
1904+
} else if (auto * ak = dynamic_cast<llm_graph_input_attn_k *>(base)) {
1905+
force(ak->self_k_idxs); force(ak->self_kq_mask);
1906+
} else if (auto * dsa = dynamic_cast<llm_graph_input_attn_k_dsa *>(base)) {
1907+
force(dsa->self_k_idxs_mla); force(dsa->self_k_idxs_lid);
1908+
force(dsa->self_kq_mask_mla); force(dsa->self_kq_mask_lid);
1909+
} else if (auto * iswa = dynamic_cast<llm_graph_input_attn_kv_iswa *>(base)) {
1910+
force(iswa->self_k_idxs); force(iswa->self_v_idxs);
1911+
force(iswa->self_k_idxs_swa); force(iswa->self_v_idxs_swa);
1912+
force(iswa->self_kq_mask); force(iswa->self_kq_mask_swa);
1913+
} else if (auto * hyb = dynamic_cast<llm_graph_input_mem_hybrid *>(base)) {
1914+
force(hyb->inp_attn->self_k_idxs); force(hyb->inp_attn->self_v_idxs);
1915+
force(hyb->inp_attn->self_kq_mask);
1916+
} else if (auto * hybk = dynamic_cast<llm_graph_input_mem_hybrid_k *>(base)) {
1917+
force(hybk->inp_attn->self_k_idxs); force(hybk->inp_attn->self_kq_mask);
1918+
} else if (auto * hybiswa = dynamic_cast<llm_graph_input_mem_hybrid_iswa *>(base)) {
1919+
force(hybiswa->inp_attn->self_k_idxs); force(hybiswa->inp_attn->self_v_idxs);
1920+
force(hybiswa->inp_attn->self_k_idxs_swa); force(hybiswa->inp_attn->self_v_idxs_swa);
1921+
force(hybiswa->inp_attn->self_kq_mask); force(hybiswa->inp_attn->self_kq_mask_swa);
1922+
} else if (auto * oid = dynamic_cast<llm_graph_input_out_ids *>(base)) {
1923+
force(oid->out_ids);
1924+
}
1925+
}
1926+
}
1927+
moe::deep_copy_phase2_graph(phase2_cache, phase2_gf);
1928+
if (!phase2_cache.persistent_gf) { ret = GGML_STATUS_FAILED; return nullptr; }
18871929
h2_hijack.captured = true;
18881930
}
18891931

1932+
// alloc_graph: Token 1 = full reserve+alloc; Token 2+ = galloc fast-path
1933+
if (!ggml_backend_sched_alloc_graph(sched_phase2.get(), phase2_gf)) { ret = GGML_STATUS_ALLOC_FAILED; return nullptr; }
1934+
18901935
ggml_backend_t be = ggml_backend_sched_get_backend(sched_phase2.get(), 0);
18911936
if (res->t_logits) ggml_backend_sched_set_tensor_backend(sched_phase2.get(), res->t_logits, be);
18921937
if (res->t_embd) ggml_backend_sched_set_tensor_backend(sched_phase2.get(), res->t_embd, be);

‎src/moe-hijacker.cpp‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,12 @@
44
#include "moe-static-bunker.h"
55

66
#include "ggml.h"
7+
#include "../ggml/src/ggml-impl.h"
78
#include "ggml-backend.h"
89

910
#include <cstdio>
1011
#include <cstring>
12+
#include <unordered_map>
1113
#include <utility>
1214
#include <vector>
1315

@@ -185,6 +187,96 @@ void cascade_force_moe_consumers(
185187
}
186188
}
187189

190+
// ---- Phase 2 Graph Deep-Copy ----
191+
192+
void phase2_graph_cache::release() {
193+
if (persistent_gf) { delete[] (char *)persistent_gf; persistent_gf = nullptr; }
194+
if (persistent_ctx) { ggml_free(persistent_ctx); persistent_ctx = nullptr; }
195+
persistent_tensors.clear();
196+
valid = false;
197+
}
198+
199+
void deep_copy_phase2_graph(
200+
phase2_graph_cache & cache,
201+
ggml_cgraph * src_gf)
202+
{
203+
cache.release();
204+
205+
int n_nodes = ggml_graph_n_nodes(src_gf);
206+
int n_leafs = ggml_graph_n_leafs(src_gf);
207+
int total_t = n_nodes + n_leafs;
208+
209+
const size_t tensor_bytes = (size_t)total_t * 1024;
210+
struct ggml_init_params params = { tensor_bytes, nullptr, true };
211+
cache.persistent_ctx = ggml_init(params);
212+
if (!cache.persistent_ctx) {
213+
fprintf(stderr, "deep_copy_phase2_graph: ggml_init(%zu) failed\n", tensor_bytes);
214+
return;
215+
}
216+
217+
std::unordered_map<const ggml_tensor *, ggml_tensor *> map;
218+
219+
auto dup_tensor = [&](const ggml_tensor * src) -> ggml_tensor * {
220+
auto * dst = ggml_dup_tensor(cache.persistent_ctx, src);
221+
if (!dst) return nullptr;
222+
strncpy(dst->name, src->name, GGML_MAX_NAME - 1);
223+
dst->name[GGML_MAX_NAME - 1] = '\0';
224+
dst->op = src->op;
225+
dst->data = src->data;
226+
dst->flags = src->flags;
227+
dst->view_offs = src->view_offs;
228+
dst->extra = src->extra;
229+
memcpy(dst->op_params, src->op_params, sizeof(dst->op_params));
230+
for (int d = 0; d < GGML_MAX_DIMS; d++) dst->nb[d] = src->nb[d];
231+
map[src] = dst;
232+
cache.persistent_tensors.push_back(dst);
233+
return dst;
234+
};
235+
236+
for (int i = 0; i < n_leafs; i++) dup_tensor(ggml_graph_leaf(src_gf, i));
237+
for (int i = 0; i < n_nodes; i++) dup_tensor(ggml_graph_node(src_gf, i));
238+
239+
for (auto * dst : cache.persistent_tensors) {
240+
for (int s = 0; s < GGML_MAX_SRC; s++) {
241+
if (dst->src[s]) {
242+
auto it = map.find(dst->src[s]);
243+
if (it != map.end()) dst->src[s] = it->second;
244+
}
245+
}
246+
if (dst->view_src) {
247+
auto it = map.find(dst->view_src);
248+
if (it != map.end()) dst->view_src = it->second;
249+
}
250+
}
251+
252+
size_t graph_bytes = sizeof(ggml_cgraph) + (size_t)total_t * sizeof(ggml_tensor *) * 2;
253+
char * buf = new char[graph_bytes]();
254+
cache.persistent_gf = (ggml_cgraph *)buf;
255+
cache.persistent_gf->size = total_t;
256+
cache.persistent_gf->n_nodes = 0;
257+
cache.persistent_gf->n_leafs = 0;
258+
cache.persistent_gf->nodes = (ggml_tensor **)(buf + sizeof(ggml_cgraph));
259+
cache.persistent_gf->leafs = cache.persistent_gf->nodes + total_t;
260+
cache.persistent_gf->grads = nullptr;
261+
cache.persistent_gf->grad_accs = nullptr;
262+
cache.persistent_gf->use_counts = nullptr;
263+
cache.persistent_gf->order = GGML_CGRAPH_EVAL_ORDER_LEFT_TO_RIGHT;
264+
cache.persistent_gf->uid = 0;
265+
266+
for (int i = 0; i < n_leafs; i++) {
267+
auto it = map.find(ggml_graph_leaf(src_gf, i));
268+
cache.persistent_gf->leafs[cache.persistent_gf->n_leafs++] = it->second;
269+
}
270+
for (int i = 0; i < n_nodes; i++) {
271+
auto it = map.find(ggml_graph_node(src_gf, i));
272+
cache.persistent_gf->nodes[cache.persistent_gf->n_nodes++] = it->second;
273+
}
274+
275+
cache.valid = true;
276+
fprintf(stderr, "deep_copy_phase2_graph: copied %d nodes + %d leafs (%zu KB ctx, %zu B graph)\n",
277+
n_nodes, n_leafs, tensor_bytes / 1024, graph_bytes);
278+
}
279+
188280
} // namespace moe
189281

190282
#endif // GGML_USE_CUDA

‎src/moe-hijacker.h‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,21 @@
11
#pragma once
22

3-
// Graph scanning and cascade forcing for MoE Phase 2.
3+
// Graph scanning, cascade forcing, and deep-copy caching for MoE Phase 2.
44
// Operates on phase2_hijack structs via reference parameters.
55
// All functions are in namespace moe.
66

77
#ifdef GGML_USE_CUDA
88

99
#include <utility>
10+
#include <unordered_map>
11+
#include <vector>
1012

1113
#include "ggml-backend.h"
1214

1315
struct phase2_hijack;
1416
struct ggml_cgraph;
17+
struct ggml_context;
18+
struct ggml_tensor;
1519

1620
namespace moe {
1721

@@ -33,11 +37,18 @@ void cascade_force_moe_consumers(
3337

3438
struct phase2_graph_cache {
3539
bool valid = false;
40+
ggml_cgraph * persistent_gf = nullptr;
41+
struct ggml_context * persistent_ctx = nullptr;
42+
std::vector<ggml_tensor *> persistent_tensors;
3643

3744
void capture() { valid = true; }
38-
void release() { valid = false; }
45+
void release();
3946
};
4047

48+
void deep_copy_phase2_graph(
49+
phase2_graph_cache & cache,
50+
ggml_cgraph * src_gf);
51+
4152
} // namespace moe
4253

4354
#endif // GGML_USE_CUDA

0 commit comments

Comments
 (0)