DSpark speculation (K=7): drafter chain, verify, and acceptance - #223
Merged
Merged
Conversation
|
Important
This repository does not receive automatic reviews because it has fewer than 10 stars. ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Advanced Run ID: Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
ndleslx
force-pushed
the
dspark-m2-speculation
branch
4 times, most recently
from
September 9, 2026 08:11
136a611 to
1ef8991
Compare
ndleslx
force-pushed
the
dspark-m2-speculation
branch
from
September 10, 2026 03:38
1ef8991 to
f244b1b
Compare
Adopt pypto-lib origin/main through #1184: the drafter backbone tap, the native idle-group skip, and the owner-row right-sizing of the decode head tail. - Swap prefill's hidden_workspace for the rank-owned dspark_target_hidden output ([rows, 3*D] BF16, layers 40/41/42 through one hc_head projection each) and register it on the local dynamic axis; decode gains the same [16-row tile, 3*D] output between pre_hc_hidden_out and x_out (pypto-lib#1084). - Trim the decode ABI's four dead local RoPE tables (freqs_cos_local/sin_local, compressed_freqs_cos_local/sin_local, 113 -> 109 args): the remaining freqs tables absorbed their role and now ride the owner-token T_DYN axis, so each rank stages its own local-token slice of the group stream instead of the whole stream. - Right-size the LM-head and greedy-sampling windows from the step's 512 group rows to each owner's MOE_TOKENS = 128 (DSPARK_MAX_LOGIT_ROWS, pypto-lib#1182); serving's packed logit-row indices and sampled readbacks were already owner-local, so only the window extents change. - Pin the AST name/order parity of both target programs (101/109) plus their dynamic-axis, dtype, and Out-annotation contracts against the kernel sources. - Drop the M1 mirror replay into inactive TP groups: pypto-lib#1161 skips their attention and sampling tails natively from a zero query_start_loc terminal while keeping every EP MoE wave, so idle-group staging stays zero-initialized (-1 logit rows and cache mappings). - Carry the decode/HCA/MoE perf batch, the prefill compressor state fixes, and the prefill-to-decode token-loop validation session between f069bc7 and this pin (pypto-lib#1150, #1158, #1159, #1163, #1165, #1168-#1172, #1175, #1181, #1183, #1184, #1092).
Pack the speculative drafter's checkpoint tensors into the exact bank shapes l3_dspark_drafter and l3_distributed_markov_sample declare. - Flatten the mtp.0/1/2 layers along each bank's first rank-local axis with per-name transforms the target banks do not need: the checkpoint's [out, in] projections (wq_a/wq_b/wkv) transpose into the kernel's matmul orientation, the router gate and confidence head cast to FP32, and tid2eid stacks from the target's hash layers as INT32. - TP-shard the drafter o-projection (wo_a groups / wo_b lora columns) and EP-shard its routed experts per rank, mirroring the target policies. - Replicate the drafter heads verbatim: mtp.0 main_proj/main_norm, the mtp.2 final norm, markov tables, and hc_head constants. - Cross-check every model dimension from the loaded tensors so a mismatched checkpoint fails at load time, not on device. - Gate the whole drafter module set behind validate_drafter_startup_contract: the K=0 path requires, loads, and validates none of it.
Run the K=7 speculative chain behind the target programs
(l3_dspark_drafter + l3_distributed_markov_sample), with host
linear-chain acceptance and eager verify publication.
- Compile the two drafter programs only at num_speculative_tokens=7
(DSPARK_QUERY_WIDTH; anything outside {0, 7} is rejected), pass --dp
through the kernel import context so the markov program covers the
full 16-rank world, and validate the checkpoint's drafter config
(target layers, noise token, markov rank) at startup.
- Give each live request a stable group-local drafter lease (64 per
group, six rotated ring blocks per draft layer) so drafter storage
survives batch reshuffles; a shared read-only zeroed filler range
backs the padded batch rows, which publish nothing.
- Stage the drafter through per-extent shared buffers (batch paddings
4/8/12/16, context extents 32/64/96/128) because the group-context
tensors' dynamic axis is not a per-rank storage prefix; selectors,
hidden taps, and draft readbacks go through packed prefix views.
- Capture a rolling 128-row prompt tail from the prefill tap each chunk
and seed the first draft chain in the worker's terminal-prefill hook
(num_sampled=0, the first sampled token as next_prefill_tokens).
- Verify eagerly: drafts occupy rows 1..7 with all eight rows published
(kv_seq_lens at the full window), so a truncated acceptance's stale
rows are rewritten by the next window before any read; requests whose
window cannot fit the position ceiling fall back per step to the
single-anchor K=0 contract and skip their next drafter query.
- Build the CSA state rings at the window end so every written page
resolves through its absolute page id and survives into later steps'
rebuilt tables.
- Correct async scheduling's optimistic seq_lens from the runner's
committed count before any metadata is built, mirroring the MTP
runner's _correct_mtp_seq_lens.
- Accept on host (longest matching prefix plus the bonus prediction,
1..8 tokens per step), redraft from the just-committed tap rows, and
report acceptance progress plus a per-request summary on release.
- Materialize the drafter's persistent buffers before the KV-capacity
snapshot, release every speculative resource in close(), and keep K=0
loading, allocating, and compiling nothing drafter-related.
Extend the DSpark HTTP guard and its coverage for the K=7 chain. - Parameterize the guard over the speculative depth: the same Palace 64-token prompt / 128-token gate now runs the K=0 target-only contract and the K=7 chain, each on its own server with isolated logs. - Keep the assertions contract-only (model, finish reason, exact token accounting) and require the runner's acceptance-progress line for the K=7 case as proof the drafter/markov/verify path really dispatched; acceptance changes speed, never text, so no exact-text comparison. - Pin both server commands' speculative-config in the host contract test. - Run both cases in the existing 16-card CI job under a wider task-submit budget. - Wait for per-device HBM to return to baseline between the two server boots (npu-smi poll, 8 GiB threshold): each boot needs nearly the whole card, and the driver reclaims a torn-down server's pooled arenas asynchronously after process exit, so an immediate next boot dies in its weight upload with a device OOM. - Document the speculative design in the developer guide: program and weight contracts, the lease/ring cache model, eager verify publication and its fallback, seeding, async position correction, and the acceptance metrics.
ndleslx
force-pushed
the
dspark-m2-speculation
branch
from
September 10, 2026 05:20
f244b1b to
00da7dc
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Implements milestone 2 of #211: the DSpark speculative chain (K=7) behind the merged target-model serving (#213).
What this adds
Select with
--speculative-config '{"method": "dspark", "num_speculative_tokens": 7}';0keeps the target-only path and loads, requires, allocates, and compiles nothing drafter-related.f069bc7 → 216b497: the target programs gain the rank-owneddspark_target_hiddentap (pypto-lib#1084), idle DP groups skip natively (pypto-lib#1161, replacing the M1 mirror replay), and the decode head tail right-sizes to owner rows (pypto-lib#1182): the decode ABI drops the four dead local RoPE tables (113 → 109 args, the remaining tables ride the owner-tokenT_DYNaxis) and the LM-head/sampling windows shrink from the 512-row group stream to each owner's 128 rows. The perf batch between the pins (#1150, #1158-#1159, #1163, #1165, #1168-#1172, #1175, #1181, #1183-#1184) rides along.mtp.0/1/2plus three replicated heads packed into the exactl3_dspark_drafter/l3_distributed_markov_samplebank shapes — projection transposes, FP32 gate/confidence casts, INT32tid2eidfrom the hash layers, TP-sharded o-projection, EP-sharded experts. All banks upload before the KV-capacity snapshot._correct_mtp_seq_lensequivalent); CSA state rings are built at the window end so every written page resolves through its absolute page id across steps.Commits
DSpark submodule bump to 216b497: owner-row head tail and ABI trimDSpark drafter weights: mtp.* checkpoint to stacked device banksDSpark speculation: drafter programs, verify, acceptance, and routingDSpark speculation tests, docs, and CIVerification
tests/unit/model/deepseek_dspark/, 7 tests) covers prefill/decode/drafter staging contracts, AST ABI parity for all four programs, acceptance semantics for every m in 0..7, and the end-to-end accept-then-redraft state updates. Repo-wide unit tests pass except two known pre-existing base failures (MTP owner-outputs fixture, worker-step-protocolKeyError: 'req', reproducible on pristine main). Ruff, header, English, docs-nav, and public-docs checks clean.finish_reason: length, exact token accounting; K=7 acceptance runs at a mean accepted length ≈ 2.0 with zero fallback steps and theDSpark speculation progressline asserted from the server log. Measured KV capacity is unchanged by the bump (K=0: 42 slots, K=7: 37 slots per partition). Note: pypto-lib main past this pin (through #1187's attention pipelining) deadlocks the serving dispatch path — the standalone kernel harness passes, serving hangs at a ring-3 producer — so this pin is also the serving-known-good boundary.Closes #211.