Skip to content

refactor: unify DSpark heads and load the DeepSeek-V4 confidence head. - #2271

Open
weizhehuang0827 wants to merge 1 commit into
xLLM-AI:mainfrom
weizhehuang0827:spec-dspark-confidence-refactor
Open

refactor: unify DSpark heads and load the DeepSeek-V4 confidence head.#2271
weizhehuang0827 wants to merge 1 commit into
xLLM-AI:mainfrom
weizhehuang0827:spec-dspark-confidence-refactor

Conversation

@weizhehuang0827

Copy link
Copy Markdown
Collaborator

Description

Two related DSpark head cleanups plus a fix that were split out of the DSpark
adaptive-decode work:

  1. Relocate the confidence head. 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, so it
    does not belong under npu/.

  2. De-duplicate DSparkMarkovHead. qwen3_dspark.h carried its own inline
    copy of the Markov head that had diverged from the shared
    models/llm/dspark_markov_head.h. Both the Qwen3 and DeepSeek-V4 drafts now
    use the shared head (which gains the markov_embed() helper the confidence
    head reuses); the Qwen3 loader strips the markov_head. prefix so the shared
    unprefixed keys match.

  3. Load the DeepSeek-V4 DSpark confidence head. The DeepSeek-V4-Flash
    checkpoint ships mtp.<last>.confidence_head.proj.weight (with-markov,
    bias-less), but DeepseekV4DSpark never loaded it, so adaptive pruning
    silently fell back to sampler proposal probabilities instead of the trained
    confidence. DSparkConfidenceHead now treats the bias as optional, and the
    DeepSeek-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 static
decoding its weights are never loaded, since only the adaptive controller
consumes them.

Test

  • Builds cleanly (including the dspark_confidence_head_test numerical-parity
    unit test, which was updated for the shared Markov head's constructor and
    unprefixed keys).
  • Runtime on DeepSeek-V4-Flash, 8 NPUs, dp_size=2, ShareGPT:
    • static (adaptive off): confidence head is not loaded (no verify failure)
      and serving is unaffected.
    • adaptive (adaptive on): the confidence head loads and is consumed by the
      controller; the run serves all requests with no crash or HCCL hang.

@weizhehuang0827 weizhehuang0827 changed the title refactor: unify DSpark heads and load the DeepSeek-V4 confidence head refactor: unify DSpark heads and load the DeepSeek-V4 confidence head. Aug 19, 2026
ustcfy
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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants