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 (NPU only), DSpark. Default is MTP.");
 
 DEFINE_int32(speculative_suffix_cache_max_depth,
              64,

xllm/core/framework/config/speculative_config.h

@@ -29,6 +29,8 @@ class JsonReader;
 
 class SpeculativeConfig final {
  public:
+  inline static constexpr std::string_view kDFlash2Algorithm = "DFlash2";
+
   SpeculativeConfig() = default;
   ~SpeculativeConfig() = default;
 
@@ -44,7 +46,11 @@ 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";
+           is_dflash2_algorithm(algorithm) || algorithm == "DSpark";
+  }
+
+  static constexpr bool is_dflash2_algorithm(std::string_view algorithm) {
+    return algorithm == kDFlash2Algorithm;
   }
 
   static bool is_mtp_algorithm(std::string_view algorithm) {
@@ -60,7 +66,9 @@ 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, kDFlash2Algorithm) ||
+           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

@@ -30,6 +30,12 @@ limitations under the License.
 
 namespace xllm {
 
+inline constexpr std::string_view kDFlash2DraftModelType = "DFlash2DraftModel";
+
+inline constexpr bool is_dflash2_draft_model_type(std::string_view model_type) {
+  return model_type == kDFlash2DraftModelType;
+}
+
 struct ModelArgs {
   // Expose every plain-data field to the generic property reflection layer so
   // the embedded Python model executor can receive the full, already-parsed
@@ -84,6 +90,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/kernels/npu/npu_fused_infer_attention.cpp

@@ -166,7 +166,9 @@ std::tuple<torch::Tensor, torch::Tensor> npu_fused_infer_attention(
     int64_t sparse_mode,
     const std::string& input_layout,
     bool softmax_lse_flag,
-    bool is_causal) {
+    bool is_causal,
+    int64_t pre_tokens_override,
+    int64_t next_tokens_override) {
   check_tensor(query, "query", "npu_fused_infer_attention");
   check_tensor(key, "key", "npu_fused_infer_attention");
   check_tensor(value, "value", "npu_fused_infer_attention");
@@ -215,8 +217,11 @@ std::tuple<torch::Tensor, torch::Tensor> npu_fused_infer_attention(
 
   std::string layout = input_layout;
   char* input_layout_ptr = const_cast<char*>(layout.c_str());
-  int64_t pre_tokens = kSwaIntMax;
-  int64_t next_tokens = is_causal ? 0 : kSwaIntMax;
+  int64_t pre_tokens =
+      pre_tokens_override >= 0 ? pre_tokens_override : kSwaIntMax;
+  int64_t next_tokens = next_tokens_override >= 0
+                            ? next_tokens_override
+                            : (is_causal ? 0 : kSwaIntMax);
   int64_t inner_precise = 0;
   int64_t antiquant_mode = 0;
   int64_t key_antiquant_mode = 0;

xllm/core/kernels/npu/npu_ops_api.h

@@ -72,7 +72,9 @@ std::tuple<torch::Tensor, torch::Tensor> npu_fused_infer_attention(
     int64_t sparse_mode,
     const std::string& input_layout,
     bool softmax_lse_flag = false,
-    bool is_causal = true);
+    bool is_causal = true,
+    int64_t pre_tokens = -1,
+    int64_t next_tokens = -1);
 
 void batch_chunked_paged_prefill(const torch::Tensor& query,
                                  const torch::Tensor& k_cache,

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/attention_metadata.h

@@ -214,6 +214,11 @@ struct AttentionMetadata {
   torch::Tensor paged_attention_tiling_data;
   // Pre-computed attention mask for npu_fused_infer_attention.
   torch::Tensor fia_attn_mask;
+  // Optional FIA band-mode overrides. Negative values retain the default
+  // causal/full-attention behavior selected by AttentionImpl.
+  int64_t fia_sparse_mode = -1;
+  int64_t fia_pre_tokens = -1;
+  int64_t fia_next_tokens = -1;
   // Host vectors for npu_fused_infer_attention (kernel requires host memory).
   std::vector<int64_t> q_cu_seq_lens_host_vec;
   std::vector<int64_t> kv_cu_seq_lens_host_vec;

xllm/core/layers/common/dense_mlp.cpp

@@ -171,6 +171,10 @@ torch::Tensor DenseMLPImpl::forward(const torch::Tensor& hidden_states) {
 }
 
 void DenseMLPImpl::load_state_dict(const StateDict& state_dict) {
+  gate_proj_weight_seen_ =
+      gate_proj_weight_seen_ || state_dict.has("gate_proj.weight");
+  up_proj_weight_seen_ =
+      up_proj_weight_seen_ || state_dict.has("up_proj.weight");
   gate_up_proj_->load_state_dict(state_dict, {"gate_proj.", "up_proj."});
   down_proj_->load_state_dict(state_dict.get_dict_with_prefix("down_proj."));
 }
@@ -180,15 +184,37 @@ void DenseMLPImpl::load_state_dict(const StateDict& state_dict,
                                    const std::string& down_name) {
   if (is_gated_) {
     CHECK_EQ(gate_up_name.size(), 2);
+    gate_proj_weight_seen_ =
+        gate_proj_weight_seen_ || state_dict.has(gate_up_name[0] + "weight");
+    up_proj_weight_seen_ =
+        up_proj_weight_seen_ || state_dict.has(gate_up_name[1] + "weight");
     gate_up_proj_->load_state_dict(state_dict, gate_up_name);
   } else {
     CHECK_EQ(gate_up_name.size(), 1);
+    up_proj_weight_seen_ =
+        up_proj_weight_seen_ || state_dict.has(gate_up_name[0] + "weight");
     gate_up_proj_->load_state_dict(
         state_dict.get_dict_with_prefix(gate_up_name[0]));
   }
   down_proj_->load_state_dict(state_dict.get_dict_with_prefix(down_name));
 }
 
+void DenseMLPImpl::verify_loaded_weights(const std::string& prefix) const {
+  if (!gate_up_proj_->is_weight_loaded()) {
+    if (is_gated_) {
+      CHECK(gate_proj_weight_seen_)
+          << "weight is not loaded for " << prefix + "gate_proj.weight";
+    }
+    CHECK(up_proj_weight_seen_)
+        << "weight is not loaded for " << prefix + "up_proj.weight";
+  }
+  CHECK(gate_up_proj_->is_weight_loaded())
+      << "weight is not loaded for " << prefix
+      << (is_gated_ ? "{gate_proj,up_proj}.weight" : "up_proj.weight");
+  CHECK(down_proj_->is_weight_loaded())
+      << "weight is not loaded for " << prefix + "down_proj.weight";
+}
+
 std::optional<torch::Tensor> DenseMLPImpl::get_fp8_input_scale() const {
   if (gate_up_proj_) {
     return gate_up_proj_->get_input_scale();

xllm/core/layers/common/dense_mlp.h

@@ -52,6 +52,8 @@ class DenseMLPImpl : public torch::nn::Module {
                        const std::vector<std::string>& gate_up_name,
                        const std::string& down_name);
 
+  void verify_loaded_weights(const std::string& prefix) const;
+
   // Get FP8 input scale from gate_up_proj for fused RMSNorm+FP8 quantization
   std::optional<torch::Tensor> get_fp8_input_scale() const;
 
@@ -66,6 +68,10 @@ class DenseMLPImpl : public torch::nn::Module {
   std::string hidden_act_;
   double swiglu_limit_ = 0.0;
   bool apply_fc1_sequence_parallel_ = true;
+  // gate/up are fused at runtime; retain their logical checkpoint presence so
+  // a partial fused load reports the exact missing projection.
+  bool gate_proj_weight_seen_ = false;
+  bool up_proj_weight_seen_ = false;
 };
 TORCH_MODULE(DenseMLP);