diff --git a/docs/paged_attention_engine.md b/docs/paged_attention_engine.md index 07a3c336e9..3d45a5b176 100644 --- a/docs/paged_attention_engine.md +++ b/docs/paged_attention_engine.md @@ -851,6 +851,13 @@ Graph buffers are allocated once at configured limits and reshaped as static vie Prefill and mixed-token steps use graph id `-1`, which tells the CUDA execution provider to run eagerly. +An Engine-hosted DFlash 2 drafter also forces eager target execution. Its target decoder output is +the packed auxiliary hidden-state tensor named by `model.dflash2.main_aux_hidden_states`; that +variable-size output does not have persistent graph buffers. Engine construction validates its +rank, element type, and static width against the drafter input before allocating cache resources. +The drafter run is synchronous because its packed inputs and outputs are owned by one proposal +call, so `model.dflash2.run_options` cannot disable execution-provider synchronization. + ## Backpressure and fairness Continuous batching does not mean every pending request runs on every step. diff --git a/src/config.cpp b/src/config.cpp index 1a0141ec49..3b4ca6a719 100644 --- a/src/config.cpp +++ b/src/config.cpp @@ -459,6 +459,8 @@ struct DecoderOutputs_Element : JSON::Element { v_.state_update_recurrent_capsule_names = JSON::Get(value); } else if (name == "hidden_states") { v_.hidden_states = JSON::Get(value); + } else if (name == "aux_hidden_states") { + v_.aux_hidden_states = JSON::Get(value); } else if (name == "outputs") { v_.outputs = JSON::Get(value); } else if (name == "lstm_hidden_state") { @@ -1093,6 +1095,137 @@ struct Mtp_Element : JSON::Element { SharedInitializers_Element shared_initializers_{v_.shared_initializers}; }; +struct Dflash2Inputs_Element : JSON::Element { + explicit Dflash2Inputs_Element(Config::Model::Dflash2::Inputs& v) : v_{v} {} + + void OnValue(std::string_view name, JSON::Value value) override { + if (name == "aux_hidden_states") { + v_.aux_hidden_states = JSON::Get(value); + } else if (name == "input_ids") { + v_.input_ids = JSON::Get(value); + } else if (name == "q_row_map") { + v_.q_row_map = JSON::Get(value); + } else if (name == "qkv_row_map") { + v_.qkv_row_map = JSON::Get(value); + } else if (name == "block_row_index") { + v_.block_row_index = JSON::Get(value); + } else if (name == "cumulative_sequence_lengths") { + v_.cumulative_sequence_lengths = JSON::Get(value); + } else if (name == "past_sequence_lengths") { + v_.past_sequence_lengths = JSON::Get(value); + } else if (name == "block_table") { + v_.block_table = JSON::Get(value); + } else if (name == "attention_metadata") { + v_.attention_metadata = JSON::Get(value); + } else if (name == "past_key_names") { + v_.past_key_names = JSON::Get(value); + } else if (name == "past_value_names") { + v_.past_value_names = JSON::Get(value); + } else { + throw JSON::unknown_value_error{}; + } + } + + private: + Config::Model::Dflash2::Inputs& v_; +}; + +struct Dflash2Outputs_Element : JSON::Element { + explicit Dflash2Outputs_Element(Config::Model::Dflash2::Outputs& v) : v_{v} {} + + void OnValue(std::string_view name, JSON::Value value) override { + if (name == "candidate_ids") { + v_.candidate_ids = JSON::Get(value); + } else if (name == "scores") { + v_.scores = JSON::Get(value); + } else if (name == "present_key_names") { + v_.present_key_names = JSON::Get(value); + } else if (name == "present_value_names") { + v_.present_value_names = JSON::Get(value); + } else { + throw JSON::unknown_value_error{}; + } + } + + private: + Config::Model::Dflash2::Outputs& v_; +}; + +struct Dflash2_Element : JSON::Element { + explicit Dflash2_Element(Config::Model::Dflash2& v) : v_{v} {} + + void OnValue(std::string_view name, JSON::Value value) override { + if (name == "filename") { + v_.filename = JSON::Get(value); + } else if (name == "num_hidden_layers") { + v_.num_hidden_layers = SafeDoubleToInt(JSON::Get(value), name); + if (v_.num_hidden_layers <= 0) throw std::out_of_range("num_hidden_layers must be > 0"); + } else if (name == "num_key_value_heads") { + v_.num_key_value_heads = SafeDoubleToInt(JSON::Get(value), name); + if (v_.num_key_value_heads <= 0) throw std::out_of_range("num_key_value_heads must be > 0"); + } else if (name == "head_size") { + v_.head_size = SafeDoubleToInt(JSON::Get(value), name); + if (v_.head_size <= 0) throw std::out_of_range("head_size must be > 0"); + } else if (name == "block_size") { + v_.block_size = SafeDoubleToInt(JSON::Get(value), name); + if (v_.block_size <= 1) throw std::out_of_range("block_size must be > 1"); + } else if (name == "num_draft_tokens") { + v_.num_draft_tokens = SafeDoubleToInt(JSON::Get(value), name); + if (v_.num_draft_tokens <= 0) throw std::out_of_range("num_draft_tokens must be > 0"); + } else if (name == "selector_top_k") { + v_.selector_top_k = SafeDoubleToInt(JSON::Get(value), name); + if (v_.selector_top_k <= 0) throw std::out_of_range("selector_top_k must be > 0"); + } else if (name == "mask_token_id") { + v_.mask_token_id = SafeDoubleToInt(JSON::Get(value), name); + } else if (name == "sliding_window") { + v_.sliding_window = SafeDoubleToInt(JSON::Get(value), name); + } else if (name == "main_aux_hidden_states") { + v_.main_aux_hidden_states = JSON::Get(value); + } else { + throw JSON::unknown_value_error{}; + } + } + + Element& OnObject(std::string_view name) override { + if (name == "session_options") { + v_.session_options = Config::SessionOptions{}; + session_options_ = std::make_unique(*v_.session_options); + return *session_options_; + } + if (name == "run_options") { + v_.run_options = Config::RunOptions{}; + run_options_ = std::make_unique(*v_.run_options); + return *run_options_; + } + if (name == "inputs") { + return inputs_; + } + if (name == "outputs") { + return outputs_; + } + throw JSON::unknown_value_error{}; + } + + Element& OnArray(std::string_view name) override { + if (name == "shared_initializers") { + return shared_initializers_; + } + if (name == "aux_hidden_state_layers") { + return aux_hidden_state_layers_; + } + throw JSON::unknown_value_error{}; + } + + private: + Config::Model::Dflash2& v_; + std::unique_ptr session_options_; + std::unique_ptr run_options_; + Dflash2Inputs_Element inputs_{v_.inputs}; + Dflash2Outputs_Element outputs_{v_.outputs}; + SharedInitializers_Element shared_initializers_{v_.shared_initializers}; + IntArray_Element aux_hidden_state_layers_{v_.aux_hidden_state_layers}; +}; + struct VisionInputs_Element : JSON::Element { explicit VisionInputs_Element(Config::Model::Vision::Inputs& v) : v_{v} {} @@ -1652,6 +1785,9 @@ struct Model_Element : JSON::Element { if (name == "mtp") { return mtp_; } + if (name == "dflash2") { + return dflash2_; + } throw JSON::unknown_value_error{}; } @@ -1668,6 +1804,7 @@ struct Model_Element : JSON::Element { Joiner_Element joiner_{v_.joiner}; VAD_Element vad_{v_.vad}; Mtp_Element mtp_{v_.mtp}; + Dflash2_Element dflash2_{v_.dflash2}; }; // Throws std::runtime_error (rather than std::overflow_error/std::invalid_argument) on failure. diff --git a/src/config.h b/src/config.h index 8763f63c3e..5f821d2609 100644 --- a/src/config.h +++ b/src/config.h @@ -488,6 +488,9 @@ struct Config { std::string state_update_conv_value_names{Defaults::StateUpdateConvValueName}; std::string state_update_recurrent_capsule_names{Defaults::StateUpdateRecurrentCapsuleName}; std::string hidden_states; // Last hidden state output (when exported with include_hidden_states; e.g. fed to the MTP head) + // Residual streams tapped at model.dflash2.aux_hidden_state_layers, concatenated on the + // last axis. Empty unless the model was exported with aux_hidden_state_layers. + std::string aux_hidden_states; // RNNT decoder outputs std::string outputs; @@ -557,6 +560,51 @@ struct Config { } outputs; } mtp; + // DFlash 2 block-drafter metadata. Unlike MTP the drafter is not decoder-shaped: it reads the + // main model's auxiliary hidden states, predicts a whole block of tokens at once, and returns + // a candidate lattice (top-k ids per slot plus the pairwise edge scores) that the engine walks + // greedily. The Engine drives its session directly rather than through a Model. + struct Dflash2 { + std::string filename; // e.g. "dflash2.onnx" + std::optional session_options; + std::optional run_options; + std::vector shared_initializers; + + int num_hidden_layers{}; + int num_key_value_heads{}; + int head_size{}; + int block_size{}; // Query rows per request: the anchor token plus one mask per draft. + int num_draft_tokens{}; // block_size - 1 + int selector_top_k{}; + int mask_token_id{}; + int sliding_window{-1}; + std::vector aux_hidden_state_layers; + + // Name of the main decoder's auxiliary hidden-states output that feeds the drafter. + std::string main_aux_hidden_states{"aux_hidden_states"}; + + struct Inputs { + std::string aux_hidden_states{"aux_hidden_states"}; + std::string input_ids{Defaults::InputIdsName}; + std::string q_row_map{"q_row_map"}; + std::string qkv_row_map{"qkv_row_map"}; + std::string block_row_index{"block_row_index"}; + std::string cumulative_sequence_lengths{Defaults::CumulativeSequenceLengthsName}; + std::string past_sequence_lengths{Defaults::PastSequenceLengthsName}; + std::string block_table{Defaults::BlockTableName}; + std::string attention_metadata{Defaults::AttentionMetadataName}; + std::string past_key_names{Defaults::PastKeyName}; + std::string past_value_names{Defaults::PastValueName}; + } inputs; + + struct Outputs { + std::string candidate_ids{"draft_candidate_ids"}; + std::string scores{"draft_scores"}; + std::string present_key_names{Defaults::PresentKeyName}; + std::string present_value_names{Defaults::PresentValueName}; + } outputs; + } dflash2; + std::optional draft; } model; diff --git a/src/dflash2_drafter.cpp b/src/dflash2_drafter.cpp new file mode 100644 index 0000000000..bd8949f9a2 --- /dev/null +++ b/src/dflash2_drafter.cpp @@ -0,0 +1,531 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "generator/generators.h" +#include "dflash2_drafter.h" +#include "engine/request.h" +#include "models/io/kv_cache.h" + +#include +#include +#include +#include +#include + +namespace Generators { + +namespace { + +// The drafter's packed stream is [context rows of request 0, its block rows, context rows of +// request 1, ...]. PagedAttention derives each row's position from past_sequence_lengths plus its +// offset within the request, so the block rows land at the anchor's position and the ones after it. +struct PackedLayout { + std::vector q_row_map; + std::vector qkv_row_map; + std::vector block_row_index; + std::vector cumulative_sequence_lengths{0}; + std::vector past_sequence_lengths; + int32_t max_query_len{}; + int32_t max_kv_len{}; + int32_t min_kv_len{std::numeric_limits::max()}; +}; + +size_t CheckedAdd(size_t left, size_t right, std::string_view description) { + if (right > std::numeric_limits::max() - left) { + throw std::runtime_error(std::string{description} + " exceeds the supported size range."); + } + return left + right; +} + +size_t CheckedMultiply(size_t left, size_t right, std::string_view description) { + if (left != 0 && right > std::numeric_limits::max() / left) { + throw std::runtime_error(std::string{description} + " exceeds the supported size range."); + } + return left * right; +} + +int32_t CheckedMetadataValue(size_t value, std::string_view description) { + if (value > static_cast(std::numeric_limits::max())) { + throw std::runtime_error( + std::string{description} + " exceeds the int32 attention metadata range."); + } + return static_cast(value); +} + +} // namespace + +std::unique_ptr CreateDflash2Config(const Config& config) { + const auto& dflash2 = config.model.dflash2; + if (dflash2.filename.empty()) { + throw std::runtime_error("model.dflash2.filename is required to create a DFlash 2 drafter."); + } + if (dflash2.num_hidden_layers <= 0 || dflash2.num_key_value_heads <= 0 || dflash2.head_size <= 0 || + dflash2.block_size <= 1 || dflash2.num_draft_tokens <= 0 || dflash2.selector_top_k <= 0) { + throw std::runtime_error("model.dflash2 geometry must be positive and describe a block of >1 token."); + } + if (dflash2.num_draft_tokens != dflash2.block_size - 1) { + throw std::runtime_error("model.dflash2.num_draft_tokens must be block_size - 1."); + } + if (dflash2.run_options) { + for (const auto& [name, value] : *dflash2.run_options) { + if (name == "disable_synchronize_execution_providers" && value == "1") { + throw std::runtime_error( + "model.dflash2.run_options cannot disable execution-provider synchronization."); + } + } + } + + auto projected = std::make_unique(config); + auto& decoder = projected->model.decoder; + decoder.filename = dflash2.filename; + if (dflash2.session_options) { + decoder.session_options = *dflash2.session_options; + } + decoder.run_options = dflash2.run_options; + decoder.shared_initializers = dflash2.shared_initializers; + decoder.num_hidden_layers = dflash2.num_hidden_layers; + decoder.num_key_value_heads = dflash2.num_key_value_heads; + decoder.head_size = dflash2.head_size; + decoder.state_groups.reset(); + decoder.sliding_window.reset(); + return projected; +} + +void ValidateDflash2ModelCompatibility(const Config& config, + const ModelStateMetadata& target_metadata, + const ModelStateMetadata& drafter_metadata) { + const auto& dflash2 = config.model.dflash2; + const auto& target_aux_output = dflash2.main_aux_hidden_states; + if (target_aux_output.empty() || !target_metadata.HasOutput(target_aux_output)) { + throw std::runtime_error( + "model.dflash2.main_aux_hidden_states must name a main-model output."); + } + + const auto& drafter_aux_input = dflash2.inputs.aux_hidden_states; + if (drafter_aux_input.empty() || !drafter_metadata.HasInput(drafter_aux_input)) { + throw std::runtime_error( + "model.dflash2.inputs.aux_hidden_states must name a drafter-model input."); + } + + const auto target_aux_shape = target_metadata.GetOutputShape(target_aux_output); + const auto drafter_aux_shape = drafter_metadata.GetInputShape(drafter_aux_input); + if (target_aux_shape.size() != 2 || target_aux_shape[1] <= 0 || + drafter_aux_shape.size() != 2 || drafter_aux_shape[1] <= 0 || + target_aux_shape[1] != drafter_aux_shape[1]) { + throw std::runtime_error( + "DFlash 2 requires matching 2-D auxiliary hidden-state tensors with a static width."); + } + if (target_metadata.GetOutputDataType(target_aux_output) != + drafter_metadata.GetInputDataType(drafter_aux_input)) { + throw std::runtime_error( + "DFlash 2 requires matching auxiliary hidden-state tensor types."); + } +} + +Dflash2Model::Dflash2Model(std::unique_ptr config, OrtEnv& ort_env) + : Model{std::move(config)} { + session_ = CreateSession(ort_env, config_->model.decoder.filename, session_options_.get()); + session_info_.Add(*session_); +} + +std::unique_ptr Dflash2Model::CreateState(DeviceSpan, const GeneratorParams&) const { + throw std::logic_error("The DFlash 2 drafter is driven by the Engine and has no State."); +} + +size_t Dflash2Drafter::BytesPerBlock(const Config& config, size_t paged_block_size) { + const auto& dflash2 = config.model.dflash2; + if (dflash2.filename.empty()) { + return 0; + } + // K and V, for every layer, for every slot in a block. The cache element width follows the + // drafter body, which is bfloat16. + return size_t{2} * static_cast(dflash2.num_hidden_layers) * paged_block_size * + static_cast(dflash2.num_key_value_heads) * static_cast(dflash2.head_size) * + sizeof(uint16_t); +} + +size_t Dflash2Drafter::PoolBlocks(const Config& config, size_t paged_block_size, + size_t max_batch_size) { + const auto& dflash2 = config.model.dflash2; + if (dflash2.filename.empty()) { + return 0; + } + if (dflash2.sliding_window <= 0) { + throw std::runtime_error( + "The Engine-hosted DFlash 2 drafter requires a sliding window; a full-attention drafter " + "would need a cache as large as the target's."); + } + // The window bounds what a query block can ever read, so a request only needs a ring long enough + // to hold that window plus the block itself, whatever its context length. + const size_t positions = static_cast(dflash2.sliding_window) + + 2 * static_cast(dflash2.block_size); + const size_t ring = (positions + paged_block_size - 1) / paged_block_size + 1; + return std::max(max_batch_size, size_t{1}) * ring; +} + +Dflash2Drafter::Dflash2Drafter(std::shared_ptr model, size_t paged_block_size, + size_t num_blocks) + : model_{std::move(model)}, + config_{model_->config_->model.dflash2}, + paged_block_size_{paged_block_size}, + num_blocks_{num_blocks} { + if (paged_block_size_ == 0 || num_blocks_ == 0) { + throw std::runtime_error("The DFlash 2 drafter needs a non-empty paged cache pool."); + } + if (config_.sliding_window > 0) { + // Positions older than this are masked out of every query row, so they are never ingested and + // the ring may alias them. + context_window_ = static_cast(config_.sliding_window) + + static_cast(config_.block_size); + ring_blocks_ = + (context_window_ + static_cast(config_.block_size) + paged_block_size_ - 1) / + paged_block_size_ + + 1; + } + + const auto& inputs = config_.inputs; + aux_type_ = model_->session_info_.GetInputDataType(inputs.aux_hidden_states); + const auto aux_shape = model_->session_info_.GetInputShape(inputs.aux_hidden_states); + if (aux_shape.size() != 2 || aux_shape[1] <= 0) { + throw std::runtime_error("model.dflash2 expects a 2-D aux_hidden_states input with a static width."); + } + aux_hidden_size_ = static_cast(aux_shape[1]); + + AllocateCache(); + + free_blocks_.resize(num_blocks_); + std::iota(free_blocks_.rbegin(), free_blocks_.rend(), int32_t{0}); + + run_options_ = OrtRunOptions::Create(); + if (model_->config_->model.decoder.run_options) { + for (const auto& entry : *model_->config_->model.decoder.run_options) { + run_options_->AddConfigEntry(entry.first.c_str(), entry.second.c_str()); + } + } +} + +void Dflash2Drafter::AllocateCache() { + const size_t layers = static_cast(config_.num_hidden_layers); + const std::vector shape{static_cast(num_blocks_), + static_cast(paged_block_size_), + config_.num_key_value_heads, + config_.head_size}; + cache_type_ = model_->session_info_.GetInputDataType( + ComposeKeyValueName(config_.inputs.past_key_names, 0)); + for (size_t layer = 0; layer < layers; ++layer) { + for (const auto* pattern : {&config_.inputs.past_key_names, &config_.inputs.past_value_names}) { + cache_input_names_.push_back(ComposeKeyValueName(*pattern, static_cast(layer))); + auto cache = std::make_unique(model_->p_device_kvcache_, cache_type_); + cache->CreateTensor(shape); + cache->GetByteSpan().Zero(); + caches_.push_back(std::move(cache)); + } + cache_output_names_.push_back(ComposeKeyValueName(config_.outputs.present_key_names, static_cast(layer))); + cache_output_names_.push_back(ComposeKeyValueName(config_.outputs.present_value_names, static_cast(layer))); + } +} + +Dflash2Drafter::RequestState& Dflash2Drafter::StateFor(const Request* request) { + return requests_[request]; +} + +void Dflash2Drafter::EnsureBlocks(RequestState& state, size_t positions) { + const size_t needed = + ring_blocks_ != 0 ? ring_blocks_ : (positions + paged_block_size_ - 1) / paged_block_size_; + while (state.blocks.size() < needed) { + if (free_blocks_.empty()) { + throw std::runtime_error("The DFlash 2 drafter's paged cache pool is exhausted."); + } + state.blocks.push_back(free_blocks_.back()); + free_blocks_.pop_back(); + } +} + +void Dflash2Drafter::Release(const Request* request) { + auto entry = requests_.find(request); + if (entry == requests_.end()) { + return; + } + free_blocks_.insert(free_blocks_.end(), entry->second.blocks.begin(), entry->second.blocks.end()); + requests_.erase(entry); +} + +void Dflash2Drafter::Propose(Tensor& aux_hidden_states, std::span feeds, + std::vector>& drafts) { + drafts.assign(feeds.size(), {}); + if (feeds.empty()) { + return; + } + + const size_t block_size = static_cast(config_.block_size); + const size_t num_spec = static_cast(config_.num_draft_tokens); + const size_t top_k = static_cast(config_.selector_top_k); + + // Batch layout. Every feed contributes its context rows so the drafter cache never develops a + // hole; only the feeds that asked also contribute a query block. + std::vector block_feed_indices; + for (size_t i = 0; i < feeds.size(); ++i) { + if (feeds[i].wants_drafts) { + block_feed_indices.push_back(i); + } + } + // The graph reshapes the query rows to [batch, block_size, hidden], so it needs at least one + // block. When nothing is eligible, borrow the first feed's slot and drop its lattice: the block + // rows only write scratch K/V at positions a later step overwrites. + const bool drafts_wanted = !block_feed_indices.empty(); + if (!drafts_wanted) { + block_feed_indices.push_back(0); + } + std::vector block_slot_of_feed(feeds.size(), block_feed_indices.size()); + for (size_t slot = 0; slot < block_feed_indices.size(); ++slot) { + block_slot_of_feed[block_feed_indices[slot]] = slot; + } + + const size_t num_block_rows = + CheckedMultiply(block_feed_indices.size(), block_size, "DFlash 2 block rows"); + CheckedMetadataValue(num_block_rows, "DFlash 2 block rows"); + size_t num_ctx_rows = 0; + for (const auto& feed : feeds) { + num_ctx_rows = CheckedAdd(num_ctx_rows, feed.aux_row_count, "DFlash 2 context rows"); + } + + PackedLayout layout; + const size_t reserved_rows = + CheckedAdd(num_ctx_rows, num_block_rows, "DFlash 2 packed rows"); + CheckedMetadataValue(reserved_rows, "DFlash 2 packed rows"); + layout.q_row_map.reserve(reserved_rows); + layout.qkv_row_map.reserve(reserved_rows); + layout.block_row_index.reserve(num_block_rows); + layout.cumulative_sequence_lengths.reserve(feeds.size() + 1); + layout.past_sequence_lengths.reserve(feeds.size()); + + // Rows this step actually ingests, after a windowed drafter drops the ones its query block can + // never read. + std::vector ingest_begin(feeds.size()); + std::vector ingest_count(feeds.size()); + size_t ctx_row = 0; + size_t max_blocks = 0; + num_ctx_rows = 0; + for (size_t i = 0; i < feeds.size(); ++i) { + const auto& feed = feeds[i]; + auto& state = StateFor(feed.request); + if (state.cached_positions != feed.first_position) { + throw std::logic_error( + "The DFlash 2 drafter's cached context is not contiguous with the target's step."); + } + + size_t dropped = 0; + if (context_window_ != 0 && feed.aux_row_count > context_window_) { + dropped = feed.aux_row_count - context_window_; + } + ingest_begin[i] = CheckedAdd(feed.aux_row_begin, dropped, "DFlash 2 auxiliary row offset"); + ingest_count[i] = feed.aux_row_count - dropped; + num_ctx_rows = CheckedAdd(num_ctx_rows, ingest_count[i], "DFlash 2 context rows"); + } + CheckedMetadataValue( + CheckedAdd(num_ctx_rows, num_block_rows, "DFlash 2 packed rows"), + "DFlash 2 packed rows"); + + for (size_t i = 0; i < feeds.size(); ++i) { + const auto& feed = feeds[i]; + auto& state = requests_[feed.request]; + const size_t first_position = CheckedAdd( + feed.first_position, feed.aux_row_count - ingest_count[i], + "DFlash 2 first position"); + + const size_t slot = block_slot_of_feed[i]; + const bool has_block = slot < block_feed_indices.size(); + const size_t block_rows = has_block ? block_size : 0; + const size_t query_len = + CheckedAdd(ingest_count[i], block_rows, "DFlash 2 query length"); + if (query_len == 0) { + throw std::logic_error("A DFlash 2 feed carries neither context nor a query block."); + } + const size_t total_positions = + CheckedAdd(first_position, query_len, "DFlash 2 KV length"); + EnsureBlocks(state, total_positions); + max_blocks = std::max( + max_blocks, (total_positions - 1) / paged_block_size_ + 1); + + // Context rows borrow the block's first query row; their attention output is discarded. + const size_t first_block_row = + CheckedMultiply(slot, block_size, "DFlash 2 block row index"); + const int32_t borrowed_q_row = has_block + ? CheckedMetadataValue(first_block_row, "DFlash 2 query row") + : 0; + for (size_t row = 0; row < ingest_count[i]; ++row) { + layout.q_row_map.push_back(borrowed_q_row); + layout.qkv_row_map.push_back(CheckedMetadataValue( + CheckedAdd(num_block_rows, CheckedAdd(ctx_row, row, "DFlash 2 context row"), + "DFlash 2 QKV row"), + "DFlash 2 QKV row")); + } + for (size_t row = 0; row < block_rows; ++row) { + layout.block_row_index.push_back( + CheckedMetadataValue(layout.q_row_map.size(), "DFlash 2 packed row")); + const int32_t block_row = CheckedMetadataValue( + CheckedAdd(first_block_row, row, "DFlash 2 block row"), "DFlash 2 block row"); + layout.q_row_map.push_back(block_row); + layout.qkv_row_map.push_back(block_row); + } + ctx_row = CheckedAdd(ctx_row, ingest_count[i], "DFlash 2 context row"); + + layout.cumulative_sequence_lengths.push_back( + CheckedMetadataValue(layout.q_row_map.size(), "DFlash 2 cumulative sequence length")); + layout.past_sequence_lengths.push_back( + CheckedMetadataValue(first_position, "DFlash 2 past sequence length")); + const int32_t query_length = CheckedMetadataValue(query_len, "DFlash 2 query length"); + const int32_t kv_length = CheckedMetadataValue(total_positions, "DFlash 2 KV length"); + layout.max_query_len = std::max(layout.max_query_len, query_length); + layout.max_kv_len = std::max(layout.max_kv_len, kv_length); + layout.min_kv_len = std::min(layout.min_kv_len, kv_length); + + state.cached_positions = CheckedAdd( + feed.first_position, feed.aux_row_count, "DFlash 2 cached positions"); + } + + const size_t num_tokens = layout.q_row_map.size(); + auto device = model_->p_device_inputs_; + + auto make = [&](ONNXTensorElementDataType type, std::vector shape) { + auto tensor = std::make_unique(device, type); + tensor->CreateTensor(shape); + return tensor; + }; + auto fill_int32 = [](Tensor& tensor, const std::vector& values) { + auto span = tensor.GetDeviceSpan(); + std::copy(values.begin(), values.end(), span.CpuSpan().begin()); + span.CopyCpuToDevice(); + }; + + auto packed_aux = make(aux_type_, {static_cast(num_ctx_rows), + static_cast(aux_hidden_size_)}); + const size_t aux_row_bytes = + CheckedMultiply(aux_hidden_size_, Ort::SizeOf(aux_type_), "DFlash 2 auxiliary row bytes"); + auto source_bytes = aux_hidden_states.GetByteSpan(); + auto destination_bytes = packed_aux->GetByteSpan(); + size_t destination_row = 0; + for (size_t i = 0; i < feeds.size(); ++i) { + if (ingest_count[i] == 0) { + continue; + } + destination_bytes.subspan(destination_row * aux_row_bytes, ingest_count[i] * aux_row_bytes) + .CopyFrom(source_bytes.subspan(ingest_begin[i] * aux_row_bytes, + ingest_count[i] * aux_row_bytes)); + destination_row += ingest_count[i]; + } + + auto input_ids = make(Ort::TypeToTensorType, {static_cast(num_block_rows)}); + { + auto span = input_ids->GetDeviceSpan(); + auto cpu = span.CpuSpan(); + for (size_t slot = 0; slot < block_feed_indices.size(); ++slot) { + const auto& feed = feeds[block_feed_indices[slot]]; + cpu[slot * block_size] = feed.wants_drafts ? feed.anchor_token : config_.mask_token_id; + for (size_t row = 1; row < block_size; ++row) { + cpu[slot * block_size + row] = config_.mask_token_id; + } + } + span.CopyCpuToDevice(); + } + + auto q_row_map = make(Ort::TypeToTensorType, {static_cast(num_tokens)}); + fill_int32(*q_row_map, layout.q_row_map); + auto qkv_row_map = make(Ort::TypeToTensorType, {static_cast(num_tokens)}); + fill_int32(*qkv_row_map, layout.qkv_row_map); + auto block_row_index = make(Ort::TypeToTensorType, {static_cast(num_block_rows)}); + fill_int32(*block_row_index, layout.block_row_index); + auto cumulative = make(Ort::TypeToTensorType, {static_cast(feeds.size() + 1)}); + fill_int32(*cumulative, layout.cumulative_sequence_lengths); + auto past_lengths = make(Ort::TypeToTensorType, {static_cast(feeds.size())}); + fill_int32(*past_lengths, layout.past_sequence_lengths); + + auto block_table = make(Ort::TypeToTensorType, + {static_cast(feeds.size()), static_cast(max_blocks)}); + { + auto span = block_table->GetDeviceSpan(); + auto cpu = span.CpuSpan(); + std::fill(cpu.begin(), cpu.end(), int32_t{-1}); + for (size_t i = 0; i < feeds.size(); ++i) { + const auto& blocks = requests_[feeds[i].request].blocks; + if (ring_blocks_ == 0) { + std::copy(blocks.begin(), blocks.end(), cpu.begin() + i * max_blocks); + continue; + } + // A windowed drafter repeats its ring across every column: column j holds the block that + // owns position j * block_size, which is ring[j % ring_blocks]. + for (size_t column = 0; column < max_blocks; ++column) { + cpu[i * max_blocks + column] = blocks[column % blocks.size()]; + } + } + span.CopyCpuToDevice(); + } + + auto metadata = std::make_unique(GetDeviceInterface(DeviceType::CPU), + Ort::TypeToTensorType); + metadata->CreateTensor(std::vector{3}); + { + auto span = metadata->GetDeviceSpan(); + auto cpu = span.CpuSpan(); + cpu[0] = layout.max_query_len; + cpu[1] = layout.max_kv_len; + cpu[2] = layout.min_kv_len; + } + + const size_t batch = block_feed_indices.size(); + auto candidate_ids = make(Ort::TypeToTensorType, + {static_cast(batch), static_cast(num_spec), + static_cast(top_k)}); + auto scores = make(Ort::TypeToTensorType, + {static_cast(batch), static_cast(num_spec), + static_cast(top_k), static_cast(top_k)}); + + std::vector input_names{ + config_.inputs.aux_hidden_states.c_str(), config_.inputs.input_ids.c_str(), + config_.inputs.q_row_map.c_str(), config_.inputs.qkv_row_map.c_str(), + config_.inputs.block_row_index.c_str(), config_.inputs.cumulative_sequence_lengths.c_str(), + config_.inputs.past_sequence_lengths.c_str(), config_.inputs.block_table.c_str(), + config_.inputs.attention_metadata.c_str()}; + std::vector inputs{packed_aux->GetOrtTensor(), input_ids->GetOrtTensor(), + q_row_map->GetOrtTensor(), qkv_row_map->GetOrtTensor(), + block_row_index->GetOrtTensor(), cumulative->GetOrtTensor(), + past_lengths->GetOrtTensor(), block_table->GetOrtTensor(), + metadata->GetOrtTensor()}; + std::vector output_names{config_.outputs.candidate_ids.c_str(), + config_.outputs.scores.c_str()}; + std::vector outputs{candidate_ids->GetOrtTensor(), scores->GetOrtTensor()}; + for (size_t i = 0; i < caches_.size(); ++i) { + input_names.push_back(cache_input_names_[i].c_str()); + inputs.push_back(caches_[i]->GetOrtTensor()); + output_names.push_back(cache_output_names_[i].c_str()); + outputs.push_back(caches_[i]->GetOrtTensor()); + } + + model_->session_->Run(run_options_.get(), input_names.data(), inputs.data(), input_names.size(), + output_names.data(), outputs.data(), output_names.size()); + + if (!drafts_wanted) { + return; + } + + // The spans own the host mirrors these point into, so they must outlive the reads below. + auto candidate_span = candidate_ids->GetDeviceSpan(); + auto scores_span = scores->GetDeviceSpan(); + auto candidate_cpu = candidate_span.CopyDeviceToCpu(); + auto scores_cpu = scores_span.CopyDeviceToCpu(); + for (size_t slot = 0; slot < block_feed_indices.size(); ++slot) { + // Greedy walk of the lattice: slot l's chosen candidate index selects the row of slot l+1's + // score matrix, so the drafted block is one coherent path rather than seven independent argmaxes. + auto& out = drafts[block_feed_indices[slot]]; + out.reserve(num_spec); + size_t previous = 0; + for (size_t step = 0; step < num_spec; ++step) { + const float* row = scores_cpu.data() + ((slot * num_spec + step) * top_k + previous) * top_k; + const size_t best = static_cast(std::max_element(row, row + top_k) - row); + out.push_back(candidate_cpu[(slot * num_spec + step) * top_k + best]); + previous = best; + } + } +} + +} // namespace Generators diff --git a/src/dflash2_drafter.h b/src/dflash2_drafter.h new file mode 100644 index 0000000000..8f6e69d695 --- /dev/null +++ b/src/dflash2_drafter.h @@ -0,0 +1,124 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include +#include + +#include "models/model.h" + +namespace Generators { + +struct Request; + +/** + * @brief Hosts the ``dflash2.onnx`` session. + * + * The drafter graph is not decoder-shaped, so this only borrows Model for its session options, + * shared initializers and device interfaces. It never produces a State. + */ +struct Dflash2Model : Model { + Dflash2Model(std::unique_ptr config, OrtEnv& ort_env); + + std::unique_ptr CreateState(DeviceSpan, const GeneratorParams&) const override; + + std::unique_ptr session_; +}; + +// Decoder-shaped view of model.dflash2, so Model's session-option and shared-initializer plumbing +// applies to the drafter session unchanged. +std::unique_ptr CreateDflash2Config(const Config& config); + +// Validates the auxiliary hidden-state tensor passed from the target to the drafter. +void ValidateDflash2ModelCompatibility(const Config& config, + const ModelStateMetadata& target_metadata, + const ModelStateMetadata& drafter_metadata); + +/** + * @brief Runs the DFlash 2 block drafter for one engine step. + * + * DFlash 2 never re-runs the target's layers. Each step it turns the target's auxiliary hidden + * states into per-layer K/V for the tokens the target just committed, writes them into its own + * paged cache, and then attends a block of `block_size` query rows (the committed token plus + * `num_draft_tokens` mask tokens) over that cache. The block rows produce a lattice -- top-k + * candidates per slot plus pairwise edge scores -- which is walked greedily on the host. + * + * Context rows and query rows travel through one packed PagedAttention call: `qkv_row_map` picks + * each packed row's K/V out of `concat(query, context)` and the context rows' attention output is + * discarded. That is also what fills the cache, so there is no separate cache-store op. + */ +struct Dflash2Drafter { + /** + * @brief One request's contribution to a step. + * + * Rows [aux_row_begin, aux_row_begin + aux_row_count) of the target's packed auxiliary hidden + * states hold positions [first_position, first_position + aux_row_count). Rejected draft rows + * are excluded by the caller, so every row named here is committed context. + */ + struct Feed { + Request* request{}; + size_t aux_row_begin{}; + size_t aux_row_count{}; + size_t first_position{}; + int32_t anchor_token{}; + bool wants_drafts{}; + }; + + Dflash2Drafter(std::shared_ptr model, size_t paged_block_size, size_t num_blocks); + + // Bytes of drafter K/V per paged block, so the main cache pool can budget for it up front. + static size_t BytesPerBlock(const Config& config, size_t paged_block_size); + + // Blocks the pool needs for `max_batch_size` concurrent requests. The drafter is windowed, so a + // request only needs a fixed ring however long its context grows. + static size_t PoolBlocks(const Config& config, size_t paged_block_size, size_t max_batch_size); + + size_t NumDraftTokens() const { return static_cast(config_.num_draft_tokens); } + + /** + * @brief Ingests every feed's context and drafts for the feeds that asked. + * @param aux_hidden_states The target's packed [token_count, aux_hidden_size] output. + * @param drafts Resized to feeds.size(); entry i is empty unless feeds[i].wants_drafts. + */ + void Propose(Tensor& aux_hidden_states, std::span feeds, + std::vector>& drafts); + + // Returns a request's blocks to the pool. Safe for requests the drafter never saw. + void Release(const Request* request); + + private: + struct RequestState { + std::vector blocks; + size_t cached_positions{}; // Positions [0, cached_positions) hold committed context K/V. + }; + + RequestState& StateFor(const Request* request); + // Grows a request's block list so positions [0, positions) are addressable. A windowed drafter + // gets a fixed ring instead, which its block table repeats across every column. + void EnsureBlocks(RequestState& state, size_t positions); + void AllocateCache(); + + std::shared_ptr model_; + const Config::Model::Dflash2& config_; + size_t paged_block_size_{}; + size_t num_blocks_{}; + // Context positions the drafter must keep behind the query block. Zero when it is not windowed, + // in which case the whole sequence stays resident. + size_t context_window_{}; + size_t ring_blocks_{}; + size_t aux_hidden_size_{}; + ONNXTensorElementDataType aux_type_{ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED}; + ONNXTensorElementDataType cache_type_{ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED}; + + std::vector> caches_; // 2 * num_hidden_layers, key then value per layer + std::vector cache_input_names_, cache_output_names_; + std::vector free_blocks_; + std::unordered_map requests_; + + std::unique_ptr run_options_; +}; + +} // namespace Generators diff --git a/src/engine/decoders/decoder.h b/src/engine/decoders/decoder.h index 76b86367d2..8114aaab97 100644 --- a/src/engine/decoders/decoder.h +++ b/src/engine/decoders/decoder.h @@ -56,6 +56,10 @@ struct DecoderIO : ModelIO { virtual Tensor* HiddenStates() const { return nullptr; } + // The step's packed [total_num_tokens, aux_hidden_size] auxiliary hidden states, or null when + // the model was not exported with aux_hidden_state_layers. Row order matches HiddenStates(). + virtual Tensor* AuxHiddenStates() const { return nullptr; } + protected: ScheduledRequests& scheduled_requests_; std::shared_ptr cache_manager_; diff --git a/src/engine/decoders/hybrid_decoder_io.h b/src/engine/decoders/hybrid_decoder_io.h index 0967b1bc76..4f5a728923 100644 --- a/src/engine/decoders/hybrid_decoder_io.h +++ b/src/engine/decoders/hybrid_decoder_io.h @@ -20,6 +20,7 @@ struct HybridDecoderIO : DecoderIO { std::vector> ProcessLogits() override; Tensor* HiddenStates() const override { return varlen_io_.HiddenStates(); } + Tensor* AuxHiddenStates() const override { return varlen_io_.AuxHiddenStates(); } private: // Takes the context by parameter: this IO is moved into ScheduledRequests and outlives the diff --git a/src/engine/decoders/simple_decoder.cpp b/src/engine/decoders/simple_decoder.cpp index 5f86806cb1..78d8d5589d 100644 --- a/src/engine/decoders/simple_decoder.cpp +++ b/src/engine/decoders/simple_decoder.cpp @@ -38,7 +38,8 @@ SimpleDecoder::SimpleDecoder(std::shared_ptr model, if (IsGraphCaptureEnabled(model_->config_->model.decoder.session_options) && cache_manager_->SupportsDynamicBatching() && !has_fixed_state_groups_ && - !has_position_ids) { + !has_position_ids && + model_->config_->model.decoder.outputs.aux_hidden_states.empty()) { graph_buffers_ = std::make_unique(*model_); } } diff --git a/src/engine/decoders/varlen_decoder_io.cpp b/src/engine/decoders/varlen_decoder_io.cpp index 9b15fd3526..14b23bf97d 100644 --- a/src/engine/decoders/varlen_decoder_io.cpp +++ b/src/engine/decoders/varlen_decoder_io.cpp @@ -209,6 +209,7 @@ VarlenDecoderIO::VarlenDecoderIO(std::shared_ptr model, PrepareHiddenStatesInput(model, scheduled_requests); PrepareLogits(model, scheduled_requests); PrepareHiddenStates(model, scheduled_requests); + PrepareAuxHiddenStates(model, scheduled_requests); auto cache = cache_manager->Cache(); for (size_t i = 0; i < cache->input_names_.size(); ++i) { @@ -551,6 +552,34 @@ void VarlenDecoderIO::PrepareHiddenStates(std::shared_ptr mod outputs_.push_back(active_hidden_states_->GetOrtTensor()); } +void VarlenDecoderIO::PrepareAuxHiddenStates(std::shared_ptr model, + ScheduledRequests& scheduled_requests) { + // Only models exported with aux_hidden_state_layers expose this; it is what a DFlash 2 drafter + // turns into its own per-layer K/V. + const auto& dflash2 = model->config_->model.dflash2; + const auto& name = dflash2.filename.empty() + ? model->config_->model.decoder.outputs.aux_hidden_states + : dflash2.main_aux_hidden_states; + if (name.empty() || !model->session_info_.HasOutput(name)) { + return; + } + if (graph_buffers_ != nullptr) { + throw std::logic_error("Auxiliary hidden states are not supported by CUDA graph capture."); + } + + const auto shape = model->session_info_.GetOutputShape(name); + if (shape.size() != 2 || shape[1] <= 0) { + throw std::runtime_error("aux_hidden_states must be 2-D with a static width."); + } + aux_hidden_states_ = std::make_unique(model->p_device_inputs_, + model->session_info_.GetOutputDataType(name)); + aux_hidden_states_->CreateTensor( + std::vector{static_cast(TokenCount(scheduled_requests)), shape[1]}); + + output_names_.push_back(name.c_str()); + outputs_.push_back(aux_hidden_states_->GetOrtTensor()); +} + std::vector> VarlenDecoderIO::ProcessLogits() { // One row per request, plus the extra rows a speculative step needs to verify its drafts: the // request's whole packed range ends with the row that predicts the token after the last draft. diff --git a/src/engine/decoders/varlen_decoder_io.h b/src/engine/decoders/varlen_decoder_io.h index 861a4a9fab..fa146e2a25 100644 --- a/src/engine/decoders/varlen_decoder_io.h +++ b/src/engine/decoders/varlen_decoder_io.h @@ -99,6 +99,10 @@ struct VarlenDecoderIO : DecoderIO { // the model emits one logits row per packed token. Tensor* HiddenStates() const override { return active_hidden_states_; } + // The step's packed [total_num_tokens, aux_hidden_size] auxiliary hidden states, or null when + // the model was not exported with aux_hidden_state_layers. This is what a DFlash 2 drafter reads. + Tensor* AuxHiddenStates() const override { return aux_hidden_states_.get(); } + private: void PrepareInputIds(std::shared_ptr model, ScheduledRequests& scheduled_requests); void PreparePositionIds(std::shared_ptr model, ScheduledRequests& scheduled_requests); @@ -106,6 +110,7 @@ struct VarlenDecoderIO : DecoderIO { void PrepareHiddenStatesInput(std::shared_ptr model, ScheduledRequests& scheduled_requests); void PrepareLogits(std::shared_ptr model, ScheduledRequests& scheduled_requests); void PrepareHiddenStates(std::shared_ptr model, ScheduledRequests& scheduled_requests); + void PrepareAuxHiddenStates(std::shared_ptr model, ScheduledRequests& scheduled_requests); // Number of packed token rows in this step, which is what both the logits and the hidden states // are indexed by. @@ -126,6 +131,7 @@ struct VarlenDecoderIO : DecoderIO { std::unique_ptr logits_fp32_; std::unique_ptr hidden_states_; Tensor* active_hidden_states_{}; + std::unique_ptr aux_hidden_states_; bool logits_are_per_token_{true}; }; diff --git a/src/engine/engine.cpp b/src/engine/engine.cpp index 7ccd532236..2b8649347b 100644 --- a/src/engine/engine.cpp +++ b/src/engine/engine.cpp @@ -100,7 +100,8 @@ Engine::Engine(std::shared_ptr model, EngineDependencies dependencies) model_executor_{std::move(dependencies.model_executor)}, mtp_model_{std::move(dependencies.mtp_model)}, mtp_cache_manager_{std::move(dependencies.mtp_cache_manager)}, - mtp_model_executor_{std::move(dependencies.mtp_model_executor)} { + mtp_model_executor_{std::move(dependencies.mtp_model_executor)}, + dflash2_drafter_{std::move(dependencies.dflash2_drafter)} { // Fail fast on a missing collaborator rather than crashing later on first use. if (!cache_manager_) { throw std::runtime_error("Engine requires a non-null cache manager."); @@ -132,6 +133,40 @@ EngineDependencies Engine::CreateDependencies(std::shared_ptr model) { mtp_bytes_per_block = PagedKeyValueCacheBytesPerBlock(mtp_model); } + std::unique_ptr dflash2_drafter; + if (!model->config_->model.dflash2.filename.empty()) { + if (!model->config_->engine.dynamic_batching) { + throw std::runtime_error("An Engine-hosted DFlash 2 drafter requires dynamic batching."); + } + auto& dflash2 = model->config_->model.dflash2; + auto& target_aux_output = model->config_->model.decoder.outputs.aux_hidden_states; + if (dflash2.main_aux_hidden_states.empty()) { + dflash2.main_aux_hidden_states = target_aux_output; + } + if (dflash2.main_aux_hidden_states.empty()) { + throw std::runtime_error( + "model.dflash2.main_aux_hidden_states must name a main-model output."); + } + target_aux_output = dflash2.main_aux_hidden_states; + if (!model->session_info_.HasOutput(target_aux_output)) { + throw std::runtime_error( + "model.dflash2.main_aux_hidden_states must name a main-model output."); + } + + const auto& batching = *model->config_->engine.dynamic_batching; + const size_t paged_block_size = static_cast(batching.block_size); + auto dflash2_model = std::make_shared( + CreateDflash2Config(*model->config_), GetOrtEnv()); + ValidateDflash2ModelCompatibility( + *model->config_, model->session_info_, dflash2_model->session_info_); + // Built before the main pool so the pool's free-memory measurement already excludes it. The + // drafter is windowed, so its footprint depends on the batch size, not the context length. + dflash2_drafter = std::make_unique( + dflash2_model, paged_block_size, + Dflash2Drafter::PoolBlocks(*model->config_, paged_block_size, + static_cast(batching.max_batch_size))); + } + std::shared_ptr cache_manager = CacheManager::Create(model, mtp_bytes_per_block); auto scheduler = Scheduler::Create(model, cache_manager); @@ -151,7 +186,77 @@ EngineDependencies Engine::CreateDependencies(std::shared_ptr model) { return EngineDependencies{ std::move(cache_manager), std::move(scheduler), std::move(model_executor), - std::move(mtp_model), std::move(mtp_cache_manager), std::move(mtp_model_executor)}; + std::move(mtp_model), std::move(mtp_cache_manager), std::move(mtp_model_executor), + std::move(dflash2_drafter)}; +} + +void Engine::PrepareDflash2Feeds(const StepPlan& plan, + const std::vector& results) { + const size_t max_drafts = std::min(MaxDraftTokensPerStep(), dflash2_drafter_->NumDraftTokens()); + dflash2_feeds_.clear(); + dflash2_feeds_.reserve(plan.requests.size()); + dflash2_draft_widths_.clear(); + dflash2_draft_widths_.reserve(plan.requests.size()); + for (size_t i = 0; i < plan.requests.size(); ++i) { + const auto& entry = plan.requests[i]; + const size_t accepted = entry.request->AcceptedDraftTokenCount(); + if (accepted > entry.draft_token_count) { + throw std::logic_error("DFlash 2 observed more accepted drafts than the target planned."); + } + // Rejected draft rows carry hidden states for tokens that were never committed; drop them. + const size_t valid_rows = + entry.unprocessed_token_count - (entry.draft_token_count - accepted); + // The step's rows start at the processed cursor, which is also what the decoder passes as + // past_sequence_lengths. Deriving it from sequence_length_before instead would be wrong for a + // prefill chunk, whose unprocessed count is the chunk rather than the whole pending suffix. + const size_t first_position = static_cast(entry.request->ProcessedSequenceLength()); + + Dflash2Drafter::Feed feed; + feed.request = entry.request.get(); + feed.aux_row_begin = entry.packed_token_offset; + feed.aux_row_count = valid_rows; + feed.first_position = first_position; + + const auto& search = entry.request->SearchOptions(); + const bool greedy = !search.do_sample || search.top_k == 1 || search.temperature == 0; + // The committed length this step ends at: the accepted prefix plus the token just sampled. + const int64_t length_after_step = static_cast(first_position + valid_rows) + + (results[i].token_appended ? 1 : 0); + const size_t width = std::min( + {max_drafts, static_cast(entry.request->params_->speculative.max_draft_tokens), + length_after_step + 1 < search.max_length + ? static_cast(search.max_length - length_after_step - 1) + : size_t{0}}); + feed.wants_drafts = width > 0 && results[i].token_appended && !results[i].done && greedy && + search.repetition_penalty == 1.0f && search.no_repeat_ngram_size == 0 && + search.min_length <= length_after_step; + feed.anchor_token = results[i].token; + dflash2_feeds_.push_back(feed); + dflash2_draft_widths_.push_back(width); + } +} + +void Engine::PublishDflash2Drafts(const StepPlan& plan, ScheduledRequests& scheduled_requests) { + if (dflash2_feeds_.empty()) { + return; + } + Tensor* aux_hidden_states = scheduled_requests.AuxHiddenStates(); + if (!aux_hidden_states) { + throw std::logic_error("The main decoder did not expose auxiliary hidden states for DFlash 2."); + } + + dflash2_drafter_->Propose(*aux_hidden_states, dflash2_feeds_, dflash2_drafts_); + ++speculative_stats_.draft_forward_passes; + for (size_t i = 0; i < dflash2_feeds_.size(); ++i) { + auto& drafts = dflash2_drafts_[i]; + if (drafts.empty()) { + continue; + } + // The drafter always emits its full block; a request with a narrower budget takes the prefix + // of the same greedy path. + drafts.resize(std::min(drafts.size(), dflash2_draft_widths_[i])); + plan.requests[i].request->SetDraftTokens(drafts); + } } std::unique_ptr Engine::PrepareMtpStep( @@ -671,6 +776,10 @@ void Engine::RemoveRequest(std::shared_ptr request) { mtp_requests_.erase(mtp_it); } + if (dflash2_drafter_) { + dflash2_drafter_->Release(request.get()); + } + ready_requests_.erase( ready_requests_.begin(), ready_requests_.begin() + static_cast(ready_request_index_)); @@ -1108,6 +1217,10 @@ std::shared_ptr Engine::StepDynamic() { CommitMtpStep(*mtp_step); } RecordSpeculativeCommit(step_plan_); + if (dflash2_drafter_ && MaxDraftTokensPerStep() > 0) { + // Reads the accepted-draft counts, which CommitStep clears below. + PrepareDflash2Feeds(step_plan_, step_results_); + } for (size_t i = 0; i < step_plan_.requests.size(); ++i) { step_plan_.requests[i].request->CommitStep( step_plan_.requests[i], step_results_[i]); @@ -1115,6 +1228,9 @@ std::shared_ptr Engine::StepDynamic() { if (mtp_step) { PublishMtpDrafts(*mtp_step); } + if (dflash2_drafter_ && MaxDraftTokensPerStep() > 0) { + PublishDflash2Drafts(step_plan_, scheduled_requests); + } } catch (...) { MarkUnhealthyAndThrow( StepOutcomeKind::ExecutionContractFailure, diff --git a/src/engine/engine.h b/src/engine/engine.h index 09d3ac6106..4b4c7ca67c 100644 --- a/src/engine/engine.h +++ b/src/engine/engine.h @@ -7,6 +7,7 @@ #include "model_executor.h" #include "scheduler.h" #include "../decoding/speculative_stats.h" +#include "../dflash2_drafter.h" /** * @file engine.h @@ -41,6 +42,7 @@ struct EngineDependencies { std::shared_ptr mtp_model; std::shared_ptr mtp_cache_manager; std::unique_ptr mtp_model_executor; + std::unique_ptr dflash2_drafter; }; struct EngineTransactionMetrics { @@ -153,6 +155,10 @@ struct Engine : std::enable_shared_from_this, void RollbackMtpStep(MtpStep& step); void CommitMtpStep(MtpStep& step); void PublishMtpDrafts(MtpStep& step); + // Runs the DFlash 2 drafter on a committed step and attaches its block to each request. The + // feeds are captured before Request::CommitStep clears the accepted-draft counts they depend on. + void PrepareDflash2Feeds(const StepPlan& plan, const std::vector& results); + void PublishDflash2Drafts(const StepPlan& plan, ScheduledRequests& scheduled_requests); void RecordSpeculativeCommit(const StepPlan& plan) noexcept; void ValidateRequestCanContinue(const std::shared_ptr& request) const; [[noreturn]] void HandleContinuationRestoreFailure( @@ -175,6 +181,11 @@ struct Engine : std::enable_shared_from_this, std::shared_ptr mtp_cache_manager_; std::unique_ptr mtp_model_executor_; std::unordered_map> mtp_requests_; + // Present only when model.dflash2 names a block drafter. Owns its own session and paged cache. + std::unique_ptr dflash2_drafter_; + std::vector dflash2_feeds_; + std::vector> dflash2_drafts_; + std::vector dflash2_draft_widths_; DeviceSpan mtp_device_drafts_; DeviceSpan mtp_device_chain_inputs_; EngineHealth health_{EngineHealth::Healthy}; diff --git a/src/engine/scheduled_requests.cpp b/src/engine/scheduled_requests.cpp index 98c04d98e6..b114169000 100644 --- a/src/engine/scheduled_requests.cpp +++ b/src/engine/scheduled_requests.cpp @@ -292,6 +292,10 @@ Tensor* ScheduledRequests::HiddenStates() const { return decoder_state_ ? decoder_state_->HiddenStates() : nullptr; } +Tensor* ScheduledRequests::AuxHiddenStates() const { + return decoder_state_ ? decoder_state_->AuxHiddenStates() : nullptr; +} + std::vector> ScheduledRequests::SelectSampledRows( std::vector>& verify_rows, std::vector>& selected_tokens, diff --git a/src/engine/scheduled_requests.h b/src/engine/scheduled_requests.h index ff7416b8c4..99d457f012 100644 --- a/src/engine/scheduled_requests.h +++ b/src/engine/scheduled_requests.h @@ -71,6 +71,7 @@ struct ScheduledRequests { std::vector> ProcessLogits(); Tensor* HiddenStates() const; + Tensor* AuxHiddenStates() const; void GenerateNextTokens(); void BeginTransaction(); diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index 8c378e37bb..4a5d5b4855 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -13,6 +13,7 @@ file(GLOB test_srcs CONFIGURE_DEPENDS ) # Keep tests with dedicated white-box executables out of the public-API unit_tests target. +list(REMOVE_ITEM test_srcs "${CMAKE_CURRENT_SOURCE_DIR}/cpp/dflash2_config_test.cpp") list(REMOVE_ITEM test_srcs "${CMAKE_CURRENT_SOURCE_DIR}/cpp/reinit_tests.cpp") list(REMOVE_ITEM test_srcs "${CMAKE_CURRENT_SOURCE_DIR}/cpp/search_checkpoint_tests.cpp") @@ -166,6 +167,7 @@ file(GLOB engine_unit_test_srcs CONFIGURE_DEPENDS "${CMAKE_CURRENT_SOURCE_DIR}/cpp/engine/*.cpp" ) list(APPEND engine_unit_test_srcs + "${CMAKE_CURRENT_SOURCE_DIR}/cpp/dflash2_config_test.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/cpp/search_checkpoint_tests.cpp" ) add_executable(engine_unit_tests ${engine_unit_test_srcs}) diff --git a/test/cpp/dflash2_config_test.cpp b/test/cpp/dflash2_config_test.cpp new file mode 100644 index 0000000000..1637928eca --- /dev/null +++ b/test/cpp/dflash2_config_test.cpp @@ -0,0 +1,200 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include + +#include +#include +#include + +#include "dflash2_drafter.h" + +namespace Generators::test { +namespace { + +Config MakeDflash2Config() { + Config config; + auto& decoder = config.model.decoder; + decoder.filename = "target.onnx"; + decoder.sliding_window = Config::Model::Decoder::SlidingWindow{4096}; + decoder.state_groups.emplace(); + decoder.state_groups->push_back(Config::Model::Decoder::StateGroup{ + Config::Model::Decoder::StateGroupKind::Fixed}); + + auto& dflash2 = config.model.dflash2; + dflash2.filename = "dflash2.onnx"; + dflash2.num_hidden_layers = 3; + dflash2.num_key_value_heads = 2; + dflash2.head_size = 8; + dflash2.block_size = 4; + dflash2.num_draft_tokens = 3; + dflash2.selector_top_k = 2; + dflash2.sliding_window = 17; + return config; +} + +struct TensorMetadata { + ONNXTensorElementDataType data_type; + std::vector shape; +}; + +class FakeModelStateMetadata final : public ModelStateMetadata { + public: + void AddInput(std::string name, ONNXTensorElementDataType data_type, + std::vector shape) { + inputs_.insert_or_assign( + std::move(name), TensorMetadata{data_type, std::move(shape)}); + } + + void AddOutput(std::string name, ONNXTensorElementDataType data_type, + std::vector shape) { + outputs_.insert_or_assign( + std::move(name), TensorMetadata{data_type, std::move(shape)}); + } + + bool HasInput(const std::string& name) const override { return inputs_.contains(name); } + bool HasOutput(const std::string& name) const override { return outputs_.contains(name); } + + ONNXTensorElementDataType GetInputDataType(const std::string& name) const override { + return inputs_.at(name).data_type; + } + + ONNXTensorElementDataType GetOutputDataType(const std::string& name) const override { + return outputs_.at(name).data_type; + } + + std::vector GetInputShape(const std::string& name) const override { + return inputs_.at(name).shape; + } + + std::vector GetOutputShape(const std::string& name) const override { + return outputs_.at(name).shape; + } + + private: + std::unordered_map inputs_; + std::unordered_map outputs_; +}; + +std::pair MakeCompatibleMetadata() { + FakeModelStateMetadata target; + target.AddOutput("aux_hidden_states", ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16, {-1, 64}); + FakeModelStateMetadata drafter; + drafter.AddInput("aux_hidden_states", ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16, {-1, 64}); + return {std::move(target), std::move(drafter)}; +} + +} // namespace + +TEST(Dflash2ConfigTest, RequiresDrafterFilename) { + auto config = MakeDflash2Config(); + config.model.dflash2.filename.clear(); + EXPECT_THROW(CreateDflash2Config(config), std::runtime_error); +} + +TEST(Dflash2ConfigTest, RequiresCompleteGeometry) { + auto config = MakeDflash2Config(); + config.model.dflash2.num_key_value_heads = 0; + EXPECT_THROW(CreateDflash2Config(config), std::runtime_error); +} + +TEST(Dflash2ConfigTest, RequiresPositiveSelectorTopK) { + auto config = MakeDflash2Config(); + config.model.dflash2.selector_top_k = 0; + EXPECT_THROW(CreateDflash2Config(config), std::runtime_error); +} + +TEST(Dflash2ConfigTest, RejectsAsynchronousExecution) { + auto config = MakeDflash2Config(); + config.model.dflash2.run_options = Config::RunOptions{ + {"disable_synchronize_execution_providers", "1"}}; + EXPECT_THROW(CreateDflash2Config(config), std::runtime_error); +} + +TEST(Dflash2ConfigTest, AcceptsCompatibleAuxiliaryHiddenStates) { + const auto config = MakeDflash2Config(); + const auto [target, drafter] = MakeCompatibleMetadata(); + EXPECT_NO_THROW(ValidateDflash2ModelCompatibility(config, target, drafter)); +} + +TEST(Dflash2ConfigTest, UsesConfiguredTargetOutput) { + auto config = MakeDflash2Config(); + config.model.dflash2.main_aux_hidden_states = "custom_aux_hidden_states"; + FakeModelStateMetadata target; + target.AddOutput("custom_aux_hidden_states", ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16, {-1, 64}); + FakeModelStateMetadata drafter; + drafter.AddInput("aux_hidden_states", ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16, {-1, 64}); + EXPECT_NO_THROW(ValidateDflash2ModelCompatibility(config, target, drafter)); +} + +TEST(Dflash2ConfigTest, RequiresConfiguredTargetOutput) { + const auto config = MakeDflash2Config(); + FakeModelStateMetadata target; + FakeModelStateMetadata drafter; + drafter.AddInput("aux_hidden_states", ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16, {-1, 64}); + EXPECT_THROW(ValidateDflash2ModelCompatibility(config, target, drafter), std::runtime_error); +} + +TEST(Dflash2ConfigTest, RequiresConfiguredDrafterInput) { + const auto config = MakeDflash2Config(); + FakeModelStateMetadata target; + target.AddOutput("aux_hidden_states", ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16, {-1, 64}); + const FakeModelStateMetadata drafter; + EXPECT_THROW(ValidateDflash2ModelCompatibility(config, target, drafter), std::runtime_error); +} + +TEST(Dflash2ConfigTest, RequiresTwoDimensionalAuxiliaryTensors) { + const auto config = MakeDflash2Config(); + auto [target, drafter] = MakeCompatibleMetadata(); + target.AddOutput("aux_hidden_states", ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16, {-1, 4, 16}); + EXPECT_THROW(ValidateDflash2ModelCompatibility(config, target, drafter), std::runtime_error); + + std::tie(target, drafter) = MakeCompatibleMetadata(); + drafter.AddInput("aux_hidden_states", ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16, {-1, 4, 16}); + EXPECT_THROW(ValidateDflash2ModelCompatibility(config, target, drafter), std::runtime_error); +} + +TEST(Dflash2ConfigTest, RequiresMatchingAuxiliaryWidth) { + const auto config = MakeDflash2Config(); + auto [target, drafter] = MakeCompatibleMetadata(); + drafter.AddInput("aux_hidden_states", ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16, {-1, 32}); + EXPECT_THROW(ValidateDflash2ModelCompatibility(config, target, drafter), std::runtime_error); +} + +TEST(Dflash2ConfigTest, RequiresMatchingAuxiliaryType) { + const auto config = MakeDflash2Config(); + auto [target, drafter] = MakeCompatibleMetadata(); + drafter.AddInput("aux_hidden_states", ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT, {-1, 64}); + EXPECT_THROW(ValidateDflash2ModelCompatibility(config, target, drafter), std::runtime_error); +} + +TEST(Dflash2ConfigTest, RequiresOneDraftPerNonAnchorBlockRow) { + auto config = MakeDflash2Config(); + config.model.dflash2.num_draft_tokens = 2; + EXPECT_THROW(CreateDflash2Config(config), std::runtime_error); +} + +TEST(Dflash2ConfigTest, ProjectsDrafterWithoutTargetState) { + const auto projected = CreateDflash2Config(MakeDflash2Config()); + const auto& decoder = projected->model.decoder; + EXPECT_EQ(decoder.filename, "dflash2.onnx"); + EXPECT_EQ(decoder.num_hidden_layers, 3); + EXPECT_EQ(decoder.num_key_value_heads, 2); + EXPECT_EQ(decoder.head_size, 8); + EXPECT_FALSE(decoder.sliding_window.has_value()); + EXPECT_FALSE(decoder.state_groups.has_value()); +} + +TEST(Dflash2ConfigTest, AccountsForWindowedPagedCache) { + const auto config = MakeDflash2Config(); + EXPECT_EQ(Dflash2Drafter::BytesPerBlock(config, 16), 3072u); + EXPECT_EQ(Dflash2Drafter::PoolBlocks(config, 8, 3), 15u); +} + +TEST(Dflash2ConfigTest, RejectsUnboundedCachePool) { + auto config = MakeDflash2Config(); + config.model.dflash2.sliding_window = 0; + EXPECT_THROW(Dflash2Drafter::PoolBlocks(config, 8, 3), std::runtime_error); +} + +} // namespace Generators::test