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.
- DFlash always runs eager (both target validate and draft query). ACL graph
  was measured net-negative for DFlash (validate flattens to DECODE, so its
  captured-shape count grows with concurrency: -2% at C=1, -38% at C=16), so
  the target and draft engines force enable_graph off; the shared graph engine
  and the MTP/Eagle3 spec-verify graph path are untouched.
- 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, eager):
- Qwen3-8B: accept rate 34.67% / accept length 6.20, matching 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

tests/core/runtime/spec_input_builder_test.cpp

@@ -529,6 +529,54 @@ TEST(DraftProbsBuilderTest, BuildValidateTensorsDenseInputFallback) {
       draft_probs, torch::tensor({{0.7f}, {0.6f}}, torch::kFloat32)));
 }
 
+TEST(DraftProbsBuilderTest, BuildValidateTensorsFromBlockDenseDefault) {
+  // DFlash default path: selected-only cached probs are scattered back into a
+  // dense [batch, n_spec, vocab] tensor so the shared rejection sampler (and
+  // its fused kernel) always receives dim==3 draft probs, matching MTP.
+  auto token_ids_block = torch::tensor({{1, 2}, {0, 3}}, torch::kInt64);
+  auto probs_block =
+      torch::tensor({{0.2f, 0.7f}, {0.9f, 0.1f}}, torch::kFloat32);
+
+  auto [draft_token_ids, draft_probs] =
+      draftProbs::build_validate_tensors_from_block(
+          token_ids_block,
+          probs_block,
+          /*vocab_size=*/5,
+          /*enable_opt_validate_probs=*/false);
+
+  EXPECT_EQ(draft_token_ids.dim(), 2);
+  EXPECT_EQ(draft_probs.dim(), 3);
+  EXPECT_EQ(draft_probs.size(0), 2);
+  EXPECT_EQ(draft_probs.size(1), 2);
+  EXPECT_EQ(draft_probs.size(2), 5);
+
+  // The scattered dense holds the selected prob at the drafted token and zero
+  // elsewhere; gather recovers the selected values and each row sums to them.
+  auto selected =
+      draft_probs.gather(/*dim=*/-1, draft_token_ids.unsqueeze(-1)).squeeze(-1);
+  auto expected = torch::tensor({{0.2f, 0.7f}, {0.9f, 0.1f}}, torch::kFloat32);
+  EXPECT_TRUE(torch::allclose(selected, expected));
+  EXPECT_TRUE(torch::allclose(draft_probs.sum(/*dim=*/-1), expected));
+}
+
+TEST(DraftProbsBuilderTest, BuildValidateTensorsFromBlockSelectedOnlyOpt) {
+  // Opt-in selected-only path returns the [batch, n_spec] block as-is.
+  auto token_ids_block = torch::tensor({{1, 2}, {0, 3}}, torch::kInt64);
+  auto probs_block =
+      torch::tensor({{0.2f, 0.7f}, {0.9f, 0.1f}}, torch::kFloat32);
+
+  auto [draft_token_ids, draft_probs] =
+      draftProbs::build_validate_tensors_from_block(
+          token_ids_block,
+          probs_block,
+          /*vocab_size=*/5,
+          /*enable_opt_validate_probs=*/true);
+
+  EXPECT_EQ(draft_token_ids.dim(), 2);
+  EXPECT_EQ(draft_probs.dim(), 2);
+  EXPECT_TRUE(torch::allclose(draft_probs, probs_block));
+}
+
 TEST(SpecDecodeInputBuilderTest, MultiBlockDraftSingleRowPerSeq) {
   std::vector<int32_t> kv_seq_lens = to_layout_seq_lens({5, 9});
   torch::Tensor positions = torch::tensor({4, 8}, torch::kInt);

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(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);

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

xllm/core/runtime/acl_graph_executor_impl.cpp

@@ -59,6 +59,24 @@ std::pair<torch::Tensor, torch::Tensor> find_attention_plan_kv_cache(
   return {torch::Tensor(), torch::Tensor()};
 }
 
+bool is_write_context_kv_phase(const ModelInputParams& params) {
+  return params.embedding.write_context_kv;
+}
+
+bool is_chunked_prefill_graph(const ModelInputParams& params) {
+  return params.meta.batch_forward_type.is_chunked_prefill() &&
+         params.is_spec_verify && !is_write_context_kv_phase(params);
+}
+
+uint32_t get_graph_sequence_width(const ModelInputParams& params,
+                                  const runtime::Options& options) {
+  if (is_chunked_prefill_graph(params)) {
+    return static_cast<uint32_t>(
+        std::max<int32_t>(params.meta.q_max_seq_len, 1));
+  }
+  return static_cast<uint32_t>(options.num_decoding_tokens());
+}
+
 }  // namespace
 
 bool AclGraph::capture(CausalLM* model,
@@ -346,6 +364,14 @@ ModelOutput AclGraph::replay(CausalLM* model,
   // Get current NPU stream from libtorch NPU API
   aclrtStream stream = c10_npu::getCurrentNPUStream().stream();
 
+  // Ensure the non-blocking persistent-buffer updates above (tokens/positions/
+  // block_tables/seq_lens) are visible before the graph replays. Without this
+  // ordering edge the replay can read stale inputs. This is a full host-device
+  // sync (a correctness requirement, not a tunable): a lighter device-side
+  // event wait would be preferable but the persistent buffers are updated on
+  // the host prepare stream, so the replay stream must observe them first.
+  aclrtSynchronizeStream(stream);
+
   graph_.replay();
   if (model->is_hybrid_linear_attention()) {
     CHECK(graph_params.has_value())
@@ -386,8 +412,8 @@ AclGraphExecutorImpl::AclGraphExecutorImpl(CausalLM* model,
                                            const torch::Device& device,
                                            const runtime::Options& options)
     : model_(model), args_(args), device_(device), options_(options) {
-  const bool need_update_attn_mask = model->is_hybrid_linear_attention();
   const bool is_hybrid_linear_attn = model->is_hybrid_linear_attention();
+  const bool need_update_attn_mask = is_hybrid_linear_attn;
   graph_slot_count_ =
       ::xllm::ExecutionConfig::get_instance().enable_graph_double_buffer() ? 2
                                                                            : 1;
@@ -428,6 +454,12 @@ ModelOutput AclGraphExecutorImpl::run(const torch::Tensor& tokens,
            << " in_spec_verify_phase: " << in_spec_verify_phase
            << " q_max_seq_len: " << params_single.meta.q_max_seq_len
            << " n_layers: " << args_.n_layers();
+  if (is_write_context_kv_phase(params_single)) {
+    VLOG(kGraphExecutorLogVerboseLevel)
+        << "AclGraphExecutorImpl::run() in eager mode for DFlash context K/V";
+    COUNTER_INC(num_model_execution_total_eager);
+    return model_->forward(tokens, positions, kv_caches, params);
+  }
   if ((!in_decoding_phase && !in_spec_verify_phase) || args_.n_layers() == 1) {
     VLOG(kGraphExecutorLogVerboseLevel)
         << "AclGraphExecutorImpl::run() in eager mode";
@@ -484,9 +516,21 @@ ModelOutput AclGraphExecutorImpl::run(const torch::Tensor& tokens,
   }
   // Keep actual n_tokens for replay output slicing.
   const uint32_t n_tokens = tokens_tensor.size(/*dim=*/0);
-  const uint32_t local_batch_size = n_tokens / options_.num_decoding_tokens();
-  const uint32_t global_batch_size =
-      graph_num_tokens / options_.num_decoding_tokens();
+  const uint32_t graph_sequence_width =
+      get_graph_sequence_width(params_single, options_);
+  CHECK_GT(graph_sequence_width, 0) << "graph sequence width must be positive";
+  if (n_tokens % graph_sequence_width != 0 ||
+      graph_num_tokens % graph_sequence_width != 0) {
+    LOG_FIRST_N(WARNING, 1)
+        << "Falling back to eager mode because graph token count is not "
+           "divisible by graph sequence width. n_tokens="
+        << n_tokens << ", graph_num_tokens=" << graph_num_tokens
+        << ", graph_sequence_width=" << graph_sequence_width;
+    COUNTER_INC(num_model_execution_total_eager);
+    return model_->forward(tokens, positions, kv_caches, params);
+  }
+  const uint32_t local_batch_size = n_tokens / graph_sequence_width;
+  const uint32_t global_batch_size = graph_num_tokens / graph_sequence_width;
 
   // Large decode batches create too many/too large ACL graphs and may OOM.
   // Fall back to eager mode when batch size exceeds the safety threshold.
@@ -498,7 +542,7 @@ ModelOutput AclGraphExecutorImpl::run(const torch::Tensor& tokens,
                             .acl_graph_decode_batch_size_limit()));
   if (global_batch_size > decode_batch_size_limit) {
     LOG_FIRST_N(WARNING, 1)
-        << "Falling back to eager mode because decode batch_size (global="
+        << "Falling back to eager mode because graph batch_size (global="
         << global_batch_size << ", local=" << local_batch_size << ") > "
         << decode_batch_size_limit
         << "; ACL graph is disabled for this request size to avoid OOM. "
@@ -508,7 +552,20 @@ ModelOutput AclGraphExecutorImpl::run(const torch::Tensor& tokens,
     return model_->forward(tokens, positions, kv_caches, params);
   }
 
-  const uint32_t bucket_num_tokens = get_bucket_num_tokens(graph_num_tokens);
+  uint32_t bucket_num_tokens = get_bucket_num_tokens(graph_num_tokens);
+  if (is_chunked_prefill_graph(params_single)) {
+    const uint32_t bucket_batch_size = get_bucket_num_tokens(global_batch_size);
+    bucket_num_tokens = bucket_batch_size * graph_sequence_width;
+  }
+  if (static_cast<int64_t>(bucket_num_tokens) >
+      options_.max_tokens_per_batch()) {
+    LOG_FIRST_N(WARNING, 1)
+        << "Falling back to eager mode because ACL graph bucket tokens ("
+        << bucket_num_tokens << ") exceed max_tokens_per_batch ("
+        << options_.max_tokens_per_batch() << ").";
+    COUNTER_INC(num_model_execution_total_eager);
+    return model_->forward(tokens, positions, kv_caches, params);
+  }
 
   // Check if conditions are suitable for graph execution (replay or capture)
   const auto max_seq_len = args_.max_position_embeddings();