speculative: window dspark drafter staging to its trained position range

[thadreber-web]

Overview

draft-dspark stages every context row since the drafter's cache position plus a full draft block in one batch. To make that fit, the server raised the draft context's n_batch to the target n_ctx + block size. Causing two problems:

1. At large -c, the draft context reserves a compute buffer sized for a full-context micro-batch. This is the ~5.9 GB allocation and boot OOM reported in Misc. bug: DSpark drafter inheriting main model's context size #74 (raising draft ctx n_batch 1024 -> 32772).
2. Rows are staged at absolute target positions, but the drafter's trained context is only 4096. Past that, the drafter runs attention/RoPE at positions it never saw in training. Measured on Ternary-Bonsai-27B (Q2_0 target, Q4_1 drafter): acceptance falls from 59% on short prompts to 4.6% at ~28k tokens, making speculation slower than no speculation.

This PR windows the drafter's staging instead:

• The drafter 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).
• The server now caps the draft context's n_batch at the drafter's n_ctx_train instead of the target n_ctx, which removes the large compute buffer from (1).

No new CLI flags, no API changes.

Measured (ctx 131072, DGX Spark GB10):

• short/moderate context: ~70-78% acceptance, 54-58 tok/s decode
• 60k context: ~69% acceptance, 32 tok/s
• 24k-deep prompt: 1.36x over no-speculation, where the old absolute-position staging measured 0.69x (slower than no speculation)

Additional information

Fixes #74. The failed-decode and failed-crop recovery paths reset the window bookkeeping (pos_base, n_retained) along with the existing cache reset, so a rebase that fails mid-round
cannot leave the drafter cache and the staging buffer out of sync. Buffer head trimming is done lazily in 512-row chunks to keep the erase off the per-round hot path.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES - analysis done by Fable. I supervised and reviewed the implementation and benchmarking.

[bri-prism]

Thanks for this. The writeup and the measured before/after numbers made it easy to review. I traced the windowing invariants through common/speculative.cpp and built llama-common and llama-server from the PR head (build CI does not run on fork PRs here); both compile clean, and the core change looks correct.

What I verified

• Batch bounds hold on both paths: non-rebase rounds satisfy n_stage + block <= window <= n_batch via the rebase trigger condition, and rebase rounds satisfy keep + block <= w_keep + block <= window by construction of w_keep. The retained "structurally unreachable" guard is honest.
• The [n_retained | pending] buffer layout stays consistent with accept()'s tail trim (it can only drop rows appended after the last draft() snapshot, never the retained head), and the contiguity asserts are correctly ordered after the recoverable mismatch check, so a post-reset state hits the soft skip path instead of aborting.
• All drafter-cache position uses are rebase-aware: the tail crop at start - base, and full-range removals in begin() and both recovery paths.
• The logits harvest is unaffected, since the bulk llama_get_logits() buffer only contains the block rows.
• Rebase cadence is sane: post-rebase headroom is window/2 - block tokens, so restaging amortizes to roughly one extra drafter row per generated token, consistent with the 60k-context numbers you measured.
• On the server side, the n_ubatch = n_batch raise in the dspark block confirms the compute-buffer mechanism behind #74, so capping n_batch at the drafter's n_ctx_train is the right lever rather than a symptom patch.

Asks before merge

1. Please run, and ideally extend, tests/test-dspark-loop.cpp. Its header still documents "continuous absolute RoPE positions across rounds" as the loop contract, which this PR changes, and it has no coverage of window crossing, rebase restaging, or the failure resets. If the synthetic rounds stay inside the window it passes vacuously; if they cross it, it will diverge from the Python reference. Driving at least one rebase in the gate (with the matching change in dspark_phase2_py_ref.py) would lock in the new semantics.
2. GGML_ASSERT(window > 2 * block_size) aborts the whole process on a reachable configuration. The constructor already throws std::runtime_error for invalid pairings a few lines up; a small CLI -b (for example -b 32 with block size 16 through llama-speculative, which bypasses the server's n_batch raise) now kills the process at init where the old code degraded gracefully. A throw, or a clamp with a warning, would fit better. The bound also permits near-thrash configs where window barely exceeds 2 * block and a rebase fires every round or two; worth a word in the comment.

Non-blocking notes

3. The window is derived from the GGUF context_length metadata, which makes that field silently load-bearing: a drafter GGUF that inherits the target's context length gets window = n_batch, which disables the windowing and reintroduces the full-size compute buffer. Worth verifying published drafter GGUFs carry the trained range, and possibly adding an explicit dspark.* KV or a CLI override later so the trained range is stated rather than inferred.
4. This fixes the boot OOM in #74, but the draft context still inherits the target's n_ctx (params_dft = params_base, never overridden), so the drafter KV cache is still allocated at full target context even though no position at or beyond window is addressable anymore. At -c 131072 with -ngld 99 that can still be significant VRAM. Fine as a follow-up issue so it is not lost when #74 closes.
5. The failed-rebase-decode recovery permanently disables drafting for the sequence: after the reset, ctx_len counts from zero while the buffer restarts at the current position, so every later round hits the mismatch branch and logs an error until the next begin(). That mirrors the pre-existing crop-failure semantics, so it is not a regression, but a "drafting disabled for this seq" flag with a single log line would be kinder than per-round error spam.
6. Nits: the lazy-trim threshold 512 deserves a named constant; after the first long-prompt trim the feature buffer's capacity is never released (a shrink_to_fit after a large cut would return real host memory for long prompts, though the old code behaved the same); and given the assert, w_keep = min(window/2, window - block) always resolves to window/2, fine as a safety net but worth noting in the comment.

Nice work overall. The invariant comments that explain the why (including the measured 59% to 4.6% decay), the consistent state resets on both failure paths, and keeping the overflow guard instead of deleting it as unreachable all made this a pleasant one to review.

[Copilot]

🟡 Not ready to approve

Batch capping, failed-rebase recovery, metadata compatibility, and regression coverage need correction.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.

Pull request overview

Windows dspark staging to the drafter’s trained context range, reducing oversized draft-context allocations.

Changes:

• Rebase retained staging rows within a bounded window.
• Size draft batches using trained context length.
• Document the new staging behavior and benchmarks.

File summaries

File	Description
common/speculative.cpp	Implements windowed staging and recovery bookkeeping.
tools/server/server-context.cpp	Adjusts dspark draft batch sizing.
docs/dspark-scope.md	Documents windowing and performance results.

Review details

Comments suppressed due to low confidence (1)

common/speculative.cpp:1264

• Clearing both the cache and all buffered rows makes recovery impossible after a failed rebase. On the next round n_cache is 0, so ctx_len covers the entire absolute history, but process() can only append rows decoded after this failure; consequently n_buf - n_ret != ctx_len and every later draft is skipped. Preserve the retained window (and reset the physical cache again in case decode partially wrote it) so it can be restaged on the next round.

                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;

• Files reviewed: 3/3 changed files
• Comments generated: 4
• Review effort level: Balanced

We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.