Rebuild each forward into a persistent arena per graph shape class (one for the talker, two for the code predictor prefill and step flavors that alternate within a frame) so nodes keep stable addresses and the CUDA graph cache replays its executable instead of reinstantiating. Pad the talker attention window to 256 and fix the predictor window to the frame cache size so decode shapes hold across steps, with the causal mask carrying neg inf over the padded tail. Drop the per step ggml context churn and the trailing sched resets: one talker step plus 15 predictor micro steps per frame no longer pay a full build/alloc/free cycle each.
42 lines
1.3 KiB
C
42 lines
1.3 KiB
C
#pragma once
|
|
// graph-arena.h: persistent no_alloc ggml context reused across graph
|
|
// builds. Rebuilding an identical graph into the same arena lands every
|
|
// node at a stable address, so the backend CUDA graph cache (keyed on
|
|
// the first node pointer) resolves to the same executable instance at
|
|
// every decode step instead of thrashing on fresh allocations.
|
|
|
|
#include "ggml.h"
|
|
|
|
#include <cstddef>
|
|
#include <cstdio>
|
|
|
|
struct GraphArena {
|
|
struct ggml_context * ctx = nullptr;
|
|
};
|
|
|
|
// Allocate the arena once, sized for max_nodes tensors plus one graph.
|
|
static bool graph_arena_init(GraphArena * a, int max_nodes) {
|
|
const size_t bytes =
|
|
ggml_tensor_overhead() * (size_t) max_nodes + ggml_graph_overhead_custom((size_t) max_nodes, false);
|
|
struct ggml_init_params gp = { bytes, NULL, true };
|
|
a->ctx = ggml_init(gp);
|
|
if (!a->ctx) {
|
|
fprintf(stderr, "[GraphArena] FATAL: ggml_init failed (%zu bytes)\n", bytes);
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
// Rewind the arena: the next build sequence reuses the same addresses.
|
|
static struct ggml_context * graph_arena_begin(GraphArena * a) {
|
|
ggml_reset(a->ctx);
|
|
return a->ctx;
|
|
}
|
|
|
|
static void graph_arena_free(GraphArena * a) {
|
|
if (a->ctx) {
|
|
ggml_free(a->ctx);
|
|
a->ctx = NULL;
|
|
}
|
|
}
|