diff --git a/docs/paged_attention_engine.md b/docs/paged_attention_engine.md index 6ebbb0656a..82dab79586 100644 --- a/docs/paged_attention_engine.md +++ b/docs/paged_attention_engine.md @@ -2321,6 +2321,62 @@ tombstone. Changes to shared types such as `Request`, `ScheduledRequests`, `ModelExecutor`, or `SimpleDecoder` should be checked against both paths. This document should be updated only where behavior is shared or where the dynamic path changes. +### Static single-slot backend that selects its own tokens + +The `qwen38_flash_graph` model type (Qwen3.8-Flash CUDA-graph decode with a built-in MTP(1) loop) runs on the static path +through a `ModelExecutor` that chooses its own tokens, instead of attaching decoder state with logits. It is a single-slot +backend: fixed in-place KV, a recurrent-state bank ring and a device-side draft -> verify -> commit loop that cannot be +paged, batched or cropped to an arbitrary token. `CreateEngineBackend` (`src/engine/backend_factory.h`) and +`CreateQwen38GraphDependencies` (`src/engine/qwen38_graph_executor.h`) assemble it: + +- `FixedSlotCacheManager` (`Qwen38GraphCacheManager`, `src/engine/fixed_slot_scheduler.h`) holds one row. `FixedSlotScheduler` + queues further Requests and recycles the slot when the resident turn is complete. `RewindToStartOfTurn` releases the slot; a + continued turn (`BeginTurn` on the resident Request) and the replay after a rewind re-ingest the whole retained history + (through the same chunked prefill as a new prompt), because the look-ahead cycle and the bank ring cannot be rolled back. + There is no prefix cache. +- `Qwen38GraphExecutor::SelectsOwnTokens()` is true. One Engine step is one prompt chunk (the size the device asks for, + `Qwen38DecodeDevice::PrefillChunkTokens()`: `model.qwen38_graph.prefill.chunk` when set, otherwise 8192 tokens with the + dense prefill (`prefill.enabled`) and 2048 per step through the decode graphs, see `EffectiveQwen38PrefillChunkTokens`; + the decode-graph path needs an even chunk; a step that only advances the prompt delivers no event, so `CancelRequest` + stays responsive; the dense prefill state continues across the chunks and is handed over to the decode state in the step that + ingests the last one) or one whole cycle (with `fold.enabled`: one CUDA-graph launch). The executor tells the device the + prompt length of the turn (`Qwen38DecodeDevice::ExpectPrefillTokens`, after `Reset()` and before the first chunk); the + device hands it to the ORT QMoE kernels with every chunk of that turn and drops it at `Reset()`/`BeginDecode()`. The + executor itself stores the expected length on the device instance rather than writing environment variables directly; the + production device then publishes `ORT_QMOE_PREFILL_TOTAL_TOKENS` through `Qwen38OrtEnvScope` around each ORT `Run` under a + process-wide mutex (isolating participating Qwen runs from each other, though unrelated concurrent ORT sessions in the same + process can still observe the process environment; see `src/models/qwen38_ort_env.h`). The + executor declares its step plan via `ScheduledRequests::BeginBackendStep` and stages owned `BackendRowResult` records + (`PrefillProgress` and/or `SelectedTokenBlock`) via `ScheduledRequests::StageBackendRow`; `GenerateNextTokens` validates the + batch before mutating any Request, commits candidates through `Request::CommitExternalTokens` (observing stop strings, EOS, + turn limit, and context limit in order), and records `TokenCommitAck` counts (consumed, appended, and staged for events) + settled through `TokenProducingSession` before events are published. Stop strings are therefore accepted on this static path; + every other static executor still rejects them. +- The next cycle is enqueued before a step returns, so the GPU runs it while the Engine delivers the previous tokens. A stop + string can only be known after delivery, so one cycle may be wasted; the next request, a rewind and Engine teardown wait for + it with `Qwen38DecodeDevice::Drain()`. The record of such a cycle is never consumed, so its status is read from the device + instead: `CheckDecodeStatus()` right after `Drain()` reports the error bits accumulated by every cycle since the last + `Reset()`, also when a record was left unconsumed. The executor always runs `Drain()` and then `CheckDecodeStatus()` before + `Reset()` (which clears the bits): a cancellation or stop-string match that leaves a cycle behind therefore still fails the + next request or the rewind when that cycle went wrong. At Engine teardown the executor drains the device and reads its status + too, but nobody is left to report an error to, so a failure poisons the device context instead (the next Engine over the + model fails). A fatal execution failure poisons the mutable device context across Engine leases until the model is reloaded. +- `ModelExecutor::ValidateTurn` and `BackendCapabilities` reject, before any Request mutation, what the device loop cannot do: + sampling, repetition penalty, no-repeat n-grams, `min_generated_tokens`, seeds and guidance. `MaxDraftTokensPerStep()` is + zero (the caller cannot attach drafts). `GetCapabilities()` reports one row and `max_request_length` = KV capacity - 4, + clamped to `model.context_length` when that is configured, so it is the smaller of the two. +- `GetSpeculativeStats()` includes the executor's counters (`ModelExecutor::AddSpeculativeStats`): one round per cycle, one + or more proposed drafts per round (the fold verifies 1..3), accepted drafts, bonus and correction tokens, the + acceptance-length histogram. The counters describe what the Engine consumed, as `Engine::RecordSpeculativeCommit` does for + sampled MTP: a turn that ends inside a cycle (turn or session limit, stop string, EOS) only accepts the drafts it delivered, + `draft_tokens_evaluated` counts the draft positions examined up to the terminal token (every accepted draft consumed, an EOS + draft included, plus the rejected position only when the turn got as far as the correction token), so `acceptance_rate` is not + diluted by drafts after the end of the turn. Rounds are classified like `Engine::RecordSpeculativeCommit` does, by the number + of accepted drafts that were delivered: none is a zero-accept round, all of the proposed ones a full-accept round and anything + in between a partial-accept round (`partial_accept_rounds`, counted whether or not a correction token was delivered). +- The model's device state is leased to one Engine at a time (`Qwen38GraphModel::AcquireDevice`); a second Engine over the same + model fails with a clear error until the first one is destroyed. The model has no Generator API support. + ## Experimental low-level public API shape The language bindings currently expose the same basic low-level loop. Production hosts are expected to wrap this surface or use its replacement rather than expose it as their stable API: diff --git a/src/engine/backend_factory.cpp b/src/engine/backend_factory.cpp new file mode 100644 index 0000000000..e54380b79b --- /dev/null +++ b/src/engine/backend_factory.cpp @@ -0,0 +1,409 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "backend_factory.h" +#include "../logging.h" +#include "../models/qwen38_graph.h" +#include "qwen38_graph_executor.h" +#include "decoders/varlen_decoder_io.h" + +#include +#include +#include +#include +#include + +namespace Generators { +namespace { + +size_t CheckedMtpAdd(size_t left, size_t right) { + if (right > std::numeric_limits::max() - left) + throw std::runtime_error("MTP scratch memory bytes overflow size_t."); + return left + right; +} + +size_t CheckedMtpMultiply(size_t left, size_t right) { + if (left != 0 && right > std::numeric_limits::max() / left) + throw std::runtime_error("MTP scratch memory bytes overflow size_t."); + return left * right; +} + +size_t IndexShareScratchBytes(const DecoderOnly_Model& model, size_t max_batch) { + const auto& indexshare = model.index_share_config_; + const size_t stages = static_cast(indexshare.max_draft_tokens - 1); + const size_t extend_rows = CheckedMtpMultiply(max_batch, static_cast(indexshare.max_draft_tokens) + 1); + const size_t output_capacity = CheckedMtpAdd(static_cast(indexshare.base_capacity), stages); + const size_t batch_ints = CheckedMtpAdd( + CheckedMtpAdd(CheckedMtpAdd(static_cast(indexshare.base_capacity), 2), std::max(size_t{1}, stages)), + CheckedMtpMultiply(stages, CheckedMtpAdd(output_capacity, 5))); + const size_t ints = CheckedMtpAdd( + CheckedMtpMultiply(extend_rows, CheckedMtpAdd(output_capacity, 3)), + CheckedMtpMultiply(max_batch, batch_ints)); + const size_t projection_rows = CheckedMtpAdd(extend_rows, CheckedMtpMultiply(stages, max_batch)); + size_t bytes = CheckedMtpAdd(CheckedMtpMultiply(ints, sizeof(int32_t)), + CheckedMtpMultiply(projection_rows, sizeof(int64_t))); + const auto& decoder = model.config_->model.decoder; + size_t update_bank_bytes = 0; + if (decoder.state_groups) { + for (const auto& group : *decoder.state_groups) { + if (!group.state_update) continue; + std::vector templates; + using Kind = Config::Model::Decoder::StateGroupKind; + if (group.kind == Kind::FixedIndexer) { + templates.push_back(decoder.outputs.state_update_indexer_names); + } else if (group.kind == Kind::FixedConv) { + templates.push_back(decoder.outputs.state_update_conv_value_names); + } else if (group.kind == Kind::FixedPle) { + templates.push_back(decoder.outputs.state_update_ple_token_names); + templates.push_back(decoder.outputs.state_update_ple_conv_value_names); + } + for (const auto& output_template : templates) { + if (output_template.empty()) continue; + for (int layer : group.layer_ids) { + std::string name = output_template; + const auto placeholder = name.find("%d"); + if (placeholder != std::string::npos) name.replace(placeholder, 2, std::to_string(layer)); + const auto shape = model.session_info_.GetOutputShape(name); + size_t row_elements = 1; + for (size_t axis = 1; axis < shape.size(); ++axis) { + if (shape[axis] <= 0) throw std::runtime_error("IndexShare state-update scratch requires static non-batch geometry."); + row_elements = CheckedMtpMultiply(row_elements, static_cast(shape[axis])); + } + update_bank_bytes = CheckedMtpAdd(update_bank_bytes, CheckedMtpMultiply( + max_batch, CheckedMtpMultiply(row_elements, Ort::SizeOf(model.session_info_.GetOutputDataType(name))))); + } + } + } + } + const size_t state_bytes = CheckedMtpAdd( + CheckedMtpMultiply(2, FixedStatePool::StateBankBytes(model, max_batch)), update_bank_bytes); + return CheckedMtpAdd(bytes, CheckedMtpMultiply(stages, state_bytes)); +} + +EngineDependencies CreateDecoderDependencies(std::shared_ptr model) { + // The static batch decoder rebuilds its contiguous cache from the whole sequence every step, so + // it cannot resume a half-written prompt. Reject a model that configures chunking here rather + // than surfacing it once a Request is already being admitted. + if (!model->config_->engine.dynamic_batching && + model->config_->search.chunk_size.value_or(0) != 0) { + throw std::runtime_error( + "search.chunk_size requires dynamic batching; the static batch scheduler cannot chunk a " + "prefill."); + } + std::shared_ptr mtp_model; + size_t mtp_bytes_per_block = 0; + if (model->config_->model.mtp.IsEnabled()) { + if (!model->config_->engine.dynamic_batching) { + throw std::runtime_error("An Engine-hosted MTP head requires dynamic batching."); + } + auto& target_hidden_states = + model->config_->model.decoder.outputs.hidden_states; + target_hidden_states = model->config_->model.mtp.main_hidden_states; + mtp_model = std::make_shared( + CreateMtpDecoderConfig(*model->config_), GetOrtEnv()); + mtp_model->InitializeIndexShare(model->config_->model.mtp.index_share, GetOrtEnv()); + ValidateMtpModelCompatibility( + *model->config_, model->session_info_, mtp_model->session_info_); + mtp_bytes_per_block = PagedKeyValueCacheBytesPerBlock(mtp_model); + // The head consumes the target decoder's packed hidden states, so the target must emit them. + model->config_->engine.hidden_states_output_required = true; + } + + std::unique_ptr dflash2_drafter; + std::shared_ptr dflash2_model; + size_t dflash2_bytes_per_block = 0; + size_t dflash2_reserved_memory_bytes = 0; + size_t dflash2_max_batch_size = 0; + size_t dflash2_pool_blocks = 0; + size_t dflash2_prefix_checkpoint_bytes = 0; + bool dflash2_prefix_checkpoints_enabled = false; + 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); + dflash2_max_batch_size = static_cast(batching.max_batch_size); + auto decoder_model = std::dynamic_pointer_cast(model); + dflash2_model = std::make_shared( + CreateDflash2Config(*model->config_), GetOrtEnv(), + decoder_model ? decoder_model->cpu_embedding_ : nullptr); + const auto dflash2_cache_type = ValidateDflash2ModelCompatibility( + *model->config_, model->session_info_, dflash2_model->session_info_, paged_block_size); + model->config_->engine.aux_hidden_states_output_required = true; + dflash2_pool_blocks = Dflash2Drafter::PoolBlocks( + *model->config_, paged_block_size, dflash2_max_batch_size); + if (dflash2_pool_blocks != 0) { + if (batching.num_blocks.has_value()) { + // Explicit num_blocks bypasses the free-memory probe. CacheManager validates and deducts + // this fixed pool before either cache pool is allocated. + dflash2_reserved_memory_bytes = Dflash2Drafter::PoolBytes( + *model->config_, paged_block_size, dflash2_pool_blocks, dflash2_cache_type); + } else { + // Automatic sizing measures free memory after allocating the fixed drafter pool. + dflash2_drafter = std::make_unique( + dflash2_model, paged_block_size, dflash2_pool_blocks, dflash2_max_batch_size); + } + } else { + // A full-attention drafter mirrors the target pool, so it is billed per target block instead + // and built below once the target pool's block count is known. + dflash2_bytes_per_block = + Dflash2Drafter::BytesPerBlock(*model->config_, paged_block_size, dflash2_cache_type); + dflash2_reserved_memory_bytes = Dflash2Drafter::FullAttentionReservedBytes( + paged_block_size, static_cast(dflash2.block_size), + dflash2_max_batch_size, dflash2_bytes_per_block); + } + if (dflash2_model->cpu_embedding_) { + const size_t embedding_bytes = Dflash2Drafter::EmbeddingReservedBytes( + dflash2_max_batch_size, static_cast(dflash2.block_size), + static_cast(dflash2_model->cpu_embedding_->hidden_size_), + dflash2_model->cpu_embedding_->type_); + if (embedding_bytes > std::numeric_limits::max() - dflash2_reserved_memory_bytes) { + throw std::runtime_error("DFlash 2 reserved memory bytes overflow size_t."); + } + dflash2_reserved_memory_bytes += embedding_bytes; + } + if (dflash2_pool_blocks != 0 && mtp_bytes_per_block == 0 && !dflash2.is_dspark && + ModelStateManifest{model->config_->model.decoder}.HasFixedStateGroups() && + ResolvePrefixCachingEnabled(model, /*auxiliary_bytes_per_block=*/0)) { + dflash2_prefix_checkpoint_bytes = Dflash2Drafter::PrefixCheckpointBytes( + *model->config_, paged_block_size, dflash2_cache_type); + } + } + + if (dflash2_bytes_per_block > std::numeric_limits::max() - mtp_bytes_per_block) { + throw std::runtime_error("Engine auxiliary cache bytes per block overflow size_t."); + } + // SimpleDecoder allocates its persistent capture buffers only after the cache has chosen its + // block count, so the automatic free-memory probe has to know about them or the paged pool claims + // the memory they still need. IndexShare's selection and per-stage state scratch are also priced + // here, even without graph capture. Sizing uses the engine-wide draft ceiling, which is at least what + // each decoder will ask for. An explicit num_blocks is the caller pinning the pool, so it is left + // alone. + size_t graph_buffer_reserved_bytes = 0; + if (model->config_->engine.dynamic_batching && + !model->config_->engine.dynamic_batching->num_blocks.has_value()) { + const auto reserve_for = [](const std::shared_ptr& decoder) -> size_t { + if (!decoder || !IsGraphCaptureEnabled(decoder->config_->model.decoder.session_options)) { + return 0; + } + return VarlenGraphBufferBytes(*decoder, PackedPositionIdPlanes(*decoder), + kMaxDraftTokensPerStep + 1); + }; + graph_buffer_reserved_bytes = reserve_for(model) + reserve_for(mtp_model); + if (mtp_model) { + const size_t max_batch = model->config_->engine.dynamic_batching->max_batch_size; + const size_t scratch_bytes = mtp_model->index_share_config_.base_capacity > 0 + ? IndexShareScratchBytes(*mtp_model, max_batch) + : CheckedMtpMultiply(2, FixedStatePool::StateBankBytes(*mtp_model, max_batch)); + graph_buffer_reserved_bytes = CheckedMtpAdd(graph_buffer_reserved_bytes, scratch_bytes); + } + } + if (graph_buffer_reserved_bytes > + std::numeric_limits::max() - dflash2_reserved_memory_bytes) { + throw std::runtime_error("Engine reserved memory bytes overflow size_t."); + } + std::shared_ptr cache_manager = + CacheManager::Create(model, mtp_bytes_per_block + dflash2_bytes_per_block, + dflash2_reserved_memory_bytes + graph_buffer_reserved_bytes, + dflash2_prefix_checkpoint_bytes, + &dflash2_prefix_checkpoints_enabled); + std::shared_ptr snapshot_cache; + const auto& batching = model->config_->engine.dynamic_batching; + const bool webgpu = model->p_device_kvcache_->GetType() == DeviceType::WEBGPU; + const bool device_resident = webgpu || (batching && !batching->snapshot_cache.Configured()); + if (webgpu && batching && batching->snapshot_cache.Enabled() && g_log.enabled && g_log.warning) { + Log("warning", + "WebGPU does not support tiered prefix snapshot byte budgets; ignoring these limits " + "and retaining device-resident snapshots subject to count limits."); + } + if (batching && batching->prefix_caching.value_or(true) && batching->dflash2_max_snapshots != 0 && + (!batching->max_snapshots || *batching->max_snapshots != 0) && + (device_resident || batching->snapshot_cache.Enabled())) { + const auto& limits = batching->snapshot_cache; + snapshot_cache = std::make_shared( + std::array{cache_manager->DraftCheckpointVramBytes(), + limits.ram_bytes, limits.disk_bytes}, + device_resident ? cache_manager->DraftCheckpointCapacity() + : batching->max_snapshots.value_or(std::numeric_limits::max()), + device_resident ? std::filesystem::path{} + : std::filesystem::path{fs::path{limits.directory}.c_str()}, + device_resident); + } + if (dflash2_model && !dflash2_drafter) { + const size_t paged_block_size = + static_cast(model->config_->engine.dynamic_batching->block_size); + if (dflash2_pool_blocks != 0) { + dflash2_drafter = std::make_unique( + dflash2_model, paged_block_size, dflash2_pool_blocks, dflash2_max_batch_size); + } else { + const auto& dflash2 = model->config_->model.dflash2; + dflash2_drafter = std::make_unique( + dflash2_model, paged_block_size, + Dflash2Drafter::FullAttentionPoolBlocks( + cache_manager->Snapshot().total_blocks, paged_block_size, + static_cast(dflash2.block_size), dflash2_max_batch_size), + dflash2_max_batch_size); + } + } + auto scheduler = Scheduler::Create(model, cache_manager); + auto model_executor = ModelExecutor::Create(model, cache_manager); + + std::shared_ptr mtp_cache_manager; + std::unique_ptr mtp_model_executor; + if (mtp_model) { + // Both decoders cover the same resident request set. Fixing the head to the main pool's block + // count makes the auxiliary bytes included above exact rather than letting it independently + // consume another gpu_utilization_factor share of currently free memory. + mtp_model->config_->engine.dynamic_batching->num_blocks = + cache_manager->Snapshot().total_blocks; + mtp_cache_manager = CacheManager::Create(mtp_model); + mtp_model_executor = ModelExecutor::Create(mtp_model, mtp_cache_manager); + } + + 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(dflash2_drafter), dflash2_prefix_checkpoints_enabled, + nullptr, std::move(snapshot_cache)}; +} + +} // namespace + +BackendCapabilities ResolveBackendCapabilities( + const Model& model, const CacheManager& cache_manager, const Scheduler& scheduler, + const ModelExecutor& executor, std::optional declared) { + BackendCapabilities resolved; + const bool transactions = cache_manager.SupportsDynamicBatching(); + const bool own_tokens = executor.SelectsOwnTokens(); + if (transactions && own_tokens) { + throw std::runtime_error("Backend-selected tokens require the nontransactional Engine execution protocol."); + } + resolved.execution_protocol = transactions ? BackendExecutionProtocol::Transactional + : BackendExecutionProtocol::Nontransactional; + resolved.token_selection = own_tokens ? BackendTokenSelection::Backend + : BackendTokenSelection::Engine; + resolved.admission = (cache_manager.SupportsContinuousMultiSlotDecoding() && + scheduler.SupportsIndependentSlotAdmission()) + ? BackendAdmission::IndependentSlots + : BackendAdmission::WholeBatch; + resolved.kv_storage = transactions ? BackendKvStorage::Paged + : (own_tokens ? BackendKvStorage::PersistentContiguous + : BackendKvStorage::RebuiltContiguous); + if (own_tokens) { + resolved.additional_state = BackendAdditionalState::PrivateBanks; + resolved.continuation = BackendContinuation::ReplayHistory; + resolved.supports_sampling = false; + resolved.supports_penalties = false; + resolved.supports_min_generated_tokens = false; + resolved.supports_guidance = false; + } else if (ModelStateManifest{model.config_->model.decoder}.HasFixedStateGroups()) { + resolved.additional_state = BackendAdditionalState::FixedPools; + } + resolved.recovery = transactions ? BackendRecovery::TransactionalRollback + : BackendRecovery::FatalAfterExecution; + resolved.prefix_reuse = cache_manager.PrefixMetrics() != nullptr; + resolved.supports_stop_strings = transactions || own_tokens; + resolved.supports_per_turn_seed = transactions && scheduler.SupportsTransactionalSamplerState(); + resolved.supports_delimited_guidance = transactions && !own_tokens; + resolved.max_slots = cache_manager.MaxBatchSize(); + resolved.max_result_tokens = own_tokens || transactions ? kMaxGeneratedTokensPerStep : 1; + + const auto& engine = model.config_->engine; + if (engine.dynamic_batching) { + const auto& batching = *engine.dynamic_batching; + resolved.configured_max_batch_size = batching.max_batch_size; + resolved.max_scheduled_tokens = batching.max_scheduled_tokens; + const size_t block_count = cache_manager.TargetBlockCount(); + const size_t block_size = cache_manager.TargetBlockSize(); + if (block_count != 0 && block_size != 0) { + if (block_count > (std::numeric_limits::max() - 1) / block_size) { + throw std::overflow_error("Engine max_request_length exceeds uint64_t."); + } + const uint64_t cache_limit = static_cast(block_count) * block_size + 1; + resolved.max_request_length = std::min( + static_cast(model.config_->model.context_length), cache_limit); + } + } else { + resolved.configured_max_batch_size = engine.static_batching + ? engine.static_batching->max_batch_size + : kDefaultStaticBatchSize; + } + if (!transactions) { + if (const size_t rows = cache_manager.MaxStaticBatchSize()) { + resolved.configured_max_batch_size = rows; + } + resolved.max_request_length = cache_manager.MaxStaticRequestLength(); + if (resolved.max_request_length != 0 && model.config_->model.context_length > 0) { + resolved.max_request_length = std::min( + resolved.max_request_length, + static_cast(model.config_->model.context_length)); + } + } + + if (!declared) return resolved; + const auto& explicit_capabilities = *declared; + if (explicit_capabilities.UsesTransactions() != transactions || + explicit_capabilities.SelectsOwnTokens() != own_tokens) { + throw std::runtime_error("Engine backend capabilities contradict its cache/executor execution protocol."); + } + if (explicit_capabilities.max_slots == 0 || + explicit_capabilities.max_slots != cache_manager.MaxBatchSize() || + explicit_capabilities.max_result_tokens == 0 || + explicit_capabilities.max_result_tokens > kMaxGeneratedTokensPerStep || + (!own_tokens && !transactions && explicit_capabilities.max_result_tokens != 1)) { + throw std::runtime_error("Engine backend capabilities contradict its effective execution limits."); + } + if (explicit_capabilities.SupportsIndependentAdmission() && + explicit_capabilities.max_slots > 1 && + (!cache_manager.SupportsContinuousMultiSlotDecoding() || + !scheduler.SupportsIndependentSlotAdmission())) { + throw std::runtime_error("Engine backend capabilities require unsupported independent slot admission."); + } + if ((explicit_capabilities.supports_per_turn_seed && !resolved.supports_per_turn_seed) || + (explicit_capabilities.supports_stop_strings && !resolved.supports_stop_strings) || + (explicit_capabilities.supports_delimited_guidance && !resolved.supports_delimited_guidance) || + (explicit_capabilities.prefix_reuse && !resolved.prefix_reuse) || + (!transactions && explicit_capabilities.recovery == BackendRecovery::TransactionalRollback) || + // Zero declares no limit of its own (the model's search limit applies), which a finite cache cannot honor. + (resolved.max_request_length != 0 && + (explicit_capabilities.max_request_length == 0 || + explicit_capabilities.max_request_length > resolved.max_request_length))) { + throw std::runtime_error("Engine backend capabilities exceed its collaborators' supported features or capacity."); + } + return explicit_capabilities; +} + +EngineBackendBundle CreateEngineBackend(std::shared_ptr model) { + if (!model) throw std::invalid_argument("Engine backend construction requires a model."); + EngineDependencies dependencies; + std::optional declared; + if (auto graph_model = std::dynamic_pointer_cast(model)) { + declared = graph_model->ResolvedBackendCapabilities(); + dependencies = CreateQwen38GraphDependencies(std::move(graph_model)); + } else { + dependencies = CreateDecoderDependencies(model); + } + auto capabilities = ResolveBackendCapabilities( + *model, *dependencies.cache_manager, *dependencies.scheduler, *dependencies.model_executor, + declared ? declared : dependencies.backend_capabilities); + dependencies.backend_capabilities = capabilities; + return {std::move(dependencies), capabilities}; +} + +} // namespace Generators diff --git a/src/engine/backend_factory.h b/src/engine/backend_factory.h new file mode 100644 index 0000000000..57a891023e --- /dev/null +++ b/src/engine/backend_factory.h @@ -0,0 +1,25 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "engine.h" + +namespace Generators { + +struct EngineBackendBundle { + EngineDependencies dependencies; + BackendCapabilities capabilities; +}; + +// Resources/profile selection belongs to Model initialization. This factory only assembles +// collaborators for that already-selected model and leases its execution context. +EngineBackendBundle CreateEngineBackend(std::shared_ptr model); + +// Compatibility for injected collaborators. Explicit descriptors are checked against the actual +// execution interfaces; old test/custom bundles can omit one and retain their existing behavior. +BackendCapabilities ResolveBackendCapabilities( + const Model& model, const CacheManager& cache_manager, const Scheduler& scheduler, + const ModelExecutor& executor, std::optional declared = std::nullopt); + +} // namespace Generators diff --git a/src/engine/cache_manager.h b/src/engine/cache_manager.h index 8b771c6e07..5ecd7a17c8 100644 --- a/src/engine/cache_manager.h +++ b/src/engine/cache_manager.h @@ -90,8 +90,17 @@ struct CacheManager { virtual bool SupportsDynamicBatching() const = 0; + virtual bool SupportsContinuousMultiSlotDecoding() const { return SupportsDynamicBatching(); } + virtual size_t MaxBatchSize() const { return 4; } + // Rows a static (non-paged) cache supports at most, or 0 when the configured static batch size applies. + virtual size_t MaxStaticBatchSize() const { return 0; } + + // Longest session (prompt plus generated tokens) a static (non-paged) cache can hold, or 0 when only the + // model-configured search.max_length limits it. Feeds EngineCapabilities::max_request_length. + virtual size_t MaxStaticRequestLength() const { return 0; } + virtual std::vector> AllocatedRequests() const = 0; virtual bool IsResident(const std::shared_ptr& request) const = 0; diff --git a/src/engine/engine.cpp b/src/engine/engine.cpp index eb88c632ba..804eb9fa0c 100644 --- a/src/engine/engine.cpp +++ b/src/engine/engine.cpp @@ -8,6 +8,7 @@ #include "../models/preprocessing/genai_tokenizer.h" #include "../stop_string_controller.h" #include "decoders/varlen_decoder_io.h" +#include "backend_factory.h" #include #include @@ -19,70 +20,6 @@ static_assert(kMaxDraftTokensPerStep < kSpeculativeAcceptanceLengthBins); namespace { -size_t CheckedMtpAdd(size_t left, size_t right) { - if (right > std::numeric_limits::max() - left) - throw std::runtime_error("MTP scratch memory bytes overflow size_t."); - return left + right; -} - -size_t CheckedMtpMultiply(size_t left, size_t right) { - if (left != 0 && right > std::numeric_limits::max() / left) - throw std::runtime_error("MTP scratch memory bytes overflow size_t."); - return left * right; -} - -size_t IndexShareScratchBytes(const DecoderOnly_Model& model, size_t max_batch) { - const auto& indexshare = model.index_share_config_; - const size_t stages = static_cast(indexshare.max_draft_tokens - 1); - const size_t extend_rows = CheckedMtpMultiply(max_batch, static_cast(indexshare.max_draft_tokens) + 1); - const size_t output_capacity = CheckedMtpAdd(static_cast(indexshare.base_capacity), stages); - const size_t batch_ints = CheckedMtpAdd( - CheckedMtpAdd(CheckedMtpAdd(static_cast(indexshare.base_capacity), 2), std::max(size_t{1}, stages)), - CheckedMtpMultiply(stages, CheckedMtpAdd(output_capacity, 5))); - const size_t ints = CheckedMtpAdd( - CheckedMtpMultiply(extend_rows, CheckedMtpAdd(output_capacity, 3)), - CheckedMtpMultiply(max_batch, batch_ints)); - const size_t projection_rows = CheckedMtpAdd(extend_rows, CheckedMtpMultiply(stages, max_batch)); - size_t bytes = CheckedMtpAdd(CheckedMtpMultiply(ints, sizeof(int32_t)), - CheckedMtpMultiply(projection_rows, sizeof(int64_t))); - const auto& decoder = model.config_->model.decoder; - size_t update_bank_bytes = 0; - if (decoder.state_groups) { - for (const auto& group : *decoder.state_groups) { - if (!group.state_update) continue; - std::vector templates; - using Kind = Config::Model::Decoder::StateGroupKind; - if (group.kind == Kind::FixedIndexer) { - templates.push_back(decoder.outputs.state_update_indexer_names); - } else if (group.kind == Kind::FixedConv) { - templates.push_back(decoder.outputs.state_update_conv_value_names); - } else if (group.kind == Kind::FixedPle) { - templates.push_back(decoder.outputs.state_update_ple_token_names); - templates.push_back(decoder.outputs.state_update_ple_conv_value_names); - } - for (const auto& output_template : templates) { - if (output_template.empty()) continue; - for (int layer : group.layer_ids) { - std::string name = output_template; - const auto placeholder = name.find("%d"); - if (placeholder != std::string::npos) name.replace(placeholder, 2, std::to_string(layer)); - const auto shape = model.session_info_.GetOutputShape(name); - size_t row_elements = 1; - for (size_t axis = 1; axis < shape.size(); ++axis) { - if (shape[axis] <= 0) throw std::runtime_error("IndexShare state-update scratch requires static non-batch geometry."); - row_elements = CheckedMtpMultiply(row_elements, static_cast(shape[axis])); - } - update_bank_bytes = CheckedMtpAdd(update_bank_bytes, CheckedMtpMultiply( - max_batch, CheckedMtpMultiply(row_elements, Ort::SizeOf(model.session_info_.GetOutputDataType(name))))); - } - } - } - } - const size_t state_bytes = CheckedMtpAdd( - CheckedMtpMultiply(2, FixedStatePool::StateBankBytes(model, max_batch)), update_bank_bytes); - return CheckedMtpAdd(bytes, CheckedMtpMultiply(stages, state_bytes)); -} - // A throwing decode/copy must fence queued work before its local tensor owners unwind. struct MtpFailureFence { DeviceInterface* device; @@ -267,6 +204,9 @@ Engine::Engine(std::shared_ptr model, EngineDependencies dependencies) ? dependencies.make_step_error : MakeEngineStepError} { // Fail fast on a missing collaborator rather than crashing later on first use. + if (!model_) { + throw std::runtime_error("Engine requires a non-null model."); + } if (!cache_manager_) { throw std::runtime_error("Engine requires a non-null cache manager."); } @@ -282,6 +222,8 @@ Engine::Engine(std::shared_ptr model, EngineDependencies dependencies) if (dflash2_drafter_ && snapshot_cache_) { dflash2_drafter_->InitializeSnapshotCache(snapshot_cache_); } + backend_capabilities_ = ResolveBackendCapabilities( + *model_, *cache_manager_, *scheduler_, *model_executor_, dependencies.backend_capabilities); fatal_contract_fallback_error_ = MakeFatalFallbackError(StepOutcomeKind::ExecutionContractFailure); fatal_execution_fallback_error_ = @@ -387,206 +329,7 @@ void ValidateMtpModelCompatibility(const Config& config, } EngineDependencies Engine::CreateDependencies(std::shared_ptr model) { - // The static batch decoder rebuilds its contiguous cache from the whole sequence every step, so - // it cannot resume a half-written prompt. Reject a model that configures chunking here rather - // than surfacing it once a Request is already being admitted. - if (!model->config_->engine.dynamic_batching && - model->config_->search.chunk_size.value_or(0) != 0) { - throw std::runtime_error( - "search.chunk_size requires dynamic batching; the static batch scheduler cannot chunk a " - "prefill."); - } - std::shared_ptr mtp_model; - size_t mtp_bytes_per_block = 0; - if (model->config_->model.mtp.IsEnabled()) { - if (!model->config_->engine.dynamic_batching) { - throw std::runtime_error("An Engine-hosted MTP head requires dynamic batching."); - } - auto& target_hidden_states = - model->config_->model.decoder.outputs.hidden_states; - target_hidden_states = model->config_->model.mtp.main_hidden_states; - mtp_model = std::make_shared( - CreateMtpDecoderConfig(*model->config_), GetOrtEnv()); - mtp_model->InitializeIndexShare(model->config_->model.mtp.index_share, GetOrtEnv()); - ValidateMtpModelCompatibility( - *model->config_, model->session_info_, mtp_model->session_info_); - mtp_bytes_per_block = PagedKeyValueCacheBytesPerBlock(mtp_model); - // The head consumes the target decoder's packed hidden states, so the target must emit them. - model->config_->engine.hidden_states_output_required = true; - } - - std::unique_ptr dflash2_drafter; - std::shared_ptr dflash2_model; - size_t dflash2_bytes_per_block = 0; - size_t dflash2_reserved_memory_bytes = 0; - size_t dflash2_max_batch_size = 0; - size_t dflash2_pool_blocks = 0; - size_t dflash2_prefix_checkpoint_bytes = 0; - bool dflash2_prefix_checkpoints_enabled = false; - 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); - dflash2_max_batch_size = static_cast(batching.max_batch_size); - auto decoder_model = std::dynamic_pointer_cast(model); - dflash2_model = std::make_shared( - CreateDflash2Config(*model->config_), GetOrtEnv(), - decoder_model ? decoder_model->cpu_embedding_ : nullptr); - const auto dflash2_cache_type = ValidateDflash2ModelCompatibility( - *model->config_, model->session_info_, dflash2_model->session_info_, paged_block_size); - model->config_->engine.aux_hidden_states_output_required = true; - dflash2_pool_blocks = Dflash2Drafter::PoolBlocks( - *model->config_, paged_block_size, dflash2_max_batch_size); - if (dflash2_pool_blocks != 0) { - if (batching.num_blocks.has_value()) { - // Explicit num_blocks bypasses the free-memory probe. CacheManager validates and deducts - // this fixed pool before either cache pool is allocated. - dflash2_reserved_memory_bytes = Dflash2Drafter::PoolBytes( - *model->config_, paged_block_size, dflash2_pool_blocks, dflash2_cache_type); - } else { - // Automatic sizing measures free memory after allocating the fixed drafter pool. - dflash2_drafter = std::make_unique( - dflash2_model, paged_block_size, dflash2_pool_blocks, dflash2_max_batch_size); - } - } else { - // A full-attention drafter mirrors the target pool, so it is billed per target block instead - // and built below once the target pool's block count is known. - dflash2_bytes_per_block = - Dflash2Drafter::BytesPerBlock(*model->config_, paged_block_size, dflash2_cache_type); - dflash2_reserved_memory_bytes = Dflash2Drafter::FullAttentionReservedBytes( - paged_block_size, static_cast(dflash2.block_size), - dflash2_max_batch_size, dflash2_bytes_per_block); - } - if (dflash2_model->cpu_embedding_) { - const size_t embedding_bytes = Dflash2Drafter::EmbeddingReservedBytes( - dflash2_max_batch_size, static_cast(dflash2.block_size), - static_cast(dflash2_model->cpu_embedding_->hidden_size_), - dflash2_model->cpu_embedding_->type_); - if (embedding_bytes > std::numeric_limits::max() - dflash2_reserved_memory_bytes) { - throw std::runtime_error("DFlash 2 reserved memory bytes overflow size_t."); - } - dflash2_reserved_memory_bytes += embedding_bytes; - } - if (dflash2_pool_blocks != 0 && mtp_bytes_per_block == 0 && !dflash2.is_dspark && - ModelStateManifest{model->config_->model.decoder}.HasFixedStateGroups() && - ResolvePrefixCachingEnabled(model, /*auxiliary_bytes_per_block=*/0)) { - dflash2_prefix_checkpoint_bytes = Dflash2Drafter::PrefixCheckpointBytes( - *model->config_, paged_block_size, dflash2_cache_type); - } - } - - if (dflash2_bytes_per_block > std::numeric_limits::max() - mtp_bytes_per_block) { - throw std::runtime_error("Engine auxiliary cache bytes per block overflow size_t."); - } - // SimpleDecoder allocates its persistent capture buffers only after the cache has chosen its - // block count, so the automatic free-memory probe has to know about them or the paged pool claims - // the memory they still need. IndexShare's selection and per-stage state scratch are also priced - // here, even without graph capture. Sizing uses the engine-wide draft ceiling, which is at least what - // each decoder will ask for. An explicit num_blocks is the caller pinning the pool, so it is left - // alone. - size_t graph_buffer_reserved_bytes = 0; - if (model->config_->engine.dynamic_batching && - !model->config_->engine.dynamic_batching->num_blocks.has_value()) { - const auto reserve_for = [](const std::shared_ptr& decoder) -> size_t { - if (!decoder || !IsGraphCaptureEnabled(decoder->config_->model.decoder.session_options)) { - return 0; - } - return VarlenGraphBufferBytes(*decoder, PackedPositionIdPlanes(*decoder), - kMaxDraftTokensPerStep + 1); - }; - graph_buffer_reserved_bytes = reserve_for(model) + reserve_for(mtp_model); - if (mtp_model) { - const size_t max_batch = model->config_->engine.dynamic_batching->max_batch_size; - const size_t scratch_bytes = mtp_model->index_share_config_.base_capacity > 0 - ? IndexShareScratchBytes(*mtp_model, max_batch) - : CheckedMtpMultiply(2, FixedStatePool::StateBankBytes(*mtp_model, max_batch)); - graph_buffer_reserved_bytes = CheckedMtpAdd(graph_buffer_reserved_bytes, scratch_bytes); - } - } - if (graph_buffer_reserved_bytes > - std::numeric_limits::max() - dflash2_reserved_memory_bytes) { - throw std::runtime_error("Engine reserved memory bytes overflow size_t."); - } - std::shared_ptr cache_manager = - CacheManager::Create(model, mtp_bytes_per_block + dflash2_bytes_per_block, - dflash2_reserved_memory_bytes + graph_buffer_reserved_bytes, - dflash2_prefix_checkpoint_bytes, - &dflash2_prefix_checkpoints_enabled); - std::shared_ptr snapshot_cache; - const auto& batching = model->config_->engine.dynamic_batching; - const bool webgpu = model->p_device_kvcache_->GetType() == DeviceType::WEBGPU; - const bool device_resident = webgpu || (batching && !batching->snapshot_cache.Configured()); - if (webgpu && batching && batching->snapshot_cache.Enabled() && g_log.enabled && g_log.warning) { - Log("warning", - "WebGPU does not support tiered prefix snapshot byte budgets; ignoring these limits " - "and retaining device-resident snapshots subject to count limits."); - } - if (batching && batching->prefix_caching.value_or(true) && batching->dflash2_max_snapshots != 0 && - (!batching->max_snapshots || *batching->max_snapshots != 0) && - (device_resident || batching->snapshot_cache.Enabled())) { - const auto& limits = batching->snapshot_cache; - snapshot_cache = std::make_shared( - std::array{cache_manager->DraftCheckpointVramBytes(), - limits.ram_bytes, limits.disk_bytes}, - device_resident ? cache_manager->DraftCheckpointCapacity() - : batching->max_snapshots.value_or(std::numeric_limits::max()), - device_resident ? std::filesystem::path{} - : std::filesystem::path{fs::path{limits.directory}.c_str()}, - device_resident); - } - if (dflash2_model && !dflash2_drafter) { - const size_t paged_block_size = - static_cast(model->config_->engine.dynamic_batching->block_size); - if (dflash2_pool_blocks != 0) { - dflash2_drafter = std::make_unique( - dflash2_model, paged_block_size, dflash2_pool_blocks, dflash2_max_batch_size); - } else { - const auto& dflash2 = model->config_->model.dflash2; - dflash2_drafter = std::make_unique( - dflash2_model, paged_block_size, - Dflash2Drafter::FullAttentionPoolBlocks( - cache_manager->Snapshot().total_blocks, paged_block_size, - static_cast(dflash2.block_size), dflash2_max_batch_size), - dflash2_max_batch_size); - } - } - auto scheduler = Scheduler::Create(model, cache_manager); - auto model_executor = ModelExecutor::Create(model, cache_manager); - - std::shared_ptr mtp_cache_manager; - std::unique_ptr mtp_model_executor; - if (mtp_model) { - // Both decoders cover the same resident request set. Fixing the head to the main pool's block - // count makes the auxiliary bytes included above exact rather than letting it independently - // consume another gpu_utilization_factor share of currently free memory. - mtp_model->config_->engine.dynamic_batching->num_blocks = - cache_manager->Snapshot().total_blocks; - mtp_cache_manager = CacheManager::Create(mtp_model); - mtp_model_executor = ModelExecutor::Create(mtp_model, mtp_cache_manager); - } - - 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(dflash2_drafter), dflash2_prefix_checkpoints_enabled, - nullptr, std::move(snapshot_cache)}; + return CreateEngineBackend(std::move(model)).dependencies; } void Engine::PrepareDflash2Feeds(const StepPlan& plan, @@ -1842,10 +1585,20 @@ uint64_t Engine::BeginTurn(const std::shared_ptr& request, "Cannot begin a turn for a request that does not belong to this engine."); } request->ValidateTurnAdmission(tokens, options); + if (backend_capabilities_.SelectsOwnTokens()) { + const int vocab_size = model_->config_->model.vocab_size; + for (size_t i = 0; i < tokens.size(); ++i) { + if (tokens[i] < 0 || tokens[i] >= vocab_size) { + throw std::invalid_argument("Prompt token " + std::to_string(i) + " has id " + + std::to_string(tokens[i]) + ", outside [0, " + std::to_string(vocab_size) + + ") (model.vocab_size)."); + } + } + } // Resolve and validate this turn's complete policy before any Request mutation, so an // unsatisfiable combination leaves the previous completed turn entirely reusable. - const bool dynamic_batching = cache_manager_->SupportsDynamicBatching(); + const bool dynamic_batching = backend_capabilities_.UsesTransactions(); const auto policy = ResolveTurnPolicy(model_->config_->search, options); const size_t turn_prompt_length = static_cast(request->IsAwaitingFirstTurn() @@ -1855,6 +1608,15 @@ uint64_t Engine::BeginTurn(const std::shared_ptr& request, ValidateTurnPolicy(policy, options, model_->p_device_scoring_->GetType(), model_->config_->model.vocab_size, turn_prompt_length, request->MaxSessionTokens()); + model_executor_->ValidateTurn(policy, options); + if ((!backend_capabilities_.supports_sampling && !policy.IsGreedy()) || + (!backend_capabilities_.supports_penalties && + (policy.repetition_penalty != 1.0f || policy.no_repeat_ngram_size != 0)) || + (!backend_capabilities_.supports_min_generated_tokens && policy.min_generated_tokens != 0) || + (!backend_capabilities_.supports_guidance && + (!options.guidance_type.empty() || !options.guidance_data.empty() || options.delimited_guidance))) { + throw std::runtime_error("The selected execution backend does not support the requested turn policy."); + } // Static batching completes generation through a non-transactional path (RunStatic) that cannot // stage, roll back, or replay a stop match, and its sampler RNG state is created once when the // Request joins the batch rather than inside a rollback-capable step. Reject both before any @@ -1863,7 +1625,9 @@ uint64_t Engine::BeginTurn(const std::shared_ptr& request, if (options.delimited_guidance) { throw std::runtime_error("Delimited guidance requires an Engine configured for dynamic batching."); } - if (!options.stop_strings.empty()) { + // An executor that commits its own tokens runs them through Request::CommitExternalTokens, which + // observes stop strings token by token and only ever ends the turn, so it has nothing to roll back. + if (!options.stop_strings.empty() && !backend_capabilities_.supports_stop_strings) { throw std::runtime_error( "Stop strings require an Engine configured for dynamic batching."); } @@ -1872,12 +1636,18 @@ uint64_t Engine::BeginTurn(const std::shared_ptr& request, "A per-turn seed requires an Engine configured for dynamic batching."); } } - if (options.seed && !scheduler_->SupportsTransactionalSamplerState()) { + if (options.seed && !backend_capabilities_.supports_per_turn_seed) { throw std::runtime_error( "A per-turn seed requires a batched sampler that supports transaction rollback."); } + if ((!options.stop_strings.empty() && !backend_capabilities_.supports_stop_strings) || + (options.delimited_guidance && !backend_capabilities_.supports_delimited_guidance)) { + throw std::runtime_error("The selected execution backend does not support the requested stop or guidance policy."); + } + const bool first_turn = request->IsAwaitingFirstTurn(); + bool allow_nonresident = false; if (!first_turn) { const bool restartable_canceled_turn = request->IsRestartableCanceledTurn() && @@ -1885,7 +1655,7 @@ uint64_t Engine::BeginTurn(const std::shared_ptr& request, const bool replay_rewound_request = request->NeedsReplayAfterRewind() && !cache_manager_->IsResident(request); - const bool allow_nonresident = + allow_nonresident = restartable_canceled_turn || replay_rewound_request; ValidateRequestCanContinue(request, allow_nonresident); if (!allow_nonresident) { @@ -1918,9 +1688,9 @@ uint64_t Engine::BeginTurn(const std::shared_ptr& request, try { request->PrepareTurnAdmission(tokens, admission); - if (first_turn) { + if (first_turn || allow_nonresident) { scheduler_->AddRequest(request); - added_to_scheduler = true; + added_to_scheduler = first_turn; } return request->CommitTurnAdmission(options, admission); } catch (...) { @@ -2029,6 +1799,9 @@ void Engine::RewindRequestToStartOfTurn( "call Engine::Run() to drain the event before rewinding."); } + if (backend_capabilities_.rewind == BackendRewind::Unsupported) { + throw std::runtime_error("The selected execution backend does not support request rewind."); + } const bool resident = cache_manager_->IsResident(request); if (!resident && !request->IsRestartableCanceledTurn() && !request->NeedsReplayAfterRewind()) { @@ -2062,7 +1835,7 @@ void Engine::CloseRequest(const std::shared_ptr& request) { } scheduler_->RemoveRequest(request); const bool retain_static_runtime_state = - !cache_manager_->SupportsDynamicBatching() && + !backend_capabilities_.UsesTransactions() && cache_manager_->IsResident(request); CloseMtpRequest(request); @@ -2244,7 +2017,7 @@ void Engine::ValidateRequestCanContinue( } if (!allow_nonresident && - !cache_manager_->SupportsDynamicBatching() && + !backend_capabilities_.SupportsIndependentAdmission() && cache_manager_->ResidentRequestCount() > 1) { throw std::runtime_error( "Continuous decoding requires exactly one resident request in a " @@ -2319,7 +2092,7 @@ size_t Engine::Run(std::span events) { std::rethrow_exception(fatal_error_); } try { - if (cache_manager_->SupportsDynamicBatching()) { + if (backend_capabilities_.UsesTransactions()) { RunDynamic(); } else { RunStatic(); @@ -2335,9 +2108,9 @@ size_t Engine::Run(std::span events) { void Engine::RunStatic() { if (scheduler_->HasPendingRequests()) { - auto scheduled_requests = [&]() -> ScheduledRequests { + auto scheduling = [&]() -> NontransactionalSchedulingResult { try { - return scheduler_->Schedule(); + return scheduler_->ScheduleStep(); } catch (...) { MarkUnhealthyAndThrow( StepOutcomeKind::ExecutionContractFailure, @@ -2347,15 +2120,38 @@ void Engine::RunStatic() { std::current_exception()); } }(); + if (!scheduling.scheduled) { + if (scheduling.outcome.kind == StepOutcomeKind::UnserviceableRequest) { + RetainEvent(FailUnserviceableRequest(scheduling.outcome.request_id)); + } else if (scheduling.outcome.kind == StepOutcomeKind::CapacityDeferred) { + ++transaction_metrics_.capacity_deferrals; + } else if (scheduling.outcome.kind != StepOutcomeKind::NoWork) { + MarkUnhealthyAndThrow(StepOutcomeKind::ExecutionContractFailure, 0, nullptr, + "Static scheduler returned an invalid non-executable outcome.", nullptr); + } + return; + } + auto& scheduled_requests = *scheduling.scheduled; + if (scheduling.outcome.kind != StepOutcomeKind::Committed || !scheduled_requests) { + MarkUnhealthyAndThrow(StepOutcomeKind::ExecutionContractFailure, 0, nullptr, + "Static scheduler returned contradictory executable work.", nullptr); + } // Scheduling may recycle an all-terminal static batch. Closed rows retain their Search, // parameters, host tokens, and length metadata only until that shared allocation is gone; // finish their physical close before executing the replacement batch. CompleteNonresidentClosedRequests(); try { + scheduled_requests.SetBackendLimits(backend_capabilities_.max_slots, backend_capabilities_.max_result_tokens); model_executor_->Decode(scheduled_requests); - ++speculative_stats_.target_forward_passes; + if (backend_capabilities_.SelectsOwnTokens() != + (scheduled_requests.OutputKind() == StepOutputKind::BackendRows)) + throw std::logic_error("Executor output protocol contradicts the resolved backend capabilities."); + if (!model_executor_->SelectsOwnTokens()) { + ++speculative_stats_.target_forward_passes; + } scheduled_requests.GenerateNextTokens(step_results_); + model_executor_->OnStepCommitted(scheduled_requests, step_results_); } catch (...) { MarkUnhealthyAndThrow( StepOutcomeKind::FatalExecutionFailure, @@ -2960,6 +2756,7 @@ EngineEvent Engine::EventFromStepError( if (health_ == EngineHealth::Unhealthy) { std::rethrow_exception(fatal_error_); } + model_executor_->OnExecutionFailure(); std::exception_ptr durable_error = outcome == StepOutcomeKind::ExecutionContractFailure @@ -3068,7 +2865,7 @@ bool Engine::HasPendingRequests() { size_t Engine::MaxDraftTokensPerStep() const { ValidateOwnerThread(); - return cache_manager_->SupportsDynamicBatching() && + return backend_capabilities_.UsesTransactions() && model_executor_->SupportsDraftVerification() ? cache_manager_->MaxDraftTokensPerStep() : 0; @@ -3113,6 +2910,7 @@ SpeculativeStats Engine::GetSpeculativeStats() const { // would be a data race. ValidateOwnerThread(); auto stats = speculative_stats_; + model_executor_->AddSpeculativeStats(stats); if (dflash2_drafter_) { stats.dflash2_admission_misses = dflash2_drafter_->AdmissionMisses(); } @@ -3130,32 +2928,14 @@ SpeculativeStats Engine::GetSpeculativeStats() const { EngineCapabilities Engine::GetCapabilities() const { ValidateOwnerThread(); - EngineCapabilities capabilities; - if (model_->config_->engine.dynamic_batching) { - const auto& batching = *model_->config_->engine.dynamic_batching; - capabilities.configured_max_batch_size = batching.max_batch_size; - capabilities.max_scheduled_tokens = batching.max_scheduled_tokens; - const size_t block_count = cache_manager_->TargetBlockCount(); - const size_t block_size = cache_manager_->TargetBlockSize(); - if (block_count != 0 && block_size != 0) { - if (block_count > (std::numeric_limits::max() - 1) / block_size) { - throw std::overflow_error("Engine max_request_length exceeds uint64_t."); - } - const uint64_t cache_limit = static_cast(block_count) * block_size + 1; - capabilities.max_request_length = - std::min(static_cast(model_->config_->model.context_length), cache_limit); - } - } else if (model_->config_->engine.static_batching) { - capabilities.configured_max_batch_size = - model_->config_->engine.static_batching->max_batch_size; - } else { - capabilities.configured_max_batch_size = kDefaultStaticBatchSize; - } - return capabilities; + return {backend_capabilities_.configured_max_batch_size, + backend_capabilities_.max_scheduled_tokens, + backend_capabilities_.max_request_length}; } std::optional Engine::PrefixCacheStats() const { ValidateOwnerThread(); + if (!backend_capabilities_.prefix_reuse) return std::nullopt; const auto* metrics = cache_manager_->PrefixMetrics(); return metrics ? std::optional{*metrics} : std::nullopt; } diff --git a/src/engine/engine.h b/src/engine/engine.h index 40822f2e78..de410c504b 100644 --- a/src/engine/engine.h +++ b/src/engine/engine.h @@ -6,6 +6,7 @@ #include "request.h" #include "model_executor.h" #include "scheduler.h" +#include "backend_capabilities.h" #include "../decoding/speculative_stats.h" #include "../dflash2_drafter.h" @@ -86,7 +87,7 @@ using EngineStepErrorFactory = * @struct EngineDependencies * @brief Bundle of the collaborators the Engine drives. * - * The Engine constructor assembles these from the model via CreateDependencies and forwards them + * The Engine constructor assembles these from the model via CreateEngineBackend and forwards them * to the dependency-injecting constructor, so there is a single construction flow. The injecting * constructor accepts a pre-assembled bundle, which also lets a caller supply alternative * collaborators. The cache manager, scheduler, and model executor are all abstract interfaces, so @@ -106,6 +107,7 @@ struct EngineDependencies { // Test-only fault injection for allocation-sensitive durable error construction. EngineStepErrorFactory make_step_error{}; std::shared_ptr snapshot_cache; + std::optional backend_capabilities; }; void ValidateMtpModelCompatibility(const Config& config, @@ -140,7 +142,7 @@ struct Engine : std::enable_shared_from_this, * and its components. * * Assembles the Engine's scheduler, cache manager, and model executor for the model via - * CreateDependencies and delegates to the dependency-injecting constructor. Behavior is identical + * CreateEngineBackend and delegates to the dependency-injecting constructor. Behavior is identical * to constructing those collaborators inline. */ Engine(std::shared_ptr model); @@ -295,6 +297,7 @@ struct Engine : std::enable_shared_from_this, std::shared_ptr snapshot_cache_; std::unique_ptr scheduler_; // The scheduler responsible for managing execution order. std::unique_ptr model_executor_; // The executor responsible for running the model. + BackendCapabilities backend_capabilities_; // Lazily created on the first stop-enabled BeginTurn (the no-stop fast path never touches this). // Shared by every Request's StopStringController so the tokenizer's underlying vocabulary/config // is loaded once per Engine rather than once per Request. diff --git a/src/engine/executor_result.h b/src/engine/executor_result.h new file mode 100644 index 0000000000..b42d2d0e2f --- /dev/null +++ b/src/engine/executor_result.h @@ -0,0 +1,98 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include +#include +#include + +#include "request_status.h" +#include "step_plan.h" + +namespace Generators { + +enum class StepOutputKind { DecoderOutputs, + BackendRows }; + +struct SlotHandle { + uint32_t index{}; + uint64_t epoch{}; + bool operator==(const SlotHandle& other) const { return index == other.index && epoch == other.epoch; } +}; + +struct CompletionKey { + uint64_t context_id{}; + uint64_t lease_epoch{}; + SlotHandle slot; + uint64_t turn_id{}; + uint64_t submission_id{}; + bool operator==(const CompletionKey& other) const { + return context_id == other.context_id && lease_epoch == other.lease_epoch && + slot == other.slot && turn_id == other.turn_id && submission_id == other.submission_id; + } +}; + +// Raw completed backend work, distinct from API-visible speculative accounting. Decode target forwards +// count once per underlying batch, even when every row's delivery is discarded. Prefill work is tokens +// actually reprocessed, including replay; no drafter-forward count is inferred from proposal width. +struct BackendWorkMetrics { + uint64_t work_records{}; + uint64_t target_forward_passes{}; + uint64_t prefill_tokens{}; + uint64_t discarded_completions{}; +}; + +struct SpeculativeRowMetrics { + size_t proposed{}; + size_t accepted{}; +}; + +// The reference keeps backend continuation resources alive through settlement. Engine never +// interprets its payload and does not assume that it is a restorable state snapshot. +using BackendContinuationCheckpoint = std::shared_ptr; + +struct SelectedTokenBlock { + std::array candidates{}; + size_t candidate_count{}; + SpeculativeRowMetrics speculation; + BackendContinuationCheckpoint checkpoint; +}; + +struct PrefillProgress { + size_t prompt_tokens_processed_this_step{}; + bool prompt_complete{}; +}; + +struct BackendRowResult { + CompletionKey key; + std::optional prefill; + std::optional selected_tokens; +}; + +// Only participating rows appear in this plan; an admitted resident may be idle at a phase barrier. +struct BackendPlannedRow { + size_t request_row{}; + CompletionKey key; + size_t max_prompt_tokens{}; + bool allow_phase_completion{}; +}; + +enum class ContinuationDisposition { Continue, + Invalidate, + Release }; + +struct TokenCommitAck { + CompletionKey key; + size_t candidates_consumed{}; + size_t tokens_appended{}; + size_t tokens_staged_for_events{}; + uint64_t request_sequence_length_after{}; + bool turn_finished{}; + GenerationFinishReason finish_reason{GenerationFinishReason::None}; + ContinuationDisposition disposition{ContinuationDisposition::Continue}; +}; + +} // namespace Generators diff --git a/src/engine/fixed_slot_scheduler.cpp b/src/engine/fixed_slot_scheduler.cpp new file mode 100644 index 0000000000..436634b075 --- /dev/null +++ b/src/engine/fixed_slot_scheduler.cpp @@ -0,0 +1,134 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#include "fixed_slot_scheduler.h" + +#include +#include + +namespace Generators { + +// ------------------------------------------------------------------------------------------------ +// Cache manager +// ------------------------------------------------------------------------------------------------ +FixedSlotCacheManager::FixedSlotCacheManager(std::shared_ptr model, std::shared_ptr session) + : CacheManager{std::move(model)}, + session_{std::move(session)}, + max_slots_{session_ ? session_->MaxBatchSize() : 1}, + slots_(max_slots_, nullptr) { + if (!session_) { + throw std::invalid_argument("FixedSlotCacheManager requires a session."); + } +} + +bool FixedSlotCacheManager::CanAllocate(const std::vector>& requests) const { + if (requests.empty() || requests.size() > max_slots_) { + return false; + } + // A slot is free, or holds a row whose turn is over and can be replaced. + size_t available = 0; + for (const auto& resident : slots_) { + if (!resident || IsTurnComplete(resident->status_) || IsClosed(resident->status_)) { + ++available; + } + } + return requests.size() <= available; +} + +void FixedSlotCacheManager::Allocate(const std::vector>& requests) { + if (!CanAllocate(requests)) { + throw std::runtime_error(max_slots_ == 1 ? "The fixed-slot backend serves one Request at a time." + : "The fixed-slot backend cannot allocate the requested batch."); + } + for (const auto& req : requests) { + if (IsResident(req)) { + continue; + } + size_t target_slot = max_slots_; + for (size_t s = 0; s < max_slots_; ++s) { + if (!slots_[s] || IsClosed(slots_[s]->status_)) { + target_slot = s; + break; + } + } + if (target_slot == max_slots_) { + for (size_t s = 0; s < max_slots_; ++s) { + if (IsTurnComplete(slots_[s]->status_)) { + target_slot = s; + break; + } + } + } + if (target_slot == max_slots_) { + throw std::runtime_error("The fixed-slot backend has no free slot."); + } + if (slots_[target_slot]) { + session_->ReleaseSlotForTeardown(target_slot); + } + slots_[target_slot] = req; + } +} + +void FixedSlotCacheManager::Deallocate(std::vector>& requests) { + if (requests.empty()) { + throw std::runtime_error("Cannot dynamically deallocate statically batched requests."); + } + for (const auto& req : requests) { + if (!IsResident(req)) { + throw std::runtime_error("Cannot dynamically deallocate statically batched requests."); + } + } + for (const auto& req : requests) { + const size_t s = SlotOf(req); + session_->ReleaseSlot(s); + slots_[s].reset(); + } +} + +void FixedSlotCacheManager::ValidateRewind(const std::shared_ptr&) const { + // Each resident row owns its dedicated slot, so releasing it for a replay never affects a peer. +} + +void FixedSlotCacheManager::ReleaseForRewind(const std::shared_ptr& request) { + if (!IsResident(request)) { + return; + } + const size_t s = SlotOf(request); + session_->ReleaseSlot(s); + slots_[s].reset(); +} + +void FixedSlotCacheManager::DetachRequestForTeardown(const std::shared_ptr& request) noexcept { + if (!IsResident(request)) { + return; + } + const size_t s = SlotOf(request); + session_->ReleaseSlotForTeardown(s); + slots_[s].reset(); +} + +std::vector> FixedSlotCacheManager::AllocatedRequests() const { + std::vector> allocated; + allocated.reserve(max_slots_); + for (const auto& req : slots_) { + if (req) { + allocated.push_back(req); + } + } + return allocated; +} + +bool FixedSlotCacheManager::IsResident(const std::shared_ptr& request) const { + return request && std::find(slots_.begin(), slots_.end(), request) != slots_.end(); +} + +size_t FixedSlotCacheManager::ResidentRequestCount() const { + return static_cast(std::count_if(slots_.begin(), slots_.end(), [](const auto& req) { return req != nullptr; })); +} + +size_t FixedSlotCacheManager::SlotOf(const std::shared_ptr& request) const { + const auto it = std::find(slots_.begin(), slots_.end(), request); + if (it == slots_.end()) throw std::logic_error("Request has no fixed-slot residency lease."); + return static_cast(it - slots_.begin()); +} + +} // namespace Generators diff --git a/src/engine/fixed_slot_scheduler.h b/src/engine/fixed_slot_scheduler.h new file mode 100644 index 0000000000..6b15338f94 --- /dev/null +++ b/src/engine/fixed_slot_scheduler.h @@ -0,0 +1,153 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#pragma once + +#include +#include + +#include "cache_manager.h" +#include "scheduler.h" + +namespace Generators { + +// Fixed storage lifecycle only. Implementations retain graph bindings, buffers, and phase constraints. +struct FixedSlotSession { + virtual ~FixedSlotSession() = default; + virtual size_t MaxRequestLength() const = 0; + virtual size_t MaxBatchSize() const = 0; + // Quiesce all work referencing a slot and inspect pending status before returning it for reuse. + virtual void ReleaseSlot(size_t slot) = 0; + virtual void ReleaseSlotForTeardown(size_t slot) noexcept = 0; +}; + +// Static-cache contract for 1..MaxBatchSize() slots backed by a FixedSlotSession. +struct FixedSlotCacheManager final : CacheManager { + FixedSlotCacheManager(std::shared_ptr model, std::shared_ptr session); + + bool CanAllocate(const std::vector>& requests) const override; + void Allocate(const std::vector>& requests) override; + void Step() override {} + void Deallocate(std::vector>& requests) override; + void ValidateRewind(const std::shared_ptr& request) const override; + void ReleaseForRewind(const std::shared_ptr& request) override; + void DetachRequestForTeardown(const std::shared_ptr& request) noexcept override; + bool SupportsDynamicBatching() const override { return false; } + bool SupportsContinuousMultiSlotDecoding() const override { return max_slots_ > 1; } + size_t MaxBatchSize() const override { return max_slots_; } + size_t MaxStaticBatchSize() const override { return max_slots_; } + size_t MaxStaticRequestLength() const override { return session_->MaxRequestLength(); } + std::vector> AllocatedRequests() const override; + bool IsResident(const std::shared_ptr& request) const override; + size_t ResidentRequestCount() const override; + size_t SlotOf(const std::shared_ptr& request) const; + + private: + std::shared_ptr session_; + size_t max_slots_{1}; + std::vector> slots_; +}; + +struct FixedSlotScheduler final : Scheduler { + FixedSlotScheduler(std::shared_ptr model, std::shared_ptr cache_manager) + : Scheduler{model}, model_{std::move(model)}, cache_manager_{std::move(cache_manager)} {} + + bool SupportsIndependentSlotAdmission() const override { + return cache_manager_->SupportsContinuousMultiSlotDecoding(); + } + + void AddRequest(std::shared_ptr request) override { + if (request->PrefillChunkSize().value_or(0) != 0) { + throw std::runtime_error( + "search.chunk_size requires dynamic batching; the static batch scheduler cannot chunk a " + "prefill."); + } + if (std::find(requests_pool_.begin(), requests_pool_.end(), request) != requests_pool_.end()) { + return; + } + requests_pool_.reserve(requests_pool_.size() + 1); + auto sampling_state = CreateSamplingState(*request); + request->CommitSamplingState(std::move(sampling_state)); + requests_pool_.push_back(std::move(request)); + } + + void RemoveRequest(std::shared_ptr request) override { + if (cache_manager_->IsResident(request)) { + std::vector> to_deallocate{request}; + cache_manager_->Deallocate(to_deallocate); + } + Forget(request); + } + + void DetachRequestForTeardown(const std::shared_ptr& request) noexcept override { + cache_manager_->DetachRequestForTeardown(request); + Forget(request); + } + + ScheduledRequests Schedule() override { + for (const auto& req : cache_manager_->AllocatedRequests()) { + if (IsClosed(req->status_)) { + cache_manager_->DetachRequestForTeardown(req); + Forget(req); + } else if (IsQueued(req->status_)) { + req->Schedule(); + } + } + + std::vector> requests_to_schedule; + for (const auto& request : requests_pool_) { + if (IsQueued(request->status_) && !cache_manager_->IsResident(request)) { + requests_to_schedule.push_back(request); + } + } + + for (const auto& request : requests_to_schedule) { + if (!cache_manager_->CanAllocate({request})) { + break; + } + if (cache_manager_->ResidentRequestCount() >= cache_manager_->MaxBatchSize()) { + for (const auto& resident : cache_manager_->AllocatedRequests()) { + if (IsTurnComplete(resident->status_) || IsClosed(resident->status_)) { + cache_manager_->DetachRequestForTeardown(resident); + if (!StaysEnrolledWhenDisplaced(*resident)) Forget(resident); + break; + } + } + } + cache_manager_->Allocate({request}); + request->Schedule(); + } + + ScheduledRequests scheduled_requests(cache_manager_->AllocatedRequests(), model_, GetBatchedSampler(), + GetBatchedSamplingPlan()); + return scheduled_requests; + } + + NontransactionalSchedulingResult ScheduleStep() override { + auto scheduled = Schedule(); + if (!scheduled) { + return {{HasPendingRequests() ? StepOutcomeKind::CapacityDeferred : StepOutcomeKind::NoWork, 0, nullptr}, + std::nullopt}; + } + return {{StepOutcomeKind::Committed, 0, nullptr}, std::move(scheduled)}; + } + + bool HasPendingRequests() const override { + for (const auto& request : requests_pool_) { + if (IsExecutable(request->status_)) { + return true; + } + } + return false; + } + + private: + void Forget(const std::shared_ptr& request) { + requests_pool_.erase(std::remove(requests_pool_.begin(), requests_pool_.end(), request), requests_pool_.end()); + } + + std::shared_ptr model_; + std::shared_ptr cache_manager_; + std::vector> requests_pool_; +}; + +} // namespace Generators diff --git a/src/engine/model_executor.h b/src/engine/model_executor.h index fde9ce55db..b8ff84e3c8 100644 --- a/src/engine/model_executor.h +++ b/src/engine/model_executor.h @@ -7,12 +7,14 @@ #include #include +#include "../decoding/speculative_stats.h" #include "../models/decoder_only.h" #include "decoders/decoder.h" #include "scheduled_requests.h" #include "cache_manager.h" #include "execution_context.h" #include "step_plan.h" +#include "turn_policy.h" namespace Generators { @@ -54,6 +56,36 @@ struct ModelExecutor { virtual bool SupportsDraftVerification() const { return false; } + /** + * @brief True when Decode chooses the next tokens itself instead of attaching decoder state with logits. + * + * Such an executor (for example one whose graph selects greedy tokens and verifies its own speculative + * drafts on the device) declares its step with ScheduledRequests::BeginBackendStep and hands each row's tokens to + * ScheduledRequests::StageBackendRow; ScheduledRequests::GenerateNextTokens validates them and commits them through + * Request::CommitExternalTokens instead of sampling. Stop strings are valid for such a turn on the static path: + * committing is token-wise and a match only ends the turn, so nothing has to be rolled back. + */ + virtual bool SelectsOwnTokens() const { return false; } + + /** + * @brief Rejects, before any Request mutation, a turn this executor cannot serve. + * + * Called by Engine::BeginTurn after the generic policy validation. Throw std::runtime_error with a + * message that names the unsupported option. + */ + virtual void ValidateTurn(const EffectiveTurnPolicy& /*policy*/, const TurnOptions& /*options*/) const {} + + /** @brief Adds this executor's own speculative-decoding counters to `stats` (Engine::GetSpeculativeStats). */ + virtual void AddSpeculativeStats(SpeculativeStats& /*stats*/) const {} + + /** @brief Notifies an executor that selects its own tokens of the committed step outcomes. */ + virtual void OnStepCommitted(const ScheduledRequests& /*scheduled_requests*/, + std::span /*results*/) {} + + // A nontransactional failure may have changed device state. Mutable contexts must remain + // poisoned after this Engine releases its lease. + virtual void OnExecutionFailure() noexcept {} + virtual ~ModelExecutor() = default; }; diff --git a/src/engine/qwen38_graph_executor.cpp b/src/engine/qwen38_graph_executor.cpp new file mode 100644 index 0000000000..2da0b42c3f --- /dev/null +++ b/src/engine/qwen38_graph_executor.cpp @@ -0,0 +1,427 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "qwen38_graph_executor.h" + +#include +#include +#include + +#include "../models/qwen38_graph.h" +#include "scheduler.h" + +namespace Generators { + +// ------------------------------------------------------------------------------------------------ +// Session +// ------------------------------------------------------------------------------------------------ +Qwen38GraphSession::Qwen38GraphSession(std::shared_ptr device, Options options) + : device_{std::move(device)}, options_{std::move(options)} { + if (!device_) { + throw std::invalid_argument("Qwen38GraphSession requires a device."); + } + if (!device_->ContextIsHealthy()) { + throw std::runtime_error("Cannot lease a poisoned token-producing execution context."); + } + // S=2 target steps consume token pairs: an odd chunk would shift the pairing of everything after it. + options_.prefill_chunk_tokens = std::max(2, options_.prefill_chunk_tokens & ~size_t{1}); + // A device with a fast prefill names its own chunk (any size is valid there); a longer chunk is one step without events, + // so the preferred size also bounds how long cancellation and Run wait. + if (const size_t preferred = device_->PrefillChunkTokens(); preferred > 0) { + options_.prefill_chunk_tokens = preferred; + } +} + +Qwen38GraphSession::~Qwen38GraphSession() { + ReleaseForTeardown(); +} + +size_t Qwen38GraphSession::MaxRequestLength() const { + return static_cast(device_->Capacity() - device_->ReservedTail()); +} + +void Qwen38GraphSession::Poison() noexcept { + device_->MarkContextFailed(); +} + +void Qwen38GraphSession::Acknowledge(const TokenCommitAck& ack) { + completions_.Acknowledge(ack); +} + +BackendPlannedRow Qwen38GraphSession::Plan(size_t request_row, const Qwen38StepRequest& request) { + if (request.tokens.empty() || request.tokens.size() > MaxRequestLength() || + std::min(request.token_budget, request.session_budget) == 0) { + throw std::logic_error("Fixed-slot request is outside the negotiated capacity or token budget."); + } + const bool same_turn = request.owner == owner_ && request.turn_id == turn_id_; + const bool needs_prefill = !same_turn || !decoding_; + const size_t prefilled = same_turn ? prefilled_ : 0; + const size_t remaining = needs_prefill ? request.tokens.size() - prefilled : 0; + const size_t chunk = std::min(options_.prefill_chunk_tokens, remaining); + return completions_.Plan(request_row, 0, request.owner, request.turn_id, chunk, needs_prefill && chunk == remaining); +} + +bool Qwen38GraphSession::IsStopToken(int32_t token) const { + return std::find(options_.stop_tokens.begin(), options_.stop_tokens.end(), token) != options_.stop_tokens.end(); +} + +void Qwen38GraphSession::Begin(const Qwen38StepRequest& request) { + if (request.tokens.empty()) { + throw std::logic_error("A Request must have at least one token before it can be decoded."); + } + if (options_.vocab_size > 0) { + ValidateQwen38PromptTokens(request.tokens, options_.vocab_size); + } + // A cycle may still be in flight from a turn that stopped early (stop string, cancellation); drain and check its + // decode status before Reset clears the device error bits. + owner_ = nullptr; + pending_step_ = {}; + if (deferred_error_) { + std::rethrow_exception(std::exchange(deferred_error_, nullptr)); + } + if (!status_checked_) { + device_->Drain(); + status_checked_ = true; + device_->CheckDecodeStatus(); + } + if (device_dirty_) { + device_->Reset(); + device_dirty_ = false; + } + owner_ = request.owner; + turn_id_ = request.turn_id; + prefilled_ = 0; + decoding_ = false; +} + +void Qwen38GraphSession::Release() { + device_->Drain(); + owner_ = nullptr; + prefilled_ = 0; + decoding_ = false; + pending_step_ = {}; + completions_.Retire(0); + if (deferred_error_) { + // Only a step consumes the failure (the Engine then turns it into a fatal error). A retried rewind or retirement must + // keep failing: clearing it here would let a later retry succeed on a poisoned device. + std::rethrow_exception(deferred_error_); + } + if (!status_checked_) { + status_checked_ = true; + try { + device_->CheckDecodeStatus(); + } catch (...) { + // The check also runs outside a step (a rewind, the retirement of a Request), where RunStatic's fatal handler does + // not see the error. The next Reset() erases the error bits, so keep the failure for the next step, as + // ReleaseForTeardown() does, and mark the context failed: the caller still gets the exception. + Poison(); + if (!deferred_error_) deferred_error_ = std::current_exception(); + throw; + } + } +} + +void Qwen38GraphSession::ReleaseForTeardown() noexcept { + try { + Release(); + } catch (...) { + Poison(); + if (!deferred_error_) { + deferred_error_ = std::current_exception(); + } + } +} + +void Qwen38GraphSession::ReleaseSlot(size_t slot) { + if (slot != 0) throw std::out_of_range("Fixed-slot retirement index out of range."); + Release(); +} + +void Qwen38GraphSession::ReleaseSlotForTeardown(size_t slot) noexcept { + if (slot != 0) return; + ReleaseForTeardown(); +} + +Qwen38StepResult Qwen38GraphSession::Step(const Qwen38StepRequest& request) { + if (deferred_error_) { + std::rethrow_exception(std::exchange(deferred_error_, nullptr)); + } + if (request.owner != owner_ || request.turn_id != turn_id_) { + Begin(request); + } + const size_t limit = std::min(request.token_budget, request.session_budget); + if (limit == 0) { + throw std::logic_error("A request with no remaining budget was scheduled for decoding."); + } + + Qwen38StepResult result; + pending_step_ = {}; + if (!decoding_) { + const size_t remaining = request.tokens.size() - prefilled_; + const size_t chunk = std::min(options_.prefill_chunk_tokens, remaining); + device_dirty_ = true; + if (prefilled_ == 0) { + // The device keeps the turn's prompt length for its chunks (Reset() and BeginDecode() drop it): the value belongs + // to this device, so Engines running in one process never see each other's. + device_->ExpectPrefillTokens(request.tokens.size()); + } + device_->PrefillChunk(request.tokens.subspan(prefilled_, chunk)); + prefilled_ += chunk; + ++work_metrics_.work_records; + work_metrics_.prefill_tokens += chunk; + result.prompt_tokens_processed = chunk; + result.prompt_complete = prefilled_ == request.tokens.size(); + if (prefilled_ < request.tokens.size()) { + return result; + } + // The first token is the target's prediction after the prompt. The head prefill, the first draft and the first verify + // are enqueued before it is returned so that the next step finds the GPU already busy. + const int32_t first = device_->PendingToken(); + const bool run_cycle = !IsStopToken(first) && limit > 1; + device_->BeginDecode(run_cycle); + if (run_cycle) { + status_checked_ = false; + } + decoding_ = true; + result.tokens[0] = first; + result.count = 1; + pending_step_ = PendingStep{/*active=*/true, /*is_decode_cycle=*/false, /*mtp=*/false, + /*proposed=*/0, /*accepted=*/0, /*produced=*/1}; + return result; + } + + const Qwen38CycleRecord record = device_->WaitCycle(); + ++work_metrics_.work_records; + ++work_metrics_.target_forward_passes; + const std::array predictions{ + static_cast(record.pred0), + static_cast(record.pred1), + static_cast(record.pred2), + static_cast(record.pred3), + }; + const size_t produced = record.TokenCount(); + if (produced == 0 || produced > predictions.size() || record.accepted_count < 0 || + record.proposed_count < 0 || record.proposed_count > 3 || + (record.accepted_count > 0 && record.accepted_count > record.proposed_count)) { + throw std::logic_error("Qwen backend returned a malformed cycle record."); + } + const std::array raw{record.pred0, record.pred1, record.pred2, record.pred3}; + for (size_t i = 0; i < produced; ++i) { + if (raw[i] < 0 || raw[i] > std::numeric_limits::max()) { + throw std::logic_error("Qwen backend returned an invalid token representation."); + } + } + size_t count = produced; + bool stop_in_cycle = false; + for (size_t i = 0; i < count; ++i) { + if (IsStopToken(predictions[i])) { + count = i + 1; + stop_in_cycle = true; + break; + } + } + const size_t committed = std::min(count, limit); + // A stop token beyond the limit is never committed, so it does not end the loop by itself. + const bool ended = (stop_in_cycle && committed == count) || committed == limit; + // Enqueue the next cycle before returning: the GPU runs it while the Engine delivers these tokens. A stop string can + // only be known after delivery, so one cycle may be wasted; Reset/Release wait for it. + device_->Advance(record, /*want_more=*/!ended); + if (ended) { + status_checked_ = true; + device_->CheckDecodeStatus(); // nothing is in flight now: a device that accumulated an error reports it here + } + + for (size_t i = 0; i < produced; ++i) { + result.tokens[i] = predictions[i]; + } + result.count = produced; + const size_t n_prop = record.proposed_count > 1 ? static_cast(record.proposed_count) : 1; + const size_t n_acc = + record.accepted_count > 0 ? static_cast(record.accepted_count) : (record.accepted ? 1 : 0); + result.speculation = SpeculativeRowMetrics{n_prop, n_acc}; + pending_step_ = PendingStep{/*active=*/true, /*is_decode_cycle=*/true, /*mtp=*/device_->MtpEnabled(), + /*proposed=*/n_prop, /*accepted=*/n_acc, /*produced=*/produced}; + return result; +} + +void Qwen38GraphSession::OnStepCommitted(const RequestStepResult& result) { + if (!pending_step_.active) { + return; + } + const PendingStep step = pending_step_; + pending_step_ = {}; + const size_t emitted = result.visible_token_count; + stats_.tokens_emitted += emitted; + if (!step.is_decode_cycle) { + return; + } + + ++stats_.rounds; + ++stats_.completed_rounds; + ++stats_.target_forward_passes; + if (step.mtp) { + const size_t n_prop = step.proposed > 0 ? step.proposed : 1; + const size_t effective_accepted = std::min(step.accepted, emitted); + // Count examined draft positions, including a terminal EOS which Request consumes without appending. + const size_t consumed = result.candidates_consumed; + const size_t n_eval = + std::min(consumed, step.accepted) + ((consumed > step.accepted && step.accepted < n_prop) ? 1 : 0); + ++stats_.target_verify_forward_passes; + stats_.draft_forward_passes += n_prop; + stats_.draft_tokens_proposed += n_prop; + stats_.draft_tokens_evaluated += n_eval; + if (effective_accepted < stats_.acceptance_length_histogram.size()) { + ++stats_.acceptance_length_histogram[effective_accepted]; + } + stats_.draft_tokens_accepted += effective_accepted; + // The same classification as Engine::RecordSpeculativeCommit, by the drafts that were delivered; whether a correction + // or bonus token followed is counted on its own. + if (effective_accepted == 0) { + ++stats_.zero_accept_rounds; + } else if (effective_accepted == n_prop) { + ++stats_.full_accept_rounds; + } else { + ++stats_.partial_accept_rounds; + } + if (effective_accepted == n_prop) { + if (emitted > n_prop) { + ++stats_.bonus_tokens; + } + } else if (emitted > effective_accepted) { + ++stats_.correction_tokens; + } + } + if (step.produced >= emitted) { + stats_.tokens_discarded += step.produced - emitted; + } +} + +void Qwen38GraphSession::AddStats(SpeculativeStats& stats) const { + const auto add = [](size_t& total, size_t value) { total += value; }; + add(stats.rounds, stats_.rounds); + add(stats.completed_rounds, stats_.completed_rounds); + add(stats.draft_tokens_proposed, stats_.draft_tokens_proposed); + add(stats.draft_tokens_evaluated, stats_.draft_tokens_evaluated); + add(stats.draft_tokens_accepted, stats_.draft_tokens_accepted); + add(stats.correction_tokens, stats_.correction_tokens); + add(stats.bonus_tokens, stats_.bonus_tokens); + add(stats.tokens_emitted, stats_.tokens_emitted); + add(stats.tokens_discarded, stats_.tokens_discarded); + add(stats.draft_forward_passes, stats_.draft_forward_passes); + add(stats.target_forward_passes, stats_.target_forward_passes); + add(stats.target_verify_forward_passes, stats_.target_verify_forward_passes); + add(stats.full_accept_rounds, stats_.full_accept_rounds); + add(stats.partial_accept_rounds, stats_.partial_accept_rounds); + add(stats.zero_accept_rounds, stats_.zero_accept_rounds); + for (size_t i = 0; i < stats.acceptance_length_histogram.size(); ++i) { + add(stats.acceptance_length_histogram[i], stats_.acceptance_length_histogram[i]); + } +} + +// ------------------------------------------------------------------------------------------------ +// Executor +// ------------------------------------------------------------------------------------------------ +void Qwen38GraphExecutor::ValidateTurn(const EffectiveTurnPolicy& policy, const TurnOptions& options) const { + constexpr const char* prefix = "The qwen38_flash_graph backend "; + if (!policy.IsGreedy()) { + throw std::runtime_error(std::string{prefix} + + "decodes greedily only (its device loop verifies drafts by argmax); set do_sample to false for the turn."); + } + if (policy.repetition_penalty != 1.0f) { + throw std::runtime_error(std::string{prefix} + "does not support repetition_penalty."); + } + if (policy.no_repeat_ngram_size != 0) { + throw std::runtime_error(std::string{prefix} + "does not support no_repeat_ngram_size."); + } + if (policy.min_generated_tokens != 0) { + throw std::runtime_error(std::string{prefix} + + "does not support min_generated_tokens (the device loop cannot mask the end-of-sequence token)."); + } + if (options.seed) { + throw std::runtime_error(std::string{prefix} + "does not support a per-turn seed (decoding is greedy)."); + } + if (!options.guidance_type.empty() || !options.guidance_data.empty() || options.delimited_guidance) { + throw std::runtime_error(std::string{prefix} + "does not support constrained decoding (guidance)."); + } +} + +void Qwen38GraphExecutor::Decode(ScheduledRequests& scheduled_requests, ExecutionContext&) { + for (size_t row = 0; row < scheduled_requests.size(); ++row) { + const auto& request = scheduled_requests[row]; + if (!IsExecuting(request->status_)) { + continue; + } + const int64_t used = request->CurrentSequenceLength(); + const size_t session_limit = request->MaxSessionTokens(); + Qwen38StepRequest step; + step.owner = request.get(); + step.turn_id = request->CurrentTurnId(); + step.tokens = request->TokensCpu(); + step.token_budget = request->RemainingTurnTokenBudget(); + step.session_budget = static_cast(used) < session_limit ? session_limit - static_cast(used) : 0; + const auto operation = session_->Plan(row, step); + scheduled_requests.BeginBackendStep({operation}); + const Qwen38StepResult result = session_->Step(step); + BackendRowResult backend_row; + backend_row.key = operation.key; + if (result.prompt_tokens_processed > 0 || result.prompt_complete) { + backend_row.prefill = PrefillProgress{result.prompt_tokens_processed, result.prompt_complete}; + } + if (result.count > 0) { + SelectedTokenBlock tokens; + if (result.count > tokens.candidates.size()) + throw std::logic_error("Backend candidate block exceeds its negotiated width."); + tokens.candidate_count = result.count; + tokens.speculation = result.speculation; + std::copy_n(result.tokens.begin(), result.count, tokens.candidates.begin()); + backend_row.selected_tokens = std::move(tokens); + } + session_->Collected(backend_row); + scheduled_requests.StageBackendRow(row, std::move(backend_row)); + return; // one resident row + } + scheduled_requests.BeginBackendStep({}); +} + +void Qwen38GraphExecutor::OnStepCommitted(const ScheduledRequests& scheduled_requests, + std::span results) { + for (const auto& ack : scheduled_requests.CommitAcks()) { + session_->Acknowledge(ack); + } + for (size_t row = 0; row < scheduled_requests.size() && row < results.size(); ++row) { + if (scheduled_requests[row].get() == session_->Owner()) { + session_->OnStepCommitted(results[row]); + return; + } + } +} + +// ------------------------------------------------------------------------------------------------ +// Assembly +// ------------------------------------------------------------------------------------------------ +EngineDependencies CreateQwen38GraphDependencies(std::shared_ptr model, + std::shared_ptr device) { + const auto& config = *model->config_; + if (config.engine.dynamic_batching) { + throw std::runtime_error( + "The qwen38_flash_graph model type runs on the Engine's static single-slot path; engine.dynamic_batching is not supported."); + } + const auto& qwen38 = config.model.qwen38_graph; + ValidateQwen38PrefillChunk(qwen38); + Qwen38GraphSession::Options options; + options.prefill_chunk_tokens = EffectiveQwen38PrefillChunkTokens(qwen38); + options.stop_tokens.assign(config.model.eos_token_id.begin(), config.model.eos_token_id.end()); + options.vocab_size = config.model.vocab_size; + auto session = std::make_shared(std::move(device), std::move(options)); + auto cache_manager = std::make_shared(model, session); + auto scheduler = std::make_unique(model, cache_manager); + auto executor = std::make_unique(std::move(session)); + return EngineDependencies{std::move(cache_manager), std::move(scheduler), std::move(executor)}; +} + +EngineDependencies CreateQwen38GraphDependencies(std::shared_ptr model) { + auto device = model->AcquireDevice(); + return CreateQwen38GraphDependencies(std::shared_ptr{std::move(model)}, std::move(device)); +} + +} // namespace Generators diff --git a/src/engine/qwen38_graph_executor.h b/src/engine/qwen38_graph_executor.h new file mode 100644 index 0000000000..0787fbed84 --- /dev/null +++ b/src/engine/qwen38_graph_executor.h @@ -0,0 +1,151 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include +#include + +#include "../decoding/speculative_stats.h" +#include "../models/qwen38_decode_device.h" +#include "../span.h" +#include "engine.h" +#include "fixed_slot_scheduler.h" +#include "token_producing_session.h" + +/** + * @file qwen38_graph_executor.h + * @brief Hosts the Qwen3.8-Flash CUDA-graph decode loop (MTP(1) draft -> verify -> commit cycles on the device) in the + * Engine's static path. + * + * The loop owns a fixed in-place KV store, a recurrent-state bank ring and a device-side accept/commit step, so it is a + * single-slot backend: one resident Request, greedy decoding, no prefix cache. Reusable residency lives in + * FixedSlotCacheManager; completion identity and settlement live in TokenProducingSession. + */ + +namespace Generators { + +struct Qwen38GraphModel; + +// What the session needs to know about the Request it serves, captured at the start of each Engine step. +struct Qwen38StepRequest { + const void* owner{}; // the Request, as an identity + uint64_t turn_id{}; // changes whenever the Request begins a new turn + std::span tokens; // the Request's whole token history (prompt of every turn plus generated tokens) + size_t token_budget{}; // generated tokens the turn may still produce + size_t session_budget{}; // tokens the session may still hold +}; + +struct Qwen38StepResult { + std::array tokens{}; + size_t count{}; // 0 for a step that only advanced the prompt + size_t prompt_tokens_processed{}; + bool prompt_complete{}; + SpeculativeRowMetrics speculation; +}; + +class Qwen38GraphSession : public FixedSlotSession { + public: + struct Options { + size_t prefill_chunk_tokens{2048}; // S=2 step prefill: rounded down to an even number (pairs stay aligned), at least 2; + // replaced by Qwen38DecodeDevice::PrefillChunkTokens() when the device names one + // (the assembly passes EffectiveQwen38PrefillChunkTokens of the configuration) + std::vector stop_tokens; // EOS ids of the model: a cycle is not followed by another one after one of them + int vocab_size{0}; // > 0: validates all request token ids at turn start before mutating device state + }; + + Qwen38GraphSession(std::shared_ptr device, Options options); + ~Qwen38GraphSession() override; + Qwen38GraphSession(const Qwen38GraphSession&) = delete; + Qwen38GraphSession& operator=(const Qwen38GraphSession&) = delete; + + // Longest session (prompt plus generated tokens) the device can hold. + size_t MaxRequestLength() const override; + size_t MaxBatchSize() const override { return 1; } + + // Runs one step for `request`: a new owner or turn first clears the device, announces the history length to it + // (Qwen38DecodeDevice::ExpectPrefillTokens) and starts ingesting the whole history; the step that ingests its last + // chunk also returns the first generated token and enqueues the first cycle; every later step waits for the cycle in + // flight, commits it and enqueues the next one before returning its tokens. + Qwen38StepResult Step(const Qwen38StepRequest& request); + + BackendPlannedRow Plan(size_t request_row, const Qwen38StepRequest& request); + void Collected(const BackendRowResult& result) { completions_.Collected(result); } + void Acknowledge(const TokenCommitAck& ack); + void Poison() noexcept; + + // Records the committed token counts after Request::CommitExternalTokens applies EOS, limits, and stop strings. + void OnStepCommitted(const RequestStepResult& result); + + const void* Owner() const noexcept { return owner_; } + + // The owner leaves (finished, rewound, replaced, or deallocated): drains work in flight, checks any unchecked decode + // status, and forgets the owner. + void Release(); + + // No-throw variant used by destructors and teardown: captures any status-check failure to report on the next operation. + void ReleaseForTeardown() noexcept; + + void ReleaseSlot(size_t slot) override; + void ReleaseSlotForTeardown(size_t slot) noexcept override; + + void AddStats(SpeculativeStats& stats) const; + const BackendWorkMetrics& WorkMetrics() const { return work_metrics_; } + + private: + struct PendingStep { + bool active{}; + bool is_decode_cycle{}; + bool mtp{}; + size_t proposed{}; + size_t accepted{}; + size_t produced{}; + }; + + bool IsStopToken(int32_t token) const; + void Begin(const Qwen38StepRequest& request); + + std::shared_ptr device_; + Options options_; + const void* owner_{}; + uint64_t turn_id_{}; + size_t prefilled_{}; + bool decoding_{}; + // Unknown until proven otherwise: the device may hold a previous Engine's state (it outlives its Engine), so the first + // request always starts with a Reset. + bool device_dirty_{true}; + bool status_checked_{true}; + std::exception_ptr deferred_error_; + PendingStep pending_step_{}; + SpeculativeStats stats_; + TokenProducingSession completions_{1}; + BackendWorkMetrics work_metrics_; +}; + +using Qwen38GraphCacheManager = FixedSlotCacheManager; + +struct Qwen38GraphExecutor final : ModelExecutor { + explicit Qwen38GraphExecutor(std::shared_ptr session) : session_{std::move(session)} {} + + void Decode(ScheduledRequests& scheduled_requests, ExecutionContext& context) override; + bool SelectsOwnTokens() const override { return true; } + void ValidateTurn(const EffectiveTurnPolicy& policy, const TurnOptions& options) const override; + void AddSpeculativeStats(SpeculativeStats& stats) const override { session_->AddStats(stats); } + void OnExecutionFailure() noexcept override { session_->Poison(); } + void OnStepCommitted(const ScheduledRequests& scheduled_requests, + std::span results) override; + + private: + std::shared_ptr session_; +}; + +// Assembles the Engine's collaborators for a qwen38_flash_graph model (leases the model's device engine). +EngineDependencies CreateQwen38GraphDependencies(std::shared_ptr model); + +// The same assembly over an arbitrary device; used by tests with a scripted device. +EngineDependencies CreateQwen38GraphDependencies(std::shared_ptr model, + std::shared_ptr device); + +} // namespace Generators diff --git a/src/engine/request.cpp b/src/engine/request.cpp index d1e83bd3a5..59ed65cc33 100644 --- a/src/engine/request.cpp +++ b/src/engine/request.cpp @@ -1147,6 +1147,9 @@ void Request::CommitStep(const RequestStepPlan& plan, void Request::StageVisibleTokens(RequestStepResult& result, size_t committed_count, std::optional sampled_token) const { + result.tokens_appended = committed_count + static_cast(sampled_token.has_value()); + result.candidates_consumed = result.tokens_appended + + static_cast(result.finish_reason == GenerationFinishReason::EosToken); if (committed_count + (sampled_token ? 1u : 0u) > result.visible_tokens.size()) { throw std::logic_error( "A single engine step produced more visible tokens than it can report."); @@ -1388,6 +1391,78 @@ RequestStepResult Request::CompleteGeneration() { return result; } +RequestStepResult Request::CommitExternalTokens(std::span tokens) { + if (tokens.empty() || tokens.size() > kMaxGeneratedTokensPerStep) { + throw std::logic_error("An executor must commit between one and the per-step maximum of tokens."); + } + if (!IsExecuting(status_)) { + throw std::logic_error("Tokens can only be committed to an active request."); + } + // Everything before this step has been processed by the model, exactly as for sampled steps. + processed_sequence_length_ = CurrentSequenceLength(); + needs_replay_after_rewind_ = false; + + size_t appended_count = 0; + size_t consumed_count = 0; + int32_t last_token = 0; + GenerationFinishReason finish_reason = GenerationFinishReason::None; + int32_t matched_stop_string_index = -1; + for (const int32_t token : tokens) { + ++consumed_count; + last_token = token; + const int64_t sequence_length_before = CurrentSequenceLength(); + search_->CommitToken(token); + const bool appended = CurrentSequenceLength() > sequence_length_before; + if (appended) { + tokens_host_.push_back(token); + ++turn_generated_tokens_; + ++appended_count; + if (stop_controller_) { + if (const auto& match = stop_controller_->ObserveToken(token)) { + finish_reason = GenerationFinishReason::StopString; + matched_stop_string_index = static_cast(match->index); + break; + } + } + } + const bool search_done = search_->IsDone(); + const bool turn_limit_reached = + turn_policy_.max_generated_tokens && + turn_generated_tokens_ >= *turn_policy_.max_generated_tokens; + const bool context_limit_reached = + static_cast(CurrentSequenceLength()) >= max_session_tokens_; + if (search_done && contains(params_->config.model.eos_token_id, token)) { + finish_reason = GenerationFinishReason::EosToken; + } else if (turn_limit_reached) { + finish_reason = GenerationFinishReason::TurnLimit; + } else if (context_limit_reached || search_done) { + finish_reason = GenerationFinishReason::ContextLimit; + } + if (finish_reason != GenerationFinishReason::None) { + break; + } + } + + const bool done = finish_reason != GenerationFinishReason::None; + if (done) { + status_ = RequestStatus::TurnComplete; + finish_reason_ = finish_reason; + matched_stop_string_index_ = matched_stop_string_index; + ReleaseTurnResources(); + } + RequestStepResult result{ + last_token, + appended_count != 0, + done, + finish_reason, + matched_stop_string_index, + }; + StageVisibleTokens(result, appended_count, std::nullopt); + result.candidates_consumed = consumed_count; + result.tokens_appended = appended_count; + return result; +} + RequestRewindState Request::PrepareRewindToStartOfTurn( uint64_t turn_id) const { const auto boundary = std::find_if( diff --git a/src/engine/request.h b/src/engine/request.h index ad69a4ed66..c4dfc7e19b 100644 --- a/src/engine/request.h +++ b/src/engine/request.h @@ -90,6 +90,9 @@ struct RequestStepResult { int32_t matched_stop_string_index{-1}; std::array visible_tokens{}; size_t visible_token_count{}; + // Internal commit feedback, independent of stop buffering and caller event-drain rate. + size_t candidates_consumed{}; + size_t tokens_appended{}; }; struct RequestTurnAdmission { @@ -420,6 +423,19 @@ struct Request : std::enable_shared_from_this, */ RequestStepResult CompleteGeneration(); + /** + * @brief Commits tokens an executor selected itself, in sequence order, and reports the step. + * @param tokens One to kMaxGeneratedTokensPerStep tokens the executor decided to commit. + * + * Each token goes through Search::CommitToken, so EOS and the session limit follow the ordinary + * bookkeeping (an EOS token ends the turn and is not appended). Stop strings observe every + * appended token. The turn limit, the context limit and a stop match end the turn at the first + * token that reaches them, with the same precedence as StageGeneration (stop string, EOS, turn + * limit, context limit); tokens after that one are not committed. The returned step lists the + * appended tokens only. + */ + RequestStepResult CommitExternalTokens(std::span tokens); + /** * @brief Checks if the current generation turn reached a stopping condition. * @return True in TurnComplete; the request may still be continued or closed. diff --git a/src/engine/scheduled_requests.cpp b/src/engine/scheduled_requests.cpp index 91fc271b5b..493fd46174 100644 --- a/src/engine/scheduled_requests.cpp +++ b/src/engine/scheduled_requests.cpp @@ -16,6 +16,10 @@ namespace Generators { +ScheduledRequests::~ScheduledRequests() = default; +ScheduledRequests::ScheduledRequests(ScheduledRequests&&) noexcept = default; +ScheduledRequests& ScheduledRequests::operator=(ScheduledRequests&&) noexcept = default; + TargetTokenSelection BuildTopKTargetSelection( std::span tokens, std::span scores, const EffectiveTurnPolicy& policy) { @@ -214,7 +218,107 @@ Tensor* ScheduledRequests::HiddenStates() const { return decoder_state_ ? decoder_state_->HiddenStates() : nullptr; } +void ScheduledRequests::BeginBackendStep(std::vector plan) { + if (output_kind_ != StepOutputKind::DecoderOutputs || decoder_state_) + throw std::logic_error("An executor step can select its output protocol only once."); + if (plan.empty() && std::any_of(requests_.begin(), requests_.end(), + [](const auto& request) { return IsExecuting(request->status_); })) + throw std::logic_error("Executable backend work requires a nonempty plan; defer idle work in the scheduler."); + for (size_t i = 0; i < plan.size(); ++i) { + const auto& row = plan[i]; + if (row.request_row >= requests_.size() || !IsExecuting(requests_[row.request_row]->status_) || + row.key.turn_id != requests_[row.request_row]->CurrentTurnId() || + row.key.context_id == 0 || row.key.lease_epoch == 0 || row.key.slot.epoch == 0 || + row.key.submission_id == 0 || row.key.slot.index >= backend_max_slots_) + throw std::logic_error("A backend plan contains an invalid live row identity."); + for (size_t j = 0; j < i; ++j) { + if (plan[j].request_row == row.request_row || plan[j].key.slot.index == row.key.slot.index || + plan[j].key.context_id != row.key.context_id || plan[j].key.lease_epoch != row.key.lease_epoch) + throw std::logic_error("A backend plan schedules the same request or slot twice."); + } + } + output_kind_ = StepOutputKind::BackendRows; + backend_plan_ = std::move(plan); + backend_rows_.reserve(backend_plan_.size()); + backend_request_rows_.reserve(backend_plan_.size()); + commit_acks_.reserve(backend_plan_.size()); +} + +void ScheduledRequests::StageBackendRow(size_t request_row, BackendRowResult result) { + if (output_kind_ != StepOutputKind::BackendRows) + throw std::logic_error("Backend rows require an explicit backend step plan."); + for (size_t row : backend_request_rows_) { + if (row == request_row) + throw std::logic_error("A backend returned duplicate completion rows."); + } + backend_request_rows_.push_back(request_row); + backend_rows_.push_back(std::move(result)); +} + +void ScheduledRequests::ValidateBackendRows() const { + if (backend_rows_.size() != backend_plan_.size()) + throw std::logic_error("Backend completion coverage does not match its submitted plan."); + for (size_t i = 0; i < backend_rows_.size(); ++i) { + const auto& row = backend_rows_[i]; + const size_t request_row = backend_request_rows_[i]; + const auto planned = std::find_if(backend_plan_.begin(), backend_plan_.end(), + [request_row](const BackendPlannedRow& p) { return p.request_row == request_row; }); + if (planned == backend_plan_.end() || !(planned->key == row.key) || + !IsExecuting(requests_[request_row]->status_) || + requests_[request_row]->CurrentTurnId() != row.key.turn_id) + throw std::logic_error("Backend completion does not belong to the planned live request."); + if (!row.prefill && !row.selected_tokens) + throw std::logic_error("A backend row must report progress or selected tokens."); + if (row.prefill) { + const auto& progress = *row.prefill; + if (progress.prompt_tokens_processed_this_step > planned->max_prompt_tokens || + (progress.prompt_tokens_processed_this_step == 0 && + !(planned->allow_phase_completion && progress.prompt_complete))) + throw std::logic_error("Backend prefill progress is outside its planned prompt range."); + } + if (row.selected_tokens) { + const auto& block = *row.selected_tokens; + if (block.candidate_count == 0 || block.candidate_count > kMaxGeneratedTokensPerStep || + block.candidate_count > backend_max_result_tokens_ || + block.speculation.accepted > block.speculation.proposed) + throw std::logic_error("A backend returned an invalid selected-token block."); + for (size_t j = 0; j < block.candidate_count; ++j) { + if (block.candidates[j] < 0 || block.candidates[j] >= model_->config_->model.vocab_size) + throw std::logic_error("A backend returned a token outside the configured vocabulary."); + } + } + } +} + void ScheduledRequests::GenerateNextTokens(std::vector& results) { + if (output_kind_ == StepOutputKind::BackendRows) { + ValidateBackendRows(); // The entire batch is checked before the first Request changes. + if (backend_committed_) + throw std::logic_error("A backend completion batch was committed more than once."); + backend_committed_ = true; // A partially failed Request commit must never be retried. + results.assign(requests_.size(), RequestStepResult{}); + for (size_t i = 0; i < backend_rows_.size(); ++i) { + const auto& row = backend_rows_[i]; + const size_t request_row = backend_request_rows_[i]; + auto& request = requests_[request_row]; + auto& result = results[request_row]; + if (row.selected_tokens) { + const auto& block = *row.selected_tokens; + result = request->CommitExternalTokens({block.candidates.data(), block.candidate_count}); + } + TokenCommitAck ack; + ack.key = row.key; + ack.candidates_consumed = result.candidates_consumed; + ack.tokens_appended = result.tokens_appended; + ack.tokens_staged_for_events = result.visible_token_count; + ack.request_sequence_length_after = static_cast(request->CurrentSequenceLength()); + ack.turn_finished = result.done; + ack.finish_reason = result.finish_reason; + ack.disposition = result.done ? ContinuationDisposition::Invalidate : ContinuationDisposition::Continue; + commit_acks_.push_back(std::move(ack)); + } + return; + } if (!decoder_state_) { throw std::runtime_error("Cannot generate next tokens without the decoder state."); } diff --git a/src/engine/scheduled_requests.h b/src/engine/scheduled_requests.h index d0073863c0..0836f1b886 100644 --- a/src/engine/scheduled_requests.h +++ b/src/engine/scheduled_requests.h @@ -6,6 +6,7 @@ #include #include "execution_context.h" +#include "executor_result.h" #include "request.h" namespace Generators { @@ -62,6 +63,11 @@ struct BatchedSamplingPlan { }; struct ScheduledRequests { + ~ScheduledRequests(); + ScheduledRequests(ScheduledRequests&&) noexcept; + ScheduledRequests& operator=(ScheduledRequests&&) noexcept; + ScheduledRequests(const ScheduledRequests&) = delete; + ScheduledRequests& operator=(const ScheduledRequests&) = delete; ScheduledRequests(std::vector> requests, std::shared_ptr model, BatchedSampler* batched_sampler, @@ -105,6 +111,18 @@ struct ScheduledRequests { std::vector> ProcessLogits(); void GenerateNextTokens(std::vector& results); + // For a ModelExecutor that selects its own tokens (ModelExecutor::SelectsOwnTokens): BeginBackendStep declares the + // rows of the step and StageBackendRow hands over each row's result. GenerateNextTokens validates the whole batch + // and commits the selected tokens through Request::CommitExternalTokens instead of sampling logits. + void BeginBackendStep(std::vector plan); + void StageBackendRow(size_t request_row, BackendRowResult result); + void SetBackendLimits(size_t max_slots, size_t max_result_tokens) { + backend_max_slots_ = max_slots; + backend_max_result_tokens_ = max_result_tokens; + } + StepOutputKind OutputKind() const { return output_kind_; } + const std::vector& BackendRows() const { return backend_rows_; } + const std::vector& CommitAcks() const { return commit_acks_; } void ScheduleGuidanceMasks() noexcept; void BeginTransaction(); void GenerateNextTokensForTransaction( @@ -117,6 +135,7 @@ struct ScheduledRequests { friend struct test::EngineRunTestAccess; bool PrepareBatchedSamplingPlan(bool require_transaction_support); + void ValidateBackendRows() const; bool TryGenerateNextTokensBatched(std::vector>& logits, bool guidance_applied, std::vector* results = nullptr); @@ -141,6 +160,14 @@ struct ScheduledRequests { // Drafts the transaction stages onto each request's sequence, in scheduled row order. Empty // outside a dynamic step plan. std::vector draft_token_counts_; + StepOutputKind output_kind_{StepOutputKind::DecoderOutputs}; + std::vector backend_plan_; + std::vector backend_rows_; + std::vector backend_request_rows_; + std::vector commit_acks_; + size_t backend_max_slots_{std::numeric_limits::max()}; + size_t backend_max_result_tokens_{kMaxGeneratedTokensPerStep}; + bool backend_committed_{}; std::shared_ptr model_; std::unique_ptr decoder_state_; std::unique_ptr execution_context_; diff --git a/src/engine/scheduler.cpp b/src/engine/scheduler.cpp index b03f190732..a2bb757d66 100644 --- a/src/engine/scheduler.cpp +++ b/src/engine/scheduler.cpp @@ -4,6 +4,7 @@ #include "engine.h" #include "admission.h" #include "decode_first_scheduler_policy.h" +#include "fixed_slot_scheduler.h" #include "request_index.h" #include "sequence_positions.h" @@ -55,6 +56,9 @@ void StaticBatchScheduler::AddRequest(std::shared_ptr request) { "search.chunk_size requires dynamic batching; the static batch scheduler cannot chunk a " "prefill."); } + if (std::find(requests_pool_.begin(), requests_pool_.end(), request) != requests_pool_.end()) { + return; + } requests_pool_.reserve(requests_pool_.size() + 1); auto sampling_state = CreateSamplingState(*request); request->CommitSamplingState(std::move(sampling_state)); @@ -105,6 +109,7 @@ ScheduledRequests StaticBatchScheduler::Schedule() { // Before allocating, we need to ensure that the existing requests in the cache manager // are terminal and no longer need to remain in the scheduler pool. for (auto& request : allocated_requests) { + if (StaysEnrolledWhenDisplaced(*request)) continue; requests_pool_.erase(std::remove(requests_pool_.begin(), requests_pool_.end(), request), requests_pool_.end()); } @@ -140,6 +145,11 @@ DynamicBatchScheduler::DynamicBatchScheduler(std::shared_ptr model, std:: : Scheduler{model}, model_{model}, cache_manager_{cache_manager} {} void DynamicBatchScheduler::AddRequest(std::shared_ptr request) { + // A rewound request, or one restarted after a cancellation, is still in the pool (it leaves only through RemoveRequest); + // enrolling it again would plan it twice in one step. + if (std::find(requests_pool_.begin(), requests_pool_.end(), request) != requests_pool_.end()) { + return; + } requests_pool_.reserve(requests_pool_.size() + 1); auto sampling_state = CreateSamplingState(*request); request->CommitSamplingState(std::move(sampling_state)); @@ -461,6 +471,9 @@ std::unique_ptr Scheduler::Create(std::shared_ptr model, std:: if (cache_manager->SupportsDynamicBatching()) { return std::make_unique(model, cache_manager); } + if (cache_manager->SupportsContinuousMultiSlotDecoding()) { + return std::make_unique(model, cache_manager); + } return std::make_unique(model, cache_manager); } diff --git a/src/engine/scheduler.h b/src/engine/scheduler.h index 61c96e50e6..9602811ec4 100644 --- a/src/engine/scheduler.h +++ b/src/engine/scheduler.h @@ -18,6 +18,19 @@ namespace Generators { inline constexpr size_t kDefaultStaticBatchSize = 4; +// Engine::BeginTurn restarts a turn that was canceled before the request made any progress, even after another request +// took the request's device state, and it enrolls only the first turn of a request. A scheduler that displaces such a +// request from its slot must therefore keep it in its pool. +inline bool StaysEnrolledWhenDisplaced(const Request& request) { + return request.IsTurnComplete() && + (request.IsRestartableCanceledTurn() || request.NeedsReplayAfterRewind()); +} + +struct NontransactionalSchedulingResult { + StepOutcome outcome; + std::optional scheduled; +}; + struct Scheduler { /** * @brief Constructs a Scheduler instance with the specified model and cache manager. @@ -68,6 +81,15 @@ struct Scheduler { */ virtual ScheduledRequests Schedule() = 0; + // Expected capacity and phase deferral use a value, never the fatal exception path. + // Existing schedulers retain their Schedule implementation through this adapter. + virtual NontransactionalSchedulingResult ScheduleStep() { + auto scheduled = Schedule(); + if (!scheduled) + return {{StepOutcomeKind::NoWork, 0, nullptr}, std::nullopt}; + return {{StepOutcomeKind::Committed, 0, nullptr}, std::move(scheduled)}; + } + virtual StepPlanningResult PlanStep(StepPlan&) { throw std::logic_error("Scheduler does not support transactional step planning."); } @@ -78,6 +100,8 @@ struct Scheduler { return !batched_sampler_ || batched_sampler_->SupportsTransactions(); } + virtual bool SupportsIndependentSlotAdmission() const { return false; } + /** * @brief Checks if the Scheduler has any pending requests. * @return True if there are pending requests, false otherwise. @@ -137,6 +161,8 @@ struct DynamicBatchScheduler : Scheduler { bool HasPendingRequests() const override; + bool SupportsIndependentSlotAdmission() const override { return true; } + private: std::shared_ptr model_; std::shared_ptr cache_manager_; diff --git a/src/engine/token_producing_session.h b/src/engine/token_producing_session.h new file mode 100644 index 0000000000..5e1a98f41a --- /dev/null +++ b/src/engine/token_producing_session.h @@ -0,0 +1,111 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#pragma once + +#include +#include +#include +#include + +#include "executor_result.h" + +namespace Generators { + +// Owner-thread host bookkeeping for token-producing backends. It contains no device state or graph layout. +// An identity namespace is never reused across Engine leases, including when local submission counters restart. +class TokenProducingSession { + public: + explicit TokenProducingSession(size_t slots) : context_id_{NextIdentity()}, slots_(slots) { + retired_.reserve(slots); + } + + BackendPlannedRow Plan(size_t request_row, size_t slot, const void* owner, uint64_t turn_id, + size_t max_prompt_tokens, bool allow_phase_completion) { + auto& state = slots_.at(slot); + for (const auto& retired : retired_) { + if (retired.slot.index == slot) throw std::logic_error("A retiring backend slot is not physically reusable."); + } + if (state.pending) throw std::logic_error("A backend completion must be settled before another submission."); + if (state.owner != owner || state.turn_id != turn_id) { + Retire(slot); + state.owner = owner; + state.turn_id = turn_id; + } + if (!owner) throw std::logic_error("A backend submission requires a live owner."); + if (state.submission == std::numeric_limits::max()) throw std::overflow_error("Backend submission identity exhausted."); + state.pending = true; + state.key = CompletionKey{context_id_, 1, SlotHandle{static_cast(slot), state.epoch}, turn_id, ++state.submission}; + return BackendPlannedRow{request_row, state.key, max_prompt_tokens, allow_phase_completion}; + } + + // False means a tracked retiring submission settled without public delivery. + bool Collected(const BackendRowResult& row) { + for (auto it = retired_.begin(); it != retired_.end(); ++it) { + if (*it == row.key) { + retired_.erase(it); + return false; + } + } + auto& state = RequirePending(row.key); + if (state.collected) throw std::logic_error("Duplicate backend completion."); + state.collected = true; + state.candidates = row.selected_tokens ? row.selected_tokens->candidate_count : 0; + return true; + } + + void Acknowledge(const TokenCommitAck& ack) { + auto& state = RequirePending(ack.key); + if (!state.collected || ack.candidates_consumed > state.candidates || + ack.tokens_appended > ack.candidates_consumed || ack.tokens_staged_for_events > ack.tokens_appended || + (ack.candidates_consumed != state.candidates && ack.disposition == ContinuationDisposition::Continue) || + (ack.turn_finished && ack.disposition == ContinuationDisposition::Continue)) { + throw std::logic_error("Invalid backend commit acknowledgement."); + } + state.pending = false; + state.collected = false; + state.candidates = 0; + } + + void Retire(size_t slot) { + auto& state = slots_.at(slot); + if (state.epoch == std::numeric_limits::max()) throw std::overflow_error("Backend slot epoch exhausted."); + if (state.pending && !state.collected) { + if (retired_.size() >= slots_.size()) throw std::logic_error("Backend retirement capacity exhausted."); + retired_.push_back(state.key); + } + const uint64_t epoch = state.epoch + 1; + state = {}; + state.epoch = epoch; + } + + private: + struct Slot { + const void* owner{}; + uint64_t turn_id{}; + uint64_t epoch{}; + uint64_t submission{}; + CompletionKey key{}; + bool pending{}; + bool collected{}; + size_t candidates{}; + }; + + static uint64_t NextIdentity() { + static std::atomic next{1}; + const uint64_t id = next.fetch_add(1, std::memory_order_relaxed); + if (id == 0) throw std::overflow_error("Backend context identity exhausted."); + return id; + } + + Slot& RequirePending(const CompletionKey& key) { + auto& state = slots_.at(key.slot.index); + if (!state.pending || !(state.key == key)) throw std::logic_error("Unknown or stale backend completion identity."); + return state; + } + + uint64_t context_id_; + std::vector slots_; + std::vector retired_; +}; + +} // namespace Generators diff --git a/src/models/qwen38_graph.cpp b/src/models/qwen38_graph.cpp index 27bd1dfd17..af922cc7f8 100644 --- a/src/models/qwen38_graph.cpp +++ b/src/models/qwen38_graph.cpp @@ -1548,7 +1548,11 @@ Qwen38GraphModel::Qwen38GraphModel(std::unique_ptr config, OrtEnv& ort_e capabilities.supports_stop_strings = true; capabilities.max_slots = 1; capabilities.configured_max_batch_size = 1; - capabilities.max_request_length = static_cast(engine_->Capacity() - engine_->ReservedTail()); + const uint64_t kv_limit = static_cast(engine_->Capacity() - engine_->ReservedTail()); + capabilities.max_request_length = + config_->model.context_length > 0 + ? std::min(kv_limit, static_cast(config_->model.context_length)) + : kv_limit; capabilities.max_result_tokens = kMaxGeneratedTokensPerStep; if (config_->engine.dynamic_batching) capabilities.max_scheduled_tokens = config_->engine.dynamic_batching->max_scheduled_tokens; diff --git a/test/cpp/engine/backend_engine_run_tests.cpp b/test/cpp/engine/backend_engine_run_tests.cpp new file mode 100644 index 0000000000..a1eb2e655e --- /dev/null +++ b/test/cpp/engine/backend_engine_run_tests.cpp @@ -0,0 +1,246 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include +#include + +#include + +#include "engine_test_doubles.h" +#include "engine_test_helpers.h" + +namespace Generators::test { +namespace { + +class ScriptedStaticScheduler final : public Scheduler { + public: + ScriptedStaticScheduler(std::shared_ptr model, std::shared_ptr cache, + std::vector outcomes) + : Scheduler{model}, inner_{model, std::move(cache)}, outcomes_{std::move(outcomes)} {} + + void AddRequest(std::shared_ptr request) override { + requests_.push_back(request); + inner_.AddRequest(std::move(request)); + } + void RemoveRequest(std::shared_ptr request) override { + inner_.RemoveRequest(request); + requests_.erase(std::remove(requests_.begin(), requests_.end(), request), requests_.end()); + } + ScheduledRequests Schedule() override { + ++execute_plans; + return inner_.Schedule(); + } + NontransactionalSchedulingResult ScheduleStep() override { + if (next_outcome_ == outcomes_.size()) return Scheduler::ScheduleStep(); + const auto kind = outcomes_[next_outcome_++]; + const void* request_id = kind == StepOutcomeKind::UnserviceableRequest && !requests_.empty() + ? requests_.front().get() + : nullptr; + return {{kind, 0, request_id}, std::nullopt}; + } + bool HasPendingRequests() const override { return inner_.HasPendingRequests(); } + + size_t execute_plans{}; + + private: + StaticBatchScheduler inner_; + std::vector outcomes_; + std::vector> requests_; + size_t next_outcome_{}; +}; + +struct ObservedExecutor final : RecordingModelExecutor { + ObservedExecutor(std::shared_ptr model, std::shared_ptr cache, + bool selects_tokens, bool produces_tokens) + : RecordingModelExecutor(std::move(model), std::move(cache), 5), + selects_tokens_{selects_tokens}, + produces_tokens_{produces_tokens} {} + + bool SelectsOwnTokens() const override { return selects_tokens_; } + void Decode(ScheduledRequests& scheduled, ExecutionContext& context) override { + if (!produces_tokens_) { + RecordingModelExecutor::Decode(scheduled, context); + return; + } + ++decode_calls; + std::vector plan; + for (size_t row = 0; row < scheduled.size(); ++row) { + CompletionKey key{42, 1, {static_cast(row), 1}, scheduled[row]->CurrentTurnId(), ++submission_}; + plan.push_back({row, key, 0, false}); + } + scheduled.BeginBackendStep(plan); + for (size_t row = 0; row < scheduled.size(); ++row) { + BackendRowResult result; + result.key = plan[row].key; + SelectedTokenBlock tokens; + for (size_t i = 0; i < candidate_width; ++i) { + tokens.candidates[i] = static_cast(5 + i); + } + tokens.candidate_count = candidate_width; + result.selected_tokens = tokens; + scheduled.StageBackendRow(row, std::move(result)); + } + } + void OnStepCommitted(const ScheduledRequests& scheduled, + std::span) override { + ++commit_calls; + acknowledged_rows = scheduled.CommitAcks().size(); + if (!scheduled.CommitAcks().empty()) { + last_staged_for_events = scheduled.CommitAcks().front().tokens_staged_for_events; + } + if (fail_acknowledgement) throw std::runtime_error("Injected backend acknowledgement failure."); + } + void OnExecutionFailure() noexcept override { ++failure_notifications; } + + bool fail_acknowledgement{}; + size_t candidate_width{1}; + size_t commit_calls{}; + size_t acknowledged_rows{}; + size_t last_staged_for_events{}; + size_t failure_notifications{}; + + private: + bool selects_tokens_; + bool produces_tokens_; + uint64_t submission_{}; +}; + +struct StaticEngineFixture { + std::shared_ptr engine; + ScriptedStaticScheduler* scheduler; + ObservedExecutor* executor; +}; + +StaticEngineFixture MakeStaticEngine(std::vector outcomes, + bool selects_tokens = false, bool produces_tokens = false) { + auto model = LoadDummyDecoderModel(); + auto cache = std::make_shared(model, 2, nullptr, false); + auto scheduler = std::make_unique(model, cache, std::move(outcomes)); + auto executor = std::make_unique(model, cache, selects_tokens, produces_tokens); + auto* scheduler_observer = scheduler.get(); + auto* executor_observer = executor.get(); + auto engine = std::make_shared( + model, EngineDependencies{cache, std::move(scheduler), std::move(executor)}); + return {std::move(engine), scheduler_observer, executor_observer}; +} + +TEST(BackendEngineRunTest, NoWorkAndCapacityDeferralKeepEngineHealthyWithoutExecuting) { + auto fixture = MakeStaticEngine({StepOutcomeKind::NoWork, StepOutcomeKind::CapacityDeferred}); + auto request = CreateEngineRequest(fixture.engine); + request->BeginTurn(std::array{3}, std::optional{1}); + const auto prompt_length = request->CurrentSequenceLength(); + std::array events; + for (size_t step = 0; step < 2; ++step) { + const size_t count = fixture.engine->Run(events); + for (size_t event = 0; event < count; ++event) { + EXPECT_EQ(events[event].flags & (EngineEventFlagToken | EngineEventFlagFailed), 0u); + } + EXPECT_EQ(request->CurrentSequenceLength(), prompt_length); + EXPECT_EQ(fixture.executor->decode_calls, 0); + EXPECT_EQ(fixture.executor->commit_calls, 0u); + EXPECT_EQ(fixture.executor->failure_notifications, 0u); + EXPECT_TRUE(fixture.engine->HasPendingRequests()); + } + const size_t count = fixture.engine->Run(events); + ASSERT_GT(count, 0u); + EXPECT_EQ(events[0].request, request); + EXPECT_NE(events[0].flags & EngineEventFlagToken, 0u); + EXPECT_EQ(fixture.executor->decode_calls, 1); + EXPECT_EQ(fixture.executor->failure_notifications, 0u); +} + +TEST(BackendEngineRunTest, UnserviceableRequestFailsWithoutPoisoningEngineOrExecutingIt) { + auto fixture = MakeStaticEngine({StepOutcomeKind::UnserviceableRequest}); + auto oversized = CreateEngineRequest(fixture.engine); + oversized->BeginTurn(std::array{3}, std::optional{1}); + std::array events; + ASSERT_EQ(fixture.engine->Run(events), 1u); + EXPECT_EQ(events[0].request, oversized); + EXPECT_EQ(events[0].error_code, EngineErrorCode::RequestUnserviceable); + EXPECT_EQ(events[0].flags & EngineEventFlagToken, 0u); + EXPECT_EQ(fixture.executor->decode_calls, 0); + EXPECT_EQ(fixture.executor->failure_notifications, 0u); + + auto successor = CreateEngineRequest(fixture.engine); + successor->BeginTurn(std::array{4}, std::optional{1}); + ASSERT_GT(fixture.engine->Run(events), 0u); + EXPECT_EQ(events[0].request, successor); + EXPECT_NE(events[0].flags & EngineEventFlagToken, 0u); + EXPECT_EQ(fixture.executor->decode_calls, 1); +} + +TEST(BackendEngineRunTest, OutputProtocolMismatchFailsBeforeRequestCommit) { + for (bool selects_tokens : {false, true}) { + auto fixture = MakeStaticEngine({}, selects_tokens, !selects_tokens); + auto request = CreateEngineRequest(fixture.engine); + request->BeginTurn(std::array{3}, std::optional{2}); + const auto prompt_length = request->CurrentSequenceLength(); + std::array events; + const size_t count = fixture.engine->Run(events); + ASSERT_GT(count, 0u); + for (size_t event = 0; event < count; ++event) { + EXPECT_EQ(events[event].flags & EngineEventFlagToken, 0u); + EXPECT_NE(events[event].flags & EngineEventFlagFailed, 0u); + } + EXPECT_EQ(request->CurrentSequenceLength(), prompt_length); + EXPECT_EQ(fixture.executor->commit_calls, 0u); + EXPECT_EQ(fixture.executor->failure_notifications, 1u); + EXPECT_THROW(fixture.engine->Run(events), EngineStepError); + EXPECT_EQ(fixture.executor->failure_notifications, 1u); + } +} + +TEST(BackendEngineRunTest, AcknowledgementFailurePublishesNoSuccessAndPoisonsOnce) { + auto fixture = MakeStaticEngine({}, true, true); + fixture.executor->fail_acknowledgement = true; + auto request = CreateEngineRequest(fixture.engine); + request->BeginTurn(std::array{3}, std::optional{2}); + const auto prompt_length = request->CurrentSequenceLength(); + std::array events; + const size_t count = fixture.engine->Run(events); + ASSERT_GT(count, 0u); + for (size_t event = 0; event < count; ++event) { + EXPECT_EQ(events[event].flags & EngineEventFlagToken, 0u); + EXPECT_NE(events[event].flags & EngineEventFlagFailed, 0u); + } + EXPECT_EQ(request->CurrentSequenceLength(), prompt_length + 1); + EXPECT_EQ(fixture.executor->commit_calls, 1u); + EXPECT_EQ(fixture.executor->acknowledged_rows, 1u); + EXPECT_EQ(fixture.executor->failure_notifications, 1u); + EXPECT_THROW(fixture.engine->Run(events), EngineStepError); + EXPECT_THROW(fixture.engine->CreateRequest(), EngineStepError); + EXPECT_EQ(fixture.executor->commit_calls, 1u); + EXPECT_EQ(fixture.executor->failure_notifications, 1u); +} + +TEST(BackendEngineRunTest, StagedOutputAckDoesNotChangeWhenCallerDrainsEventsAcrossMultipleRuns) { + auto fixture = MakeStaticEngine({}, true, true); + fixture.executor->candidate_width = 3; + auto request = CreateEngineRequest(fixture.engine); + request->BeginTurn(std::array{3}, std::optional{3}); + + std::array single_event; + ASSERT_EQ(fixture.engine->Run(single_event), 1u); + EXPECT_EQ(fixture.executor->decode_calls, 1); + EXPECT_EQ(fixture.executor->commit_calls, 1u); + EXPECT_EQ(fixture.executor->last_staged_for_events, 3u); + EXPECT_EQ(single_event[0].token, 5); + EXPECT_EQ(single_event[0].flags & EngineEventFlagTurnFinished, 0u); + + ASSERT_EQ(fixture.engine->Run(single_event), 1u); + EXPECT_EQ(fixture.executor->decode_calls, 1); + EXPECT_EQ(fixture.executor->commit_calls, 1u); + EXPECT_EQ(single_event[0].token, 6); + EXPECT_EQ(single_event[0].flags & EngineEventFlagTurnFinished, 0u); + + ASSERT_EQ(fixture.engine->Run(single_event), 1u); + EXPECT_EQ(fixture.executor->decode_calls, 1); + EXPECT_EQ(fixture.executor->commit_calls, 1u); + EXPECT_EQ(single_event[0].token, 7); + EXPECT_NE(single_event[0].flags & EngineEventFlagTurnFinished, 0u); +} + +} // namespace +} // namespace Generators::test diff --git a/test/cpp/engine/backend_factory_tests.cpp b/test/cpp/engine/backend_factory_tests.cpp new file mode 100644 index 0000000000..4c110467a7 --- /dev/null +++ b/test/cpp/engine/backend_factory_tests.cpp @@ -0,0 +1,169 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include + +#include "engine/backend_factory.h" +#include "engine_test_doubles.h" +#include "engine_test_helpers.h" + +namespace Generators::test { +namespace { + +struct IndependentSlotCache final : RecordingCacheManager { + explicit IndependentSlotCache(std::shared_ptr model) + : RecordingCacheManager(std::move(model), 2, nullptr, false) {} + bool SupportsContinuousMultiSlotDecoding() const override { return true; } + size_t MaxStaticBatchSize() const override { return 2; } + size_t MaxStaticRequestLength() const override { return 64; } +}; + +struct SelectedTokenExecutor final : ModelExecutor { + bool SelectsOwnTokens() const override { return true; } + void Decode(ScheduledRequests&, ExecutionContext&) override { + throw std::logic_error("This capability test must not execute device work."); + } +}; + +std::shared_ptr LoadStaticDecoderOnlyModel() { + auto config = CreateConfig(GetOrtEnv(), MODEL_PATH "engine/dummy-decoder"); + config->model.type = "decoder"; + return CreateModel(GetOrtEnv(), std::move(config)); +} + +TEST(EngineBackendFactoryTest, DefaultDecoderBundleKeepsItsExecutionProtocol) { + auto model = LoadStaticDecoderOnlyModel(); + auto bundle = CreateEngineBackend(model); + EXPECT_FALSE(bundle.capabilities.UsesTransactions()); + EXPECT_FALSE(bundle.capabilities.SelectsOwnTokens()); + EXPECT_EQ(bundle.capabilities.kv_storage, BackendKvStorage::RebuiltContiguous); + ASSERT_TRUE(bundle.dependencies.backend_capabilities); + EXPECT_EQ(bundle.dependencies.backend_capabilities->max_slots, + bundle.dependencies.cache_manager->MaxBatchSize()); + auto engine = std::make_shared(model, std::move(bundle.dependencies)); + EXPECT_EQ(engine->GetCapabilities().configured_max_batch_size, + bundle.capabilities.configured_max_batch_size); +} + +TEST(EngineBackendFactoryTest, InjectedIndependentSlotsRemainNontransactional) { + auto model = LoadDummyDecoderModel(); + auto cache = std::make_shared(model); + auto scheduler = Scheduler::Create(model, cache); + auto executor = std::make_unique(); + const auto capabilities = ResolveBackendCapabilities(*model, *cache, *scheduler, *executor); + EXPECT_TRUE(capabilities.SupportsIndependentAdmission()); + EXPECT_TRUE(capabilities.SelectsOwnTokens()); + EXPECT_FALSE(capabilities.UsesTransactions()); + EXPECT_TRUE(capabilities.supports_stop_strings); + EXPECT_FALSE(capabilities.supports_sampling); + EXPECT_EQ(capabilities.max_request_length, 64u); + auto engine = std::make_shared( + model, EngineDependencies{cache, std::move(scheduler), std::move(executor)}); + EXPECT_EQ(engine->GetCapabilities().configured_max_batch_size, 2u); + EXPECT_EQ(engine->GetCapabilities().max_request_length, 64u); + RequestOptions oversized; + oversized.max_session_tokens = 65; + EXPECT_THROW(engine->CreateRequest(oversized), std::runtime_error); +} + +TEST(EngineBackendFactoryTest, ExplicitContradictionsFailBeforeAdmission) { + auto model = LoadDummyDecoderModel(); + auto cache = std::make_shared(model); + auto scheduler = Scheduler::Create(model, cache); + auto executor = std::make_unique(); + auto capabilities = ResolveBackendCapabilities(*model, *cache, *scheduler, *executor); + capabilities.execution_protocol = BackendExecutionProtocol::Transactional; + EXPECT_THROW(ResolveBackendCapabilities(*model, *cache, *scheduler, *executor, capabilities), + std::runtime_error); + capabilities = ResolveBackendCapabilities(*model, *cache, *scheduler, *executor); + capabilities.max_request_length = 65; + EXPECT_THROW(ResolveBackendCapabilities(*model, *cache, *scheduler, *executor, capabilities), + std::runtime_error); + // Zero is no limit: a cache that holds 64 positions cannot declare it. + capabilities.max_request_length = 0; + EXPECT_THROW(ResolveBackendCapabilities(*model, *cache, *scheduler, *executor, capabilities), + std::runtime_error); + capabilities.max_request_length = 64; + EXPECT_NO_THROW(ResolveBackendCapabilities(*model, *cache, *scheduler, *executor, capabilities)); + capabilities = ResolveBackendCapabilities(*model, *cache, *scheduler, *executor); + capabilities.token_selection = BackendTokenSelection::Engine; + EngineDependencies dependencies{cache, std::move(scheduler), std::move(executor)}; + dependencies.backend_capabilities = capabilities; + EXPECT_THROW((void)std::make_shared(model, std::move(dependencies)), std::runtime_error); +} + +TEST(EngineBackendFactoryTest, SelectedTokensCannotClaimPagedTransactions) { + auto model = LoadDummyDecoderModel(); + model->config_->engine.dynamic_batching = Config::Engine::DynamicBatching{}; + auto cache = std::make_shared(model, 2); + auto scheduler = Scheduler::Create(model, cache); + SelectedTokenExecutor executor; + EXPECT_THROW(ResolveBackendCapabilities(*model, *cache, *scheduler, executor), std::runtime_error); +} + +TEST(EngineBackendFactoryTest, StaticDecoderCannotAdvertiseUnsupportedStopHandling) { + auto model = LoadDummyDecoderModel(); + auto cache = std::make_shared(model, 2, nullptr, false); + auto scheduler = Scheduler::Create(model, cache); + RecordingModelExecutor executor{model, cache, 5}; + auto capabilities = ResolveBackendCapabilities(*model, *cache, *scheduler, executor); + capabilities.supports_stop_strings = true; + EXPECT_THROW(ResolveBackendCapabilities(*model, *cache, *scheduler, executor, capabilities), + std::runtime_error); +} + +TEST(EngineBackendFactoryTest, ResolvedCapabilitiesDoNotChangeAfterAssembly) { + auto model = LoadStaticDecoderOnlyModel(); + auto bundle = CreateEngineBackend(model); + const auto configured_rows = bundle.capabilities.configured_max_batch_size; + auto engine = std::make_shared(model, std::move(bundle.dependencies)); + model->config_->engine.static_batching = Config::Engine::StaticBatching{}; + model->config_->engine.static_batching->max_batch_size = 17; + EXPECT_EQ(engine->GetCapabilities().configured_max_batch_size, configured_rows); +} + +TEST(EngineBackendFactoryTest, InjectedSelectedTokenPolicyIsValidatedBeforeMutation) { + auto model = LoadDummyDecoderModel(); + auto cache = std::make_shared(model); + auto scheduler = Scheduler::Create(model, cache); + auto engine = std::make_shared( + model, EngineDependencies{cache, std::move(scheduler), std::make_unique()}); + auto request = engine->CreateRequest(); + const std::array prompt{3}; + TurnOptions options; + options.do_sample = true; + options.top_k = 4; + EXPECT_THROW(engine->BeginTurn(request, prompt, options), std::runtime_error); + EXPECT_TRUE(request->IsAwaitingFirstTurn()); + EXPECT_EQ(request->CurrentSequenceLength(), 0); +} + +TEST(EngineBackendFactoryTest, StaticBatchSchedulerCannotAdvertiseIndependentAdmission) { + auto model = LoadDummyDecoderModel(); + auto cache = std::make_shared(model); + StaticBatchScheduler static_scheduler{model, cache}; + SelectedTokenExecutor executor; + const auto resolved = ResolveBackendCapabilities(*model, *cache, static_scheduler, executor); + EXPECT_FALSE(resolved.SupportsIndependentAdmission()); + auto declared = resolved; + declared.admission = BackendAdmission::IndependentSlots; + EXPECT_THROW(ResolveBackendCapabilities(*model, *cache, static_scheduler, executor, declared), + std::runtime_error); +} + +TEST(EngineBackendFactoryTest, StaticMaxRequestLengthIsClampedToModelContextLength) { + auto model = LoadDummyDecoderModel(); + model->config_->model.context_length = 32; + auto cache = std::make_shared(model); // MaxStaticRequestLength() == 64 + auto scheduler = Scheduler::Create(model, cache); + SelectedTokenExecutor executor; + const auto resolved = ResolveBackendCapabilities(*model, *cache, *scheduler, executor); + EXPECT_EQ(resolved.max_request_length, 32u); + auto declared = resolved; + declared.max_request_length = 64; + EXPECT_THROW(ResolveBackendCapabilities(*model, *cache, *scheduler, executor, declared), + std::runtime_error); +} + +} // namespace +} // namespace Generators::test diff --git a/test/cpp/engine/engine_run_tests.cpp b/test/cpp/engine/engine_run_tests.cpp index d159b754a1..c0d8492762 100644 --- a/test/cpp/engine/engine_run_tests.cpp +++ b/test/cpp/engine/engine_run_tests.cpp @@ -4338,6 +4338,28 @@ TEST_F(EngineRunTest, FailedReplayCannotBeginAnotherTurn) { EXPECT_FALSE(engine.engine->HasPendingRequests()); } +// A turn canceled before it made progress restarts through Engine::BeginTurn, which enrolls the request with the scheduler +// again. The dynamic scheduler still holds it, so the restarted turn must run once and not fail on a duplicate candidate. +TEST_F(EngineRunTest, DynamicCanceledUnstartedTurnRestartsAndRunsOnce) { + auto engine = MakeDoublesEngine(model_, /*capacity=*/8, /*forced_token=*/5); + auto request = CreateRequestWithPrompt(engine.engine, Prompt(10)); + ASSERT_TRUE(request->Cancel(request->CurrentTurnId())); + const auto canceled = RunOne(*engine.engine); + ASSERT_EQ(canceled.request, request); + ASSERT_EQ(canceled.finish_reason, GenerationFinishReason::Canceled); + ASSERT_TRUE(request->IsRestartableCanceledTurn()); + + request->BeginTurn(Prompt(20), std::optional{1}); + + const auto event = RunOne(*engine.engine); + ASSERT_EQ(event.request, request); + EXPECT_EQ(event.flags & EngineEventFlagFailed, 0u); + EXPECT_EQ(event.token, 5); + EXPECT_NE(event.flags & EngineEventFlagTurnFinished, 0u); + EXPECT_EQ(engine.cache->AllocatedCount(), 1u); + EXPECT_FALSE(engine.engine->HasPendingRequests()); +} + TEST_F(EngineRunTest, DynamicCompositeRewindReleasesAndRebuildsAllState) { model_ = LoadSyntheticCompositeModel(); auto engine = diff --git a/test/cpp/engine/executor_result_tests.cpp b/test/cpp/engine/executor_result_tests.cpp new file mode 100644 index 0000000000..da6b067627 --- /dev/null +++ b/test/cpp/engine/executor_result_tests.cpp @@ -0,0 +1,183 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include + +#include "engine/scheduled_requests.h" +#include "engine_test_doubles.h" +#include "engine_test_helpers.h" + +namespace Generators::test { +namespace { + +class ExecutorResultTest : public ::testing::Test { + protected: + void SetUp() override { + model_ = LoadDummyDecoderModel(); + engine_ = MakeStaticDoublesEngine(model_, 4, EosToken(*model_)).engine; + } + + std::shared_ptr ActiveRequest() { + auto request = CreateRequestWithPrompt(engine_, std::array{2, 3}); + request->Schedule(); + return request; + } + + BackendPlannedRow Plan(const std::shared_ptr& request, size_t row = 0) { + return {row, {17, 2, {static_cast(row), 1}, request->CurrentTurnId(), 1}, 2, true}; + } + + BackendRowResult Tokens(const BackendPlannedRow& plan, std::initializer_list values) { + BackendRowResult row; + row.key = plan.key; + row.selected_tokens.emplace(); + std::copy(values.begin(), values.end(), row.selected_tokens->candidates.begin()); + row.selected_tokens->candidate_count = values.size(); + return row; + } + + std::shared_ptr model_; + std::shared_ptr engine_; +}; + +TEST_F(ExecutorResultTest, ExplicitPrefillProgressNeverSamplesOrAppends) { + auto request = ActiveRequest(); + auto plan = Plan(request); + ScheduledRequests batch{{request}, model_, nullptr, nullptr}; + batch.BeginBackendStep({plan}); + batch.StageBackendRow(0, {plan.key, PrefillProgress{2, false}, std::nullopt}); + std::vector results; + batch.GenerateNextTokens(results); + EXPECT_EQ(request->CurrentSequenceLength(), 2); + ASSERT_EQ(batch.CommitAcks().size(), 1); + EXPECT_EQ(batch.CommitAcks()[0].candidates_consumed, 0); + EXPECT_EQ(batch.CommitAcks()[0].request_sequence_length_after, 2); + EXPECT_FALSE(batch.CommitAcks()[0].turn_finished); + EXPECT_THROW(batch.GenerateNextTokens(results), std::logic_error); +} + +TEST_F(ExecutorResultTest, ProgressAndTokensCanShareOneRowInAMixedBatch) { + auto first = ActiveRequest(); + auto second = ActiveRequest(); + const auto a = Plan(first); + const auto b = Plan(second, 1); + ScheduledRequests batch{{first, second}, model_, nullptr, nullptr}; + batch.BeginBackendStep({a, b}); + auto completed = Tokens(b, {4}); + completed.prefill = PrefillProgress{2, true}; + batch.StageBackendRow(1, completed); + batch.StageBackendRow(0, {a.key, PrefillProgress{1, false}, std::nullopt}); + // The result owns its small token block through commit and settlement. + completed.selected_tokens->candidates[0] = 7; + std::vector results; + batch.GenerateNextTokens(results); + EXPECT_EQ(first->CurrentSequenceLength(), 2); + EXPECT_EQ(second->TokensCpu().back(), 4); + ASSERT_EQ(batch.CommitAcks().size(), 2); + EXPECT_EQ(batch.CommitAcks()[0].tokens_appended, 1); + EXPECT_EQ(batch.CommitAcks()[0].tokens_staged_for_events, 1); +} + +TEST_F(ExecutorResultTest, MalformedSecondRowCannotMutateFirstRequest) { + auto first = ActiveRequest(); + auto second = ActiveRequest(); + auto a = Plan(first); + auto b = Plan(second, 1); + for (int violation = 0; violation < 7; ++violation) { + ScheduledRequests batch{{first, second}, model_, nullptr, nullptr}; + batch.BeginBackendStep({a, b}); + batch.StageBackendRow(0, Tokens(a, {4})); + auto bad = Tokens(b, {5}); + switch (violation) { + case 0: + ++bad.key.lease_epoch; + break; + case 1: + bad.selected_tokens->candidates[0] = -1; + break; + case 2: + bad.selected_tokens->candidates[0] = model_->config_->model.vocab_size; + break; + case 3: + bad.selected_tokens->candidate_count = kMaxGeneratedTokensPerStep + 1; + break; + case 4: + bad.selected_tokens.reset(); + break; + case 5: + bad.prefill = PrefillProgress{3, true}; + break; + case 6: + bad.prefill = PrefillProgress{0, false}; + break; + } + batch.StageBackendRow(1, std::move(bad)); + std::vector results; + EXPECT_THROW(batch.GenerateNextTokens(results), std::logic_error) << violation; + EXPECT_EQ(first->CurrentSequenceLength(), 2) << violation; + EXPECT_EQ(second->CurrentSequenceLength(), 2) << violation; + EXPECT_TRUE(batch.CommitAcks().empty()); + } +} + +TEST_F(ExecutorResultTest, CoverageAndDuplicateRowsAreValidated) { + auto first = ActiveRequest(); + auto second = ActiveRequest(); + auto a = Plan(first); + auto b = Plan(second, 1); + ScheduledRequests missing{{first, second}, model_, nullptr, nullptr}; + missing.BeginBackendStep({a, b}); + missing.StageBackendRow(0, Tokens(a, {4})); + EXPECT_THROW(missing.StageBackendRow(0, Tokens(a, {4})), std::logic_error); + std::vector results; + EXPECT_THROW(missing.GenerateNextTokens(results), std::logic_error); + EXPECT_EQ(first->CurrentSequenceLength(), 2); + ScheduledRequests unexpected{{first, second}, model_, nullptr, nullptr}; + unexpected.BeginBackendStep({a}); + unexpected.StageBackendRow(1, Tokens(b, {4})); + EXPECT_THROW(unexpected.GenerateNextTokens(results), std::logic_error); + EXPECT_EQ(second->CurrentSequenceLength(), 2); +} + +TEST_F(ExecutorResultTest, NegotiatedResultAndSlotLimitsAreEnforced) { + auto request = ActiveRequest(); + auto plan = Plan(request); + ScheduledRequests batch{{request}, model_, nullptr, nullptr}; + batch.SetBackendLimits(1, 1); + batch.BeginBackendStep({plan}); + batch.StageBackendRow(0, Tokens(plan, {4, 5})); + std::vector results; + EXPECT_THROW(batch.GenerateNextTokens(results), std::logic_error); + ScheduledRequests bad_slot{{request}, model_, nullptr, nullptr}; + bad_slot.SetBackendLimits(1, 1); + plan.key.slot.index = 1; + EXPECT_THROW(bad_slot.BeginBackendStep({plan}), std::logic_error); + EXPECT_EQ(request->CurrentSequenceLength(), 2); +} + +TEST_F(ExecutorResultTest, EosConsumptionDiffersFromAppendedTokensAtEveryPosition) { + for (size_t eos_position = 0; eos_position < kMaxGeneratedTokensPerStep; ++eos_position) { + auto request = ActiveRequest(); + auto plan = Plan(request); + auto row = Tokens(plan, {}); + row.selected_tokens->candidates.fill(4); + row.selected_tokens->candidates[eos_position] = EosToken(*model_); + row.selected_tokens->candidate_count = kMaxGeneratedTokensPerStep; + ScheduledRequests batch{{request}, model_, nullptr, nullptr}; + batch.BeginBackendStep({plan}); + batch.StageBackendRow(0, std::move(row)); + std::vector results; + batch.GenerateNextTokens(results); + const auto& ack = batch.CommitAcks().at(0); + EXPECT_EQ(ack.candidates_consumed, eos_position + 1); + EXPECT_EQ(ack.tokens_appended, eos_position); + EXPECT_EQ(ack.tokens_staged_for_events, eos_position); + EXPECT_EQ(ack.request_sequence_length_after, eos_position + 2); + EXPECT_TRUE(ack.turn_finished); + EXPECT_EQ(ack.disposition, ContinuationDisposition::Invalidate); + EXPECT_EQ(ack.finish_reason, GenerationFinishReason::EosToken); + } +} + +} // namespace +} // namespace Generators::test diff --git a/test/cpp/engine/qwen38_graph_executor_tests.cpp b/test/cpp/engine/qwen38_graph_executor_tests.cpp new file mode 100644 index 0000000000..da01ce61db --- /dev/null +++ b/test/cpp/engine/qwen38_graph_executor_tests.cpp @@ -0,0 +1,1804 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// Engine-level tests of the static single-slot Qwen3.8-Flash CUDA-graph backend. A scripted device replaces the CUDA +// loop (see qwen38_graph_test_doubles.h); everything else is the real Engine: StaticBatchScheduler, the single-slot +// cache manager, the executor, Request::CommitExternalTokens and the Engine event path. The fixture model is the +// synthetic-paged one because it carries a tiny tokenizer for the stop-string cases (token 5 -> "ST", 6 -> "OP", +// 0/1 special, EOS is 1). + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "engine/qwen38_graph_executor.h" +#include "engine_test_doubles.h" +#include "engine_test_helpers.h" +#include "models/qwen38_ort_env.h" +#include "qwen38_graph_test_doubles.h" + +namespace Generators::test { +namespace { + +constexpr int32_t kEos = 1; + +// A reusable barrier for a fixed number of threads. It gives up (instead of hanging the test) when a peer is gone. +class Rendezvous { + public: + explicit Rendezvous(size_t parties) : parties_{parties} {} + + // Blocks until every party has arrived; false when the rendezvous was abandoned or the peers never came. + bool ArriveAndWait() { + std::unique_lock lock{mutex_}; + if (abandoned_) return false; + const size_t generation = generation_; + if (++arrived_ == parties_) { + arrived_ = 0; + ++generation_; + cv_.notify_all(); + return true; + } + cv_.wait_for(lock, std::chrono::seconds(30), [&] { return generation_ != generation || abandoned_; }); + return generation_ != generation; + } + + void Abandon() { + { + std::lock_guard lock{mutex_}; + abandoned_ = true; + } + cv_.notify_all(); + } + + size_t Completed() { + std::lock_guard lock{mutex_}; + return generation_; + } + + private: + const size_t parties_; + std::mutex mutex_; + std::condition_variable cv_; + size_t arrived_{}; + size_t generation_{}; + bool abandoned_{}; +}; + +// Sets a process environment variable for its lifetime and restores the previous state: the stand-in for the variable a +// process-wide publication of the prompt length would write. +class ScopedProcessEnv { + public: + ScopedProcessEnv(const char* name, const char* value) : name_{name} { + if (const char* previous = std::getenv(name)) previous_ = previous; + Write(value); + } + ~ScopedProcessEnv() { + if (previous_) { + Write(previous_->c_str()); + } else { + Unset(); + } + } + ScopedProcessEnv(const ScopedProcessEnv&) = delete; + ScopedProcessEnv& operator=(const ScopedProcessEnv&) = delete; + + private: + void Write(const char* value) const { +#if defined(_WIN32) + _putenv_s(name_, value); +#else + setenv(name_, value, 1); +#endif + } + void Unset() const { +#if defined(_WIN32) + _putenv_s(name_, ""); // an empty value removes the variable +#else + unsetenv(name_); +#endif + } + + const char* name_; + std::optional previous_; +}; + +struct Outcome { + std::vector tokens; + GenerationFinishReason finish_reason{GenerationFinishReason::None}; + int32_t matched_stop_string_index{-1}; + TurnUsage usage{}; + size_t runs{}; + size_t finished_events{}; + uint32_t failure_flags{}; +}; + +class Qwen38GraphEngineTest : public ::testing::Test { + protected: + void SetUp() override { + model_ = LoadSyntheticPagedModel(); + model_->config_->engine.dynamic_batching.reset(); + model_->config_->model.qwen38_graph.prefill.chunk = 4; + device_ = std::make_shared(); + } + + std::shared_ptr MakeEngine() { + return std::make_shared(model_, CreateQwen38GraphDependencies(model_, device_)); + } + + static std::vector Prompt(size_t count, int32_t first = 12) { + std::vector prompt(count); + for (size_t i = 0; i < count; ++i) prompt[i] = first + static_cast(i % 20); + return prompt; + } + + // Runs the Engine until it has no work and returns what `request` produced. + static Outcome RunToCompletion(Engine& engine, const std::shared_ptr& request, size_t max_runs = 10000) { + Outcome outcome; + std::array buffer; + while (engine.HasPendingRequests()) { + if (++outcome.runs > max_runs) { + ADD_FAILURE() << "the Engine did not finish"; + break; + } + const size_t count = engine.Run(buffer); + for (size_t i = 0; i < count; ++i) { + const EngineEvent& event = buffer[i]; + if (event.request != request) continue; + if (event.flags & EngineEventFlagToken) outcome.tokens.push_back(event.token); + if (event.flags & (EngineEventFlagFailed | EngineEventFlagRetryable | EngineEventFlagCapacityBlocked)) { + outcome.failure_flags |= event.flags; + } + if (event.flags & EngineEventFlagTurnFinished) { + ++outcome.finished_events; + outcome.finish_reason = event.finish_reason; + outcome.matched_stop_string_index = event.matched_stop_string_index; + outcome.usage = event.usage; + } + } + } + return outcome; + } + + std::vector Expected(size_t first_position, size_t count) const { + std::vector tokens; + for (size_t i = 0; i < count; ++i) tokens.push_back(device_->TokenAt(first_position + i)); + return tokens; + } + + // Runs the Engine until an event reports a failure and returns its error code (nullopt: none did). + static std::optional RunUntilFailure(Engine& engine, size_t max_runs = 100) { + std::array buffer; + for (size_t run = 0; run < max_runs && engine.HasPendingRequests(); ++run) { + const size_t count = engine.Run(buffer); + for (size_t i = 0; i < count; ++i) { + if (buffer[i].flags & EngineEventFlagFailed) return buffer[i].error_code; + } + } + return std::nullopt; + } + + static TurnOptions Limit(size_t max_generated_tokens) { + TurnOptions options; + options.max_generated_tokens = max_generated_tokens; + return options; + } + + std::shared_ptr Begin(Engine& engine, std::span prompt, size_t max_generated_tokens, + std::optional max_session_tokens = std::nullopt) { + RequestOptions request_options; + request_options.max_session_tokens = max_session_tokens; + auto request = engine.CreateRequest(request_options); + auto options = Limit(max_generated_tokens); + options.request = request; + request->BeginTurn(prompt, options); + return request; + } + + std::shared_ptr BeginWithStopString(Engine& engine, std::span prompt, size_t max_generated_tokens) { + auto request = engine.CreateRequest(); + auto options = Limit(max_generated_tokens); + options.request = request; + options.stop_strings = {"STOP"}; + request->BeginTurn(prompt, options); + return request; + } + + std::shared_ptr model_; + std::shared_ptr device_; +}; + +// ---- token path ------------------------------------------------------------------------------ + +TEST_F(Qwen38GraphEngineTest, GeneratesTheScriptedStreamAndEndsAtTheTurnLimit) { + auto engine = MakeEngine(); + const auto prompt = Prompt(5); + auto request = Begin(*engine, prompt, /*max_generated_tokens=*/7); + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(outcome.tokens, Expected(5, 7)); + EXPECT_EQ(outcome.finish_reason, GenerationFinishReason::TurnLimit); + EXPECT_EQ(outcome.finished_events, 1u); + EXPECT_EQ(outcome.usage.prompt_tokens, 5u); + EXPECT_EQ(outcome.usage.generated_tokens, 7u); + EXPECT_EQ(request->TokensCpu().size(), 12u); + + // Accept pattern {accept, reject}: tokens 1 | 2 | 1 | 2 | 1 = 7 over 4 cycles; the last cycle was waited for and + // no further cycle was enqueued. + EXPECT_EQ(device_->CallCount("wait"), 4u); + EXPECT_EQ(device_->cycles_enqueued, 4u); + EXPECT_FALSE(device_->inflight()); + EXPECT_EQ(device_->calls.back(), "advance:0"); + EXPECT_EQ(device_->ingested, prompt); + + const auto stats = engine->GetSpeculativeStats(); + EXPECT_EQ(stats.rounds, 4u); + EXPECT_EQ(stats.draft_tokens_proposed, 4u); + EXPECT_EQ(stats.draft_tokens_accepted, 2u); + EXPECT_EQ(stats.bonus_tokens, 2u); + EXPECT_EQ(stats.correction_tokens, 2u); + EXPECT_EQ(stats.tokens_emitted, 7u); + EXPECT_EQ(stats.acceptance_length_histogram[0], 2u); + EXPECT_EQ(stats.acceptance_length_histogram[1], 2u); + EXPECT_FLOAT_EQ(stats.acceptance_rate, 0.5f); + EXPECT_FLOAT_EQ(stats.avg_draft_tokens_per_round, 1.0f); +} + +TEST_F(Qwen38GraphEngineTest, TheTurnLimitCutsAnAcceptedCycleAndEndsTheLoop) { + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(5), /*max_generated_tokens=*/2); + + const auto outcome = RunToCompletion(*engine, request); + + // First token + the accepted first cycle would commit 3 tokens; only 2 are allowed. + EXPECT_EQ(outcome.tokens, Expected(5, 2)); + EXPECT_EQ(outcome.finish_reason, GenerationFinishReason::TurnLimit); + EXPECT_EQ(device_->calls.back(), "advance:0"); + EXPECT_FALSE(device_->inflight()); + EXPECT_EQ(engine->GetSpeculativeStats().tokens_discarded, 1u); +} + +TEST_F(Qwen38GraphEngineTest, ATurnLimitOfOneEndsOnThePromptPredictionWithoutStartingACycle) { + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(5), /*max_generated_tokens=*/1); + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(outcome.tokens, Expected(5, 1)); + EXPECT_EQ(outcome.finish_reason, GenerationFinishReason::TurnLimit); + EXPECT_EQ(outcome.finished_events, 1u); + EXPECT_EQ(device_->cycles_enqueued, 0u); + EXPECT_EQ(device_->CallCount("wait"), 0u); + EXPECT_EQ(engine->GetSpeculativeStats().rounds, 0u); +} + +TEST_F(Qwen38GraphEngineTest, AnEndOfSequenceTokenFirstInACycleEndsTheTurnWithoutBeingDelivered) { + device_->overrides[6] = kEos; + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(5), /*max_generated_tokens=*/20); + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(outcome.tokens, Expected(5, 1)); + EXPECT_EQ(outcome.finish_reason, GenerationFinishReason::EosToken); + EXPECT_EQ(outcome.usage.generated_tokens, 1u); + EXPECT_EQ(device_->calls.back(), "advance:0"); + EXPECT_FALSE(device_->inflight()); + EXPECT_EQ(device_->CallCount("wait"), 1u); + const auto stats = engine->GetSpeculativeStats(); + EXPECT_EQ(stats.tokens_emitted, 1u); + EXPECT_EQ(stats.tokens_discarded, 2u); + EXPECT_EQ(stats.draft_tokens_accepted, 0u); + EXPECT_EQ(stats.bonus_tokens, 0u); +} + +TEST_F(Qwen38GraphEngineTest, AnEndOfSequenceTokenSecondInAnAcceptedCycleKeepsTheFirstToken) { + device_->overrides[7] = kEos; + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(5), /*max_generated_tokens=*/20); + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(outcome.tokens, Expected(5, 2)); // positions 5 and 6; the EOS at 7 is not delivered + EXPECT_EQ(outcome.finish_reason, GenerationFinishReason::EosToken); + EXPECT_EQ(outcome.usage.generated_tokens, 2u); + EXPECT_EQ(device_->calls.back(), "advance:0"); + const auto stats = engine->GetSpeculativeStats(); + EXPECT_EQ(stats.tokens_emitted, 2u); + EXPECT_EQ(stats.tokens_discarded, 1u); + EXPECT_EQ(stats.draft_tokens_accepted, 1u); + EXPECT_EQ(stats.bonus_tokens, 0u); +} + +TEST_F(Qwen38GraphEngineTest, AnEndOfSequenceFirstTokenNeverStartsACycle) { + device_->overrides[5] = kEos; + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(5), /*max_generated_tokens=*/20); + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_TRUE(outcome.tokens.empty()); + EXPECT_EQ(outcome.finish_reason, GenerationFinishReason::EosToken); + EXPECT_EQ(outcome.finished_events, 1u); + EXPECT_EQ(device_->CallCount("begin:0"), 1u); + EXPECT_EQ(device_->CallCount("wait"), 0u); + EXPECT_EQ(device_->cycles_enqueued, 0u); + const auto stats = engine->GetSpeculativeStats(); + EXPECT_EQ(stats.tokens_emitted, 0u); + EXPECT_EQ(stats.tokens_discarded, 0u); + EXPECT_EQ(stats.rounds, 0u); +} + +TEST_F(Qwen38GraphEngineTest, TheSessionLimitEndsTheTurnWithTheContextLimitReason) { + auto engine = MakeEngine(); + // 5 prompt tokens, room for 3 more. + auto request = Begin(*engine, Prompt(5), /*max_generated_tokens=*/50, /*max_session_tokens=*/8); + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(outcome.tokens, Expected(5, 3)); + EXPECT_EQ(outcome.finish_reason, GenerationFinishReason::ContextLimit); + EXPECT_EQ(request->CurrentSequenceLength(), 8); + EXPECT_EQ(device_->calls.back(), "advance:0"); +} + +TEST_F(Qwen38GraphEngineTest, ASessionThatFillsTheWholeKvCapacityEndsWithTheContextLimitReason) { + auto engine = MakeEngine(); + const size_t limit = engine->GetCapabilities().max_request_length; + auto request = Begin(*engine, Prompt(5), /*max_generated_tokens=*/limit, /*max_session_tokens=*/limit); + + const auto outcome = RunToCompletion(*engine, request); + + // The device ledger throws when a cycle does not fit, so reaching the limit cleanly proves the reserved tail is honored. + EXPECT_EQ(outcome.failure_flags, 0u); + EXPECT_EQ(outcome.finish_reason, GenerationFinishReason::ContextLimit); + EXPECT_EQ(outcome.tokens, Expected(5, limit - 5)); + EXPECT_EQ(request->CurrentSequenceLength(), static_cast(limit)); +} + +TEST_F(Qwen38GraphEngineTest, AStopStringEndsTheTurnAndTheWastedLookAheadCycleIsHarmless) { + // Tokens 5 and 6 decode to "ST" and "OP": the stop string "STOP" completes on the first token of the first cycle. + device_->overrides[5] = 5; + device_->overrides[6] = 6; + auto engine = MakeEngine(); + auto request = engine->CreateRequest(); + TurnOptions options; + options.request = request; + options.max_generated_tokens = 20; + options.stop_strings = {"STOP"}; + request->BeginTurn(Prompt(5), options); + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(outcome.finish_reason, GenerationFinishReason::StopString); + EXPECT_EQ(outcome.matched_stop_string_index, 0); + EXPECT_EQ(outcome.tokens, (std::vector{5, 6})); + const auto stop_stats = engine->GetSpeculativeStats(); + EXPECT_EQ(stop_stats.tokens_emitted, 2u); + EXPECT_EQ(stop_stats.tokens_discarded, 1u); + EXPECT_EQ(stop_stats.draft_tokens_accepted, 1u); + EXPECT_EQ(stop_stats.bonus_tokens, 0u); + // The device cannot know about stop strings: the cycle after the one that matched is still in flight... + EXPECT_TRUE(device_->inflight()); + EXPECT_TRUE(device_->running()); + + // ...and a new request waits for it (Reset) before reusing the device. + const size_t resets = device_->CallCount("reset"); + device_->overrides.clear(); + auto next = Begin(*engine, Prompt(5), /*max_generated_tokens=*/3); + const auto next_outcome = RunToCompletion(*engine, next); + EXPECT_EQ(device_->CallCount("reset"), resets + 1); + EXPECT_EQ(device_->status_checks, 2u); + EXPECT_EQ(next_outcome.tokens, Expected(5, 3)); +} + +TEST_F(Qwen38GraphEngineTest, StopStringsSpanningTheTwoTokensOfACycleAreMatched) { + device_->overrides[6] = 5; + device_->overrides[7] = 6; + auto engine = MakeEngine(); + auto request = engine->CreateRequest(); + TurnOptions options; + options.request = request; + options.stop_strings = {"STOP"}; + request->BeginTurn(Prompt(5), options); + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(outcome.finish_reason, GenerationFinishReason::StopString); + EXPECT_EQ(outcome.tokens, (std::vector{device_->TokenAt(5), 5, 6})); +} + +TEST_F(Qwen38GraphEngineTest, EveryDraftRejectedCommitsTheGreedyTokenOfEachCycleOnly) { + device_->accept_pattern = {false}; // the fold.force_reject hook: no draft ever survives its verify + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(5), /*max_generated_tokens=*/6); + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(outcome.tokens, Expected(5, 6)); // the same stream as with any other accept pattern + EXPECT_EQ(device_->CallCount("wait"), 5u); // the first token comes from the prompt, then one token per cycle + const auto stats = engine->GetSpeculativeStats(); + EXPECT_EQ(stats.rounds, 5u); + EXPECT_EQ(stats.draft_tokens_accepted, 0u); + EXPECT_EQ(stats.zero_accept_rounds, 5u); + EXPECT_EQ(stats.acceptance_length_histogram[0], 5u); + EXPECT_EQ(device_->ledger_length(), 5u + 5u); +} + +TEST_F(Qwen38GraphEngineTest, EveryDraftAcceptedCommitsTwoTokensPerCycle) { + device_->accept_pattern = {true}; + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(5), /*max_generated_tokens=*/9); + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(outcome.tokens, Expected(5, 9)); + EXPECT_EQ(device_->CallCount("wait"), 4u); // 1 + 4 * 2 tokens + const auto stats = engine->GetSpeculativeStats(); + EXPECT_EQ(stats.draft_tokens_accepted, 4u); + EXPECT_EQ(stats.full_accept_rounds, 4u); + EXPECT_EQ(device_->ledger_length(), 5u + 8u); + EXPECT_EQ(device_->ledger_bank(), static_cast(device_->ledger_length() % 3)); +} + +TEST_F(Qwen38GraphEngineTest, WithoutMtpEveryCycleCommitsOneToken) { + device_ = std::make_shared(/*capacity=*/68, /*mtp=*/false); + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(4), /*max_generated_tokens=*/5); + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(outcome.tokens, Expected(4, 5)); + EXPECT_EQ(device_->CallCount("wait"), 4u); + const auto stats = engine->GetSpeculativeStats(); + EXPECT_EQ(stats.draft_tokens_proposed, 0u); + EXPECT_EQ(stats.target_forward_passes, 4u); + EXPECT_EQ(stats.tokens_emitted, 5u); +} + +// ---- speculative statistics of multi-draft cycles -------------------------------------------- +// +// The verify of every cycle proposes 3 drafts and accepts the next entry of the given counts of them. The prompt has 5 +// tokens: position 5 is the first token (it comes from the prompt), the first cycle delivers positions 6.. in order +// (accepted drafts, then the correction or, with all three accepted, the bonus). Engine::RecordSpeculativeCommit is the +// reference: a draft counts as evaluated when the turn examined it before its terminal token. + +TEST_F(Qwen38GraphEngineTest, ATurnLimitInsideTheAcceptedDraftsOnlyCountsTheDraftItConsumed) { + device_->SetDrafts(3, {3}); + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(5), /*max_generated_tokens=*/2); // room for one token of the first cycle + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(outcome.tokens, Expected(5, 2)); + EXPECT_EQ(outcome.finish_reason, GenerationFinishReason::TurnLimit); + const auto stats = engine->GetSpeculativeStats(); + EXPECT_EQ(stats.rounds, 1u); + EXPECT_EQ(stats.draft_tokens_proposed, 3u); + EXPECT_EQ(stats.draft_tokens_evaluated, 1u); // not accepted + 1: the second draft was never examined + EXPECT_EQ(stats.draft_tokens_accepted, 1u); + EXPECT_EQ(stats.acceptance_length_histogram[1], 1u); + EXPECT_FLOAT_EQ(stats.acceptance_rate, 1.0f); + EXPECT_EQ(stats.tokens_discarded, 3u); +} + +TEST_F(Qwen38GraphEngineTest, ATurnLimitAfterSeveralAcceptedDraftsCountsExactlyTheDraftsItConsumed) { + device_->SetDrafts(3, {3}); + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(5), /*max_generated_tokens=*/3); // room for two tokens of the first cycle + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(outcome.tokens, Expected(5, 3)); + const auto stats = engine->GetSpeculativeStats(); + EXPECT_EQ(stats.draft_tokens_evaluated, 2u); + EXPECT_EQ(stats.draft_tokens_accepted, 2u); + EXPECT_EQ(stats.acceptance_length_histogram[2], 1u); + EXPECT_FLOAT_EQ(stats.acceptance_rate, 1.0f); +} + +TEST_F(Qwen38GraphEngineTest, ATurnLimitOnTheCorrectionTokenCountsTheRejectedPositionToo) { + device_->SetDrafts(3, {1}); + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(5), /*max_generated_tokens=*/3); // the accepted draft and the correction + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(outcome.tokens, Expected(5, 3)); + const auto stats = engine->GetSpeculativeStats(); + EXPECT_EQ(stats.draft_tokens_evaluated, 2u); // the accepted draft and the position that was rejected + EXPECT_EQ(stats.draft_tokens_accepted, 1u); + EXPECT_FLOAT_EQ(stats.acceptance_rate, 0.5f); +} + +TEST_F(Qwen38GraphEngineTest, ATurnLimitRightAfterTheLastAcceptedDraftDoesNotReachTheRejectedPosition) { + device_->SetDrafts(3, {1}); + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(5), /*max_generated_tokens=*/2); // only the accepted draft fits + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(outcome.tokens, Expected(5, 2)); + const auto stats = engine->GetSpeculativeStats(); + EXPECT_EQ(stats.draft_tokens_evaluated, 1u); + EXPECT_EQ(stats.draft_tokens_accepted, 1u); + EXPECT_FLOAT_EQ(stats.acceptance_rate, 1.0f); +} + +TEST_F(Qwen38GraphEngineTest, TheSessionLimitTruncatesTheEvaluatedDraftsLikeTheTurnLimit) { + device_->SetDrafts(3, {3}); + auto engine = MakeEngine(); + // 5 prompt tokens, room for 2 more: the first token and one token of the cycle. + auto request = Begin(*engine, Prompt(5), /*max_generated_tokens=*/50, /*max_session_tokens=*/7); + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(outcome.tokens, Expected(5, 2)); + EXPECT_EQ(outcome.finish_reason, GenerationFinishReason::ContextLimit); + const auto stats = engine->GetSpeculativeStats(); + EXPECT_EQ(stats.draft_tokens_evaluated, 1u); + EXPECT_EQ(stats.draft_tokens_accepted, 1u); + EXPECT_FLOAT_EQ(stats.acceptance_rate, 1.0f); +} + +TEST_F(Qwen38GraphEngineTest, AStopStringInsideTheAcceptedDraftsOnlyCountsTheDraftsUpToTheMatch) { + device_->SetDrafts(3, {3}); + device_->overrides[6] = 5; // "ST", "OP": the stop string completes on the second accepted draft + device_->overrides[7] = 6; + auto engine = MakeEngine(); + auto request = BeginWithStopString(*engine, Prompt(5), /*max_generated_tokens=*/20); + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(outcome.finish_reason, GenerationFinishReason::StopString); + EXPECT_EQ(outcome.tokens, (std::vector{device_->TokenAt(5), 5, 6})); + const auto stats = engine->GetSpeculativeStats(); + EXPECT_EQ(stats.draft_tokens_evaluated, 2u); + EXPECT_EQ(stats.draft_tokens_accepted, 2u); + EXPECT_FLOAT_EQ(stats.acceptance_rate, 1.0f); + EXPECT_EQ(stats.tokens_discarded, 2u); +} + +TEST_F(Qwen38GraphEngineTest, AStopStringOnTheCorrectionTokenCountsTheRejectedPositionToo) { + device_->SetDrafts(3, {1}); + device_->overrides[6] = 5; // the accepted draft + device_->overrides[7] = 6; // the correction completes the stop string + auto engine = MakeEngine(); + auto request = BeginWithStopString(*engine, Prompt(5), /*max_generated_tokens=*/20); + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(outcome.finish_reason, GenerationFinishReason::StopString); + const auto stats = engine->GetSpeculativeStats(); + EXPECT_EQ(stats.draft_tokens_evaluated, 2u); + EXPECT_EQ(stats.draft_tokens_accepted, 1u); + EXPECT_FLOAT_EQ(stats.acceptance_rate, 0.5f); +} + +TEST_F(Qwen38GraphEngineTest, AnEndOfSequenceDraftIsExaminedButNotAccepted) { + device_->SetDrafts(3, {3}); + device_->overrides[7] = kEos; // the second accepted draft + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(5), /*max_generated_tokens=*/20); + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(outcome.tokens, Expected(5, 2)); + EXPECT_EQ(outcome.finish_reason, GenerationFinishReason::EosToken); + const auto stats = engine->GetSpeculativeStats(); + EXPECT_EQ(stats.draft_tokens_evaluated, 2u); // like Request::CommitAcceptedDraftsForTransaction: the EOS was examined + EXPECT_EQ(stats.draft_tokens_accepted, 1u); // ...but nothing was appended for it + EXPECT_EQ(stats.acceptance_length_histogram[1], 1u); + EXPECT_FLOAT_EQ(stats.acceptance_rate, 0.5f); +} + +TEST_F(Qwen38GraphEngineTest, AnEndOfSequenceCorrectionTokenCountsTheRejectedPositionToo) { + device_->SetDrafts(3, {1}); + device_->overrides[7] = kEos; // the correction + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(5), /*max_generated_tokens=*/20); + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(outcome.tokens, Expected(5, 2)); + EXPECT_EQ(outcome.finish_reason, GenerationFinishReason::EosToken); + const auto stats = engine->GetSpeculativeStats(); + EXPECT_EQ(stats.draft_tokens_evaluated, 2u); + EXPECT_EQ(stats.draft_tokens_accepted, 1u); + EXPECT_EQ(stats.correction_tokens, 0u); // it was not delivered +} + +TEST_F(Qwen38GraphEngineTest, AnEndOfSequenceBonusTokenIsNotAnExaminedDraft) { + device_->SetDrafts(3, {3}); + device_->overrides[9] = kEos; // the bonus token after three accepted drafts + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(5), /*max_generated_tokens=*/20); + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(outcome.tokens, Expected(5, 4)); + EXPECT_EQ(outcome.finish_reason, GenerationFinishReason::EosToken); + const auto stats = engine->GetSpeculativeStats(); + EXPECT_EQ(stats.draft_tokens_evaluated, 3u); + EXPECT_EQ(stats.draft_tokens_accepted, 3u); + EXPECT_FLOAT_EQ(stats.acceptance_rate, 1.0f); +} + +TEST_F(Qwen38GraphEngineTest, RoundsAreClassifiedByTheNumberOfDraftsThatWereAccepted) { + device_->SetDrafts(3, {0, 1, 2, 3}); + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(5), /*max_generated_tokens=*/11); // the first token, then 1 + 2 + 3 + 4 tokens + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(outcome.tokens, Expected(5, 11)); + const auto stats = engine->GetSpeculativeStats(); + EXPECT_EQ(stats.rounds, 4u); + EXPECT_EQ(stats.draft_tokens_proposed, 12u); + EXPECT_EQ(stats.draft_tokens_evaluated, 9u); // 1 + 2 + 3 + 3: every round but the full one examined its rejected draft + EXPECT_EQ(stats.draft_tokens_accepted, 6u); + EXPECT_FLOAT_EQ(stats.acceptance_rate, 6.0f / 9.0f); + EXPECT_EQ(stats.zero_accept_rounds, 1u); + EXPECT_EQ(stats.partial_accept_rounds, 2u); + EXPECT_EQ(stats.full_accept_rounds, 1u); + for (size_t accepted = 0; accepted <= 3; ++accepted) { + EXPECT_EQ(stats.acceptance_length_histogram[accepted], 1u) << "accepted=" << accepted; + } + EXPECT_EQ(stats.correction_tokens, 3u); + EXPECT_EQ(stats.bonus_tokens, 1u); + EXPECT_EQ(stats.tokens_emitted, 11u); + EXPECT_EQ(stats.tokens_discarded, 0u); +} + +TEST_F(Qwen38GraphEngineTest, APartialAcceptRoundIsCountedEvenWhenTheTurnEndsBeforeItsCorrectionToken) { + device_->SetDrafts(3, {1}); + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(5), /*max_generated_tokens=*/2); // the accepted draft fits, the correction does not + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(outcome.tokens, Expected(5, 2)); + const auto stats = engine->GetSpeculativeStats(); + EXPECT_EQ(stats.partial_accept_rounds, 1u); + EXPECT_EQ(stats.zero_accept_rounds, 0u); + EXPECT_EQ(stats.full_accept_rounds, 0u); + EXPECT_EQ(stats.correction_tokens, 0u); // no correction was delivered, but the round is partial all the same + EXPECT_EQ(stats.bonus_tokens, 0u); +} + +TEST_F(Qwen38GraphEngineTest, ARoundTruncatedInsideItsAcceptedDraftsIsClassifiedByTheDraftsItDelivered) { + device_->SetDrafts(3, {3}); + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(5), /*max_generated_tokens=*/2); // one of the three accepted drafts fits + + RunToCompletion(*engine, request); + + const auto stats = engine->GetSpeculativeStats(); + EXPECT_EQ(stats.draft_tokens_accepted, 1u); + EXPECT_EQ(stats.partial_accept_rounds, 1u); + EXPECT_EQ(stats.full_accept_rounds, 0u); + EXPECT_EQ(stats.bonus_tokens, 0u); +} + +TEST_F(Qwen38GraphEngineTest, AnEndOfSequenceFirstDraftIsAZeroAcceptRound) { + device_->SetDrafts(3, {3}); + device_->overrides[6] = kEos; + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(5), /*max_generated_tokens=*/20); + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(outcome.tokens, Expected(5, 1)); + EXPECT_EQ(outcome.finish_reason, GenerationFinishReason::EosToken); + const auto stats = engine->GetSpeculativeStats(); + EXPECT_EQ(stats.draft_tokens_evaluated, 1u); + EXPECT_EQ(stats.draft_tokens_accepted, 0u); + EXPECT_EQ(stats.zero_accept_rounds, 1u); + EXPECT_EQ(stats.partial_accept_rounds, 0u); + EXPECT_EQ(stats.acceptance_length_histogram[0], 1u); +} + +TEST_F(Qwen38GraphEngineTest, RoundCountersOfSeveralRequestsAddUp) { + device_->SetDrafts(3, {1}); + auto engine = MakeEngine(); + for (int i = 0; i < 2; ++i) { + auto request = Begin(*engine, Prompt(5, 12 + i), /*max_generated_tokens=*/3); // first token, accepted draft, correction + ASSERT_EQ(RunToCompletion(*engine, request).tokens, Expected(5, 3)); + } + + const auto stats = engine->GetSpeculativeStats(); + EXPECT_EQ(stats.rounds, 2u); + EXPECT_EQ(stats.partial_accept_rounds, 2u); + EXPECT_EQ(stats.acceptance_length_histogram[1], 2u); + EXPECT_EQ(stats.correction_tokens, 2u); + EXPECT_EQ(stats.draft_tokens_evaluated, 4u); + EXPECT_EQ(stats.draft_tokens_accepted, 2u); +} + +// ---- prefill, cancellation, queueing --------------------------------------------------------- + +TEST_F(Qwen38GraphEngineTest, ThePromptIsIngestedInChunksOverSeveralSteps) { + auto engine = MakeEngine(); // prefill.chunk = 4 + const auto prompt = Prompt(11); + auto request = Begin(*engine, prompt, /*max_generated_tokens=*/2); + + std::array buffer; + EXPECT_EQ(engine->Run(buffer), 0u); // chunk 1: no event + EXPECT_EQ(engine->Run(buffer), 0u); // chunk 2 + EXPECT_TRUE(engine->HasPendingRequests()); + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(device_->CallCount("prefill:4"), 2u); + EXPECT_EQ(device_->CallCount("prefill:3"), 1u); + EXPECT_EQ(device_->ingested, prompt); + EXPECT_EQ(outcome.tokens, Expected(11, 2)); +} + +TEST_F(Qwen38GraphEngineTest, AnUnsetChunkUsesTheDefaultOfTheDecodeGraphPath) { + device_ = std::make_shared(/*capacity=*/4200); + model_->config_->model.context_length = 4200; + model_->config_->model.qwen38_graph.prefill.chunk = 0; + auto engine = MakeEngine(); + const auto prompt = Prompt(2050); + auto request = Begin(*engine, prompt, /*max_generated_tokens=*/2, /*max_session_tokens=*/4100); + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(device_->CallCount("prefill:2048"), 1u); // EffectiveQwen38PrefillChunkTokens of the unset key + EXPECT_EQ(device_->CallCount("prefill:2"), 1u); + EXPECT_EQ(outcome.tokens, Expected(2050, 2)); +} + +TEST_F(Qwen38GraphEngineTest, AnInvalidChunkConfigurationIsRejectedWhenTheBackendIsAssembled) { + auto& prefill = model_->config_->model.qwen38_graph.prefill; + for (const int chunk : {1, 5, -2}) { // an odd chunk would shift the S=2 pairs of everything after it + prefill.chunk = chunk; + EXPECT_THROW(CreateQwen38GraphDependencies(model_, device_), std::runtime_error) << "prefill.chunk=" << chunk; + } + prefill.chunk = 0; + EXPECT_NO_THROW(CreateQwen38GraphDependencies(model_, device_)); +} + +TEST(Qwen38GraphSessionTest, AnOddChunkOfTheOptionsIsRoundedDownToAnEvenOne) { + auto device = std::make_shared(); + Qwen38GraphSession::Options options; + options.prefill_chunk_tokens = 5; // reachable by direct construction only: the configuration rejects odd chunks + Qwen38GraphSession session{device, options}; + std::vector prompt(11); + for (size_t i = 0; i < prompt.size(); ++i) prompt[i] = 12 + static_cast(i); + Qwen38StepRequest request; + request.owner = &prompt; + request.turn_id = 1; + request.tokens = prompt; + request.token_budget = 2; + request.session_budget = 50; + + EXPECT_EQ(session.Step(request).count, 0u); // chunk 1 + EXPECT_EQ(session.Step(request).count, 0u); // chunk 2 + EXPECT_EQ(session.Step(request).count, 1u); // the last chunk returns the first token + + EXPECT_EQ(device->CallCount("prefill:4"), 2u); + EXPECT_EQ(device->CallCount("prefill:3"), 1u); +} + +TEST_F(Qwen38GraphEngineTest, TheDeviceIsToldThePromptLengthAfterTheResetAndBeforeTheFirstChunkOfEveryTurn) { + auto engine = MakeEngine(); // prefill.chunk = 4 + const auto prompt = Prompt(11); + auto request = Begin(*engine, prompt, /*max_generated_tokens=*/3); + ASSERT_EQ(RunToCompletion(*engine, request).tokens, Expected(11, 3)); + + EXPECT_EQ(device_->CallCount("expect:11"), 1u); + EXPECT_EQ(device_->prefill_totals, (std::vector{11, 11, 11})); // chunks 4, 4, 3 all see the whole length + const auto position = [&](const std::string& call) { + return std::find(device_->calls.begin(), device_->calls.end(), call) - device_->calls.begin(); + }; + EXPECT_LT(position("reset"), position("expect:11")); + EXPECT_LT(position("expect:11"), position("prefill:4")); + + // The continued turn re-ingests the whole history (11 + 3 + 2 tokens) and announces its new length. + const std::vector follow_up{30, 31}; + TurnOptions options = Limit(2); + options.request = request; + request->BeginTurn(follow_up, options); + ASSERT_EQ(RunToCompletion(*engine, request).tokens, Expected(16, 2)); + EXPECT_EQ(device_->CallCount("expect:16"), 1u); + EXPECT_EQ(device_->prefill_totals, (std::vector{11, 11, 11, 16, 16, 16, 16})); +} + +// The prompt length of a turn belongs to its device. Two Engines on two owner threads (each with its own model and device) +// are held between the publication of the length and the first chunk of the turn, so the schedule "A publishes, B +// publishes, A ingests" of a process-wide variable happens at every turn start. Each device must see the length of its own +// turn in every chunk, after a Reset (continued turn) and after a cancellation, and nothing may travel through the process +// environment. +TEST_F(Qwen38GraphEngineTest, EnginesOnTwoOwnerThreadsKeepTheirOwnExpectedPrefillTotals) { + struct Owner { + size_t prompt; // the first turn; a continued turn follows it + size_t canceled_prompt; // canceled after its first chunk + size_t last_prompt; + std::shared_ptr model; + std::shared_ptr device; + std::vector> tokens; + std::vector finish_reasons; + std::string error; + }; + const auto make_owner = [](size_t prompt, size_t canceled_prompt, size_t last_prompt) { + Owner owner{prompt, canceled_prompt, last_prompt, LoadSyntheticPagedModel(), std::make_shared(/*capacity=*/128), {}, {}, {}}; + owner.model->config_->engine.dynamic_batching.reset(); + owner.model->config_->model.qwen38_graph.prefill.chunk = 4; + return owner; + }; + Owner a = make_owner(21, 15, 6); + Owner b = make_owner(7, 9, 10); + + Rendezvous rendezvous{2}; + for (Owner* owner : {&a, &b}) { + owner->device->on_expect_prefill = [&rendezvous](size_t) { + if (!rendezvous.ArriveAndWait()) throw std::runtime_error("the other owner thread never reached its publication"); + }; + } + + const auto run_owner = [&rendezvous](Owner& owner) { + try { + auto engine = std::make_shared(owner.model, CreateQwen38GraphDependencies(owner.model, owner.device)); + const auto drive = [&](const std::shared_ptr& request) { + std::vector tokens; + auto reason = GenerationFinishReason::None; + std::array buffer; + for (size_t runs = 0; engine->HasPendingRequests(); ++runs) { + if (runs > 10000) throw std::runtime_error("the Engine did not finish"); + const size_t count = engine->Run(buffer); + for (size_t i = 0; i < count; ++i) { + if (buffer[i].request != request) continue; + if (buffer[i].flags & EngineEventFlagFailed) throw std::runtime_error("the request failed"); + if (buffer[i].flags & EngineEventFlagToken) tokens.push_back(buffer[i].token); + if (buffer[i].flags & EngineEventFlagTurnFinished) reason = buffer[i].finish_reason; + } + } + owner.tokens.push_back(std::move(tokens)); + owner.finish_reasons.push_back(reason); + }; + + auto request = engine->CreateRequest(); + TurnOptions first = Limit(3); + first.request = request; + request->BeginTurn(Prompt(owner.prompt), first); + drive(request); + + const std::vector follow_up{30, 31}; // the continued turn re-ingests the whole history after a Reset + TurnOptions second = Limit(2); + second.request = request; + request->BeginTurn(follow_up, second); + drive(request); + + auto canceled = engine->CreateRequest(); + TurnOptions canceled_options; + canceled_options.request = canceled; + const uint64_t turn = canceled->BeginTurn(Prompt(owner.canceled_prompt), canceled_options); + std::array buffer; + engine->Run(buffer); // the first chunk of the prompt only + engine->CancelRequest(canceled, turn); + drive(canceled); + + auto last = engine->CreateRequest(); + TurnOptions last_options = Limit(2); + last_options.request = last; + last->BeginTurn(Prompt(owner.last_prompt), last_options); + drive(last); + } catch (const std::exception& error) { + owner.error = error.what(); + rendezvous.Abandon(); + } + }; + + ScopedProcessEnv sentinel{Qwen38OrtEnv::kPrefillTotalTokens, "424242"}; + std::thread thread_a{[&] { run_owner(a); }}; + std::thread thread_b{[&] { run_owner(b); }}; + thread_a.join(); + thread_b.join(); + + EXPECT_EQ(rendezvous.Completed(), 4u); // one publication per turn on each device + const auto totals = [](size_t total, size_t chunk) { return std::vector((total + chunk - 1) / chunk, total); }; + for (Owner* owner : {&a, &b}) { + SCOPED_TRACE("prompt of " + std::to_string(owner->prompt) + " tokens"); + EXPECT_EQ(owner->error, ""); + auto expected = totals(owner->prompt, 4); + for (const auto& part : {totals(owner->prompt + 3 + 2, 4), std::vector{owner->canceled_prompt}, + totals(owner->last_prompt, 4)}) { + expected.insert(expected.end(), part.begin(), part.end()); + } + EXPECT_EQ(owner->device->prefill_totals, expected); + + const auto& device = *owner->device; + ASSERT_EQ(owner->tokens.size(), 4u); + EXPECT_EQ(owner->tokens[0], (std::vector{device.TokenAt(owner->prompt), device.TokenAt(owner->prompt + 1), + device.TokenAt(owner->prompt + 2)})); + EXPECT_EQ(owner->tokens[1], + (std::vector{device.TokenAt(owner->prompt + 5), device.TokenAt(owner->prompt + 6)})); + EXPECT_TRUE(owner->tokens[2].empty()); + EXPECT_EQ(owner->tokens[3], + (std::vector{device.TokenAt(owner->last_prompt), device.TokenAt(owner->last_prompt + 1)})); + EXPECT_EQ(owner->finish_reasons, + (std::vector{GenerationFinishReason::TurnLimit, GenerationFinishReason::TurnLimit, + GenerationFinishReason::Canceled, GenerationFinishReason::TurnLimit})); + } + + const char* published = std::getenv(Qwen38OrtEnv::kPrefillTotalTokens); + ASSERT_NE(published, nullptr); + EXPECT_STREQ(published, "424242") << "the executor must not publish the prompt length in the process environment"; +} + +TEST_F(Qwen38GraphEngineTest, CancelDuringPrefillStopsBeforeDecodingAndTheNextRequestStartsClean) { + auto engine = MakeEngine(); + auto request = engine->CreateRequest(); + TurnOptions options; + options.request = request; + const uint64_t turn = request->BeginTurn(Prompt(11), options); + std::array buffer; + ASSERT_EQ(engine->Run(buffer), 0u); // first chunk only + + EXPECT_TRUE(engine->CancelRequest(request, turn)); + const size_t count = engine->Run(buffer); + ASSERT_EQ(count, 1u); + EXPECT_EQ(buffer[0].finish_reason, GenerationFinishReason::Canceled); + EXPECT_FALSE(engine->HasPendingRequests()); + EXPECT_EQ(device_->CallCount("begin:1") + device_->CallCount("begin:0"), 0u); + + auto next = Begin(*engine, Prompt(6, 20), /*max_generated_tokens=*/3); + const auto outcome = RunToCompletion(*engine, next); + EXPECT_EQ(outcome.tokens, Expected(6, 3)); + EXPECT_EQ(device_->ingested, Prompt(6, 20)); // the canceled prompt is gone from the device +} + +// A request canceled while it was still ingesting its prompt made no progress, so its next turn restarts it from scratch, +// possibly after another request took the device. The Engine does not enroll such a request in the scheduler again, so the +// scheduler that displaced it must have kept it. +TEST_F(Qwen38GraphEngineTest, ACanceledPrefillRequestRunsAgainAfterAnotherRequestTookTheSlot) { + auto engine = MakeEngine(); + auto first = engine->CreateRequest(); + TurnOptions first_options; + first_options.request = first; + const uint64_t turn = first->BeginTurn(Prompt(11), first_options); + std::array buffer; + ASSERT_EQ(engine->Run(buffer), 0u); // first chunk only + ASSERT_TRUE(engine->CancelRequest(first, turn)); + while (engine->HasPendingRequests()) engine->Run(buffer); + + auto second = Begin(*engine, Prompt(6, 20), /*max_generated_tokens=*/3); + EXPECT_EQ(RunToCompletion(*engine, second).tokens, Expected(6, 3)); + + // The restarted turn replays the canceled prompt ahead of its own input. + auto restart = Limit(3); + restart.request = first; + first->BeginTurn(Prompt(5, 40), restart); + ASSERT_TRUE(engine->HasPendingRequests()) << "the restarted request is not known to the scheduler"; + const auto outcome = RunToCompletion(*engine, first); + EXPECT_EQ(outcome.finish_reason, GenerationFinishReason::TurnLimit); + EXPECT_EQ(outcome.tokens, Expected(16, 3)); + auto replayed = Prompt(11); + const auto input = Prompt(5, 40); + replayed.insert(replayed.end(), input.begin(), input.end()); + EXPECT_EQ(device_->ingested, replayed); +} + +// While the slot is still the canceled request's own the next turn continues from what the request kept, which is the +// documented behavior of cancellation: the canceled prompt is followed by the new input. +TEST_F(Qwen38GraphEngineTest, ACanceledPrefillRequestRunsAgainWhileItsSlotIsStillItsOwn) { + auto engine = MakeEngine(); + auto first = engine->CreateRequest(); + TurnOptions first_options; + first_options.request = first; + const uint64_t turn = first->BeginTurn(Prompt(11), first_options); + std::array buffer; + ASSERT_EQ(engine->Run(buffer), 0u); // first chunk only + ASSERT_TRUE(engine->CancelRequest(first, turn)); + while (engine->HasPendingRequests()) engine->Run(buffer); + + auto restart = Limit(3); + restart.request = first; + first->BeginTurn(Prompt(5, 40), restart); + const auto outcome = RunToCompletion(*engine, first); + EXPECT_EQ(outcome.finish_reason, GenerationFinishReason::TurnLimit); + EXPECT_EQ(outcome.tokens, Expected(16, 3)); + auto replayed = Prompt(11); + const auto input = Prompt(5, 40); + replayed.insert(replayed.end(), input.begin(), input.end()); + EXPECT_EQ(device_->ingested, replayed); +} + +TEST_F(Qwen38GraphEngineTest, CancelWithACycleInFlightIsSafeAndTheNextRequestWaitsForIt) { + model_->config_->model.qwen38_graph.prefill.chunk = 64; // the first step ingests the whole prompt + auto engine = MakeEngine(); + auto request = engine->CreateRequest(); + TurnOptions options; + options.request = request; + const uint64_t turn = request->BeginTurn(Prompt(5), options); + std::array buffer; + ASSERT_GE(engine->Run(buffer), 1u); // prefill + first token (the first cycle is now in flight) + ASSERT_TRUE(device_->inflight()); + + EXPECT_TRUE(engine->CancelRequest(request, turn)); + ASSERT_EQ(engine->Run(buffer), 1u); + EXPECT_EQ(buffer[0].finish_reason, GenerationFinishReason::Canceled); + EXPECT_TRUE(device_->inflight()); // still running on the device + + auto next = Begin(*engine, Prompt(5, 25), /*max_generated_tokens=*/4); + const auto outcome = RunToCompletion(*engine, next); + EXPECT_EQ(outcome.tokens, Expected(5, 4)); + EXPECT_GE(device_->CallCount("reset"), 1u); +} + +TEST_F(Qwen38GraphEngineTest, AFurtherRequestQueuesBehindTheResidentOneAndRunsAfterIt) { + auto engine = MakeEngine(); + auto first = Begin(*engine, Prompt(5), /*max_generated_tokens=*/4); + auto second = Begin(*engine, Prompt(6, 25), /*max_generated_tokens=*/3); + + std::array buffer; + std::vector first_tokens, second_tokens; + std::vector order; + while (engine->HasPendingRequests()) { + const size_t count = engine->Run(buffer); + for (size_t i = 0; i < count; ++i) { + if (!(buffer[i].flags & EngineEventFlagToken)) continue; + (buffer[i].request == first ? first_tokens : second_tokens).push_back(buffer[i].token); + order.push_back(buffer[i].request.get()); + } + } + + EXPECT_EQ(first_tokens, Expected(5, 4)); + EXPECT_EQ(second_tokens, Expected(6, 3)); + // Every event of the first request precedes every event of the second. + const auto last_first = std::find(order.rbegin(), order.rend(), first.get()); + const auto first_second = std::find(order.begin(), order.end(), second.get()); + EXPECT_LT(std::distance(order.begin(), last_first.base()), std::distance(order.begin(), first_second) + 1); + EXPECT_GE(device_->CallCount("reset"), 1u); + EXPECT_EQ(device_->ingested, Prompt(6, 25)); + EXPECT_EQ(engine->GetCapabilities().configured_max_batch_size, 1u); +} + +// ---- continuation, rewind -------------------------------------------------------------------- + +TEST_F(Qwen38GraphEngineTest, AContinuedTurnReIngestsTheWholeHistory) { + auto engine = MakeEngine(); + const auto prompt = Prompt(5); + auto request = Begin(*engine, prompt, /*max_generated_tokens=*/3); + ASSERT_EQ(RunToCompletion(*engine, request).tokens, Expected(5, 3)); + const size_t resets = device_->CallCount("reset"); + + const std::vector follow_up{30, 31}; + TurnOptions options = Limit(2); + options.request = request; + request->BeginTurn(follow_up, options); + const auto outcome = RunToCompletion(*engine, request); + + std::vector history = prompt; + for (const int32_t token : Expected(5, 3)) history.push_back(token); + history.insert(history.end(), follow_up.begin(), follow_up.end()); + EXPECT_EQ(device_->CallCount("reset"), resets + 1); + EXPECT_EQ(device_->ingested, history); + EXPECT_EQ(outcome.tokens, Expected(history.size(), 2)); + EXPECT_EQ(outcome.finish_reason, GenerationFinishReason::TurnLimit); +} + +TEST_F(Qwen38GraphEngineTest, RewindToTheStartOfATurnReleasesTheSlotAndReplaysTheRetainedPrefix) { + auto engine = MakeEngine(); + const auto prompt = Prompt(5); + auto request = Begin(*engine, prompt, /*max_generated_tokens=*/3); + ASSERT_EQ(RunToCompletion(*engine, request).tokens, Expected(5, 3)); + TurnOptions options = Limit(2); + options.request = request; + const uint64_t discarded_turn = request->BeginTurn(std::array{30, 31}, options); + ASSERT_EQ(RunToCompletion(*engine, request).tokens, Expected(10, 2)); + + request->RewindToStartOfTurn(discarded_turn); + EXPECT_FALSE(engine->HasPendingRequests()); + EXPECT_EQ(request->CurrentSequenceLength(), 8); + + const std::vector replacement{40}; + options = Limit(2); + options.request = request; + request->BeginTurn(replacement, options); + const auto outcome = RunToCompletion(*engine, request); + + std::vector history = prompt; + for (const int32_t token : Expected(5, 3)) history.push_back(token); + history.push_back(40); + EXPECT_EQ(device_->ingested, history); + EXPECT_EQ(outcome.tokens, Expected(history.size(), 2)); +} + +// ---- dense (fast) prefill: chunk streaming and hand-off --------------------------------------- + +TEST_F(Qwen38GraphEngineTest, ADeviceWithAPreferredChunkReceivesThePromptInChunksOfThatSizeAndHandsOffOnce) { + device_->dense_prefill = true; + device_->preferred_chunk = 8; // replaces the chunk of the configuration (prefill.chunk = 4) + auto engine = MakeEngine(); + const auto prompt = Prompt(20); + auto request = Begin(*engine, prompt, /*max_generated_tokens=*/5); + + std::array buffer; + EXPECT_EQ(engine->Run(buffer), 0u); // chunk 1 of 3: no event, no hand-off yet + EXPECT_EQ(device_->handoffs, 0u); + EXPECT_EQ(engine->Run(buffer), 0u); // chunk 2 + EXPECT_EQ(device_->handoffs, 0u); + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(device_->CallCount("prefill:8"), 2u); + EXPECT_EQ(device_->CallCount("prefill:4"), 1u); + EXPECT_EQ(device_->ingested, prompt); + EXPECT_EQ(device_->handoffs, 1u); + EXPECT_EQ(device_->handoff_length, 20u); + EXPECT_EQ(outcome.tokens, Expected(20, 5)); +} + +TEST_F(Qwen38GraphEngineTest, DenseChunksMayHaveAnOddSizeAndTheLastChunkEndsTheStreamExactly) { + device_->dense_prefill = true; + device_->preferred_chunk = 7; + auto engine = MakeEngine(); + const auto prompt = Prompt(15); + auto request = Begin(*engine, prompt, /*max_generated_tokens=*/3); + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(device_->CallCount("prefill:7"), 2u); + EXPECT_EQ(device_->CallCount("prefill:1"), 1u); + EXPECT_EQ(device_->ingested, prompt); + EXPECT_EQ(outcome.tokens, Expected(15, 3)); +} + +TEST_F(Qwen38GraphEngineTest, ADensePromptShorterThanTheChunkIsOneChunk) { + device_->dense_prefill = true; + device_->preferred_chunk = 64; + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(5), /*max_generated_tokens=*/2); + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(device_->CallCount("prefill:5"), 1u); + EXPECT_EQ(outcome.tokens, Expected(5, 2)); +} + +TEST_F(Qwen38GraphEngineTest, ADensePromptOfExactlyTheChunkSizeHasNoEmptyTailChunk) { + device_->dense_prefill = true; + device_->preferred_chunk = 8; + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(16), /*max_generated_tokens=*/2); + + RunToCompletion(*engine, request); + + EXPECT_EQ(device_->CallCount("prefill:8"), 2u); + EXPECT_EQ(device_->CallCount("prefill:0"), 0u); +} + +TEST_F(Qwen38GraphEngineTest, ContinuationAndRewindReIngestTheHistoryThroughTheDenseChunksAndHandOffAgain) { + device_->dense_prefill = true; + device_->preferred_chunk = 4; + auto engine = MakeEngine(); + const auto prompt = Prompt(5); + auto request = Begin(*engine, prompt, /*max_generated_tokens=*/3); + ASSERT_EQ(RunToCompletion(*engine, request).tokens, Expected(5, 3)); + ASSERT_EQ(device_->handoffs, 1u); + + TurnOptions options = Limit(2); + options.request = request; + const uint64_t turn = request->BeginTurn(std::array{30, 31}, options); + ASSERT_EQ(RunToCompletion(*engine, request).tokens, Expected(10, 2)); + EXPECT_EQ(device_->handoffs, 2u); + EXPECT_EQ(device_->handoff_length, 10u); // 5 prompt + 3 generated + 2 follow-up + + request->RewindToStartOfTurn(turn); + options = Limit(2); + options.request = request; + request->BeginTurn(std::array{40}, options); + const auto outcome = RunToCompletion(*engine, request); + std::vector history = prompt; + for (const int32_t token : Expected(5, 3)) history.push_back(token); + history.push_back(40); + EXPECT_EQ(device_->ingested, history); + EXPECT_EQ(device_->handoffs, 3u); + EXPECT_EQ(device_->handoff_length, history.size()); + EXPECT_EQ(outcome.tokens, Expected(history.size(), 2)); +} + +TEST_F(Qwen38GraphEngineTest, AFirstTokenThatEndsTheTurnSkipsTheHandOff) { + device_->dense_prefill = true; + device_->preferred_chunk = 4; + device_->overrides[5] = kEos; + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(5), /*max_generated_tokens=*/20); + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_TRUE(outcome.tokens.empty()); + EXPECT_EQ(device_->handoffs, 0u); + EXPECT_EQ(device_->CallCount("begin:0"), 1u); + EXPECT_EQ(device_->status_checks, 0u); +} + +TEST_F(Qwen38GraphEngineTest, CancelBetweenDenseChunksNeverHandsOffAndTheNextRequestStartsClean) { + device_->dense_prefill = true; + device_->preferred_chunk = 4; + auto engine = MakeEngine(); + auto first = engine->CreateRequest(); + TurnOptions first_options; + first_options.request = first; + const uint64_t turn = first->BeginTurn(Prompt(12), first_options); + std::array buffer; + ASSERT_EQ(engine->Run(buffer), 0u); // chunk 1 of 3 + EXPECT_TRUE(engine->CancelRequest(first, turn)); + while (engine->HasPendingRequests()) engine->Run(buffer); + EXPECT_EQ(device_->handoffs, 0u); + + const auto prompt = Prompt(9, 20); + auto second = Begin(*engine, prompt, /*max_generated_tokens=*/3); + const auto outcome = RunToCompletion(*engine, second); + EXPECT_EQ(device_->ingested, prompt); // the chunk of the cancelled request is gone + EXPECT_EQ(device_->handoffs, 1u); + EXPECT_EQ(outcome.tokens, Expected(9, 3)); +} + +// ---- end-of-decode device status --------------------------------------------------------------- + +TEST_F(Qwen38GraphEngineTest, TheDeviceStatusIsCheckedOnceWhenATurnEndsWithNothingInFlight) { + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(5), /*max_generated_tokens=*/7); + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(outcome.tokens, Expected(5, 7)); + EXPECT_EQ(device_->status_checks, 1u); +} + +TEST_F(Qwen38GraphEngineTest, DeviceErrorBitsAtTheEndOfATurnAreReportedAsAnExecutionFailure) { + device_->fail_status_check = true; + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(5), /*max_generated_tokens=*/3); + + std::array buffer; + bool saw_failure = false; + for (int i = 0; i < 100 && engine->HasPendingRequests() && !saw_failure; ++i) { + const size_t count = engine->Run(buffer); + for (size_t j = 0; j < count; ++j) { + if (buffer[j].flags & EngineEventFlagFailed) { + saw_failure = true; + EXPECT_EQ(buffer[j].error_code, EngineErrorCode::EngineExecutionFailure); + } + } + } + EXPECT_TRUE(saw_failure); +} + +TEST_F(Qwen38GraphEngineTest, DeviceErrorBitsAfterAStopStringTurnAreSurfacedBeforeReset) { + device_->overrides[5] = 5; + device_->overrides[6] = 6; + device_->fail_status_check = true; + auto engine = MakeEngine(); + auto request = engine->CreateRequest(); + TurnOptions options; + options.request = request; + options.max_generated_tokens = 20; + options.stop_strings = {"STOP"}; + const uint64_t turn = request->BeginTurn(Prompt(5), options); + + const auto outcome = RunToCompletion(*engine, request); + EXPECT_EQ(outcome.finish_reason, GenerationFinishReason::StopString); + + // Releasing the slot for rewind (or starting a new turn/request) drains the look-ahead cycle and checks decode status. + EXPECT_THROW(request->RewindToStartOfTurn(turn), std::runtime_error); + EXPECT_EQ(device_->status_checks, 1u); +} + +// The status check of a rewind runs outside a step, where RunStatic's fatal handler does not see its failure. The error bits +// are erased by the next Reset(), so the failure must outlive the exception the caller caught: the context is marked failed +// and the next step fails, instead of the context resuming as if the bits had never been read. +TEST_F(Qwen38GraphEngineTest, ADeviceErrorFoundWhileRewindingStaysFatalAfterTheCallerCaughtIt) { + device_->overrides[5] = 5; + device_->overrides[6] = 6; + device_->fail_status_check = true; + auto engine = MakeEngine(); + auto request = engine->CreateRequest(); + TurnOptions options; + options.request = request; + options.max_generated_tokens = 20; + options.stop_strings = {"STOP"}; + const uint64_t turn = request->BeginTurn(Prompt(5), options); + ASSERT_EQ(RunToCompletion(*engine, request).finish_reason, GenerationFinishReason::StopString); + + ASSERT_THROW(request->RewindToStartOfTurn(turn), std::runtime_error); + EXPECT_FALSE(device_->context_healthy); + + // What a caller that went on would find: bits that no longer exist and a request that is ready to run. + device_->fail_status_check = false; + auto next = Begin(*engine, Prompt(5, 25), /*max_generated_tokens=*/4); + EXPECT_EQ(RunUntilFailure(*engine), EngineErrorCode::EngineExecutionFailure); +} + +// The failure belongs to the device, not to the attempt that found it: a caller that retries the rewind must keep failing +// instead of consuming the error, or a later retry would succeed on a poisoned device and the next step would reset away +// the error bits. +TEST_F(Qwen38GraphEngineTest, RetryingAFailedRewindCannotConsumeTheDeviceFailure) { + device_->overrides[5] = 5; + device_->overrides[6] = 6; + device_->fail_status_check = true; + auto engine = MakeEngine(); + auto request = engine->CreateRequest(); + TurnOptions options; + options.request = request; + options.max_generated_tokens = 20; + options.stop_strings = {"STOP"}; + const uint64_t turn = request->BeginTurn(Prompt(5), options); + ASSERT_EQ(RunToCompletion(*engine, request).finish_reason, GenerationFinishReason::StopString); + + ASSERT_THROW(request->RewindToStartOfTurn(turn), std::runtime_error); + device_->fail_status_check = false; + EXPECT_THROW(request->RewindToStartOfTurn(turn), std::runtime_error); + EXPECT_THROW(request->RewindToStartOfTurn(turn), std::runtime_error); + EXPECT_FALSE(device_->context_healthy); + + auto next = Begin(*engine, Prompt(5, 25), /*max_generated_tokens=*/4); + EXPECT_EQ(RunUntilFailure(*engine), EngineErrorCode::EngineExecutionFailure); +} + +TEST_F(Qwen38GraphEngineTest, DeviceErrorBitsAfterACanceledTurnAreSurfacedOnTheNextRun) { + model_->config_->model.qwen38_graph.prefill.chunk = 64; + auto engine = MakeEngine(); + auto request = engine->CreateRequest(); + TurnOptions options; + options.request = request; + const uint64_t turn = request->BeginTurn(Prompt(5), options); + std::array buffer; + ASSERT_GE(engine->Run(buffer), 1u); + ASSERT_TRUE(device_->inflight()); + + device_->fail_status_check = true; + EXPECT_TRUE(engine->CancelRequest(request, turn)); + ASSERT_EQ(engine->Run(buffer), 1u); + EXPECT_EQ(buffer[0].finish_reason, GenerationFinishReason::Canceled); + + // When the next request claims the slot, Release() drains the canceled cycle and surfaces the device status error. + auto next = Begin(*engine, Prompt(5, 25), /*max_generated_tokens=*/4); + bool saw_failure = false; + for (int i = 0; i < 10 && engine->HasPendingRequests() && !saw_failure; ++i) { + const size_t count = engine->Run(buffer); + for (size_t j = 0; j < count; ++j) { + if (buffer[j].flags & EngineEventFlagFailed) { + saw_failure = true; + EXPECT_EQ(buffer[j].error_code, EngineErrorCode::EngineExecutionFailure); + } + } + } + EXPECT_TRUE(saw_failure); + EXPECT_EQ(device_->status_checks, 1u); +} + +// The wasted look-ahead cycle of a stop-string or canceled turn is never waited for, so its record stays unconsumed: the +// device reports its validation failure only once it was drained, and a Reset() would erase it. The executor must +// therefore drain and check before it lets anything reset the device. +TEST_F(Qwen38GraphEngineTest, AFailureOfTheLookAheadCycleVisibleOnlyAfterDrainSurfacesWhenAStopStringTurnIsRewound) { + device_->overrides[5] = 5; + device_->overrides[6] = 6; + device_->fail_cycle = 2; // cycle 1 completes the stop string; cycle 2 is the look-ahead nobody waits for + auto engine = MakeEngine(); + auto request = engine->CreateRequest(); + TurnOptions options; + options.request = request; + options.max_generated_tokens = 20; + options.stop_strings = {"STOP"}; + const uint64_t turn = request->BeginTurn(Prompt(5), options); + + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(outcome.finish_reason, GenerationFinishReason::StopString); // the turn itself is fine + EXPECT_EQ(outcome.failure_flags, 0u); + EXPECT_TRUE(device_->inflight()); // the record of the look-ahead cycle is unconsumed... + EXPECT_TRUE(device_->running()); // ...and nothing has synchronized the stream yet + EXPECT_EQ(device_->status_checks, 0u); + + EXPECT_THROW(request->RewindToStartOfTurn(turn), std::runtime_error); + EXPECT_TRUE(device_->inflight()); // Drain() completed the cycle without consuming its record + EXPECT_FALSE(device_->running()); + EXPECT_EQ(device_->order, (std::vector{"reset", "drain", "status"})); // the only Reset was the first turn's +} + +TEST_F(Qwen38GraphEngineTest, AFailureOfTheLookAheadCycleVisibleOnlyAfterDrainFailsTheNextRequestInsteadOfBeingResetAway) { + device_->overrides[5] = 5; + device_->overrides[6] = 6; + device_->fail_cycle = 2; + auto engine = MakeEngine(); + auto request = engine->CreateRequest(); + TurnOptions options; + options.request = request; + options.max_generated_tokens = 20; + options.stop_strings = {"STOP"}; + request->BeginTurn(Prompt(5), options); + ASSERT_EQ(RunToCompletion(*engine, request).finish_reason, GenerationFinishReason::StopString); + ASSERT_TRUE(device_->inflight()); + + // The next request claims the slot: Release() drains the unconsumed cycle and checks the device before any Reset(). + device_->overrides.clear(); + auto next = Begin(*engine, Prompt(5, 25), /*max_generated_tokens=*/4); + const auto failure = RunUntilFailure(*engine); + + ASSERT_TRUE(failure.has_value()) << "the error bits of the drained cycle were never reported"; + EXPECT_EQ(*failure, EngineErrorCode::EngineExecutionFailure); + EXPECT_EQ(device_->order, (std::vector{"reset", "drain", "status"})); // no Reset() erased the bits first + EXPECT_EQ(device_->CallCount("expect:5"), 1u); // only the first request's announcement +} + +TEST_F(Qwen38GraphEngineTest, AFailureOfTheFirstCycleVisibleOnlyAfterDrainSurfacesAfterACancellation) { + model_->config_->model.qwen38_graph.prefill.chunk = 64; + device_->fail_cycle = 1; // the cycle BeginDecode enqueued + auto engine = MakeEngine(); + auto request = engine->CreateRequest(); + TurnOptions options; + options.request = request; + const uint64_t turn = request->BeginTurn(Prompt(5), options); + std::array buffer; + ASSERT_GE(engine->Run(buffer), 1u); // prefill and the first token; cycle 1 is running + ASSERT_TRUE(device_->running()); + + EXPECT_TRUE(engine->CancelRequest(request, turn)); + ASSERT_EQ(engine->Run(buffer), 1u); + EXPECT_EQ(buffer[0].finish_reason, GenerationFinishReason::Canceled); + EXPECT_EQ(buffer[0].flags & EngineEventFlagFailed, 0u); // cancellation does not wait for the device + EXPECT_TRUE(device_->running()); + + auto next = Begin(*engine, Prompt(5, 25), /*max_generated_tokens=*/4); + const auto failure = RunUntilFailure(*engine); + + ASSERT_TRUE(failure.has_value()) << "the error bits of the canceled cycle were never reported"; + EXPECT_EQ(*failure, EngineErrorCode::EngineExecutionFailure); + EXPECT_EQ(device_->order, (std::vector{"reset", "drain", "status"})); +} + +TEST_F(Qwen38GraphEngineTest, ADrainedLookAheadCycleWithoutErrorBitsIsCheckedAndTheNextRequestRuns) { + device_->overrides[5] = 5; + device_->overrides[6] = 6; + auto engine = MakeEngine(); + auto request = engine->CreateRequest(); + TurnOptions options; + options.request = request; + options.max_generated_tokens = 20; + options.stop_strings = {"STOP"}; + request->BeginTurn(Prompt(5), options); + ASSERT_EQ(RunToCompletion(*engine, request).finish_reason, GenerationFinishReason::StopString); + + device_->overrides.clear(); + auto next = Begin(*engine, Prompt(5, 25), /*max_generated_tokens=*/3); + const auto outcome = RunToCompletion(*engine, next); + + EXPECT_EQ(outcome.failure_flags, 0u); + EXPECT_EQ(outcome.tokens, Expected(5, 3)); + EXPECT_EQ(device_->order, (std::vector{"reset", "drain", "status", "reset", "status"})); +} + +TEST_F(Qwen38GraphEngineTest, CapabilitiesAreHonest) { + auto engine = MakeEngine(); + const auto capabilities = engine->GetCapabilities(); + EXPECT_EQ(capabilities.configured_max_batch_size, 1u); + EXPECT_EQ(capabilities.max_request_length, static_cast(device_->Capacity() - Qwen38DecodeDevice::kReservedTail)); + EXPECT_EQ(engine->MaxDraftTokensPerStep(), 0u); + EXPECT_FALSE(engine->PrefixCacheStats().has_value()); + + RequestOptions too_long; + too_long.max_session_tokens = capabilities.max_request_length + 1; + try { + engine->CreateRequest(too_long); + FAIL() << "a session longer than the KV capacity must be rejected"; + } catch (const std::runtime_error& error) { + EXPECT_NE(std::string{error.what()}.find("max_request_length"), std::string::npos) << error.what(); + } + RequestOptions at_limit; + at_limit.max_session_tokens = capabilities.max_request_length; + EXPECT_NO_THROW(engine->CreateRequest(at_limit)); +} + +TEST_F(Qwen38GraphEngineTest, MaxRequestLengthIsTheSmallerOfTheModelContextLengthAndTheKvCapacity) { + const int kv_limit = device_->Capacity() - Qwen38DecodeDevice::kReservedTail; + + model_->config_->model.context_length = kv_limit - 10; + EXPECT_EQ(MakeEngine()->GetCapabilities().max_request_length, static_cast(kv_limit - 10)); + + model_->config_->model.context_length = kv_limit + 10; + EXPECT_EQ(MakeEngine()->GetCapabilities().max_request_length, static_cast(kv_limit)); +} + +TEST_F(Qwen38GraphEngineTest, UnsupportedTurnOptionsAreRejectedBeforeTheRequestChanges) { + auto engine = MakeEngine(); + auto request = engine->CreateRequest(); + const auto prompt = Prompt(4); + const auto expect_rejected = [&](const char* label, TurnOptions options, const char* expected_fragment) { + options.request = request; + try { + request->BeginTurn(prompt, options); + ADD_FAILURE() << label << " must be rejected"; + } catch (const std::runtime_error& error) { + EXPECT_NE(std::string{error.what()}.find(expected_fragment), std::string::npos) << label << ": " << error.what(); + } + EXPECT_TRUE(request->IsAwaitingFirstTurn()) << label << " mutated the request"; + }; + + TurnOptions sampled; + sampled.do_sample = true; + sampled.top_k = 5; + sampled.temperature = 0.7f; + expect_rejected("sampling", sampled, "greedily"); + + TurnOptions penalty; + penalty.repetition_penalty = 1.2f; + expect_rejected("repetition_penalty", penalty, "repetition_penalty"); + + TurnOptions ngram; + ngram.no_repeat_ngram_size = 3; + expect_rejected("no_repeat_ngram_size", ngram, "no_repeat_ngram_size"); + + TurnOptions minimum; + minimum.min_generated_tokens = 4; + expect_rejected("min_generated_tokens", minimum, "min_generated_tokens"); + + TurnOptions seeded; + seeded.seed = 7; + expect_rejected("seed", seeded, "seed"); + + TurnOptions guided; + guided.guidance_type = "lark_grammar"; + guided.guidance_data = "start: \"a\""; + expect_rejected("guidance", guided, "guidance"); + + // The same request still starts a plain greedy turn. + TurnOptions greedy = Limit(2); + greedy.request = request; + EXPECT_NO_THROW(request->BeginTurn(prompt, greedy)); + EXPECT_EQ(RunToCompletion(*engine, request).tokens, Expected(4, 2)); +} + +TEST_F(Qwen38GraphEngineTest, ADeviceFailureIsFatalAndReportedThroughTheEngineErrorTaxonomy) { + model_->config_->model.qwen38_graph.prefill.chunk = 64; // the first step ingests the whole prompt + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(5), /*max_generated_tokens=*/10); + std::array buffer; + ASSERT_GE(engine->Run(buffer), 1u); // prefill and the first token succeed + device_->fail_wait = true; + + const size_t count = engine->Run(buffer); + ASSERT_GE(count, 1u); + bool saw_failure = false; + for (size_t i = 0; i < count; ++i) { + if (buffer[i].flags & EngineEventFlagFailed) { + saw_failure = true; + EXPECT_EQ(buffer[i].error_code, EngineErrorCode::EngineExecutionFailure); + } + } + EXPECT_TRUE(saw_failure); + EXPECT_THROW(engine->Run(buffer), std::exception); // the Engine is no longer healthy +} + +TEST_F(Qwen38GraphEngineTest, DestroyingTheEngineDrainsTheDevice) { + model_->config_->model.qwen38_graph.prefill.chunk = 64; // the first step ingests the whole prompt + { + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(5), /*max_generated_tokens=*/10); + std::array buffer; + ASSERT_GE(engine->Run(buffer), 1u); + ASSERT_TRUE(device_->running()); + } + EXPECT_FALSE(device_->running()); // the cycle was waited for (its record stays unconsumed) + EXPECT_GE(device_->CallCount("drain"), 1u); +} + +TEST_F(Qwen38GraphEngineTest, ANewEngineOverAUsedDeviceStartsFromACleanDevice) { + model_->config_->model.qwen38_graph.prefill.chunk = 64; + { + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(5), /*max_generated_tokens=*/3); + ASSERT_EQ(RunToCompletion(*engine, request).tokens, Expected(5, 3)); + } + const size_t resets = device_->CallCount("reset"); + + // The device kept the first Engine's decode state; the next Engine must not prefill on top of it. + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(6, 20), /*max_generated_tokens=*/3); + const auto outcome = RunToCompletion(*engine, request); + + EXPECT_EQ(device_->CallCount("reset"), resets + 1); + EXPECT_EQ(device_->ingested, Prompt(6, 20)); + EXPECT_EQ(outcome.tokens, Expected(6, 3)); +} + +TEST_F(Qwen38GraphEngineTest, MalformedCycleWidthFailsBeforeAnyCandidateCanMutateTheRequest) { + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(3), 10); + std::array events; + ASSERT_GT(engine->Run(events), 0u); // prompt prediction + const auto retained = request->TokensCpu(); + const std::vector before(retained.begin(), retained.end()); + device_->malformed_cycle_count = true; + engine->Run(events); + const auto after = request->TokensCpu(); + EXPECT_EQ(std::vector(after.begin(), after.end()), before); + EXPECT_FALSE(device_->context_healthy); +} + +TEST_F(Qwen38GraphEngineTest, FailureHealthPersistsWhenANewEngineTriesToLeaseTheSameContext) { + auto engine = MakeEngine(); + auto request = Begin(*engine, Prompt(3), 10); + std::array events; + engine->Run(events); + device_->fail_wait = true; + engine->Run(events); + EXPECT_FALSE(device_->context_healthy); + engine.reset(); + device_->fail_wait = false; + EXPECT_THROW(MakeEngine(), std::runtime_error); +} + +// ---- cache manager contract ------------------------------------------------------------------ + +TEST_F(Qwen38GraphEngineTest, TheCacheManagerHoldsExactlyOneRow) { + auto engine = MakeEngine(); + auto session = std::make_shared(device_, Qwen38GraphSession::Options{}); + Qwen38GraphCacheManager cache{model_, session}; + auto a = engine->CreateRequest(); + auto b = engine->CreateRequest(); + a->BeginTurn(Prompt(3)); + b->BeginTurn(Prompt(3)); + + EXPECT_FALSE(cache.CanAllocate({})); + EXPECT_FALSE(cache.CanAllocate({a, b})); + ASSERT_TRUE(cache.CanAllocate({a})); + cache.Allocate({a}); + EXPECT_EQ(cache.ResidentRequestCount(), 1u); + EXPECT_TRUE(cache.IsResident(a)); + EXPECT_FALSE(cache.IsResident(b)); + EXPECT_FALSE(cache.CanAllocate({b})) << "an active resident row cannot be replaced"; + EXPECT_THROW(cache.Allocate({b}), std::runtime_error); + + std::vector> wrong{b}; + EXPECT_THROW(cache.Deallocate(wrong), std::runtime_error); + std::vector> right{a}; + cache.Deallocate(right); + EXPECT_EQ(cache.ResidentRequestCount(), 0u); + EXPECT_TRUE(cache.CanAllocate({b})); + EXPECT_EQ(cache.MaxBatchSize(), 1u); + EXPECT_EQ(cache.MaxStaticRequestLength(), static_cast(device_->Capacity() - Qwen38DecodeDevice::kReservedTail)); + EXPECT_FALSE(cache.SupportsDynamicBatching()); +} + +TEST_F(Qwen38GraphEngineTest, DynamicBatchingConfigurationIsRejected) { + model_->config_->engine.dynamic_batching = Config::Engine::DynamicBatching{}; + EXPECT_THROW(CreateQwen38GraphDependencies(model_, device_), std::runtime_error); +} + +TEST_F(Qwen38GraphEngineTest, InvalidPromptTokenIdsAreRejectedAtTurnAdmissionWithoutPoisoningTheDevice) { + auto engine = MakeEngine(); + auto request = engine->CreateRequest(); + auto options = Limit(3); + options.request = request; + const int32_t vocab_size = model_->config_->model.vocab_size; + + for (const int32_t bad_id : {int32_t{-1}, vocab_size}) { + const std::array bad_prompt = {12, bad_id, 14}; + EXPECT_THROW(request->BeginTurn(bad_prompt, options), std::invalid_argument) << bad_id; + EXPECT_TRUE(request->IsAwaitingFirstTurn()); + EXPECT_EQ(request->CurrentSequenceLength(), 0); + EXPECT_TRUE(device_->ContextIsHealthy()); + EXPECT_TRUE(device_->calls.empty()); + } + + // The caller can immediately retry with a valid prompt on the same Request and Engine. + request->BeginTurn(Prompt(5), options); + const auto outcome = RunToCompletion(*engine, request); + EXPECT_EQ(outcome.finish_reason, GenerationFinishReason::TurnLimit); + EXPECT_EQ(outcome.tokens, Expected(5, 3)); +} + +// Host cost of one Engine step (Run) with a device that answers instantly: everything the Engine adds around a cycle. Run +// manually: engine_unit_tests --gtest_also_run_disabled_tests --gtest_filter='*RunOverhead*'. +TEST_F(Qwen38GraphEngineTest, DISABLED_RunOverheadPerCycle) { + device_ = std::make_shared(/*capacity=*/33028); + model_->config_->model.context_length = 33028; + model_->config_->model.qwen38_graph.prefill.chunk = 2048; + auto engine = MakeEngine(); + const size_t generated = 20000; + auto request = Begin(*engine, Prompt(16), generated, /*max_session_tokens=*/33024); + std::array buffer; + size_t runs = 0, tokens = 0; + const auto start = std::chrono::steady_clock::now(); + while (engine->HasPendingRequests()) { + const size_t count = engine->Run(buffer); + ++runs; + tokens += count; + } + const double seconds = std::chrono::duration(std::chrono::steady_clock::now() - start).count(); + std::cout << "Engine::Run: " << runs << " runs, " << tokens << " events in " << seconds * 1e3 << " ms = " + << seconds / static_cast(runs) * 1e6 << " us per run" << std::endl; + EXPECT_EQ(tokens, generated); +} + +} // namespace + +// Shared host conformance runs independently of graph names and device bank conventions. +TEST(TokenProducingSessionTest, ProgressSettlesWithoutATokenAndDuplicateAcknowledgementFails) { + TokenProducingSession host{1}; + int owner; + const auto plan = host.Plan(0, 0, &owner, 1, 4, false); + BackendRowResult progress; + progress.key = plan.key; + progress.prefill = PrefillProgress{4, false}; + EXPECT_TRUE(host.Collected(progress)); + TokenCommitAck ack; + ack.key = plan.key; + EXPECT_NO_THROW(host.Acknowledge(ack)); + EXPECT_THROW(host.Acknowledge(ack), std::logic_error); +} + +TEST(TokenProducingSessionTest, RetiringCompletionDiscardsBeforeTheSlotCanBeReused) { + TokenProducingSession host{1}; + int old_owner, new_owner; + const auto old = host.Plan(0, 0, &old_owner, 1, 0, false); + host.Retire(0); + EXPECT_THROW(host.Plan(0, 0, &new_owner, 1, 0, false), std::logic_error); + BackendRowResult late; + late.key = old.key; + late.selected_tokens.emplace(); + late.selected_tokens->candidate_count = 1; + late.selected_tokens->candidates[0] = 12; + EXPECT_FALSE(host.Collected(late)); + const auto current = host.Plan(0, 0, &new_owner, 1, 0, false); + EXPECT_FALSE(old.key == current.key); + EXPECT_THROW(host.Collected(late), std::logic_error); + TokenCommitAck stale; + stale.key = old.key; + EXPECT_THROW(host.Acknowledge(stale), std::logic_error); +} + +TEST(TokenProducingSessionTest, NewLeasesCannotAliasLocalSlotAndSubmissionCounters) { + int owner; + TokenProducingSession first{1}; + TokenProducingSession second{1}; + const auto old = first.Plan(0, 0, &owner, 1, 0, false); + const auto current = second.Plan(0, 0, &owner, 1, 0, false); + EXPECT_NE(old.key.context_id, current.key.context_id); + EXPECT_EQ(old.key.submission_id, current.key.submission_id); + BackendRowResult stale; + stale.key = old.key; + EXPECT_THROW(second.Collected(stale), std::logic_error); +} + +TEST(TokenProducingSessionTest, ExactCandidateCountsAndDuplicateCompletionsAreValidated) { + int owner; + TokenProducingSession host{1}; + const auto plan = host.Plan(0, 0, &owner, 1, 0, false); + BackendRowResult result; + result.key = plan.key; + result.selected_tokens.emplace(); + result.selected_tokens->candidate_count = 2; + result.selected_tokens->candidates[0] = 12; + result.selected_tokens->candidates[1] = 1; + EXPECT_TRUE(host.Collected(result)); + EXPECT_THROW(host.Collected(result), std::logic_error); + TokenCommitAck ack; + ack.key = plan.key; + ack.candidates_consumed = 3; + EXPECT_THROW(host.Acknowledge(ack), std::logic_error); + ack.candidates_consumed = 2; + ack.tokens_appended = 1; + ack.tokens_staged_for_events = 1; + ack.turn_finished = true; + ack.finish_reason = GenerationFinishReason::EosToken; + ack.disposition = ContinuationDisposition::Invalidate; + EXPECT_NO_THROW(host.Acknowledge(ack)); +} + +} // namespace Generators::test diff --git a/test/cpp/engine/qwen38_graph_test_doubles.h b/test/cpp/engine/qwen38_graph_test_doubles.h new file mode 100644 index 0000000000..b2df580fcb --- /dev/null +++ b/test/cpp/engine/qwen38_graph_test_doubles.h @@ -0,0 +1,245 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "span.h" + +namespace Generators::test { + +// A scripted stand-in for the CUDA-graph decode loop. It reproduces the loop's observable contract without a model: +// +// - the "model" is a function of the position: the greedy token at sequence position p is TokenAt(p), so a Request that +// runs it sees the sequence prompt ++ TokenAt(H), TokenAt(H + 1), ... (H = prompt length) regardless of how the steps +// were chunked, and re-ingesting a longer history (a continued turn) keeps producing TokenAt of the next position; +// - the drafter is right on the cycles named by `accept_pattern` (cycled), so cycles commit 1 or 2 tokens; SetDrafts() +// switches to the fold's multi-draft cycles (2..3 drafts per verify, any accepted prefix 0..drafts per cycle); +// - `dense_prefill` models the fast prefill: chunks of any size advance the ledger at bank 0 (no S=2 pairs), the hand-off +// to the decode state happens in BeginDecode(run_cycle = true) and only once all chunks arrived, `preferred_chunk` is +// the chunk size the device asks the executor for; +// - the pending token / in-flight cycle / ledger bookkeeping follows Qwen38GraphEngine, and misuse (a wait without a +// cycle in flight, an advance of something that was not waited for, prefill after decoding started, ...) throws, so the +// tests also prove the executor drives the device in a legal order; +// - a cycle runs asynchronously: it is enqueued by BeginDecode/Advance and completes once something synchronizes the +// stream (WaitCycle, Drain, CheckDecodeStatus, Reset). Like the real device, Drain() completes it without consuming its +// record (the cycle stays "in flight" until WaitCycle), and CheckDecodeStatus() synchronizes itself and reports the +// error bits accumulated by every completed cycle, consumed or not, until Reset() clears them. `fail_cycle` makes a cycle +// accumulate such bits when it completes, so a failure only becomes visible after the stream was drained. +class ScriptedQwen38Device final : public Qwen38DecodeDevice { + public: + explicit ScriptedQwen38Device(int capacity = 68, bool mtp = true) : capacity_{capacity}, mtp_{mtp}, ledger_{capacity} {} + + // ---- script ---- + std::vector accept_pattern{true, false}; + std::map overrides; // position -> token (for example the EOS id) + int32_t TokenAt(size_t position) const { + const auto found = overrides.find(position); + if (found != overrides.end()) return found->second; + return static_cast(10 + (position * 7) % 40); // never the EOS id 1 or a stop-string token (< 10) + } + + bool dense_prefill{}; + size_t preferred_chunk{}; + + // Called at the end of ExpectPrefillTokens(), that is between the publication of a turn's prompt length and its first + // PrefillChunk(): a test that drives several Engines blocks here to interleave their owner threads the worst way. + std::function on_expect_prefill; + + // Every verify proposes `drafts` tokens and accepts the next entry of `accepted_counts` (cycled, each 0..drafts) of them + // in order, so the cycle commits that many tokens + 1 (the correction, or the bonus when all were accepted). The record + // then reports them like the production device does, through `accepted_count`. Call it before the first Reset(). + void SetDrafts(int drafts, std::vector accepted_counts) { + if (drafts < 1 || drafts > 3) throw std::invalid_argument("a verify proposes 1..3 drafts"); + for (const int count : accepted_counts) { + if (count < 0 || count > drafts) throw std::invalid_argument("a verify accepts 0..drafts drafts"); + } + drafts_ = drafts; + accepted_counts_ = std::move(accepted_counts); + ledger_.SetNumBanks(drafts + 2); // the bank ring of the fold: one more than the rows of a verify + } + + // ---- observation ---- + std::vector calls; // "reset", "expect:", "prefill:", "begin:", "wait", "advance:", "drain" + std::vector ingested; // tokens received since the last reset, in order + std::vector prefill_totals; // for every PrefillChunk, the prompt length the device was told to expect (0 = none) + size_t cycles_enqueued{}; + size_t handoffs{}; // dense prefill: hand-offs performed (one per BeginDecode with a cycle) + size_t handoff_length{}; // prompt length at the last hand-off + size_t status_checks{}; // CheckDecodeStatus calls (not in `calls`, which tests compare by suffix) + std::vector order; // "drain", "status" and "reset" in the order the device saw them + bool inflight() const { return inflight_; } // a cycle was enqueued and its record was not consumed by WaitCycle + bool running() const { return running_; } // ... and it has not completed on the device yet (nothing synchronized) + size_t ledger_length() const { return ledger_.Length(); } + int ledger_bank() const { return ledger_.Bank(); } + size_t CallCount(const std::string& name) const { return std::count(calls.begin(), calls.end(), name); } + + // ---- fault injection ---- + bool fail_wait{}; + bool malformed_cycle_count{}; + bool context_healthy{true}; + bool fail_status_check{}; // every CheckDecodeStatus() of a started decode reports error bits + size_t fail_cycle{}; // 1-based cycle of a decode that accumulates error bits when it completes (0 = none) + + // ---- Qwen38DecodeDevice ---- + void MarkContextFailed() noexcept override { context_healthy = false; } + bool ContextIsHealthy() const noexcept override { return context_healthy; } + int Capacity() const override { return capacity_; } + bool MtpEnabled() const override { return mtp_; } + + void Reset() override { + calls.push_back("reset"); + order.push_back("reset"); + Complete(); // Reset() waits for every enqueued cycle... + ledger_.Reset(); + ingested.clear(); + inflight_ = false; + decoding_ = false; + pending_ = -1; + cycle_ = 0; + expected_prefill_ = 0; + error_bits_ = 0; // ...then clears the device state, error bits included + } + + void Drain() noexcept override { + calls.push_back("drain"); + order.push_back("drain"); + Complete(); // the cycle in flight finishes, but its record stays unconsumed: inflight_ is not cleared + } + + void ExpectPrefillTokens(size_t total_tokens) override { + calls.push_back("expect:" + std::to_string(total_tokens)); + if (decoding_ || ledger_.Length() != 0) throw std::logic_error("the prompt length was announced after its ingestion started"); + expected_prefill_ = total_tokens; + if (on_expect_prefill) on_expect_prefill(total_tokens); + } + + void PrefillChunk(std::span tokens) override { + calls.push_back("prefill:" + std::to_string(tokens.size())); + if (decoding_) throw std::logic_error("prefill after decoding started"); + if (tokens.empty()) throw std::logic_error("empty prefill chunk"); + prefill_totals.push_back(expected_prefill_); + ledger_.RequireRoom(tokens.size(), "the prompt", 3); + if (dense_prefill) { + ledger_.AdvancePrompt(tokens.size()); + } else { + for (size_t i = 0; i < tokens.size();) { + const size_t step = i + 1 < tokens.size() ? 2 : 1; + ledger_.Advance(step); + i += step; + } + } + ingested.insert(ingested.end(), tokens.begin(), tokens.end()); + pending_ = TokenAt(ledger_.Length()); + } + + int32_t PendingToken() const override { return pending_; } + size_t PrefillChunkTokens() const override { return preferred_chunk; } + + void BeginDecode(bool run_cycle) override { + calls.push_back(std::string("begin:") + (run_cycle ? "1" : "0")); + if (decoding_) throw std::logic_error("decoding already started"); + if (ledger_.Length() == 0) throw std::logic_error("BeginDecode before the prompt was ingested"); + decoding_ = true; + expected_prefill_ = 0; + if (run_cycle && dense_prefill) { + if (ingested.size() != ledger_.Length()) throw std::logic_error("hand-off before every chunk was ingested"); + ++handoffs; + handoff_length = ledger_.Length(); + } + if (run_cycle) Enqueue(); + } + + Qwen38CycleRecord WaitCycle() override { + calls.push_back("wait"); + if (fail_wait) throw std::runtime_error("injected device failure"); + if (!inflight_) throw std::logic_error("no cycle in flight"); + Complete(); + inflight_ = false; + if (fail_cycle != 0 && cycle_ == fail_cycle) throw std::runtime_error("injected device validation failure"); + Qwen38CycleRecord record; + if (malformed_cycle_count) { + record.accepted_count = 100; + record.proposed_count = 100; + return record; + } + const size_t length = ledger_.Length(); + size_t accepted_drafts = 0; + if (accepted_counts_.empty()) { + record.accepted = mtp_ && accept_pattern[(cycle_ - 1) % accept_pattern.size()]; // the legacy single-draft flag + accepted_drafts = record.accepted ? 1 : 0; + } else { + accepted_drafts = mtp_ ? static_cast(accepted_counts_[(cycle_ - 1) % accepted_counts_.size()]) : 0; + record.accepted = accepted_drafts > 0; + record.accepted_count = static_cast(accepted_drafts); + } + record.proposed_count = mtp_ ? drafts_ : 1; + // Rows after the first mismatch are meaningless; the cycle commits the accepted drafts and one token more. + const size_t step = accepted_drafts + 1; + const std::array predictions{&record.pred0, &record.pred1, &record.pred2, &record.pred3}; + for (size_t i = 0; i < predictions.size(); ++i) { + *predictions[i] = i < step ? TokenAt(length + 1 + i) : -7; + } + if (record.TokenCount() != step) throw std::logic_error("the scripted record disagrees with its own token count"); + record.new_pending = *predictions[step - 1]; + record.new_length = static_cast(length + step); + return record; + } + + void Advance(const Qwen38CycleRecord& record, bool want_more) override { + calls.push_back(std::string("advance:") + (want_more ? "1" : "0")); + if (inflight_) throw std::logic_error("advance with a cycle still in flight"); + ledger_.Advance(record.TokenCount()); + ledger_.CheckDeviceLength(record.new_length); + pending_ = static_cast(record.new_pending); + if (want_more) Enqueue(); + } + + void CheckDecodeStatus() override { + ++status_checks; + order.push_back("status"); + if (!decoding_) return; // like the real device: nothing to check before the decode started + Complete(); // it synchronizes the stream itself, also while a cycle is in flight or its record unconsumed + if (fail_status_check || error_bits_ != 0) throw std::runtime_error("injected device error bits"); + } + + private: + void Enqueue() { + ledger_.RequireRoom(static_cast(drafts_) + 1, "the next verify step"); + inflight_ = true; + running_ = true; + ++cycle_; + ++cycles_enqueued; + } + + // The stream was synchronized: the cycle that was running completes and leaves its validation bits behind. + void Complete() { + if (!running_) return; + running_ = false; + if (fail_cycle != 0 && cycle_ == fail_cycle) error_bits_ |= 2; + } + + int capacity_; + bool mtp_; + int drafts_{1}; + std::vector accepted_counts_; + Qwen38DecodeLedger ledger_; + bool inflight_{}; + bool running_{}; + bool decoding_{}; + int32_t pending_{-1}; + size_t cycle_{}; + int64_t error_bits_{}; + size_t expected_prefill_{}; // ExpectPrefillTokens: device-local, dropped by Reset() and BeginDecode() +}; + +} // namespace Generators::test diff --git a/test/cpp/engine/scheduler_contract_tests.cpp b/test/cpp/engine/scheduler_contract_tests.cpp index babfb0cae7..27d5842494 100644 --- a/test/cpp/engine/scheduler_contract_tests.cpp +++ b/test/cpp/engine/scheduler_contract_tests.cpp @@ -262,6 +262,24 @@ TEST_F(SchedulerContractTest, DynamicResidentQueuedRequestIsNotReadmitted) { EXPECT_EQ(cache->AllocatedCount(), 1u); } +// Engine::BeginTurn enrolls a request again when a turn restarts after a rewind or a cancellation. The request never left +// the pool, so the second enrollment must not queue it twice. +TEST_F(SchedulerContractTest, DynamicEnrollingAnEnrolledRequestAgainPlansItOnce) { + auto cache = std::make_shared(model_, /*capacity=*/8); + DynamicBatchScheduler scheduler(model_, cache); + + auto request = Assigned(10); + scheduler.AddRequest(request); + scheduler.AddRequest(request); + StepPlan plan; + + const auto result = scheduler.PlanStep(plan); + + ASSERT_TRUE(result.executable); + ASSERT_EQ(plan.requests.size(), 1u); + EXPECT_EQ(plan.requests[0].request, request); +} + // Removing a request from the dynamic scheduler deallocates its cache resources immediately. TEST_F(SchedulerContractTest, DynamicRemoveReleasesCacheResources) { auto cache = std::make_shared(model_, /*capacity=*/8);