feat: add DFlash block-diffusion speculative decoding on npu device.

[ustcfy]

描述

在 NPU/ATB 后端引入 DFlash（block-diffusion）投机解码。draft 用 target 捕获的中间层 hidden，一次非因果前向并行产出一整块候选 token；target 用共享的 rejection sampler 一次性验证整块，取最长匹配前缀 + 1 个 bonus token。

支持的 target：

• Qwen3 dense（如 Qwen3-4B/8B）与 Qwen3-MoE（如 Qwen3-Coder-30B-A3B），均为标准 attention、head_dim=128、4D dense KV cache；aux-hidden 捕获在 qwen3.h / qwen3_moe.h 打开。
• draft 恒为小 dense Qwen3（model_type=DFlashDraftModel），不论 target 是 dense 还是 MoE，draft 自身从不跑 expert。
• 本 PR 不支持 hybrid / linear-attention target（如 Qwen3.5/Qwen3-Next 系）。Qwen3.5 走 NPU_TORCH backend + Qwen3Next hybrid（linear + full attn 交替），本 PR 的 ATB-only DFlash draft 无法直接消费，作为独立 followup。

核心组件：

• DFlashWorkerImpl（runtime）：prepare_query_inputs / run_decode_draft / run_validate / write_context_kv / write_target_context_to_cache，与 MTP 解耦。
• context-KV 投影内联在 DFlashQwen3ModelImpl：fc → RMSNorm → fused per-layer K/V linear → per-layer k_norm + RoPE → scatter。权重加载跨 checkpoint 分片累积后再融合成 layer-major 的 fused K/V 权重。
• DFlashQwen3 draft 模型 + REGISTER_CAUSAL_MODEL(DFlashDraftModel, ...) 注册。
• 守卫：DFlash（dflash_worker_impl.cpp）与 EAGLE-3（eagle3_worker_impl.cpp）拒绝 context parallelism（cp_size > 1），否则会把错误的（lm_head-gathered）hidden 喂给 draft。
• Draft config use_sliding_window=true 被 xLLM NPU FIA silently 忽略（aclnn FIA hard-code pre_tokens=INT_MAX）。本 PR 不加 fail-fast 阻断，与主流框架一致——真正的 SWA 支持是独立 track。

顺带修复的正确性问题：

• qwen3_moe.h forward 返回 aux hidden 前补 CHECK_EQ(capture_idx, layers_to_capture_set_.size())（与 qwen3.h 对齐），使越界的 target_layer_id 快速失败，而不是把未初始化内存喂给 draft。
• DFlash validate 默认重建 dense draft probs（像 MTP 一样响应 enable_opt_validate_probs），使共享 rejection sampler 及其 fused kernel 恒收到 dim==3 的 probs；与 selected-only 输出分布等价。dense scatter 在 build_validate_tensors_from_block 中内联，与 MTP build_validate_tensors 语义一致。

相关 Issue

变更类型

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

验证

数据集 gsm8k、greedy、num_speculative_tokens=15、block_size=16、20 prompts、concurrency=1：

Target	Draft (z-lab)	Draft Accept Rate	Avg Accepted Length	Throughput (tok/s)
Qwen3-4B	Qwen3-4B-DFlash-b16	32.46%	5.869	299.17
Qwen3-8B	Qwen3-8B-DFlash-b16	33.36%	6.004	218.05
Qwen3-Coder-30B-A3B-Instruct (TP=2)	Qwen3-Coder-30B-A3B-DFlash	24.70%	4.704	150.59

主二进制 python setup.py build --device npu 编译通过。

Reviewer Notes

• graph 默认关闭。实测 DFlash 开 graph 当前非盈利（Qwen3-8B, gsm8k, num_spec=15）：validate 走 DECODE + graph-paged attention 后，摊平的 num_seq × (num_spec+1) 行数会桶化成多个 shape 档，capture + padding 开销盖过 launch 节省。近期实测吞吐相对 eager 大致为 C=1/4/16 −0.5% / −4% / −3%，中并发 C=32/64/128 −3% / −13% / −7%（acl_graph_decode_batch_size_limit 默认 16，batch>16 触发 fallback eager 但仍付 capture-attempt 开销）。盈利化需要固定-shape 的 verify capture（uniform-decode 做法），是独立工作，不在本 PR。
• draft query 因非因果 chunked-prefill 强制 eager（ACL graph 无法在 replay 时刷新 ATB attention plan 的递增 kv_len），这是正确性要求的 fallback。
• 单测（spec_input_builder 等）暂留本地，未包含在本 PR。

Followups

• Qwen3Next hybrid backend（TORCH）的 DFlash draft path，解锁 Qwen3.5/Qwen3.6 系（4B / 9B / 27B / 35B-A3B / 122B-A10B / 397B-A17B）。
• 固定-shape verify graph capture（提升 graph 盈利性）。
• xLLM NPU SWA 端到端支持（长 context 场景 Qwen3.5-* draft 精度）。

[gemini-code-assist]

Code Review

This pull request introduces the 'DFlash' speculative decoding algorithm, implementing the DFlashWorkerImpl and DFlashQwen3Model to support block-diffusion speculative decoding. Key changes include skipping tokenizer initialization for draft engines, handling non-causal draft queries eagerly, and materializing target hidden states into the draft's KV cache. The code review feedback highlights opportunities to improve robustness and performance, specifically by resolving inconsistent null-pointer checks in DFlashWorkerImpl::init_model, investigating the performance bottleneck of stream synchronization in AclGraphExecutorImpl::replay, marking DFlashQwen3ModelImpl as final per the style guide, and fixing a signed/unsigned comparison warning in Qwen3MoeModelImpl.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[Copilot]

Pull request overview

This PR introduces NPU-side DFlash (block-diffusion) speculative decoding for Qwen3 targets, including a dedicated draft model path that can materialize context K/V from captured target hidden states, plus ACL-graph support for DFlash chunked-prefill verify and several correctness guards/fixes around auxiliary hidden capture and persistent attention masks.

Changes:

• Add DFlash draft model + context-K/V materialization helper, and wire a new DFlashWorkerImpl speculative worker into runtime/engine/model registry.
• Extend ACL graph execution and persistent-parameter handling to support DFlash chunked-prefill verification and correct persistent attention-mask reuse.
• Update speculative config/options/tokenizer handling to support draft engines and the new "DFlash" algorithm.

Reviewed changes

Copilot reviewed 25 out of 25 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
xllm/pybind/args.py	Expose DFlash as a CLI speculative algorithm option.
xllm/models/models.h	Register/include the new NPU DFlash Qwen3 model header.
xllm/models/llm/npu/qwen3.h	Enable aux-hidden capture for Eagle3/DFlash targets; add chunked-prefill attn-mask sourcing; add capture count guard; add overridable mask builder.
xllm/models/llm/npu/qwen3_moe.h	Same aux-hidden capture gating/guarding for Qwen3-MoE targets.
xllm/models/llm/npu/qwen3_dflash.h	New DFlash draft model impl (Qwen3-based) + registry/args loader; custom non-causal mask override.
xllm/models/llm/npu/dflash_context_kv.h	New helper to project captured target hidden into per-layer draft K/V cache (fc→norm→fused KV linear→k_norm+RoPE→scatter).
xllm/models/CMakeLists.txt	Ensure models link NPU xllm_ops needed by new scatter ops.
xllm/core/util/utils.h	Add helpers for DFlash draft model-type detection and “aux-hidden-capturing” algorithm whitelist.
xllm/core/runtime/worker.cpp	Instantiate DFlashWorkerImpl when --speculative_algorithm=DFlash.
xllm/core/runtime/worker_impl.cpp	Fix graph-prepare guard for speculative workers; adjust tokenizer usage for draft engines; add DFlash layers_to_capture plumbing from draft config.
xllm/core/runtime/spec_input_builder.h	Declare build_validate_tensors_from_block for block drafters (DFlash).
xllm/core/runtime/spec_input_builder.cpp	Add shared scatter_selected_to_dense helper; implement build_validate_tensors_from_block.
xllm/core/runtime/options.h	Add Options::is_dflash() helper keyed on the configured algorithm string.
xllm/core/runtime/eagle3_worker_impl.cpp	Add cp_size>1 guard for aux-hidden-based draft driving (EAGLE-3).
xllm/core/runtime/dflash_worker_impl.h	New DFlash speculative worker interface.
xllm/core/runtime/dflash_worker_impl.cpp	New DFlash speculative worker implementation (draft query/build/validate/materialize context K/V/caching).
xllm/core/runtime/CMakeLists.txt	Build and link the new DFlash worker implementation.
xllm/core/runtime/acl_graph_persistent_param.cpp	Add DFlash chunked-prefill support and fix persistent attn-mask tail reuse correctness.
xllm/core/runtime/acl_graph_executor_impl.cpp	Add DFlash graph keys + chunked-prefill handling + eager fallbacks for non-causal draft query and context-K/V phase.
xllm/core/layers/common/attention_metadata_builder.cpp	Allow ACL graph path for block-diffusion chunked-prefill metadata.
xllm/core/framework/model/model_input_params.h	Add write_context_kv embedding flag and DFlash graph flags (is_block_diffusion_graph, causal).
xllm/core/framework/config/speculative_config.cpp	Document DFlash in the supported algorithm flag help text.
xllm/core/distributed_runtime/worker_server.cpp	Block spawned-worker mode for aux-hidden speculative workers (Eagle3/DFlash).
xllm/core/distributed_runtime/speculative_engine.cpp	Allow graph on draft engine for DFlash; remove draft/target tokenizer compatibility checks.
xllm/core/distributed_runtime/llm_engine.cpp	Skip tokenizer initialization for draft engines; guard vocab-size checks accordingly.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[ustcfy]

感谢 @gemini-code-assist 的 review。已在 commit 0540cf75 处理，逐条回复：

1. DFlashWorkerImpl::init_model 空指针检查不一致 — 已采纳。 在函数开头加 CHECK(impl_ != nullptr) / CHECK(draft_impl_ != nullptr)，并移除后续 if (impl_ != nullptr ...) / if (draft_impl_ != nullptr ...) 的冗余检查，与 estimate_kv_cache_capacity / allocate_kv_cache 保持一致。

2. AclGraphExecutorImpl::replay 的 aclrtSynchronizeStream 阻塞 — 已加注释说明，暂不改逻辑。 这个全流同步是正确性要求（persistent buffer 在 host prepare stream 更新，replay stream 必须先看到），非可调项。但经实测，它不是 DFlash graph 非盈利的主因：真因是 validate 走 DECODE 摊平后 num_seq×16 的行数被桶化成约 15 个 shape 档，capture + padding 开销盖过 launch 节省（graph 全程生效、无 fallback，实测已确认）。已在注释中写明这一点。用设备侧 event wait 替代是合理的后续优化方向，但需单独验证，不在本 PR。

3. DFlashQwen3ModelImpl 未标 final — 已采纳。 已加 final，与同文件的 DFlashQwen3ForCausalLMImpl 及 style guide §4 对齐。

4. Qwen3MoeModelImpl 的 signed/unsigned 比较 — 经核实不成立，未改。 该处 capture_idx 声明为 size_t（qwen3_moe.h:330），layers_to_capture_set_.size() 返回 size_t，两者同为无符号，CHECK_EQ 无 sign-compare 问题。qwen3.h 需要 static_cast<int64_t> 是因为它那边 capture_idx 是 int64_t（有符号）。若在此加 static_cast<int64_t> 反而会引入 int64_t vs size_t 的混合比较。故保持不转换。

[gemini-code-assist]

感谢你的详细回复和对这些技术细节的深入说明。你对 DFlashWorkerImpl::init_model 中空指针检查一致性的采纳，以及对 AclGraphExecutorImpl::replay 中 aclrtSynchronizeStream 阻塞原因的解释（即正确性要求而非性能瓶颈），非常清晰且专业。

关于 Qwen3MoeModelImpl 中 size_t 与 size_t 比较的分析也非常准确，确实避免了不必要的类型转换带来的潜在风险。这些调整对于提升代码的健壮性和可维护性都非常有价值。

既然该 PR 已经涵盖了核心的 DFlash 投机解码逻辑及相关的正确性修复，期待后续在设备侧 event wait 等优化方向上的进一步探索。感谢你对代码质量的严谨把控。

[weizhehuang0827]

The vocabulary in eagle3 might be inconsistent with the original model? Therefore, eagle3_worker_impl now has a map operation.

[weizhehuang0827]

Is it necessary to create a separate file in this model directory specifically for DFlash?