feat: add Qwen3.8 DFlash2 decoding support.

xllm/core/framework/config/speculative_config.cpp

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

xllm/core/framework/config/speculative_config.h

@@ -44,7 +44,7 @@ class SpeculativeConfig final {
   // classify without an initialized singleton.
   static bool requires_aux_hidden_capture(std::string_view algorithm) {
     return algorithm == "Eagle3" || algorithm == "DFlash" ||
-           algorithm == "DSpark";
+           algorithm == "DFlash2" || algorithm == "DSpark";
   }
 
   static bool is_mtp_algorithm(std::string_view algorithm) {
@@ -60,7 +60,8 @@ class SpeculativeConfig final {
   // classified separately via is_mtp_algorithm; callers that also accept MTP
   // must OR the two.
   static bool is_block_diffusion_algorithm(std::string_view algorithm) {
-    return iequals(algorithm, "dflash") || iequals(algorithm, "dspark");
+    return iequals(algorithm, "dflash") || iequals(algorithm, "dflash2") ||
+           iequals(algorithm, "dspark");
   }
 
   void from_flags();

xllm/core/framework/model/causal_lm.h

@@ -53,6 +53,14 @@ struct ModelGraphMetadataState {
   virtual ~ModelGraphMetadataState() = default;
 };
 
+struct DFlash2CandidateOutput {
+  // Candidate vocabulary ids [batch, draft_steps, top_k]. Edge logits use
+  // [batch, draft_steps, predecessor_top_k, successor_top_k]; at step zero
+  // every predecessor entry represents the same anchor token.
+  torch::Tensor candidate_ids;
+  torch::Tensor edge_logits;
+};
+
 class CausalLM : public torch::nn::Module {
  public:
   ~CausalLM() override = default;
@@ -183,6 +191,14 @@ class CausalLM : public torch::nn::Module {
     return {};
   }
 
+  virtual DFlash2CandidateOutput dflash2_candidates(
+      const torch::Tensor& hidden_states,
+      const torch::Tensor& unary_logits,
+      const torch::Tensor& anchor_token_ids) {
+    NOT_IMPLEMENTED();
+    return {};
+  }
+
   // DSpark-specific low-rank Markov projection. The draft worker owns the
   // sequential sampling lifecycle; the model owns only the trained weights and
   // bias computation.
@@ -319,6 +335,18 @@ class CausalLMImpl : public CausalLM {
     }
   }
 
+  DFlash2CandidateOutput dflash2_candidates(
+      const torch::Tensor& hidden_states,
+      const torch::Tensor& unary_logits,
+      const torch::Tensor& anchor_token_ids) override {
+    if constexpr (detail::has_dflash2_candidates<Model>::value) {
+      return model_->dflash2_candidates(
+          hidden_states, unary_logits, anchor_token_ids);
+    }
+    return CausalLM::dflash2_candidates(
+        hidden_states, unary_logits, anchor_token_ids);
+  }
+
   torch::Tensor dspark_markov_bias(
       const torch::Tensor& previous_token_ids) override {
     if constexpr (detail::has_dspark_markov_bias<Model>::value) {

xllm/core/framework/model/model_args.h

@@ -84,6 +84,13 @@ struct ModelArgs {
   PROPERTY(bool, enable_confidence_head) = false;
   PROPERTY(bool, confidence_head_with_markov) = false;
 
+  // DFlash2 local-convolution and candidate-selector geometry.
+  PROPERTY(int32_t, dflash2_block_size) = 0;
+  PROPERTY(int32_t, dflash2_conv_group_size) = 0;
+  PROPERTY(int32_t, dflash2_conv_kernel_size) = 0;
+  PROPERTY(int32_t, dflash2_selector_rank) = 0;
+  PROPERTY(int32_t, dflash2_selector_top_k) = 0;
+
   PROPERTY(bool, use_qk_norm) = false;
   PROPERTY(float, rms_norm_eps) = 0.0f;
 

xllm/core/framework/model/model_traits.h

@@ -261,6 +261,17 @@ struct has_write_context_kv<
         std::declval<std::vector<KVCache>&>(),
         std::declval<const ModelInputParams&>()))>> : std::true_type {};
 
+template <typename T, typename = void>
+struct has_dflash2_candidates : std::false_type {};
+
+template <typename T>
+struct has_dflash2_candidates<
+    T,
+    std::void_t<decltype(std::declval<T>()->dflash2_candidates(
+        std::declval<const torch::Tensor&>(),
+        std::declval<const torch::Tensor&>(),
+        std::declval<const torch::Tensor&>()))>> : std::true_type {};
+
 template <typename T, typename = void>
 struct has_dspark_markov_bias : std::false_type {};
 

xllm/core/layers/common/CMakeLists.txt

@@ -31,6 +31,7 @@ cc_library(
     dsa_topk_share_plan.h
     dp_utils.h
     add_matmul.h
+    dflash2_grouped_conv.h
     moe_fused_topk.h
   SRCS
     oxygen_vision_attention.cpp
@@ -57,6 +58,7 @@ cc_library(
     dsa_metadata_builder.cpp
     dp_utils.cpp
     add_matmul.cpp
+    dflash2_grouped_conv.cpp
     moe_fused_topk.cpp
   DEPS
     "-Wl,--whole-archive"

xllm/core/layers/common/dflash2_grouped_conv.cpp

@@ -0,0 +1,138 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "core/layers/common/dflash2_grouped_conv.h"
+
+#include <glog/logging.h>
+
+#include "core/framework/state_dict/utils.h"
+
+namespace xllm::layer {
+
+torch::Tensor dflash2_grouped_conv(const torch::Tensor& hidden_states,
+                                   const torch::Tensor& delta,
+                                   const torch::Tensor& base,
+                                   int32_t block_size,
+                                   int32_t num_groups,
+                                   int32_t group_size,
+                                   int32_t taps) {
+  CHECK_EQ(hidden_states.dim(), 2);
+  CHECK_EQ(delta.dim(), 3);
+  CHECK_EQ(base.dim(), 2);
+  CHECK_GT(block_size, 0);
+  CHECK_EQ(hidden_states.size(1),
+           static_cast<int64_t>(num_groups) * group_size);
+  CHECK_EQ(delta.size(0), hidden_states.size(0));
+  CHECK_EQ(delta.size(1), taps);
+  CHECK_EQ(delta.size(2), num_groups);
+  CHECK_EQ(base.size(0), taps);
+  CHECK_EQ(base.size(1), hidden_states.size(1));
+
+  const int64_t num_tokens = hidden_states.size(0);
+  torch::Tensor blocks =
+      hidden_states.view({num_tokens, num_groups, group_size});
+  torch::Tensor coefficients =
+      base.view({1, taps, num_groups, group_size}) + delta.unsqueeze(-1);
+  torch::Tensor output = coefficients.select(/*dim=*/1, /*index=*/0) * blocks;
+  torch::Tensor positions = torch::arange(num_tokens,
+                                          torch::TensorOptions()
+                                              .dtype(torch::kLong)
+                                              .device(hidden_states.device())) %
+                            block_size;
+
+  for (int32_t tap = 1; tap < taps; ++tap) {
+    CHECK_GT(num_tokens, tap)
+        << "DFlash2 convolution token count must exceed its tap offset.";
+    torch::Tensor padding =
+        torch::zeros({tap, num_groups, group_size}, hidden_states.options());
+    torch::Tensor shifted = torch::cat(
+        {padding, blocks.slice(/*dim=*/0, /*start=*/0, num_tokens - tap)},
+        /*dim=*/0);
+    torch::Tensor valid =
+        positions.ge(tap).view({num_tokens, 1, 1}).to(hidden_states.dtype());
+    output.add_(coefficients.select(/*dim=*/1, /*index=*/tap) * shifted *
+                valid);
+  }
+  return output.flatten(/*start_dim=*/1);
+}
+
+DFlash2GroupedConvImpl::DFlash2GroupedConvImpl(
+    int64_t hidden_size,
+    int32_t taps,
+    int32_t group_size,
+    int32_t block_size,
+    const torch::TensorOptions& options)
+    : block_size_(block_size), taps_(taps), group_size_(group_size) {
+  CHECK_GT(hidden_size, 0);
+  CHECK_GT(taps_, 0);
+  CHECK_GT(group_size_, 0);
+  CHECK_GT(block_size_, 0);
+  CHECK_EQ(hidden_size % group_size_, 0)
+      << "DFlash2 conv_group_size must divide hidden_size.";
+  num_groups_ = static_cast<int32_t>(hidden_size / group_size_);
+  base_kernel_ = register_parameter(
+      "base_kernel", torch::empty({2, taps_, hidden_size}, options), false);
+  kernel_projection_ = register_module("kernel_projection",
+                                       AddMatmul(hidden_size,
+                                                 2LL * taps_ * num_groups_,
+                                                 /*with_bias=*/false,
+                                                 options));
+}
+
+std::tuple<torch::Tensor, torch::Tensor> DFlash2GroupedConvImpl::prepare(
+    const torch::Tensor& hidden_states) {
+  torch::Tensor coefficients =
+      kernel_projection_->forward(hidden_states)
+          .view({hidden_states.size(0), 2, taps_, num_groups_});
+  return {convolve(hidden_states,
+                   coefficients.select(/*dim=*/1, /*index=*/0),
+                   /*side=*/0),
+          coefficients.select(/*dim=*/1, /*index=*/1)};
+}
+
+torch::Tensor DFlash2GroupedConvImpl::finish(
+    const torch::Tensor& hidden_states,
+    const torch::Tensor& coefficients) {
+  return convolve(hidden_states, coefficients, /*side=*/1);
+}
+
+void DFlash2GroupedConvImpl::load_state_dict(const StateDict& state_dict) {
+  weight::load_weight(
+      state_dict, "base_kernel", base_kernel_, base_kernel_is_loaded_);
+  kernel_projection_->load_state_dict(
+      state_dict.get_dict_with_prefix("kernel_projection."));
+}
+
+void DFlash2GroupedConvImpl::verify_loaded_weights(
+    const std::string& prefix) const {
+  CHECK(base_kernel_is_loaded_)
+      << "weight is not loaded for " << prefix + "base_kernel";
+  kernel_projection_->verify_loaded_weights(prefix + "kernel_projection.");
+}
+
+torch::Tensor DFlash2GroupedConvImpl::convolve(
+    const torch::Tensor& hidden_states,
+    const torch::Tensor& delta,
+    int32_t side) const {
+  return dflash2_grouped_conv(hidden_states,
+                              delta,
+                              base_kernel_.select(/*dim=*/0, side),
+                              block_size_,
+                              num_groups_,
+                              group_size_,
+                              taps_);
+}
+
+}  // namespace xllm::layer

xllm/core/layers/common/dflash2_grouped_conv.h

@@ -0,0 +1,68 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#pragma once
+
+#include <torch/torch.h>
+
+#include <cstdint>
+#include <tuple>
+
+#include "core/framework/state_dict/state_dict.h"
+#include "core/layers/common/add_matmul.h"
+
+namespace xllm::layer {
+
+torch::Tensor dflash2_grouped_conv(const torch::Tensor& hidden_states,
+                                   const torch::Tensor& delta,
+                                   const torch::Tensor& base,
+                                   int32_t block_size,
+                                   int32_t num_groups,
+                                   int32_t group_size,
+                                   int32_t taps);
+
+class DFlash2GroupedConvImpl final : public torch::nn::Module {
+ public:
+  DFlash2GroupedConvImpl(int64_t hidden_size,
+                         int32_t taps,
+                         int32_t group_size,
+                         int32_t block_size,
+                         const torch::TensorOptions& options);
+
+  std::tuple<torch::Tensor, torch::Tensor> prepare(
+      const torch::Tensor& hidden_states);
+
+  torch::Tensor finish(const torch::Tensor& hidden_states,
+                       const torch::Tensor& coefficients);
+
+  void load_state_dict(const StateDict& state_dict);
+  void verify_loaded_weights(const std::string& prefix) const;
+
+ private:
+  torch::Tensor convolve(const torch::Tensor& hidden_states,
+                         const torch::Tensor& delta,
+                         int32_t side) const;
+
+  AddMatmul kernel_projection_{nullptr};
+  torch::Tensor base_kernel_;
+  bool base_kernel_is_loaded_ = false;
+  int32_t block_size_ = 0;
+  int32_t taps_ = 0;
+  int32_t group_size_ = 0;
+  int32_t num_groups_ = 0;
+};
+TORCH_MODULE(DFlash2GroupedConv);
+
+}  // namespace xllm::layer

xllm/core/layers/common/qwen2_attention.cpp

@@ -25,7 +25,12 @@ limitations under the License.
 namespace {
 inline bool is_qwen3_model(const std::string& model_type) {
   static const std::unordered_set<std::string> qwen3_type_set = {
-      "qwen3", "qwen3_vl", "qwen3_moe", "qwen3_vl_moe", "oxygenvlm"};
+      "qwen3",
+      "qwen3_vl",
+      "qwen3_moe",
+      "qwen3_vl_moe",
+      "oxygenvlm",
+      "DFlash2DraftModel"};
   return qwen3_type_set.contains(model_type);
 }
 

xllm/core/layers/npu/CMakeLists.txt

@@ -173,3 +173,15 @@ cc_test(
     glog::glog
     GTest::gtest_main
 )
+
+cc_test(
+  NAME
+    dflash2_grouped_conv_test
+  SRCS
+    dflash2_grouped_conv_tests.cpp
+  DEPS
+    :common_layers
+    torch
+    glog::glog
+    GTest::gtest_main
+)