refactor: unify DSpark heads and load the DeepSeek-V4 confidence head. - #2271
Open
weizhehuang0827 wants to merge 1 commit into
Open
refactor: unify DSpark heads and load the DeepSeek-V4 confidence head.#2271weizhehuang0827 wants to merge 1 commit into
weizhehuang0827 wants to merge 1 commit into
Conversation
weizhehuang0827
requested review from
DongheJin,
DragonFive,
JimHsiung,
Kang-Meng,
liujinguang0125,
liutongxuan,
ustcfy,
xiao-yu-chen,
yingxudeng,
yinjiawei01 and
zhang-minchao
as code owners
August 19, 2026 09:44
ustcfy
previously approved these changes
Aug 19, 2026
Two related DSpark cleanups plus a fix: - Move dspark_confidence_head.h from models/llm/npu/ to models/llm/ (namespace xllm), next to dspark_markov_head.h. It is plain torch code with no NPU dependency and belongs beside the other shared DSpark head. - Remove the duplicate DSparkMarkovHead defined inline in qwen3_dspark.h and use the shared models/llm/dspark_markov_head.h from both the Qwen3 and DeepSeek-V4 drafts. The shared head gains the markov_embed() helper the ConfidenceHead reuses; the Qwen3 loader strips the "markov_head." prefix so the shared unprefixed keys match. - Load the DeepSeek-V4 DSpark confidence head. The checkpoint ships mtp.<last>.confidence_head.proj.weight (with_markov, bias-less) but the model never loaded it, so adaptive pruning silently fell back to sampler proposal probs. DSparkConfidenceHead now treats the bias as optional, and DeepseekV4DSpark loads + exposes the head via the same interface as Qwen3. The ConfidenceHead is only stood up when adaptive speculative decode is on (SpeculativeConfig::enable_adaptive_speculative_decode()); under static decoding its weights are never loaded, since only the adaptive controller consumes them.
weizhehuang0827
force-pushed
the
spec-dspark-confidence-refactor
branch
from
August 19, 2026 11:00
8d4ef74 to
7bd6400
Compare
ustcfy
approved these changes
Aug 19, 2026
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
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.
Description
Two related DSpark head cleanups plus a fix that were split out of the DSpark
adaptive-decode work:
Relocate the confidence head. Move
dspark_confidence_head.hfrommodels/llm/npu/tomodels/llm/(namespacexllm), next todspark_markov_head.h. It is plain torch code with no NPU dependency, so itdoes not belong under
npu/.De-duplicate
DSparkMarkovHead.qwen3_dspark.hcarried its own inlinecopy of the Markov head that had diverged from the shared
models/llm/dspark_markov_head.h. Both the Qwen3 and DeepSeek-V4 drafts nowuse the shared head (which gains the
markov_embed()helper the confidencehead reuses); the Qwen3 loader strips the
markov_head.prefix so the sharedunprefixed keys match.
Load the DeepSeek-V4 DSpark confidence head. The DeepSeek-V4-Flash
checkpoint ships
mtp.<last>.confidence_head.proj.weight(with-markov,bias-less), but
DeepseekV4DSparknever loaded it, so adaptive pruningsilently fell back to sampler proposal probabilities instead of the trained
confidence.
DSparkConfidenceHeadnow treats the bias as optional, and theDeepSeek-V4 model loads and exposes the head through the same interface as
Qwen3.
The confidence head is only stood up when adaptive speculative decode is enabled
(
SpeculativeConfig::enable_adaptive_speculative_decode()); under staticdecoding its weights are never loaded, since only the adaptive controller
consumes them.
Test
dspark_confidence_head_testnumerical-parityunit test, which was updated for the shared Markov head's constructor and
unprefixed keys).
dp_size=2, ShareGPT:and serving is unaffected.
controller; the run serves all requests with no crash or HCCL hang.