[2/3] feat(dsv4): implement exact DSpark speculative decoding

[calvarado2004]

Outcome

FreeToken can now load the DSpark drafter shipped inside DeepSeek-V4-Flash-0731, generate its five-token proposal block, verify that block with the TP target, and commit or reject it without changing the target model's output distribution.

This is the correctness and operational milestone: it connects the checkpoint architecture to FreeToken's scheduler, KV/SWA pools, recurrent compressor state, hybrid CPU/GPU MoE backend, TP collectives, OpenAI response path, and service lifecycle. PR 3 adds hardware-aware adaptive width selection and removes remaining sampler transfers.

Stack and review boundary

This is PR 2 of 3 and depends on PR 1 (pr/dsv4-tp-runtime). GitHub will show the TP foundation until PR 1 merges; the new review delta here is:

Commit	Purpose
423edd1	Load the checkpoint drafter and implement exact draft/verify semantics
d42be9b	Make hybrid TP serving operational: addressing, rollback, allocation cleanup, NUMA placement, and lifecycle

Delta from PR 1: 45 files, 4,466 additions, 79 deletions.

Design references

This implementation follows, rather than reinterprets:

• DSpark: Confidence-Scheduled Speculative Decoding with Semi-Autoregressive Generation
• FreeToken: Efficient Edge-Native MoE Serving with Bandwidth-Adaptive Execution
• the working DeepSeek-V4 DSpark path in the vLLM reference fork, used to cross-check checkpoint names, anchor/noise layout, Markov sampling, and speculative state flow

The checkpoint remains authoritative for gamma, the noise token, draft layer count, Markov rank, target feature layers, and all tensor dimensions.

End-to-end speculative cycle

committed target frontier
        │
        ├─ target hidden taps at checkpoint-declared layers (40, 41, 42)
        │          └─ projection supplies each draft layer's context KV
        │
        └─ anchor token + checkpoint noise tokens
                   │
                   ▼
          parallel three-layer draft backbone
                   │
                   ▼
      left-to-right low-rank Markov sampling → proposals + q + confidence
                   │
                   ▼
       one target pass over anchor + five proposals → target p
                   │
                   ▼
        exact contiguous-prefix acceptance / residual recovery
                   │
                   ▼
       commit accepted frontier; restore and release rejected suffix

Checkpoint-faithful DSpark architecture

• Load the checkpoint's three full MoE draft blocks under mtp.*, continuing the target's layer-ID space so target and draft experts share one cache address space without collision.
• Project the target's hidden state from the declared auxiliary layers into the draft hidden width. The layers and ordering come from the checkpoint, not a guessed “last hidden state.”
• Derive each draft layer's context KV from that projection instead of rerunning the prompt through the drafter.
• Write context rows with absolute positions and per-row rotary frequencies. Reusing position-zero frequencies after the first step silently destroys proposal quality.
• Feed the block exactly as trained: the last committed token as anchor, followed by the checkpoint's noise token.
• Run the expensive draft backbone over the block in parallel.
• Sample the lightweight Markov head left-to-right so proposal k conditions on proposal k-1; the first proposal conditions on the anchor.
• Preserve the full proposal distribution q and confidence scores. Keeping only argmax tokens would make correct sampled verification impossible.
• Keep the Markov and confidence heads replicated under TP; sharding these per-position lightweight heads would add a full-vocabulary collective to every proposal position.

The trained maximum is gamma=5. This PR does not manufacture a longer block that the checkpoint never learned.

Exact decoding semantics

The implementation supports both greedy and probabilistic requests.

Greedy

Accept a contiguous prefix while proposal tokens match the target argmax. Stop at the first disagreement and emit the target token already computed at that position. A later accidental match cannot be accepted after an earlier rejection.

Probabilistic

Both target p and drafter q are formed after applying the request's temperature, top-p, and top-k policy. For each proposed token x:

accept x with probability min(1, p(x) / q(x))

At the first rejection, sample the replacement token from normalized:

max(p - q, 0)

If the complete proposal prefix survives, sample the bonus token from the target's next-position distribution. This is standard rejection sampling: the emitted sequence remains distributed exactly as target-only decoding. An argmax comparison for a sampled request would be faster but biased, so it is not used.

State ownership and rollback

A verify writes speculative state before acceptance is known. Every state owner therefore has an explicit commit/release rule:

State	Commit behavior	Rejection behavior
Request token buffer	Retain accepted proposals plus target bonus	Truncate at accepted frontier
Target paged KV	Keep accepted rows	Return unused tail pages immediately
Draft sliding-window KV	Catch up only from target rows that committed	Do not advance from abandoned target rows
SWA slots	Keep accepted span	Release rejected tail slots immediately
Target auxiliary features	Select by physical speculative position	Never commit the feature row belonging to a rejected suffix
Compressor/indexer carry	Journal per physical row before overwrite	Restore the carry belonging to the accepted frontier

The carry journal is required because several speculative positions may share one 128-token recurrent ring page. Releasing a page is insufficient when the rejected row overwrote state inside a page that an accepted row still owns.

Disconnect and abort paths use the same cleanup contracts, so a client that disappears during verification cannot leak speculative pages until process teardown.

FreeToken hybrid execution

A target verify contains only anchor + gamma rows per request. Treating it as bulk prefill would stream whole expert layers and erase the value of speculation. It instead uses the short-decode hybrid path:

1. GPU-resident expert hits execute on GPU.
2. A bandwidth-matched subset of misses is fetched over PCIe and executed on GPU.
3. Remaining misses execute in rank-local CPU worker pools.
4. CPU work is launched before GPU fetch/GEMM so both paths overlap.
5. The two exact rank-local partials are merged, then normal TP completion applies.

The split follows FreeToken Eq. 4 using measured overlapping bandwidth, not a hand-tuned expert count. The validation profile resolved to approximately 28% PCIe fetch / 72% CPU computation.

Draft MoE layers stay on the GPU/cache path because the draft is a serial dependency of every cycle; assigning those layers to the CPU pool would put the slowest resource before every target verify.

Serving and operational work

• Bind TP ranks and their CPU worker pools to the appropriate NUMA nodes before large allocations.
• Size CPU scratch for the six-row target verify span.
• Increase the distributed startup timeout for legitimately skewed rank-local expert loading.
• Add speculative acceptance accounting to decode logs.
• Add a repeatable DSV4 server manager for graceful startup, readiness, logs, status, test requests, and conservative cleanup of wedged ranks.
• Restore the FreeToken ASCII-art startup banner and report TP/GPU/CPU/NUMA/DSpark configuration at startup.

The manager treats ready to serve as readiness; an early HTTP bind is not considered ready while the expert banks are still loading.

Suggested review order

Area	Files	What to verify
Checkpoint architecture	deepseek_v4/dspark.py, args.py, weight.py	mtp.* mapping, anchor/noise inputs, target taps, Markov conditioning
Model integration	model.py, attention.py, compress.py, sparse backends	Auxiliary capture, absolute rotary positions, draft layer IDs
Exact sampling	dspark.py, engine.py	Prefix rule, p/q test, residual recovery, target bonus
State safety	rollback.py, scheduler/cache and KV/SWA pools	Accepted-frontier ownership and suffix release
Hybrid/NUMA	offload_cache.py, cpu_executor.py, utils/numa.py	Concurrent GPU/CPU partials and rank placement
Serving	API/launch/status files, utils/banner.py, scripts/freetoken-dsv4.sh	lifecycle, abort cleanup, useful startup diagnostics

Validation

Focused coverage includes:

• checkpoint wiring and draft-head shapes;
• anchor/noise layout, non-causal draft backbone, and causal Markov order;
• absolute-position rotary behavior and target feature layer IDs;
• greedy contiguous-prefix acceptance;
• probabilistic distribution preservation across agreeing, disagreeing, and nearly disjoint p/q distributions;
• residual fallback and complete-prefix target bonus;
• multi-token request advancement;
• KV/SWA allocation, tail release, and recurrent carry restoration;
• target auxiliary row addressing;
• TP checkpoint loading and vocabulary-head integration;
• NUMA placement and FreeToken hybrid bandwidth balance.

The completed stack passes on the TP4 validation server:

192 passed in 36.14s

Runtime validation uses temperature-1 probabilistic generation, not only a greedy diagnostic path.

Scope and limitations

• Verification is fixed at the checkpoint's full gamma=5 in this milestone. Hardware-aware prefix admission is isolated in PR 3.
• The PR preserves requested sampling semantics; it does not force greedy decoding to inflate acceptance.
• It does not increase the checkpoint-trained block width.
• It does not claim workload-independent speedup. Acceptance depends on the task, which is why performance results are reported in PR 3 with both code and open-ended chat workloads.

[gdevenyi]

Flagging one thing here rather than only on #71, since this is the PR that introduces the file.

python/freetoken/utils/numa.py collides with a branch of mine

This PR adds utils/numa.py for NUMA topology and how TP ranks spread across nodes. I have a branch not yet up as a PR that adds a file at the same path, solving the other half of the problem: where the expert banks' pages actually land, via mbind(MPOL_PREFERRED). Different API, same filename — so whichever lands first, the second one has to be reworked rather than merged.

Worth settling early because the two halves interact, and in a direction that matters for this PR specifically. Your docstring already names the failure:

If the loader's threads are unbound they scatter across every socket, and the CPU expert pool, pinned to one socket's cores, then reads much of its data across the interconnect.

What I measured on a 2× Xeon Gold 6526Y is that confining the workers without also placing the banks can make throughput less predictable, not more: a steady ~69 GB/s on the CPU MoE GEMV unconfined, versus 56–123 GB/s run-to-run with workers pinned but banks left to first-touch. Binding ranks to nodes — which this PR does — is exactly what puts a machine into that regime, so explicit placement is what keeps the binding from turning into a coin flip on a two-socket target.

Your topology/rank-spread half is the one that has to exist first, and placement() is already load-bearing for _tp_core_share, so I'd rather not ship a competing file. Happy to rebase mine on top of yours as an additive place_banks() once this lands, if the module API is stable enough to build on — or to adjust now if you'd prefer it arrive together.

Fuller context on the other overlaps (the bs=1 hybrid/offload calibration vs DSpark's width + 1 verification block, and the effect on the adaptive width cost curve) is in my comment on #71: #71 (comment)

Measured on a two-socket Xeon for the numbers above; my other measurements in the #71 comment are from a single-socket consumer box and are indicative rather than predictive for your 4× A4000 target.

[gdevenyi]

Measured this on 2× RTX 6000 Ada + 2× Xeon Gold 6526Y, DeepSeek-V4-Flash-0731, offloaded experts, TP=2 (with #70). It works — drafter builds, verify graphs capture, output is correct — but on this hardware it is a net slowdown, and I think the reason is structural rather than a tuning problem.

dSpark drafter enabled: 3 draft layers, block size 5, target layers (40, 41, 42)
Capturing DSpark target verify graphs: requests=[1, 2, 4], rows/request=[6]
accepted (42%)

config	median tok/s
TP=2, no speculation	39.25 (n=10)
TP=2, --speculative-dspark	35.91 (n=12)

~8% slower at a 42% acceptance rate.

Why I think it loses here specifically. This deployment is bound by expert traffic over PCIe, not by per-step latency. Speculation trades more work per step for fewer steps, but with offloaded experts "more tokens per step" means more expert traffic per step — so it pushes on the exact resource that is already the constraint. Two compounding effects:

• The drafter's own routed experts enlarge the host pool from 143 → 153 GiB, which shrinks the GPU expert cache (moe_cache_size 5622 → 5551) and therefore raises the miss rate for the target model too.
• Each verify step fetches the union of the experts routed by all speculated positions, so a block of 5 draft tokens touches materially more unique experts than one token would.

That is not an argument against the PR. On a machine where the expert pool is resident (or nearly so) and decode is latency-bound rather than bandwidth-bound, the same 42% acceptance should pay off — it is the offload regime specifically that inverts it. Worth noting in the docs, since --speculative-dspark and --moe-backend offload are both things an edge user reaches for, and together they can subtract.

If you have an acceptance-rate/step-time breakdown for a resident-expert configuration, that would make the tradeoff legible. And if there is a knob for block size, a smaller block would fetch fewer experts per verify and might land better under offload — I did not see one exposed.

Merged cleanly on top of #70 with no conflicts against my CPU MoE work. No test regressions attributable to it.

[calvarado2004]

Thank you for validating both correctness and the negative performance result. I agree with the diagnosis: at 42% proposal acceptance, exact verification expands the offloaded expert working set without producing enough accepted prefix to amortize it. That is a real losing regime, not something we should hide behind the aggregate A4000 result.

The fixed five-token block in this PR is deliberate: #69 establishes the exact DSpark sampling and rollback semantics first. #71 adds the hardware-measured width choice over 0–5 proposals, so it can decline or shorten speculation when the measured D + V(k) cost is not justified. A static smaller-block knob would be useful for experiments, but it is not a substitute for that measured decision.

Our workload split supports your result:

• on the 4x A4000 hybrid CPU/GPU system, structured code at roughly 80–82% proposal acceptance reaches about 28–30 tok/s from an ~18 tok/s baseline;
• a recent open-ended poem/reasoning run at roughly 37% acceptance reaches only about 11–14 tok/s, below baseline;
• your TP=2 pure-offload result at 42% acceptance (35.91 versus 39.25 tok/s) is therefore exactly the kind of case where speculation should be disabled or shortened.

I agree the expected-regime note should be explicit: exactness does not imply a workload-independent speedup. Resident/hybrid expert placement and high prefix survival are favorable; offload-dominated verification and low survival can lose. The proposal acceptance, selected width, draft cost, and verification cost should be visible together so users can tell which regime they are in.

On NUMA, the approaches are complementary. The current intent is rank binding before expert-bank allocation, allowing loader first-touch to follow that affinity. Explicit mbind(MPOL_PREFERRED) is stronger than relying on first-touch behavior alone, and an additive place_banks() built on the existing utils.numa topology/placement API sounds like the right merge shape. I am happy to avoid duplicating that work and compose the stacks during rebase.

Thanks for running the full stack on a materially different topology. This is exactly the kind of result needed to define the boundary of the implementation honestly.

[calvarado2004]

Follow-up from the DSpark TP>1 hybrid CPU+GPU integration work in #71.

First, thanks again for the pure-offload measurements and for calling out the autoregressive dependency. We took that result seriously. Our measurements suggest that it is important to separate two regimes:

• unconditional/pure speculative CPU offload can lose when draft latency sits directly on the decode critical path;
• a hybrid, device-resident DSpark verifier with high acceptance can still win because the accepted target tokens amortize that cost.

On our 4x GPU + dual-NUMA workstation, after adopting Gabriel Devenyi's fused HC inverse-RMS improvement and retaining the exact speculative acceptance/rejection semantics, the current isolated branch produced:

• long coding: ~29.6 tok/s sustained, 31.7 tok/s final checkpoint;
• repeated-route-aware variant: 33.07 tok/s peak, with a 3,000-token coding run remaining coherent and ~29.6 tok/s sustained;
• explicit reasoning_effort=high coding: 30.36 tok/s mean after warm-up, 28.31–32.10 range, ~86.9% isolated acceptance;
• explicit high structured reasoning completed the Bayes and knights/knaves tests correctly, commonly around 16–20 tok/s and briefly 22 tok/s.

The reprofile is also useful: the old symmetric NCCL path accounted for 878.0 ms / 92.1% of CUDA time over 87 calls. With direct standard NCCL, those same 87 calls were 27.2 ms / 15.8%; the new largest kernel category is repeated expert-cache index copying at 65.9 ms / 38.3%. That moves the next optimization target to the paper's hybrid fetch fraction/q* balance, rather than removing exact verification or assuming more CPU work is free.

So our proposed direction remains: keep #69's findings as the warning against universal pure-offload claims, keep exact probabilistic verification and block/rejection behavior, and let #71 adapt speculative width/fetch policy to the measured target curve. The repeated-token issue we saw during integration was our shared detokenizer state regression, not caused by this PR, and is fixed in the current #71 head.

Detailed integration measurements: #71 (comment)