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

docs/src/content/docs/en/cli_reference.md

@@ -183,7 +183,7 @@ xLLM uses gflags to manage service startup parameters. `--model <PATH>` is the o
 | `draft_model` | `string` | `""` | Draft model path. See [MTP](/en/features/mtp/) for MTP usage. |
 | `draft_devices` | `string` | `""` | Devices used by the draft model, for example `npu:0` or `npu:0,npu:1`. If omitted, uses the target model devices when speculative decoding is enabled. |
 | `num_speculative_tokens` | `int32` | `0` | Number of speculative tokens generated per speculative decoding step. |
-| `speculative_algorithm` | `string` | `"MTP"` | Speculative decoding algorithm. Supported values: `MTP`, `Eagle3`, `Suffix`. |
+| `speculative_algorithm` | `string` | `"MTP"` | Speculative decoding algorithm. Supported values: `MTP`, `Eagle3`, `Suffix`, `DFlash`. |
 | `speculative_suffix_cache_max_depth` | `int32` | `64` | Maximum suffix-tree depth for suffix speculative decoding. |
 | `speculative_suffix_max_spec_factor` | `double` | `1.0` | Maximum suffix speculation token factor relative to match length. |
 | `speculative_suffix_max_spec_offset` | `double` | `0.0` | Maximum additive token offset for suffix speculation. |

docs/src/content/docs/zh/cli_reference.md

@@ -183,7 +183,7 @@ xLLM 使用 gflags 管理服务启动参数。`--model <PATH>` 是唯一必填
 | `draft_model` | `string` | `""` | draft 模型路径；MTP 使用方式详见 [MTP](/zh/features/mtp/)。 |
 | `draft_devices` | `string` | `""` | draft 模型使用的设备，例如 `npu:0`、`npu:0,npu:1`；未指定时，启用 speculative decoding 会使用 target 模型的设备。 |
 | `num_speculative_tokens` | `int32` | `0` | 每轮 speculative decoding 生成的 speculative token 数。 |
-| `speculative_algorithm` | `string` | `"MTP"` | Speculative decoding 算法，支持 `MTP`、`Eagle3`、`Suffix`。 |
+| `speculative_algorithm` | `string` | `"MTP"` | Speculative decoding 算法，支持 `MTP`、`Eagle3`、`Suffix`、`DFlash`。 |
 | `speculative_suffix_cache_max_depth` | `int32` | `64` | Suffix speculative decoding 的后缀树最大深度。 |
 | `speculative_suffix_max_spec_factor` | `double` | `1.0` | Suffix speculation 相对于匹配长度的最大 token 系数。 |
 | `speculative_suffix_max_spec_offset` | `double` | `0.0` | Suffix speculation 的最大 token 加性偏移。 |

tests/core/scheduler/profile/CMakeLists.txt

@@ -6,8 +6,11 @@ cc_test(
   SRCS
     profile_graph_warmup_test.cpp
   DEPS
+    :xllm_server
     :block
     :profile
     :request
     GTest::gtest_main
 )
+target_link_libraries(profile_graph_warmup_test PRIVATE
+  "$<LINK_GROUP:RESCAN,xtensor,xllm_server>")

xllm/core/distributed_runtime/llm_engine.cpp

@@ -190,12 +190,16 @@ 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();
+
+  // A draft engine is fed token ids and detokenized by the target, so it
+  // shares the target vocabulary and loads no tokenizer of its own.
+  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_;
@@ -221,21 +225,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

@@ -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;
-      }
-    }
+    // Draft tokens are detokenized by the target, so the draft engine needs no
+    // tokenizer and no draft/target vocab-compatibility check (not universal;
+    // see llm_engine.cpp).
 
     // 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

@@ -42,6 +42,7 @@ limitations under the License.
 #include "core/framework/config/kv_cache_config.h"
 #include "core/framework/config/parallel_config.h"
 #include "core/framework/config/service_config.h"
+#include "core/framework/config/speculative_config.h"
 #include "core/platform/platform.h"
 #if defined(USE_CUDA) || defined(USE_MLU) || defined(USE_DCU)
 #include "core/platform/numa_utils.h"
@@ -418,6 +419,15 @@ WorkerServer::WorkerServer(int32_t local_worker_idx,
                            bool use_spawn_worker)
     : server_name_("DistributeWorkerServer") {
   server_name_.append(std::to_string(options.server_idx()));
+  // Eagle3/DFlash targets capture aux hidden states and don't support spawned
+  // workers yet.
+  const std::string& speculative_algorithm = options.speculative_algorithm();
+  if (use_spawn_worker && options.enable_speculative_decode() &&
+      SpeculativeConfig::requires_aux_hidden_capture(speculative_algorithm)) {
+    LOG(FATAL) << 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/config/speculative_config.h

@@ -18,6 +18,7 @@ limitations under the License.
 #include <cstdint>
 #include <nlohmann/json_fwd.hpp>
 #include <string>
+#include <string_view>
 
 #include "core/common/macros.h"
 #include "core/framework/config/option_category.h"
@@ -33,6 +34,18 @@ class SpeculativeConfig final {
 
   static SpeculativeConfig& get_instance();
 
+  // Whether a speculative algorithm requires the target model to capture
+  // intermediate-layer aux hidden states to drive the draft (Eagle3 and
+  // DFlash). Centralizes the algorithm-string classification in one place. The
+  // worker consults it to decide whether to populate the target's
+  // layers_to_capture; the model then keys off that list alone, never the
+  // algorithm string. Takes the algorithm explicitly so a spawned worker
+  // process, which reads its own Options rather than this global config, can
+  // classify without an initialized singleton.
+  static bool requires_aux_hidden_capture(std::string_view algorithm) {
+    return algorithm == "Eagle3" || algorithm == "DFlash";
+  }
+
   void from_flags();
   void from_json(const JsonReader& json);
   void append_config_json(nlohmann::ordered_json& config_json) const;

xllm/core/framework/model/causal_lm.h

@@ -156,6 +156,20 @@ class CausalLM : public torch::nn::Module {
     NOT_IMPLEMENTED();
   }
 
+  // DFlash-specific interface. Attention-free scatter-only pass that projects
+  // target hidden into the draft's per-layer KV cache. Not part of forward()
+  // because it has no attention and its shape doesn't match forward's decode
+  // graph; sits outside the executor to avoid a per-flag branch in the eager
+  // gate. Default: NOT_IMPLEMENTED. Overridden by DFlashDraftModel.
+  virtual ModelOutput write_context_kv(const torch::Tensor& target_hidden,
+                                       const torch::Tensor& positions,
+                                       const torch::Tensor& device_cache_slots,
+                                       std::vector<KVCache>& kv_caches,
+                                       const ModelInputParams& input_params) {
+    NOT_IMPLEMENTED();
+    return {};
+  }
+
   virtual void lazy_load_model(std::unique_ptr<ModelLoader> loader) {
     NOT_IMPLEMENTED();
   }
@@ -238,6 +252,26 @@ class CausalLMImpl : public CausalLM {
     model_->load_model(std::move(loader));
   }
 
+  ModelOutput write_context_kv(const torch::Tensor& target_hidden,
+                               const torch::Tensor& positions,
+                               const torch::Tensor& device_cache_slots,
+                               std::vector<KVCache>& kv_caches,
+                               const ModelInputParams& input_params) override {
+    if constexpr (detail::has_write_context_kv<Model>::value) {
+      return model_->write_context_kv(target_hidden,
+                                      positions,
+                                      device_cache_slots,
+                                      kv_caches,
+                                      input_params);
+    } else {
+      return CausalLM::write_context_kv(target_hidden,
+                                        positions,
+                                        device_cache_slots,
+                                        kv_caches,
+                                        input_params);
+    }
+  }
+
   void lazy_load_model(std::unique_ptr<ModelLoader> loader) override {
     if constexpr (detail::has_lazy_load_model<Model>::value) {
       model_->lazy_load_model(std::move(loader));

xllm/core/framework/model/model_traits.h

@@ -20,10 +20,12 @@ limitations under the License.
 #include <string>
 #include <type_traits>
 #include <utility>
+#include <vector>
 
 namespace xllm {
 struct ModelInputParams;
 struct ModelGraphMetadataState;
+class KVCache;
 
 namespace layer {
 class LmHead;
@@ -218,5 +220,19 @@ struct has_init_or_refresh_rolling_runtime<
         std::declval<const std::string&>()))>> : std::true_type {};
 
 #endif
+
+template <typename T, typename = void>
+struct has_write_context_kv : std::false_type {};
+
+template <typename T>
+struct has_write_context_kv<
+    T,
+    std::void_t<decltype(std::declval<T>()->write_context_kv(
+        std::declval<const torch::Tensor&>(),
+        std::declval<const torch::Tensor&>(),
+        std::declval<const torch::Tensor&>(),
+        std::declval<std::vector<KVCache>&>(),
+        std::declval<const ModelInputParams&>()))>> : std::true_type {};
+
 }  // namespace detail
 }  // namespace xllm

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