From a46e17052d56eccdf175e4f4f8dac22eb04af0a4 Mon Sep 17 00:00:00 2001 From: Thad Reber Date: Wed, 22 Jul 2026 11:14:00 -0700 Subject: [PATCH 1/2] 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; } From 4180cfd610863107863c79e531a62aab0e0efa22 Mon Sep 17 00:00:00 2001 From: Thad Reber Date: Fri, 31 Jul 2026 04:54:53 -0700 Subject: [PATCH 2/2] dspark : cap the draft batch, add the phase 2 gate fixtures The draft context only ever raised n_batch toward the drafter's trained range, so an explicit -b above that range passed through untouched and kept reserving a compute batch it can never address. Set it in both directions and align n_ubatch. Measured it with a 4096-range drafter at -b 8192: draft ctx n_batch was 8192, now 4096. test-dspark-loop takes a tiny GGUF and a reference JSON on the command line, but nothing in the tree produced them: - scripts/dspark/build_tiny.py: deterministic tiny drafter export, converted to GGUF by conversion/dspark.py - scripts/dspark/phase2_py_ref.py: drives the same rounds through the upstream DeepSpec reference, rebasing positions the way common/speculative.cpp does, and writes ref.json The gate now counts rounds that cross the staging window and fails if none did. It passes 7/7 token-for-token with 3 rebases. Also correct the test header, which still documented the pre-windowing "continuous absolute RoPE positions" contract. speculative: throw instead of GGML_ASSERT when the staging window is too small for a draft block, which a small -b can reach. Drop the claim that a drafter without context_length metadata falls back to the batch size: context_length is a required GGUF key, so such a model never loads. I have read and agree with the contributing guidelines. AI usage disclosure: YES - analysis done by Claude Opus 5. I supervised and reviewed the implementation and bench marking. --- common/speculative.cpp | 13 ++- scripts/dspark/README.md | 106 +++++++++++++++++ scripts/dspark/build_tiny.py | 130 +++++++++++++++++++++ scripts/dspark/phase2_py_ref.py | 200 ++++++++++++++++++++++++++++++++ tests/test-dspark-loop.cpp | 32 ++++- tools/server/server-context.cpp | 11 +- 6 files changed, 479 insertions(+), 13 deletions(-) create mode 100644 scripts/dspark/README.md create mode 100644 scripts/dspark/build_tiny.py create mode 100644 scripts/dspark/phase2_py_ref.py diff --git a/common/speculative.cpp b/common/speculative.cpp index 4510cc4b4072..d6768c179fbc 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -1025,11 +1025,18 @@ struct common_speculative_impl_draft_dspark : public common_speculative_impl { // 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). + // 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(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"); + 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(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); diff --git a/scripts/dspark/README.md b/scripts/dspark/README.md new file mode 100644 index 000000000000..2545758f96c7 --- /dev/null +++ b/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 + +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. diff --git a/scripts/dspark/build_tiny.py b/scripts/dspark/build_tiny.py new file mode 100644 index 000000000000..dc1ea4db43f0 --- /dev/null +++ b/scripts/dspark/build_tiny.py @@ -0,0 +1,130 @@ +#!/usr/bin/env python3 +"""Build the tiny synthetic dspark drafter export used by the Phase 2 gate. + +Writes an HF-format export (config.json + model.safetensors). Convert it with +the repo's own converter so the GGUF and the torch model hold identical weights: + + python convert_hf_to_gguf.py --outfile tiny.gguf --outtype f32 + +The result feeds tests/test-dspark-loop.cpp as argv[1], and the same export dir +feeds scripts/dspark/phase2_py_ref.py, which loads it through the upstream +DeepSpec Qwen3DSparkModel. See scripts/dspark/README.md. + +Weights are drawn from a fixed seed, so the export is reproducible. They are not +the original fixture's weights (those were lost); the gate compares C++ against +Python on the SAME checkpoint, so only self-consistency matters. + +No torch: numpy + safetensors only, so this runs on the host. Only the reference +driver needs a torch container. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np +from safetensors.numpy import save_file + +# max_position_embeddings becomes the GGUF context_length, which becomes the +# drafter's n_ctx_train and therefore `window` in common/speculative.cpp +# (window = min(n_ctx_train, n_batch)). test-dspark-loop.cpp runs at +# n_batch = 2*block_size + 9 = 23, so 20 keeps window at 20 with w_keep = 10 and +# forces a rebase on round 3 of the ACCEPT_SCHEDULE. Do not raise it without +# rechecking that arithmetic -- above 23 the window stops binding and the gate +# never exercises the rebase path. See scripts/dspark/README.md. +CONFIG = { + "architectures": ["Qwen3DSparkModel"], + "model_type": "qwen3", + "hidden_size": 32, + "intermediate_size": 64, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 8, + "vocab_size": 64, + "max_position_embeddings": 20, + "rms_norm_eps": 1e-6, + "rope_theta": 10000.0, + "attention_bias": False, + "attention_dropout": 0.0, + "sliding_window": None, + "layer_types": ["full_attention", "full_attention"], + "tie_word_embeddings": False, + "torch_dtype": "float32", + # dspark-specific (read by Qwen3DSparkModel.__init__ and conversion/dspark.py) + "num_target_layers": 4, + "target_layer_ids": [0, 1, 2], + "block_size": 7, + "mask_token_id": 63, + "num_anchors": 1, + "enable_confidence_head": False, + "markov_rank": 4, + "markov_head_type": "vanilla", +} + + +def build_tensors(cfg: dict, seed: int) -> dict[str, np.ndarray]: + rng = np.random.default_rng(seed) + h = cfg["hidden_size"] + inter = cfg["intermediate_size"] + vocab = cfg["vocab_size"] + head_dim = cfg["head_dim"] + n_q = cfg["num_attention_heads"] * head_dim + n_kv = cfg["num_key_value_heads"] * head_dim + n_cap = len(cfg["target_layer_ids"]) + rank = cfg["markov_rank"] + + def lin(*shape): + return rng.normal(0.0, 0.02, size=shape).astype(np.float32) + + def norm(n): + # jittered rather than exactly 1.0: an ignored norm weight would be + # invisible against an all-ones init + return (1.0 + rng.normal(0.0, 0.02, size=(n,))).astype(np.float32) + + t = { + "embed_tokens.weight": lin(vocab, h), + "norm.weight": norm(h), + "fc.weight": lin(h, n_cap * h), + "hidden_norm.weight": norm(h), + "lm_head.weight": lin(vocab, h), + "markov_head.markov_w1.weight": lin(vocab, rank), + "markov_head.markov_w2.weight": lin(vocab, rank), + } + for i in range(cfg["num_hidden_layers"]): + p = f"layers.{i}." + t[p + "self_attn.q_proj.weight"] = lin(n_q, h) + t[p + "self_attn.k_proj.weight"] = lin(n_kv, h) + t[p + "self_attn.v_proj.weight"] = lin(n_kv, h) + t[p + "self_attn.o_proj.weight"] = lin(h, n_q) + t[p + "self_attn.q_norm.weight"] = norm(head_dim) + t[p + "self_attn.k_norm.weight"] = norm(head_dim) + t[p + "mlp.gate_proj.weight"] = lin(inter, h) + t[p + "mlp.up_proj.weight"] = lin(inter, h) + t[p + "mlp.down_proj.weight"] = lin(h, inter) + t[p + "input_layernorm.weight"] = norm(h) + t[p + "post_attention_layernorm.weight"] = norm(h) + return t + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("outdir", type=Path, help="export directory to create") + ap.add_argument("--seed", type=int, default=0) + args = ap.parse_args() + + args.outdir.mkdir(parents=True, exist_ok=True) + tensors = build_tensors(CONFIG, args.seed) + save_file(tensors, str(args.outdir / "model.safetensors")) + (args.outdir / "config.json").write_text(json.dumps(CONFIG, indent=2) + "\n") + + n_params = sum(int(v.size) for v in tensors.values()) + print(f"wrote {len(tensors)} tensors ({n_params} params, seed={args.seed}) to {args.outdir}") + print(f"convert with:\n python convert_hf_to_gguf.py {args.outdir} " + f"--outfile {args.outdir / 'tiny.gguf'} --outtype f32") + + +if __name__ == "__main__": + main() diff --git a/scripts/dspark/phase2_py_ref.py b/scripts/dspark/phase2_py_ref.py new file mode 100644 index 000000000000..c1476948646d --- /dev/null +++ b/scripts/dspark/phase2_py_ref.py @@ -0,0 +1,200 @@ +#!/usr/bin/env python3 +"""Python reference for the dspark Phase 2 block-draft loop gate. + +Drives the same synthetic rounds as tests/test-dspark-loop.cpp through the real +upstream reference ops (deepseek-ai/DeepSpec) and dumps the expected per-round +drafted block to ref.json, which the C++ test diffs against round for round. + + test-dspark-loop + +torch must run in a container on this host. See scripts/dspark/README.md for the +exact docker invocation, the fixture builder (build_tiny.py), and the derivation +of every constant duplicated below. + +WINDOWING: upstream's forward_dspark_draft_block 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 this driver +applies the same window/rebase rule and calls model._forward_backbone directly +rather than forward_dspark_draft_block. Everything else (embed_tokens, the +backbone, compute_logits, sample_draft_token_step) is the reference's own code. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np +import torch +from safetensors.torch import load_file +from transformers import AutoConfig +from transformers.cache_utils import DynamicCache + +from deepspec.modeling.dspark.qwen3.modeling import Qwen3DSparkModel + +# mirrored verbatim in tests/test-dspark-loop.cpp -- keep both in sync +ACCEPT_SCHEDULE = [7, 3, 0, 7, 5, 1, 4] +PROMPT = [1, 2, 3, 4, 5] + +U32 = 0xFFFFFFFF + + +def hash_u32(x: int) -> int: + x &= U32 + x ^= x >> 16 + x = (x * 0x7FEB352D) & U32 + x ^= x >> 15 + x = (x * 0x846CA68B) & U32 + x ^= x >> 16 + return x + + +def synth_feat(pos: int, d: int) -> float: + h = hash_u32((pos * 131071 + d * 97 + 12345) & U32) + return float(np.float32((h % 2000) - 1000) / np.float32(500.0)) + + +def synth_bonus_token(round_idx: int, vocab_size: int, mask_token_id: int) -> int: + h = hash_u32(((round_idx & U32) * 2654435761 + 999983) & U32) + v = h % (vocab_size - 1) + if v == mask_token_id: + v = (v + 1) % vocab_size + return v + + +def synth_feat_rows(pos_beg: int, n_rows: int, n_embd_cap: int) -> np.ndarray: + return np.array( + [[synth_feat(pos_beg + i, d) for d in range(n_embd_cap)] for i in range(n_rows)], + dtype=np.float32, + ) + + +def load_model(export_dir: Path, device: torch.device) -> Qwen3DSparkModel: + cfg = AutoConfig.from_pretrained(export_dir) + for k, v in json.loads((export_dir / "config.json").read_text()).items(): + if not hasattr(cfg, k): + setattr(cfg, k, v) + # eager keeps the gate deterministic; the trained config uses flex_attention + cfg._attn_implementation = "eager" + model = Qwen3DSparkModel(cfg) + model.load_state_dict(load_file(export_dir / "model.safetensors"), strict=True) + return model.to(device=device, dtype=torch.float32).eval() + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("export_dir", type=Path, help="tiny export from build_tiny.py") + ap.add_argument("out", type=Path, help="ref.json to write") + ap.add_argument("--rounds", type=int, default=len(ACCEPT_SCHEDULE)) + ap.add_argument("--n-batch", type=int, default=None, + help="drafter n_batch; defaults to test-dspark-loop.cpp's 2*block_size+9") + args = ap.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model = load_model(args.export_dir, device) + cfg = model.config + + block_size = int(cfg.block_size) + vocab_size = int(cfg.vocab_size) + mask_token_id = int(cfg.mask_token_id) + n_embd_cap = len(cfg.target_layer_ids) * int(cfg.hidden_size) + + # window/w_keep exactly as common/speculative.cpp derives them + n_batch = args.n_batch if args.n_batch is not None else 2 * block_size + 9 + window = min(int(cfg.max_position_embeddings), n_batch) + if window <= 2 * block_size: + raise SystemExit(f"window={window} too small for block_size={block_size}") + w_keep = min(window // 2, window - block_size) + + start = len(PROMPT) + base = 0 + id_last = synth_bonus_token(-1, vocab_size, mask_token_id) + + rounds = [] + n_rebases = 0 + for r in range(args.rounds): + n_accepted = ACCEPT_SCHEDULE[r % len(ACCEPT_SCHEDULE)] + + # the feature buffer is a pure function of position: after the prefill + # and every accept(), it covers absolute positions [0, start) + n_buf = start + rebase = start - base + block_size > window + if rebase: + base = start - min(n_buf, w_keep) + n_rebases += 1 + + # attended context is exactly [base, start): a rebase wiped everything + # below `base` out of the drafter cache, so restaging from 0 would hand + # the rows negative rebased positions + ctx_pos = list(range(base, n_buf)) + feat = synth_feat_rows(base, len(ctx_pos), n_embd_cap) + target_hidden = torch.from_numpy(feat).unsqueeze(0).to(device) + + draft_input_ids = torch.full((1, block_size), mask_token_id, + dtype=torch.long, device=device) + draft_input_ids[0, 0] = id_last + + # rebased positions, matching the C++ batch: context rows at pos-base, + # block rows at start-base+k + position_ids = torch.tensor( + [[p - base for p in ctx_pos] + [start - base + k for k in range(block_size)]], + dtype=torch.long, device=device, + ) + + with torch.no_grad(): + block_hidden = model._forward_backbone( + target_hidden_states=target_hidden, + noise_embedding=model.embed_tokens(draft_input_ids), + position_ids=position_ids, + attention_mask=None, + past_key_values=DynamicCache(), + use_cache=False, + is_causal=False, + ) + base_logits = model.compute_logits(block_hidden).float() + + # sequential Markov resample: never batched over the block + prev = torch.tensor([id_last], dtype=torch.long, device=device) + sampled = [] + for k in range(block_size): + tok, _ = model.sample_draft_token_step( + base_logits[:, k, :], prev_token_ids=prev, temperature=0.0, + ) + sampled.append(int(tok.item())) + prev = tok + + bonus = synth_bonus_token(r, vocab_size, mask_token_id) + rounds.append({ + "n_accepted": n_accepted, + "ctx_len": n_buf, + "start": start, + "sampled": sampled, + "bonus": bonus, + "rebase": rebase, + }) + print(f"round {r}: start={start} base={base} ctx={len(ctx_pos)} " + f"rebase={rebase} sampled={sampled}") + + id_last = bonus + start += n_accepted + 1 + + ref = { + "n_embd_cap": n_embd_cap, + "block_size": block_size, + "vocab_size": vocab_size, + "mask_token_id": mask_token_id, + "prefill_bonus": synth_bonus_token(-1, vocab_size, mask_token_id), + "window": window, + "w_keep": w_keep, + "rounds": rounds, + } + args.out.write_text(json.dumps(ref, indent=2) + "\n") + print(f"wrote {args.out} ({len(rounds)} rounds, {n_rebases} rebase(s), " + f"window={window}, w_keep={w_keep})") + if n_rebases == 0: + print("WARNING: no rebase fired -- the gate would pass vacuously") + + +if __name__ == "__main__": + main() diff --git a/tests/test-dspark-loop.cpp b/tests/test-dspark-loop.cpp index 09ead03a1266..347090f1d869 100644 --- a/tests/test-dspark-loop.cpp +++ b/tests/test-dspark-loop.cpp @@ -3,9 +3,14 @@ // for a single call) is already gated bit-accurate against the Python reference's real // Qwen3DSparkModel -- see test-dspark-forward.cpp. This test exercises the NEW // Phase 2 piece: the repeated draft/verify loop around that graph -- persistent -// KV-cache growth/crop, block seeding (anchor + mask_token_id), continuous -// absolute RoPE positions across rounds, and the sequential (never-batched) -// Markov resample. +// KV-cache growth/crop, block seeding (anchor + mask_token_id), windowed staging +// positions (drafter positions stay inside [0, window) and rebase to 0 when the +// next block would cross it -- they are NOT continuous across rounds), and the +// sequential (never-batched) Markov resample. +// +// NOTE: the rounds below do not deliberately drive a rebase. If they cross +// `window`, the Python reference must rebase identically or the comparison +// diverges. // // There is no real target model available for this gate, so this drives a // small, fully-synthetic checkpoint (using the real Python reference @@ -189,6 +194,7 @@ int main(int argc, char ** argv) { if (id_last != prefill_bonus_ref) fail("C++/python prefill bonus token disagree -- synth_bonus_token drifted"); int32_t n_mismatch_rounds = 0; + int32_t n_rebase_rounds = 0; for (size_t r = 0; r < rounds_ref.size(); ++r) { const auto & rr = rounds_ref[r]; @@ -197,6 +203,13 @@ int main(int argc, char ** argv) { const int64_t start_ref = rr.at("start").get(); const std::vector sampled_ref = rr.at("sampled").get>(); + // rounds the reference marked as crossing the staging window. Both sides + // derive window/w_keep from the same GGUF context_length and n_batch, and + // `start` is asserted equal below, so a flagged round rebases on both. + if (rr.value("rebase", false)) { + n_rebase_rounds++; + } + if ((int32_t) ACCEPT_SCHEDULE[r % ACCEPT_SCHEDULE.size()] != n_accepted) { fail("ACCEPT_SCHEDULE drifted out of sync with ref.json at round " + std::to_string(r)); } @@ -275,8 +288,17 @@ int main(int argc, char ** argv) { " rounds mismatched the Python reference"); } + // without a window crossing the rounds never touch the rebase path, so the + // run would pass while covering none of the windowed staging logic + if (n_rebase_rounds == 0) { + fail("no round crossed the staging window -- this gate is vacuous. Regenerate " + "ref.json from a fixture whose context_length forces a rebase " + "(see scripts/dspark/README.md)"); + } + printf("\nPhase 2 gate PASSED: %zu/%zu rounds token-for-token identical to the Python reference implementation " - "(cache growth/crop, block seeding, RoPE positions, sequential markov resample).\n", - rounds_ref.size(), rounds_ref.size()); + "(cache growth/crop, block seeding, windowed RoPE positions across %d rebase(s), " + "sequential markov resample).\n", + rounds_ref.size(), rounds_ref.size(), n_rebase_rounds); return 0; } diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 48f639b56b91..b2b89378391d 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -1050,14 +1050,15 @@ struct server_context_impl { 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 (windowed staging + block)\n", + // set it, do not only raise it: an explicit -b above the drafter's + // trained range would otherwise keep reserving a compute batch the + // draft context can never address + if (cparams.n_batch != n_batch_dspark) { + SRV_INF("draft-dspark: draft ctx n_batch %u -> %u (windowed staging + block)\n", cparams.n_batch, n_batch_dspark); cparams.n_batch = n_batch_dspark; } - if (cparams.n_ubatch < cparams.n_batch) { - cparams.n_ubatch = cparams.n_batch; - } + cparams.n_ubatch = cparams.n_batch; } ctx_dft.reset(llama_init_from_model(model_dft.get(), cparams));