diff --git a/xllm/core/framework/config/speculative_config.cpp b/xllm/core/framework/config/speculative_config.cpp index 7cce59c83c..d9f1d7009a 100644 --- a/xllm/core/framework/config/speculative_config.cpp +++ b/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, diff --git a/xllm/core/framework/config/speculative_config.h b/xllm/core/framework/config/speculative_config.h index bb9370c323..ce3da576b0 100644 --- a/xllm/core/framework/config/speculative_config.h +++ b/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(); diff --git a/xllm/core/framework/model/causal_lm.h b/xllm/core/framework/model/causal_lm.h index 7aeace1c19..201caf07b0 100644 --- a/xllm/core/framework/model/causal_lm.h +++ b/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::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::value) { diff --git a/xllm/core/framework/model/model_args.h b/xllm/core/framework/model/model_args.h index c23983bc5d..716c3e604a 100644 --- a/xllm/core/framework/model/model_args.h +++ b/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; diff --git a/xllm/core/framework/model/model_traits.h b/xllm/core/framework/model/model_traits.h index 45c224bf85..49d9214833 100644 --- a/xllm/core/framework/model/model_traits.h +++ b/xllm/core/framework/model/model_traits.h @@ -261,6 +261,17 @@ struct has_write_context_kv< std::declval&>(), std::declval()))>> : std::true_type {}; +template +struct has_dflash2_candidates : std::false_type {}; + +template +struct has_dflash2_candidates< + T, + std::void_t()->dflash2_candidates( + std::declval(), + std::declval(), + std::declval()))>> : std::true_type {}; + template struct has_dspark_markov_bias : std::false_type {}; diff --git a/xllm/core/kernels/npu/npu_fused_infer_attention.cpp b/xllm/core/kernels/npu/npu_fused_infer_attention.cpp index 097c029721..6e23964a65 100644 --- a/xllm/core/kernels/npu/npu_fused_infer_attention.cpp +++ b/xllm/core/kernels/npu/npu_fused_infer_attention.cpp @@ -166,7 +166,9 @@ std::tuple 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 npu_fused_infer_attention( std::string layout = input_layout; char* input_layout_ptr = const_cast(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; diff --git a/xllm/core/kernels/npu/npu_ops_api.h b/xllm/core/kernels/npu/npu_ops_api.h index e757e41f3a..ad9719f513 100644 --- a/xllm/core/kernels/npu/npu_ops_api.h +++ b/xllm/core/kernels/npu/npu_ops_api.h @@ -72,7 +72,9 @@ std::tuple 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, diff --git a/xllm/core/layers/common/CMakeLists.txt b/xllm/core/layers/common/CMakeLists.txt index 5b978c9c53..8ad36d1286 100755 --- a/xllm/core/layers/common/CMakeLists.txt +++ b/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" diff --git a/xllm/core/layers/common/attention_metadata.h b/xllm/core/layers/common/attention_metadata.h index 2b02963dcb..d98e8cedc3 100644 --- a/xllm/core/layers/common/attention_metadata.h +++ b/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 q_cu_seq_lens_host_vec; std::vector kv_cu_seq_lens_host_vec; diff --git a/xllm/core/layers/common/dense_mlp.cpp b/xllm/core/layers/common/dense_mlp.cpp index c3b3137240..ba1ae4c007 100644 --- a/xllm/core/layers/common/dense_mlp.cpp +++ b/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 DenseMLPImpl::get_fp8_input_scale() const { if (gate_up_proj_) { return gate_up_proj_->get_input_scale(); diff --git a/xllm/core/layers/common/dense_mlp.h b/xllm/core/layers/common/dense_mlp.h index 0f9324c25b..0269e054d4 100644 --- a/xllm/core/layers/common/dense_mlp.h +++ b/xllm/core/layers/common/dense_mlp.h @@ -52,6 +52,8 @@ class DenseMLPImpl : public torch::nn::Module { const std::vector& 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 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); diff --git a/xllm/core/layers/common/dflash2_grouped_conv.cpp b/xllm/core/layers/common/dflash2_grouped_conv.cpp new file mode 100644 index 0000000000..29daac79c9 --- /dev/null +++ b/xllm/core/layers/common/dflash2_grouped_conv.cpp @@ -0,0 +1,143 @@ +/* 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/xLLM-AI/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 + +#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(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); + // Rows must be laid out as contiguous, fixed-width blocks per sequence. + // The modulo position below is only valid when every sequence begins on a + // block_size boundary. + CHECK_EQ(num_tokens % block_size, 0) + << "DFlash2 convolution rows must contain complete sequence blocks."; + 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(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 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 diff --git a/xllm/core/layers/common/dflash2_grouped_conv.h b/xllm/core/layers/common/dflash2_grouped_conv.h new file mode 100644 index 0000000000..a430bda8f9 --- /dev/null +++ b/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/xLLM-AI/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 + +#include +#include + +#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 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 diff --git a/xllm/core/layers/common/qwen2_attention.cpp b/xllm/core/layers/common/qwen2_attention.cpp index 6ae8a3596d..6d75c6bd3e 100644 --- a/xllm/core/layers/common/qwen2_attention.cpp +++ b/xllm/core/layers/common/qwen2_attention.cpp @@ -26,7 +26,8 @@ namespace { inline bool is_qwen3_model(const std::string& model_type) { static const std::unordered_set qwen3_type_set = { "qwen3", "qwen3_vl", "qwen3_moe", "qwen3_vl_moe", "oxygenvlm"}; - return qwen3_type_set.contains(model_type); + return qwen3_type_set.contains(model_type) || + xllm::is_dflash2_draft_model_type(model_type); } #if defined(USE_CUDA) || defined(USE_DCU) @@ -193,6 +194,9 @@ torch::Tensor Qwen2AttentionImpl::forward( } void Qwen2AttentionImpl::load_state_dict(const StateDict& state_dict) { + q_proj_weight_seen_ = q_proj_weight_seen_ || state_dict.has("q_proj.weight"); + k_proj_weight_seen_ = k_proj_weight_seen_ || state_dict.has("k_proj.weight"); + v_proj_weight_seen_ = v_proj_weight_seen_ || state_dict.has("v_proj.weight"); qkv_proj_->load_state_dict(state_dict, {"q_proj.", "k_proj.", "v_proj."}); o_proj_->load_state_dict(state_dict.get_dict_with_prefix("o_proj.")); if (is_qwen3_style_) { @@ -201,6 +205,26 @@ void Qwen2AttentionImpl::load_state_dict(const StateDict& state_dict) { } } +void Qwen2AttentionImpl::verify_loaded_weights( + const std::string& prefix) const { + if (!qkv_proj_->is_weight_loaded()) { + CHECK(q_proj_weight_seen_) + << "weight is not loaded for " << prefix + "q_proj.weight"; + CHECK(k_proj_weight_seen_) + << "weight is not loaded for " << prefix + "k_proj.weight"; + CHECK(v_proj_weight_seen_) + << "weight is not loaded for " << prefix + "v_proj.weight"; + } + CHECK(qkv_proj_->is_weight_loaded()) << "weight is not loaded for " << prefix + << "{q_proj,k_proj,v_proj}.weight"; + CHECK(o_proj_->is_weight_loaded()) + << "weight is not loaded for " << prefix + "o_proj.weight"; + if (is_qwen3_style_) { + q_norm_->verify_loaded_weights(prefix + "q_norm."); + k_norm_->verify_loaded_weights(prefix + "k_norm."); + } +} + std::optional Qwen2AttentionImpl::get_fp8_input_scale() const { if (qkv_proj_) { return qkv_proj_->get_input_scale(); diff --git a/xllm/core/layers/common/qwen2_attention.h b/xllm/core/layers/common/qwen2_attention.h index 44744fd918..41cbc7d044 100644 --- a/xllm/core/layers/common/qwen2_attention.h +++ b/xllm/core/layers/common/qwen2_attention.h @@ -42,6 +42,8 @@ class Qwen2AttentionImpl : public torch::nn::Module { void load_state_dict(const StateDict& state_dict); + void verify_loaded_weights(const std::string& prefix) const; + // Get FP8 input scale from qkv_proj for fused RMSNorm+FP8 quantization std::optional get_fp8_input_scale() const; @@ -55,6 +57,12 @@ class Qwen2AttentionImpl : public torch::nn::Module { float scaling_; bool is_qwen3_style_; bool can_use_fused_qk_norm_rope_; + // The fused QKV loader only exposes an aggregate loaded flag. Preserve the + // logical checkpoint names across shards so verification can identify the + // missing projection precisely. + bool q_proj_weight_seen_ = false; + bool k_proj_weight_seen_ = false; + bool v_proj_weight_seen_ = false; QKVParallelLinear qkv_proj_{nullptr}; RowParallelLinear o_proj_{nullptr}; diff --git a/xllm/core/layers/npu/CMakeLists.txt b/xllm/core/layers/npu/CMakeLists.txt index b84aa3cf82..3273184511 100644 --- a/xllm/core/layers/npu/CMakeLists.txt +++ b/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 +) diff --git a/xllm/core/layers/npu/dflash2_grouped_conv_tests.cpp b/xllm/core/layers/npu/dflash2_grouped_conv_tests.cpp new file mode 100644 index 0000000000..b0a8078165 --- /dev/null +++ b/xllm/core/layers/npu/dflash2_grouped_conv_tests.cpp @@ -0,0 +1,61 @@ +/* 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/xLLM-AI/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 +#include + +#include "core/layers/common/dflash2_grouped_conv.h" + +namespace xllm::layer { + +TEST(DFlash2GroupedConvTest, ResetsHistoryAtEachBlockBoundary) { + torch::Tensor hidden = torch::tensor({{1.0f, 2.0f, 3.0f, 4.0f}, + {5.0f, 6.0f, 7.0f, 8.0f}, + {9.0f, 10.0f, 11.0f, 12.0f}, + {13.0f, 14.0f, 15.0f, 16.0f}}); + torch::Tensor delta = torch::zeros({4, 2, 2}); + torch::Tensor base = torch::ones({2, 4}); + + torch::Tensor output = dflash2_grouped_conv(hidden, + delta, + base, + /*block_size=*/2, + /*num_groups=*/2, + /*group_size=*/2, + /*taps=*/2); + torch::Tensor expected = torch::stack( + {hidden[0], hidden[1] + hidden[0], hidden[2], hidden[3] + hidden[2]}); + EXPECT_TRUE(torch::allclose(output, expected)); +} + +TEST(DFlash2GroupedConvTest, BroadcastsDynamicKernelWithinEachGroup) { + torch::Tensor hidden = + torch::tensor({{1.0f, 2.0f, 3.0f, 4.0f}, {5.0f, 6.0f, 7.0f, 8.0f}}); + torch::Tensor delta = torch::tensor({{{1.0f, 2.0f}}, {{3.0f, 4.0f}}}); + torch::Tensor base = torch::zeros({1, 4}); + + torch::Tensor output = dflash2_grouped_conv(hidden, + delta, + base, + /*block_size=*/2, + /*num_groups=*/2, + /*group_size=*/2, + /*taps=*/1); + torch::Tensor expected = + torch::tensor({{1.0f, 2.0f, 6.0f, 8.0f}, {15.0f, 18.0f, 28.0f, 32.0f}}); + EXPECT_TRUE(torch::allclose(output, expected)); +} + +} // namespace xllm::layer diff --git a/xllm/core/layers/npu_torch/attention.cpp b/xllm/core/layers/npu_torch/attention.cpp index f881966661..f866811d59 100644 --- a/xllm/core/layers/npu_torch/attention.cpp +++ b/xllm/core/layers/npu_torch/attention.cpp @@ -106,8 +106,19 @@ void AttentionImpl::prefill_forward(torch::Tensor& query, num_kv_heads_, scale_, /*block_size=*/0, - /*sparse_mode=*/3, - "TND"); + // FIA requires sparse mode 3 when an explicit attention mask is used. + // DFlash2 supplies a non-causal block mask, so the mask and causality + // controls must remain independent. + /*sparse_mode=*/attn_metadata.fia_sparse_mode >= 0 + ? attn_metadata.fia_sparse_mode + : (attn_metadata.fia_attn_mask.defined() + ? 3 + : (attn_metadata.is_causal ? 3 : 0)), + "TND", + /*softmax_lse_flag=*/false, + /*is_causal=*/attn_metadata.is_causal, + /*pre_tokens=*/attn_metadata.fia_pre_tokens, + /*next_tokens=*/attn_metadata.fia_next_tokens); output.copy_(std::get<0>(fia_result).view_as(output)); } else if (attn_metadata.is_chunked_prefill) { torch::Tensor k = k_cache.view({k_cache.size(0), k_cache.size(1), -1}); @@ -129,8 +140,16 @@ void AttentionImpl::prefill_forward(torch::Tensor& query, num_kv_heads_, scale_, /*block_size=*/k_cache.size(1), - /*sparse_mode=*/3, - "TND"); + /*sparse_mode=*/attn_metadata.fia_sparse_mode >= 0 + ? attn_metadata.fia_sparse_mode + : (attn_metadata.fia_attn_mask.defined() + ? 3 + : (attn_metadata.is_causal ? 3 : 0)), + "TND", + /*softmax_lse_flag=*/false, + /*is_causal=*/attn_metadata.is_causal, + /*pre_tokens=*/attn_metadata.fia_pre_tokens, + /*next_tokens=*/attn_metadata.fia_next_tokens); output.copy_(std::get<0>(fia_result).view_as(output)); } } diff --git a/xllm/core/layers/qwen2_decoder_layer.cpp b/xllm/core/layers/qwen2_decoder_layer.cpp index 67b84f5c85..aea4d6316b 100644 --- a/xllm/core/layers/qwen2_decoder_layer.cpp +++ b/xllm/core/layers/qwen2_decoder_layer.cpp @@ -51,6 +51,24 @@ Qwen2DecoderLayerImpl::Qwen2DecoderLayerImpl(const ModelContext& context, parallel_args_.tp_group_, options, mlp_module_prefix)); + + use_dflash2_conv_ = is_dflash2_draft_model_type(model_args.model_type()); + if (use_dflash2_conv_) { + attention_conv_ = register_module( + "attention_conv", + DFlash2GroupedConv(model_args.hidden_size(), + model_args.dflash2_conv_kernel_size(), + model_args.dflash2_conv_group_size(), + model_args.dflash2_block_size(), + options)); + mlp_conv_ = register_module( + "mlp_conv", + DFlash2GroupedConv(model_args.hidden_size(), + model_args.dflash2_conv_kernel_size(), + model_args.dflash2_conv_group_size(), + model_args.dflash2_block_size(), + options)); + } } void Qwen2DecoderLayerImpl::load_state_dict(const StateDict& state_dict) { @@ -60,6 +78,23 @@ void Qwen2DecoderLayerImpl::load_state_dict(const StateDict& state_dict) { post_norm_->load_state_dict( state_dict.get_dict_with_prefix("post_attention_layernorm.")); mlp_->load_state_dict(state_dict.get_dict_with_prefix("mlp.")); + if (use_dflash2_conv_) { + attention_conv_->load_state_dict( + state_dict.get_dict_with_prefix("attention_conv.")); + mlp_conv_->load_state_dict(state_dict.get_dict_with_prefix("mlp_conv.")); + } +} + +void Qwen2DecoderLayerImpl::verify_loaded_weights( + const std::string& prefix) const { + attention_->verify_loaded_weights(prefix + "self_attn."); + mlp_->verify_loaded_weights(prefix + "mlp."); + input_norm_->verify_loaded_weights(prefix + "input_layernorm."); + post_norm_->verify_loaded_weights(prefix + "post_attention_layernorm."); + if (use_dflash2_conv_) { + attention_conv_->verify_loaded_weights(prefix + "attention_conv."); + mlp_conv_->verify_loaded_weights(prefix + "mlp_conv."); + } } std::tuple> @@ -97,14 +132,30 @@ torch::Tensor Qwen2DecoderLayerImpl::forward( // Pre-attention norm std::tie(x, residual) = apply_norm(input_norm_, x, residual, pre_fp8_scale); + torch::Tensor attention_coefficients; + if (use_dflash2_conv_) { + std::tie(x, attention_coefficients) = attention_conv_->prepare(x); + } + // Attention x = attention_->forward(positions, x, attn_metadata, kv_cache); + if (use_dflash2_conv_) { + x = attention_conv_->finish(x, attention_coefficients); + } // Post-attention norm std::tie(x, residual) = apply_norm(post_norm_, x, residual, post_fp8_scale); + torch::Tensor mlp_coefficients; + if (use_dflash2_conv_) { + std::tie(x, mlp_coefficients) = mlp_conv_->prepare(x); + } + // MLP x = mlp_->forward(x); + if (use_dflash2_conv_) { + x = mlp_conv_->finish(x, mlp_coefficients); + } return x; } diff --git a/xllm/core/layers/qwen2_decoder_layer.h b/xllm/core/layers/qwen2_decoder_layer.h index 86c14046aa..e1cda76192 100644 --- a/xllm/core/layers/qwen2_decoder_layer.h +++ b/xllm/core/layers/qwen2_decoder_layer.h @@ -21,6 +21,7 @@ limitations under the License. #include #include "common/dense_mlp.h" +#include "common/dflash2_grouped_conv.h" #include "common/qwen2_attention.h" #include "common/rms_norm.h" #include "framework/kv_cache/kv_cache.h" @@ -40,6 +41,8 @@ class Qwen2DecoderLayerImpl : public torch::nn::Module { void load_state_dict(const StateDict& state_dict); + void verify_loaded_weights(const std::string& prefix) const; + torch::Tensor forward(torch::Tensor& x, std::optional& residual, torch::Tensor& positions, @@ -52,6 +55,9 @@ class Qwen2DecoderLayerImpl : public torch::nn::Module { DenseMLP mlp_{nullptr}; RMSNorm input_norm_{nullptr}; RMSNorm post_norm_{nullptr}; + DFlash2GroupedConv attention_conv_{nullptr}; + DFlash2GroupedConv mlp_conv_{nullptr}; + bool use_dflash2_conv_ = false; ParallelArgs parallel_args_; diff --git a/xllm/core/runtime/CMakeLists.txt b/xllm/core/runtime/CMakeLists.txt index ee70ba1905..98620caaaa 100644 --- a/xllm/core/runtime/CMakeLists.txt +++ b/xllm/core/runtime/CMakeLists.txt @@ -41,6 +41,7 @@ cc_library( suffix_worker_impl.h eagle3_worker_impl.h dflash_worker_impl.h + dflash2_worker_impl.h dspark_worker_impl.h forward_shared_memory_manager.h worker_rendezvous.h @@ -77,6 +78,7 @@ cc_library( suffix_worker_impl.cpp eagle3_worker_impl.cpp dflash_worker_impl.cpp + dflash2_worker_impl.cpp dspark_worker_impl.cpp forward_shared_memory_manager.cpp DEPS diff --git a/xllm/core/runtime/dflash2_worker_impl.cpp b/xllm/core/runtime/dflash2_worker_impl.cpp new file mode 100644 index 0000000000..c6956e0540 --- /dev/null +++ b/xllm/core/runtime/dflash2_worker_impl.cpp @@ -0,0 +1,233 @@ +/* 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/xLLM-AI/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 "runtime/dflash2_worker_impl.h" + +#include + +#include "common/metrics.h" +#include "core/framework/parallel_state/process_group.h" +#include "core/framework/sampling/sampler.h" +#include "util/timer.h" + +namespace xllm { + +namespace { + +void clear_target_linear_state_metadata(ModelInputParams& input_params) { + // Qwen3.8 targets own recurrent GDN state, while the official Qwen3 + // DFlash2 draft is a pure full-attention model. prepare_query_inputs starts + // from a copy of the target input, so do not let target-only slot ids make + // the draft attention metadata builder classify these rows as recurrent. + // The target validation input is prepared separately and keeps this state. + input_params.embedding.linear_state_ids.clear(); + input_params.embedding.linear_state_indices = torch::Tensor(); + input_params.linear_state_cache_ops.clear(); + input_params.linear_state_validity_mask.clear(); +} + +} // namespace + +DFlash2WorkerImpl::DFlash2WorkerImpl(const ParallelArgs& parallel_args, + const torch::Device& device, + const runtime::Options& options) + : DFlashWorkerImpl(parallel_args, device, options), + sampling_process_group_(parallel_args.tp_group_ != nullptr + ? parallel_args.tp_group_ + : parallel_args.process_group_) {} + +DFlashWorkerImpl::DraftBlock DFlash2WorkerImpl::run_decode_draft( + const ForwardInput& input, + ForwardInput& validate_input) { + Timer timer; + ForwardInput query_input; + prepare_query_inputs(input, query_input); + clear_target_linear_state_metadata(query_input.input_params); + + const int32_t batch_size = input.input_params.meta.num_sequences; + const int32_t num_speculative_tokens = options_.num_speculative_tokens(); + CHECK_GT(batch_size, 0); + CHECK_GT(num_speculative_tokens, 0); + CHECK(input.token_ids_host.defined()); + CHECK_GE(input.token_ids_host.numel(), batch_size); + torch::Tensor anchor_token_ids = + input.token_ids_host.slice(/*dim=*/0, /*start=*/0, /*end=*/batch_size) + .to(draft_impl_->device(), torch::kLong); + + query_input.skip_sampling_for_logits_only = true; + query_input.return_selected_hidden = true; + ForwardInput processed_input; + draft_impl_->prepare_work_before_execute_on_stream( + query_input, + processed_input, + *prepare_stream_, + /*record_ready_event=*/prepare_stream_.get() != compute_stream_.get()); + draft_impl_->set_hierarchy_layer_synchronizer(processed_input.input_params); + std::optional draft_output = + draft_impl_->execute_no_sync_on_stream(processed_input, + *compute_stream_, + /*record_ready_event=*/false); + CHECK(draft_output.has_value()); + CHECK(draft_output->logits.defined()); + CHECK(draft_output->selected_hidden.defined()) + << "DFlash2 requires selected pre-lm-head hidden states."; + prepare_validate_inputs(input, validate_input); + + const int64_t num_rows = draft_output->logits.size(0); + CHECK_EQ(num_rows, static_cast(batch_size) * num_speculative_tokens); + torch::Tensor unary_logits = draft_output->logits.view( + {batch_size, num_speculative_tokens, draft_output->logits.size(-1)}); + torch::Tensor hidden_states = draft_output->selected_hidden.view( + {batch_size, + num_speculative_tokens, + draft_output->selected_hidden.size(-1)}); + + BlockSampleOutput sampled; + { + c10::StreamGuard stream_guard = compute_stream_->set_stream_guard(); + DFlash2CandidateOutput candidates = draft_impl_->dflash2_candidates( + hidden_states, unary_logits, anchor_token_ids); + SamplingParameters sampling_params = input.sampling_params.to( + unary_logits.device(), unary_logits.scalar_type()); + sampled = + sample_path(candidates, sampling_params, unary_logits.size(/*dim=*/-1)); + } + + DraftBlock draft_block; + draft_block.token_ids = std::move(sampled.token_ids); + draft_block.probs = std::move(sampled.selected_probs); + draft_block.dense_probs = std::move(sampled.dense_probs); + draft_block.retained_inputs = take_retained_inputs(*draft_output); + COUNTER_ADD(speculative_execution_latency_seconds_draft, + timer.elapsed_seconds()); + return draft_block; +} + +DFlash2WorkerImpl::BlockSampleOutput DFlash2WorkerImpl::sample_path( + const DFlash2CandidateOutput& candidates, + const SamplingParameters& sampling_params, + int64_t vocab_size) const { + CHECK_EQ(candidates.candidate_ids.dim(), 3); + CHECK_EQ(candidates.edge_logits.dim(), 4); + const int64_t batch_size = candidates.candidate_ids.size(0); + const int64_t num_steps = candidates.candidate_ids.size(1); + const int64_t top_k = candidates.candidate_ids.size(2); + CHECK_EQ(candidates.edge_logits.sizes(), + torch::IntArrayRef({batch_size, num_steps, top_k, top_k})); + + SamplingParameters step_params = sampling_params; + const torch::TensorOptions index_options = + torch::TensorOptions() + .dtype(torch::kInt) + .device(candidates.edge_logits.device()); + step_params.selected_token_idxes = torch::empty({0}, index_options); + step_params.sample_idxes = torch::empty({0}, index_options); + step_params.return_probs = !step_params.all_greedy_sample; + step_params.logprobs = false; + step_params.max_top_logprobs = 0; + step_params.use_beam_search = false; + // DFlash2 samples selector edges from softmax(edge_logits / temperature). + // Target-side truncation, penalties, and grammar constraints are applied by + // verification and must not be applied a second time to the selector's + // top-k candidate distribution. + step_params.top_p = torch::Tensor(); + step_params.top_k = torch::Tensor(); + step_params.frequency_penalties = torch::Tensor(); + step_params.presence_penalties = torch::Tensor(); + step_params.repetition_penalties = torch::Tensor(); + step_params.filter_mask = torch::Tensor(); + step_params.filter_bitmask = torch::Tensor(); + + torch::Tensor token_ids = + torch::empty({batch_size, num_steps}, candidates.candidate_ids.options()); + torch::Tensor selected_probs = + torch::empty({batch_size, num_steps}, + torch::TensorOptions() + .dtype(torch::kFloat32) + .device(candidates.edge_logits.device())); + torch::Tensor candidate_probs = + torch::empty({batch_size, num_steps, top_k}, selected_probs.options()); + torch::Tensor previous_indices = + torch::zeros({batch_size}, candidates.candidate_ids.options()); + Sampler sampler; + + using ISlice = torch::indexing::Slice; + for (int64_t step = 0; step < num_steps; ++step) { + torch::Tensor edge = + candidates.edge_logits.select(/*dim=*/1, /*index=*/step); + torch::Tensor gather_indices = previous_indices.view({batch_size, 1, 1}) + .expand({batch_size, 1, top_k}); + torch::Tensor step_logits = + edge.gather(/*dim=*/1, gather_indices).squeeze(/*dim=*/1); + SampleOutput output = sampler.forward(step_logits, step_params); + torch::Tensor sampled_indices = output.next_tokens.to(torch::kLong); + synchronize_sampled_indices(sampled_indices, step_params); + + torch::Tensor step_candidates = + candidates.candidate_ids.select(/*dim=*/1, /*index=*/step); + torch::Tensor sampled_tokens = + step_candidates.gather(/*dim=*/1, sampled_indices.view({-1, 1})) + .view({-1}); + torch::Tensor step_probs; + if (step_params.all_greedy_sample) { + step_probs = torch::zeros({batch_size, top_k}, selected_probs.options()); + step_probs.scatter_( + /*dim=*/1, + sampled_indices.view({-1, 1}), + torch::ones({batch_size, 1}, selected_probs.options())); + } else { + CHECK_EQ(output.probs.sizes(), step_logits.sizes()); + step_probs = output.probs.to(torch::kFloat32); + if (!step_params.all_random_sample) { + torch::Tensor greedy_probs = + torch::zeros({batch_size, top_k}, selected_probs.options()); + greedy_probs.scatter_( + /*dim=*/1, + sampled_indices.view({-1, 1}), + torch::ones({batch_size, 1}, selected_probs.options())); + step_probs = torch::where(step_params.do_sample.view({batch_size, 1}), + step_probs, + greedy_probs); + } + } + torch::Tensor chosen_probs = + step_probs.gather(/*dim=*/1, sampled_indices.view({-1, 1})).view({-1}); + token_ids.index_put_({ISlice(), step}, sampled_tokens); + selected_probs.index_put_({ISlice(), step}, chosen_probs); + candidate_probs.index_put_({ISlice(), step, ISlice()}, step_probs); + previous_indices = sampled_indices; + } + + torch::Tensor dense_probs = torch::zeros({batch_size, num_steps, vocab_size}, + selected_probs.options()); + dense_probs.scatter_( + /*dim=*/-1, candidates.candidate_ids, candidate_probs); + return {.token_ids = std::move(token_ids), + .selected_probs = std::move(selected_probs), + .dense_probs = std::move(dense_probs)}; +} + +void DFlash2WorkerImpl::synchronize_sampled_indices( + torch::Tensor& sampled_indices, + const SamplingParameters& sampling_params) const { + if (sampling_params.all_greedy_sample || sampling_process_group_ == nullptr || + sampling_process_group_->world_size() <= 1) { + return; + } + sampled_indices = sampled_indices.contiguous(); + sampling_process_group_->broadcast(sampled_indices, /*root_rank=*/0); +} + +} // namespace xllm diff --git a/xllm/core/runtime/dflash2_worker_impl.h b/xllm/core/runtime/dflash2_worker_impl.h new file mode 100644 index 0000000000..f4eb591c63 --- /dev/null +++ b/xllm/core/runtime/dflash2_worker_impl.h @@ -0,0 +1,54 @@ +/* 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/xLLM-AI/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 "runtime/dflash_worker_impl.h" + +namespace xllm { + +class ProcessGroup; + +class DFlash2WorkerImpl final : public DFlashWorkerImpl { + public: + DFlash2WorkerImpl(const ParallelArgs& parallel_args, + const torch::Device& device, + const runtime::Options& options); + + ~DFlash2WorkerImpl() override = default; + + protected: + DraftBlock run_decode_draft(const ForwardInput& input, + ForwardInput& validate_input) override; + + private: + struct BlockSampleOutput { + torch::Tensor token_ids; + torch::Tensor selected_probs; + torch::Tensor dense_probs; + }; + + BlockSampleOutput sample_path(const DFlash2CandidateOutput& candidates, + const SamplingParameters& sampling_params, + int64_t vocab_size) const; + + void synchronize_sampled_indices( + torch::Tensor& sampled_indices, + const SamplingParameters& sampling_params) const; + + ProcessGroup* sampling_process_group_ = nullptr; +}; + +} // namespace xllm diff --git a/xllm/core/runtime/dflash_worker_impl.cpp b/xllm/core/runtime/dflash_worker_impl.cpp index 320ddbb474..b77d972d2a 100644 --- a/xllm/core/runtime/dflash_worker_impl.cpp +++ b/xllm/core/runtime/dflash_worker_impl.cpp @@ -15,16 +15,18 @@ limitations under the License. #include "runtime/dflash_worker_impl.h" +#include #include #include -#include #include #include +#include #include #include #include "common/metrics.h" +#include "core/framework/config/execution_config.h" #include "core/framework/config/kernel_config.h" #include "core/framework/config/scheduler_config.h" #include "core/framework/config/speculative_config.h" @@ -37,6 +39,7 @@ limitations under the License. #include "framework/kv_cache_transfer/mooncake_kv_cache_transfer.h" #endif #if defined(USE_NPU) +#include "core/layers/common/expanded_decode_metadata_builder.h" #include "core/layers/npu_torch/deepseek_sparse_attention.h" #include "framework/kv_cache_transfer/kv_transfer_completion.h" #endif @@ -91,6 +94,22 @@ runtime::Options draft_options(const runtime::Options& options) { return opts; } +KVCacheShape dflash2_draft_kv_cache_shape(const KVCacheShape& target_shape, + const ModelArgs& draft_model_args, + const ParallelArgs& parallel_args, + int64_t block_size) { + CHECK_GT(parallel_args.dp_size(), 0); + CHECK_GT(parallel_args.cp_size(), 0); + const int64_t dp_local_tp_size = + parallel_args.world_size() / + (parallel_args.dp_size() * parallel_args.cp_size()); + CHECK_GT(dp_local_tp_size, 0); + KVCacheCapacity draft_capacity; + draft_capacity.n_blocks(target_shape.key_cache_shape()[0]) + .block_size(block_size); + return KVCacheShape(draft_capacity, draft_model_args, dp_local_tp_size); +} + void expand_block_parallel_sequence_rows(ModelInputParams& input_params, int32_t query_width) { input_params.meta.num_sequences *= query_width; @@ -146,6 +165,59 @@ void wait_metadata_ready_event(const ForwardInput& input, Stream& stream) { << "failed to wait DFlash metadata ready event"; } +#if defined(USE_NPU) +void build_dflash_expanded_spec_verify_graph_input( + ModelInputParams& input_params, + const torch::Device& device, + int32_t block_size) { + if (!::xllm::ExecutionConfig::get_instance().enable_graph() || + !input_params.is_spec_verify || + !input_params.meta.batch_forward_type.is_chunked_prefill()) { + return; + } + const auto& q_seq_lens = input_params.attention.host.q_seq_lens; + const auto& kv_seq_lens = input_params.attention.host.kv_seq_lens; + CHECK(!q_seq_lens.empty()); + CHECK_EQ(q_seq_lens.size(), kv_seq_lens.size()); + CHECK(input_params.attention.device.block_tables.defined()); + CHECK_GE(input_params.attention.device.block_tables.size(0), + static_cast(q_seq_lens.size())); + + std::vector expanded_kv_seq_lens = + layer::ExpandedDecodeMetadataBuilder::build_tokenwise_kv_seq_lens( + q_seq_lens, kv_seq_lens); + torch::Tensor expanded_kv_seq_lens_device = + torch::tensor(expanded_kv_seq_lens, + torch::TensorOptions() + .dtype(torch::kInt) + .device(torch::kCPU) + .pinned_memory(true)) + .to(device, /*non_blocking=*/true); + + std::vector expanded_block_rows; + expanded_block_rows.reserve(expanded_kv_seq_lens.size()); + for (int64_t seq_idx = 0; seq_idx < static_cast(q_seq_lens.size()); + ++seq_idx) { + for (int32_t token_idx = 0; + token_idx < q_seq_lens[static_cast(seq_idx)]; + ++token_idx) { + expanded_block_rows.emplace_back( + input_params.attention.device.block_tables.select(/*dim=*/0, + seq_idx)); + } + } + CHECK(!expanded_block_rows.empty()); + torch::Tensor expanded_block_tables = + torch::stack(expanded_block_rows, /*dim=*/0).contiguous(); + layer::ExpandedDecodeMetadataBuilder::populate_expanded_layout( + input_params, + expanded_kv_seq_lens_device, + expanded_block_tables, + std::move(expanded_kv_seq_lens), + block_size); +} +#endif + std::optional run_llm_no_sync_impl( LLMWorkerImpl& worker, const ForwardInput& input, @@ -154,9 +226,13 @@ std::optional run_llm_no_sync_impl( ForwardInput* processed_output = nullptr) { ForwardInput processed_input; worker.prepare_work_before_execute_on_stream( - input, processed_input, prepare_stream); - std::optional output = - worker.execute_no_sync_on_stream(processed_input, compute_stream); + input, + processed_input, + prepare_stream, + /*record_ready_event=*/&prepare_stream != &compute_stream); + worker.set_hierarchy_layer_synchronizer(processed_input.input_params); + std::optional output = worker.execute_no_sync_on_stream( + processed_input, compute_stream, /*record_ready_event=*/false); if (processed_output != nullptr) { *processed_output = std::move(processed_input); } @@ -200,7 +276,6 @@ void build_query_rows(const ForwardInput& input, buf.out_kv_seq_lens.reserve(metadata_rows); buf.out_q_seq_lens.reserve(metadata_rows); buf.out_q_cu_seq_lens.reserve(metadata_rows + 1); - buf.out_q_cu_seq_lens.emplace_back(0); selected_idxes.reserve(num_sequences * query_width); @@ -311,6 +386,12 @@ DFlashWorkerImpl::DFlashWorkerImpl(const ParallelArgs& parallel_args, "parallelism (cp_size > 1)."; draft_impl_ = std::make_unique( parallel_args, device, draft_options(options)); + speculative_position_labels_.reserve( + static_cast(options.num_speculative_tokens())); + for (int32_t position = 0; position < options.num_speculative_tokens(); + ++position) { + speculative_position_labels_.emplace_back(std::to_string(position)); + } // Adaptive per-seq validate pruning. DP is supported: the worker gathers // each rank's true validate token count over the DP group before the target @@ -389,10 +470,18 @@ bool DFlashWorkerImpl::init_model(const std::string& model_weights_path, // basis and reduces acceptance to near-random levels. } else { #if defined(USE_NPU) - auto head = impl_->get_npu_lm_head(); - draft_impl_->set_npu_lm_head(head); - auto word_embedding = impl_->get_npu_word_embedding(); - draft_impl_->set_npu_word_embedding(word_embedding); + if (SpeculativeConfig::is_dflash2_algorithm( + options_.speculative_algorithm())) { + auto head = impl_->get_lm_head(); + draft_impl_->set_lm_head(head); + auto word_embedding = impl_->get_word_embedding(); + draft_impl_->set_word_embedding(word_embedding); + } else { + auto head = impl_->get_npu_lm_head(); + draft_impl_->set_npu_lm_head(head); + auto word_embedding = impl_->get_npu_word_embedding(); + draft_impl_->set_npu_word_embedding(word_embedding); + } #else auto head = impl_->get_lm_head(); draft_impl_->set_lm_head(head); @@ -422,6 +511,20 @@ bool DFlashWorkerImpl::init_model(const std::string& model_weights_path, CHECK_LT(mask_token_id_, draft_vocab_size) << "Block-diffusion mask_token_id (" << mask_token_id_ << ") must be < draft vocab_size (" << draft_vocab_size << ")."; + if (SpeculativeConfig::is_dflash2_algorithm( + options_.speculative_algorithm())) { + const int32_t requested_block_size = + options_.num_speculative_tokens() + 1; + CHECK_EQ(requested_block_size, draft_args.dflash2_block_size()) + << "DFlash2 runtime block size must match the checkpoint's trained " + "dflash_config.block_size."; + const ModelArgs& target_args = impl_->context_.get_model_args(); + CHECK(!has_linear_attention_layers(target_args) || + ::xllm::ExecutionConfig::get_instance().enable_graph()) + << "DFlash2 with a hybrid recurrent target requires ACL Graph: the " + "expanded spec-verify replay path preserves the accepted GDN " + "checkpoint, while eager validation is not lossless."; + } // Context hidden comes from the target. const ModelArgs& target_args = impl_->context_.get_model_args(); const int64_t num_target_layers = @@ -459,7 +562,17 @@ bool DFlashWorkerImpl::allocate_kv_cache(const KVCacheShape& kv_cache_shape) { bool draft_allocated = true; const WorkerImpl::Status draft_status = draft_impl_->get_status(); if (draft_status == WorkerImpl::Status::LOADED) { - draft_allocated = draft_impl_->allocate_kv_cache(kv_cache_shape); + if (SpeculativeConfig::is_dflash2_algorithm( + options_.speculative_algorithm())) { + const KVCacheShape draft_shape = + dflash2_draft_kv_cache_shape(kv_cache_shape, + draft_impl_->context_.get_model_args(), + parallel_args_, + options_.block_size()); + draft_allocated = draft_impl_->allocate_kv_cache(draft_shape); + } else { + draft_allocated = draft_impl_->allocate_kv_cache(kv_cache_shape); + } } else { CHECK_EQ(draft_status, WorkerImpl::Status::READY); } @@ -495,8 +608,19 @@ bool DFlashWorkerImpl::allocate_kv_cache_with_transfer( bool draft_allocated = true; const WorkerImpl::Status draft_status = draft_impl_->get_status(); if (draft_status == WorkerImpl::Status::LOADED) { - draft_allocated = draft_impl_->allocate_kv_cache_with_transfer( - kv_cache_transfer_, kv_cache_shape); + if (SpeculativeConfig::is_dflash2_algorithm( + options_.speculative_algorithm())) { + const KVCacheShape draft_shape = + dflash2_draft_kv_cache_shape(kv_cache_shape, + draft_impl_->context_.get_model_args(), + parallel_args_, + options_.block_size()); + draft_allocated = draft_impl_->allocate_kv_cache_with_transfer( + kv_cache_transfer_, draft_shape); + } else { + draft_allocated = draft_impl_->allocate_kv_cache_with_transfer( + kv_cache_transfer_, kv_cache_shape); + } } else { CHECK_EQ(draft_status, WorkerImpl::Status::READY); } @@ -1072,12 +1196,27 @@ SampleOutput DFlashWorkerImpl::validate( .index({ISlice(), ISlice(/*start=*/0, /*end=*/effective_speculative_tokens)}) .contiguous(); - auto [draft_token_ids, draft_probs] = - specBuilder::draftProbs::build_validate_tensors_from_block( - pruned_token_ids, - pruned_probs, - vocab_size, - enable_opt_validate_probs); + torch::Tensor draft_token_ids = pruned_token_ids.to(torch::kLong); + torch::Tensor draft_probs; + if (draft_block.dense_probs.defined()) { + CHECK_EQ(draft_block.dense_probs.dim(), 3) + << "DFlash2 dense proposal probs must be [batch, width, vocab]."; + CHECK_EQ(draft_block.dense_probs.size(2), vocab_size) + << "DFlash2 proposal vocab must match target vocab."; + draft_probs = draft_block.dense_probs + .index({ISlice(), + ISlice(/*start=*/0, + /*end=*/effective_speculative_tokens), + ISlice()}) + .contiguous(); + } else { + std::tie(draft_token_ids, draft_probs) = + specBuilder::draftProbs::build_validate_tensors_from_block( + pruned_token_ids, + pruned_probs, + vocab_size, + enable_opt_validate_probs); + } return validate(sampling_params, draft_token_ids, draft_probs, @@ -1231,9 +1370,46 @@ void DFlashWorkerImpl::prepare_validate_inputs(const ForwardInput& input, c10::StreamGuard stream_guard = prepare_stream_->set_stream_guard(); ForwardInput prepared_input = input; prepared_input.metadata_ready_event.reset(); + const bool use_linear_spec_verify = + impl_ != nullptr && + has_linear_attention_layers(impl_->context_.get_model_args()); + prepared_input.input_params.is_spec_verify = use_linear_spec_verify; SpeculativeWorkerImpl::prepare_validate_inputs(prepared_input, validate_input); validate_input.input_params.embedding.input_embedding = torch::Tensor(); + if (use_linear_spec_verify) { + ModelInputParams& input_params = validate_input.input_params; + std::vector accepted_prefix_lengths( + input.input_params.meta.num_sequences, 1); + if (embedding_cache_ != nullptr && + !input.input_params.embedding.embedding_ids.empty()) { + accepted_prefix_lengths = embedding_cache_->read_accepted_prefix_lengths( + input.input_params.embedding.embedding_ids, + input.input_params.embedding.request_ids); + } + input_params.num_accepted_tokens_host.assign( + accepted_prefix_lengths.begin(), accepted_prefix_lengths.end()); + input_params.num_accepted_tokens = + torch::tensor(accepted_prefix_lengths, + validate_input.token_ids.options().dtype(torch::kInt32)); + if (!input_params.attention.host.q_seq_lens.empty()) { + std::vector q_cu_seq_lens; + q_cu_seq_lens.reserve(input_params.attention.host.q_seq_lens.size() + 1); + q_cu_seq_lens.emplace_back(0); + for (int32_t q_len : input_params.attention.host.q_seq_lens) { + q_cu_seq_lens.emplace_back(q_cu_seq_lens.back() + q_len); + } + input_params.attention.host.q_cu_seq_lens = std::move(q_cu_seq_lens); + } + // The generic builder materializes the device buffer before the hybrid + // cumulative lengths above are canonicalized. Rebind it so GDN sees + // [0, q_len_0, ...] rather than the stale pre-verify host layout. + input_params.attention.rebuild_device_buffer(device_); +#if defined(USE_NPU) + build_dflash_expanded_spec_verify_graph_input( + input_params, device_.unwrap(), options_.block_size()); +#endif + } record_metadata_ready_event(*prepare_stream_, validate_input); } @@ -1463,6 +1639,14 @@ void DFlashWorkerImpl::write_target_context_to_cache( std::vector DFlashWorkerImpl::compute_adaptive_prefix_lengths( const DraftBlock& draft_block, const ForwardInput& input) { + // The current Qwen3.8 GDN spec-verify kernel requires a uniform validation + // width. Keep DFlash adaptive pruning off for hybrid targets until the + // generic varlen builder carries the same recurrent checkpoint contract as + // MTP's specialized path. + if (impl_ != nullptr && + has_linear_attention_layers(impl_->context_.get_model_args())) { + return {}; + } const int32_t num_speculative_tokens = options_.num_speculative_tokens(); if (adaptive_spec_controller_ == nullptr || !adaptive_spec_controller_->enabled()) { @@ -1562,6 +1746,8 @@ void DFlashWorkerImpl::record_validate_metrics( const int64_t* token_data = next_tokens_cpu.const_data_ptr(); int64_t num_draft_tokens = 0; int64_t accepted_count = 0; + c10::SmallVector accepted_per_position( + static_cast(num_speculative_tokens), 0); for (int32_t seq_id = 0; seq_id < batch_size; ++seq_id) { // seq_width = target-side validate width for this seq (anchor + drafts). // Under adaptive per-seq varlen prune it is per_seq_val_tokens[i], else @@ -1580,19 +1766,30 @@ void DFlashWorkerImpl::record_validate_metrics( const int32_t prefix_len = seq_width - 1; num_draft_tokens += prefix_len; - // Count accepted drafts by walking columns [0, prefix_len) — the first - // -1 marks the boundary where the sampler rejected. Padding tail past - // prefix_len is ignored so it never counts as rejection. + // next_tokens column 0 is the token committed by the target; accepted + // draft position i is represented by column i + 1. Walk the complete + // per-seq output (draft prefix plus target bonus/replacement), then remove + // that guaranteed first token. Padding past seq_width is ignored. const int64_t row_offset = static_cast(seq_id) * static_cast(width); - int32_t emitted = 0; - for (int32_t token_idx = 0; token_idx < prefix_len; ++token_idx) { + int32_t emitted_len = 0; + for (int32_t token_idx = 0; token_idx < seq_width; ++token_idx) { if (token_data[row_offset + token_idx] < 0) { break; } - ++emitted; + ++emitted_len; } - accepted_count += emitted; + const int32_t accepted = std::min(prefix_len, std::max(emitted_len - 1, 0)); + accepted_count += accepted; + for (int32_t position = 0; position < accepted; ++position) { + ++accepted_per_position[static_cast(position)]; + } + } + for (int32_t position = 0; position < num_speculative_tokens; ++position) { + MULTI_COUNTER_ADD( + speculative_num_accepted_tokens_per_pos, + speculative_position_labels_[static_cast(position)], + accepted_per_position[static_cast(position)]); } COUNTER_ADD(speculative_num_draft_tokens_total, num_draft_tokens); COUNTER_ADD(speculative_num_accepted_tokens_total, accepted_count); diff --git a/xllm/core/runtime/dflash_worker_impl.h b/xllm/core/runtime/dflash_worker_impl.h index fadfbdc0f3..a2508c2d82 100644 --- a/xllm/core/runtime/dflash_worker_impl.h +++ b/xllm/core/runtime/dflash_worker_impl.h @@ -99,6 +99,11 @@ class DFlashWorkerImpl : public SpeculativeWorkerImpl { struct DraftBlock { torch::Tensor token_ids; torch::Tensor probs; + // Optional exact proposal distribution [batch, num_speculative_tokens, + // vocab]. DFlash2 populates this because its selector samples from a sparse + // top-k distribution; retaining only the chosen probability would make + // rejection recovery inexact. + torch::Tensor dense_probs; // Optional acceptance-probability estimate, [batch, num_speculative_tokens] // fp32 in [0, 1]. Populated by DSpark's ConfidenceHead when available; // consumed by the adaptive-speculative pruning controller. When undefined, @@ -228,6 +233,9 @@ class DFlashWorkerImpl : public SpeculativeWorkerImpl { #endif int32_t mask_token_id_ = -1; int64_t expected_context_hidden_size_ = 0; + // Preformatted labels keep per-position acceptance telemetry allocation-free + // on the decode hot path. + std::vector speculative_position_labels_; dflash_detail::DSparkSasMode draft_sas_mode_ = dflash_detail::DSparkSasMode::NOT_DSPARK; }; diff --git a/xllm/core/runtime/llm_worker_impl.h b/xllm/core/runtime/llm_worker_impl.h index fa1b6d6ffe..7b1c78ef45 100644 --- a/xllm/core/runtime/llm_worker_impl.h +++ b/xllm/core/runtime/llm_worker_impl.h @@ -118,6 +118,14 @@ class LLMWorkerImpl : public WorkerImpl { return model_->has_dspark_confidence_head(); } + DFlash2CandidateOutput dflash2_candidates( + const torch::Tensor& hidden_states, + const torch::Tensor& unary_logits, + const torch::Tensor& anchor_token_ids) { + return model_->dflash2_candidates( + hidden_states, unary_logits, anchor_token_ids); + } + bool share_weights_from(LLMWorkerImpl& source) { return model_->share_weights_from(*source.model_); } diff --git a/xllm/core/runtime/mtp_worker_impl.cpp b/xllm/core/runtime/mtp_worker_impl.cpp index 720e44bc7b..f1b2a4e9d5 100644 --- a/xllm/core/runtime/mtp_worker_impl.cpp +++ b/xllm/core/runtime/mtp_worker_impl.cpp @@ -1352,6 +1352,15 @@ void MTPWorkerImpl::prepare_prefill_inputs(const ForwardInput& input, prefill_input.sampling_params.return_probs = true; clear_ready_events(prefill_input); auto& input_params = prefill_input.input_params; + // Block-diffusion and MTP workers both reuse a hybrid target's ForwardInput, + // but their Qwen draft models are pure full-attention models. Keep this + // cleanup at the MTP draft seam as well: otherwise target-only recurrent + // slot metadata makes MTP prefill independently enter a stateful path for + // which it has neither a validity mask nor a recurrent cache. + input_params.embedding.linear_state_ids.clear(); + input_params.embedding.linear_state_indices = torch::Tensor(); + input_params.linear_state_cache_ops.clear(); + input_params.linear_state_validity_mask.clear(); auto& extra_token_ids = input_params.embedding.extra_token_ids; const torch::Tensor& token_ids = input.token_ids_host; diff --git a/xllm/core/runtime/speculative_worker_impl.cpp b/xllm/core/runtime/speculative_worker_impl.cpp index 178402bfa2..458e2088eb 100644 --- a/xllm/core/runtime/speculative_worker_impl.cpp +++ b/xllm/core/runtime/speculative_worker_impl.cpp @@ -411,6 +411,14 @@ void SpeculativeWorkerImpl::prepare_validate_inputs( const int32_t num_val_tokens = num_speculative_tokens + 1; const int32_t total_num_val_tokens = num_sequences * num_val_tokens; const int32_t block_size = options_.block_size(); + // Hybrid targets (for example Qwen3.8 GDN) mark validation as spec-verify + // before entering this generic builder. They must keep one sequence row + // with an N+1-token query so recurrent state is checkpointed and committed + // by the model's spec-verify kernel instead of being expanded into N+1 + // independent decode rows. + const bool use_chunked_spec_verify = + ::xllm::SpeculativeConfig::get_instance().enable_atb_spec_kernel() || + input.input_params.is_spec_verify; specBuilder::DecodeRowContext row_ctx = specBuilder::make_decode_row_context(input); @@ -421,7 +429,7 @@ void SpeculativeWorkerImpl::prepare_validate_inputs( buf.out_token_ids.reserve(total_num_val_tokens); buf.out_positions.reserve(total_num_val_tokens); buf.out_new_cache_slots.reserve(total_num_val_tokens); - if (!::xllm::SpeculativeConfig::get_instance().enable_atb_spec_kernel()) { + if (!use_chunked_spec_verify) { buf.out_kv_seq_lens.reserve(total_num_val_tokens); buf.out_q_seq_lens.reserve(total_num_val_tokens); buf.out_q_cu_seq_lens.reserve(total_num_val_tokens); @@ -450,16 +458,13 @@ void SpeculativeWorkerImpl::prepare_validate_inputs( row.token_id = -val_idx; } row.position_offset = val_idx; - row.append_kv_len = - !::xllm::SpeculativeConfig::get_instance().enable_atb_spec_kernel(); - row.append_q_len_one = - !::xllm::SpeculativeConfig::get_instance().enable_atb_spec_kernel(); - row.append_block_table = - !::xllm::SpeculativeConfig::get_instance().enable_atb_spec_kernel(); + row.append_kv_len = !use_chunked_spec_verify; + row.append_q_len_one = !use_chunked_spec_verify; + row.append_block_table = !use_chunked_spec_verify; specBuilder::append_decode_row(row_ctx, row, block_size, buf); } - if (::xllm::SpeculativeConfig::get_instance().enable_atb_spec_kernel()) { + if (use_chunked_spec_verify) { const int32_t kv_len_after_validation = kv_len + num_speculative_tokens; specBuilder::update_kv_seq_lens_and_max( atb_kv_seq_lens_vec, kv_len_after_validation, atb_kv_max_seq_len); @@ -479,7 +484,7 @@ void SpeculativeWorkerImpl::prepare_validate_inputs( token_options, position_options); // update the input_params - if (!::xllm::SpeculativeConfig::get_instance().enable_atb_spec_kernel()) { + if (!use_chunked_spec_verify) { input_params.meta.num_sequences = total_num_val_tokens; input_params.meta.q_max_seq_len = 1; input_params.meta.batch_forward_type = BatchForwardType::DECODE; @@ -487,7 +492,7 @@ void SpeculativeWorkerImpl::prepare_validate_inputs( input_params.meta.q_max_seq_len = num_val_tokens; input_params.meta.batch_forward_type = BatchForwardType::CHUNKED_PREFILL; } - if (::xllm::SpeculativeConfig::get_instance().enable_atb_spec_kernel()) { + if (use_chunked_spec_verify) { specBuilder::update_input_params(input_params, buf, num_val_tokens, diff --git a/xllm/core/runtime/worker.cpp b/xllm/core/runtime/worker.cpp index 5e089e32b1..ec79360a63 100644 --- a/xllm/core/runtime/worker.cpp +++ b/xllm/core/runtime/worker.cpp @@ -30,6 +30,7 @@ limitations under the License. #include "framework/kv_cache/kv_cache.h" #include "framework/model/model_input_params.h" #include "framework/state_dict/state_dict.h" +#include "runtime/dflash2_worker_impl.h" #include "runtime/dflash_worker_impl.h" #include "runtime/dit_worker_impl.h" #include "runtime/dspark_worker_impl.h" @@ -56,6 +57,11 @@ Worker::Worker(const ParallelArgs& parallel_args, impl_ = new Eagle3WorkerImpl(parallel_args, device, options); } else if (algorithm == "DFlash") { impl_ = new DFlashWorkerImpl(parallel_args, device, options); + } else if (SpeculativeConfig::is_dflash2_algorithm(algorithm)) { +#if !defined(USE_NPU) + LOG(FATAL) << "DFlash2 speculative decoding is only supported on NPU."; +#endif + impl_ = new DFlash2WorkerImpl(parallel_args, device, options); } else if (algorithm == "DSpark") { impl_ = new DSparkWorkerImpl(parallel_args, device, options); } else if (algorithm == "Suffix") { diff --git a/xllm/core/runtime/worker_impl.cpp b/xllm/core/runtime/worker_impl.cpp index 798c10c032..3055e8462b 100644 --- a/xllm/core/runtime/worker_impl.cpp +++ b/xllm/core/runtime/worker_impl.cpp @@ -1294,7 +1294,15 @@ void WorkerImpl::prepare_work_before_execute_on_stream( #if defined(USE_NPU) || defined(USE_MLU) || defined(USE_CUDA) || \ defined(USE_MUSA) - if (has_linear_attention_layers(context_.get_model_args())) { + // SpeculativeWorkerImpl carries the target model context but delegates + // target/draft execution to its inner workers and therefore owns no KV + // cache itself. In particular, a Qwen3.5/3.8 target advertises linear + // attention while the outer DFlash worker's kv_caches_ is empty. Let the + // inner target worker (which owns the recurrent cache) perform preparation + // and restore; attempting it here would treat target cache operations as + // draft operations and fail discover_num_slots(). + if (!kv_caches_.empty() && + has_linear_attention_layers(context_.get_model_args())) { prepare_input_params_for_linear_attention(input_params); // Under schedule_overlap chunked prefill the previous chunk's forward // runs on compute_stream_ from a worker thread that may not have @@ -1947,7 +1955,9 @@ bool WorkerImpl::init_model(const std::string& model_weights_path, const std::string& speculative_algorithm = options_.speculative_algorithm(); const bool is_block_diffusion = - speculative_algorithm == "DFlash" || speculative_algorithm == "DSpark"; + speculative_algorithm == "DFlash" || + SpeculativeConfig::is_dflash2_algorithm(speculative_algorithm) || + speculative_algorithm == "DSpark"; #if defined(USE_NPU) if (options_.enable_speculative_decode() && @@ -1961,8 +1971,13 @@ bool WorkerImpl::init_model(const std::string& model_weights_path, const bool is_dspark = speculative_algorithm == "DSpark"; const bool is_deepseek_v4_dspark = is_dspark && util::is_deepseek_v4_model_type(args.model_type()); - std::string draft_model_type = - is_dspark ? "DSparkDraftModel" : "DFlashDraftModel"; + std::string draft_model_type = "DFlashDraftModel"; + if (is_dspark) { + draft_model_type = "DSparkDraftModel"; + } else if (SpeculativeConfig::is_dflash2_algorithm( + speculative_algorithm)) { + draft_model_type = std::string(kDFlash2DraftModelType); + } if (is_deepseek_v4_dspark) { draft_model_type = std::string(util::kDeepseekV4DSparkModelType); } diff --git a/xllm/models/llm/llm_model_base.h b/xllm/models/llm/llm_model_base.h index fa1e65e1a6..61bb1643b6 100644 --- a/xllm/models/llm/llm_model_base.h +++ b/xllm/models/llm/llm_model_base.h @@ -177,10 +177,19 @@ class LlmForCausalLMImplBase : public torch::nn::Module { embedding_mode_ = context.get_model_args().embedding_mode(); // register submodules model_ = register_module("model", LlmModelType(context)); - if (!embedding_mode_) { + // DFlash2 shares the target lm-head after loading. Avoid registering a + // draft-local copy: assigning lm_head_ later does not remove an already + // registered module, so the unused vocabulary projection would otherwise + // remain resident on the device. + const bool shares_target_lm_head = + is_dflash2_draft_model_type(context.get_model_args().model_type()); + if (!embedding_mode_ && !shares_target_lm_head) { lm_head_ = register_module("lm_head", layer::LmHead(context)); - } else { + } else if (embedding_mode_) { LOG(INFO) << "Skip registering lm_head in embedding mode."; + } else { + LOG(INFO) << "Skip registering draft-local lm_head for DFlash2; the " + "target lm_head will be shared after model loading."; } } diff --git a/xllm/models/llm/npu/qwen3_dflash2.h b/xllm/models/llm/npu/qwen3_dflash2.h new file mode 100644 index 0000000000..ce540a6de5 --- /dev/null +++ b/xllm/models/llm/npu/qwen3_dflash2.h @@ -0,0 +1,465 @@ +/* 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/xLLM-AI/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 +#include + +#include +#include +#include +#include +#include + +#include "core/kernels/ops_api.h" +#include "core/layers/common/add_matmul.h" +#include "core/layers/common/rms_norm.h" +#include "core/layers/npu/rotary_embedding.h" +#include "framework/model_loader.h" +#include "models/llm/qwen3.h" +#include "models/model_registry.h" + +namespace xllm::npu::model { + +class DFlash2CandidateSelector final { + public: + DFlash2CandidateSelector(const ModelArgs& args, + const torch::TensorOptions& options) + : options_(options), + hidden_size_(args.hidden_size()), + vocab_size_(args.vocab_size()), + rank_(args.dflash2_selector_rank()), + top_k_(args.dflash2_selector_top_k()) { + CHECK_GT(rank_, 0) << "DFlash2 selector_rank must be positive."; + CHECK_GT(top_k_, 0) << "DFlash2 selector_top_k must be positive."; + } + + void load_state_dict(const StateDict& state_dict) { + torch::Tensor hidden_projection = + state_dict.get_tensor("hidden_projection.weight"); + torch::Tensor predecessor = state_dict.get_tensor("predecessor_codebook"); + torch::Tensor successor = state_dict.get_tensor("successor_codebook"); + if (hidden_projection.defined()) { + hidden_projection_ = hidden_projection.to(options_); + } + if (predecessor.defined()) { + predecessor_codebook_ = predecessor.to(options_); + } + if (successor.defined()) { + successor_codebook_ = successor.to(options_); + } + } + + void verify_loaded_weights() const { + CHECK(hidden_projection_.defined()) + << "Missing DFlash2 candidate_selector.hidden_projection.weight."; + CHECK(predecessor_codebook_.defined()) + << "Missing DFlash2 candidate_selector.predecessor_codebook."; + CHECK(successor_codebook_.defined()) + << "Missing DFlash2 candidate_selector.successor_codebook."; + CHECK_EQ(hidden_projection_.sizes(), + torch::IntArrayRef({rank_, hidden_size_})); + CHECK_EQ(predecessor_codebook_.sizes(), + torch::IntArrayRef({vocab_size_, rank_})); + CHECK_EQ(successor_codebook_.sizes(), + torch::IntArrayRef({vocab_size_, rank_})); + } + + DFlash2CandidateOutput forward(const torch::Tensor& hidden_states, + const torch::Tensor& unary_logits, + const torch::Tensor& anchor_token_ids) const { + CHECK_EQ(hidden_states.dim(), 3); + CHECK_EQ(unary_logits.dim(), 3); + CHECK_EQ(hidden_states.size(0), unary_logits.size(0)); + CHECK_EQ(hidden_states.size(1), unary_logits.size(1)); + CHECK_EQ(unary_logits.size(2), vocab_size_); + CHECK_EQ(anchor_token_ids.dim(), 1); + CHECK_EQ(anchor_token_ids.size(0), hidden_states.size(0)); + + auto topk = torch::topk(unary_logits, top_k_, /*dim=*/-1); + torch::Tensor values = std::get<0>(topk).to(torch::kFloat32); + torch::Tensor candidate_ids = std::get<1>(topk).to(torch::kLong); + + namespace F = torch::nn::functional; + torch::Tensor hidden = F::linear(hidden_states, hidden_projection_); + torch::Tensor successors = F::embedding(candidate_ids, successor_codebook_); + torch::Tensor anchor = + anchor_token_ids.view({-1, 1, 1}).expand({-1, 1, top_k_}); + torch::Tensor predecessor_ids = + torch::cat({anchor, + candidate_ids.slice(/*dim=*/1, + /*start=*/0, + candidate_ids.size(1) - 1)}, + /*dim=*/1); + torch::Tensor predecessors = + F::embedding(predecessor_ids, predecessor_codebook_); + torch::Tensor pair_scores = + torch::einsum("blpr,blcr->blpc", + {predecessors * hidden.unsqueeze(/*dim=*/2), successors}); + + DFlash2CandidateOutput output; + output.candidate_ids = candidate_ids; + output.edge_logits = values.unsqueeze(/*dim=*/2) + pair_scores; + return output; + } + + private: + torch::Tensor hidden_projection_; + torch::Tensor predecessor_codebook_; + torch::Tensor successor_codebook_; + torch::TensorOptions options_; + int64_t hidden_size_ = 0; + int64_t vocab_size_ = 0; + int64_t rank_ = 0; + int64_t top_k_ = 0; +}; + +class DFlash2Qwen3ModelImpl final : public ::xllm::QWen3ModelImpl { + public: + explicit DFlash2Qwen3ModelImpl(const ModelContext& context) + : ::xllm::QWen3ModelImpl(context), + selector_(context.get_model_args(), context.get_tensor_options()) { + const ModelArgs& args = context.get_model_args(); + const ParallelArgs& parallel_args = context.get_parallel_args(); + tensor_options_ = context.get_tensor_options(); + head_dim_ = args.head_dim(); + rms_norm_eps_ = args.rms_norm_eps(); + sliding_window_ = args.sliding_window(); + block_size_ = args.dflash2_block_size(); + CHECK_GT(sliding_window_, 0); + CHECK_GT(block_size_, 0); + + const int32_t dp_size = parallel_args.dp_size(); + const int32_t cp_size = parallel_args.cp_size(); + CHECK_GT(dp_size, 0); + CHECK_GT(cp_size, 0); + CHECK_EQ(parallel_args.world_size() % (dp_size * cp_size), 0); + tp_size_ = parallel_args.world_size() / (dp_size * cp_size); + tp_rank_ = parallel_args.rank() % tp_size_; + + fc_ = register_module("fc", + layer::AddMatmul(args.hidden_size() * args.n_layers(), + args.hidden_size(), + /*with_bias=*/false, + tensor_options_)); + hidden_norm_ = register_module( + "hidden_norm", + layer::RMSNorm( + args.hidden_size(), args.rms_norm_eps(), tensor_options_)); + rotary_embedding_ = std::make_shared( + head_dim_, + args.max_position_embeddings(), + layer::rotary::compute_inv_freq( + head_dim_, args.rope_theta(), tensor_options_), + /*interleaved=*/false, + tensor_options_); + } + + void load_state_dict(const StateDict& state_dict) override { + fc_->load_state_dict(state_dict.get_dict_with_prefix("fc.")); + hidden_norm_->load_state_dict( + state_dict.get_dict_with_prefix("hidden_norm.")); + selector_.load_state_dict( + state_dict.get_dict_with_prefix("candidate_selector.")); + load_context_kv_weights(state_dict); + for (int32_t i = 0; i < static_cast(layers_.size()); ++i) { + layers_[i]->load_state_dict( + state_dict.get_dict_with_prefix("layers." + std::to_string(i) + ".")); + } + norm_->load_state_dict(state_dict.get_dict_with_prefix("norm.")); + } + + void verify_loaded_weights() const { + fc_->verify_loaded_weights("fc."); + hidden_norm_->verify_loaded_weights("hidden_norm."); + norm_->verify_loaded_weights("norm."); + selector_.verify_loaded_weights(); + verify_context_kv_weights(); + for (int32_t i = 0; i < static_cast(layers_.size()); ++i) { + layers_[i]->verify_loaded_weights("layers." + std::to_string(i) + "."); + } + } + + void finalize_loaded_weights() { build_fused_context_kv_weights(); } + + DFlash2CandidateOutput candidates( + const torch::Tensor& hidden_states, + const torch::Tensor& unary_logits, + const torch::Tensor& anchor_token_ids) const { + return selector_.forward(hidden_states, unary_logits, anchor_token_ids); + } + + ModelOutput write_context_kv(const torch::Tensor& target_hidden, + const torch::Tensor& positions, + const torch::Tensor& device_cache_slots, + std::vector& kv_caches, + const ModelInputParams& input_params) { + const int64_t num_layers = static_cast(layers_.size()); + CHECK_EQ(static_cast(kv_caches.size()), num_layers); + CHECK_EQ(device_cache_slots.numel(), target_hidden.size(0)); + torch::Tensor projected_hidden = fc_->forward(target_hidden); + projected_hidden = std::get<0>(hidden_norm_->forward(projected_hidden)); + CHECK(fused_kv_weight_.defined()); + + const int64_t num_context = projected_hidden.size(0); + torch::Tensor all_kv = + torch::nn::functional::linear(projected_hidden, fused_kv_weight_) + .view({num_context, num_layers, 2, local_kv_heads_, head_dim_}) + .permute({2, 1, 0, 3, 4}) + .contiguous(); + torch::Tensor all_key = + apply_k_norm(all_kv.select(/*dim=*/0, /*index=*/0), k_norm_weight_); + torch::Tensor all_value = all_kv.select(/*dim=*/0, /*index=*/1); + torch::Tensor flat_key = + all_key.reshape({num_layers * num_context, local_kv_heads_, head_dim_}); + flat_key = apply_rope(flat_key, positions.repeat({num_layers})); + all_key = + flat_key.view({num_layers, num_context, local_kv_heads_, head_dim_}); + + const int32_t device_index = all_key.device().index(); + for (int64_t i = 0; i < num_layers; ++i) { + kernel::ReshapePagedCacheParams params; + params.key = all_key[i]; + params.value = all_value[i]; + params.k_cache = kv_caches[i].get_k_cache(); + params.v_cache = kv_caches[i].get_v_cache(); + params.slot_mapping = device_cache_slots; + CHECK_EQ(params.key.dim(), 3); + CHECK_EQ(params.value->dim(), 3); + CHECK_EQ(params.k_cache.dim(), 4); + CHECK_EQ(params.v_cache->dim(), 4); + CHECK_EQ(params.key.size(0), params.slot_mapping.numel()); + CHECK_EQ(params.key.size(1), params.k_cache.size(2)) + << "DFlash2 context key/cache KV-head mismatch; key=" + << params.key.sizes() << ", cache=" << params.k_cache.sizes(); + CHECK_EQ(params.key.size(2), params.k_cache.size(3)) + << "DFlash2 context key/cache head-dim mismatch; key=" + << params.key.sizes() << ", cache=" << params.k_cache.sizes(); + CHECK_EQ(params.value->sizes(), params.key.sizes()); + CHECK_EQ(params.v_cache->sizes(), params.k_cache.sizes()); + kernel::reshape_paged_cache(params); + if (input_params.parallel.layer_synchronizer != nullptr && + !input_params.parallel.layer_synchronizer->record_event( + i, device_index)) { + return ModelOutput(); + } + } + return ModelOutput(projected_hidden); + } + + protected: + layer::AttentionMetadata get_attention_metadata( + const ModelInputParams& params, + const torch::Tensor& h) override { + layer::AttentionMetadata metadata = + QWen3ModelImpl::get_attention_metadata(params, h); + // DFlash2 jointly denoises the whole query block. Keep the 2048x2048 FIA + // optimized mask built by the base metadata path and use band mode to + // expose the full proposal block within the checkpoint's sliding window. + metadata.is_causal = false; +#if defined(USE_NPU) + metadata.fia_sparse_mode = 4; + metadata.fia_pre_tokens = sliding_window_ - 1; + metadata.fia_next_tokens = block_size_ - 1; +#endif + return metadata; + } + + torch::Tensor gen_append_attn_mask(int32_t q_len, + int32_t kv_len, + int32_t max_kv_len, + torch::Dtype dtype, + torch::Device device) override { + CHECK_GT(q_len, 0); + CHECK_GE(kv_len, q_len); + CHECK_GE(max_kv_len, kv_len); + + const torch::TensorOptions index_options = + torch::TensorOptions().dtype(torch::kLong).device(device); + const torch::Tensor key_positions = + torch::arange(max_kv_len, index_options).view({1, max_kv_len}); + const torch::Tensor query_offsets = + torch::arange(q_len, index_options).view({q_len, 1}); + const int64_t first_query_position = kv_len - q_len; + const torch::Tensor first_visible_key = + query_offsets + first_query_position - (sliding_window_ - 1); + const torch::Tensor last_visible_key = + query_offsets + first_query_position + (block_size_ - 1); + const torch::Tensor masked = key_positions.lt(first_visible_key) | + key_positions.gt(last_visible_key) | + key_positions.ge(kv_len); + + // Match AttentionMask's numeric convention: fp16 uses -inf while other + // dtypes use the FIA-compatible finite sentinel. + const float mask_value = dtype == torch::kFloat16 + ? -std::numeric_limits::infinity() + : -9984.0f; + const torch::TensorOptions mask_options = + torch::TensorOptions().dtype(dtype).device(device); + return torch::zeros({q_len, max_kv_len}, mask_options) + .masked_fill(masked, mask_value); + } + + private: + torch::Tensor apply_k_norm(const torch::Tensor& key, + const torch::Tensor& weight) const { + torch::Tensor key_fp32 = key.to(torch::kFloat32); + torch::Tensor variance = key_fp32.pow(2).mean(/*dim=*/-1, /*keepdim=*/true); + return (key_fp32 * torch::rsqrt(variance + rms_norm_eps_) * weight) + .to(key.scalar_type()); + } + + torch::Tensor apply_rope(const torch::Tensor& key, + const torch::Tensor& positions) const { + CHECK(rotary_embedding_ != nullptr); + return std::get<1>(rotary_embedding_->forward(key, key, positions)); + } + + void load_context_kv_weights(const StateDict& state_dict) { + const int32_t num_layers = static_cast(layers_.size()); + if (per_layer_k_proj_.empty()) { + per_layer_k_proj_.resize(num_layers); + per_layer_v_proj_.resize(num_layers); + per_layer_k_norm_.resize(num_layers); + } + for (int32_t i = 0; i < num_layers; ++i) { + StateDict layer_dict = + state_dict.get_dict_with_prefix("layers." + std::to_string(i) + "."); + torch::Tensor k_proj = layer_dict.get_sharded_tensor( + "self_attn.k_proj.weight", /*dim=*/0, tp_rank_, tp_size_); + torch::Tensor v_proj = layer_dict.get_sharded_tensor( + "self_attn.v_proj.weight", /*dim=*/0, tp_rank_, tp_size_); + torch::Tensor k_norm = layer_dict.get_tensor("self_attn.k_norm.weight"); + if (!k_proj.defined() && !v_proj.defined() && !k_norm.defined()) { + continue; + } + CHECK(k_proj.defined()); + CHECK(v_proj.defined()); + CHECK(k_norm.defined()); + const int64_t local_kv_heads = k_proj.size(0) / head_dim_; + if (local_kv_heads_ == 0) { + local_kv_heads_ = local_kv_heads; + } + CHECK_EQ(local_kv_heads_, local_kv_heads); + per_layer_k_proj_[i] = k_proj.to(tensor_options_); + per_layer_v_proj_[i] = v_proj.to(tensor_options_); + per_layer_k_norm_[i] = k_norm.to(tensor_options_).to(torch::kFloat32); + } + } + + void verify_context_kv_weights() const { + const int32_t num_layers = static_cast(layers_.size()); + CHECK_EQ(static_cast(per_layer_k_proj_.size()), num_layers); + CHECK_GT(local_kv_heads_, 0); + for (int32_t i = 0; i < num_layers; ++i) { + CHECK(per_layer_k_proj_[i].defined()); + CHECK(per_layer_v_proj_[i].defined()); + CHECK(per_layer_k_norm_[i].defined()); + } + } + + void build_fused_context_kv_weights() { + const int32_t num_layers = static_cast(layers_.size()); + std::vector kv_weights; + std::vector k_norm_weights; + kv_weights.reserve(num_layers * 2); + k_norm_weights.reserve(num_layers); + for (int32_t i = 0; i < num_layers; ++i) { + kv_weights.emplace_back(per_layer_k_proj_[i]); + kv_weights.emplace_back(per_layer_v_proj_[i]); + k_norm_weights.emplace_back(per_layer_k_norm_[i]); + } + fused_kv_weight_ = torch::cat(kv_weights, /*dim=*/0).contiguous(); + k_norm_weight_ = + torch::stack(k_norm_weights, /*dim=*/0).view({num_layers, 1, 1, -1}); + per_layer_k_proj_.clear(); + per_layer_v_proj_.clear(); + per_layer_k_norm_.clear(); + } + + layer::AddMatmul fc_{nullptr}; + layer::RMSNorm hidden_norm_{nullptr}; + DFlash2CandidateSelector selector_; + std::shared_ptr rotary_embedding_; + std::vector per_layer_k_proj_; + std::vector per_layer_v_proj_; + std::vector per_layer_k_norm_; + torch::Tensor fused_kv_weight_; + torch::Tensor k_norm_weight_; + torch::TensorOptions tensor_options_; + int64_t head_dim_ = 0; + int64_t local_kv_heads_ = 0; + double rms_norm_eps_ = 1e-6; + int32_t sliding_window_ = -1; + int32_t block_size_ = 0; + int32_t tp_rank_ = 0; + int32_t tp_size_ = 1; +}; +TORCH_MODULE(DFlash2Qwen3Model); + +class DFlash2Qwen3ForCausalLMImpl final + : public ::xllm::LlmForCausalLMImplBase { + public: + using Base = ::xllm::LlmForCausalLMImplBase; + + explicit DFlash2Qwen3ForCausalLMImpl(const ModelContext& context) + : Base(context) {} + + using Base::logits; + + torch::Tensor logits(const torch::Tensor& hidden_states, + const torch::Tensor& selected_idxes, + torch::Tensor& out_hidden) { + out_hidden = selected_idxes.defined() + ? hidden_states.index_select( + /*dim=*/0, selected_idxes.to(torch::kLong)) + : hidden_states; + return lm_head_(out_hidden); + } + + void load_model(std::unique_ptr loader, + std::string prefix = "") override { + for (const std::unique_ptr& state_dict : + loader->get_state_dicts()) { + model_->load_state_dict(state_dict->get_dict_with_prefix(prefix)); + } + model_->verify_loaded_weights(); + model_->finalize_loaded_weights(); + } + + ModelOutput write_context_kv(const torch::Tensor& target_hidden, + const torch::Tensor& positions, + const torch::Tensor& device_cache_slots, + std::vector& kv_caches, + const ModelInputParams& input_params) { + return model_->write_context_kv( + target_hidden, positions, device_cache_slots, kv_caches, input_params); + } + + DFlash2CandidateOutput dflash2_candidates( + const torch::Tensor& hidden_states, + const torch::Tensor& unary_logits, + const torch::Tensor& anchor_token_ids) { + return model_->candidates(hidden_states, unary_logits, anchor_token_ids); + } +}; +TORCH_MODULE(DFlash2Qwen3ForCausalLM); + +REGISTER_CAUSAL_MODEL_WITH_VARNAME(dflash2_draft_model, + DFlash2DraftModel, + DFlash2Qwen3ForCausalLM); + +} // namespace xllm::npu::model diff --git a/xllm/models/llm/qwen3.h b/xllm/models/llm/qwen3.h index 13079f6dea..5e6cea8c6d 100644 --- a/xllm/models/llm/qwen3.h +++ b/xllm/models/llm/qwen3.h @@ -46,8 +46,15 @@ class QWen3ModelImpl : public LlmModelImplBase { layers_.reserve(model_args.n_layers()); norm_ = register_module("norm", layer::RMSNorm(context)); - embed_tokens_ = - register_module("embed_tokens", layer::WordEmbedding(context)); + // DFlash2 consumes the target model's embedding table. Do not allocate a + // draft-local copy here: set_word_embedding() binds the shared target + // module after both workers finish loading. Merely replacing the holder + // after register_module() would leave the original large-vocabulary module + // owned by torch::nn::Module and keep its device memory alive. + if (!is_dflash2_draft_model_type(model_args.model_type())) { + embed_tokens_ = + register_module("embed_tokens", layer::WordEmbedding(context)); + } #if defined(USE_NPU) attn_mask_ = layer::AttentionMask( options.device(), options.dtype().toScalarType(), /*mask_value=*/-9984); @@ -182,7 +189,7 @@ class QWen3ModelImpl : public LlmModelImplBase { } protected: - layer::AttentionMetadata get_attention_metadata( + virtual layer::AttentionMetadata get_attention_metadata( const ModelInputParams& params, const torch::Tensor& h) { #if defined(USE_NPU) @@ -216,7 +223,7 @@ class QWen3ModelImpl : public LlmModelImplBase { for (int32_t j = 0; j < num_sequences; ++j) { const int32_t q_len = params.attention.host.q_seq_lens[j]; const int32_t kv_len = params.attention.host.kv_seq_lens[j]; - req_mask_vec.emplace_back(attn_mask_.gen_append_mask( + req_mask_vec.emplace_back(gen_append_attn_mask( q_len, kv_len, max_kv_seq, h.dtype().toScalarType(), h.device())); } return layer::AttentionMetadataBuilder::build( @@ -227,6 +234,16 @@ class QWen3ModelImpl : public LlmModelImplBase { #endif } +#if defined(USE_NPU) + virtual torch::Tensor gen_append_attn_mask(int32_t q_len, + int32_t kv_len, + int32_t max_kv_len, + torch::Dtype dtype, + torch::Device device) { + return attn_mask_.gen_append_mask(q_len, kv_len, max_kv_len, dtype, device); + } +#endif + private: #if defined(USE_NPU) layer::AttentionMask attn_mask_; @@ -259,6 +276,7 @@ REGISTER_CAUSAL_MODEL(qwen3, QWen3ForCausalLM); REGISTER_MODEL_ARGS(qwen3, [&] { LOAD_ARG_OR(model_type, "model_type", "qwen3"); LOAD_ARG_OR(dtype, "torch_dtype", ""); + LOAD_ARG_OR(dtype, "dtype", args->dtype()); LOAD_ARG_OR(vocab_size, "vocab_size", 152064); LOAD_ARG_OR(hidden_size, "hidden_size", 3584); LOAD_ARG_OR(hidden_act, "hidden_act", "silu"); @@ -271,13 +289,21 @@ REGISTER_MODEL_ARGS(qwen3, [&] { LOAD_ARG_OR(rms_norm_eps, "rms_norm_eps", 1e-6); LOAD_ARG_OR(eos_token_id, "eos_token_id", 151643); LOAD_ARG_OR(rope_theta, "rope_theta", 1000000.0f); + LOAD_ARG_OR(rope_theta, "rope_parameters.rope_theta", args->rope_theta()); // For qwen3/2.5 model < 7B, tie_word_embeddings = true LOAD_ARG_OR(tie_word_embeddings, "tie_word_embeddings", false); LOAD_ARG_OR(use_sliding_window, "use_sliding_window", false); + LOAD_ARG_OR(sliding_window, "sliding_window", -1); LOAD_ARG_OR(max_window_layers, "max_window_layers", 28); + LOAD_ARG_OR(dflash2_block_size, "dflash_config.block_size", 0); + LOAD_ARG_OR(dflash2_conv_group_size, "dflash_config.conv_group_size", 0); + LOAD_ARG_OR(dflash2_conv_kernel_size, "dflash_config.conv_kernel_size", 0); + LOAD_ARG_OR(dflash2_selector_rank, "dflash_config.selector_rank", 0); + LOAD_ARG_OR(dflash2_selector_top_k, "dflash_config.selector_top_k", 0); + LOAD_ARG_OR_FUNC(head_dim, "head_dim", [&] { return args->hidden_size() / args->n_heads(); }); diff --git a/xllm/models/llm/qwen3_next_hybrid_base.h b/xllm/models/llm/qwen3_next_hybrid_base.h index 45d17b71d7..22d9949b19 100644 --- a/xllm/models/llm/qwen3_next_hybrid_base.h +++ b/xllm/models/llm/qwen3_next_hybrid_base.h @@ -24,7 +24,9 @@ limitations under the License. #include #include "core/common/flash_comm1_context.h" +#include "core/framework/config/scheduler_config.h" #include "core/framework/kv_cache/kv_cache.h" +#include "core/framework/model/aux_hidden_capture.h" #include "core/framework/model/model_input_params.h" #include "core/framework/model/model_output.h" #include "core/framework/model_context.h" @@ -65,7 +67,10 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule { : device_(context.get_tensor_options().device()), model_args_(context.get_model_args()), parallel_args_(context.get_parallel_args()), - flash_comm1_options_(context.get_flash_comm1_options()) { + flash_comm1_options_(context.get_flash_comm1_options()), + aux_capture_(context.get_model_args(), + context.get_tensor_options(), + SchedulerConfig::get_instance().max_tokens_per_batch()) { if (model_args_.n_routed_experts() > 0) { flash_comm1_options_.enable_flashcomm1 = false; flash_comm1_options_.enable_mmrs_fusion = false; @@ -135,8 +140,14 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule { const int32_t num_tokens = static_cast(tokens.size(0)); const auto& batch_forward_type = input_params.meta.batch_forward_type; const bool is_prefill_side = batch_forward_type.no_decode(); - FlashComm1Context fc1_ctx = build_flash_comm1_context( - num_tokens, is_prefill_side, parallel_args_, flash_comm1_options_); + FlashComm1Context fc1_ctx; + // Sequence sharding changes the token-row layout. Per-layer auxiliary + // capture must keep the full DP-local token set so every capture slot and + // the final hidden states have identical row ordering. + if (!aux_capture_.enabled()) { + fc1_ctx = build_flash_comm1_context( + num_tokens, is_prefill_side, parallel_args_, flash_comm1_options_); + } FlashComm1ContextScope fc1_scope(&fc1_ctx); torch::Tensor h; @@ -169,6 +180,11 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule { attn_metadata.unshared_plan_info->layer_id = static_cast(i); } #endif + // Capture hooks run before a layer. Worker-side draft config loading + // converts target output layer L to capture index L + 1, matching the + // established Qwen3 DFlash convention. Hybrid layers keep the residual + // stream split between h and residual, so capture their sum. + aux_capture_.capture_layer(static_cast(i), h, residual); auto& layer = layers_[i]; h = layer->forward(h, residual, @@ -190,7 +206,7 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule { if (is_sequence_sharded(fc1_ctx)) { h = gather_sequence(h, fc1_ctx); } - return ModelOutput(h); + return aux_capture_.finalize(h); } // load the weight from the checkpoint @@ -285,6 +301,7 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule { layer::AttentionMask attn_mask_; layer::AttentionMask dense_attn_mask_; layer::WordEmbedding embed_tokens_{nullptr}; + AuxHiddenCapture aux_capture_; }; class Qwen3HybridForCausalLMImplBase : public torch::nn::Module { diff --git a/xllm/models/model_registry.cpp b/xllm/models/model_registry.cpp index 0f9021687f..f8af3d0df7 100644 --- a/xllm/models/model_registry.cpp +++ b/xllm/models/model_registry.cpp @@ -24,6 +24,7 @@ limitations under the License. #include "core/framework/config/kernel_config.h" #include "core/framework/config/model_config.h" +#include "core/framework/model/model_args.h" #include "core/util/dit_model_discovery.h" #include "llm/py_causal_lm.h" #include "models.h" @@ -83,7 +84,8 @@ bool is_torch_only_model_type(const std::string& model_type) { "qwen3_5_moe_mtp", "qwen3_next", "minimax_m2"}; - return kTorchOnlyModelTypes.count(model_type) != 0; + return kTorchOnlyModelTypes.count(model_type) != 0 || + is_dflash2_draft_model_type(model_type); } #endif diff --git a/xllm/models/models.h b/xllm/models/models.h index a07ba8c0ad..9548545252 100644 --- a/xllm/models/models.h +++ b/xllm/models/models.h @@ -47,6 +47,7 @@ limitations under the License. #include "llm/npu/qwen2.h" // IWYU pragma: keep #include "llm/npu/qwen3.h" // IWYU pragma: keep #include "llm/npu/qwen3_dflash.h" // IWYU pragma: keep +#include "llm/npu/qwen3_dflash2.h" // IWYU pragma: keep #include "llm/npu/qwen3_dspark.h" // IWYU pragma: keep #include "llm/npu/qwen3_eagle3.h" // IWYU pragma: keep #include "llm/npu/qwen3_moe.h" // IWYU pragma: keep diff --git a/xllm/pybind/args.py b/xllm/pybind/args.py index be09c14796..a2110d60ff 100644 --- a/xllm/pybind/args.py +++ b/xllm/pybind/args.py @@ -86,7 +86,7 @@ def __init__(self) -> None: "--speculative_algorithm", type=str, default="MTP", - help="Speculative decoding algorithm. Supported options: MTP, Eagle3, Suffix, DFlash, DSpark.", + help="Speculative decoding algorithm. Supported options: MTP, Eagle3, Suffix, DFlash, DFlash2 (NPU only), DSpark.", ) self.parser.add_argument( "--num_request_handling_threads", type=int, default=4, help="Number of handling threads."