feat: add NPU DFlash block-diffusion speculative decoding.

Introduce DFlash (block-diffusion) speculative decoding on the NPU/ATB
backend. The draft proposes a full block of candidate tokens in a single
non-causal forward from the target's captured intermediate-layer hidden
states, and the target verifies the block via the shared rejection sampler.

Supported targets:
- Qwen3 dense (e.g. Qwen3-8B) and Qwen3-MoE (e.g. Qwen3-Coder-30B-A3B) as
  the verify target. Both are standard-attention, head_dim=128, 4D dense KV
  cache; aux-hidden capture is enabled on qwen3.h and qwen3_moe.h.
- The draft is always a small dense Qwen3 (model_type=DFlashDraftModel),
  regardless of whether the target is dense or MoE; it never runs experts.
- Hybrid / linear-attention targets (e.g. Qwen3.5/Next) are NOT supported in
  this PR (deferred: they need a TORCH-backend draft and target-side capture).

Core pieces:
- DFlashWorkerImpl (runtime): prepare draft query / run draft / validate /
  materialize context K/V / write accepted context, decoupled from MTP.
- DFlashContextKV helper: backend-independent fc -> RMSNorm -> fused per-layer
  K/V linear -> per-layer k_norm + RoPE -> scatter, parameterized by
  DFlashContextKVConfig (head_dim, rotary_dim from partial_rotary_factor,
  rope_theta, k_norm) so it can be reused across draft bodies.
- DFlashQwen3 draft model + registration.
- ACL graph integration: DFlash chunked-prefill graph keys, non-causal draft
  eager fallback, and persistent attn-mask refresh.
- Guard: DFlash and EAGLE-3 reject context parallelism (cp_size > 1), which
  would feed the draft the wrong (lm_head-gathered) hidden states.

Correctness fixes folded in:
- chunked-prefill graph mask reset stale tail columns [kv_len, max_seq_len)
  across replays; a heterogeneous-kv_len batch otherwise let a shorter
  sequence attend to another sequence's KV (graph accept rate 1.11% ->
  34.10%, matching eager).
- qwen3-moe forward asserts capture_idx == layers_to_capture_set_.size()
  before returning aux hidden, mirroring qwen3.h, so a misconfigured
  target_layer_id fails fast instead of feeding uninitialized memory.
- dflash validate reconstructs dense draft probs by default (honoring
  enable_opt_validate_probs like MTP) so the shared rejection sampler and its
  fused kernel always receive dim==3 probs; output-distribution-equivalent to
  selected-only. scatter_selected_to_dense() is shared with
  build_validate_tensors.

Validated (gsm8k, greedy, num_spec=15, graph vs eager):
- Qwen3-8B: accept rate 34.67% / accept length 6.20, graph == eager
  bit-identical; matches the transformers reference (34.87% / 6.23).
- Qwen3-Coder-30B-A3B (MoE, pure TP=2): accept rate 25.0% / 4.75 end-to-end.

Author: fengyan.119

xllm/core/distributed_runtime/llm_engine.cpp

@@ -182,12 +182,17 @@ bool LLMEngine::init_model(MasterStatus master_status) {
   auto model_loader = ModelLoader::create(model_path);
   LOG(INFO) << "Initializing model from: " << model_path;
 
-  tokenizer_ = model_loader->tokenizer();
-  CHECK(tokenizer_ != nullptr);
-
   args_ = model_loader->model_args();
   quant_args_ = model_loader->quant_args();
-  tokenizer_args_ = model_loader->tokenizer_args();
+
+  // Every speculative algorithm drives its draft from the target's hidden
+  // states and shares the target vocabulary, so a draft engine never needs its
+  // own tokenizer.
+  if (!options_.is_draft_engine()) {
+    tokenizer_ = model_loader->tokenizer();
+    CHECK(tokenizer_ != nullptr);
+    tokenizer_args_ = model_loader->tokenizer_args();
+  }
 
   // compute the number of local kv heads and head dim
   const uint32_t world_size = dp_local_tp_size_;
@@ -213,21 +218,23 @@ bool LLMEngine::init_model(MasterStatus master_status) {
             << ", dtype: " << dtype_
             << ", kv_cache_dtype: " << options_.kv_cache_dtype();
 
-  const int64_t tokenizer_vocab_size = tokenizer_->vocab_size();
-  int64_t model_vocab_size = args_.vocab_size();
-  if (tokenizer_vocab_size != model_vocab_size) {
-    // use tokenizer vocab size if model vocab size is not set
-    if (model_vocab_size <= 0) {
-      LOG(WARNING) << "Model vocab size is not set, using tokenizer vocab "
-                      "size: "
-                   << tokenizer_vocab_size;
-      args_.vocab_size(tokenizer_vocab_size);
-    } else if (tokenizer_vocab_size > model_vocab_size) {
-      LOG(WARNING) << "Unsafe vocab mismatch: tokenizer: "
-                   << tokenizer_vocab_size << ", model: " << model_vocab_size;
-    } else {
-      LOG(INFO) << "Tokenizer/model vocab differ: tokenizer="
-                << tokenizer_vocab_size << ", model=" << model_vocab_size;
+  if (tokenizer_ != nullptr) {
+    const int64_t tokenizer_vocab_size = tokenizer_->vocab_size();
+    int64_t model_vocab_size = args_.vocab_size();
+    if (tokenizer_vocab_size != model_vocab_size) {
+      // use tokenizer vocab size if model vocab size is not set
+      if (model_vocab_size <= 0) {
+        LOG(WARNING) << "Model vocab size is not set, using tokenizer vocab "
+                        "size: "
+                     << tokenizer_vocab_size;
+        args_.vocab_size(tokenizer_vocab_size);
+      } else if (tokenizer_vocab_size > model_vocab_size) {
+        LOG(WARNING) << "Unsafe vocab mismatch: tokenizer: "
+                     << tokenizer_vocab_size << ", model: " << model_vocab_size;
+      } else {
+        LOG(INFO) << "Tokenizer/model vocab differ: tokenizer="
+                  << tokenizer_vocab_size << ", model=" << model_vocab_size;
+      }
     }
   }
 

xllm/core/distributed_runtime/speculative_engine.cpp

@@ -55,7 +55,7 @@ SpeculativeEngine::SpeculativeEngine(const runtime::Options& options,
         .devices(options.draft_devices())
         .num_decoding_tokens(1)
         .enable_speculative_decode(/*enable_speculative_decode=*/false)
-        .enable_graph(/*enable_graph=*/false)
+        .enable_graph(options_.is_dflash() ? options_.enable_graph() : false)
         .is_draft_engine(true);
     draft_engine_ =
         std::make_unique<LLMEngine>(draft_engine_options, dist_manager_);
@@ -97,27 +97,9 @@ bool SpeculativeEngine::init_model() {
       return false;
     }
 
-    // check if the tokenizers are compatible
-    const auto* draft_tokenizer = draft_engine_->tokenizer();
-    const auto* target_tokenizer = engine_->tokenizer();
-    if (draft_tokenizer->vocab_size() != target_tokenizer->vocab_size()) {
-      LOG(ERROR) << "draft and target tokenizers have different vocab sizes, "
-                    "draft vocab_size: "
-                 << draft_tokenizer->vocab_size()
-                 << ", target vocab_size: " << target_tokenizer->vocab_size();
-      return false;
-    }
-
-    const std::string test_text = "hello from xllm!";
-    std::vector<int32_t> draft_token_ids;
-    std::vector<int32_t> target_token_ids;
-    if (!draft_tokenizer->encode(test_text, &draft_token_ids) ||
-        !target_tokenizer->encode(test_text, &target_token_ids)) {
-      if (draft_token_ids != target_token_ids) {
-        LOG(ERROR) << "draft and target tokenizers are not compatible";
-        return false;
-      }
-    }
+    // Modern speculative algorithms (MTP / EAGLE / EAGLE3 / DFlash) drive the
+    // draft from the target's hidden states and share the target vocabulary, so
+    // there is no independent draft tokenizer to validate for compatibility.
 
     // check if the max context length are the same
     const auto& draft_model_args = draft_engine_->model_args();

xllm/core/distributed_runtime/worker_server.cpp

@@ -54,6 +54,7 @@ limitations under the License.
 #include "util/net.h"
 #include "util/threadpool.h"
 #include "util/timer.h"
+#include "util/utils.h"
 #if defined(USE_NPU)
 #include "hccl/hccl.h"
 #include "xllm_atb_layers/core/include/atb_speed/base/external_comm_manager.h"
@@ -66,6 +67,13 @@ extern char** environ;
 namespace xllm {
 namespace {
 void handle_signal(int signum) { _exit(0); }
+
+bool uses_aux_hidden_speculative_worker(const runtime::Options& options) {
+  if (!options.enable_speculative_decode()) {
+    return false;
+  }
+  return util::algorithm_captures_aux_hidden(options.speculative_algorithm());
+}
 }  // namespace
 
 void WorkerServer::create_server(const runtime::Options& options,
@@ -364,6 +372,11 @@ WorkerServer::WorkerServer(int32_t local_worker_idx,
                            bool use_spawn_worker)
     : server_name_("DistributeWorkerServer") {
   server_name_.append(std::to_string(options.server_idx()));
+  if (use_spawn_worker && uses_aux_hidden_speculative_worker(options)) {
+    LOG(FATAL) << options.speculative_algorithm()
+               << " does not support spawned workers yet. Disable offline "
+                  "spawn workers or use --speculative_algorithm=MTP.";
+  }
 
   if (worker_type == WorkerType::LLM || worker_type == WorkerType::ELM ||
       worker_type == WorkerType::VLM || worker_type == WorkerType::EVLM ||

xllm/core/framework/config/speculative_config.cpp

@@ -33,7 +33,7 @@ DEFINE_int32(num_speculative_tokens, 0, "Number of speculative tokens.");
 DEFINE_string(speculative_algorithm,
               "MTP",
               "Speculative decoding algorithm. Supported options: MTP, Eagle3, "
-              "Suffix. Default is MTP.");
+              "Suffix, DFlash. Default is MTP.");
 
 DEFINE_int32(speculative_suffix_cache_max_depth,
              64,

xllm/core/framework/model/model_input_params.h

@@ -713,6 +713,10 @@ struct ModelEmbeddingInput {
   // input embedding
   mutable torch::Tensor input_embedding;
 
+  // Input is target hidden fed through the context-K/V materialization path
+  // (the model reads input_embedding, not token_ids).
+  bool write_context_kv = false;
+
   // embedding ids of each sequence
   std::vector<int32_t> embedding_ids;
 
@@ -739,6 +743,7 @@ struct ModelEmbeddingInput {
   ModelEmbeddingInput to(const torch::Device& device) const {
     ModelEmbeddingInput out;
     out.input_embedding = safe_to(input_embedding, device);
+    out.write_context_kv = write_context_kv;
     out.embedding_ids = embedding_ids;
     out.linear_state_ids = linear_state_ids;
     out.linear_state_indices = safe_to(linear_state_indices, device, true);
@@ -883,6 +888,12 @@ struct ExpertInput {
 struct GraphInput {
   torch::Tensor attn_mask;
   torch::Tensor tiling_data;
+  bool is_block_diffusion_graph = false;
+  // Whether the attention is causal. Defaults to true (standard autoregressive
+  // masking). Block-diffusion drafts set this false: the block's candidate
+  // tokens attend to the full context and not to each other, while the target
+  // validate path stays causal.
+  bool causal = true;
   bool use_expanded_decode_for_spec_verify_attention = false;
   torch::Tensor expanded_kv_seq_lens;
   torch::Tensor expanded_block_tables;
@@ -897,6 +908,8 @@ struct GraphInput {
     GraphInput out;
     out.attn_mask = safe_to(attn_mask, device, true);
     out.tiling_data = safe_to(tiling_data, device, true);
+    out.is_block_diffusion_graph = is_block_diffusion_graph;
+    out.causal = causal;
     out.use_expanded_decode_for_spec_verify_attention =
         use_expanded_decode_for_spec_verify_attention;
     out.expanded_kv_seq_lens = safe_to(expanded_kv_seq_lens, device, true);

xllm/core/layers/common/attention_metadata_builder.cpp

@@ -121,9 +121,13 @@ AttentionMetadata build_attention_metadata(
   bool is_spec_verify_chunked_prefill =
       params.is_spec_verify &&
       params.meta.batch_forward_type.is_chunked_prefill();
+  bool is_block_diffusion_chunked_prefill =
+      params.graph.is_block_diffusion_graph &&
+      params.meta.batch_forward_type.is_chunked_prefill();
   bool use_acl_graph = ::xllm::ExecutionConfig::get_instance().enable_graph() &&
                        params.enable_graph &&
-                       (is_decode || is_spec_verify_chunked_prefill) &&
+                       (is_decode || is_spec_verify_chunked_prefill ||
+                        is_block_diffusion_chunked_prefill) &&
                        params.graph.tiling_data.defined();
   if (use_acl_graph) {
     // ACL graph mode: use CustomPagedAttention with tiling_data on device

xllm/core/runtime/CMakeLists.txt

@@ -36,6 +36,7 @@ cc_library(
     mtp_worker_impl.h
     suffix_worker_impl.h
     eagle3_worker_impl.h
+    dflash_worker_impl.h
     spec_input_builder.h
     forward_shared_memory_manager.h
     worker_rendezvous.h
@@ -68,6 +69,7 @@ cc_library(
     mtp_worker_impl.cpp
     suffix_worker_impl.cpp
     eagle3_worker_impl.cpp
+    dflash_worker_impl.cpp
     spec_input_builder.cpp
     forward_shared_memory_manager.cpp
     cp_input_partition.cpp