tts: persistent graph arenas and padded attention windows

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.

src/code-predictor-forward.h

@@ -25,15 +25,17 @@
 //   - one private embedding table and one private linear head per
 //     acoustic codebook (1..15)
 //
-// The single-frame loop here recomputes the full graph at every step g
-// (0..14) over a sequence of length g+2. With 5 layers and at most 16
-// tokens per recompute this is sub-millisecond on modern GPUs.
+// Graph metadata lives in two caller owned persistent arenas, one for
+// the T=2 prefill and one for the T=1 steps: each shape class keeps a
+// stable first node address so the CUDA graph cache replays instead of
+// reinstantiating when the two alternate within a frame.
 
 #include "code-predictor-weights.h"
 #include "debug.h"
 #include "ggml-alloc.h"
 #include "ggml-backend.h"
 #include "ggml.h"
+#include "graph-arena.h"
 #include "kv-cache.h"
 #include "qt-error.h"
 #include "sampling.h"
@@ -67,11 +69,17 @@ static struct ggml_tensor * code_predictor_attn_f32(struct ggml_context * ctx,
     return ggml_cont(ctx, ggml_permute(ctx, out, 0, 2, 1, 3));
 }
 
+// Node budget for one predictor graph, same accounting as the talker.
+static int code_predictor_graph_max_nodes(int n_layers) {
+    return 48 * n_layers + 64;
+}
+
 // One Qwen3 decoder block, KV cached. K and V for the T fresh positions
 // are written into the cache at [n_past, n_past+T) on dim 1; the
-// attention reads the contiguous slice [0, n_past+T). Returns the layer
-// output [hidden, T]. use_flash_attn and clamp_fp16 follow the same
-// contract as in talker-forward.h.
+// attention reads the fixed [0, n_kv_pad) window with the mask carrying
+// neg inf beyond n_past+T. Returns the layer output [hidden, T].
+// use_flash_attn and clamp_fp16 follow the same contract as in
+// talker-forward.h.
 static struct ggml_tensor * code_predictor_layer_forward(struct ggml_context *        ctx,
                                                          const CodePredictorWeights * cw,
                                                          const TalkerLayer &          layer,
@@ -82,6 +90,7 @@ static struct ggml_tensor * code_predictor_layer_forward(struct ggml_context *
                                                          struct ggml_tensor *         v_cache,
                                                          int                          n_past,
                                                          int                          T,
+                                                         int                          n_kv_pad,
                                                          bool                         use_flash_attn,
                                                          bool                         clamp_fp16,
                                                          struct ggml_cgraph *         gf) {
@@ -126,9 +135,8 @@ static struct ggml_tensor * code_predictor_layer_forward(struct ggml_context *
     ggml_build_forward_expand(gf, k_cpy);
     ggml_build_forward_expand(gf, v_cpy);
 
-    const int            T_full = n_past + T;
-    struct ggml_tensor * k_full = ggml_view_3d(ctx, k_cache, hd, T_full, n_kv, k_cache->nb[1], k_cache->nb[2], 0);
-    struct ggml_tensor * v_full = ggml_view_3d(ctx, v_cache, hd, T_full, n_kv, v_cache->nb[1], v_cache->nb[2], 0);
+    struct ggml_tensor * k_full = ggml_view_3d(ctx, k_cache, hd, n_kv_pad, n_kv, k_cache->nb[1], k_cache->nb[2], 0);
+    struct ggml_tensor * v_full = ggml_view_3d(ctx, v_cache, hd, n_kv_pad, n_kv, v_cache->nb[1], v_cache->nb[2], 0);
 
     // Q permute [hd, n_q_heads, T] -> [hd, T, n_q_heads] for flash_attn_ext.
     struct ggml_tensor * q_p = ggml_permute(ctx, q, 0, 2, 1, 3);
@@ -184,6 +192,7 @@ static struct ggml_tensor * code_predictor_layer_forward(struct ggml_context *
 static bool code_predictor_run(const CodePredictorWeights * cw,
                                KVCache *                    kv,
                                ggml_backend_sched_t         sched,
+                               GraphArena *                 arena,
                                const float *                fresh_input,
                                int                          T,
                                int                          n_past,
@@ -196,20 +205,18 @@ static bool code_predictor_run(const CodePredictorWeights * cw,
     const int n_layers = cw->num_hidden_layers;
     const int T_full   = n_past + T;
 
-    const int    max_nodes   = 48 * n_layers + 64;
-    const size_t arena_bytes = ggml_tensor_overhead() * max_nodes + ggml_graph_overhead_custom(max_nodes, false);
+    // The attention window spans the whole frame cache (16 slots): a
+    // constant width keeps prefill and step graph shapes fixed across
+    // frames so the CUDA graph cache replays each of the two flavors.
+    const int n_kv_pad = kv->max_seq_len;
 
-    struct ggml_init_params gp   = { arena_bytes, NULL, true };
-    struct ggml_context *   gctx = ggml_init(gp);
-    if (!gctx) {
-        fprintf(stderr, "[CodePredictor] FATAL: ggml_init failed\n");
-        return false;
-    }
+    const int             max_nodes = code_predictor_graph_max_nodes(n_layers);
+    struct ggml_context * gctx      = graph_arena_begin(arena);
 
     // Inputs: fresh embeddings (talker_hidden), positions, attention mask
     struct ggml_tensor * x_in    = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, talker_hidden, T);
     struct ggml_tensor * pos_in  = ggml_new_tensor_1d(gctx, GGML_TYPE_I32, T);
-    struct ggml_tensor * mask_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F16, T_full, T);
+    struct ggml_tensor * mask_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F16, n_kv_pad, T);
     ggml_set_name(x_in, "sub_input");
     ggml_set_name(pos_in, "positions");
     ggml_set_name(mask_in, "causal_mask");
@@ -229,7 +236,7 @@ static bool code_predictor_run(const CodePredictorWeights * cw,
 
     for (int l = 0; l < n_layers; l++) {
         h = code_predictor_layer_forward(gctx, cw, cw->layers[(size_t) l], h, pos_in, mask_in, kv->k[(size_t) l],
-                                         kv->v[(size_t) l], n_past, T, use_flash_attn, clamp_fp16, gf);
+                                         kv->v[(size_t) l], n_past, T, n_kv_pad, use_flash_attn, clamp_fp16, gf);
     }
 
     struct ggml_tensor * h_final = ggml_rms_norm(gctx, h, cw->rms_norm_eps);
@@ -244,7 +251,6 @@ static bool code_predictor_run(const CodePredictorWeights * cw,
     if (!ggml_backend_sched_alloc_graph(sched, gf)) {
         fprintf(stderr, "[CodePredictor] FATAL: graph allocation failed\n");
         ggml_backend_sched_reset(sched);
-        ggml_free(gctx);
         return false;
     }
 
@@ -259,7 +265,7 @@ static bool code_predictor_run(const CodePredictorWeights * cw,
     }
 
     {
-        std::vector<ggml_fp16_t> mask((size_t) T * (size_t) T_full);
+        std::vector<ggml_fp16_t> mask((size_t) T * (size_t) n_kv_pad);
         const ggml_fp16_t        zero    = ggml_fp32_to_fp16(0.0f);
         const ggml_fp16_t        neg_inf = ggml_fp32_to_fp16(-INFINITY);
         for (size_t i = 0; i < mask.size(); i++) {
@@ -268,7 +274,7 @@ static bool code_predictor_run(const CodePredictorWeights * cw,
         for (int q = 0; q < T; q++) {
             const int q_pos = n_past + q;
             for (int k = 0; k <= q_pos; k++) {
-                mask[(size_t) q * (size_t) T_full + (size_t) k] = zero;
+                mask[(size_t) q * (size_t) n_kv_pad + (size_t) k] = zero;
             }
         }
         ggml_backend_tensor_set(mask_in, mask.data(), 0, mask.size() * sizeof(ggml_fp16_t));
@@ -277,7 +283,6 @@ static bool code_predictor_run(const CodePredictorWeights * cw,
     if (ggml_backend_sched_graph_compute(sched, gf) != GGML_STATUS_SUCCESS) {
         fprintf(stderr, "[CodePredictor] FATAL: graph compute failed\n");
         ggml_backend_sched_reset(sched);
-        ggml_free(gctx);
         return false;
     }
 
@@ -286,9 +291,6 @@ static bool code_predictor_run(const CodePredictorWeights * cw,
     ggml_backend_tensor_get(logits, logits_out->data(), (size_t) (T - 1) * row_bytes, row_bytes);
 
     kv->cur_len = T_full;
-
-    ggml_backend_sched_reset(sched);
-    ggml_free(gctx);
     return true;
 }
 
@@ -327,6 +329,8 @@ static bool code_predictor_step(const TalkerWeights *        tw,
                                 const CodePredictorWeights * cw,
                                 KVCache *                    kv,
                                 ggml_backend_sched_t         sched,
+                                GraphArena *                 arena_prefill,
+                                GraphArena *                 arena_step,
                                 const float *                talker_hidden_last,
                                 int                          c0,
                                 float                        temperature,
@@ -361,8 +365,8 @@ static bool code_predictor_step(const TalkerWeights *        tw,
     embed_row_from_backend(tw->codec_embedding, c0, talker_hidden, prefill_input.data() + (size_t) talker_hidden);
 
     std::vector<float> logits;
-    if (!code_predictor_run(cw, kv, sched, prefill_input.data(), 2, 0, talker_hidden, 0, use_flash_attn, clamp_fp16,
-                            &logits)) {
+    if (!code_predictor_run(cw, kv, sched, arena_prefill, prefill_input.data(), 2, 0, talker_hidden, 0, use_flash_attn,
+                            clamp_fp16, &logits)) {
         return false;
     }
     {
@@ -386,8 +390,8 @@ static bool code_predictor_step(const TalkerWeights *        tw,
     for (int g = 1; g < n_acoustic; g++) {
         embed_row_from_backend(cw->codec_embedding[(size_t) (g - 1)], out->codes[(size_t) g], talker_hidden,
                                step_input.data());
-        if (!code_predictor_run(cw, kv, sched, step_input.data(), 1, kv->cur_len, talker_hidden, g, use_flash_attn,
-                                clamp_fp16, &logits)) {
+        if (!code_predictor_run(cw, kv, sched, arena_step, step_input.data(), 1, kv->cur_len, talker_hidden, g,
+                                use_flash_attn, clamp_fp16, &logits)) {
             return false;
         }
         float u_g = 0.0f;

src/graph-arena.h

@@ -0,0 +1,41 @@
+#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;
+    }
+}

src/pipeline-tts.cpp

@@ -245,6 +245,29 @@ bool pipeline_tts_load(PipelineTTS * pt,
         return false;
     }
 
+    // Persistent graph arenas: one shape class each so the backend CUDA
+    // graph cache keeps a stable executable per flavor across steps.
+    if (!graph_arena_init(&pt->talker_arena, talker_graph_max_nodes(pt->talker.num_hidden_layers)) ||
+        !graph_arena_init(&pt->cp_prefill_arena,
+                          code_predictor_graph_max_nodes(pt->code_predictor.num_hidden_layers)) ||
+        !graph_arena_init(&pt->cp_step_arena, code_predictor_graph_max_nodes(pt->code_predictor.num_hidden_layers))) {
+        graph_arena_free(&pt->talker_arena);
+        graph_arena_free(&pt->cp_prefill_arena);
+        graph_arena_free(&pt->cp_step_arena);
+        kv_cache_free(&pt->code_predictor_kv);
+        kv_cache_free(&pt->talker_kv);
+        ggml_backend_sched_free(pt->sched);
+        pt->sched = NULL;
+        pipeline_codec_free(&pt->codec);
+        if (pt->has_speaker_encoder) {
+            speaker_encoder_weights_free(&pt->speaker_encoder);
+        }
+        code_predictor_weights_free(&pt->code_predictor);
+        talker_weights_free(&pt->talker);
+        gf_close(&pt->gguf_talker);
+        return false;
+    }
+
     qt_log(QT_LOG_INFO,
            "[Pipeline] Loaded: arch=%s variant=%s tokenizer=%s codebooks=%d speaker_encoder=%s speakers=%zu fa=%s "
            "clamp_fp16=%s",
@@ -255,6 +278,9 @@ bool pipeline_tts_load(PipelineTTS * pt,
 }
 
 void pipeline_tts_free(PipelineTTS * pt) {
+    graph_arena_free(&pt->cp_step_arena);
+    graph_arena_free(&pt->cp_prefill_arena);
+    graph_arena_free(&pt->talker_arena);
     kv_cache_free(&pt->code_predictor_kv);
     kv_cache_free(&pt->talker_kv);
     if (pt->sched) {
@@ -579,11 +605,11 @@ qt_status pipeline_tts_synthesize(PipelineTTS *                pt,
         bool                ok;
         Timer               t_talker;
         if (step == 0) {
-            ok = talker_forward_prefill(&pt->talker, &pt->talker_kv, pt->sched, prompt.input_embed.data(), prompt.T_ctx,
-                                        use_fa, clamp_fp16, step_dump, &fw);
+            ok = talker_forward_prefill(&pt->talker, &pt->talker_kv, pt->sched, &pt->talker_arena,
+                                        prompt.input_embed.data(), prompt.T_ctx, use_fa, clamp_fp16, step_dump, &fw);
         } else {
-            ok =
-                talker_forward_decode(&pt->talker, &pt->talker_kv, pt->sched, next_emb.data(), use_fa, clamp_fp16, &fw);
+            ok = talker_forward_decode(&pt->talker, &pt->talker_kv, pt->sched, &pt->talker_arena, next_emb.data(),
+                                       use_fa, clamp_fp16, &fw);
         }
         if (!ok) {
             return QT_STATUS_GENERATE_FAILED;
@@ -637,8 +663,9 @@ qt_status pipeline_tts_synthesize(PipelineTTS *                pt,
         const char *        cp_dump = (params->dump_dir && step == 0) ? params->dump_dir : NULL;
         Timer               t_pred;
         if (!code_predictor_step(&pt->talker, &pt->code_predictor, &pt->code_predictor_kv, pt->sched,
-                                 fw.hidden_last.data(), c0, subtk_T, params->subtalker_top_k, params->subtalker_top_p,
-                                 resolved_seed, subseq_counter - 1, use_fa, clamp_fp16, cp_dump, &cp)) {
+                                 &pt->cp_prefill_arena, &pt->cp_step_arena, fw.hidden_last.data(), c0, subtk_T,
+                                 params->subtalker_top_k, params->subtalker_top_p, resolved_seed, subseq_counter - 1,
+                                 use_fa, clamp_fp16, cp_dump, &cp)) {
             return QT_STATUS_GENERATE_FAILED;
         }
         perf.predictor_ms += t_pred.ms();

src/pipeline-tts.h

@@ -14,6 +14,7 @@
 #include "code-predictor-weights.h"
 #include "ggml-backend.h"
 #include "gguf-weights.h"
+#include "graph-arena.h"
 #include "kv-cache.h"
 #include "pipeline-codec.h"
 #include "qwen.h"
@@ -126,6 +127,15 @@ struct PipelineTTS {
     // frame in code_predictor_step.
     KVCache talker_kv;
     KVCache code_predictor_kv;
+
+    // Persistent graph arenas, one per graph shape class. Stable node
+    // addresses across rebuilds keep the backend CUDA graph cache hot:
+    // the talker shares one arena for prefill and decode, the predictor
+    // splits prefill (T=2) and step (T=1) so the two flavors that
+    // alternate within a frame each keep their own executable.
+    GraphArena talker_arena;
+    GraphArena cp_prefill_arena;
+    GraphArena cp_step_arena;
 };
 
 // Open the talker GGUF and the codec GGUF, load every module on the

src/talker-forward.h

@@ -40,6 +40,7 @@
 #include "ggml-alloc.h"
 #include "ggml-backend.h"
 #include "ggml.h"
+#include "graph-arena.h"
 #include "kv-cache.h"
 #include "talker-weights.h"
 
@@ -75,6 +76,14 @@ static bool talker_is_bisect_layer(int l) {
     return false;
 }
 
+// Node budget for one talker graph. Counts approximate: ~38 ops per
+// layer plus 2 cpy and 4 views for the KV write, IO tensors, final
+// norm, codec_head and bisect dump branches. Sizes both the persistent
+// arena in the pipeline and the graph allocated per forward.
+static int talker_graph_max_nodes(int n_layers) {
+    return 48 * n_layers + 256;
+}
+
 // Manual F32 attention chain. Used when use_flash_attn is false: GQA
 // scaled dot product with explicit mul_mat / soft_max_ext / mul_mat,
 // FP32 accumulators end to end. Mirrors the qwen3_attn_f32 helper in
@@ -117,6 +126,7 @@ static struct ggml_tensor * talker_layer_forward(struct ggml_context * ctx,
                                                  struct ggml_tensor *  v_cache,
                                                  int                   n_past,
                                                  int                   T,
+                                                 int                   n_kv_pad,
                                                  bool                  use_flash_attn,
                                                  bool                  clamp_fp16,
                                                  struct ggml_cgraph *  gf) {
@@ -172,12 +182,15 @@ static struct ggml_tensor * talker_layer_forward(struct ggml_context * ctx,
     ggml_build_forward_expand(gf, k_cpy);
     ggml_build_forward_expand(gf, v_cpy);
 
-    // Read the [0, n_past + T) window for attention. The cache slice is
-    // already in the [hd, T_full, n_kv] layout flash_attn_ext expects for
-    // K and V (n_embd, n_kv, n_head_kv, ne3), passed directly as views.
-    const int            T_full = n_past + T;
-    struct ggml_tensor * k_full = ggml_view_3d(ctx, k_cache, hd, T_full, n_kv, k_cache->nb[1], k_cache->nb[2], 0);
-    struct ggml_tensor * v_full = ggml_view_3d(ctx, v_cache, hd, T_full, n_kv, v_cache->nb[1], v_cache->nb[2], 0);
+    // Read the [0, n_kv_pad) window for attention. n_kv_pad covers the
+    // causal context and rounds it up so the view shape stays constant
+    // across consecutive decode steps: the mask carries neg inf beyond
+    // n_past + T and the cache buffer is zero initialized, so the padded
+    // tail contributes nothing. The cache slice is already in the
+    // [hd, n_kv_pad, n_kv] layout flash_attn_ext expects for K and V
+    // (n_embd, n_kv, n_head_kv, ne3), passed directly as views.
+    struct ggml_tensor * k_full = ggml_view_3d(ctx, k_cache, hd, n_kv_pad, n_kv, k_cache->nb[1], k_cache->nb[2], 0);
+    struct ggml_tensor * v_full = ggml_view_3d(ctx, v_cache, hd, n_kv_pad, n_kv, v_cache->nb[1], v_cache->nb[2], 0);
 
     // Q permute [hd, n_q_heads, T] -> [hd, T, n_q_heads]. flash_attn_ext
     // expects ne[1] = n_batch and ne[2] = n_head; the GQA broadcast
@@ -243,10 +256,12 @@ static struct ggml_tensor * talker_layer_forward(struct ggml_context * ctx,
 // it and pulls out the last position hidden + logits. T tokens are
 // appended to the cache starting at n_past. When n_past == 0 and
 // dump_dir is set, the bisect taps fire on the prefill path. use_fa /
-// clamp_fp16 are forwarded as is to every layer.
+// clamp_fp16 are forwarded as is to every layer. The graph metadata
+// lives in the caller owned persistent arena.
 static bool talker_forward_core(const TalkerWeights * tw,
                                 KVCache *             kv,
                                 ggml_backend_sched_t  sched,
+                                GraphArena *          arena,
                                 const float *         input_embed,
                                 int                   T,
                                 int                   n_past,
@@ -259,31 +274,24 @@ static bool talker_forward_core(const TalkerWeights * tw,
     const int vocab    = tw->vocab_size;
     const int T_full   = n_past + T;
 
-    // Dedicated context for graph + IO tensors. Counts approximate:
-    //   per layer: ~38 ops (added 2 cpy + 4 views per layer for KV)
-    //   IO        : 4 tensors (input embed, positions, mask, output norm)
-    //   final     : norm + codec_head + dump branches
-    const int    max_nodes         = 48 * n_layers + 256;
-    const size_t graph_arena_bytes = ggml_tensor_overhead() * max_nodes + ggml_graph_overhead_custom(max_nodes, false);
+    // Attention window rounded up to 256 and clamped to the cache size.
+    // Fixed shapes over spans of 256 decode steps let the CUDA graph
+    // executable update in place (pointer patch) instead of rebuilding.
+    // Ternary instead of std::min: windows.h min/max macros break the
+    // latter in headers on MSVC.
+    const int kv_pad_raw = (int) GGML_PAD(T_full, 256);
+    const int n_kv_pad   = kv_pad_raw < kv->max_seq_len ? kv_pad_raw : kv->max_seq_len;
 
-    struct ggml_init_params gparams = {
-        graph_arena_bytes,
-        NULL,
-        true,
-    };
-    struct ggml_context * gctx = ggml_init(gparams);
-    if (!gctx) {
-        fprintf(stderr, "[TalkerForward] FATAL: ggml_init failed\n");
-        return false;
-    }
+    const int             max_nodes = talker_graph_max_nodes(n_layers);
+    struct ggml_context * gctx      = graph_arena_begin(arena);
 
     // IO tensors: input embedding, positions, causal mask. The mask
-    // spans [T_full, T]: for each fresh query q in [0, T) we allow
-    // keys k in [0, n_past + q]. In the decode case (T=1) this is a
-    // single row of zeros of length T_full.
+    // spans [n_kv_pad, T]: for each fresh query q in [0, T) keys k in
+    // [0, n_past + q] carry 0 and every other slot carries neg inf,
+    // including the padded tail beyond T_full.
     struct ggml_tensor * x_in    = ggml_new_tensor_2d(gctx, GGML_TYPE_F32, hidden, T);
     struct ggml_tensor * pos_in  = ggml_new_tensor_1d(gctx, GGML_TYPE_I32, T);
-    struct ggml_tensor * mask_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F16, T_full, T);
+    struct ggml_tensor * mask_in = ggml_new_tensor_2d(gctx, GGML_TYPE_F16, n_kv_pad, T);
     ggml_set_name(x_in, "input_embed");
     ggml_set_name(pos_in, "positions");
     ggml_set_name(mask_in, "causal_mask");
@@ -296,7 +304,7 @@ static bool talker_forward_core(const TalkerWeights * tw,
     std::vector<struct ggml_tensor *> taps(TALKER_N_BISECT_LAYERS, NULL);
     for (int l = 0; l < n_layers; l++) {
         h = talker_layer_forward(gctx, tw, tw->layers[(size_t) l], h, pos_in, mask_in, kv->k[(size_t) l],
-                                 kv->v[(size_t) l], n_past, T, use_flash_attn, clamp_fp16, gf);
+                                 kv->v[(size_t) l], n_past, T, n_kv_pad, use_flash_attn, clamp_fp16, gf);
         if (record_taps && talker_is_bisect_layer(l)) {
             for (int i = 0; i < TALKER_N_BISECT_LAYERS; i++) {
                 if (TALKER_BISECT_LAYERS[i] == l) {
@@ -334,7 +342,6 @@ static bool talker_forward_core(const TalkerWeights * tw,
     if (!ggml_backend_sched_alloc_graph(sched, gf)) {
         fprintf(stderr, "[TalkerForward] FATAL: graph allocation failed\n");
         ggml_backend_sched_reset(sched);
-        ggml_free(gctx);
         return false;
     }
 
@@ -350,11 +357,12 @@ static bool talker_forward_core(const TalkerWeights * tw,
         ggml_backend_tensor_set(pos_in, pos.data(), 0, (size_t) T * sizeof(int32_t));
     }
 
-    // Causal mask: 0 where k <= n_past + q, -inf otherwise. Stored
-    // row-major [T_q, T_k] with T_k as the fast axis (ne[0]). F16 dtype
-    // matches ggml_flash_attn_ext convention used by the attention path.
+    // Causal mask: 0 where k <= n_past + q, neg inf otherwise. Stored
+    // row major [T_q, n_kv_pad] with n_kv_pad as the fast axis (ne[0]).
+    // F16 dtype matches the ggml_flash_attn_ext convention used by the
+    // attention path.
     {
-        std::vector<ggml_fp16_t> mask((size_t) T * (size_t) T_full);
+        std::vector<ggml_fp16_t> mask((size_t) T * (size_t) n_kv_pad);
         const ggml_fp16_t        zero    = ggml_fp32_to_fp16(0.0f);
         const ggml_fp16_t        neg_inf = ggml_fp32_to_fp16(-INFINITY);
         for (size_t i = 0; i < mask.size(); i++) {
@@ -363,7 +371,7 @@ static bool talker_forward_core(const TalkerWeights * tw,
         for (int q = 0; q < T; q++) {
             const int q_pos = n_past + q;
             for (int k = 0; k <= q_pos; k++) {
-                mask[(size_t) q * (size_t) T_full + (size_t) k] = zero;
+                mask[(size_t) q * (size_t) n_kv_pad + (size_t) k] = zero;
             }
         }
         ggml_backend_tensor_set(mask_in, mask.data(), 0, mask.size() * sizeof(ggml_fp16_t));
@@ -372,7 +380,6 @@ static bool talker_forward_core(const TalkerWeights * tw,
     if (ggml_backend_sched_graph_compute(sched, gf) != GGML_STATUS_SUCCESS) {
         fprintf(stderr, "[TalkerForward] FATAL: graph compute failed\n");
         ggml_backend_sched_reset(sched);
-        ggml_free(gctx);
         return false;
     }
 
@@ -414,11 +421,9 @@ static bool talker_forward_core(const TalkerWeights * tw,
     }
 
     // Advance the cache write head. The graph already executed the cpy
-    // nodes so positions [n_past, n_past + T) are now populated.
+    // nodes so positions [n_past, n_past + T) are now populated. The
+    // arena and the sched allocation persist into the next forward.
     kv->cur_len = T_full;
-
-    ggml_backend_sched_reset(sched);
-    ggml_free(gctx);
     return true;
 }
 
@@ -427,6 +432,7 @@ static bool talker_forward_core(const TalkerWeights * tw,
 static bool talker_forward_prefill(const TalkerWeights * tw,
                                    KVCache *             kv,
                                    ggml_backend_sched_t  sched,
+                                   GraphArena *          arena,
                                    const float *         input_embed,
                                    int                   T,
                                    bool                  use_flash_attn,
@@ -438,7 +444,7 @@ static bool talker_forward_prefill(const TalkerWeights * tw,
         fprintf(stderr, "[TalkerForward] FATAL: prefill T=%d exceeds cache max_seq_len=%d\n", T, kv->max_seq_len);
         return false;
     }
-    return talker_forward_core(tw, kv, sched, input_embed, T, 0, use_flash_attn, clamp_fp16, dump_dir, out);
+    return talker_forward_core(tw, kv, sched, arena, input_embed, T, 0, use_flash_attn, clamp_fp16, dump_dir, out);
 }
 
 // Decode: feed exactly one embedding and append one position to the
@@ -447,6 +453,7 @@ static bool talker_forward_prefill(const TalkerWeights * tw,
 static bool talker_forward_decode(const TalkerWeights * tw,
                                   KVCache *             kv,
                                   ggml_backend_sched_t  sched,
+                                  GraphArena *          arena,
                                   const float *         input_embed_1,
                                   bool                  use_flash_attn,
                                   bool                  clamp_fp16,
@@ -456,5 +463,6 @@ static bool talker_forward_decode(const TalkerWeights * tw,
                 kv->max_seq_len);
         return false;
     }
-    return talker_forward_core(tw, kv, sched, input_embed_1, 1, kv->cur_len, use_flash_attn, clamp_fp16, NULL, out);
+    return talker_forward_core(tw, kv, sched, arena, input_embed_1, 1, kv->cur_len, use_flash_attn, clamp_fp16, NULL,
+                               out);
 }