From a46e17052d56eccdf175e4f4f8dac22eb04af0a4 Mon Sep 17 00:00:00 2001 From: Thad Reber Date: Wed, 22 Jul 2026 11:14:00 -0700 Subject: [PATCH] speculative: window dspark drafter staging to its trained position range The drafter was staged at absolute target positions, so past its 4096 n_ctx_train its acceptance collapsed (59% -> 4.6% between short and 28k prompts on Ternary-Bonsai-27B). Stage rows at rebased positions inside window = min(drafter n_ctx_train, n_batch) and slide the window (wipe + restage last w_keep rows) when a block would cross it. Cap the draft context n_batch at the drafter's n_ctx_train instead of the target n_ctx, which forced a full-context micro-batch that stalled startup at large -c. Measured (target Q2_0, drafter Q4_1, ctx 131072): acceptance holds ~70-78% at short context (54-58 tok/s decode) and ~69% at 60k (32 tok/s); 1.36x over no-speculation at 24k depth vs 0.69x for the old staging. Assisted-by: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01PhRPgjy3fSwQ2dhdxdWkzW --- common/speculative.cpp | 111 +++++++++++++++++++++++++++----- docs/dspark-scope.md | 14 ++++ tools/server/server-context.cpp | 20 +++--- 3 files changed, 121 insertions(+), 24 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index fc4b8f8ded41..4510cc4b4072 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -903,6 +903,26 @@ struct common_speculative_impl_draft_dspark : public common_speculative_impl { // regardless of who proposed a given round's tokens. std::vector rows_since_accept; + // Windowed staging: the drafter's own attention/RoPE is only ever run at + // positions [0, window), window = min(drafter n_ctx_train, drafter n_batch). + // Positions the drafter sees are REBASED: drafter_pos = abs_pos - pos_base. + // When the next block would cross `window`, the drafter cache is wiped and + // the last `w_keep` buffered rows are restaged at positions 0..w_keep-1 + // (pos_base slides forward). The tap features already encode the target's + // full-context state at each row, so long-range information still reaches + // the drafter through them; what the window preserves is the drafter + // staying inside its trained position range at ANY target n_ctx (measured + // on Ternary-Bonsai-27B: acceptance decays 59% -> 4.6% between short and + // 28k-token prompts when staged at absolute positions). + int64_t window = 0; + int64_t w_keep = 0; + std::vector pos_base; + + // rows at the head of ctx_feat/ctx_pos that were already consumed by a + // past draft() round but are retained for the next rebase restage. + // Buffer layout: [n_retained consumed rows][pending rows]. + std::vector n_retained; + // process()'s per-seq contiguous-range bookkeeping (mirrors draft-mtp). std::vector i_batch_beg; std::vector i_batch_end; @@ -1003,10 +1023,23 @@ struct common_speculative_impl_draft_dspark : public common_speculative_impl { const int32_t n_b = (int32_t) llama_n_batch(ctx_dft); batch = llama_batch_init(/* n_tokens = */ n_b, /* embd = */ 0, /* n_seq_max = */ 1); + // windowed staging bounds (see the member comment on `window`): keep + // the drafter's positions inside its trained context regardless of the + // target's n_ctx. A drafter GGUF without context_length metadata + // falls back to the batch size (i.e. the old behavior's ceiling). + const int32_t n_ctx_train_dft = llama_model_n_ctx_train(model_dft); + window = std::min(n_ctx_train_dft > 0 ? n_ctx_train_dft : n_b, n_b); + GGML_ASSERT(window > 2 * (int64_t) block_size && "dspark: staging window too small for a draft block"); + w_keep = std::min(window / 2, window - block_size); + LOG_INF("%s: - staging window=%lld (drafter n_ctx_train=%d, n_batch=%d), w_keep=%lld\n", + __func__, (long long) window, n_ctx_train_dft, n_b, (long long) w_keep); + n_cache.assign(n_seq, 0); ctx_feat.assign(n_seq, {}); ctx_pos.assign(n_seq, {}); rows_since_accept.assign(n_seq, 0); + pos_base.assign(n_seq, 0); + n_retained.assign(n_seq, 0); i_batch_beg.assign(n_seq, -1); i_batch_end.assign(n_seq, -1); } @@ -1033,6 +1066,8 @@ struct common_speculative_impl_draft_dspark : public common_speculative_impl { ctx_feat[seq_id].clear(); ctx_pos[seq_id].clear(); rows_since_accept[seq_id] = 0; + pos_base[seq_id] = 0; + n_retained[seq_id] = 0; llama_memory_seq_rm(llama_get_memory(params.ctx_dft), seq_id, 0, -1); } @@ -1154,41 +1189,64 @@ struct common_speculative_impl_draft_dspark : public common_speculative_impl { "skipping this round\n", __func__, (int) seq_id, (long long) start, (long long) L); continue; } - if ((int64_t) pos.size() != ctx_len) { - LOG_ERR("%s: seq %d staged context rows (%zu) != expected ctx_len (%lld) -- " + const int64_t n_ret = n_retained[seq_id]; + const int64_t n_buf = (int64_t) pos.size(); + if (n_buf - n_ret != ctx_len) { + LOG_ERR("%s: seq %d staged context rows (%lld) != expected ctx_len (%lld) -- " "n_past bookkeeping is out of sync with process()/accept(); " "aborting draft for this seq this round\n", - __func__, (int) seq_id, pos.size(), (long long) ctx_len); + __func__, (int) seq_id, (long long)(n_buf - n_ret), (long long) ctx_len); continue; } - GGML_ASSERT(pos.front() == (int32_t) L && "dspark: staged rows do not start at the drafter's cache position"); - GGML_ASSERT(pos.back() == (int32_t) start - 1 && "dspark: staged rows do not end just before the anchor position"); + GGML_ASSERT(pos[n_ret] == (int32_t) L && "dspark: staged rows do not start at the drafter's cache position"); + GGML_ASSERT(pos.back() == (int32_t) start - 1 && "dspark: staged rows do not end just before the anchor position"); - const int64_t n_tokens = ctx_len + block_size; + // windowed staging (see the member comment on `window`): append the + // pending rows incrementally at rebased positions, or -- when the + // block would cross the window -- slide it: wipe the drafter seq and + // restage the last `keep` buffered rows at positions 0..keep-1. + int64_t base = pos_base[seq_id]; + int64_t stage0 = n_ret; + const bool rebase = start - base + (int64_t) block_size > window; + if (rebase) { + const int64_t keep = std::min(n_buf, w_keep); + stage0 = n_buf - keep; + base = start - keep; + } + + const int64_t n_stage = n_buf - stage0; + const int64_t n_tokens = n_stage + block_size; if (n_tokens > n_batch_max) { + // structurally unreachable (n_stage + block_size <= window <= n_batch + // on both paths above); kept as a guard so a bookkeeping bug skips + // the round instead of overflowing the batch. LOG_ERR("%s: seq %d round needs %lld tokens > n_batch=%lld -- skipping\n", __func__, (int) seq_id, (long long) n_tokens, (long long) n_batch_max); continue; } - llama_set_dspark_ctx(ctx_dft, feat.data(), ctx_len, n_embd_cap); + if (rebase) { + llama_memory_seq_rm(llama_get_memory(ctx_dft), seq_id, -1, -1); + } + + llama_set_dspark_ctx(ctx_dft, feat.data() + (size_t) stage0 * (size_t) n_embd_cap, n_stage, n_embd_cap); common_batch_clear(batch); - for (int64_t i = 0; i < ctx_len; ++i) { + for (int64_t i = 0; i < n_stage; ++i) { // dummy token id: this row's real content comes from the // staged dspark ctx feature above, not the token embedding // (see src/models/dspark.cpp -- these columns are sliced away // before the residual stream even forms). logits=false: this // impl never reads output for context rows. - common_batch_add(batch, /* token = */ 0, (llama_pos)(L + i), { seq_id }, /* logits = */ false); + common_batch_add(batch, /* token = */ 0, (llama_pos)(pos[stage0 + i] - base), { seq_id }, /* logits = */ false); } // block position 0 is seeded with the REAL last-accepted token // (the "anchor"), NOT mask_token_id -- matches the Python reference // reference's evaluator._propose (draft_input_ids[:,0] = // output_ids[:,start]). Positions 1..block_size-1 are masked. - common_batch_add(batch, dp.id_last, (llama_pos) start, { seq_id }, /* logits = */ true); + common_batch_add(batch, dp.id_last, (llama_pos)(start - base), { seq_id }, /* logits = */ true); for (int32_t k = 1; k < block_size; ++k) { - common_batch_add(batch, mask_token_id, (llama_pos)(start + k), { seq_id }, /* logits = */ true); + common_batch_add(batch, mask_token_id, (llama_pos)(start - base + k), { seq_id }, /* logits = */ true); } const int32_t rc = llama_decode(ctx_dft, batch); @@ -1198,6 +1256,16 @@ struct common_speculative_impl_draft_dspark : public common_speculative_impl { if (rc != 0) { LOG_WRN("%s: llama_decode(ctx_dft) failed rc=%d for seq %d\n", __func__, rc, (int) seq_id); + if (rebase) { + // the cache was already wiped for the rebase: bookkeeping must + // not pretend the pre-rebase rows are still resident. + n_cache[seq_id] = 0; + pos_base[seq_id] = 0; + n_retained[seq_id] = 0; + feat.clear(); + pos.clear(); + rows_since_accept[seq_id] = 0; + } continue; } @@ -1207,7 +1275,8 @@ struct common_speculative_impl_draft_dspark : public common_speculative_impl { // The speculative tail is discarded every round regardless of // what the target ultimately accepts; only accept()/process() // decide what becomes real context for the NEXT round. - if (!llama_memory_seq_rm(llama_get_memory(ctx_dft), seq_id, (llama_pos) start, -1)) { + // NOTE: cache positions are drafter-rebased, hence `start - base`. + if (!llama_memory_seq_rm(llama_get_memory(ctx_dft), seq_id, (llama_pos)(start - base), -1)) { // Could not crop just the speculative tail (e.g. the backend // rejected the partial removal): the physical drafter cache still // contains the draft rows, so advancing n_cache to `start` would @@ -1219,16 +1288,26 @@ struct common_speculative_impl_draft_dspark : public common_speculative_impl { "resetting the drafter sequence to recover\n", __func__, (int) seq_id, (long long) start); llama_memory_seq_rm(llama_get_memory(ctx_dft), seq_id, -1, -1); - n_cache[seq_id] = 0; + n_cache[seq_id] = 0; + pos_base[seq_id] = 0; + n_retained[seq_id] = 0; feat.clear(); pos.clear(); rows_since_accept[seq_id] = 0; continue; } - n_cache[seq_id] = start; + n_cache[seq_id] = start; + pos_base[seq_id] = base; - feat.clear(); - pos.clear(); + // this round's pending rows are consumed; retain the buffer tail for + // future rebase restaging. Trim the head lazily (in 512-row chunks) + // to keep the O(w_keep) erase off the per-round path. + if (n_buf > w_keep + 512) { + const int64_t cut = n_buf - w_keep; + feat.erase(feat.begin(), feat.begin() + (size_t) cut * (size_t) n_embd_cap); + pos.erase(pos.begin(), pos.begin() + cut); + } + n_retained[seq_id] = (int64_t) pos.size(); rows_since_accept[seq_id] = 0; // this round's rows were just consumed // --- sequential Markov resample ------------------------------- diff --git a/docs/dspark-scope.md b/docs/dspark-scope.md index 25d4a7701bcd..6d00dd8ea9e4 100644 --- a/docs/dspark-scope.md +++ b/docs/dspark-scope.md @@ -25,6 +25,20 @@ work below is gated on that result. masked block forward, and the `draft-dspark` speculative impl (`common/speculative.cpp`) runs the block-diffusion propose plus the sequential Markov resample (host BLAS/scalar, with an optional CUDA path). +4. Windowed drafter staging. The drafter's own attention/RoPE only ever runs at + positions `[0, window)`, `window = min(drafter n_ctx_train, n_batch)`. Staged + rows are rebased (`drafter_pos = abs_pos - pos_base`); when the next block + would cross the window, the drafter cache is wiped and the last `w_keep` + buffered rows are restaged at positions `0..w_keep-1`. The tap features + already encode the target's full-context state per row, so long-range + information still reaches the drafter; the window keeps it inside its trained + position range at any target `n_ctx`. Measured on Ternary-Bonsai-27B (target + Q2_0, drafter Q4_1, ctx 131072): staged at absolute positions, acceptance + decays 59% -> 4.6% between short and 28k-token prompts; windowed, it holds + ~70-78% (54-58 tok/s decode) at short context and ~69% (32 tok/s) at 60k. + The server also caps the draft context `n_batch` at the drafter's + `n_ctx_train` instead of the full target `n_ctx`, which previously forced a + full-context micro-batch that stalled context creation at large `-c`. ## Open / deferred diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 45f668d7c1d9..48f639b56b91 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1039,15 +1039,19 @@ struct server_context_impl { } } - // dspark stages all context rows since its cache position PLUS a - // full block in ONE batch -- worst case ctx_len == n_ctx right - // after begin() (e.g. a follow-up request with a long history). - // If the drafter's batch cannot fit ctx_len + block_size, the - // round is skipped ("round needs N tokens > n_batch") and - // speculation silently degrades to plain AR. - const uint32_t n_batch_dspark = cparams.n_ctx + (block_size > 0 ? block_size : 64); + // dspark stages context rows PLUS a full block in ONE batch, but + // the impl windows its staging to the drafter's trained position + // range (common/speculative.cpp draft-dspark), so the batch only + // needs to cover that window -- NOT the full target n_ctx. + // (Sizing it to n_ctx forced a full-context micro-batch on the + // draft context, which stalled context creation at large -c.) + const int32_t n_ctx_train_dft = llama_model_n_ctx_train(model_dft.get()); + uint32_t n_batch_dspark = cparams.n_ctx + (block_size > 0 ? block_size : 64); + if (n_ctx_train_dft > 0) { + n_batch_dspark = std::min(n_batch_dspark, (uint32_t) n_ctx_train_dft); + } if (cparams.n_batch < n_batch_dspark) { - SRV_INF("draft-dspark: raising draft ctx n_batch %u -> %u (full-context staging + block)\n", + SRV_INF("draft-dspark: raising draft ctx n_batch %u -> %u (windowed staging + block)\n", cparams.n_batch, n_batch_dspark); cparams.n_batch = n_batch_dspark; }