diff --git a/docs/src/content/docs/en/cli_reference.md b/docs/src/content/docs/en/cli_reference.md index 2733fe8676..66d8c75bdf 100644 --- a/docs/src/content/docs/en/cli_reference.md +++ b/docs/src/content/docs/en/cli_reference.md @@ -183,7 +183,7 @@ xLLM uses gflags to manage service startup parameters. `--model ` is the o | `draft_model` | `string` | `""` | Draft model path. See [MTP](/en/features/mtp/) for MTP usage. | | `draft_devices` | `string` | `""` | Devices used by the draft model, for example `npu:0` or `npu:0,npu:1`. If omitted, uses the target model devices when speculative decoding is enabled. | | `num_speculative_tokens` | `int32` | `0` | Number of speculative tokens generated per speculative decoding step. | -| `speculative_algorithm` | `string` | `"MTP"` | Speculative decoding algorithm. Supported values: `MTP`, `Eagle3`, `Suffix`. | +| `speculative_algorithm` | `string` | `"MTP"` | Speculative decoding algorithm. Supported values: `MTP`, `Eagle3`, `Suffix`, `DFlash`. | | `speculative_suffix_cache_max_depth` | `int32` | `64` | Maximum suffix-tree depth for suffix speculative decoding. | | `speculative_suffix_max_spec_factor` | `double` | `1.0` | Maximum suffix speculation token factor relative to match length. | | `speculative_suffix_max_spec_offset` | `double` | `0.0` | Maximum additive token offset for suffix speculation. | diff --git a/docs/src/content/docs/zh/cli_reference.md b/docs/src/content/docs/zh/cli_reference.md index c98421b6f0..368bf908e5 100644 --- a/docs/src/content/docs/zh/cli_reference.md +++ b/docs/src/content/docs/zh/cli_reference.md @@ -183,7 +183,7 @@ xLLM 使用 gflags 管理服务启动参数。`--model ` 是唯一必填 | `draft_model` | `string` | `""` | draft 模型路径；MTP 使用方式详见 [MTP](/zh/features/mtp/)。 | | `draft_devices` | `string` | `""` | draft 模型使用的设备，例如 `npu:0`、`npu:0,npu:1`；未指定时，启用 speculative decoding 会使用 target 模型的设备。 | | `num_speculative_tokens` | `int32` | `0` | 每轮 speculative decoding 生成的 speculative token 数。 | -| `speculative_algorithm` | `string` | `"MTP"` | Speculative decoding 算法，支持 `MTP`、`Eagle3`、`Suffix`。 | +| `speculative_algorithm` | `string` | `"MTP"` | Speculative decoding 算法，支持 `MTP`、`Eagle3`、`Suffix`、`DFlash`。 | | `speculative_suffix_cache_max_depth` | `int32` | `64` | Suffix speculative decoding 的后缀树最大深度。 | | `speculative_suffix_max_spec_factor` | `double` | `1.0` | Suffix speculation 相对于匹配长度的最大 token 系数。 | | `speculative_suffix_max_spec_offset` | `double` | `0.0` | Suffix speculation 的最大 token 加性偏移。 | diff --git a/tests/core/scheduler/profile/CMakeLists.txt b/tests/core/scheduler/profile/CMakeLists.txt index 037587006a..5bfaf259bf 100644 --- a/tests/core/scheduler/profile/CMakeLists.txt +++ b/tests/core/scheduler/profile/CMakeLists.txt @@ -6,8 +6,11 @@ cc_test( SRCS profile_graph_warmup_test.cpp DEPS + :xllm_server :block :profile :request GTest::gtest_main ) +target_link_libraries(profile_graph_warmup_test PRIVATE + "$") diff --git a/xllm/core/distributed_runtime/llm_engine.cpp b/xllm/core/distributed_runtime/llm_engine.cpp index 4559fab46c..ecd550782e 100644 --- a/xllm/core/distributed_runtime/llm_engine.cpp +++ b/xllm/core/distributed_runtime/llm_engine.cpp @@ -190,12 +190,16 @@ bool LLMEngine::init_model(MasterStatus master_status) { auto model_loader = ModelLoader::create(model_path); LOG(INFO) << "Initializing model from: " << model_path; - tokenizer_ = model_loader->tokenizer(); - CHECK(tokenizer_ != nullptr); - args_ = model_loader->model_args(); quant_args_ = model_loader->quant_args(); - tokenizer_args_ = model_loader->tokenizer_args(); + + // A draft engine is fed token ids and detokenized by the target, so it + // shares the target vocabulary and loads no tokenizer of its own. + if (!options_.is_draft_engine()) { + tokenizer_ = model_loader->tokenizer(); + CHECK(tokenizer_ != nullptr); + tokenizer_args_ = model_loader->tokenizer_args(); + } // compute the number of local kv heads and head dim const uint32_t world_size = dp_local_tp_size_; @@ -221,21 +225,23 @@ bool LLMEngine::init_model(MasterStatus master_status) { << ", dtype: " << dtype_ << ", kv_cache_dtype: " << options_.kv_cache_dtype(); - const int64_t tokenizer_vocab_size = tokenizer_->vocab_size(); - int64_t model_vocab_size = args_.vocab_size(); - if (tokenizer_vocab_size != model_vocab_size) { - // use tokenizer vocab size if model vocab size is not set - if (model_vocab_size <= 0) { - LOG(WARNING) << "Model vocab size is not set, using tokenizer vocab " - "size: " - << tokenizer_vocab_size; - args_.vocab_size(tokenizer_vocab_size); - } else if (tokenizer_vocab_size > model_vocab_size) { - LOG(WARNING) << "Unsafe vocab mismatch: tokenizer: " - << tokenizer_vocab_size << ", model: " << model_vocab_size; - } else { - LOG(INFO) << "Tokenizer/model vocab differ: tokenizer=" - << tokenizer_vocab_size << ", model=" << model_vocab_size; + if (tokenizer_ != nullptr) { + const int64_t tokenizer_vocab_size = tokenizer_->vocab_size(); + int64_t model_vocab_size = args_.vocab_size(); + if (tokenizer_vocab_size != model_vocab_size) { + // use tokenizer vocab size if model vocab size is not set + if (model_vocab_size <= 0) { + LOG(WARNING) << "Model vocab size is not set, using tokenizer vocab " + "size: " + << tokenizer_vocab_size; + args_.vocab_size(tokenizer_vocab_size); + } else if (tokenizer_vocab_size > model_vocab_size) { + LOG(WARNING) << "Unsafe vocab mismatch: tokenizer: " + << tokenizer_vocab_size << ", model: " << model_vocab_size; + } else { + LOG(INFO) << "Tokenizer/model vocab differ: tokenizer=" + << tokenizer_vocab_size << ", model=" << model_vocab_size; + } } } diff --git a/xllm/core/distributed_runtime/speculative_engine.cpp b/xllm/core/distributed_runtime/speculative_engine.cpp index bb434dd102..54c7507055 100644 --- a/xllm/core/distributed_runtime/speculative_engine.cpp +++ b/xllm/core/distributed_runtime/speculative_engine.cpp @@ -97,27 +97,9 @@ bool SpeculativeEngine::init_model() { return false; } - // check if the tokenizers are compatible - const auto* draft_tokenizer = draft_engine_->tokenizer(); - const auto* target_tokenizer = engine_->tokenizer(); - if (draft_tokenizer->vocab_size() != target_tokenizer->vocab_size()) { - LOG(ERROR) << "draft and target tokenizers have different vocab sizes, " - "draft vocab_size: " - << draft_tokenizer->vocab_size() - << ", target vocab_size: " << target_tokenizer->vocab_size(); - return false; - } - - const std::string test_text = "hello from xllm!"; - std::vector draft_token_ids; - std::vector target_token_ids; - if (!draft_tokenizer->encode(test_text, &draft_token_ids) || - !target_tokenizer->encode(test_text, &target_token_ids)) { - if (draft_token_ids != target_token_ids) { - LOG(ERROR) << "draft and target tokenizers are not compatible"; - return false; - } - } + // Draft tokens are detokenized by the target, so the draft engine needs no + // tokenizer and no draft/target vocab-compatibility check (not universal; + // see llm_engine.cpp). // check if the max context length are the same const auto& draft_model_args = draft_engine_->model_args(); diff --git a/xllm/core/distributed_runtime/worker_server.cpp b/xllm/core/distributed_runtime/worker_server.cpp index 5b3a93f403..d1632ff92c 100644 --- a/xllm/core/distributed_runtime/worker_server.cpp +++ b/xllm/core/distributed_runtime/worker_server.cpp @@ -42,6 +42,7 @@ limitations under the License. #include "core/framework/config/kv_cache_config.h" #include "core/framework/config/parallel_config.h" #include "core/framework/config/service_config.h" +#include "core/framework/config/speculative_config.h" #include "core/platform/platform.h" #if defined(USE_CUDA) || defined(USE_MLU) || defined(USE_DCU) #include "core/platform/numa_utils.h" @@ -418,6 +419,15 @@ WorkerServer::WorkerServer(int32_t local_worker_idx, bool use_spawn_worker) : server_name_("DistributeWorkerServer") { server_name_.append(std::to_string(options.server_idx())); + // Eagle3/DFlash targets capture aux hidden states and don't support spawned + // workers yet. + const std::string& speculative_algorithm = options.speculative_algorithm(); + if (use_spawn_worker && options.enable_speculative_decode() && + SpeculativeConfig::requires_aux_hidden_capture(speculative_algorithm)) { + LOG(FATAL) << speculative_algorithm + << " does not support spawned workers yet. Disable offline " + "spawn workers or use --speculative_algorithm=MTP."; + } if (worker_type == WorkerType::LLM || worker_type == WorkerType::ELM || worker_type == WorkerType::VLM || worker_type == WorkerType::EVLM || diff --git a/xllm/core/framework/config/speculative_config.cpp b/xllm/core/framework/config/speculative_config.cpp index 3e8ed709c5..d74b498e0b 100644 --- a/xllm/core/framework/config/speculative_config.cpp +++ b/xllm/core/framework/config/speculative_config.cpp @@ -33,7 +33,7 @@ DEFINE_int32(num_speculative_tokens, 0, "Number of speculative tokens."); DEFINE_string(speculative_algorithm, "MTP", "Speculative decoding algorithm. Supported options: MTP, Eagle3, " - "Suffix. Default is MTP."); + "Suffix, DFlash. Default is MTP."); DEFINE_int32(speculative_suffix_cache_max_depth, 64, diff --git a/xllm/core/framework/config/speculative_config.h b/xllm/core/framework/config/speculative_config.h index d1d9697796..54c98b9db9 100644 --- a/xllm/core/framework/config/speculative_config.h +++ b/xllm/core/framework/config/speculative_config.h @@ -18,6 +18,7 @@ limitations under the License. #include #include #include +#include #include "core/common/macros.h" #include "core/framework/config/option_category.h" @@ -33,6 +34,18 @@ class SpeculativeConfig final { static SpeculativeConfig& get_instance(); + // Whether a speculative algorithm requires the target model to capture + // intermediate-layer aux hidden states to drive the draft (Eagle3 and + // DFlash). Centralizes the algorithm-string classification in one place. The + // worker consults it to decide whether to populate the target's + // layers_to_capture; the model then keys off that list alone, never the + // algorithm string. Takes the algorithm explicitly so a spawned worker + // process, which reads its own Options rather than this global config, can + // classify without an initialized singleton. + static bool requires_aux_hidden_capture(std::string_view algorithm) { + return algorithm == "Eagle3" || algorithm == "DFlash"; + } + void from_flags(); void from_json(const JsonReader& json); void append_config_json(nlohmann::ordered_json& config_json) const; diff --git a/xllm/core/framework/model/causal_lm.h b/xllm/core/framework/model/causal_lm.h index 6f8f09edd0..8ac6fb26b7 100644 --- a/xllm/core/framework/model/causal_lm.h +++ b/xllm/core/framework/model/causal_lm.h @@ -156,6 +156,20 @@ class CausalLM : public torch::nn::Module { NOT_IMPLEMENTED(); } + // DFlash-specific interface. Attention-free scatter-only pass that projects + // target hidden into the draft's per-layer KV cache. Not part of forward() + // because it has no attention and its shape doesn't match forward's decode + // graph; sits outside the executor to avoid a per-flag branch in the eager + // gate. Default: NOT_IMPLEMENTED. Overridden by DFlashDraftModel. + virtual ModelOutput write_context_kv(const torch::Tensor& target_hidden, + const torch::Tensor& positions, + const torch::Tensor& device_cache_slots, + std::vector& kv_caches, + const ModelInputParams& input_params) { + NOT_IMPLEMENTED(); + return {}; + } + virtual void lazy_load_model(std::unique_ptr loader) { NOT_IMPLEMENTED(); } @@ -238,6 +252,26 @@ class CausalLMImpl : public CausalLM { model_->load_model(std::move(loader)); } + ModelOutput write_context_kv(const torch::Tensor& target_hidden, + const torch::Tensor& positions, + const torch::Tensor& device_cache_slots, + std::vector& kv_caches, + const ModelInputParams& input_params) override { + if constexpr (detail::has_write_context_kv::value) { + return model_->write_context_kv(target_hidden, + positions, + device_cache_slots, + kv_caches, + input_params); + } else { + return CausalLM::write_context_kv(target_hidden, + positions, + device_cache_slots, + kv_caches, + input_params); + } + } + void lazy_load_model(std::unique_ptr loader) override { if constexpr (detail::has_lazy_load_model::value) { model_->lazy_load_model(std::move(loader)); diff --git a/xllm/core/framework/model/model_traits.h b/xllm/core/framework/model/model_traits.h index 2c68690782..13d3b2d154 100644 --- a/xllm/core/framework/model/model_traits.h +++ b/xllm/core/framework/model/model_traits.h @@ -20,10 +20,12 @@ limitations under the License. #include #include #include +#include namespace xllm { struct ModelInputParams; struct ModelGraphMetadataState; +class KVCache; namespace layer { class LmHead; @@ -218,5 +220,19 @@ struct has_init_or_refresh_rolling_runtime< std::declval()))>> : std::true_type {}; #endif + +template +struct has_write_context_kv : std::false_type {}; + +template +struct has_write_context_kv< + T, + std::void_t()->write_context_kv( + std::declval(), + std::declval(), + std::declval(), + std::declval&>(), + std::declval()))>> : std::true_type {}; + } // namespace detail } // namespace xllm diff --git a/xllm/core/runtime/CMakeLists.txt b/xllm/core/runtime/CMakeLists.txt index 7aab17b434..5937cadc1b 100644 --- a/xllm/core/runtime/CMakeLists.txt +++ b/xllm/core/runtime/CMakeLists.txt @@ -36,6 +36,7 @@ cc_library( mtp_worker_impl.h suffix_worker_impl.h eagle3_worker_impl.h + dflash_worker_impl.h spec_input_builder.h forward_shared_memory_manager.h worker_rendezvous.h @@ -68,6 +69,7 @@ cc_library( mtp_worker_impl.cpp suffix_worker_impl.cpp eagle3_worker_impl.cpp + dflash_worker_impl.cpp spec_input_builder.cpp forward_shared_memory_manager.cpp cp_input_partition.cpp diff --git a/xllm/core/runtime/dflash_worker_impl.cpp b/xllm/core/runtime/dflash_worker_impl.cpp new file mode 100644 index 0000000000..a7afa89c16 --- /dev/null +++ b/xllm/core/runtime/dflash_worker_impl.cpp @@ -0,0 +1,1143 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "runtime/dflash_worker_impl.h" + +#include + +#include +#include +#include +#include +#include + +#include "common/metrics.h" +#include "core/framework/config/disagg_pd_config.h" +#include "core/framework/config/kernel_config.h" +#include "core/framework/config/scheduler_config.h" +#include "core/framework/config/speculative_config.h" +#include "framework/model/model_args.h" +#include "framework/parallel_state/process_group.h" +#include "framework/sampling/rejection_sampler.h" +#if defined(USE_MLU) +#include "framework/kv_cache_transfer/mooncake_kv_cache_transfer.h" +#endif +#if defined(USE_NPU) +#include "framework/kv_cache_transfer/kv_transfer_completion.h" +#include "framework/kv_cache_transfer/spec_kv_cache_transfer.h" +#endif +#include "runtime/spec_input_builder.h" +#include "util/json_reader.h" +#include "util/timer.h" + +namespace xllm { +namespace { + +// Under schedule-overlap the scheduler feeds placeholder (-1) tokens, so +// per-rank sampling RNG can diverge across the tensor-parallel group. +// Broadcasting the sampled draft/accepted tokens to the group's rank 0 keeps +// every rank's cached draft probs and accepted prefixes identical. No-op for a +// single rank (world_size <= 1). +ProcessGroup* spec_broadcast_group(const ParallelArgs& parallel_args) { + return parallel_args.tp_group_ != nullptr ? parallel_args.tp_group_ + : parallel_args.process_group_; +} + +void broadcast_spec_tokens(torch::Tensor& tokens, + ProcessGroup* pg, + int32_t root_rank = 0) { + if (pg == nullptr || pg->world_size() <= 1 || !tokens.defined()) { + return; + } + tokens = tokens.contiguous(); + pg->broadcast(tokens, root_rank); +} + +runtime::Options target_options(const runtime::Options& options) { + runtime::Options opts = options; + opts.enable_schedule_overlap(false) + .is_draft_engine(false) + .enable_graph_aux_hidden_states(true); + return opts; +} + +runtime::Options draft_options(const runtime::Options& options) { + runtime::Options opts = options; + opts.enable_schedule_overlap(false) + .is_draft_engine(true) + .num_decoding_tokens(1) + .num_speculative_tokens(0) + .enable_graph_aux_hidden_states(false); + return opts; +} + +// Pack a host int32 vector into a pinned CPU tensor and stage an async H2D +// copy onto the caller's active stream. Consolidates the three-line idiom +// `TensorOptions(int, device) + specBuilder::make_cpu_int_tensor(vec) + +// safe_to(...)`. +torch::Tensor cpu_int_vec_to_device(const std::vector& values, + const Device& device) { + return safe_to( + specBuilder::make_cpu_int_tensor(values), + torch::TensorOptions().dtype(torch::kInt).device(device.unwrap()), + /*non_blocking=*/true); +} + +void repeat_sampling_tensor(torch::Tensor& tensor, int32_t repeats) { + if (tensor.defined()) { + tensor = tensor.repeat_interleave(/*repeats=*/repeats, /*dim=*/0); + } +} + +void repeat_sampling_params(SamplingParameters& sampling_params, + int32_t repeats) { + repeat_sampling_tensor(sampling_params.frequency_penalties, repeats); + repeat_sampling_tensor(sampling_params.presence_penalties, repeats); + repeat_sampling_tensor(sampling_params.repetition_penalties, repeats); + repeat_sampling_tensor(sampling_params.temperatures, repeats); + repeat_sampling_tensor(sampling_params.top_p, repeats); + repeat_sampling_tensor(sampling_params.top_k, repeats); + repeat_sampling_tensor(sampling_params.unique_token_ids, repeats); + repeat_sampling_tensor(sampling_params.unique_token_counts, repeats); + repeat_sampling_tensor(sampling_params.unique_token_ids_lens, repeats); + repeat_sampling_tensor(sampling_params.do_sample, repeats); +} + +void clear_selected_embeddings(ForwardOutput& output) { + output.sample_output.selected_embeddings = torch::Tensor(); +} + +void clear_all_output_embeddings(ForwardOutput& output) { + output.sample_output.embeddings = torch::Tensor(); + clear_selected_embeddings(output); +} + +void record_metadata_ready_event(Stream& stream, ForwardInput& input) { + StreamEventPtr event = stream.record_event(); + if (event == nullptr) { + stream.synchronize(); + } + input.metadata_ready_event = event; +} + +void wait_metadata_ready_event(const ForwardInput& input, Stream& stream) { + CHECK(stream.wait_event(input.metadata_ready_event)) + << "failed to wait DFlash metadata ready event"; +} + +void scale_dp_global_token_nums(ModelInputParams& input_params, + int32_t multiplier) { + for (int32_t& token_num : input_params.parallel.dp_global_token_nums) { + token_num *= multiplier; + } +} + +std::optional run_llm_no_sync_impl( + LLMWorkerImpl& worker, + const ForwardInput& input, + Stream& prepare_stream, + Stream& compute_stream, + 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); + if (processed_output != nullptr) { + *processed_output = std::move(processed_input); + } + return output; +} + +void build_query_rows(const ForwardInput& input, + int32_t mask_token_id, + int32_t num_speculative_tokens, + int32_t block_size, + specBuilder::DecodeBuildBuffers& buf, + std::vector& selected_idxes, + std::vector& q_cu_seq_lens) { + const int32_t num_sequences = input.input_params.meta.num_sequences; + const int32_t query_width = num_speculative_tokens + 1; + specBuilder::DecodeRowContext row_ctx = + specBuilder::make_decode_row_context(input); + Slice token_ids = { + input.token_ids_host.data_ptr(), + static_cast(input.token_ids_host.numel())}; + CHECK_GE(static_cast(token_ids.size()), num_sequences) + << "DFlash input token_ids size is smaller than num_sequences."; + + buf.out_token_ids.reserve(num_sequences * query_width); + buf.out_positions.reserve(num_sequences * query_width); + buf.out_new_cache_slots.reserve(num_sequences * query_width); + buf.out_kv_seq_lens.reserve(num_sequences); + buf.out_q_seq_lens.reserve(num_sequences); + + selected_idxes.reserve(num_sequences * num_speculative_tokens); + q_cu_seq_lens.reserve(num_sequences + 1); + q_cu_seq_lens.emplace_back(0); + + for (int32_t seq_id = 0; seq_id < num_sequences; ++seq_id) { + for (int32_t query_idx = 0; query_idx < query_width; ++query_idx) { + specBuilder::RowSpec row; + row.seq_id = seq_id; + row.token_id = query_idx == 0 ? token_ids[seq_id] : mask_token_id; + row.position_offset = query_idx; + row.append_kv_len = false; + row.append_q_len_one = false; + row.append_block_table = false; + specBuilder::append_decode_row(row_ctx, row, block_size, buf); + if (query_idx > 0) { + selected_idxes.emplace_back(seq_id * query_width + query_idx); + } + } + + specBuilder::append_seq_len_by_layout(buf.out_q_seq_lens, query_width); + q_cu_seq_lens.emplace_back(q_cu_seq_lens.back() + query_width); + const int32_t kv_len = + specBuilder::calc_kv_len(input.input_params.attention.host.kv_seq_lens, + seq_id, + /*offset=*/0) + + num_speculative_tokens; + specBuilder::update_kv_seq_lens_and_max( + buf.out_kv_seq_lens, kv_len, buf.meta.kv_max_seq_len); + } +} + +std::vector build_accepted_context_rows( + const ForwardInput& input, + const torch::Tensor& accepted_tokens_cpu, + int32_t block_size, + specBuilder::DecodeBuildBuffers& buf) { + const int32_t batch_size = static_cast(accepted_tokens_cpu.size(0)); + const int32_t token_width = static_cast(accepted_tokens_cpu.size(1)); + CHECK_EQ(input.input_params.meta.num_sequences, batch_size) + << "DFlash accepted token batch mismatch."; + + specBuilder::DecodeRowContext row_ctx = + specBuilder::make_decode_row_context(input); + std::vector accepted_idxes; + accepted_idxes.reserve(static_cast(accepted_tokens_cpu.numel())); + buf.out_positions.reserve(buf.out_positions.size() + + static_cast(accepted_tokens_cpu.numel())); + buf.out_new_cache_slots.reserve( + buf.out_new_cache_slots.size() + + static_cast(accepted_tokens_cpu.numel())); + + const int64_t* accepted_tokens_data = + accepted_tokens_cpu.const_data_ptr(); + for (int32_t seq_id = 0; seq_id < batch_size; ++seq_id) { + const int64_t row_offset = static_cast(seq_id) * token_width; + for (int32_t token_idx = 0; token_idx < token_width; ++token_idx) { + if (accepted_tokens_data[row_offset + token_idx] < 0) { + break; + } + + specBuilder::RowSpec row; + row.seq_id = seq_id; + row.position_offset = token_idx; + row.append_token = false; + row.append_kv_len = false; + specBuilder::append_decode_row(row_ctx, row, block_size, buf); + accepted_idxes.emplace_back(row_offset + token_idx); + } + } + + CHECK(!accepted_idxes.empty()) + << "DFlash accepted context must not be empty."; + CHECK_EQ(buf.out_new_cache_slots.size(), buf.out_positions.size()) + << "DFlash accepted context slots/positions mismatch."; + return accepted_idxes; +} + +} // namespace + +DFlashWorkerImpl::DFlashWorkerImpl(const ParallelArgs& parallel_args, + const torch::Device& device, + const runtime::Options& options) + : SpeculativeWorkerImpl(parallel_args, + device, + options, + target_options(options)) { + // DFlash feeds the target's captured intermediate-layer aux hidden states + // into the draft's context K/V. Under context parallelism the worker only + // exposes the lm_head-gathered final hidden (see llm_worker_impl.cpp), not + // the aux hidden, so the draft would silently receive the wrong tensor. + // Reject cp_size > 1 until aux-hidden plumbing under CP is implemented. + CHECK_LE(parallel_args.cp_size(), 1) + << "DFlash speculative decoding does not support context parallelism " + "(cp_size > 1)."; + draft_impl_ = std::make_unique( + parallel_args, device, draft_options(options)); +} + +bool DFlashWorkerImpl::init_model(const std::string& model_weights_path, + int32_t random_seed, + MasterStatus master_status) { + // DFlash draft attends each block non-causally, which the shared QWen3 model + // only wires up on the chunked-prefill mask path. Without it the draft falls + // back to a causal mask and proposal quality silently degrades, so require + // the flag rather than accept a misconfigured run. + CHECK(::xllm::SchedulerConfig::get_instance().enable_chunked_prefill()) + << "DFlash requires --enable_chunked_prefill=true."; + bool result = true; + const bool loading_target = + impl_->get_status() == WorkerImpl::Status::UNINITIALIZED; + if (loading_target) { + result = SpeculativeWorkerImpl::init_model( + model_weights_path, random_seed, master_status); + } else { + CHECK_EQ(draft_impl_->get_status(), WorkerImpl::Status::UNINITIALIZED); + // Draft config's use_sliding_window / sliding_window is intentionally + // ignored: xLLM NPU FIA hard-codes pre_tokens=INT_MAX, next_tokens=0 + // and never plumbs sliding_window into the aclnn call. Attended kv_len + // on the draft path is bounded by the target sequence length, not by + // block_size, so enforcing a block_size-vs-window relationship here + // would compare the wrong quantities. + result = draft_impl_->WorkerImpl::init_model( + model_weights_path, random_seed, master_status); + } + + if (impl_->get_status() == WorkerImpl::Status::LOADED) { + context_ = impl_->context_; + } + + if (draft_impl_->get_status() == WorkerImpl::Status::LOADED) { + // Draft shares the target's lm_head and word embedding to save memory and a + // redundant matmul. + auto share_torch_head_and_embedding = [this]() { + 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); + }; +#if defined(USE_NPU) + // The DFlash draft body is registered ATB-only, so a TORCH-backend run + // aborts in create_llm_model before reaching here; the draft always uses + // the target's NPU (ATB) head and embedding. + 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 + share_torch_head_and_embedding(); +#endif + + JsonReader reader; + const std::string config_path = model_weights_path + "/config.json"; + CHECK(reader.parse(config_path)) + << "Failed to parse DFlash draft config: " << config_path; + std::optional mask_token_id = + reader.value("dflash_config.mask_token_id"); + CHECK(mask_token_id.has_value()) + << "DFlash draft config requires dflash_config.mask_token_id."; + mask_token_id_ = mask_token_id.value(); + + const ModelArgs& draft_args = draft_impl_->context_.get_model_args(); + const int64_t draft_vocab_size = draft_args.vocab_size(); + CHECK_GT(draft_vocab_size, 0) << "DFlash draft vocab_size must be set."; + CHECK_GE(mask_token_id_, 0) << "DFlash mask_token_id (" << mask_token_id_ + << ") must be a valid embedding index (>= 0)."; + CHECK_LT(mask_token_id_, draft_vocab_size) + << "DFlash mask_token_id (" << mask_token_id_ + << ") must be < draft vocab_size (" << draft_vocab_size << ")."; + const int64_t num_target_layers = + static_cast(draft_args.layers_to_capture().size()); + CHECK_GT(num_target_layers, 0) + << "DFlash requires dflash_config.target_layer_ids."; + expected_context_hidden_size_ = + static_cast(draft_args.hidden_size()) * num_target_layers; + } + return result; +} + +std::tuple DFlashWorkerImpl::estimate_kv_cache_capacity() { + const std::tuple target_memory = + impl_->estimate_kv_cache_capacity(); + const std::tuple draft_memory = + draft_impl_->estimate_kv_cache_capacity(); + const int64_t cache_size_in_bytes = + std::min(std::get<0>(target_memory), std::get<0>(draft_memory)); + const int64_t total_memory = + std::min(std::get<1>(target_memory), std::get<1>(draft_memory)); + return {cache_size_in_bytes, total_memory}; +} + +bool DFlashWorkerImpl::allocate_kv_cache(const KVCacheShape& kv_cache_shape) { + const int64_t num_blocks = kv_cache_shape.key_cache_shape()[0]; + embedding_cache_ = std::make_shared(num_blocks); + + bool target_allocated = true; + const WorkerImpl::Status target_status = impl_->get_status(); + if (target_status == WorkerImpl::Status::LOADED) { + target_allocated = impl_->allocate_kv_cache(kv_cache_shape); + } else { + CHECK_EQ(target_status, WorkerImpl::Status::READY); + } + + 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); + } else { + CHECK_EQ(draft_status, WorkerImpl::Status::READY); + } + + return target_allocated && draft_allocated; +} + +#if defined(USE_NPU) || defined(USE_MLU) +bool DFlashWorkerImpl::allocate_kv_cache_with_transfer( + const KVCacheShape& kv_cache_shape) { + const int64_t num_blocks = kv_cache_shape.key_cache_shape()[0]; + + if (kv_cache_transfer_ == nullptr) { +#if defined(USE_NPU) + kv_cache_transfer_ = std::make_shared( + options_.transfer_listen_port(), + options_.instance_role(), + context_.get_model_args().model_type(), + context_.get_model_args().index_n_heads() > 0); +#elif defined(USE_MLU) + CHECK_EQ(::xllm::DisaggPDConfig::get_instance().kv_cache_transfer_type(), + "Mooncake") + << "MLU DFlash push only supports Mooncake KV transfer."; + kv_cache_transfer_ = std::make_shared( + device_.index(), + options_.transfer_listen_port(), + device_, + context_.get_model_args().model_type()); +#endif + + const int32_t device_id = device_.index(); + kv_cache_transfer_->initialize(device_id); + } + + bool target_allocated = true; + const WorkerImpl::Status target_status = impl_->get_status(); + if (target_status == WorkerImpl::Status::LOADED) { + target_allocated = impl_->allocate_kv_cache_with_transfer( + kv_cache_transfer_, kv_cache_shape); + } else { + CHECK_EQ(target_status, WorkerImpl::Status::READY); + } + + 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); + } else { + CHECK_EQ(draft_status, WorkerImpl::Status::READY); + } + + embedding_cache_ = std::make_shared(num_blocks); + return target_allocated && draft_allocated; +} +#endif + +ForwardInput DFlashWorkerImpl::update_input_by_last_step_output( + ForwardInput& inputs) { + return inputs; +} + +std::optional DFlashWorkerImpl::step_empty( + const ForwardInput& input) { + if (!input.input_params.meta.batch_forward_type.is_decode()) { + std::optional output = + run_llm_no_sync_impl(*impl_, input, *prepare_stream_, *compute_stream_); + // Warmup only: prime the draft; its output is unused. Keep it alive until + // the sync below so the no-sync draft input is not freed while its kernel + // is still in flight. + std::optional draft_output = run_llm_no_sync_impl( + *draft_impl_, input, *prepare_stream_, *compute_stream_); + // Both forwards launched no-sync, so their staged inputs and the returned + // target output are still in flight. Sync before returning so a DP idle + // rank or graph warmup cannot reuse the input buffers, and non-overlap + // callers do not observe an unfinished target output. + compute_stream_->synchronize(); + if (output.has_value()) { + clear_all_output_embeddings(output.value()); + } + return output; + } + + const int32_t query_width = options_.num_speculative_tokens() + 1; + ForwardInput query_input = input; + query_input.input_params.meta.batch_forward_type = + BatchForwardType::CHUNKED_PREFILL; + query_input.input_params.meta.q_max_seq_len = query_width; + scale_dp_global_token_nums(query_input.input_params, query_width); + // Warmup only: prime the draft; its output is unused. Keep it alive until the + // sync below so the no-sync draft input is not freed while the target forward + // launched next can reuse the buffer. + std::optional draft_output = run_llm_no_sync_impl( + *draft_impl_, query_input, *prepare_stream_, *compute_stream_); + + ForwardInput validate_input = input; + scale_dp_global_token_nums(validate_input.input_params, query_width); + ForwardOutput output = + run_llm_no_sync_impl( + *impl_, validate_input, *prepare_stream_, *compute_stream_) + .value(); + // See above: sync the no-sync draft and target forwards before returning. + compute_stream_->synchronize(); + clear_all_output_embeddings(output); + return output; +} + +std::optional DFlashWorkerImpl::step_prefill( + const ForwardInput& input) { + Timer timer; + ForwardInput processed_target_input; + ForwardOutput output = run_llm_no_sync_impl(*impl_, + input, + *prepare_stream_, + *compute_stream_, + &processed_target_input) + .value(); + COUNTER_ADD(speculative_execution_latency_seconds_target, + timer.elapsed_seconds()); + + const torch::Tensor& embeddings = output.sample_output.embeddings; + if (embeddings.defined()) { + CHECK(processed_target_input.positions_host.defined()) + << "DFlash prefill requires processed positions_host."; + Slice positions = { + processed_target_input.positions_host.data_ptr(), + static_cast(processed_target_input.positions_host.numel())}; + CHECK_EQ(positions.size(), static_cast(embeddings.size(0))) + << "DFlash prefill hidden/position count mismatch."; + const std::vector& processed_new_cache_slots = + processed_target_input.input_params.attention.host.new_cache_slots; + CHECK_EQ(processed_new_cache_slots.size(), positions.size()) + << "DFlash prefill hidden/cache slot count mismatch."; + + timer.reset(); + write_context_kv( + processed_target_input, + embeddings, + processed_target_input.positions, + processed_target_input.input_params.attention.device.new_cache_slots); + COUNTER_ADD(speculative_execution_latency_seconds_draft, + timer.elapsed_seconds()); + } + + if (input.sampling_params.selected_token_idxes.defined()) { + embedding_cache_->write_prefill_target_context( + input.input_params.embedding.embedding_ids, + input.input_params.embedding.request_ids, + output.sample_output.next_tokens, + embeddings, + input.sampling_params.selected_token_idxes); + // PD handoff: the decode instance requires get_mtp_bootstrap_embedding() + // defined before it accepts the request (disagg_pd_scheduler). Compress the + // full prefill hidden to one row per sequence, as + // write_prefill_target_context stores it. + torch::Tensor bootstrap_embeddings = embeddings; + if (bootstrap_embeddings.size(0) != + static_cast( + input.input_params.embedding.embedding_ids.size())) { + torch::Tensor bootstrap_idxes = + input.sampling_params.selected_token_idxes.to( + torch::dtype(torch::kLong).device(bootstrap_embeddings.device())); + bootstrap_embeddings = + bootstrap_embeddings.index_select(/*dim=*/0, bootstrap_idxes); + } + output.sample_output.embeddings = bootstrap_embeddings.detach(); + clear_selected_embeddings(output); + } else { + clear_all_output_embeddings(output); + } + + if (!enable_schedule_overlap() && !driver_ && !dp_driver_) { + return std::nullopt; + } + return output; +} + +std::optional DFlashWorkerImpl::step_decode( + const ForwardInput& raw_input) { + ForwardInput input = raw_input; + ForwardInput validate_input; + + CHECK(embedding_cache_ != nullptr) + << "DFlash embedding cache is not allocated"; + + const auto& embedding = input.input_params.embedding; + if (embedding.mtp_bootstrap_embeddings.defined()) { + CHECK(input.token_ids_host.defined()) + << "DFlash bootstrap requires host token ids"; + CHECK(input.token_ids_host.device().is_cpu()) + << "DFlash bootstrap host token ids must be on CPU"; + CHECK_EQ(input.token_ids_host.scalar_type(), torch::kInt) + << "DFlash bootstrap host token ids must be int32"; + + torch::Tensor bootstrap_embeddings = + safe_to(embedding.mtp_bootstrap_embeddings, + torch::dtype(dtype_).device(device_)); + CHECK_EQ(bootstrap_embeddings.size(0), + static_cast(embedding.mtp_bootstrap_row_idxes.size())) + << "DFlash bootstrap row count mismatch"; + + Slice token_ids = { + input.token_ids_host.data_ptr(), + static_cast(input.token_ids_host.numel())}; + for (int32_t i = 0; + i < static_cast(embedding.mtp_bootstrap_row_idxes.size()); + ++i) { + const int32_t row_idx = embedding.mtp_bootstrap_row_idxes[i]; + CHECK_GE(row_idx, 0) << "DFlash bootstrap row index should be valid"; + CHECK_LT(row_idx, static_cast(embedding.embedding_ids.size())) + << "DFlash bootstrap row index exceeds embedding ids"; + CHECK_LT(row_idx, static_cast(embedding.request_ids.size())) + << "DFlash bootstrap row index exceeds request ids"; + CHECK_LT(static_cast(row_idx), input.token_ids_host.numel()) + << "DFlash bootstrap row index exceeds token ids"; + embedding_cache_->write_mtp_bootstrap_context( + embedding.embedding_ids[row_idx], + embedding.request_ids[row_idx], + token_ids[row_idx], + bootstrap_embeddings[i]); + } + } + + std::vector last_states = + embedding_cache_->read_decode_states( + input.input_params.embedding.embedding_ids, + input.input_params.embedding.request_ids); + CHECK_EQ(last_states.size(), + input.input_params.embedding.embedding_ids.size()) + << "DFlash decode target state count mismatch"; + + update_decode_step_input(input, last_states); + DraftBlock draft_block = run_decode_draft(input, validate_input); + return run_validate(input, draft_block, validate_input); +} + +DFlashWorkerImpl::DraftBlock DFlashWorkerImpl::run_decode_draft( + const ForwardInput& input, + ForwardInput& validate_input) { + Timer timer; + + ForwardInput query_input; + prepare_query_inputs(input, query_input); + + ForwardOutput draft_output = + run_llm_no_sync_impl( + *draft_impl_, query_input, *prepare_stream_, *compute_stream_) + .value(); + // Overlap validate input preparation with the async draft forward: the draft + // launch above returns immediately, so building validate_input here (it only + // reads the original input; draft tokens are injected later in + // fill_validate_input_from_draft_outputs) runs on the host while the draft + // computes on device, instead of delaying the draft launch. + prepare_validate_inputs(input, validate_input); + // Unify the draft next_tokens across the tensor-parallel group before + // process_draft_sample_output() compresses the probs into the cache, so every + // rank caches the same selected draft prob under schedule-overlap. No-op for + // a single rank. + maybe_broadcast_spec_tokens(draft_output.sample_output.next_tokens); + process_draft_sample_output(draft_output.sample_output); + COUNTER_ADD(speculative_execution_latency_seconds_draft, + timer.elapsed_seconds()); + + // Draft emits the whole block in one forward; reshape the flat outputs into + // [batch, num_speculative_tokens] instead of splitting into per-step outputs. + const int32_t num_speculative_tokens = options_.num_speculative_tokens(); + const int32_t num_draft_tokens = + static_cast(draft_output.sample_output.next_tokens.numel()); + CHECK_EQ(num_draft_tokens % num_speculative_tokens, 0) + << "DFlash draft token count mismatch."; + const int32_t batch_size = num_draft_tokens / num_speculative_tokens; + CHECK_EQ(draft_output.sample_output.probs.numel(), num_draft_tokens) + << "DFlash draft output requires selected draft probs."; + + DraftBlock draft_block; + draft_block.token_ids = draft_output.sample_output.next_tokens.view( + {batch_size, num_speculative_tokens}); + draft_block.probs = draft_output.sample_output.probs.view( + {batch_size, num_speculative_tokens}); + // Keep the draft's no-sync input alive past run_validate's compute-stream + // sync (see DraftBlock::draft_retained_input). + draft_block.draft_retained_input = std::move(draft_output.retained_input); + return draft_block; +} + +void DFlashWorkerImpl::fill_validate_input_from_draft_outputs( + const DraftBlock& draft_block, + ForwardInput& validate_input, + Stream& compute_stream) { + const int32_t num_speculative_tokens = options_.num_speculative_tokens(); + const int32_t num_val_tokens = num_speculative_tokens + 1; + CHECK(draft_block.token_ids.defined()) + << "DFlash draft token_ids must be defined for validate token fill"; + CHECK_EQ(draft_block.token_ids.dim(), 2) + << "DFlash draft token_ids must be [batch, num_speculative_tokens]"; + CHECK_EQ(draft_block.token_ids.size(1), num_speculative_tokens) + << "DFlash draft token_ids width mismatch"; + CHECK(validate_input.token_ids.defined()) + << "DFlash validate token_ids must be prepared before draft token fill"; + CHECK_EQ(validate_input.token_ids.dim(), 1) + << "DFlash validate token_ids must be flat"; + CHECK_EQ(validate_input.token_ids.numel() % num_val_tokens, 0) + << "DFlash validate token_ids size must be divisible by validation width"; + + const int64_t total_num_val_tokens = validate_input.token_ids.numel(); + const int64_t num_sequences = total_num_val_tokens / num_val_tokens; + CHECK_EQ(draft_block.token_ids.size(0), num_sequences) + << "DFlash draft batch must match validate sequence count"; + const torch::TensorOptions token_options = validate_input.token_ids.options(); + c10::StreamGuard stream_guard = compute_stream.set_stream_guard(); + wait_metadata_ready_event(validate_input, compute_stream); + torch::Tensor validate_token_rows = + validate_input.token_ids.view({num_sequences, num_val_tokens}); + + validate_input.device_tensors_ready = false; + // Column 0 keeps the real input token; the draft block fills columns + // [1, num_val_tokens) in one copy rather than per-step. + torch::Tensor draft_tokens = + safe_to(draft_block.token_ids, token_options, /*non_blocking=*/true); + using ISlice = torch::indexing::Slice; + validate_token_rows.index({ISlice(), ISlice(1, num_val_tokens)}) + .copy_(draft_tokens, /*non_blocking=*/true); + validate_input.device_tensors_ready = true; + // Publish this compute-stream write so the target's prepare stage (which + // consumes validate_input.token_ids under ACL-graph double buffering) waits + // for the copy to complete before staging into the graph's persistent + // buffer. Without this, prepare could read stale/placeholder token ids. + record_metadata_ready_event(compute_stream, validate_input); +} + +std::optional DFlashWorkerImpl::run_validate( + const ForwardInput& input, + const DraftBlock& draft_block, + ForwardInput& validate_input) { + Timer timer; + fill_validate_input_from_draft_outputs( + draft_block, validate_input, *compute_stream_); + ForwardOutput target_output = + run_llm_no_sync_impl( + *impl_, validate_input, *prepare_stream_, *compute_stream_) + .value(); + COUNTER_ADD(speculative_execution_latency_seconds_target, + timer.elapsed_seconds()); + + timer.reset(); + SampleOutput val_output = + validate(input.sampling_params, draft_block, target_output); + COUNTER_ADD(speculative_execution_latency_seconds_validation, + timer.elapsed_seconds()); + + // target forward and validate()'s reads share compute_stream_, so they + // serialize without a cross-stream wait; the sync below only makes the + // accepted tokens host-visible for the D2H copy and context-cache write. + maybe_broadcast_spec_tokens(val_output.next_tokens); + compute_stream_->synchronize(); + val_output.next_tokens = val_output.next_tokens.to(torch::kCPU); + write_target_context_to_cache(input, val_output); + + if (!enable_schedule_overlap() && !driver_ && !dp_driver_) { + return std::nullopt; + } + val_output.embeddings = torch::Tensor(); + target_output.sample_output = val_output; + return target_output; +} + +SampleOutput DFlashWorkerImpl::validate( + const SamplingParameters& sampling_params, + const DraftBlock& draft_block, + const ForwardOutput& target_output) { + // Draft already emits the whole block [batch, num_speculative_tokens]; feed + // it straight to the verifier without the per-step select/view/cat round + // trip. The shared rejection sampler uses MTP's dense contract, so + // reconstruct dense draft probs unless the selected-only optimization is on. + const int32_t vocab_size = + static_cast(target_output.logits.size(/*dim=*/-1)); + const bool enable_opt_validate_probs = + ::xllm::SpeculativeConfig::get_instance().enable_opt_validate_probs(); + auto [draft_token_ids, draft_probs] = + specBuilder::draftProbs::build_validate_tensors_from_block( + draft_block.token_ids, + draft_block.probs, + vocab_size, + enable_opt_validate_probs); + return validate(sampling_params, draft_token_ids, draft_probs, target_output); +} + +SampleOutput DFlashWorkerImpl::validate( + const SamplingParameters& sampling_params, + const torch::Tensor& draft_token_ids, + const torch::Tensor& draft_probs, + const ForwardOutput& target_output) { + const int32_t num_val_tokens = options_.num_speculative_tokens() + 1; + // Derive batch_size from the target logits rows rather than next_tokens so + // the reshape stays valid regardless of how the target was sampled. + const int32_t num_logits_rows = + static_cast(target_output.logits.size(/*dim=*/0)); + CHECK_EQ(num_logits_rows % num_val_tokens, 0) + << "DFlash validate target logits rows must be divisible by validation " + "width"; + const int32_t batch_size = num_logits_rows / num_val_tokens; + const int32_t vocab_size = + static_cast(target_output.logits.size(/*dim=*/-1)); + + using torch::indexing::None; + using ISlice = torch::indexing::Slice; + torch::Tensor bonus_token_ids = + target_output.sample_output.next_tokens + .index({"...", ISlice(num_val_tokens - 1, None, num_val_tokens)}) + .view({-1, 1}); + + torch::Tensor target_logits = + target_output.logits.view({batch_size, num_val_tokens, vocab_size}); + + auto rejection_sampler = + std::make_unique(sampling_params.do_sample, + sampling_params.all_random_sample, + sampling_params.all_greedy_sample, + target_output.logprobs, + target_output.max_top_logprobs, + enable_fused_kernel_); + + SampleOutput sample_output = + rejection_sampler->forward(draft_token_ids.to(bonus_token_ids), + draft_probs.to(target_logits.device()), + target_logits, + bonus_token_ids, + /*mask_out_rejected_tokens=*/true); + + const torch::Tensor& embeddings = target_output.sample_output.embeddings; + sample_output.embeddings = + embeddings.view({batch_size, num_val_tokens, embeddings.size(-1)}); + return sample_output; +} + +void DFlashWorkerImpl::process_draft_sample_output( + SampleOutput& sample_output) { + if (sample_output.probs.defined()) { + CHECK(sample_output.next_tokens.defined()) + << "DFlash draft sample_output.next_tokens must be defined when probs " + "exist"; + CHECK_EQ(sample_output.next_tokens.dim(), 1) + << "DFlash draft cache expects next_tokens [batch], got " + << sample_output.next_tokens.sizes(); + CHECK(sample_output.probs.dim() == 1 || sample_output.probs.dim() == 2) + << "DFlash draft cache expects probs [batch] or [batch,vocab], got " + << sample_output.probs.sizes(); + CHECK_EQ(sample_output.probs.size(0), sample_output.next_tokens.size(0)) + << "DFlash draft cache probs/token batch mismatch"; + sample_output.probs = specBuilder::draftProbs::compress_for_cache( + sample_output.probs, sample_output.next_tokens); + } +} + +void DFlashWorkerImpl::maybe_broadcast_spec_tokens(torch::Tensor& tokens) { + if (get_optimization_config().enable_spec_token_broadcast && + enable_schedule_overlap()) { + c10::StreamGuard stream_guard = compute_stream_->set_stream_guard(); + broadcast_spec_tokens(tokens, spec_broadcast_group(parallel_args_)); + } +} + +void DFlashWorkerImpl::update_decode_step_input( + ForwardInput& input, + const std::vector& last_states) const { + const int32_t num_sequences = input.input_params.meta.num_sequences; + CHECK_EQ(last_states.size(), static_cast(num_sequences)) + << "DFlash decode context state count mismatch"; + const bool enable_cache_correction = enable_schedule_overlap(); + + std::vector token_ids_vec; + std::vector positions_vec; + std::vector kv_seq_lens_vec; + token_ids_vec.reserve(num_sequences); + positions_vec.reserve(num_sequences); +#if defined(USE_NPU) + kv_seq_lens_vec.reserve(num_sequences); +#else + kv_seq_lens_vec.reserve(num_sequences + 1); +#endif + + const torch::Tensor& token_ids_cpu = input.token_ids_host; + const torch::Tensor& positions_cpu = input.positions_host; + Slice input_token_ids = {token_ids_cpu.data_ptr(), + static_cast(token_ids_cpu.numel())}; + Slice input_positions = {positions_cpu.data_ptr(), + static_cast(positions_cpu.numel())}; + + for (int32_t seq_id = 0; seq_id < num_sequences; ++seq_id) { + CHECK_LT(static_cast(seq_id), input_token_ids.size()) + << "DFlash decode context token seq_id out of range, seq_id=" << seq_id; + CHECK_LT(static_cast(seq_id), input_positions.size()) + << "DFlash decode context position seq_id out of range, seq_id=" + << seq_id; + const EmbeddingCache::DecodeState& state = last_states[seq_id]; + const int32_t input_token_id = input_token_ids[seq_id]; + const bool input_is_fake_token = input_token_id < 0; + // Rewrite fake input tokens to the last committed real token so KV cache + // scatter has a valid id; only apply the recorded position offset when the + // cached state is still valid. + const bool rewrite_fake_token = + enable_cache_correction && input_is_fake_token; + const bool use_cache_correction = rewrite_fake_token && state.valid; + const int32_t position_offset = + use_cache_correction ? state.position_offset : 0; + const int32_t current_position = input_positions[seq_id] + position_offset; + const int32_t current_kv_len = specBuilder::calc_kv_len( + input.input_params.attention.host.kv_seq_lens, seq_id, position_offset); + const int32_t expected_kv_len = current_position + 1; + + CHECK_EQ(expected_kv_len, current_kv_len) + << "DFlash decode context position/kv_len mismatch, seq_id=" << seq_id + << ", current_position=" << current_position + << ", current_kv_len=" << current_kv_len; + + token_ids_vec.emplace_back(rewrite_fake_token ? state.token_id + : input_token_id); + positions_vec.emplace_back(current_position); + specBuilder::append_seq_len_by_layout(kv_seq_lens_vec, current_kv_len); + } + + input.token_ids_host = specBuilder::make_cpu_int_tensor(token_ids_vec); + input.positions_host = specBuilder::make_cpu_int_tensor(positions_vec); + input.input_params.attention.host.kv_seq_lens = std::move(kv_seq_lens_vec); + input.device_tensors_ready = false; +} + +void DFlashWorkerImpl::prepare_validate_inputs(const ForwardInput& input, + ForwardInput& validate_input) { + c10::StreamGuard stream_guard = prepare_stream_->set_stream_guard(); + ForwardInput prepared_input = input; + prepared_input.metadata_ready_event.reset(); + SpeculativeWorkerImpl::prepare_validate_inputs(prepared_input, + validate_input); + validate_input.input_params.embedding.input_embedding = torch::Tensor(); + record_metadata_ready_event(*prepare_stream_, validate_input); +} + +void DFlashWorkerImpl::prepare_query_inputs(const ForwardInput& input, + ForwardInput& query_input) { + c10::StreamGuard stream_guard = prepare_stream_->set_stream_guard(); + query_input = input; + query_input.device_tensors_ready = false; + ModelInputParams& input_params = query_input.input_params; + input_params.embedding.input_embedding = torch::Tensor(); + + specBuilder::DecodeBuildBuffers buf; + std::vector selected_idxes; + std::vector q_cu_seq_lens; + build_query_rows(input, + mask_token_id_, + options_.num_speculative_tokens(), + options_.block_size(), + buf, + selected_idxes, + q_cu_seq_lens); + const int32_t query_width = options_.num_speculative_tokens() + 1; + // DFlash emits query_width rows per seq unconditionally, so DP shape + // symmetry holds by construction. Catch scheduler regressions that break + // the invariant (see MTP dp_enabled idle-rank branch). + const int32_t num_sequences_query = input.input_params.meta.num_sequences; + CHECK_EQ(static_cast(buf.out_positions.size()), + num_sequences_query * query_width) + << "DFlash per-seq row count must be uniform (query_width=" << query_width + << ", num_sequences=" << num_sequences_query << ")"; + + specBuilder::set_token_position_tensors(query_input, + buf.out_token_ids, + buf.out_positions, + input.token_ids.options(), + input.positions.options()); + input_params.meta.batch_forward_type = BatchForwardType::CHUNKED_PREFILL; + specBuilder::update_input_params(input_params, + buf, + query_width, + std::move(buf.out_q_seq_lens), + std::move(q_cu_seq_lens), + buf.meta.kv_max_seq_len, + std::move(buf.out_kv_seq_lens), + /*update_block_tables=*/false); + scale_dp_global_token_nums(input_params, query_width); + input_params.attention.rebuild_device_buffer(device_); + + torch::TensorOptions idx_options = + torch::TensorOptions().dtype(torch::kInt).device(device_); + // Pinned-host + async H2D on prepare_stream_ (the file's idiom), so the copy + // overlaps instead of a blocking non-pinned transfer every decode step. + query_input.sampling_params.selected_token_idxes = + safe_to(specBuilder::make_cpu_int_tensor(selected_idxes), + idx_options, + /*non_blocking=*/true); + query_input.sampling_params.sample_idxes = + torch::arange(static_cast(selected_idxes.size()), idx_options); + // Force the draft sampler to emit selected-token probabilities even on the + // greedy path (temperature=0); the rejection sampler needs them to verify + // the block. Without this the greedy sampler skips probs entirely. + query_input.sampling_params.return_probs = true; + repeat_sampling_params(query_input.sampling_params, + options_.num_speculative_tokens()); + query_input.device_tensors_ready = true; +} + +void DFlashWorkerImpl::write_context_kv( + const ForwardInput& input, + const torch::Tensor& context_hidden, + const torch::Tensor& positions_device, + const torch::Tensor& new_cache_slots_device) { + CHECK(context_hidden.defined()) << "DFlash context hidden is undefined."; + CHECK_EQ(context_hidden.dim(), 2) << "DFlash context hidden must be 2D."; + CHECK_EQ(context_hidden.size(1), expected_context_hidden_size_) + << "DFlash context hidden size must be hidden_size * " + << "target_layer_ids.size()."; + + CHECK(context_hidden.device() == device_.unwrap()) + << "DFlash context hidden must already be on the compute device."; + + // Both the target forward that produced context_hidden and this pass run + // on compute_stream_, so no explicit event dance is needed — the stream + // orders them. Model methods below use torch ops on the same stream. + c10::StreamGuard stream_guard = compute_stream_->set_stream_guard(); + +#if defined(USE_NPU) + // PD PUSH: the draft context-KV scattered below is not covered by the target + // push, so wire the draft push here. Overwrite layer_synchronizer with a + // draft-sized one (one event per draft layer); the target's was already + // waited on before this runs, so the overwrite is safe. No-op when + // transfer_kv_infos is empty. NPU-only: the scatter's record_event loop is + // NPU-gated, so a non-NPU push would stall on events that never record. + KVTransferCompletion kv_transfers; + if (options_.kv_cache_transfer_mode() == "PUSH" && + !input.transfer_kv_infos.empty()) { + std::shared_ptr layer_synchronizer = + std::make_shared( + draft_impl_->context_.get_model_args().n_layers()); + const_cast(&(input.input_params)) + ->parallel.layer_synchronizer = layer_synchronizer; + kv_transfers.add(kv_cache_transfer_->push_kv_blocks_async( + input.transfer_kv_infos, + draft_impl_->context_.get_parallel_args(), + layer_synchronizer, + /*is_spec_draft=*/true)); + } +#endif + + ModelOutput scatter_output = + draft_impl_->write_context_kv(context_hidden, + positions_device, + new_cache_slots_device, + input.input_params); + // Model returns an empty ModelOutput (hidden_states undefined) when the + // per-layer NPULayerSynchronizer::record_event fails mid-scatter under + // PD-PUSH transfer. Fail fast instead of silently continuing: subsequent + // layers won't have been scattered and the PD transfer side would block + // forever waiting on the missing per-layer event. + CHECK(scatter_output.hidden_states.defined()) + << "DFlash context-KV scatter failed (layer_synchronizer record_event " + "returned false); PD-PUSH transfer would deadlock."; + + // The context-KV scatter is the last device op of the step and is launched + // no-sync. Under schedule-overlap the scheduler dispatches the next step's + // host prep as soon as this step's host call returns, while this scatter may + // still be in flight. Sync so the next step's draft query reads a fully + // written context-KV cache instead of a partially scattered one. + compute_stream_->synchronize(); + +#if defined(USE_NPU) + // Wait for the draft KV push (if any) so the source draft cache is not + // overwritten by the next step while the transfer is still reading it + // (mirrors step_internal's wait_kv_push()). No-op when no push was issued. + CHECK(kv_transfers.wait()) << "DFlash draft context-KV push failed"; +#endif +} + +void DFlashWorkerImpl::write_target_context_to_cache( + const ForwardInput& input, + const SampleOutput& validate_output) { + const torch::Tensor& accepted_embeddings = validate_output.embeddings; + CHECK(accepted_embeddings.defined()) + << "DFlash validate target embeddings are undefined."; + CHECK_EQ(accepted_embeddings.dim(), 3) + << "DFlash validate target embeddings must be [batch,width,hidden]."; + + torch::Tensor accepted_tokens = validate_output.next_tokens; + CHECK(accepted_tokens.defined()) << "DFlash accepted tokens are undefined."; + if (accepted_tokens.scalar_type() != torch::kInt64) { + accepted_tokens = accepted_tokens.to(torch::kInt64); + } + DCHECK(accepted_tokens.is_contiguous()) + << "DFlash accepted tokens must be contiguous (guaranteed by upstream " + ".to(kCPU) and .to(kInt64) branches); check for stride changes if " + "this fires."; + + CHECK_EQ(accepted_tokens.dim(), 2) + << "DFlash accepted tokens must be [batch,width]."; + const int64_t batch_size = accepted_tokens.size(0); + const int64_t token_width = accepted_tokens.size(1); + CHECK_EQ(accepted_embeddings.size(0), batch_size) + << "DFlash accepted token/embedding batch mismatch."; + CHECK_EQ(accepted_embeddings.size(1), token_width) + << "DFlash accepted token/embedding width mismatch."; + + specBuilder::DecodeBuildBuffers buf; + std::vector accepted_idxes = build_accepted_context_rows( + input, accepted_tokens, options_.block_size(), buf); + torch::TensorOptions host_index_options = torch::TensorOptions() + .dtype(torch::kLong) + .device(torch::kCPU) + .pinned_memory(true); + torch::TensorOptions device_index_options = + torch::TensorOptions() + .dtype(torch::kLong) + .device(accepted_embeddings.device()); + c10::StreamGuard stream_guard = prepare_stream_->set_stream_guard(); + torch::Tensor accepted_index = + safe_to(torch::tensor(accepted_idxes, host_index_options), + device_index_options, + /*non_blocking=*/true); + torch::Tensor flat_embeddings = accepted_embeddings.reshape( + {batch_size * token_width, accepted_embeddings.size(/*dim=*/2)}); + torch::Tensor context_hidden = + flat_embeddings.index_select(/*dim=*/0, accepted_index); + torch::Tensor positions_device = + cpu_int_vec_to_device(buf.out_positions, device_); + torch::Tensor new_cache_slots_device = + cpu_int_vec_to_device(buf.out_new_cache_slots, device_); + // Publish the prepare_stream_ work (index_select producing context_hidden + + // pinned H2D copies for positions/slots) so compute_stream_ waits for it + // before the model reads these tensors. Without this, torch does not + // enforce cross-stream ordering: the write_context_kv scatter could launch + // before index_select finishes, producing corrupt KV cache. The prefill + // caller runs its producer on compute_stream_, so no event is needed there. + StreamEventPtr context_hidden_ready_event = prepare_stream_->record_event(); + if (context_hidden_ready_event != nullptr) { + CHECK(compute_stream_->wait_event(context_hidden_ready_event)) + << "failed to wait DFlash context hidden ready event"; + } else { + prepare_stream_->synchronize(); + } + write_context_kv( + input, context_hidden, positions_device, new_cache_slots_device); + CHECK(!input.input_params.embedding.embedding_ids.empty()) + << "DFlash target context cache write requires embedding ids"; + embedding_cache_->write_target_context( + input.input_params.embedding.embedding_ids, + input.input_params.embedding.request_ids, + validate_output.next_tokens, + validate_output.embeddings, + options_.num_speculative_tokens()); +} + +} // namespace xllm diff --git a/xllm/core/runtime/dflash_worker_impl.h b/xllm/core/runtime/dflash_worker_impl.h new file mode 100644 index 0000000000..d39dad9035 --- /dev/null +++ b/xllm/core/runtime/dflash_worker_impl.h @@ -0,0 +1,129 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "framework/kv_cache/embedding_cache.h" +#include "framework/kv_cache_transfer/kv_cache_transfer.h" +#include "runtime/speculative_worker_impl.h" + +namespace xllm { + +class DFlashWorkerImpl final : public SpeculativeWorkerImpl { + public: + DFlashWorkerImpl(const ParallelArgs& parallel_args, + const torch::Device& device, + const runtime::Options& options); + + ~DFlashWorkerImpl() override = default; + + bool init_model(const std::string& model_weights_path, + int32_t random_seed, + MasterStatus master_status) override; + + std::tuple estimate_kv_cache_capacity() override; + + bool allocate_kv_cache(const KVCacheShape& kv_cache_shape) override; + +#if defined(USE_NPU) || defined(USE_MLU) + bool allocate_kv_cache_with_transfer( + const KVCacheShape& kv_cache_shape) override; +#endif + + ForwardInput update_input_by_last_step_output(ForwardInput& inputs) override; + + protected: + std::optional step_prefill(const ForwardInput& input) override; + + std::optional step_decode(const ForwardInput& input) override; + + std::optional step_empty(const ForwardInput& input) override; + + // Draft produces all speculative tokens of a block in one forward, so its + // output is a single [batch, num_speculative_tokens] block rather than the + // per-step outputs an autoregressive drafter (e.g. MTP) yields. + struct DraftBlock { + torch::Tensor token_ids; + torch::Tensor probs; + // The draft runs no-sync, so execute_no_sync_on_stream stashes its still + // in-flight input in ForwardOutput::retained_input. Anchor it here so the + // draft input outlives run_validate's unconditional compute-stream sync; + // otherwise the buffer frees when run_decode_draft returns and the validate + // stage could reuse memory the draft kernel is still reading. + std::shared_ptr draft_retained_input; + }; + + DraftBlock run_decode_draft(const ForwardInput& input, + ForwardInput& validate_input); + + private: + void fill_validate_input_from_draft_outputs(const DraftBlock& draft_block, + ForwardInput& validate_input, + Stream& compute_stream); + + std::optional run_validate(const ForwardInput& input, + const DraftBlock& draft_block, + ForwardInput& validate_input); + + SampleOutput validate(const SamplingParameters& sampling_params, + const DraftBlock& draft_block, + const ForwardOutput& target_output); + + SampleOutput validate(const SamplingParameters& sampling_params, + const torch::Tensor& draft_token_ids, + const torch::Tensor& draft_probs, + const ForwardOutput& target_output); + + void process_draft_sample_output(SampleOutput& sample_output); + + // Mirrors sampled tokens to rank 0 under schedule-overlap so every rank + // commits the same prefix. No-op for a single rank. + void maybe_broadcast_spec_tokens(torch::Tensor& tokens); + + void update_decode_step_input( + ForwardInput& input, + const std::vector& last_states) const; + + void prepare_validate_inputs(const ForwardInput& input, + ForwardInput& validate_input); + + void prepare_query_inputs(const ForwardInput& input, + ForwardInput& query_input); + + void write_context_kv(const ForwardInput& input, + const torch::Tensor& context_hidden, + const torch::Tensor& positions_device, + const torch::Tensor& new_cache_slots_device); + + void write_target_context_to_cache(const ForwardInput& input, + const SampleOutput& validate_output); + + std::unique_ptr draft_impl_; + std::shared_ptr embedding_cache_; +#if defined(USE_NPU) || defined(USE_MLU) + std::shared_ptr kv_cache_transfer_; +#endif + int32_t mask_token_id_ = -1; + int64_t expected_context_hidden_size_ = 0; +}; + +} // namespace xllm diff --git a/xllm/core/runtime/eagle3_worker_impl.cpp b/xllm/core/runtime/eagle3_worker_impl.cpp index 7419063f1a..0d28ee81e1 100644 --- a/xllm/core/runtime/eagle3_worker_impl.cpp +++ b/xllm/core/runtime/eagle3_worker_impl.cpp @@ -54,7 +54,16 @@ Eagle3WorkerImpl::Eagle3WorkerImpl(const ParallelArgs& parallel_args, eagle3_main_options(options), eagle3_draft_options(options), ::xllm::SpeculativeConfig::get_instance() - .enable_opt_validate_probs()) {} + .enable_opt_validate_probs()) { + // EAGLE-3 drives the draft from the target's captured intermediate-layer + // aux hidden states. Context parallelism only exposes the lm_head-gathered + // final hidden (see llm_worker_impl.cpp), not the aux hidden, so the draft + // would silently receive the wrong tensor. Reject cp_size > 1 until + // aux-hidden plumbing under CP is implemented. + CHECK_LE(parallel_args.cp_size(), 1) + << "EAGLE-3 speculative decoding does not support context parallelism " + "(cp_size > 1)."; +} bool Eagle3WorkerImpl::init_model(const std::string& model_weights_path, int32_t random_seed, diff --git a/xllm/core/runtime/llm_worker_impl.h b/xllm/core/runtime/llm_worker_impl.h index 86cef56b2d..bb9757e012 100644 --- a/xllm/core/runtime/llm_worker_impl.h +++ b/xllm/core/runtime/llm_worker_impl.h @@ -97,6 +97,17 @@ class LLMWorkerImpl : public WorkerImpl { model_->set_word_embedding(embedding); }; + // DFlash-specific delegate: eagerly project target hidden into the draft's + // per-layer KV cache. Runs outside the executor because the pass has no + // attention and its shape doesn't match the decode graph. See CausalLM. + ModelOutput write_context_kv(const torch::Tensor& target_hidden, + const torch::Tensor& positions, + const torch::Tensor& device_cache_slots, + const ModelInputParams& input_params) { + return model_->write_context_kv( + target_hidden, positions, device_cache_slots, kv_caches_, input_params); + } + protected: std::unique_ptr beam_searcher_; }; diff --git a/xllm/core/runtime/mtp_worker_impl.cpp b/xllm/core/runtime/mtp_worker_impl.cpp index 705fd83e00..82812c1893 100644 --- a/xllm/core/runtime/mtp_worker_impl.cpp +++ b/xllm/core/runtime/mtp_worker_impl.cpp @@ -105,14 +105,6 @@ KVCacheEstimateOptions make_kv_cache_estimate_options( return estimate_options; } -torch::Tensor make_cpu_int_tensor(const std::vector& values) { - return torch::tensor(values, - torch::TensorOptions() - .dtype(torch::kInt) - .device(torch::kCPU) - .pinned_memory(true)); -} - void record_metadata_ready_event(Stream& stream, ForwardInput& input) { StreamEventPtr event = stream.record_event(); if (event == nullptr) { @@ -507,26 +499,11 @@ void replace_host_token_placeholders(ForwardInput& input, } } -void set_token_position_tensors(ForwardInput& input, - const std::vector& token_ids, - const std::vector& positions, - const torch::TensorOptions& token_options, - const torch::TensorOptions& position_options) { - input.device_tensors_ready = false; - input.token_ids_host = make_cpu_int_tensor(token_ids); - input.positions_host = make_cpu_int_tensor(positions); - input.token_ids = - safe_to(input.token_ids_host, token_options, /*non_blocking=*/true); - input.positions = - safe_to(input.positions_host, position_options, /*non_blocking=*/true); - input.device_tensors_ready = true; -} - void set_positions_tensor(ForwardInput& input, const std::vector& positions, const torch::TensorOptions& device_options) { input.device_tensors_ready = false; - input.positions_host = make_cpu_int_tensor(positions); + input.positions_host = specBuilder::make_cpu_int_tensor(positions); input.positions = safe_to(input.positions_host, device_options, /*non_blocking=*/true); input.device_tensors_ready = true; @@ -1090,7 +1067,8 @@ void MTPWorkerImpl::prepare_prefill_inputs(const ForwardInput& input, new_token_ids.emplace_back(extra_token_ids[i]); } prefill_input.device_tensors_ready = false; - prefill_input.token_ids_host = make_cpu_int_tensor(new_token_ids); + prefill_input.token_ids_host = + specBuilder::make_cpu_int_tensor(new_token_ids); prefill_input.token_ids = safe_to(prefill_input.token_ids_host, prefill_input.positions.options(), /*non_blocking=*/true); @@ -1464,8 +1442,8 @@ void MTPWorkerImpl::update_decode_step_input( specBuilder::append_seq_len_by_layout(kv_seq_lens_vec, current_kv_len); } - input.token_ids_host = make_cpu_int_tensor(token_ids_vec); - input.positions_host = make_cpu_int_tensor(positions_vec); + input.token_ids_host = specBuilder::make_cpu_int_tensor(token_ids_vec); + input.positions_host = specBuilder::make_cpu_int_tensor(positions_vec); input.input_params.attention.host.kv_seq_lens = std::move(kv_seq_lens_vec); input.device_tensors_ready = false; } @@ -1547,11 +1525,11 @@ void MTPWorkerImpl::prepare_validate_inputs(const ForwardInput& input, CHECK_EQ(buf.out_positions.size(), buf.out_token_ids.size()) << "validate positions/tokens mismatch"; - set_token_position_tensors(validate_input, - buf.out_token_ids, - buf.out_positions, - token_options, - position_options); + specBuilder::set_token_position_tensors(validate_input, + buf.out_token_ids, + buf.out_positions, + token_options, + position_options); if (!use_atb_spec_kernel) { input_params.meta.num_sequences = total_num_val_tokens; input_params.meta.batch_forward_type = BatchForwardType::DECODE; @@ -1740,11 +1718,11 @@ void MTPWorkerImpl::prepare_draft_extend_inputs( CHECK_EQ(expanded_embeddings.size(), buf.out_positions.size()) << "draft extend embeddings/positions mismatch"; - set_token_position_tensors(extend_input, - buf.out_token_ids, - buf.out_positions, - token_options, - position_options); + specBuilder::set_token_position_tensors(extend_input, + buf.out_token_ids, + buf.out_positions, + token_options, + position_options); if (use_chunked_prefill) { input_params.meta.num_sequences = num_sequences; input_params.meta.batch_forward_type = BatchForwardType::CHUNKED_PREFILL; @@ -1818,18 +1796,20 @@ void MTPWorkerImpl::prepare_draft_extend_inputs( params.selected_token_idxes.defined() ? params.selected_token_idxes.options() : torch::dtype(torch::kInt).device(device_); - params.selected_token_idxes = safe_to(make_cpu_int_tensor(selected_row_idx), - idx_options, - /*non_blocking=*/true); + params.selected_token_idxes = + safe_to(specBuilder::make_cpu_int_tensor(selected_row_idx), + idx_options, + /*non_blocking=*/true); if (!params.sample_idxes.defined()) { std::vector sample_idxes_vec; sample_idxes_vec.reserve(num_sequences); for (int32_t i = 0; i < num_sequences; ++i) { sample_idxes_vec.emplace_back(i); } - params.sample_idxes = safe_to(make_cpu_int_tensor(sample_idxes_vec), - idx_options, - /*non_blocking=*/true); + params.sample_idxes = + safe_to(specBuilder::make_cpu_int_tensor(sample_idxes_vec), + idx_options, + /*non_blocking=*/true); } extend_input.device_tensors_ready = true; finish_metadata_prepare(*prepare_stream_, extend_input); diff --git a/xllm/core/runtime/spec_input_builder.cpp b/xllm/core/runtime/spec_input_builder.cpp index 19bc15e071..f180a043e2 100644 --- a/xllm/core/runtime/spec_input_builder.cpp +++ b/xllm/core/runtime/spec_input_builder.cpp @@ -429,6 +429,29 @@ void update_input_params(ModelInputParams& input_params, } } +torch::Tensor make_cpu_int_tensor(const std::vector& values) { + return torch::tensor(values, + torch::TensorOptions() + .dtype(torch::kInt) + .device(torch::kCPU) + .pinned_memory(true)); +} + +void set_token_position_tensors(ForwardInput& input, + const std::vector& token_ids, + const std::vector& positions, + const torch::TensorOptions& token_options, + const torch::TensorOptions& position_options) { + input.device_tensors_ready = false; + input.token_ids_host = make_cpu_int_tensor(token_ids); + input.positions_host = make_cpu_int_tensor(positions); + input.token_ids = + safe_to(input.token_ids_host, token_options, /*non_blocking=*/true); + input.positions = + safe_to(input.positions_host, position_options, /*non_blocking=*/true); + input.device_tensors_ready = true; +} + namespace draftProbs { namespace { @@ -456,6 +479,22 @@ torch::Tensor extract_selected_probs(const torch::Tensor& draft_probs, return torch::Tensor(); } +// Scatter selected-only probs into a dense last-vocab-dim tensor. `ids` and +// `selected_probs` share shape [..., width]; the result is [..., width, vocab], +// with each selected id's prob placed at its vocab slot and zeros elsewhere. +torch::Tensor scatter_selected_to_dense(const torch::Tensor& ids, + const torch::Tensor& selected_probs, + int32_t vocab_size) { + CHECK_GT(vocab_size, 0) << "vocab_size must be > 0 for dense draft probs"; + std::vector dense_shape(ids.sizes().begin(), ids.sizes().end()); + dense_shape.emplace_back(vocab_size); + torch::Tensor dense_probs = + torch::zeros(dense_shape, selected_probs.options()); + dense_probs.scatter_( + /*dim=*/-1, ids.unsqueeze(-1), selected_probs.unsqueeze(-1)); + return dense_probs; +} + } // namespace torch::Tensor compress_for_cache(const torch::Tensor& draft_probs, @@ -495,13 +534,8 @@ std::pair build_validate_tensors( if (enable_opt_validate_probs) { probs_vec.emplace_back(selected_probs); } else { - auto dense_probs = - torch::zeros({batch_size, 1, vocab_size}, selected_probs.options()); - dense_probs.scatter_( - /*dim=*/-1, - draft_token_ids.unsqueeze(-1), - selected_probs.unsqueeze(-1)); - probs_vec.emplace_back(dense_probs); + probs_vec.emplace_back(scatter_selected_to_dense( + draft_token_ids, selected_probs, vocab_size)); } } @@ -513,6 +547,39 @@ std::pair build_validate_tensors( return {draft_token_ids, draft_probs}; } +std::pair build_validate_tensors_from_block( + const torch::Tensor& token_ids_block, + const torch::Tensor& probs_block, + int32_t vocab_size, + bool enable_opt_validate_probs) { + CHECK(token_ids_block.defined()) << "token_ids_block must be defined"; + CHECK(probs_block.defined()) << "probs_block must be defined"; + CHECK_EQ(token_ids_block.dim(), 2) + << "token_ids_block must be [batch, n_speculative_tokens], got " + << token_ids_block.sizes(); + CHECK_EQ(probs_block.dim(), 2) + << "probs_block must be selected-only [batch, n_speculative_tokens], got " + << probs_block.sizes(); + CHECK_EQ(probs_block.size(0), token_ids_block.size(0)) + << "probs/token block batch mismatch"; + CHECK_EQ(probs_block.size(1), token_ids_block.size(1)) + << "probs/token block width mismatch"; + + auto draft_token_ids = token_ids_block.to(torch::kLong); + if (enable_opt_validate_probs) { + // The draft block is already stacked as [batch, n_speculative_tokens]; pass + // the selected-only probs straight to the rejection sampler. + return {draft_token_ids, probs_block}; + } + + // Default path: scatter the selected probs into a dense + // [batch, n_speculative_tokens, vocab_size] tensor, matching MTP's + // build_validate_tensors output so the shared rejection sampler (and its + // fused kernel) always receives dense draft_probs. + return {draft_token_ids, + scatter_selected_to_dense(draft_token_ids, probs_block, vocab_size)}; +} + } // namespace draftProbs } // namespace xllm::specBuilder diff --git a/xllm/core/runtime/spec_input_builder.h b/xllm/core/runtime/spec_input_builder.h index e9f82fc697..9bf3624206 100644 --- a/xllm/core/runtime/spec_input_builder.h +++ b/xllm/core/runtime/spec_input_builder.h @@ -154,6 +154,17 @@ void update_input_params(ModelInputParams& input_params, std::vector kv_seq_lens_vec, bool update_block_tables = false); +// Packs a host int32 vector into a pinned CPU tensor for async H2D staging. +torch::Tensor make_cpu_int_tensor(const std::vector& values); + +// Stages token_ids/positions into both the host and device tensors of `input` +// with async H2D copies, toggling device_tensors_ready around the write. +void set_token_position_tensors(ForwardInput& input, + const std::vector& token_ids, + const std::vector& positions, + const torch::TensorOptions& token_options, + const torch::TensorOptions& position_options); + namespace draftProbs { // Compress draft probs to selected-only format [batch_size] for cache storage. @@ -179,6 +190,23 @@ std::pair build_validate_tensors( bool enable_opt_validate_probs, bool draft_probs_required = true); +// Build validate inputs from an already-stacked draft block. Block-diffusion +// drafters (DFlash) emit the whole block in one forward, so there are no +// per-step tensors to assemble — this avoids the per-step select/view/cat +// round trip of build_validate_tensors. +// token_ids_block / probs_block: [batch_size, n_speculative_tokens] +// Returns draft_token_ids [batch, n_spec] (int64). When +// enable_opt_validate_probs is true the selected-only probs_block is returned +// as-is [batch, n_spec]; otherwise the selected probs are scattered into a +// dense [batch, n_spec, vocab_size] tensor matching MTP's +// build_validate_tensors output, which the default path and the fused +// rejection kernel require. +std::pair build_validate_tensors_from_block( + const torch::Tensor& token_ids_block, + const torch::Tensor& probs_block, + int32_t vocab_size, + bool enable_opt_validate_probs); + } // namespace draftProbs } // namespace specBuilder diff --git a/xllm/core/runtime/speculative_worker_impl.cpp b/xllm/core/runtime/speculative_worker_impl.cpp index e7cf4b1cd5..890e73925e 100644 --- a/xllm/core/runtime/speculative_worker_impl.cpp +++ b/xllm/core/runtime/speculative_worker_impl.cpp @@ -33,26 +33,6 @@ namespace { : tensor_; \ } while (0) -torch::Tensor make_cpu_int_tensor(const std::vector& values) { - return torch::tensor(values, - torch::TensorOptions() - .dtype(torch::kInt) - .device(torch::kCPU) - .pinned_memory(true)); -} - -void set_token_position_tensors(ForwardInput& input, - const std::vector& token_ids, - const std::vector& positions, - const torch::TensorOptions& token_options, - const torch::TensorOptions& position_options) { - input.device_tensors_ready = false; - input.token_ids_host = make_cpu_int_tensor(token_ids); - input.positions_host = make_cpu_int_tensor(positions); - input.token_ids = safe_to(input.token_ids_host, token_options, true); - input.positions = safe_to(input.positions_host, position_options, true); -} - Slice tensor_slice(const torch::Tensor& tensor) { return {tensor.data_ptr(), static_cast(tensor.numel())}; } @@ -158,11 +138,11 @@ ForwardInput SpeculativeWorkerImpl::update_input_by_last_step_output( CHECK_EQ(buf.out_positions.size(), buf.out_token_ids.size()) << "step-update positions/tokens mismatch"; - set_token_position_tensors(new_inputs, - buf.out_token_ids, - buf.out_positions, - inputs.token_ids.options(), - inputs.positions.options()); + specBuilder::set_token_position_tensors(new_inputs, + buf.out_token_ids, + buf.out_positions, + inputs.token_ids.options(), + inputs.positions.options()); // update the input_params input_params.meta.kv_max_seq_len = buf.meta.kv_max_seq_len; input_params.attention.host.kv_seq_lens = std::move(buf.out_kv_seq_lens); @@ -277,11 +257,11 @@ void SpeculativeWorkerImpl::prepare_validate_inputs( CHECK_EQ(buf.out_positions.size(), buf.out_token_ids.size()) << "validate positions/tokens mismatch"; - set_token_position_tensors(validate_input, - buf.out_token_ids, - buf.out_positions, - token_options, - position_options); + specBuilder::set_token_position_tensors(validate_input, + buf.out_token_ids, + buf.out_positions, + token_options, + position_options); // update the input_params if (!::xllm::SpeculativeConfig::get_instance().enable_atb_spec_kernel()) { input_params.meta.num_sequences = total_num_val_tokens; diff --git a/xllm/core/runtime/speculative_worker_impl.h b/xllm/core/runtime/speculative_worker_impl.h index 3c4e4075a7..7ddac65593 100644 --- a/xllm/core/runtime/speculative_worker_impl.h +++ b/xllm/core/runtime/speculative_worker_impl.h @@ -27,7 +27,7 @@ namespace xllm { // Base class for all speculative decoding workers. // Provides common logic: target model management, step dispatch, and // sampling parameter updates. Subclasses implement algorithm-specific -// draft generation and validation (MTP, Eagle3, Suffix, etc.). +// draft generation and validation (MTP, Eagle3, Suffix, DFlash, etc.). class SpeculativeWorkerImpl : public WorkerImpl { public: ~SpeculativeWorkerImpl() override = default; diff --git a/xllm/core/runtime/worker.cpp b/xllm/core/runtime/worker.cpp index 0d9e614567..f5c6c41b75 100644 --- a/xllm/core/runtime/worker.cpp +++ b/xllm/core/runtime/worker.cpp @@ -29,6 +29,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/dflash_worker_impl.h" #include "runtime/dit_worker_impl.h" #include "runtime/eagle3_worker_impl.h" #include "runtime/embed_vlm_worker_impl.h" @@ -51,6 +52,8 @@ Worker::Worker(const ParallelArgs& parallel_args, LOG(INFO) << "Speculative decode is enabled, algorithm: " << algo; if (algo == "Eagle3") { impl_ = new Eagle3WorkerImpl(parallel_args, device, options); + } else if (algo == "DFlash") { + impl_ = new DFlashWorkerImpl(parallel_args, device, options); } else if (algo == "Suffix") { impl_ = new SuffixWorkerImpl(parallel_args, device, options); } else { diff --git a/xllm/core/runtime/worker_impl.cpp b/xllm/core/runtime/worker_impl.cpp index 7a672d0e04..060875fce8 100644 --- a/xllm/core/runtime/worker_impl.cpp +++ b/xllm/core/runtime/worker_impl.cpp @@ -75,6 +75,7 @@ limitations under the License. #if defined(USE_NPU) #include "layers/npu/loader/rolling_weight_buffer.h" #endif +#include "util/json_reader.h" #include "util/tensor_helper.h" #include "util/threadpool.h" #include "util/timer.h" @@ -132,6 +133,25 @@ class ScopedAtenLoadThreads { bool active_ = false; }; +// DFlash draft config lists target_layer_ids as the target-model layer indices +// (0-based) whose output the draft consumes. xLLM's capture hook fires BEFORE +// layer i runs, so capturing layer L's output means putting L+1 in the capture +// set (matched against layer index i in the forward loop). Returns those L+1 +// ids. +std::vector read_dflash_capture_layer_ids( + const std::string& model_weights_path) { + JsonReader reader; + const std::string config_path = model_weights_path + "/config.json"; + CHECK(reader.parse(config_path)) + << "Failed to parse DFlash config: " << config_path; + std::vector capture_layer_ids; + for (int32_t layer_id : reader.value_or>( + "dflash_config.target_layer_ids", std::vector{})) { + capture_layer_ids.emplace_back(layer_id + 1); + } + return capture_layer_ids; +} + void move_tensor_to_device_if_needed(torch::Tensor& tensor, const torch::Device& device) { if (tensor.defined() && tensor.device() != device) { @@ -1449,8 +1469,6 @@ bool WorkerImpl::init_model(const std::string& model_weights_path, auto model_loader = ModelLoader::create(model_weights_path); model_weights_path_ = std::move(model_weights_path); - auto tokenizer = model_loader->tokenizer(); - CHECK(tokenizer != nullptr); auto args = model_loader->model_args(); auto quant_args = model_loader->quant_args(); @@ -1458,21 +1476,47 @@ bool WorkerImpl::init_model(const std::string& model_weights_path, args.embedding_mode(embedding_mode); torch::ScalarType dtype = util::parse_dtype(args.dtype(), device_); - const int64_t tokenizer_vocab_size = tokenizer->vocab_size(); - int64_t model_vocab_size = args.vocab_size(); - // use tokenizer vocab size if model vocab size is not set - if (model_vocab_size <= 0) { - LOG(WARNING) << "Model vocab size is not set, using tokenizer vocab size: " - << tokenizer_vocab_size; - args.vocab_size(tokenizer_vocab_size); - } else if (tokenizer_vocab_size > model_vocab_size) { - LOG(WARNING) << "Unsafe vocab mismatch: tokenizer: " << tokenizer_vocab_size - << ", model: " << model_vocab_size; + // Draft engine is fed token ids and detokenized by the target, so it loads + // no tokenizer of its own (not universal: Eagle3 keeps its own draft vocab; + // see speculative_engine.cpp). + if (!options_.is_draft_engine()) { + std::unique_ptr tokenizer = model_loader->tokenizer(); + CHECK(tokenizer != nullptr); + const int64_t tokenizer_vocab_size = tokenizer->vocab_size(); + int64_t model_vocab_size = args.vocab_size(); + // use tokenizer vocab size if model vocab size is not set + if (model_vocab_size <= 0) { + LOG(WARNING) + << "Model vocab size is not set, using tokenizer vocab size: " + << tokenizer_vocab_size; + args.vocab_size(tokenizer_vocab_size); + } else if (tokenizer_vocab_size > model_vocab_size) { + LOG(WARNING) << "Unsafe vocab mismatch: tokenizer: " + << tokenizer_vocab_size << ", model: " << model_vocab_size; + } } #if defined(USE_NPU) - if (options_.enable_speculative_decode() && - ::xllm::SpeculativeConfig::get_instance().enable_atb_spec_kernel()) { + if (options_.speculative_algorithm() == "DFlash") { + // Both engines capture the same target layers, whose ids live in the draft + // config: the draft engine reads its own weights path, the target engine + // reads --draft_model. The draft engine additionally swaps in the + // DFlashDraftModel body. + std::string draft_config_path; + if (options_.is_draft_engine()) { + LOG(INFO) << "Overriding draft model_type from " << args.model_type() + << " to DFlashDraftModel for DFlash speculative decoding"; + args.model_type("DFlashDraftModel"); + draft_config_path = model_weights_path_; + } else { + CHECK(options_.draft_model_path().has_value()) + << "DFlash requires --draft_model."; + draft_config_path = options_.draft_model_path().value(); + } + args.layers_to_capture(read_dflash_capture_layer_ids(draft_config_path)); + } else if (options_.enable_speculative_decode() && + ::xllm::SpeculativeConfig::get_instance() + .enable_atb_spec_kernel()) { args.num_speculative_tokens(options_.num_speculative_tokens()); } else if (options_.enable_speculative_decode() && options_.num_speculative_tokens() == 0 && @@ -1494,6 +1538,20 @@ bool WorkerImpl::init_model(const std::string& model_weights_path, args.full_attention_interval(1); } } + // Eagle3/DFlash targets capture intermediate-layer aux hidden from the layers + // in layers_to_capture, the model's sole capture signal. Fill the default + // {2, n/2, n-3} for an Eagle3 target whose config omits the list; DFlash + // already filled it from the draft config. The DFlash draft body + // (DFlashDraftModel) consumes context-KV rather than capturing, so exclude + // it. + if (options_.enable_speculative_decode() && + SpeculativeConfig::requires_aux_hidden_capture( + options_.speculative_algorithm()) && + args.model_type() != "DFlashDraftModel" && + args.layers_to_capture().empty()) { + const int32_t num_layers = static_cast(args.n_layers()); + args.layers_to_capture({2, num_layers / 2, num_layers - 3}); + } #else if (options_.enable_speculative_decode()) { args.num_speculative_tokens(options_.num_speculative_tokens()); diff --git a/xllm/models/llm/npu/qwen3.h b/xllm/models/llm/npu/qwen3.h index a0b0c324b4..074a74a157 100644 --- a/xllm/models/llm/npu/qwen3.h +++ b/xllm/models/llm/npu/qwen3.h @@ -18,13 +18,13 @@ limitations under the License. #include #include +#include #include #include #include "core/common/global_flags.h" #include "core/framework/config/kernel_config.h" #include "core/framework/config/scheduler_config.h" -#include "core/framework/config/speculative_config.h" #include "core/framework/model/model_output.h" #include "core/layers/npu/npu_qwen3_decoder_layer_impl.h" #include "llm_model_base.h" @@ -80,17 +80,18 @@ class QWen3ModelImpl : public LlmModelImplBase { blocks_->push_back(block); } - // Eagle3: layer ids to capture (can be read from layers_to_capture in - // config.json) - if (::xllm::SpeculativeConfig::get_instance().speculative_algorithm() == - "Eagle3") { - const auto& layer_ids_from_config = model_args.layers_to_capture(); - if (!layer_ids_from_config.empty()) { - set_eagle3_layers_to_capture( - std::make_optional>(layer_ids_from_config)); - } else { - set_eagle3_layers_to_capture(); - } + // Eagle3/DFlash target captures intermediate-layer aux hidden states to + // drive the draft. The worker fills layers_to_capture (from config.json for + // Eagle3, or the draft config for DFlash) before construction, so a + // non-empty list is the capture signal. The DFlash draft model itself never + // captures. + const bool is_dflash_draft_model = + model_args.model_type() == "DFlashDraftModel"; + const auto& layer_ids_from_config = model_args.layers_to_capture(); + const bool enable_aux_hidden_capture = + !is_dflash_draft_model && !layer_ids_from_config.empty(); + if (enable_aux_hidden_capture) { + set_aux_hidden_capture_layers(layer_ids_from_config); // Pre-allocate aux output buffer [max_tokens_per_batch, hidden_size * // num_captured] const int64_t num_captured = layers_to_capture_set_.size(); @@ -102,20 +103,12 @@ class QWen3ModelImpl : public LlmModelImplBase { } } - void set_eagle3_layers_to_capture( - const std::optional>& layer_ids = std::nullopt) { + void set_aux_hidden_capture_layers(const std::vector& layer_ids) { capture_aux_hidden_states_ = true; layers_to_capture_set_.clear(); - if (!layer_ids.has_value()) { - int32_t num_layers = layers_.size(); - layers_to_capture_set_.insert(2); - layers_to_capture_set_.insert(num_layers / 2); - layers_to_capture_set_.insert(num_layers - 3); - } else { - // Config uses 0-based layer indices, same as default {2, n/2, n-3} - for (int32_t val : layer_ids.value()) { - layers_to_capture_set_.insert(val); - } + // Config uses 0-based layer indices. + for (int32_t val : layer_ids) { + layers_to_capture_set_.insert(val); } LOG(INFO) << "layers_to_capture_set_ size: " << layers_to_capture_set_.size(); @@ -196,19 +189,19 @@ class QWen3ModelImpl : public LlmModelImplBase { // for chunked prefill, generate the attn mask. if (!input_params.meta.batch_forward_type.is_decode()) { if (::xllm::SchedulerConfig::get_instance().enable_chunked_prefill()) { - int max_kv_seq = input_params.meta.kv_max_seq_len; - int num_sequences = input_params.meta.num_sequences; + int32_t max_kv_seq = input_params.meta.kv_max_seq_len; + int32_t num_sequences = input_params.meta.num_sequences; if (num_sequences > 0) { std::vector req_mask_vec; req_mask_vec.reserve(num_sequences); - for (int j = 0; j < num_sequences; j++) { - auto mask = attn_mask_.gen_append_mask( - input_params.attention.host.q_seq_lens[j], - input_params.attention.host.kv_seq_lens[j], - max_kv_seq, - cos_pos.dtype().toScalarType(), - cos_pos.device()); + for (int32_t j = 0; j < num_sequences; j++) { + torch::Tensor mask = + gen_append_attn_mask(input_params.attention.host.q_seq_lens[j], + input_params.attention.host.kv_seq_lens[j], + max_kv_seq, + cos_pos.dtype().toScalarType(), + cos_pos.device()); req_mask_vec.emplace_back(mask); } attn_mask = torch::cat(req_mask_vec, 0); @@ -273,6 +266,8 @@ class QWen3ModelImpl : public LlmModelImplBase { } auto hidden_states = norm_(h, 0); if (capture_aux_hidden_states_) { + CHECK_EQ(capture_idx, static_cast(layers_to_capture_set_.size())) + << "Captured aux hidden layer count mismatch."; torch::Tensor aux_hidden_states = aux_output_buffer_.slice(0, 0, num_tokens); return ModelOutput(hidden_states, torch::Tensor(), aux_hidden_states); @@ -280,6 +275,15 @@ class QWen3ModelImpl : public LlmModelImplBase { return ModelOutput(hidden_states); } + protected: + 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); + } + private: torch::Tensor viusal_pos_mask_; std::unordered_set layers_to_capture_set_; diff --git a/xllm/models/llm/npu/qwen3_dflash.h b/xllm/models/llm/npu/qwen3_dflash.h new file mode 100644 index 0000000000..9802f74624 --- /dev/null +++ b/xllm/models/llm/npu/qwen3_dflash.h @@ -0,0 +1,396 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include +#include + +#include +#include +#include +#include + +#include "core/kernels/ops_api.h" +#include "core/layers/npu/npu_column_parallel_linear_impl.h" +#include "core/layers/npu/npu_rms_norm_impl.h" +#include "core/layers/npu/rotary_embedding.h" +#include "framework/model_loader.h" +#include "models/llm/npu/qwen3.h" +#include "models/model_registry.h" + +namespace xllm::npu::model { + +class DFlashQwen3ModelImpl final : public QWen3ModelImpl { + public: + explicit DFlashQwen3ModelImpl(const ModelContext& context) + : QWen3ModelImpl(context) { + const ModelArgs& model_args = context.get_model_args(); + const ParallelArgs& parallel_args = context.get_parallel_args(); + const int32_t dp_size = parallel_args.dp_size(); + const int32_t cp_size = parallel_args.cp_size(); + CHECK_GT(dp_size, 0) << "DFlash dp_size must be positive."; + CHECK_GT(cp_size, 0) << "DFlash cp_size must be positive."; + CHECK_EQ(parallel_args.world_size() % (dp_size * cp_size), 0) + << "DFlash world_size must be divisible by dp_size * cp_size."; + tp_size_ = parallel_args.world_size() / (dp_size * cp_size); + CHECK_GT(tp_size_, 0) << "DFlash tp_size must be positive."; + tp_rank_ = parallel_args.rank() % tp_size_; + + tensor_options_ = context.get_tensor_options(); + head_dim_ = model_args.head_dim(); + rms_norm_eps_ = model_args.rms_norm_eps(); + CHECK_GT(head_dim_, 0) << "DFlash head_dim must be positive."; + CHECK_GT(model_args.layers_to_capture().size(), 0u) + << "DFlash requires dflash_config.target_layer_ids."; + + fc_ = register_module("fc", layer::NpuColumnParallelLinear(context)); + hidden_norm_ = register_module("hidden_norm", layer::NpuRMSNorm(context)); + rotary_embedding_ = std::make_shared( + head_dim_, + model_args.max_position_embeddings(), + layer::rotary::compute_inv_freq( + head_dim_, model_args.rope_theta(), tensor_options_), + // DFlash draft is fixed Qwen3-dense: NeoX-style, non-interleaved rope. + /*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.")); + 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 std::string& prefix) const override { + fc_->verify_loaded_weights(prefix + "fc."); + hidden_norm_->verify_loaded_weights(prefix + "hidden_norm."); + verify_context_kv_weights(); + for (int32_t i = 0; i < static_cast(layers_.size()); ++i) { + layers_[i]->verify_loaded_weights(prefix + "layers." + std::to_string(i) + + "."); + } + norm_->verify_loaded_weights(prefix + "norm."); + } + + void merge_loaded_weights() override { + fc_->merge_loaded_weights(); + hidden_norm_->merge_loaded_weights(); + build_fused_context_kv_weights(); + for (QWen3DecoderLayer& layer : layers_) { + layer->merge_loaded_weights(); + } + norm_->merge_loaded_weights(); + } + + protected: + 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 { + // Block-diffusion draft attends the full context non-causally: every draft + // position sees the whole block, so all q_len rows share one kv-only mask. + // Do not restore a causal (per-row diagonal) mask here. + torch::Tensor non_causal_mask = attn_mask_.gen_append_mask( + /*q_len=*/1, kv_len, max_kv_len, dtype, device); + return non_causal_mask.repeat({q_len, 1}); + } + + public: + // Projects the target's captured hidden states into the draft's dense KV + // cache: fc -> hidden_norm -> fused per-layer K/V linear -> per-layer k_norm + // + RoPE -> ATB ReshapeAndCache. The projection/norm/rope run here in the + // model layer (as Eagle3 keeps its fc in the model); only the KV scatter is + // delegated to the ATB operator wrapper. Sits outside forward() because it + // has no attention and its shape doesn't match the decode graph. + 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(device_cache_slots.defined()) + << "DFlash context K/V requires device new_cache_slots."; + CHECK_EQ(device_cache_slots.numel(), target_hidden.size(0)) + << "DFlash device cache slot count mismatch."; + + torch::Tensor projected_hidden = fc_(target_hidden, 0); + projected_hidden = hidden_norm_(projected_hidden, 0); + CHECK(fused_kv_weight_.defined()) + << "DFlash fused K/V weight is not initialized."; + CHECK_GT(local_kv_heads_, 0) << "DFlash local KV heads is invalid."; + + const int64_t num_context = projected_hidden.size(0); + torch::Tensor all_kv = + torch::nn::functional::linear(projected_hidden, fused_kv_weight_); + // fused output: [num_context, num_layers, 2(k/v), local_kv_heads, head_dim] + // -> permute to [2(k/v), num_layers, num_context, local_kv_heads, head_dim] + // so the selects below peel k/v first, then per-layer. The single + // contiguous() copy is intentional: the cache writer needs a contiguous + // per-layer k/v view, and doing it once beats re-materializing per layer. + all_kv = + all_kv.view({num_context, num_layers, 2, local_kv_heads_, head_dim_}) + .permute({2, 1, 0, 3, 4}) + .contiguous(); + + torch::Tensor all_key = all_kv.select(/*dim=*/0, /*index=*/0); + torch::Tensor all_value = all_kv.select(/*dim=*/0, /*index=*/1); + + // k_norm has a distinct weight per draft layer, but the normalization + // itself (fp32 RMSNorm over the per-head [head_dim] vector) is uniform, so + // apply it to all layers at once: stack the per-layer weights to broadcast + // over [num_layers, num_context, local_kv_heads, head_dim] and normalize in + // one shot instead of a per-layer loop. Mirrors minimax_rms_norm's batched + // fp32 RMSNorm. + torch::Tensor all_key_normed = apply_k_norm(all_key, k_norm_weight_); + + // RoPE depends only on positions and head_dim, not the layer index, so all + // layers share the same rotation. Flatten [num_layers, num_context, ...] to + // [num_layers * num_context, ...] and repeat positions per layer to apply + // RoPE in a single fused call instead of one per layer. + torch::Tensor flat_key = all_key_normed.reshape( + {num_layers * num_context, local_kv_heads_, head_dim_}); + torch::Tensor repeated_positions = positions.repeat({num_layers}); + flat_key = apply_rope(flat_key, repeated_positions); + all_key_normed = + flat_key.view({num_layers, num_context, local_kv_heads_, head_dim_}); + + // Cache write is inherently per-layer (each layer has its own KV cache). +#if defined(USE_NPU) + const int32_t device_index = all_key_normed.device().index(); +#endif + for (int64_t i = 0; i < num_layers; ++i) { + kernel::ReshapePagedCacheParams scatter_params; + scatter_params.key = all_key_normed[i]; + scatter_params.value = all_value[i]; + scatter_params.k_cache = kv_caches[i].get_k_cache(); + scatter_params.v_cache = kv_caches[i].get_v_cache(); + scatter_params.slot_mapping = device_cache_slots; + kernel::reshape_paged_cache(scatter_params); +#if defined(USE_NPU) + // The standard attention layer op records this event from inside ATB; + // this custom scatter path must record it explicitly so a + // PD-disaggregated PUSH transfer's per-layer synchronizer does not stall + // waiting on it. + if (input_params.parallel.layer_synchronizer != nullptr && + !input_params.parallel.layer_synchronizer->record_event( + i, device_index)) { + return ModelOutput(); + } +#endif + } + return ModelOutput(projected_hidden); + } + + private: + // Batched fp32 RMSNorm over the trailing head_dim. `weight` broadcasts over + // the leading dims (per-layer weights stacked to [num_layers,1,1,head_dim]). + 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); + torch::Tensor normed_key = + key_fp32 * torch::rsqrt(variance + rms_norm_eps_); + // Qwen3 k_norm scales by the raw weight (not gemma-style 1 + weight). + return (normed_key * weight).to(key.scalar_type()); + } + + torch::Tensor apply_rope(const torch::Tensor& key, + const torch::Tensor& positions) const { + CHECK(rotary_embedding_ != nullptr) + << "DFlash rotary embedding is not initialized."; + return std::get<1>(rotary_embedding_->forward(key, key, positions)); + } + + // Collect per-layer k/v projections and concatenate them layer-major + // ([l0_k, l0_v, l1_k, l1_v, ...]) into one fused weight, so a single matmul + // projects all draft layers. Built directly here; there is no merge step. + // Accumulates per-layer k/v/norm weights from one state dict. A sharded HF + // checkpoint delivers draft layers across multiple state dicts, so only the + // layers present in this shard are filled; missing ones stay undefined until + // a later shard. build_fused_context_kv_weights() cat/stacks them once all + // shards are loaded (called from merge_loaded_weights). + 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; // this draft layer lives in another shard + } + CHECK(k_proj.defined()) << "Failed to find DFlash draft layers." << i + << ".self_attn.k_proj.weight."; + CHECK(v_proj.defined()) << "Failed to find DFlash draft layers." << i + << ".self_attn.v_proj.weight."; + CHECK(k_norm.defined()) << "Failed to find DFlash draft layers." << i + << ".self_attn.k_norm.weight."; + CHECK_EQ(k_proj.dim(), 2) << "DFlash k_proj weight must be 2D."; + CHECK_EQ(v_proj.dim(), 2) << "DFlash v_proj weight must be 2D."; + CHECK_EQ(k_proj.size(0), v_proj.size(0)) + << "DFlash k/v projection output size mismatch."; + CHECK_EQ(k_proj.size(1), v_proj.size(1)) + << "DFlash k/v projection weight shape mismatch."; + CHECK_EQ(k_proj.size(0) % head_dim_, 0) + << "DFlash k_proj output size must align to head_dim."; + const int64_t layer_local_kv_heads = k_proj.size(0) / head_dim_; + if (local_kv_heads_ == 0) { + local_kv_heads_ = layer_local_kv_heads; + } else { + CHECK_EQ(local_kv_heads_, layer_local_kv_heads) + << "DFlash local KV heads mismatch."; + } + per_layer_k_proj_[i] = k_proj.to(tensor_options_); + per_layer_v_proj_[i] = v_proj.to(tensor_options_); + // k_norm is applied per draft layer (one weight per captured layer). A + // plain fp32 RMSNorm over the per-head [head_dim] vector matches HF Qwen3 + // semantics. + per_layer_k_norm_[i] = k_norm.to(tensor_options_).to(torch::kFloat32); + } + } + + // Fuses the accumulated per-layer weights into the layer-major fused k/v + // weight ([l0_k, l0_v, l1_k, l1_v, ...]) and stacked k_norm. Called once all + // shards are loaded; every draft layer must be present by now. + void build_fused_context_kv_weights() { + const int32_t num_layers = static_cast(layers_.size()); + std::vector kv_weights; + kv_weights.reserve(num_layers * 2); + std::vector k_norm_weights; + k_norm_weights.reserve(num_layers); + for (int32_t i = 0; i < num_layers; ++i) { + CHECK(per_layer_k_proj_[i].defined()) + << "Missing DFlash draft layers." << i << ".self_attn.k_proj.weight."; + CHECK(per_layer_v_proj_[i].defined()) + << "Missing DFlash draft layers." << i << ".self_attn.v_proj.weight."; + CHECK(per_layer_k_norm_[i].defined()) + << "Missing DFlash draft layers." << i << ".self_attn.k_norm.weight."; + 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(); + // Stack per-layer k_norm weights to [num_layers, 1, 1, head_dim] so a + // single batched RMSNorm broadcasts over [num_layers, num_context, + // kv_heads, head_dim] at once instead of a per-layer loop. + 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(); + } + + // Verifies every draft layer's context-K/V weights arrived across all shards + // before they are fused. Runs before merge_loaded_weights(), so it checks the + // accumulated per-layer tensors rather than the fused result. + 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) + << "DFlash context K/V weights were not accumulated."; + CHECK_GT(local_kv_heads_, 0) << "DFlash local KV heads is invalid."; + for (int32_t i = 0; i < num_layers; ++i) { + CHECK(per_layer_k_proj_[i].defined()) + << "Missing DFlash draft layers." << i << ".self_attn.k_proj.weight."; + CHECK(per_layer_v_proj_[i].defined()) + << "Missing DFlash draft layers." << i << ".self_attn.v_proj.weight."; + CHECK(per_layer_k_norm_[i].defined()) + << "Missing DFlash draft layers." << i << ".self_attn.k_norm.weight."; + } + } + + layer::NpuColumnParallelLinear fc_{nullptr}; + layer::NpuRMSNorm hidden_norm_{nullptr}; + std::shared_ptr rotary_embedding_; + + // Context-K/V write weights. load_context_kv_weights accumulates per-layer + // tensors across shards; build_fused_context_kv_weights fuses them into the + // layer-major fused k/v projection and the stacked per-layer k_norm. + 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; + double rms_norm_eps_ = 1e-6; + int32_t tp_rank_ = 0; + int32_t tp_size_ = 1; + int64_t local_kv_heads_ = 0; +}; +TORCH_MODULE(DFlashQwen3Model); + +class DFlashQwen3ForCausalLMImpl final + : public LlmForCausalLMImplBase { + public: + explicit DFlashQwen3ForCausalLMImpl(const ModelContext& context) + : LlmForCausalLMImplBase(context) {} + + void load_model(std::unique_ptr loader, + std::string prefix = "model.") override { + for (const std::unique_ptr& state_dict : + loader->get_state_dicts()) { + StateDict sub_dict = state_dict->get_dict_with_prefix(prefix); + if (sub_dict.size() == 0) { + sub_dict = state_dict->get_dict_with_prefix(""); + } + model_->load_state_dict(sub_dict); + } + model_->verify_loaded_weights(""); + model_->merge_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); + } +}; +TORCH_MODULE(DFlashQwen3ForCausalLM); + +// Draft config carries model_type="qwen3" and loads via the qwen3_atb args +// loader; worker_impl then overwrites args.model_type to "DFlashDraftModel" +// so this factory is picked to build the draft body. No dedicated DFlash args +// loader is registered. +REGISTER_CAUSAL_MODEL_WITH_VARNAME(dflash_draft_model, + DFlashDraftModel, + DFlashQwen3ForCausalLM); + +} // namespace xllm::npu::model diff --git a/xllm/models/llm/npu/qwen3_moe.h b/xllm/models/llm/npu/qwen3_moe.h index 2cc27e8305..cb9eb1ee13 100644 --- a/xllm/models/llm/npu/qwen3_moe.h +++ b/xllm/models/llm/npu/qwen3_moe.h @@ -17,7 +17,6 @@ limitations under the License. #include "core/framework/config/kernel_config.h" #include "core/framework/config/scheduler_config.h" -#include "core/framework/config/speculative_config.h" #include "core/framework/model/model_output.h" #include "core/framework/model_context.h" #include "core/framework/parallel_state/npu_dp_ep_padding.h" @@ -169,17 +168,18 @@ class Qwen3MoeModelImpl : public torch::nn::Module { indices.push_back(i); } - // Eagle3: layer ids to capture (can be read from layers_to_capture in - // config.json) - if (::xllm::SpeculativeConfig::get_instance().speculative_algorithm() == - "Eagle3") { - const auto& layer_ids_from_config = model_args.layers_to_capture(); - if (!layer_ids_from_config.empty()) { - set_eagle3_layers_to_capture( - std::make_optional>(layer_ids_from_config)); - } else { - set_eagle3_layers_to_capture(); - } + // Eagle3/DFlash target captures intermediate-layer aux hidden states to + // drive the draft. The worker fills layers_to_capture (from config.json for + // Eagle3, or the draft config for DFlash) before construction, so a + // non-empty list is the capture signal. The DFlash draft model itself never + // captures. + const bool is_dflash_draft_model = + model_args.model_type() == "DFlashDraftModel"; + const auto& layer_ids_from_config = model_args.layers_to_capture(); + const bool enable_aux_hidden_capture = + !is_dflash_draft_model && !layer_ids_from_config.empty(); + if (enable_aux_hidden_capture) { + set_aux_hidden_capture_layers(layer_ids_from_config); // Pre-allocate aux output buffer [max_tokens_per_batch, hidden_size * // num_captured] const size_t num_captured = layers_to_capture_set_.size(); @@ -192,20 +192,12 @@ class Qwen3MoeModelImpl : public torch::nn::Module { } } - void set_eagle3_layers_to_capture( - const std::optional>& layer_ids = std::nullopt) { + void set_aux_hidden_capture_layers(const std::vector& layer_ids) { capture_aux_hidden_states_ = true; layers_to_capture_set_.clear(); - if (!layer_ids.has_value()) { - int32_t num_layers = static_cast(layers_.size()); - layers_to_capture_set_.insert(2); - layers_to_capture_set_.insert(num_layers / 2); - layers_to_capture_set_.insert(num_layers - 3); - } else { - // Config uses 0-based layer indices, same as default {2, n/2, n-3} - for (int32_t val : layer_ids.value()) { - layers_to_capture_set_.insert(val); - } + // Config uses 0-based layer indices. + for (int32_t val : layer_ids) { + layers_to_capture_set_.insert(val); } LOG(INFO) << "layers_to_capture_set_ size: " << layers_to_capture_set_.size(); @@ -370,6 +362,8 @@ class Qwen3MoeModelImpl : public torch::nn::Module { h = h + residual.value(); auto hidden_states = norm_(h, 0); if (capture_aux_hidden_states_) { + CHECK_EQ(capture_idx, layers_to_capture_set_.size()) + << "Captured aux hidden layer count mismatch."; torch::Tensor aux_hidden_states = aux_output_buffer_.slice(0, 0, num_tokens); return ModelOutput(hidden_states, torch::Tensor(), aux_hidden_states); diff --git a/xllm/models/models.h b/xllm/models/models.h index 440b31ee23..4205f72034 100644 --- a/xllm/models/models.h +++ b/xllm/models/models.h @@ -42,6 +42,7 @@ limitations under the License. #include "llm/npu/oxygen.h" // IWYU pragma: keep #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_eagle3.h" // IWYU pragma: keep #include "llm/npu/qwen3_moe.h" // IWYU pragma: keep #include "llm/qwen3.h" // IWYU pragma: keep diff --git a/xllm/pybind/args.py b/xllm/pybind/args.py index db1e1fe5a6..7e4767c6a7 100644 --- a/xllm/pybind/args.py +++ b/xllm/pybind/args.py @@ -39,14 +39,14 @@ def __init__(self) -> None: self.parser.add_argument('--draft_devices', type=str, default='npu:0', help='Devices to run the draft model on, e.g. npu:0,npu:1.') self.parser.add_argument('--limit_image_per_prompt', type=int, default=8, help='Maximum number of images per prompt.') self.parser.add_argument('--block_size', type=int, default=128, help='Number of slots per kv cache block. Default is 128.') - self.parser.add_argument('--max_cache_size', type=int, default=0, help='Max gpu memory size for kv cache. Default is 0, which means cache size is caculated by available memory.') + self.parser.add_argument('--max_cache_size', type=int, default=0, help='Max gpu memory size for kv cache. Default is 0, which means cache size is calculated by available memory.') self.parser.add_argument('--max_memory_utilization', type=float, default=0.8, help='The fraction of GPU memory to be used for model inference, including model weights and kv cache.') self.parser.add_argument('--enable_prefix_cache', nargs='?', const=True, default=True, type=_str_to_bool, help='Whether to enable the prefix cache for the block manager.') self.parser.add_argument('--max_tokens_per_batch', type=int, default=10240, help='Max number of tokens per batch.') self.parser.add_argument('--max_seqs_per_batch', type=int, default=1024, help='Max number of sequences per batch.') self.parser.add_argument('--max_tokens_per_chunk_for_prefill', type=int, default=-1, help='Max number of tokens per chunk for request in prefill stage.') self.parser.add_argument('--num_speculative_tokens', type=int, default=0, help='Number of speculative tokens.') - self.parser.add_argument('--speculative_algorithm', type=str, default='MTP', help='Speculative decoding algorithm. Supported options: MTP, Eagle3, Suffix.') + self.parser.add_argument('--speculative_algorithm', type=str, default='MTP', help='Speculative decoding algorithm. Supported options: MTP, Eagle3, Suffix, DFlash.') self.parser.add_argument('--num_request_handling_threads', type=int, default=4, help='Number of handling threads.') self.parser.add_argument('--communication_backend', type=str, default='hccl', help='npu communication backend.') self.parser.add_argument('--rank_tablefile', type=str, default='', help='atb hccl rank table file')