feat: support DeepSeek V4 DSpark decoding. - #2172
Conversation
21f992c to
db1b5e5
Compare
ustcfy
left a comment
There was a problem hiding this comment.
Code-quality (non-blocking) suggestions from a reuse/simplification/efficiency/altitude pass. None are correctness issues; take or leave per your judgment.
a0f83f1 to
8f646d4
Compare
|
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
2:capture layer 的 +1 属于后端语义 定位/workspace/xllm-pr-2172-review/xllm/core/runtime/worker_impl.cpp:145
3:CANN/NPU 专用配置不应无条件暴露给所有后端 定位:/workspace/xllm-pr-2172-review/xllm/core/framework/config/speculative_config.cpp:69、/workspace/
|
e02784b to
a5f26d4
Compare
感谢检视,已调整: |
|
将 checkpoint target layer ID 的解析与运行时 capture point 映射拆开,个人认为是过度设计。 |
- 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.
This reverts commit bc9a35c.
This reverts commit c6b14f5.
…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.
e3172ec to
7a14f4e
Compare
…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.
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:
ori_sparse_indices;--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.
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
ppasses when the accepted draft prefix contains at leastp + 1tokens.Summary:
3,220 / (951 * 5) = 67.7%;3,220 / 951 = 3.39tokens per proposal;520 / 951 = 54.7%;153 / 951 = 16.1%;1 + 3.39 = 4.39tokens per decode step. This is the intended interpretation ofspeculative_mean_tokens_per_decode_stepoutside 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
Pull Request Checklist
PR Title and Commit Messages
<type>: <subject>.Pre-commit Checks
pre-commitby runningpip install pre-commitor an equivalent command.pre-commit install.pre-commit run --all-filesand fixed any reported issues.Self Review
.agents/skills/code-review/references/custom-code-style.md, especially code written or assisted by AI.mainbranch.Build and Test Coverage
python setup.py build testhas passed on a CUDA machine.python setup.py build testhas passed on an NPU machine.python setup.py build testhas passed on an MLU machine.