Skip to content

feat: support DeepSeek V4 DSpark decoding. - #2172

Merged
ustcfy merged 51 commits into
xLLM-AI:mainfrom
chenchuw886:feat/deepseek-v4-dspark
Aug 15, 2026
Merged

feat: support DeepSeek V4 DSpark decoding.#2172
ustcfy merged 51 commits into
xLLM-AI:mainfrom
chenchuw886:feat/deepseek-v4-dspark

Conversation

@chenchuw886

@chenchuw886 chenchuw886 commented Aug 9, 2026

Copy link
Copy Markdown
Contributor

Description

This PR adds DeepSeek-V4-Flash-0731 DSpark speculative decoding support on Ascend NPU, following the Qwen DSpark execution design introduced in #2132.

The implementation:

  • registers the DeepSeek V4 DSpark draft model on the Torch NPU backend and loads all three diffusion stages, vocabulary weights, and the Markov head;
  • supports both original and QuaRot checkpoint layouts, preferring dedicated DSpark vocabulary weights and falling back to the shared target embedding/head;
  • implements block-parallel draft attention, draft-specific DSA metadata, grouped SWA cache writes, and balanced MoE execution across DP/EP ranks;
  • provides native SparseAttnSharedkv semantics when supported and an explicit CANN 9.0-compatible fallback when the installed operator rejects ori_sparse_indices;
  • generalizes target/draft KV-cache memory estimation so both MTP and DSpark reserve their fixed and token-dependent cache storage before allocation;
  • keeps the existing DeepSeek V4 MTP path and exported MTP checkpoint layout compatible;
  • adds aggregate and per-position speculative acceptance metrics;
  • documents the 0731 W8A8 model, launch options, recommended gamma (--num_speculative_tokens=5), SAS modes, and MTP/DSpark weight preparation.

The model used for validation is Eco-Tech/DeepSeek-V4-Flash-0731-w8a8.

Atlas A3 performance

Environment: one Atlas A3 server, 16 NPUs, DP=4, Torch NPU backend, 8,192 average input tokens, 1,024 output tokens, and DSpark gamma=5. Each concurrency level uses the request count shown below.

Concurrency Requests Variant Avg TTFT Avg TPOT Output throughput
32 128 Baseline 5621.2 ms 70.1 ms 422.95 tok/s
32 128 MTP (gamma=1) 7150.8 ms 45.1 ms 613.76 tok/s
32 128 DSpark (gamma=5) 9997.1 ms 28.0 ms 845.74 tok/s
Concurrency Requests Variant Avg TTFT Avg TPOT Output throughput
16 64 Baseline 3776.0 ms 55.8 ms 268.89 tok/s
16 64 MTP (gamma=1) 5331.7 ms 32.9 ms 420.26 tok/s
16 64 DSpark (gamma=5) 6239.5 ms 18.6 ms 646.06 tok/s
Concurrency Requests Variant Avg TTFT Avg TPOT Output throughput
8 32 Baseline 2680.6 ms 44.6 ms 169.27 tok/s
8 32 MTP (gamma=1) 3394.1 ms 27.3 ms 261.11 tok/s
8 32 DSpark (gamma=5) 3807.3 ms 13.7 ms 458.78 tok/s
Concurrency Requests Variant Avg TTFT Avg TPOT Output throughput
1 4 Baseline 849.2 ms 34.8 ms 28.11 tok/s
1 4 MTP (gamma=1) 890.8 ms 21.7 ms 44.39 tok/s
1 4 DSpark (gamma=5) 1239.7 ms 10.6 ms 84.71 tok/s

Across concurrency 32/16/8/1, DSpark reduces average TPOT by 60.1%/66.7%/69.3%/69.5% versus baseline and by 37.9%/43.5%/49.8%/51.2% versus MTP. Draft setup increases TTFT, so the benefit is strongest for long outputs where decode savings amortize the startup overhead.

DSpark acceptance profile

A separate acceptance run used four requests with approximately 30K input tokens and 1K output tokens each, gamma=5, and 951 speculative proposal sequences. The per-position metric is an accepted-prefix survival probability: position p passes when the accepted draft prefix contains at least p + 1 tokens.

Draft position Accepted-prefix count Per-position acceptance
0 798 / 951 83.9%
1 708 / 951 74.4%
2 625 / 951 65.7%
3 569 / 951 59.8%
4 520 / 951 54.7%

Summary:

  • aggregate draft-token acceptance: 3,220 / (951 * 5) = 67.7%;
  • mean accepted draft-prefix length: 3,220 / 951 = 3.39 tokens per proposal;
  • full-block acceptance: 520 / 951 = 54.7%;
  • zero-draft-token acceptance: 153 / 951 = 16.1%;
  • mean committed length, including the guaranteed target token: approximately 1 + 3.39 = 4.39 tokens per decode step. This is the intended interpretation of speculative_mean_tokens_per_decode_step outside terminal padding.

The 3.39/5 mean accepted length and 54.7% full-block acceptance are consistent with the measured TPOT improvement: most target validation steps commit multiple tokens, while only 16.1% fall back to the guaranteed target token alone.

Related Issues

Change Type

  • Bug fix
  • New feature
  • Performance improvement
  • Refactor
  • Documentation
  • Test
  • Build or CI

Pull Request Checklist

PR Title and Commit Messages

  • The PR title and each commit message follow the xLLM commit format: <type>: <subject>.

Pre-commit Checks

  • I have installed pre-commit by running pip install pre-commit or an equivalent command.
  • I have installed the hooks with pre-commit install.
  • I have run pre-commit run --all-files and fixed any reported issues.

Self Review

  • I have self-reviewed the code according to .agents/skills/code-review/references/custom-code-style.md, especially code written or assisted by AI.
  • I have rebased this PR onto the latest main branch.

Build and Test Coverage

  • Tests have been added or updated as needed.
  • CUDA: python setup.py build test has passed on a CUDA machine.
  • NPU: python setup.py build test has passed on an NPU machine.
  • MLU: python setup.py build test has passed on an MLU machine.

@chenchuw886
chenchuw886 marked this pull request as ready for review August 9, 2026 15:09
@chenchuw886
chenchuw886 force-pushed the feat/deepseek-v4-dspark branch 2 times, most recently from 21f992c to db1b5e5 Compare August 10, 2026 03:16
yinjiawei01
yinjiawei01 previously approved these changes Aug 10, 2026
Comment thread xllm/core/distributed_runtime/speculative_output_metrics.h Outdated
Comment thread xllm/models/llm/dspark_weight_source.h Outdated
Comment thread tests/models/model_registry_test.cpp Outdated

@ustcfy ustcfy left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code-quality (non-blocking) suggestions from a reuse/simplification/efficiency/altitude pass. None are correctness issues; take or leave per your judgment.

Comment thread xllm/core/layers/npu_torch/deepseek_sparse_attention.cpp Outdated
Comment thread xllm/core/runtime/dflash_worker_impl.h Outdated
Comment thread xllm/core/runtime/worker_impl.cpp Outdated
Comment thread xllm/models/llm/deepseek_v4_dspark.h
Comment thread xllm/core/common/speculative_output_metrics.h Outdated
Comment thread xllm/core/distributed_runtime/worker_service.cpp Outdated
@chenchuw886
chenchuw886 force-pushed the feat/deepseek-v4-dspark branch 2 times, most recently from a0f83f1 to 8f646d4 Compare August 11, 2026 03:07
@bupt-ycq

Copy link
Copy Markdown
Contributor

1:DSpark 配置逻辑不应整体放在 USE_NPU 下

定位:/workspace/xllm-pr-2172-review/xllm/core/runtime/worker_impl.cpp:166、/workspace/xllm-pr-2172-review/xllm/core/runtime/worker_impl.cpp:1512

configure_block_diffusion_model_args() 中的大部分逻辑属于 block-diffusion/DSpark 的算法配置,而不是
NPU 特有实现。例如 draft checkpoint 选择、draft layer 数量校验、block size 设置、capture layer 解析以
及 target/draft model 处理,这些逻辑 MLU 也需要。

如果把整个函数放在 USE_NPU 下,其他后端只能重复实现 checkpoint contract,后续容易产生行为不一致。建议
将公共配置逻辑移出 NPU 宏,只把真正的平台差异放进 platform/model adapter,例如:

  • draft model 的注册名称;
  • checkpoint layer ID 到运行时 capture point 的映射;
  • native SAS capability;
  • 后端特有的 ModelArgs 调整。

2:capture layer 的 +1 属于后端语义

定位/workspace/xllm-pr-2172-review/xllm/core/runtime/worker_impl.cpp:145

read_dflash_capture_layer_ids() 当前同时负责 checkpoint 解析和 layer_id + 1 转换。但这个偏移量反映的是
当前 NPU model 在哪里触发 hidden-state capture,并不是 checkpoint 本身的公共语义。

其他后端可能在 layer 执行完成后 capture,并直接使用 checkpoint 中的 layer ID。如果直接复用这个公共
runtime helper,可能会静默 capture 到错误的 hidden states。

建议拆分成两个步骤:

  1. 从 checkpoint 解析原始 target layer IDs;
  2. 通过 model/backend policy 将 checkpoint ID 映射为运行时 capture point。

3:CANN/NPU 专用配置不应无条件暴露给所有后端

定位:/workspace/xllm-pr-2172-review/xllm/core/framework/config/speculative_config.cpp:69、/workspace/
xllm-pr-2172-review/xllm/core/framework/config/speculative_config.h:103

enable_dspark_native_sas 注册在公共 speculative 配置中,但帮助信息和实际消费逻辑都是 CANN/NPU 特有的。
这样会导致 MLU 等非 NPU 构建也暴露并接受该参数,但参数实际上不会生效。

建议考虑以下方案之一:

  • 将它移动到 NPU kernel config;
  • 在名称中明确 NPU 范围;
  • 抽象成后端无关的 DSpark attention mode,并通过平台 capability 校验。

对不支持该模式的后端,建议明确报错,不要静默忽略。

@chenchuw886
chenchuw886 force-pushed the feat/deepseek-v4-dspark branch 2 times, most recently from e02784b to a5f26d4 Compare August 11, 2026 07:05
@chenchuw886

Copy link
Copy Markdown
Contributor Author

1:DSpark 配置逻辑不应整体放在 USE_NPU 下

定位:/workspace/xllm-pr-2172-review/xllm/core/runtime/worker_impl.cpp:166、/workspace/xllm-pr-2172-review/xllm/core/runtime/worker_impl.cpp:1512

configure_block_diffusion_model_args() 中的大部分逻辑属于 block-diffusion/DSpark 的算法配置,而不是
NPU 特有实现。例如 draft checkpoint 选择、draft layer 数量校验、block size 设置、capture layer 解析以
及 target/draft model 处理,这些逻辑 MLU 也需要。
如果把整个函数放在 USE_NPU 下,其他后端只能重复实现 checkpoint contract,后续容易产生行为不一致。建议
将公共配置逻辑移出 NPU 宏,只把真正的平台差异放进 platform/model adapter,例如:

  • draft model 的注册名称;
  • checkpoint layer ID 到运行时 capture point 的映射;
  • native SAS capability;
  • 后端特有的 ModelArgs 调整。

2:capture layer 的 +1 属于后端语义

定位/workspace/xllm-pr-2172-review/xllm/core/runtime/worker_impl.cpp:145

read_dflash_capture_layer_ids() 当前同时负责 checkpoint 解析和 layer_id + 1 转换。但这个偏移量反映的是
当前 NPU model 在哪里触发 hidden-state capture,并不是 checkpoint 本身的公共语义。
其他后端可能在 layer 执行完成后 capture,并直接使用 checkpoint 中的 layer ID。如果直接复用这个公共
runtime helper,可能会静默 capture 到错误的 hidden states。
建议拆分成两个步骤:

  1. 从 checkpoint 解析原始 target layer IDs;
  2. 通过 model/backend policy 将 checkpoint ID 映射为运行时 capture point。

3:CANN/NPU 专用配置不应无条件暴露给所有后端

定位:/workspace/xllm-pr-2172-review/xllm/core/framework/config/speculative_config.cpp:69、/workspace/ xllm-pr-2172-review/xllm/core/framework/config/speculative_config.h:103

enable_dspark_native_sas 注册在公共 speculative 配置中,但帮助信息和实际消费逻辑都是 CANN/NPU 特有的。
这样会导致 MLU 等非 NPU 构建也暴露并接受该参数,但参数实际上不会生效。
建议考虑以下方案之一:

  • 将它移动到 NPU kernel config;
  • 在名称中明确 NPU 范围;
  • 抽象成后端无关的 DSpark attention mode,并通过平台 capability 校验。

对不支持该模式的后端,建议明确报错,不要静默忽略。

感谢检视,已调整:
1、将 block-diffusion/DSpark 的公共 checkpoint 解析和模型配置逻辑从 USE_NPU 条件编译中拆出,集中到 block_diffusion_model_config.cpp。平台相关行为仅保留 draft model 注册、capture point 映射和后端能力配置。
2、将 checkpoint target layer ID 的解析与运行时 capture point 映射拆开。checkpoint 保留原始 layer ID,NPU 通过 platform policy 映射为 layer_id + 1,其他后端默认保持原始 ID。
3、将 enable_dspark_native_sas 从公共 SpeculativeConfig 移至 NPU KernelConfig。非 NPU 构建若在 JSON 中显式启用该配置会直接报错,避免静默忽略。

Kang-Meng
Kang-Meng previously approved these changes Aug 11, 2026
Kang-Meng
Kang-Meng previously approved these changes Aug 12, 2026
guojinrong-nn
guojinrong-nn previously approved these changes Aug 12, 2026
@ustcfy

ustcfy commented Aug 12, 2026

Copy link
Copy Markdown
Collaborator

将 checkpoint target layer ID 的解析与运行时 capture point 映射拆开,个人认为是过度设计。

ustcfy added 20 commits August 14, 2026 17:12
- Dedup DSpark SWA window formula: drop dflash local wrapper and call
  layer::deepseek_v4_ori_window_left directly.
- Drop cached speculative_position_labels_ in WorkerService; inline
  std::to_string at the sole call site.
- Precompute an unordered_map slot lookup in AuxHiddenCapture so
  should_capture and capture_layer stop scanning layers_to_capture per
  layer per forward.
- Tighten DFlashWorkerImpl::draft_sas_mode() to CHECK draft_impl_ (only
  reached after init).
- Fold DSparkMarkovHead::initialize into its constructor; drop the
  two-phase init from DeepseekV4DSpark and DSparkQwen3.
- Cache dspark SWA layer/cache/block_size in DeepseekV4ModelImpl so
  build_dspark_swa_metadata indexes directly instead of scanning
  caches_info_ per forward.
- Wrap configure_deepseek_v4_dspark_args in #if USE_NPU and drop the
  unreachable non-NPU LOG(FATAL) branch.
- Drop redundant .to(torch::kLong) on already-kLong block_columns before
  block_table.gather in build_dspark_swa_indices.
- AuxHiddenCapture: swap unordered_map back to std::vector<int32_t>
  with std::find. For k <= 5 captured layers the cache-line scan beats
  a hash lookup; also drops the redundant num_captured_ field.
- WorkerService: restore the cached speculative_position_labels_ vector
  and drop the per-step std::to_string in the metrics loop.
Critical:
- build_dspark_swa_indices: replace clamp with remainder so the block
  column follows the SWA ring semantics used by calc_ring_slot_id and
  slot_for_position. clamp was pinning wrapped indices to the last
  block table entry once kv_len exceeded ring_size*block_size, causing
  native DSpark SAS to read from the wrong physical block.
- Add DSparkNativeSwaIndicesWrapAroundRingBuffer test covering the
  wrap.

Important:
- Move the process-wide static speculative_metrics_mutex onto
  WorkerService as a per-instance member so multi-DP-in-one-process
  runs stop contending on a single global lock.
- Skip the .to(kInt64) copy in calculate_speculative_output_stats when
  next_tokens is already int64, and add a CHECK_LE guarding
  accepted_per_position bounds.

Minor:
- /*value=*/ annotation on the compress_ratios fill literal.
- /*default=*/ annotation on the dspark_num_mtp_layers fallback.
Replace the element-wise copy of new_cache_slots into out_slots_acc with a
single std::memcpy on the block-parallel passthrough branch. Same semantics
(both sides are contiguous int32), fewer instructions in the DSA metadata
build path that runs every draft forward.
- Make record_speculative_metrics_from_output a WorkerService member so
  options_/labels_/mutex_/totals are read directly instead of passed
  through a 6-arg free-function signature at both call sites.
- Drop the .to(kInt64) ternary in calculate_speculative_output_stats;
  Tensor::to() is already a no-op when dtype matches.
- Drop three trivial tests (default bool, one-branch predicate, string
  equality) and keep the ones that verify real invariants (SAS mode
  dispatch, dedicated-vs-fallback load order, dp count sync,
  attention-metadata-only invalidation).
- Compress verbose header comments on calc_ring_slot_id,
  build_grouped_prefill_swa_slots, write_context_kv, and
  SpeculativeOutputStats / calculate_speculative_output_stats to
  single lines (or drop where the function name is self-descriptive).
- Shorten the DFlash validate scaling comment to the load-bearing
  N -> N+1 invariant.
- Cache DSparkSasMode on DFlashWorkerImpl once at init: draft_impl_ /
  sample_from_anchor() / model_type() are fixed after load, so the
  string-equality classify was firing on every decode step.
- Hoist the loop-invariant RowSpec fields in build_query_rows out of the
  inner query_idx loop so each of the num_sequences * query_width rows
  only sets the per-row seq_id / token_id / position_offset.
- Route DeepseekV4MtpModelImpl's ori_win_left through
  layer::deepseek_v4_ori_window_left so the compatibility-fallback
  formula lives in one place with the DSA layer.
Revert the coincidental capture_layer refactor (std::distance-based slot
placement, <iterator> include, CHECK_LE) back to upstream's capture_idx_
scheme. layers_to_capture is always ascending and callers iterate layers
in order, so the extra machinery was over-engineering.

Only the three additions DSpark actually needs remain: enabled(),
should_capture(), and the finalize(residual) overload. Update the aux
capture test to use an ascending {0, 2} layer list matching real usage.
Eliminate the two-phase-construction pattern (default construct, then
call init() from the model ctor body). The class now takes the same
three arguments in its constructor and each model member-initializes
aux_capture_ instead of forgetting-to-init being possible.

Also drops the redundant if (!layers_to_capture().empty()) guard around
aux_capture_ construction in DeepseekV4ModelImpl — the constructor
already short-circuits on that condition.
…reset counter.

capture_layer now computes slot_idx = std::distance(begin, find(layer_idx))
directly, so slot placement is a pure function of layer_idx rather than
capture call order. That removes the need for:
- the capture_idx_ counter member,
- the reset_capture_index() call each model forward had to issue,
- the finalize() CHECK_EQ that only existed to catch "forgot to reset".

Model forward loops (deepseek_v4, qwen3, qwen3_moe) and the test drop
their reset_capture_index() calls. The class is now stateless across
forwards; slot correctness relies solely on the ascending
layers_to_capture invariant plus the existing loop-based capture
contract, which every caller already satisfies.
--num_speculative_tokens now defaults to the checkpoint's dspark_block_size
and only overrides it when the user sets a different value (with a warning).
Also drop the tautological CHECK_EQ + misleading "must match" message in
DFlashWorkerImpl::init_model, and collapse the 3-if mask_token_id fallback
chain into a single JsonReader::value_or({keys...}, -1) call.
Drop the n_mtp_layers / dspark_num_mtp_layers fallback chain plus the
hardcoded default of 3. Neither key exists in the DeepSeek-V4-Flash-0731
config; the branch always hit the coincidental default that only matched
because len(dspark_target_layer_ids) happens to be 3. Read the array
directly and take its size so the draft layer count stays in lockstep
with the captured target layers.
- Add DFlashWorkerImpl::draft_use_block_parallel_rows() /
  draft_batch_forward_type() private helpers so step_empty and
  prepare_query_inputs stop duplicating the DSparkSasMode -> bool /
  BatchForwardType mapping. A future third SAS mode now lands in one
  place instead of drifting between the two call sites.
- Rename build_dspark_swa_indices' num_speculative_tokens parameter to
  dspark_block_size. The only caller passes ModelArgs::dspark_block_size,
  and the two happen to coincide today only because DSpark is configured
  to make them match; the semantics being computed is the diffusion
  block width, not the raw speculative-token count.
- Remove is_deepseek_v4_dspark_draft() wrapper: it forwarded to
  util::is_deepseek_v4_dspark_model_type and only served to attach a
  head/embedding-aliasing comment; move that comment to the actual call
  site and inline the predicate there. classify_dspark_sas_mode now
  reads util:: directly.
- Remove DFlashWorkerImpl::draft_sas_mode() getter: after this round's
  refactor both call sites went through draft_use_block_parallel_rows /
  draft_batch_forward_type, so the getter became dead.
- Drop the two capture_layer_ids temporaries in WorkerImpl::init_model:
  args.layers_to_capture() has an rvalue overload, so the read/move pair
  collapses to one call.
- aux_hidden_capture.h: drop "Fuse residual into preallocated slice." —
  torch::add_out(slot, h, residual) is self-explanatory.
- deepseek_v4.h: drop "Precomputed location of the first SWA cache..."
  atop the three dspark_swa_* members; the field names already say it.
- dflash_worker_impl.cpp: compress the 4-line cpu_int_vec_to_device
  comment to one line; drop the pre-refactor idiom quote.
- dflash_worker_impl.cpp: compress the DSpark draft_options comment
  from two lines to two shorter lines.
- speculative_worker_impl.h: compress estimate_kv_cache_capacity_with_draft
  header comment from four lines to two.
- deepseek_v4_indexer_tests.cpp: rename /*num_speculative_tokens=*/ to
  /*dspark_block_size=*/ to match the renamed parameter.
The raw source() getter was read only by acl_graph_executor_test; assert
the DEDICATED terminal state through the public should_load() API instead.
After reserving both the target and draft constant SWA pools, the remaining
budget fits 2 shared token units (not 1), yielding c4_count=64/c128_count=2.
The prior 32/1 expectation under-counted; the reserved-budget arithmetic and
the LE invariant both confirm 2 units.
deepseek_v4_ori_window_left is declared in the NPU-only
npu_torch/deepseek_sparse_attention.h, whose include is already USE_NPU
gated, but the warning that calls it in DFlashWorkerImpl::init_model sat
outside any guard, breaking the MLU and CUDA builds. Wrap the native-SAS
warning in #if defined(USE_NPU).
Three DeepSeek-V4 DSpark defects surfaced by the test suite:

- The model loader copied dspark_block_size into the shared target args,
  arming non-causal DSpark attention on the target model. Leave it unset in
  the loader; the draft worker now defaults it from the checkpoint's
  block_size (overridable by --num_speculative_tokens) via read_block_size().

- build_dspark_swa_indices used true division for the block-column index,
  yielding a float tensor that torch::gather rejects. Use torch::floor_divide.

- The WrapAroundRingBuffer test hard-coded start_pos = kv - window, dropping
  the q_len term; correct the expected slots to (kv - q_len) - window.
@ustcfy
ustcfy force-pushed the feat/deepseek-v4-dspark branch from e3172ec to 7a14f4e Compare August 14, 2026 09:37
ustcfy added 5 commits August 14, 2026 20:31
…p-unused-source-accessor

# Conflicts:
#	xllm/core/distributed_runtime/worker_service.cpp
#	xllm/core/runtime/speculative_worker_impl.h
#	xllm/models/llm/npu/qwen3_dspark.h
…p-unused-source-accessor

# Conflicts:
#	tests/core/runtime/acl_graph_executor_test.cpp
#	xllm/core/runtime/speculative_worker_impl.cpp
Correct the misleading /*num_speculative_tokens=*/ annotation on
build_dspark_swa_indices (the 5th arg is dspark_block_size), and note that the
base loader deliberately leaves dspark_block_size at 0 until
configure_deepseek_v4_dspark_args applies it.
Derive speculative_mean_tokens_per_decode_step from the process-global
committed/drafts counters instead of per-instance totals, so multiple
WorkerService instances (multi-DP) converge on one aggregate rather than
overwriting the shared gauge with a rank-local value. Drop the now-unused
per-instance totals and mutex, add COUNTER_VALUE, and note why FlashComm1 is
skipped during aux capture.
@yinjiawei01
yinjiawei01 self-requested a review August 15, 2026 05:05
@ustcfy
ustcfy merged commit 67459ab into xLLM-AI:main Aug 15, 2026
15 of 26 checks passed
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.

7 participants