speculative: window dspark drafter staging to its trained position range

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<int64_t> 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<int64_t> 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<int64_t> n_retained;
+
     // process()'s per-seq contiguous-range bookkeeping (mirrors draft-mtp).
     std::vector<int32_t> i_batch_beg;
     std::vector<int32_t> i_batch_end;
@@ -1003,10 +1023,30 @@ 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. context_length is a required GGUF key, so a drafter
+        // missing it fails to load before this runs -- the n_ctx_train_dft <= 0
+        // branch is only a defensive guard, not a metadata-free fallback.
+        const int32_t n_ctx_train_dft = llama_model_n_ctx_train(model_dft);
+        window = std::min<int64_t>(n_ctx_train_dft > 0 ? n_ctx_train_dft : n_b, n_b);
+        if (window <= 2 * (int64_t) block_size) {
+            LOG_ERR("%s: staging window=%lld (drafter n_ctx_train=%d, n_batch=%d) too small for block_size=%d\n",
+                    __func__, (long long) window, n_ctx_train_dft, n_b, block_size);
+            throw std::runtime_error("dspark: staging window too small for a draft block "
+                    "(raise -b above 2*block_size)");
+        }
+        // a window barely above 2*block_size rebases every round or two
+        w_keep = std::min<int64_t>(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 +1073,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 +1196,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<int64_t>(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 +1263,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 +1282,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 +1295,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 -------------------------------

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
 

scripts/dspark/README.md

@@ -0,0 +1,106 @@
+# dspark test fixtures
+
+`tests/test-dspark-loop.cpp` (the Phase 2 block-draft-loop gate) takes its
+fixtures on the command line:
+
+    test-dspark-loop <tiny.gguf> <ref.json>
+
+The two scripts here generate them. `build_tiny.py` needs only numpy and
+safetensors and runs anywhere; `phase2_py_ref.py` needs torch and the upstream
+reference, so it runs in a container.
+
+## Upstream reference
+
+The Python reference is `deepseek-ai/DeepSpec`:
+
+    git clone --depth 1 https://github.com/deepseek-ai/DeepSpec.git
+
+| symbol                       | path                                          |
+|------------------------------|-----------------------------------------------|
+| `Qwen3DSparkModel`           | `deepspec/modeling/dspark/qwen3/modeling.py`  |
+| `sample_draft_token_step`    | `deepspec/modeling/dspark/qwen3/modeling.py`  |
+| `forward_dspark_draft_block` | `deepspec/eval/dspark/draft_ops.py`           |
+| `_propose`                   | `deepspec/eval/dspark/evaluator.py`           |
+| `build_markov_head`          | `deepspec/modeling/dspark/markov_head.py`     |
+
+## Building the fixtures
+
+    python scripts/dspark/build_tiny.py $EXPORT
+    python convert_hf_to_gguf.py $EXPORT --outfile $FIX/tiny.gguf --outtype f32
+
+    docker run --rm --gpus all --ipc=host --ulimit memlock=-1 \
+      --ulimit stack=67108864 \
+      -v $DEEPSPEC:/DeepSpec -v $PWD:/work -v $FIX:/fixtures \
+      -w /work -e PYTHONPATH=/DeepSpec nvcr.io/nvidia/pytorch:25.12-py3 \
+      bash -c "pip install -q transformers && \
+        python scripts/dspark/phase2_py_ref.py /fixtures/export /fixtures/ref.json"
+
+    ./build/bin/test-dspark-loop $FIX/tiny.gguf $FIX/ref.json
+
+`transformers` is not in the NGC image, hence the install. The weights are drawn
+from a fixed seed, so the export is reproducible; their values do not matter,
+since the gate compares C++ against Python on the same checkpoint.
+
+## Why the driver does not call forward_dspark_draft_block
+
+Upstream slices `position_ids` at absolute positions and calls
+`past_key_values.crop(start)`. `common/speculative.cpp` instead keeps drafter
+positions inside `[0, window)` and rebases them, so `phase2_py_ref.py` applies
+the same window/rebase rule and calls `model._forward_backbone` directly.
+Everything else (`embed_tokens`, the backbone, `compute_logits`,
+`sample_draft_token_step`) is the reference's own code.
+
+The attended context is exactly `[base, start)`. A rebase drops everything below
+`base` from the drafter cache, so restaging from 0 would give those rows negative
+rebased positions.
+
+## RoPE positions
+
+dspark does not use the stock helper: it imports `rotate_half` and defines its
+own `apply_rotary_pos_emb` (modeling.py):
+
+    q_len = q.size(-2)
+    q_embed = (q * cos[..., -q_len:, :]) + (rotate_half(q) * sin[..., -q_len:, :])
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+
+`k` takes the full cos/sin; `q` takes the last `q_len` entries. So block queries
+get positions `start .. start+block_size-1` and context keys keep
+`past_len .. start-1`. This matches the C++ batch, which adds context rows at
+`pos[stage0+i] - base` and block rows at `start - base + k`.
+
+## Keeping the gate non-vacuous
+
+`window = min(drafter n_ctx_train, n_batch)`, where `n_ctx_train` is the GGUF
+`context_length` (the export's `max_position_embeddings`), and
+`test-dspark-loop.cpp` runs at `n_batch = 2*block_size + 9`.
+
+With `block_size = 7`, `n_batch = 23` and the constructor requires
+`window > 14`. `max_position_embeddings = 20` gives `window = 20`, `w_keep = 10`,
+and against the test's `PROMPT` and `ACCEPT_SCHEDULE`:
+
+    round 0: start=5   5 + 7 = 12 <= 20   no rebase
+    round 1: start=13  13 + 7 = 20 <= 20  no rebase
+    round 2: start=17  17 + 7 = 24 >  20  rebase
+
+Three rebases fire over one `ACCEPT_SCHEDULE` cycle. If the rounds never cross
+the window the gate covers none of the windowed staging logic, so
+`test-dspark-loop` counts the rounds `ref.json` marks as rebasing and fails if
+none did. Re-check this arithmetic when changing the fixture dims.
+
+## ref.json
+
+Top level: `n_embd_cap`, `block_size`, `vocab_size`, `mask_token_id`,
+`prefill_bonus`, `window`, `w_keep`, `rounds`. Each round: `n_accepted`,
+`ctx_len`, `start`, `sampled`, `bonus`, `rebase`.
+
+`hash_u32` / `synth_feat` / `synth_bonus_token` are duplicated between
+`phase2_py_ref.py` and `test-dspark-loop.cpp` rather than shared: they are
+closed-form functions of small integers, not a stateful RNG stream, so bit
+parity across languages falls out of using the same uint32 wraparound
+arithmetic. `ACCEPT_SCHEDULE` and `PROMPT` are likewise mirrored. Keep both
+sides in sync.
+
+## Not covered
+
+`tests/test-dspark-forward.cpp` Tier 2 needs a `ref.bin` from a separate
+generator that does not exist in the tree; that gate cannot run.