diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp index 18ed536c27f0..6a05bc7bb8df 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp @@ -849,8 +849,29 @@ std::vector BlockRadixTree::matchTokenPath( return results; } +namespace +{ +// Drop `backoff` tokens off the tail of a match, dropping whole blocks while the backoff +// outruns them. Only the tail shrinks, so pruneMatch's "leading entries are full blocks" +// invariant still holds. +void backOffMatch(std::vector& matched, int backoff) +{ + while (backoff > 0 && !matched.empty()) + { + auto& last = matched.back(); + if (last.numMatchedTokens > backoff) + { + last.numMatchedTokens -= backoff; + return; + } + backoff -= last.numMatchedTokens; + matched.pop_back(); + } +} +} // namespace + std::vector BlockRadixTree::pruneMatch( - std::vector matched, std::optional ssmLcId) const + std::vector matched, std::optional ssmLcId, int alignment) const { // All blocks except the last must be fully matched (mirrors Python: matched[:-1]). TLLM_CHECK_DEBUG(matched.size() <= 1 @@ -864,6 +885,17 @@ std::vector BlockRadixTree::pruneMatch( // shortens the match, so the loop terminates. while (!matched.empty()) { + // Align before validating pages: moving backwards can make stale SWA + // pages active. Snapshot and coverage trims retry through this boundary. + if (alignment > 1) + { + int const remainder = numMatchedTokens(matched, mTokensPerBlock) % alignment; + backOffMatch(matched, remainder); + if (matched.empty()) + { + break; + } + } // Check SSM snapshot availability first: truncating to the last reusable SSM // snapshot changes the matched length that all the attention checks use. if (ssmLcId.has_value()) @@ -889,6 +921,10 @@ std::vector BlockRadixTree::pruneMatch( break; } matched.back().numMatchedTokens = ssmMatchLen; + if (alignment > 1 && numMatchedTokens(matched, mTokensPerBlock) % alignment != 0) + { + continue; + } } // Only pages that are active at this candidate endpoint constrain attention @@ -946,30 +982,10 @@ std::vector BlockRadixTree::pruneMatch( return matched; } -namespace -{ -// Drop `backoff` tokens off the tail of a match, dropping whole blocks while the backoff -// outruns them. Only the tail shrinks, so pruneMatch's "leading entries are full blocks" -// invariant still holds. -void backOffMatch(std::vector& matched, int backoff) -{ - while (backoff > 0 && !matched.empty()) - { - auto& last = matched.back(); - if (last.numMatchedTokens > backoff) - { - last.numMatchedTokens -= backoff; - return; - } - backoff -= last.numMatchedTokens; - matched.pop_back(); - } -} -} // namespace - -BlockRadixTree::ReuseMatch BlockRadixTree::match( - ReuseScope const& reuseScope, TokenSpan tokens, bool knownNoDigest, bool enablePartialMatch, int backoff) const +BlockRadixTree::ReuseMatch BlockRadixTree::match(ReuseScope const& reuseScope, TokenSpan tokens, bool knownNoDigest, + bool enablePartialMatch, int backoff, int alignment) const { + TLLM_CHECK_DEBUG(alignment > 0 && mTokensPerBlock % alignment == 0); auto rawMatched = matchTokenPath(reuseScope, tokens, knownNoDigest, enablePartialMatch); // Content-divergence depth is measured before page or recurrent-snapshot pruning. int const numReusableTokensBeforePruning = numMatchedTokens(rawMatched, mTokensPerBlock); @@ -986,9 +1002,10 @@ BlockRadixTree::ReuseMatch BlockRadixTree::match( std::optional numReusableTokensBeforeHybridPruning; if (ssmLcId.has_value()) { - numReusableTokensBeforeHybridPruning = numMatchedTokens(pruneMatch(rawMatched, std::nullopt), mTokensPerBlock); + numReusableTokensBeforeHybridPruning + = numMatchedTokens(pruneMatch(rawMatched, std::nullopt, alignment), mTokensPerBlock); } - auto matched = pruneMatch(std::move(rawMatched), ssmLcId); + auto matched = pruneMatch(std::move(rawMatched), ssmLcId, alignment); ReuseMatch result{}; result.numTokens = numMatchedTokens(matched, mTokensPerBlock); result.numLookupTokens = static_cast(tokens.size()); diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h index b6fe8d1dfd4b..1e6a95aa3ee9 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h @@ -396,8 +396,10 @@ class BlockRadixTree // Takes a non-owning TokenSpan so a zero-copy int32 token buffer can be matched without // allocating/copying (the hot path). Callers holding a std::vector pass toSpan(vec). // backoff: tokens trimmed off the tail of the match (see KVCacheManagerConfig::reuseMatchBackoff). + // alignment: the final match is trimmed down to a multiple of this after backoff and pruning + // (see KVCacheManagerConfig::reuseMatchAlignment). ReuseMatch match(ReuseScope const& reuseScope, TokenSpan tokens, bool knownNoDigest = false, - bool enablePartialMatch = false, int backoff = 0) const; + bool enablePartialMatch = false, int backoff = 0, int alignment = 1) const; // Detach all cached blocks. ~Block() releases pages when the last owner drops a block. void clear(); @@ -439,7 +441,8 @@ class BlockRadixTree // Shorten `matched` to the prefix that is actually reusable. Passing // std::nullopt for `ssmLcId` skips the recurrent-snapshot constraint and // yields the attention-only prefix (used for numReusableTokensBeforeHybridPruning). - std::vector pruneMatch(std::vector matched, std::optional ssmLcId) const; + std::vector pruneMatch( + std::vector matched, std::optional ssmLcId, int alignment = 1) const; // Erase any pending empty root blocks from mRoots. // Const-qualified: deferred cleanup is not a logical mutation. diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.cpp index 8a6e20701d73..51a624594625 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.cpp @@ -40,6 +40,11 @@ void DiskCacheTierConfig::validate() const void KVCacheManagerConfig::validate() const { + if (reuseMatchAlignment <= 0 || tokensPerBlock % reuseMatchAlignment != 0) + { + throw AssertionError( + "KVCacheManagerConfig: reuse_match_alignment must be a positive divisor of tokens_per_block"); + } if (swaScratchReuse.has_value()) { swaScratchReuse->validate(); diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h index 19600e0fd52a..1340d94b4efb 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h @@ -420,6 +420,10 @@ struct KVCacheManagerConfig // Applied inside the match so a single tree walk yields the usable depth. int reuseMatchBackoff = 0; + // Required final prefix alignment after backoff and coverage pruning. + // Whole-block connectors use tokensPerBlock; 1 preserves token-granular reuse. + int reuseMatchAlignment = 1; + // Constraint-based memory partitioning. std::vector constraints; // batches that must always be supportable std::optional typicalStep; // typical step for initial ratio computation diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp index 1ec31e1bc0ae..5a693ea7a9f8 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp @@ -239,7 +239,8 @@ BlockRadixTree::ReuseMatch KvCacheManager::matchReuse( ReuseScope const& reuseScope, TokenSpan inputTokens, bool knownNoDigest) const { KVCM2_REJECT_IF_POISONED(); - return mRadixTree->match(reuseScope, inputTokens, knownNoDigest, enablePartialMatch(), mConfig.reuseMatchBackoff); + return mRadixTree->match(reuseScope, inputTokens, knownNoDigest, enablePartialMatch(), mConfig.reuseMatchBackoff, + mConfig.reuseMatchAlignment); } int KvCacheManager::probeReuse(ReuseScope reuseScope, TokenSpan inputTokens, bool knownNoDigest) const diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index c503e0250648..35e65e41b1c0 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -1753,7 +1753,7 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) std::optional> initialPoolRatio, std::optional swaScratchReuse, bool commitMinSnapshot, bool enableStats, bool textOnly, bool enablePartialCommit, kv::PoolRebalanceConfig poolRebalance, - kv::Priority coldTierOffloadMinPriority) + kv::Priority coldTierOffloadMinPriority, int reuseMatchAlignment) { new (cfg) kv::KVCacheManagerConfig(); cfg->tokensPerBlock = tokensPerBlock; @@ -1769,6 +1769,7 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) cfg->maxUtilForResume = maxUtilForResume; cfg->enablePartialReuse = enablePartialReuse; cfg->reuseMatchBackoff = reuseMatchBackoff; + cfg->reuseMatchAlignment = reuseMatchAlignment; cfg->typicalStep = std::move(typicalStep); cfg->constraints = std::move(constraints); cfg->initialPoolRatio = std::move(initialPoolRatio); @@ -1790,13 +1791,15 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) nb::arg("swa_scratch_reuse").none() = std::nullopt, nb::arg("commit_min_snapshot") = false, nb::arg("enable_stats") = true, nb::arg("text_only") = false, nb::arg("enable_partial_commit") = true, nb::arg("pool_rebalance") = kv::PoolRebalanceConfig{}, - nb::arg("cold_tier_offload_min_priority") = tensorrt_llm::batch_manager::kDefaultColdTierOffloadMinPriority) + nb::arg("cold_tier_offload_min_priority") = tensorrt_llm::batch_manager::kDefaultColdTierOffloadMinPriority, + nb::arg("reuse_match_alignment") = 1) .def_rw("tokens_per_block", &kv::KVCacheManagerConfig::tokensPerBlock) .def_rw("cache_tiers", &kv::KVCacheManagerConfig::cacheTiers) .def_rw("layers", &kv::KVCacheManagerConfig::layers) .def_rw("max_util_for_resume", &kv::KVCacheManagerConfig::maxUtilForResume) .def_rw("enable_partial_reuse", &kv::KVCacheManagerConfig::enablePartialReuse) .def_rw("reuse_match_backoff", &kv::KVCacheManagerConfig::reuseMatchBackoff) + .def_rw("reuse_match_alignment", &kv::KVCacheManagerConfig::reuseMatchAlignment) .def_rw("enable_partial_commit", &kv::KVCacheManagerConfig::enablePartialCommit) .def_rw("cold_tier_offload_min_priority", &kv::KVCacheManagerConfig::coldTierOffloadMinPriority) .def_rw("typical_step", &kv::KVCacheManagerConfig::typicalStep) diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp index 99756c07befc..4478d4655fbf 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp @@ -16,6 +16,7 @@ */ #include "kvCacheManagerV2TestUtils.h" +#include "kv_cache_manager_v2/exceptions.h" #include "kv_cache_manager_v2/lifeCycleRegistry.h" #include "kv_cache_manager_v2/storage/config.h" @@ -46,6 +47,23 @@ TEST(KvCacheManagerV2ConfigTest, ExistingBufferConfigurationDefaultsToDense) EXPECT_NO_THROW(makeConfig().validate()); } +TEST(KvCacheManagerV2ConfigTest, ReuseMatchAlignmentMustDivideTokensPerBlock) +{ + auto config = makeConfig(); + ASSERT_EQ(config.tokensPerBlock, 4); + EXPECT_EQ(config.reuseMatchAlignment, 1); + for (int const alignment : {1, 2, 4}) + { + config.reuseMatchAlignment = alignment; + EXPECT_NO_THROW(config.validate()) << "alignment " << alignment; + } + for (int const alignment : {0, -4, 3, 8}) + { + config.reuseMatchAlignment = alignment; + EXPECT_THROW(config.validate(), AssertionError) << "alignment " << alignment; + } +} + TEST(KvCacheManagerV2ConfigTest, SparseBuffersRequireHostMemoryAtLevelOne) { auto config = makeConfig(); diff --git a/docs/source/features/disagg-serving.md b/docs/source/features/disagg-serving.md index 23752b6d26e8..dd6e7ded9360 100644 --- a/docs/source/features/disagg-serving.md +++ b/docs/source/features/disagg-serving.md @@ -102,6 +102,8 @@ The disaggregated server generates this ID itself as a **snowflake** — a self- Global disaggregated IDs occupy the range `[1 << 40, 2**63)`; worker-local and warm-up request IDs occupy the disjoint range `[0, 1 << 40)`, so the two never collide. If a client supplies its own positive `disagg_request_id`, that value is used verbatim and must be globally unique; when unset, the server mints a snowflake ID as above. +If a request to a worker fails, the disaggregated server retries a context request under a new ID: up to `max_retries` times, or up to 5 times for a connection reset or server disconnect. A generation request is re-sent only when it cannot have reached the generation server (connection refused, or reset while the request was being written), and it keeps its ID so it still pairs with the context KV cache. Any later failure of a generation request fails the request, because the generation server may already hold the one-shot KV-cache handoff. `TRTLLM_DISAGG_NO_RETRY=1` disables all retries. A chat or completion stream that fails before its first chunk is answered with an HTTP error status; a Responses stream ends with an `error` event. + ## Usage ### Dynamo diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index 53918b857687..b011bf9a17dd 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -67,7 +67,7 @@ on the leader process; with attention DP, they run on each owning rank. * **`request_reset(self, request: LlmRequest)`** * **Description**: Optional, with a no-op default. The request's allocation has been released and it will run again from the start, unlike `request_finished`, which ends it. Any per-request state held in page indices now describes pages another request may own, so it must be dropped or rebuilt here. - * **When it fires**: on `KVCacheManagerV2`, from `free_resources`, which is the path a rollback or a failed admission takes. Direct preemption calls `request_finished` instead, so a connector keying state by request id sees exactly one of the two. + * **When it fires**: on `KVCacheManagerV2`, from `free_resources`, which is the path a rollback or a failed admission takes. A restart after a failed prefix load (see the worker's `reports_failed_loads` below) takes it too. Direct preemption calls `request_finished` instead, so a connector keying state by request id sees exactly one of the two. * **`update_state_after_alloc_by_layer_group(self, request: LlmRequest, block_ids_by_layer_group: list[list[int]])`** * **`request_finished_by_layer_group(self, request: LlmRequest, cache_block_ids_by_layer_group: list[list[int]]) -> bool`** @@ -82,6 +82,12 @@ on the leader process; with attention DP, they run on each owning rank. A set is complete when both of its methods are defined. An existing flat connector needs no change on the models it already covers: with a single layer group the base per-layer-group implementation folds back to the flat call. Hybrid / linear-attention models are not enabled for the connector yet; the per-layer-group set is the shape their cache will need. See [Running under VSWA](#running-under-vswa). +* **`supports_speculative_prefill: bool`** (class attribute) + * **Description**: Optional, `False` by default; only an explicit `True` counts, read from the class before the scheduler is constructed. Declares a scheduler whose content identity (its store keys) covers the `draft_prompt_lookahead` prompt tokens past each block's end, which never offers or saves a block whose lookahead tokens fall outside the prompt, and which never saves a block the current pass has not computed. Together with the worker's declaration it lets the connector run next to one-model speculative decoding on a disaggregated context server; see [Speculative decoding on a disaggregated context server](#speculative-decoding-on-a-disaggregated-context-server). + +* **`supports_admission_query: bool`** (class attribute) + * **Description**: Optional, `False` by default; only an explicit `True` counts, read from the class so every rank agrees. Declares that `get_num_new_matched_tokens` may be asked while the scheduler is still admitting the request; see [Admission-time queries](#admission-time-queries). + ##### Running under VSWA KVCacheManagerV2 supports multiple layer groups with a connector implementing the grouped APIs below. KVCacheManagerV1 registers a single primary pool and cannot describe VSWA to a connector. Single-window SWA remains allowed with either manager. @@ -198,13 +204,67 @@ following configuration limits at bring-up. | Configuration | Mechanism | |---|---| -| Speculative decoding | Rejected draft tokens shrink a request's page list, and the freed slot goes to whichever request allocates next. The connector is never told the tail block moved. | +| Speculative decoding, outside the [speculative-prefill-only mode](#speculative-decoding-on-a-disaggregated-context-server) | Rejected draft tokens shrink a request's page list, and the freed slot goes to whichever request allocates next. The connector is never told the tail block moved. | | A capacity scheduler policy other than `GUARANTEED_NO_EVICT` with KV cache manager V1 | V1 does not reset the connector block delta when a request is destroyed and replayed. V2 resets that state and permits replay; an active connector load retains its allocation until completion. | | A host or disk cache tier | Tier eviction reassigns the GPU slot. See [KV cache tiers are GPU-only under a connector](#kv-cache-tiers-are-gpu-only-under-a-connector). | The exact set the runtime refuses depends on your cache configuration; the bring-up error is -authoritative. Speculative decoding is unsupported with a connector wherever it is not refused — -the same page-list shrink happens there. +authoritative. Outside the speculative-prefill-only mode below, speculative decoding is unsupported +with a connector wherever it is not refused — the same page-list shrink happens there. + +##### Speculative decoding on a disaggregated context server + +A context-only request finishes after prefill and never rejects a draft token, so on a +disaggregated context server the page list the connector is told about only grows. The runtime +therefore lets one-model speculative decoding run next to a connector there, in what it calls the +speculative-prefill-only mode. The mode is on when all of the following hold: + +* The KV cache manager is `KVCacheManagerV2`. An explicit `use_kv_cache_manager_v2=False` keeps + V1's behaviour unchanged, without this mode; `"auto"` resolving to V1 is rejected once the manager + is built. +* `cache_transceiver_config.backend` is set, so the server takes part in disaggregated serving. +* Speculative decoding is one-model MTP (MTP-Eagle or vanilla MTP) or EAGLE3. Other modes, such as + NGram or the suffix automaton, keep the refusal above. +* The draft layers share the target KV cache. A separate draft KV cache is not registered with the + connector and is refused. +* The connector's scheduler and worker classes both declare `supports_speculative_prefill = True`. + +What the declaration promises: + +* The scheduler's content identity covers the `draft_prompt_lookahead(speculative_config)` prompt + tokens past each block's end: one for MTP-Eagle and EAGLE3, `max_draft_len` for vanilla MTP. A + one-model draft writes the KV at a position from the tokens after it, so a served block's + draft-layer KV then belongs to this prompt. A block whose lookahead tokens fall outside the + prompt is neither offered nor saved. +* Only blocks the current pass has computed are saved. One-model speculation allocates + `max_draft_len - 1` extra KV slots past every context chunk, so a block can own a page before it + is computed. +* Loads and saves move whole pages of every layer group in the registered `KvCacheLayout`, which + includes the draft layers. Per-layer hooks are installed on `DecoderLayer` modules only, and + whether a model's draft layers are such modules depends on the model, so a worker that saves per + layer must make sure it covers them. + +At runtime: + +* A request that is not context-only gets an error response at admission, and only that request + fails. That includes `/health_generate` sent straight to the context server. +* The connector manager refuses such a request again at the lookup, the prefix reservation and + the per-iteration scheduler output, for anything that bypasses admission. There the error is not + confined to the request. +* A disaggregated generation server should use a connector role that transfers no KV + (`capacity_only`). When its classes declare the mode, such a role is exempt from both refusals + and may itself run one-model MTP or EAGLE3. A generation server whose connector transfers KV and + which meets the conditions above refuses every generation-only request. +* A one-model draft can back a local reuse match off by its prompt lookahead, which leaves a match + of whole blocks short of a block boundary. A worker that stores whole blocks can ask + `KVCacheManagerV2` to align local matches down to a block (`requires_block_aligned_reuse`) or to + drop the backoff (`drops_reuse_match_backoff`), both described with the worker interface below. + Dropping it can lower draft acceptance but never changes the output. +* The mode composes with [attention data parallelism](#attention-data-parallelism) when both + classes also declare `supports_attention_dp`. + +The [`mooncake-store`](#mooncake-store) connector declares the mode; see +[its notes](#attention-data-parallelism-and-speculative-prefill). ##### `RequestData` fields that may not be populated @@ -302,17 +362,133 @@ worker's transfer stream. Inference does not wait for the retirement control ite ##### Legacy final-batch queries -`get_num_new_matched_tokens` is called at most once per KV allocation after batch selection. -Its offer can reduce computation for that request, but does not free token budget for another -request in the same iteration. With chunked prefill, allocation growth can limit the served prefix; -`RequestData.computed_position` reflects the load interval that the runtime honors. +`get_num_new_matched_tokens` is called at most once per KV allocation, after batch selection +unless [admission-time queries](#admission-time-queries) are enabled. Its offer can reduce +computation for that request, but does not free token budget for another request in the same +iteration unless admission-time queries are enabled. With chunked prefill, allocation growth can +limit the served prefix; `RequestData.computed_position` reflects the load interval that the runtime +honors. Destroying an allocation clears its connector state, so replay queries again. Specify `enable_block_reuse=True` alongside the connector; see [Block reuse alongside the connector](#block-reuse-alongside-the-connector). +Restarting a request whose prefix load failed (`reports_failed_loads`) is available only on +this path. **Deployment note.** Under a connector, a workload that was token-bound becomes KV-bound: the connector removes forward-pass tokens but its prefix still occupies GPU pages. Lowering `max_num_tokens` to hand memory back to the KV pool is usually the right adjustment, the opposite of the guidance for a connector-free deployment. +##### Admission-time queries + +A connector without source reservations can be asked while the KV cache manager V2 scheduler +admits a first context chunk, instead of after the batch is final. Its scheduler declares +`supports_admission_query = True`, which promises: + +- `get_num_new_matched_tokens` has no effect beyond its answer until `update_state_after_alloc`; +- asking again for the same allocation replaces the earlier answer; +- a query may be followed by `request_reset` and another query, with no `update_state_after_alloc` + in between. + +A connector that starts a transfer inside `get_num_new_matched_tokens` must not declare it. A +request asked at admission can still be dropped before `prepare_resources`: by the disaggregated +benchmark gate, by batch waiting for generation requests, by the total-KV cap, by the multimodal +scheduler wrapper, or when the batch cannot be queued. It is asked again on a later iteration. +The [`mooncake-store`](#asynchronous-loads) connector declares it. + +Enable it per server with `kv_connector_config.admission_query`: + +```yaml +kv_connector_config: + connector: + admission_query: + max_uncharged_parks: 2 # default 0 + max_offer_deferrals: 16 # default 16 +disable_overlap_scheduler: true +``` + +The scheduler asks each scheduled first chunk once its chunk is sized and its pages are allocated, +and `prepare_resources` uses that answer instead of asking again. An offer that reaches past the +chunk moves the chunk window, so the pages it needs are booked right away. When the KV cache cannot +cover them, the request gives its chunk and pages back and waits for a later iteration, up to +`max_offer_deferrals` iterations; after that, or right after it preempted a started request for +pages, it is admitted and the offer is capped to the whole blocks the chunk's pages hold. With +`max_uncharged_parks > 0`, up to that many first chunks of at least two blocks per iteration, which +the connector loads asynchronously, are admitted without charging their tokens to `max_num_tokens`: +the executor parks them before the forward pass. They still take a batch slot and keep their pages. +A request admitted without a charge that does not park stops the executor with a connector failure, +which requires a process restart. + +| Setting | Effect | +|---|---| +| unset | Final-batch queries only. | +| `admission_query: {}` | Admission-time queries and page booking. | +| `admission_query: {max_uncharged_parks: N}` | Also admits up to `N` parking chunks per iteration without a token charge. | + +Comparing the first two isolates page booking, and the last two isolate the token charge. + +Admission-time queries require KV cache manager V2 and `disable_overlap_scheduler=True`, which the +executor checks at startup, and chunked prefill, without which they have no effect. They are ignored, with a +warning, under attention data parallelism, for a capacity-only connector role, for connectors that +reserve their prefix during admission, and for connector schedulers that do not declare +`supports_admission_query`. Under `enable_swa_scratch_reuse` the pages a shifting offer needs are +not booked at admission; `prepare_resources` grows them as on the final-batch path. The settings are +configuration rather than environment variables because they change admission behavior and can +differ between the context and generation servers of one disaggregated deployment. + +For debugging, `TLLM_CTX_PARK_STATS=1` appends a flat `, ctx_park = key=value ...` suffix to every +`print_iter_log` line of the non-overlap executor loop, covering the iteration the line reports: + +| Field | Meaning | +|---|---| +| `parked`, `parked_tok` | Context requests the connector parked for an asynchronous load, and their chunk tokens. | +| `unc` | Of those, the requests admitted without a token charge. | +| `fwd_tok` | Context chunk tokens left for the forward pass. | +| `bskip` | First chunks the scheduler passed over because the token budget was spent. | +| `_ms=wall/cpu` | Wall and executor-thread CPU time of the `sched`, `prep`, `meta`, `start`, `fwd`, `upd` and `term` loop phases. | +| `offer_*` | KV cache manager V2 connector offers: queries (`admission_used` of them answered at admission), offers and asynchronous offers, offered and honoured tokens, how each was placed (`in_chunk`, `shifted`, `capped`, `not_honoured`), admission-time grows (`pregrown`, `pregrow_failed`), and host time in the query, the placement and the page report (`query_ms`, `reserve_ms`, `report_ms`). | +| `gil=max/p99/n` | With `TLLM_GIL_PROBE_US=`, how late a probe thread sleeping that many microseconds woke up, in ms, a proxy for GIL contention. | + +##### Completion headroom on disaggregated servers + +A disaggregated server with KV cache manager V2 can reserve logical GPU-page headroom for each +admitted request to finish; it does not allocate those pages early. The reservation is opt-in: a +server qualifies only when no tier sits below GPU and either its connector's worker declares +[`requires_completion_reservation`](#2-worker-interface-kvcacheconnectorworker) or its +`kv_cache_config.enable_disagg_completion_reservation` (prototype, `False` by default) is set. A +connector that transfers KV leaves no such tier, because the host tier is dropped automatically and +the others are refused at bring-up (see +[KV cache tiers](#kv-cache-tiers-are-gpu-only-under-a-connector)), so its declaration is enough. A +server with neither keeps its existing admission, `host_cache_size=0` included, whether it has no +connector, a capacity-only one, or one that does not declare it. The Mooncake store connector +declares it. The reservation covers full-attention models only, without beam search or pipeline +parallelism. Windowed, recurrent and mixed-precision caches keep their existing admission policy. + +A disaggregated generation server without a connector opts in through the configuration. Without a +connector the host tier is provisioned automatically, so set `host_cache_size: 0` and configure no +disk tier; otherwise the cache manager logs a warning and ignores the setting. Such a server +reserves decode headroom only; the connector-prefill reservation below needs a connector. + +```yaml +kv_cache_config: + host_cache_size: 0 + enable_disagg_completion_reservation: true +``` + +- **Decode.** Disaggregated generation admission reserves pages through each request's output + limit, including speculative padding. Incoming transfers may overlap decoding, but cannot + consume the pages reserved for admitted requests to finish. +- **Connector prefill.** For connectors without source reservations, the scheduler reserves + headroom through the end of each admitted prompt before allocating its first chunk; chunk + allocation stays incremental. A new request waits for headroom instead of preempting an + admitted one, so other requests' small chunks cannot consume the pages needed to restore its + matching prefix after the final-batch query. + +A reservation is released when the request finishes or is cancelled, when an unconsumed first +chunk returns to the queue, when a generation-init admission is reverted before its transfer, and +on shutdown. The check counts free and evictable GPU pages, so droppable reusable pages do not +block admission. A server whose pool was clamped below the configured `max_seq_len` may defer a +request whose output limit reaches `max_seq_len` indefinitely; size the pool above `max_seq_len`. +This GPU admission budget is separate from any store capacity behind the connector. + #### 2. Worker Interface (`KvCacheConnectorWorker`) These methods run on all workers (GPU processes) and interact with the actual GPU data. @@ -345,6 +521,41 @@ These methods run on all workers (GPU processes) and interact with the actual GP * **`get_finished_prefix_loads(self) -> list[int]`** * **Description**: Reports locally completed reservation identities for loads dispatched through `SchedulerOutput.prefix_loads`, including synchronous loads. Completion means this worker has finished all reads and writes and established the required CUDA stream visibility. This method must not wait for other workers. The runtime collects reports asynchronously and distributes retirement decisions before scheduling; parked async requests resume and shared resources are released only after every worker has reported. * **Compatibility**: Legacy request-ID load and save completions continue through `get_finished`. + +* **`reports_failed_loads: bool`** (class attribute) + * **Description**: Optional, `False` by default; only an explicit `True` counts. Declares a worker whose prefix load can come up short — a page offered at lookup is gone at load time, for example evicted from a full pool — and which reports the affected requests instead of raising. The runtime has already counted those positions as computed, so it restarts each named request from scratch: the request leaves the pass, its allocation is freed and rewound, and the scheduler readmits it and asks `get_num_new_matched_tokens` again. Other requests are unaffected. + * **Requirements**: `take_failed_load_requests` and `drop_bound_saves` must be implemented, the KV cache manager must be `KVCacheManagerV2`, and prefix reservations must be off (see [Legacy final-batch queries](#legacy-final-batch-queries)); each is checked at startup. A restart is owner-local, so it applies on a single-rank server or an attention-DP owner. On a tensor-parallel rank, whose peers hold shards of the same pages, the worker must still raise; a report from such a rank fails the executor. A `capacity_only` worker's declaration is ignored. + +* **`take_failed_load_requests(self) -> set[int]`** + * **Description**: Polled right after every `start_load_kv` when `reports_failed_loads` is set. Returns and forgets the `RequestData.request_id`s whose load in that call came up short. Each must belong to a context request in the metadata that call loaded. + +* **`take_failed_async_load_requests(self) -> set[int]`** + * **Description**: Polled right after every `get_finished` when `reports_failed_loads` is set. Returns and forgets parked requests whose asynchronous load came up short. Each must have been returned as finished loading by the `get_finished` call this poll follows: after that the scheduler may admit the request over its pages, so a later report comes too late to restart it. An ID for a request that was cancelled while parked names nothing to restart and is ignored. + * **Contract for both**: every write into a named request's pages must be complete or abandoned when its ID is returned, because the pages are freed at once and may go to another request. An ID outside the loaded batch, for an asynchronous load not yet reported finished, or for one an earlier `get_finished` already reported finished, is a contract violation and fails the executor. + +* **`drop_bound_saves(self, request_ids) -> int`** + * **Description**: Removes these requests' saves from the metadata bound for the current pass, so a restarted request does not publish pages it never computed. Called before `wait_for_save`. Must accept IDs with nothing bound and be idempotent; returns the number of entries removed. + +* **`shutdown(self)`** + * **Description**: Optional, with a no-op default. Releases whatever the worker holds — store handles, registered buffers, background threads. Called once, after the executor's worker thread has joined, so no transfer can start afterwards; implementations must be idempotent. Without it a worker's resources live as long as the process, so engines built back to back in one session accumulate them. + +* **`capacity_only(self) -> bool`** (property) + * **Description**: Optional, `False` by default. Declares a worker that contributes resources to an external store but transfers no KV of its own — for example a rank that lends host memory to a shared pool other engines read and write. Must be the same on every rank; the runtime gathers it at startup and rejects a deployment whose ranks disagree. Such a worker is never asked for a lookup, a save, a prefix reservation, or a completion, and the runtime keeps its own connector bookkeeping — page-index gathers, scheduler output, the per-layer hooks — off the forward path entirely. + +* **`supports_speculative_prefill: bool`** (class attribute) + * **Description**: Optional, `False` by default; only an explicit `True` counts, read from the class before the worker is constructed. Declares a worker whose every load and save moves whole pages of every layer group in the registered `KvCacheLayout`, including the one-model draft layers that share the target pool. A worker that saves per layer must make sure its hooks cover the draft layers, since hooks are installed on `DecoderLayer` modules only. Together with the scheduler's declaration it lets the connector run next to one-model speculative decoding on a disaggregated context server; see [Speculative decoding on a disaggregated context server](#speculative-decoding-on-a-disaggregated-context-server). + +* **`requires_block_aligned_reuse: bool`** (class attribute) + * **Description**: Optional, `False` by default; only an explicit `True` counts. Declares a connector that stores and offers KV by whole blocks starting at the local prefix match, so that match must end on a block boundary. It matters only when a one-model speculative-decoding draft (for example MTP or one-model EAGLE3) makes `KVCacheManagerV2` back every local reuse match off by the draft's prompt lookahead, which leaves a match of whole blocks a few tokens short of a boundary, and partial reuse is off. The manager then aligns every local match down to a multiple of `tokens_per_block`, trading the trimmed block for the connector serving it along with the rest of the prefix. Partial reuse keeps matches token-granular, so a connector that sets this should also turn `KvCacheConfig.enable_partial_reuse` off. + +* **`drops_reuse_match_backoff: bool`** (class attribute) + * **Description**: Optional, `False` by default; only an explicit `True` counts. Declares a connector that prefers whole local matches to backed-off ones: the target `KVCacheManagerV2` drops the one-model draft's reuse backoff, so a local match of whole blocks is used whole and needs no alignment. The draft KV at the tail of a reused block may then belong to a different next token, which can lower draft acceptance but never changes the output. A draft pool never drops its own backoff. + +* **`requires_completion_reservation: bool`** (class attribute) + * **Description**: Optional, `False` by default; only an explicit `True` counts. Turns on the completion reservation of a disaggregated server with no tier below GPU (see [Completion headroom on disaggregated servers](#completion-headroom-on-disaggregated-servers)): `KVCacheManagerV2` then admits a decode or a connector first chunk only while the pages every admitted request needs to finish still fit. Declare it when the connector restores a prefix only after the final batch is chosen, so other requests' chunks must not take the pages that restore needs. Read from the worker, which every rank holds, once when the cache manager is built after the connector, so a worker may set it per instance in `__init__`. Must be the same on every rank. A `capacity_only` worker's declaration is ignored. + +The two reuse declarations above are read once, when the cache manager is built after the connector, so a worker may set them per instance in `__init__`. A `capacity_only` worker's declarations are ignored. They must be the same on every rank: unlike `capacity_only`, they are not gathered at startup, and ranks that disagree would claim different local prefixes for the same request. + ## Attention data parallelism Set `enable_attention_dp=True` and @@ -437,14 +648,26 @@ them for the backend. Non-ADP behavior has no new transfer deadline. Workers must keep asynchronous start/poll callbacks nonblocking; a callback that itself hangs remains subject to the executor's hang detector. +A failed prefix load that a connector reports through `reports_failed_loads` is +not a connector failure. Only that request restarts, on its owner, without a +status-vote field or a deadline of its own. Exceptions, polling errors, deadline +misses and misreported request IDs remain fail-stop through the gate above. If +the restarted request was the owner's only work and no padding dummy fits, the +fleet skips one iteration, as it does whenever an owner cannot pad. + ADP requires `KVCacheManagerV2`; an explicit V1 selection is rejected before connector construction, and an automatic selection that resolves to V1 is rejected after the manager is built. V1 connector behavior remains unchanged. PP=1, CP=1 and beam width 1 are required. V2 layer-group layout and callbacks are used directly. A backend must implement both grouped callbacks for multiple layer groups; the existing layout validation remains in place. Internal host -or disk tiers, speculative decoding and Mamba/hybrid state remain unsupported. -This change adds no Mooncake backend, deployment configuration or capacity tuning. +or disk tiers and Mamba/hybrid state remain unsupported. Speculative decoding +is supported only in the +[speculative-prefill-only mode](#speculative-decoding-on-a-disaggregated-context-server), +for which both classes declare `supports_speculative_prefill` as well. +The [`mooncake-store`](#mooncake-store) connector declares `supports_attention_dp` and +`supports_speculative_prefill`; see +[its notes on attention data parallelism and speculative prefill](#attention-data-parallelism-and-speculative-prefill). Select a connector through `connector_module`, `connector_scheduler_class`, and `connector_worker_class` as usual. Both classes must opt in; enabling ADP @@ -470,11 +693,20 @@ Add `examples/llm-api` to `PYTHONPATH`. The example supports one full-attention layer group and no chunked prefill; it demonstrates shared storage, not production transfer performance. -* **`shutdown(self)`** - * **Description**: Optional, with a no-op default. Releases whatever the worker holds — store handles, registered buffers, background threads. Called once, after the executor's worker thread has joined, so no transfer can start afterwards; implementations must be idempotent. Without it a worker's resources live as long as the process, so engines built back to back in one session accumulate them. +**Routing by uncached prefill.** With a shared pool, a later request of a +conversation restores most of its prompt on whichever attention-DP rank it +lands, so prompt length says little about the prefill it brings. Setting -* **`capacity_only(self) -> bool`** (property) - * **Description**: Optional, `False` by default. Declares a worker that contributes resources to an external store but transfers no KV of its own — for example a rank that lends host memory to a shared pool other engines read and write. Must be the same on every rank; the runtime gathers it at startup and rejects a deployment whose ranks disagree. Such a worker is never asked for a lookup, a save, a prefix reservation, or a completion, and the runtime keeps its own connector bookkeeping — page-index gathers, scheduler output, the per-layer hooks — off the forward path entirely. +```yaml +attention_dp_config: + enable_work_aware_routing: true +``` + +balances ranks by each request's prompt growth since the previous request of +the same conversation (the whole prompt for a first turn). It requires a +connector that moves KV; without one, a warning is logged and load-based +routing is used. It cannot be combined with `enable_kv_cache_aware_routing` or +`kv_cache_routing_conversation_affinity`. ## Example Implementation @@ -575,6 +807,10 @@ and `org-b/model` share a directory name while meaning different weights. `segme rank against a per-node memory limit, so a node's demand is `ranks_on_node × segment_size` and is checked against available memory at startup. +A server publishes a block only once the forward pass that computes its last token has been queued +for saving, so pages it allocates ahead of computation, such as those of a context chunk trimmed to +the token budget, never reach the pool. + ### Roles Capacity and traffic are separate: every role lends the memory its config asks for, and the role @@ -595,6 +831,39 @@ keeps its own host and disk cache tiers, including the host tier [`MAX_UTILIZATION` suspends pages into](#kv-cache-tiers-are-gpu-only-under-a-connector), and pays no per-request or per-layer connector cost. +### Pool eviction + +A full pool evicts on the master, and two settings shape what it evicts. Both apply only to a role +that writes the pool (`both` or `producer`), both need `mooncake-transfer-engine-cuda13` 0.3.13 or +newer, and both are off by default. + +```yaml +kv_connector_config: + connector: mooncake-store + mooncake_store: + pool: file:///shared/pool.json + model_key: qwen2-1.5b-instruct + soft_pin_ttl_ms: 86400000 # 24 h, the master's default ceiling + group_blocks: true +``` + +* `soft_pin_ttl_ms` soft-pins every page the server writes for that many milliseconds. Mooncake + 0.3.13's first eviction pass makes up a shortfall from the shard-ordered rest of the pool when + lookups re-lease part of its oldest frontier mid-eviction, which leaves holes in prefixes still in + use: a prefix is replayed only up to its first missing block, so every stored block behind a hole + is recomputed. Soft-pinned pages are evicted only by the second, age-ordered pass. The TTL must not + exceed the master's `--max_kv_soft_pin_ttl` (24 h by default), or the master refuses every put, + which the server logs once; and the master must allow evicting soft-pinned objects, which is + Mooncake's default. +* `group_blocks` puts every page a hit on a block needs (all its layer groups and tensor-parallel + shards) into one eviction group, `//w/`, so the master + evicts a block whole and never keeps a shard that holds space but can no longer hit. It matters + most under tensor parallelism, where a block has one page per rank. + +Both are per-server settings like `transfer_batch_size`, rendered into the client config and read +back from an inherited `MOONCAKE_CONFIG_PATH`. Servers sharing a pool should agree on +`soft_pin_ttl_ms`, since the master evicts unpinned pages first. + ### Settings this connector overrides For a role that transfers KV, the pool *is* the deployment's offload tier, so the server's own @@ -611,17 +880,116 @@ Both are settled in the `LLM` constructor, before the ranks are spawned and befo reads the arguments, so `trtllm-serve`, `trtllm-bench` and a direct `LLM(...)` all see the same settings. A `capacity` server keeps both as configured. +### Attention data parallelism and speculative prefill + +Both classes declare `supports_attention_dp`, so `enable_attention_dp=True` runs a scheduler +adapter and a worker on every owner, each deciding and moving pages for its own requests only. + +* A stored page is one attention shard's KV. Under tensor parallelism each rank keys its own shard + and a prefix is a hit only once every shard of it is in the pool. Every attention-DP owner holds + all heads, so all owners key their pages as shard 0 of 1 (`w1r0` in the key): a prefix one owner + saved is a hit for every other owner, and for a single-rank server with the same `model_key`, + `tokens_per_block` and page size. Tensor-parallel keys are unchanged. +* Under attention DP a save in flight counts against `kv_connector_config.transfer_timeout_sec` + (see [Attention data parallelism](#attention-data-parallelism)); size it for the save backlog a + congested fabric can build up. +* A page that is gone between lookup and load restarts only the affected request, on its owner, + as on a single-rank server; see [Restrictions](#restrictions). +* A `role: capacity` server may run with attention DP; it transfers nothing either way. + +Both classes also declare `supports_speculative_prefill`, so a disaggregated context server can run +one-model MTP or EAGLE3 next to this connector in the +[speculative-prefill-only mode](#speculative-decoding-on-a-disaggregated-context-server), with or +without attention DP. + +* A block's key also covers the prompt tokens the draft reads past the block: one for MTP-Eagle and + EAGLE3, `max_draft_len` for vanilla MTP. Pages a speculative server writes are therefore never + shared with a server without a draft or with a different lookahead, and the keys of a server + without a draft are unchanged. A block whose lookahead runs past the prompt is neither offered + nor saved. +* Only the blocks a pass has computed are saved, although one-model speculation allocates pages + past every context chunk. +* Loads and saves move whole pages of every layer group, which includes the draft layers sharing + the target KV cache. Any other drafter, and draft layers in a separate KV cache, are refused at + startup, since their KV would be missing or stale after a restored prefix. A `role: capacity` + server is exempt, since it addresses no page. +* A disaggregated generation server should run `role: capacity`, which may itself run one-model MTP + or EAGLE3, with or without attention DP. A generation server whose role transfers KV refuses + every generation-only request in this mode. + +`mooncake_store.local_reuse_backoff` sets how a server that reads the pool (`role: both` or +`consumer`) treats its local prefix match under one-model speculative decoding. The draft backs +every local match off by its prompt lookahead, which leaves a match of whole blocks short of the +block boundary the pool serves from. + +* `off`, the default, has the worker declare + [`drops_reuse_match_backoff`](#2-worker-interface-kvcacheconnectorworker): local hits are used + whole. The draft KV at the tail of a reused block may belong to a different next token, which can + lower draft acceptance but never changes the output. +* `align` keeps the backoff, and `KVCacheManagerV2` aligns the match down to a block instead + (`requires_block_aligned_reuse`, which a role that reads the pool always declares). The pool then + serves the trimmed block along with the rest of the prefix, at the cost of that block's transfer. + +It is an engine setting read from the `mooncake_store` block on every rank rather than part of the +rendered client config, so a deployment that inherits `MOONCAKE_CONFIG_PATH` without the block gets +`off`. A `producer` or `capacity` server keeps the cache manager's own reuse either way, and a +server without a one-model draft has no backoff for the setting to act on. + +### Diagnostics + +`mooncake_store.stats_interval: 60` has every rank of a server that reads the pool (`role: both` or +`consumer`) log two info-level lines at most once per 60 seconds, each covering the interval since +its last line and skipped when there is nothing to report: + +* `mooncake-store prefix-hole stats` counts lookups (`lookups`, `lookup_blocks`, `hit_blocks`) and + the holes they stopped at. A hole is a missing block followed by still-stored blocks of the same + chain, which the pool evicted from the middle of a prefix it otherwise holds: `holes` counts the + lookups that met one, `hole_blocks` the missing blocks inside them, and `holed_tokens` and + `max_holed_tokens` the stored tokens behind a hole that are recomputed anyway. They come from the + lookup's own existence probe, so they cost no extra store call. Only a rank that looks prefixes + up reports them: rank 0 under tensor parallelism, every owner under attention DP. Frequent holes + are what [`soft_pin_ttl_ms` and `group_blocks`](#pool-eviction) address. +* `mooncake-store load stats` counts loads (`start_load_kv` passes that moved pages, or with + [`async_load`](#asynchronous-loads) each parked request's load), their `pages`, store calls + (`batches`), `loaded_gib` and `failed_pages`, the transfer time from the destinations being ready + to the last page landing (`transfer_ms`), and each load's mean wall and CPU milliseconds for + `fence`, `resolve`, `describe`, `get`, `unstage` and `sync`. With `async_load` it also reports + the time a load waited for a load thread and the allocation fence (`wait_ms`), and, at the time + of the line, the load threads (`workers`) and the loads waiting for one (`queue`) or running + (`active`); a line is then logged while loads are waiting even if none finished. + +Both are flushed at shutdown. Unset, the default, counts nothing and costs nothing. The lines are +logged at info level, so a log level above info hides them. Like the eviction settings, the +interval is rendered into the client config and read back from an inherited `MOONCAKE_CONFIG_PATH`. + ### Restrictions Checked at startup, before any request is admitted, because each one's failure mode is KV replayed without all of the state it was computed with — a wrong answer rather than a slow one. -Beyond beam search, attention data parallelism, speculative decoding and Mamba/hybrid models, which -are rejected for [every connector](#a-page-slot-must-not-be-reassigned-underneath-the-connector), +Beyond beam search, Mamba/hybrid models and speculative decoding outside the +[speculative-prefill-only mode](#speculative-decoding-on-a-disaggregated-context-server), which are +rejected for [every connector](#a-page-slot-must-not-be-reassigned-underneath-the-connector), this one also rejects KV cache manager V1, context parallelism, pipeline parallelism, -sliding-window attention, and `sparse_attention_config` with an index-V cache unless -`sparse_disable_index_value=True`; it bypasses individual requests that carry a LoRA adapter or -multimodal content without hashes, leaving the rest of the deployment served. +sliding-window attention, `sparse_attention_config` with an index-V cache unless +`sparse_disable_index_value=True`, and, in a role that transfers KV, any drafter other than +one-model MTP or EAGLE3 sharing the target KV cache; it bypasses individual requests that carry a +LoRA adapter or multimodal content without hashes, leaving the rest of the deployment served. + +A page offered at lookup can still be gone at load time, evicted from a full pool or left behind by +a put that never completed, whether the load runs synchronously or on a +[load thread](#asynchronous-loads). The worker declares +[`reports_failed_loads`](#2-worker-interface-kvcacheconnectorworker), so on a single-rank server or +an attention-DP owner only the affected request restarts, without the offered prefix, and the rest +of its batch runs. Under tensor parallelism every rank holds a shard of the same page and cannot +restart the request alone, so a missing page fails the executor there. + +Aggregated serving under attention data parallelism is not supported yet, although startup does +not refuse it. When the attention-DP gate skips a forward pass after the batch was prepared, the +executor shrinks each generation request's KV capacity again, and a page the connector was already +told about can go to another request without the connector hearing of it; a later save could then +publish that request's bytes under this one's key. A disaggregated context server never runs +generation requests and is unaffected, as is a role that saves nothing (`consumer`, `capacity`). ### GPUDirect RDMA @@ -668,3 +1036,50 @@ A host without `nvidia_peermem` usually needs only `WITH_NVIDIA_PEERMEM=0`, whic selects on its own where it can. When neither path is available, set `stage_through_host: true` to copy pages through a pinned host buffer instead, at the cost of a copy each way. A `capacity` server needs neither, since it registers no pages. + +Each direction a role drives (loads, saves) pins `staging_buffer_bytes` of host memory per rank, +1 GiB by default, given as a binary size (`512MiB`) or a byte count. It bounds how many pages one +store call stages, not how large a page may be: a budget below one page still pins one page, with +a warning. The startup node memory check counts it per rank and direction, next to +`segment_size`. + +### Asynchronous loads + +By default a prefix the pool serves is loaded inside the executor iteration that schedules the +request, before its forward pass. With `async_load: true`, a server that reads the pool +(`role: both` or `consumer`) parks such a request outside the batch instead, and +`async_load_workers` background threads per rank (default 1) load its pages while other requests +keep computing. The request rejoins the batch once every rank holds its pages. Every hit is +offered this way, including a one-block offer under `local_reuse_backoff: align`, and parking a +request costs it at least one iteration. + +```yaml +kv_connector_config: + connector: mooncake-store + mooncake_store: + pool: file:///shared/pool.json + model_key: qwen2-1.5b-instruct + async_load: true + async_load_workers: 2 +disable_overlap_scheduler: true +``` + +| Path | How a load thread fills the pages | +|---|---| +| Zero copy (default) | The NIC writes into the registered KV pools, from the load thread. | +| `stage_through_host` | Each load thread has its own CUDA stream and `staging_buffer_bytes / async_load_workers` of pinned slots, and scatters each batch on that stream. No separate synchronous load buffer is pinned. | + +A load thread first waits for the KV cache manager's allocation fence for the request's pages, +so the executor never blocks on it. A page that is gone by load time restarts only that request, +without the offered prefix, on a single-rank server or an attention-DP owner, and fails the +server under tensor parallelism, as for synchronous loads. An exception raised by Mooncake on a +load thread always fails the server, at the next executor call. Under attention DP a parked load +counts toward `kv_connector_config.transfer_timeout_sec` (see +[Attention data parallelism](#attention-data-parallelism)). Asynchronous loads are validated with +`disable_overlap_scheduler: true`. + +The scheduler adapter declares +[admission-time queries](#admission-time-queries): `kv_connector_config.admission_query: {}` books +the pages a hit reaching past the first chunk needs while the request is admitted, with or without +`async_load`. Admitting a chunk without charging its tokens (`max_uncharged_parks`) needs +`async_load`, since only an asynchronous offer parks the request. diff --git a/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py b/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py index e5f33a18ce00..84b2751578ba 100644 --- a/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py +++ b/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py @@ -120,6 +120,9 @@ class KvCacheConfigV2: # tensorrt_llm/llmapi/llm_args.py; KVCacheManagerV2._build_base_config() # reads it unconditionally. kv_pool_rebalance_config: KvPoolRebalanceConfig = field(default_factory=KvPoolRebalanceConfig) + # Mirrors KvCacheConfig.enable_disagg_completion_reservation in + # tensorrt_llm/llmapi/llm_args.py; KVCacheManagerV2.__init__() reads it. + enable_disagg_completion_reservation: bool = False class _Timeout(Exception): diff --git a/examples/llm-api/configs/trtllm_mooncake_store_connector_extra.yaml b/examples/llm-api/configs/trtllm_mooncake_store_connector_extra.yaml index bfeff513f451..22c7ad9f5015 100644 --- a/examples/llm-api/configs/trtllm_mooncake_store_connector_extra.yaml +++ b/examples/llm-api/configs/trtllm_mooncake_store_connector_extra.yaml @@ -25,3 +25,12 @@ kv_connector_config: segment_size: 32GiB # Identifies this checkpoint to the pool. Must match on every server. model_key: qwen2-1.5b-instruct + # Soft-pin written pages so a full pool evicts by age; must not exceed the + # master's --max_kv_soft_pin_ttl (24 h by default). Needs Mooncake 0.3.13+. + # soft_pin_ttl_ms: 86400000 + # Evict every page of a block (all layer groups and TP shards) together. + # group_blocks: true + # Load a prefix the pool serves on background threads while the request + # waits outside the batch; validated with disable_overlap_scheduler: true. + # async_load: true + # async_load_workers: 1 diff --git a/tensorrt_llm/_torch/disaggregation/orchestration/coordinator.py b/tensorrt_llm/_torch/disaggregation/orchestration/coordinator.py index 82185c8f3e62..a79dceac40d4 100644 --- a/tensorrt_llm/_torch/disaggregation/orchestration/coordinator.py +++ b/tensorrt_llm/_torch/disaggregation/orchestration/coordinator.py @@ -119,6 +119,10 @@ def __init__( self._pending_timed_out_requests: List[LlmRequest] = [] # Requests whose timed-out transfer was already cancelled in flight. self._timed_out_ctx_cancelled_ids: Set[int] = set() + # Timed-out requests whose cancellation was requested and whose outcome + # the transceiver reports later: accepted in-flight cancels, or, + # without in-flight cancellation, generation cancels the transceiver + # refused. self._timed_out_gen_cancelled_ids: Set[int] = set() self._inflight_cancel_unsupported_logged = False # When a forward is still in flight at admission, as under the overlap @@ -223,7 +227,12 @@ def prepare_context_schedulable(self, new_requests: List[LlmRequest]) -> None: @nvtx_range("poll_gen_transfers") def poll_gen_transfers(self) -> None: - """Poll receive-side transfers and their timeouts; rank-synchronized.""" + """Poll receive-side transfers and their timeouts; rank-synchronized. + + Without in-flight cancellation, also cancels and fails the receives + that ``check_transfer_timeouts`` flagged, so an idle server does not + wait for a forward step. + """ if not uses_async_gen_transfer(): return if self._defer_unready_gen_receives: @@ -235,6 +244,8 @@ def poll_gen_transfers(self) -> None: if self.inflight_cancel_active(): self._cancel_timed_out_gen_transfers() self._check_gen_transfer_errors_consensus() + else: + self._fail_timed_out_gen_transfers() @nvtx_range("check_transfer_timeouts") def check_transfer_timeouts(self, only_with_context_sends: bool = False) -> None: @@ -747,9 +758,19 @@ def request_cancellation(self, request: LlmRequest) -> bool: def fail_timed_out(self, requests: List[LlmRequest]) -> None: """Fail generation requests whose transfer timed out. - Under multi-rank ADP the error response enters a collective, so it is - deferred to ``handle_timeouts_synced``. + Callers pass only requests whose transfer no longer touches their + blocks, either because ``cancel_request`` returned True or because the + transfer already finished. Under multi-rank ADP the error response + enters a collective, so it is deferred to ``handle_timeouts_synced``. """ + for req in requests: + # cancel_request does not change the state. Left in progress, + # PyExecutor._handle_errors would park the failure for a transfer + # that no longer exists (_transfer_owns_request_blocks), and + # _sweep_transfer_held_failures would drop it without a response + # because the request has already left active_requests. + if req.state == LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS: + req.state = LlmRequestState.DISAGG_TRANS_ERROR if self._enable_attention_dp and self._dist.world_size != 1: self._pending_timed_out_requests.extend(requests) return @@ -856,6 +877,59 @@ def _cancel_timed_out_gen_transfers(self) -> None: "to report final cleanup" ) + @nvtx_range("fail_timed_out_gen_transfers") + def _fail_timed_out_gen_transfers(self) -> None: + """Cancel and fail generation receives flagged by ``check_transfer_timeouts``. + + Runs from ``poll_gen_transfers`` every iteration when in-flight + cancellation is inactive. The response pass does the same, but only + after a forward step, so without this an idle generation server would + keep a request whose KV never arrives until the client gives up. + + A cancel that finishes at once fails the request through + ``fail_timed_out``: inline, or buffered for ``handle_timeouts_synced`` + under multi-rank ADP. A cancel the transceiver refuses is asked only + once. The Python transceiver has then marked the session cancelled, + and the next ``reap_gen_receives`` fails the request through the + transfer-error path. The C++ transceiver cannot cancel a receive + already in flight without in-flight cancellation, so retrying would + only repeat its warning; that receive settles on its own. + User-canceled requests are left to ``_handle_canceled_requests``. + """ + if self._transceiver.kv_transfer_timeout_ms is None: + return + timed_out = [ + request + for request in self._registry.active_requests() + if request.py_kv_transfer_timed_out + and request.is_disagg_generation_transmission_in_progress + and request.py_request_id not in self._timed_out_gen_cancelled_ids + ] + if not timed_out: + return + user_canceled_ids = set(self._registry.canceled_request_ids()) + cancelled = [] + for request in timed_out: + if self._vote_id(request) in user_canceled_ids: + continue + if self.request_cancellation(request): + cancelled.append(request) + logger.warning( + f"Cancelled timed-out generation KV transfer for request " + f"{request.py_request_id}; failing it" + ) + else: + self._timed_out_gen_cancelled_ids.add(request.py_request_id) + logger.warning( + f"Requested cancellation of timed-out generation KV transfer " + f"for request {request.py_request_id}; waiting for the " + "transceiver to report it" + ) + for request in cancelled: + self._registry.remove(request) + if cancelled: + self.fail_timed_out(cancelled) + @nvtx_range("check_gen_transfer_errors_consensus") def _check_gen_transfer_errors_consensus(self) -> None: """Flush generation transfer errors through a TP-uniform path.""" diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 492d6cc6cd77..c4b3ede846ba 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -4294,6 +4294,12 @@ def create_py_executor_instance( enable_prefix_aware_scheduling = ( scheduler_config.enable_prefix_aware_scheduling if scheduler_config is not None else True) + # Admission-time connector queries. The executor checks the connector, + # the overlap scheduler and attention DP once the connector is bound; + # the fallbacks match the config defaults and are inert while it is off. + kv_connector_config = getattr(llm_args, "kv_connector_config", None) + admission = (kv_connector_config.admission_query + if kv_connector_config is not None else None) scheduler = KVCacheV2Scheduler( max_batch_size=max_batch_size, max_num_tokens=max_num_tokens, @@ -4309,6 +4315,11 @@ def create_py_executor_instance( enable_prefix_aware_scheduling=enable_prefix_aware_scheduling, # A disaggregated generation worker must not replay context locally. enable_recompute_pause=not is_disagg, + connector_admission_query=admission is not None, + max_uncharged_connector_parks=(admission.max_uncharged_parks + if admission is not None else 0), + max_connector_offer_deferrals=(admission.max_offer_deferrals + if admission is not None else 16), ) elif (scheduler_config is not None and scheduler_config.use_python_scheduler): diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py index 8c4c3815493a..e5c62345bfce 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py @@ -65,6 +65,27 @@ AllocationFence = Callable[[Iterable[int]], None] +def _is_dummy_request(request: LlmRequest) -> bool: + """Whether ``request`` is a placeholder with no prompt of its own. + + Attention DP pads idle ranks with these, and they carry nothing a connector + could load or save. Reads ``is_dummy_request``, which the C++ ``LlmRequest`` + binds and so is readable on the base class some callers hand over, and the + Python-only ``is_attention_dp_dummy``. Each is compared against True, since + a mock request answers every attribute truthily and is not a dummy. + + Args: + request: The request to classify. + + Returns: + True when either dummy flag is set. + """ + return ( + getattr(request, "is_dummy_request", False) is True + or getattr(request, "is_attention_dp_dummy", False) is True + ) + + # Used to store data for a single inflight request. @dataclass class RequestData: @@ -180,6 +201,54 @@ class KvCacheConnectorWorker(ABC): # Opt in only when transfers, endpoints and completion are owner-local, # and the backend can safely share compatible cache content across owners. supports_attention_dp: bool = False + #: Opt in to run next to one-model speculative decoding (MTP or EAGLE3) on + #: a disaggregated context server; see + #: `KvCacheConnectorManager.speculative_prefill_only`. Declare it only when + #: every load and save moves whole pages of every layer group in the + #: registered `KvCacheLayout`, which includes the one-model draft layers + #: that share the target pool. Per-layer hooks are installed on + #: `DecoderLayer` modules only, and whether a model's draft layers are such + #: modules depends on the model, so a worker that saves per layer must + #: itself make sure it covers the draft layers before opting in. The + #: scheduler must declare it too. Only an explicit True counts, and it is + #: read from the class before either half is constructed. + supports_speculative_prefill: bool = False + #: The connector stores and offers KV by whole blocks starting at the local + #: prefix match, so that match must end on a block boundary. + #: `KVCacheManagerV2` reads it only when a one-model draft's lookahead backoff + #: is in effect and partial reuse is off, and then aligns every local match + #: down to `tokens_per_block`. Only an explicit True counts. May be set per + #: instance in `__init__`, since the manager is built after the connector. + #: Must be the same on every rank. + requires_block_aligned_reuse: bool = False + #: The connector prefers whole local matches to backed-off ones. The target + #: `KVCacheManagerV2` then drops the one-model draft lookahead backoff; the + #: draft KV at the tail of a reused block may belong to a different next + #: token, which can lower draft acceptance but never changes the output. + #: Only an explicit True counts. May be set per instance. Must be the same + #: on every rank. + drops_reuse_match_backoff: bool = False + #: Opt in when a prefix load can come up short -- a page offered at lookup + #: is gone at load time -- and the worker reports the affected requests + #: through `take_failed_load_requests` / `take_failed_async_load_requests` + #: instead of raising. The runtime then restarts only those requests. On a + #: tensor-parallel rank (more than one rank without attention DP) the + #: worker must still raise, because its peers hold shards of the same + #: pages. Requires `take_failed_load_requests` and `drop_bound_saves`, + #: KVCacheManagerV2, and the final-batch query path (prefix reservations + #: off). Only an explicit True counts. + reports_failed_loads: bool = False + #: Opt in to completion headroom on a disaggregated server whose + #: `KVCacheManagerV2` has no tier below GPU. The manager then admits a + #: request only while the pages every admitted request still needs to + #: finish fit -- a decode its output, a connector prefill the rest of its + #: prompt -- so a restore chosen after the final batch is not starved by + #: other requests' chunks. Only an explicit True counts, and a + #: `capacity_only` worker's is ignored. Read from the worker, which every + #: rank holds, so every rank sees it; may be set per instance in + #: `__init__`, since the manager is built after the connector. Must be the + #: same on every rank. + requires_completion_reservation: bool = False def __init__(self, llm_args: TorchLlmArgs): self._llm_args = llm_args @@ -360,6 +429,65 @@ def get_finished_prefix_loads(self) -> List[int]: """ return [] + def take_failed_load_requests(self) -> Set[int]: + """Return and forget the requests whose last synchronous prefix load came up short. + + Polled right after every `start_load_kv` when `reports_failed_loads` is + True. Every ID must be the `RequestData.request_id` of a context + request in the metadata that call loaded; an ID outside that batch + fails the executor. The runtime removes each named request from this + pass, calls `drop_bound_saves` for it, frees its allocation and rewinds + it, and the scheduler readmits it and asks `get_num_new_matched_tokens` + again. + + Every write into a named request's pages must be complete or abandoned + when its ID is returned: the pages are freed right away and may be + handed to another request, so no DMA into them may still be in flight. + + Returns: + The request IDs to restart. Empty by default. + """ + return set() + + def take_failed_async_load_requests(self) -> Set[int]: + """Return and forget the parked requests whose asynchronous load came up short. + + Polled right after every `get_finished` when `reports_failed_loads` is + True. Each ID must have been returned as finished loading by the + `get_finished` call this poll follows: once a load is reported finished + the scheduler may admit the request over its pages, so a later report + would come too late to restart it. An ID whose load is still pending, + or that an earlier `get_finished` already returned as finished, fails + the executor. The same no-DMA-in-flight rule as + `take_failed_load_requests` applies. + + Returns: + The request IDs to restart. Empty by default. + """ + return set() + + def drop_bound_saves(self, request_ids: Iterable[int]) -> int: + """Remove these requests' saves from the metadata bound for this pass. + + Called before `wait_for_save` for a request restarted after + `bind_connector_meta`, whose saves would otherwise publish pages it + never computed. Must accept IDs with nothing bound, and be idempotent. + + Args: + request_ids: The restarted requests. + + Returns: + The number of save entries removed. + + Raises: + NotImplementedError: By default. A worker that declares + `reports_failed_loads` must implement it. + """ + raise NotImplementedError( + f"{type(self).__name__} declares reports_failed_loads but does not " + "implement drop_bound_saves" + ) + def shutdown(self) -> None: """Release whatever this worker holds, once the executor is done with it. @@ -372,6 +500,37 @@ class KvCacheConnectorScheduler(ABC): # ADP creates one adapter per rank. Adapters may connect to the same # logical storage pool; their request and allocation state stays local. supports_attention_dp: bool = False + #: Opt in to run next to one-model speculative decoding (MTP or EAGLE3) on + #: a disaggregated context server; see + #: `KvCacheConnectorManager.speculative_prefill_only`. A one-model draft + #: writes KV at each prompt position from the next + #: `draft_prompt_lookahead(llm_args.speculative_config)` prompt tokens, so + #: declare it only when the adapter's content identity (its store keys) + #: covers those tokens past each block's end; a served block's draft-layer + #: KV then belongs to this prompt. The adapter must never offer or save a + #: block whose lookahead tokens fall outside the prompt, and must never + #: save a block this pass has not computed: one-model speculation + #: allocates `max_draft_len - 1` extra KV slots past every context chunk, + #: so a page can exist before its block is computed. The worker must + #: declare it too. Only an explicit True counts, and it is read from the + #: class before either half is constructed. + supports_speculative_prefill: bool = False + #: Whether `get_num_new_matched_tokens` may be asked while the KV cache + #: manager V2 scheduler is still admitting the request, before the batch is + #: final (`kv_connector_config.admission_query`). The contract: + #: - the answer has no effect beyond itself until `update_state_after_alloc`; + #: - asking again for the same allocation replaces the earlier answer; + #: - a query may be followed by `request_reset` and another query with no + #: `update_state_after_alloc` in between, which is how a deferred first + #: chunk gives its pages back. + #: A request asked at admission can still be dropped before + #: `prepare_resources` -- by the disaggregated benchmark gate, batch waiting + #: for generation requests, the total-KV cap, the multimodal scheduler + #: wrapper, or a batch that cannot be queued -- and is then asked again on a + #: later iteration. A connector that starts a transfer inside + #: `get_num_new_matched_tokens` must not declare it. Read from the class, + #: so every tensor-parallel rank agrees. + supports_admission_query: bool = False def __init__(self, llm_args: TorchLlmArgs): self._llm_args = llm_args @@ -489,10 +648,12 @@ def request_reset(self, request: LlmRequest) -> None: """ Called when a request's allocation is released for it to run again. - Rollback and failed admission take this path, where request_finished - ends a request instead. Another request may own the released pages - before this one is readmitted, so state the connector recorded in page - indices has to be dropped or rebuilt here. Defaults to a no-op. + Rollback, failed admission and a restart after a failed prefix load + (`KvCacheConnectorWorker.reports_failed_loads`) take this path, where + request_finished ends a request instead. Another request may own the + released pages before this one is readmitted, so state the connector + recorded in page indices has to be dropped or rebuilt here. Defaults to + a no-op. Args: request: The request whose allocation was released. @@ -683,8 +844,7 @@ def update_and_build_data(self, req: LlmRequest, kv_cache_manager: "KVCacheManag class KvCacheConnectorSchedulerOutputManager: - def __init__(self, *, enable_attention_dp: bool = False) -> None: - self.enable_attention_dp = enable_attention_dp + def __init__(self) -> None: self.requests = defaultdict(KvCacheConnectorSchedulerOutputRequest) self.external_loads = dict() @@ -697,7 +857,7 @@ def build_scheduler_output( scheduler_output = SchedulerOutput() for req in scheduled_batch.context_requests: - if self.enable_attention_dp and req.is_dummy_request: + if _is_dummy_request(req): continue if req.request_id in new_async_requests.loading_ids: continue @@ -718,7 +878,7 @@ def build_scheduler_output( scheduler_output.cached_requests.append(request_data) for req in scheduled_batch.generation_requests: - if self.enable_attention_dp and req.is_dummy_request: + if _is_dummy_request(req): continue request_data = self.requests[req.request_id].update_and_build_data( req, kv_cache_manager @@ -755,6 +915,9 @@ class KvCacheConnectorManager(KvCacheConnectorManagerCpp): 2. Handling MPI communication for tensor parallelism. Under attention DP, each rank runs its own scheduler adapter and worker without connector collectives. Adapters can use a shared external storage pool. + 3. Refusing every request that is not context-only work in the + speculative-prefill-only mode, so one-model speculative decoding never + rejects draft tokens underneath the connector. Note: This class is solely an implementation detail, and is not part of the connector interface itself. When implementing a connector API, you do not need to implement this class. @@ -766,7 +929,31 @@ def __init__( scheduler: Optional[KvCacheConnectorScheduler], *, enable_attention_dp: bool = False, + speculative_prefill_only: bool = False, + admission_query: bool = False, ) -> None: + """Bind the connector halves and the modes the executor runs them in. + + Args: + worker: The connector worker, present on every rank. + scheduler: The connector scheduler: on every attention-DP owner, or + on rank 0 only under tensor parallelism. + enable_attention_dp: Run the scheduler locally on each + attention-DP owner, without connector collectives. Both halves + must declare `supports_attention_dp`. + speculative_prefill_only: Allow one-model speculative decoding by + refusing every request that is not context-only (disaggregated + prefill), which keeps the connector's append-only page deltas + valid. Both halves must declare `supports_speculative_prefill`. + admission_query: The connector scheduler declares + `KvCacheConnectorScheduler.supports_admission_query`. Decided + from the scheduler class, which every rank resolves, because + under tensor parallelism only rank 0 holds an instance. + + Raises: + NotImplementedError: When a mode is asked for and a connector half + does not declare it. + """ assert (scheduler is not None) == (enable_attention_dp or mpi_rank() == 0), ( "A scheduler is required on every attention-DP owner, or only rank 0 for TP." ) @@ -783,6 +970,34 @@ def __init__( #: because this disables most of the per-request path, which is too #: much to infer from a merely truthy value. self.capacity_only = worker.capacity_only is True + #: Whether the KV cache manager may ask the connector while the + #: scheduler admits a request; see + #: `KvCacheConnectorScheduler.supports_admission_query`. Off under + #: attention DP: owners schedule independently, and the batch-start + #: status vote and padding in `PyExecutor._kv_connector_start_batch` + #: were not validated with requests admitted without a token charge. + self.admission_query_capable = bool(admission_query) and not enable_attention_dp + #: Whether the worker hands failed prefix loads back for a restart + #: instead of raising; see `KvCacheConnectorWorker.reports_failed_loads`. + #: Compared against True like `capacity_only`, and off for a role that + #: loads nothing. + self.reports_failed_loads = ( + getattr(worker, "reports_failed_loads", False) is True and not self.capacity_only + ) + if self.reports_failed_loads: + self._validate_failed_load_reporting(type(worker)) + if speculative_prefill_only and not self.declares_speculative_prefill( + type(worker), None if scheduler is None else type(scheduler) + ): + raise NotImplementedError( + "Speculative-prefill-only mode requires a connector whose scheduler " + "and worker both declare supports_speculative_prefill=True." + ) + #: Whether only context-only requests may reach the connector, which is + #: what lets it run next to one-model speculative decoding; see + #: `_validate_speculative_prefill_request`. Compared against True like + #: `capacity_only`. + self.speculative_prefill_only = speculative_prefill_only is True # Requests that haven't yet been passed into get_finished. self.new_async_requests = AsyncRequests(dict(), dict()) @@ -798,9 +1013,7 @@ def __init__( self._scheduler_output = None self._allocation_fence: Optional[AllocationFence] = None - self.scheduler_output_manager = KvCacheConnectorSchedulerOutputManager( - enable_attention_dp=enable_attention_dp - ) + self.scheduler_output_manager = KvCacheConnectorSchedulerOutputManager() self.prefix_reservations_enabled = False self._prefix_capability = None self._next_prefix_reservation_id = 1 @@ -815,6 +1028,9 @@ def __init__( self._prefix_completion_tracker: PrefixLoadCompletionTracker | None = None self._deferred_load_terminations: Dict[int, LlmRequest] = {} self._finished_load_terminations: List[LlmRequest] = [] + # Asynchronous loads the latest `get_finished` returned as finished: the + # only ones a failed-load report may still restart. + self._last_finished_load_ids: Set[int] = set() @staticmethod def validate_attention_dp( @@ -833,6 +1049,54 @@ def validate_attention_dp( "and implement owner-local scheduling/transfers against a shared pool." ) + @staticmethod + def declares_speculative_prefill(worker_cls: type, scheduler_cls: Optional[type]) -> bool: + """Whether every present connector half opted into speculative prefill. + + Only an explicit True counts; see + `KvCacheConnectorScheduler.supports_speculative_prefill` and + `KvCacheConnectorWorker.supports_speculative_prefill`. + + Args: + worker_cls: The connector worker class. + scheduler_cls: The connector scheduler class, or None on a TP + follower rank, which holds no scheduler instance. + + Returns: + True when the worker and, if given, the scheduler declare it. + """ + return all( + getattr(cls, "supports_speculative_prefill", False) is True + for cls in (worker_cls, scheduler_cls) + if cls is not None + ) + + @staticmethod + def _validate_failed_load_reporting(worker_cls: type) -> None: + """Refuse a worker that declares `reports_failed_loads` without its hooks. + + Only the base defaults count as missing, so a worker class that does + not inherit `KvCacheConnectorWorker` is not refused here. + + Args: + worker_cls: The worker class that declares `reports_failed_loads`. + + Raises: + ValueError: When `take_failed_load_requests` or `drop_bound_saves` + is left at the base default. + """ + missing = [ + name + for name in ("take_failed_load_requests", "drop_bound_saves") + if getattr(worker_cls, name, None) is getattr(KvCacheConnectorWorker, name) + ] + if missing: + raise ValueError( + f"{worker_cls.__name__} declares reports_failed_loads=True but does not " + f"implement {', '.join(missing)}; a failed prefix load would have no way " + "back to the runtime." + ) + def _run_on_leader(self, f: Callable[[], _SchedulerResult]) -> _SchedulerResult: """ Run locally under ADP, or broadcast the TP leader's result to its workers. @@ -897,6 +1161,18 @@ def configure_prefix_reservations(self, enabled: bool) -> None: self.prefix_reservations_enabled = ( enabled and self._prefix_capability and not self.capacity_only ) + if self.prefix_reservations_enabled and self.reports_failed_loads: + # Every rank ran the collective probe above with the same worker + # class, so all of them raise here together. + raise ValueError( + "KV connector reports_failed_loads is not supported together with " + "prefix reservations. A reserved load completes through " + "get_finished_prefix_loads and the prefix-load completion tracker, " + "which have no failure outcome, and KVCacheManagerV2.free_resources " + "refuses a dispatched load. Use the final-batch query path: drop the " + "reservation methods from the connector, or set " + "scheduler_config.enable_prefix_aware_scheduling=False." + ) if self.prefix_reservations_enabled and self._prefix_completion_tracker is None: comm = ( mpi_comm().Dup() if not self.enable_attention_dp and mpi_world_size() > 1 else None @@ -905,10 +1181,11 @@ def configure_prefix_reservations(self, enabled: bool) -> None: def reserve_prefix(self, request: LlmRequest, local_end: int) -> Optional[PrefixReservation]: """Query once per pending attempt and protect the promised source range.""" - if not self.prefix_reservations_enabled or ( - self.enable_attention_dp and request.is_dummy_request - ): + if not self.prefix_reservations_enabled or _is_dummy_request(request): return None + # Ahead of the leader broadcast below, so every tensor-parallel rank + # raises for the same request instead of leaving followers waiting. + self._validate_speculative_prefill_request(request) if local_end < 0: raise ValueError("A prefix reservation cannot start before the prompt") existing = self.get_prefix_reservation(request) @@ -1097,7 +1374,11 @@ def query_num_new_matched_tokens( The connector ABC promises one query per allocation, so a caller must reach this at most once per request per allocation whatever it does with - the result. A caller that cannot consume the whole answer -- one + the result. The one exception is a scheduler that declares + `KvCacheConnectorScheduler.supports_admission_query`: a request asked + while it is admitted and then dropped before `prepare_resources` is + asked again, and the new answer replaces the old one. A caller that + cannot consume the whole answer -- one allocating per context chunk, where an offer can outrun the reserved pages -- records only the part it honours via ``commit_new_matched_tokens``. @@ -1108,13 +1389,16 @@ def query_num_new_matched_tokens( ``KVCacheManagerV2`` holds the Python subclass where it is a method, so no single spelling of the check reads correctly in both. """ - if self.enable_attention_dp and request.is_dummy_request: + if _is_dummy_request(request): return 0, False if self.capacity_only: # No page of this cache is registered, so nothing can be loaded # into it, and the query would cost the leader broadcast below on # every request. return 0, False + # Ahead of the leader broadcast below, so every tensor-parallel rank + # raises for the same request instead of leaving followers waiting. + self._validate_speculative_prefill_request(request) num_tokens, load_kv_async = self._run_on_leader( lambda: self.scheduler.get_num_new_matched_tokens(request, num_computed_tokens) @@ -1150,7 +1434,7 @@ def commit_new_matched_tokens( ``cancel_load``, and a parked request carries no honoured count into ``build_connector_meta``. """ - if self.enable_attention_dp and request.is_dummy_request: + if _is_dummy_request(request): return if self.capacity_only: # Only `build_scheduler_output` consumes `external_loads`, and this @@ -1200,6 +1484,33 @@ def get_num_new_matched_tokens(self, request: LlmRequest, num_computed_tokens: i self.commit_new_matched_tokens(request, num_tokens, load_kv_async) return num_tokens + def _validate_speculative_prefill_request(self, request: LlmRequest) -> None: + """Keep one-model speculation's append-only page lists on context-only work. + + A context-only request never rejects a draft token, so its page list + only grows; anything else could shrink it underneath the connector. + This backs up the executor's admission check, for anything that + bypasses it. Dummy requests carry nothing to load or save and pass. + + Args: + request: The request about to reach the connector. + + Raises: + NotImplementedError: In the speculative-prefill-only mode, for a + request that is neither a dummy nor context-only. + """ + if ( + self.speculative_prefill_only + and not _is_dummy_request(request) + and getattr(request, "is_context_only_request", False) is not True + ): + raise NotImplementedError( + "This KV connector runs next to one-model speculative decoding on " + "context-only (disaggregated prefill) requests only. Run " + "speculative generation on a server whose connector transfers no " + "KV (capacity_only) or without a connector." + ) + _GROUPED_SCHEDULER_METHODS = ( "update_state_after_alloc_by_layer_group", "request_finished_by_layer_group", @@ -1291,6 +1602,124 @@ def should_add_sequence(self, request: LlmRequest) -> bool: req_id = request.request_id return req_id not in self.finished_async_loading_requests + def take_failed_load_requests(self) -> Set[int]: + """Return the requests whose synchronous prefix load in this pass came up short. + + Polled right after `start_load_kv`; see + `KvCacheConnectorWorker.take_failed_load_requests` for the contract and + `PyExecutor._recover_failed_connector_loads` for the restart. + + Returns: + The worker's request IDs, or an empty set for a connector that does + not declare `reports_failed_loads`. + + Raises: + RuntimeError: When a tensor-parallel rank reports a failure. + """ + if not self.reports_failed_loads: + return set() + return self._require_restartable(set(self.worker.take_failed_load_requests())) + + def take_failed_async_load_requests(self) -> List[LlmRequest]: + """Return the parked requests whose asynchronous prefix load came up short. + + Unlike `KvCacheConnectorWorker.take_failed_async_load_requests`, which + returns IDs, this returns the requests themselves, popped from + `finished_async_loading_requests` so that `should_add_sequence` admits + them afresh. An ID for a request that finished or was cancelled while + parked names nothing to restart: its termination already went to + `take_finished_load_terminations`, so it is dropped here. + + Returns: + The requests to restart, in request-ID order. Empty for a connector + that does not declare `reports_failed_loads`. + + Raises: + RuntimeError: When a tensor-parallel rank reports a failure; when a + reported load has not been reported finished, whose destination + pages may still be written; or when an earlier `get_finished` + already returned it as finished, so the scheduler may already + have admitted the request over the missing pages. + """ + if not self.reports_failed_loads: + return [] + failed = self._require_restartable(set(self.worker.take_failed_async_load_requests())) + still_loading = sorted( + request_id + for request_id in failed + if any( + request_id in requests.loading + for requests in ( + self.new_async_requests, + self.pending_async_requests, + self.local_finished_async_requests, + ) + ) + ) + if still_loading: + raise RuntimeError( + f"KV connector reported failed asynchronous loads for requests {still_loading} " + "that it has not reported finished; the destination pages may still be written." + ) + # `finished_async_loading_requests` keeps a request until it ends, well + # after the scheduler admitted it, so only this pass's finishes are + # still parked. + late = sorted( + request_id + for request_id in failed + if request_id in self.finished_async_loading_requests + and request_id not in self._last_finished_load_ids + ) + if late: + raise RuntimeError( + f"KV connector reported failed asynchronous loads for requests {late} after an " + "earlier get_finished returned them as finished; they may already be scheduled " + "over the missing pages." + ) + restarted = [] + for request_id in sorted(failed): + request = self.finished_async_loading_requests.pop(request_id, None) + if request is not None: + restarted.append(request) + return restarted + + def forget_request(self, request: LlmRequest) -> None: + """Drop the saves bound for a request that is about to restart from scratch. + + Those saves would publish pages the request never computed. Nothing is + reset here: `KVCacheManagerV2.drop_context_allocation` releases the + allocation next, and its `free_resources` runs `reset_request_state`. + + Args: + request: The request to restart. + """ + self.worker.drop_bound_saves([request.request_id]) + + def _require_restartable(self, failed: Set[int]) -> Set[int]: + """Refuse failed-load reports from a tensor-parallel rank. + + A restart is owner-local, which holds under attention DP and on a + single rank. Tensor-parallel peers hold shards of the same pages and + would keep computing on them, so a worker there must raise instead; + this is the backstop for one that does not. + + Args: + failed: The request IDs the worker reported. + + Returns: + ``failed``, unchanged. + + Raises: + RuntimeError: When ``failed`` is not empty on a tensor-parallel rank. + """ + if failed and not self.enable_attention_dp and mpi_world_size() > 1: + raise RuntimeError( + f"KV connector reported failed prefix loads for requests {sorted(failed)} on a " + "tensor-parallel rank; restarts are owner-local (attention DP or a single " + "rank). The connector must raise instead." + ) + return failed + def has_pending_transfers(self, request_id: int) -> bool: """Keep canceled requests allocated until the worker releases DMA access.""" return request_id in self._prefix_load_requests or any( @@ -1323,6 +1752,13 @@ def build_scheduler_output( # which on a generation server is inter-token latency. # `handle_metadata` returns early when none was built. return + if self.speculative_prefill_only: + # Before the output manager records any state for this batch. + for request in ( + *scheduled_batch.context_requests, + *scheduled_batch.generation_requests, + ): + self._validate_speculative_prefill_request(request) async_requests = AsyncRequests( {}, { @@ -1425,7 +1861,7 @@ def request_finished( free_resources on the request. """ - if self.enable_attention_dp and req.is_dummy_request: + if _is_dummy_request(req): return False if req.request_id in self.finished_async_loading_requests: @@ -1538,6 +1974,7 @@ def get_finished(self) -> List[LlmRequest]: # For requests that have finished loading, move them back to the context state. for req in all_finished.loading.values(): self._finish_load(req, is_async=True) + self._last_finished_load_ids = set(all_finished.loading) if self.prefix_reservations_enabled: finished = ( @@ -1601,11 +2038,7 @@ def update_state_after_alloc( block_ids: List[int], block_ids_by_layer_group: Optional[List[List[int]]] = None, ) -> None: - if ( - self.scheduler is None - or self.capacity_only - or (self.enable_attention_dp and req.is_dummy_request) - ): + if self.scheduler is None or self.capacity_only or _is_dummy_request(req): return # A cache that reports per layer group is reported that way, whatever # the group count -- including one empty list per group for a request diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py index 4cbd8243bd11..c7213f223354 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py @@ -24,12 +24,13 @@ """ import json +import math import os import re import socket from dataclasses import dataclass from enum import Enum -from typing import TYPE_CHECKING, Any, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple if TYPE_CHECKING: from tensorrt_llm.llmapi.llm_args import MooncakeStoreConfig, TorchLlmArgs @@ -37,10 +38,12 @@ __all__ = [ "CLIENT_CONFIG_NAME", "CONFIG_PATH_ENV", + "DEFAULT_STAGING_BUFFER_BYTES", "MooncakeStoreConnectorConfig", "SEGMENTS_DIR_NAME", "StoreRole", "local_address", + "local_reuse_backoff", "parse_size", "pool_config", "provisioned_config_path", @@ -64,6 +67,10 @@ #: Per-process Mooncake transfer buffer, not pool capacity. Written into the #: rendered config so a vLLM reader of the same file sees the same value. DEFAULT_LOCAL_BUFFER_SIZE = 1073741824 +#: Pinned host memory per transfer direction under `stage_through_host`. Must +#: equal `parse_size` of the `MooncakeStoreConfig.staging_buffer_bytes` default, +#: so a client config without the key stages exactly what an unset field asks. +DEFAULT_STAGING_BUFFER_BYTES = 1 << 30 DEFAULT_NAMESPACE = "trtllm" #: Mooncake's peer-to-peer handshake, which keeps a separate metadata process #: out of the deployment. An empty connstring is not a form `store.setup` @@ -239,6 +246,18 @@ def parse_size(value: Any, *, strict_units: bool = False) -> int: return int(float(magnitude) * scale) +def _json_bool(raw: Dict[str, Any], name: str, default: bool, path: str) -> bool: + """A JSON boolean setting, refusing anything else. + + `bool("false")` is True, so coercing a quoted value would silently turn on + what the file meant to turn off. + """ + value = raw.get(name, default) + if not isinstance(value, bool): + raise ValueError(f"{name}={value!r} in {path} must be true or false") + return value + + def pool_config(llm_args: "TorchLlmArgs") -> Optional["MooncakeStoreConfig"]: """The `mooncake_store` block of `llm_args`, if the deployment set one. @@ -249,6 +268,25 @@ def pool_config(llm_args: "TorchLlmArgs") -> Optional["MooncakeStoreConfig"]: return connector.mooncake_store if connector is not None else None +def local_reuse_backoff(llm_args: "TorchLlmArgs") -> Literal["off", "align"]: + """How a local prefix match is handled ahead of the pool under a one-model draft. + + See `MooncakeStoreConfig.local_reuse_backoff`. It is an engine setting read + from `llm_args` on every rank, not part of the rendered client config, so a + deployment that inherits `MOONCAKE_CONFIG_PATH` without a `mooncake_store` + block gets the field's default. + + Args: + llm_args: The engine arguments every rank parses. + + Returns: + `"off"` to drop the draft's reuse-match backoff, or `"align"` to keep it + and align local matches down to a block boundary. + """ + pool = pool_config(llm_args) + return pool.local_reuse_backoff if pool is not None else "off" + + def provisioned_config_path(run_dir: Optional[str]) -> Optional[str]: """The client config a server rendered into `run_dir`, if there is one. @@ -290,6 +328,32 @@ class MooncakeStoreConnectorConfig: #: pools with Mooncake. Costs a copy each way, but works without GPUDirect #: RDMA, which registering device memory requires. stage_through_host: bool = False + #: Pinned host memory per transfer direction under `stage_through_host`. + #: It bounds how many pages one store call stages, not how large a page may + #: be. With `async_load` the load direction's share is split evenly across + #: the load workers. + staging_buffer_bytes: int = DEFAULT_STAGING_BUFFER_BYTES + #: Load a prefix the pool serves on background threads while the request + #: waits outside the batch, instead of inside the executor iteration. + #: Ignored unless the role reads the pool. + async_load: bool = False + #: Background load threads per rank when `async_load` is set, each with its + #: own CUDA stream and, under `stage_through_host`, an equal share of + #: `staging_buffer_bytes`. + async_load_workers: int = 1 + #: Soft-pin every page this server writes for this many milliseconds, or + #: `None` to leave pages unpinned. Mooncake's first eviction pass refills a + #: shortfall in shard order rather than by age, which punches holes into + #: prefixes still in use; soft-pinned pages are evicted only by the second, + #: age-ordered pass. Must not exceed the master's `--max_kv_soft_pin_ttl`, + #: or the master refuses every put. + soft_pin_ttl_ms: Optional[int] = None + #: Put every page of a block, across layer groups and attention shards, + #: into one eviction group, so the master evicts a block whole. + group_blocks: bool = False + #: Seconds between each rank's prefix-hole and load statistics lines, or + #: `None` to count nothing. + stats_interval: Optional[float] = None def __post_init__(self) -> None: """Reject settings that would fail later, inside a transfer.""" @@ -301,6 +365,20 @@ def __post_init__(self) -> None: raise ValueError("global_segment_size must be >= 0") if self.transfer_batch_size <= 0: raise ValueError("transfer_batch_size must be > 0") + if self.staging_buffer_bytes <= 0: + raise ValueError("staging_buffer_bytes must be > 0") + if self.async_load_workers < 1: + raise ValueError("async_load_workers must be >= 1") + if self.soft_pin_ttl_ms is not None and self.soft_pin_ttl_ms <= 0: + raise ValueError( + f"soft_pin_ttl_ms must be > 0, or absent to leave pages unpinned; " + f"got {self.soft_pin_ttl_ms}" + ) + if self.stats_interval is not None and not 0 < self.stats_interval < math.inf: + raise ValueError( + f"stats_interval must be a finite number of seconds > 0, or absent " + f"to count nothing; got {self.stats_interval}" + ) @property def capacity_only(self) -> bool: @@ -336,6 +414,20 @@ def from_file(path: str) -> "MooncakeStoreConnectorConfig": model_key=raw.get("model_key") or None, transfer_batch_size=int(raw.get("transfer_batch_size", 64)), stage_through_host=bool(raw.get("stage_through_host", False)), + # A config provisioned for vLLM carries none of these three, and + # gets today's synchronous loads through 1 GiB of staging. + staging_buffer_bytes=parse_size( + raw.get("staging_buffer_bytes", DEFAULT_STAGING_BUFFER_BYTES), strict_units=True + ), + async_load=_json_bool(raw, "async_load", False, path), + async_load_workers=int(raw.get("async_load_workers", 1)), + soft_pin_ttl_ms=( + int(raw["soft_pin_ttl_ms"]) if raw.get("soft_pin_ttl_ms") is not None else None + ), + group_blocks=_json_bool(raw, "group_blocks", False, path), + stats_interval=( + float(raw["stats_interval"]) if raw.get("stats_interval") is not None else None + ), ) @staticmethod diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py index b6ec205a85e8..f7dbb534b174 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py @@ -19,16 +19,27 @@ so a key can only be reused by a request whose prefix is byte-identical. The chain is seeded with a `ReuseScope` covering what the tokens do not record. +Under one-model speculative decoding the draft layers write the KV at a position +from the prompt tokens after it, so a block's hash then also covers the prompt +lookahead the draft reads past the block's end, and the chain is seeded apart +from the one a server without that lookahead uses. Pages a speculative server +writes are therefore never served to a server whose draft reads differently, +and the keys of a server without a draft are unchanged. + A key is `/`. The namespace pins down everything else -that decides what the stored bytes mean: the model, the shard that produced -them, the layer group inside that shard, the tokens each page holds and how many -bytes a page is. A change to any of those reads as a miss rather than as garbage. +that decides what the stored bytes mean: the model, the attention shard that +produced them, the layer group inside that shard, the tokens each page holds and +how many bytes a page is. A change to any of those reads as a miss rather than +as garbage. With `group_blocks`, the pages of one block also share an eviction +group named by `KeyNamespace.group`. """ import hashlib from dataclasses import dataclass, field from typing import List, Optional, Sequence, Tuple +import numpy as np + __all__ = [ "BlockHashChain", "KeyNamespace", @@ -88,13 +99,47 @@ class BlockHashChain: Extended in place as the token list grows, so a generation step costs one digest per newly completed block rather than a rehash of the prompt. + + With a prompt lookahead of `L` tokens, block `b` is hashed over its own + tokens and the `L` that follow it, `tokens[b * T : (b + 1) * T + L]` for `T` + tokens per block, so a block exists only once the `L` tokens after it do. """ - def __init__(self, tokens_per_block: int, scope: Optional[ReuseScope] = None): + def __init__( + self, + tokens_per_block: int, + scope: Optional[ReuseScope] = None, + *, + prompt_lookahead: int = 0, + ): + """Start an empty chain. + + Args: + tokens_per_block: Tokens each block, and so each hash, covers. + scope: What besides the tokens decides whose KV a hash names. + prompt_lookahead: Prompt tokens past a block's end that its KV + depends on: the `draft_prompt_lookahead` of a one-model draft + sharing the KV cache, 0 without one. A non-zero value seeds the + chain apart from the lookahead-0 one, whose keys it leaves as + they are. + + Raises: + ValueError: When `tokens_per_block` is not positive or + `prompt_lookahead` is negative. + """ if tokens_per_block <= 0: raise ValueError(f"tokens_per_block must be > 0, got {tokens_per_block}") + if prompt_lookahead < 0: + raise ValueError(f"prompt_lookahead must be >= 0, got {prompt_lookahead}") self._tokens_per_block = int(tokens_per_block) + self._prompt_lookahead = int(prompt_lookahead) self._seed = (scope or ReuseScope()).seed() + if self._prompt_lookahead: + # Mixed in only when non-zero, so every lookahead-0 key stays the one + # pools already hold. + self._seed = _digest( + self._seed, b"prompt-lookahead", self._prompt_lookahead.to_bytes(8, "little") + ) self._hashes: List[bytes] = [] @property @@ -102,6 +147,11 @@ def tokens_per_block(self) -> int: """Tokens covered by each entry in the chain.""" return self._tokens_per_block + @property + def prompt_lookahead(self) -> int: + """Tokens past each block's end that its hash also covers.""" + return self._prompt_lookahead + @property def hashes(self) -> Sequence[bytes]: """Hashes computed so far, indexed by block ordinal.""" @@ -110,6 +160,8 @@ def hashes(self) -> Sequence[bytes]: def extend(self, tokens: Sequence[int]) -> Sequence[bytes]: """Grow the chain to cover every full block of `tokens`. + A block counts as full only once its prompt lookahead is there too. + Args: tokens: The request's complete token list, prompt first. Must extend what was passed previously; a shorter list means the @@ -118,20 +170,39 @@ def extend(self, tokens: Sequence[int]) -> Sequence[bytes]: Returns: The full chain, indexed by block ordinal. """ - num_full_blocks = len(tokens) // self._tokens_per_block - if num_full_blocks < len(self._hashes): + tokens_per_block = self._tokens_per_block + lookahead = self._prompt_lookahead + num_full_blocks = max(0, len(tokens) - lookahead) // tokens_per_block + first_new = len(self._hashes) + if num_full_blocks < first_new: raise ValueError( - f"token list shrank from {len(self._hashes)} to {num_full_blocks} " + f"token list shrank from {first_new} to {num_full_blocks} " "full blocks; a hash chain belongs to exactly one request" ) - for ordinal in range(len(self._hashes), num_full_blocks): - start = ordinal * self._tokens_per_block - block = tokens[start : start + self._tokens_per_block] - parent = self._hashes[-1] if self._hashes else self._seed - # Fixed-width little-endian token ids: a delimiter-free encoding - # would let two different token sequences serialize identically. - payload = b"".join(int(token).to_bytes(8, "little", signed=True) for token in block) - self._hashes.append(_digest(parent, payload)) + if num_full_blocks == first_new: + return self._hashes + # Fixed-width little-endian token ids: a delimiter-free encoding would + # let two different token sequences serialize identically. ' str: def key(self, block_hash: bytes) -> str: """The store key holding one page of this namespace.""" return f"{self.prefix}/{block_hash.hex()}" + + def group(self, block_hash: bytes) -> str: + """The Mooncake eviction group every page of one block is put into. + + It leaves out the attention shard, the layer group, the tokens per + block and the page size, so every page a hit on this block needs (all + layer groups, all tensor-parallel shards) shares one group and the + master evicts them together. It keeps the shard count, so deployments + sharded differently never share a group. + + Args: + block_hash: The block's hash, as its keys carry it. + + Returns: + The group id to put the block's pages under. + """ + return f"{self.namespace}/{self.model_key}/w{self.world_size}/{block_hash.hex()}" diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py index c946c97b1682..29fe0a087d0f 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py @@ -423,6 +423,14 @@ def _client_config( "role": StoreRole(pool.role).value, "transfer_batch_size": pool.transfer_batch_size, "stage_through_host": pool.stage_through_host, + "staging_buffer_bytes": parse_size(pool.staging_buffer_bytes, strict_units=True), + "async_load": pool.async_load, + "async_load_workers": pool.async_load_workers, + "group_blocks": pool.group_blocks, + # Omitted when off, which is what a reader that does not know the key + # assumes too. + **({"soft_pin_ttl_ms": pool.soft_pin_ttl_ms} if pool.soft_pin_ttl_ms is not None else {}), + **({"stats_interval": pool.stats_interval} if pool.stats_interval is not None else {}), } return config diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py index f1eb2db7ee14..5ad3121c0c1c 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py @@ -14,10 +14,12 @@ # limitations under the License. """The per-iteration work list the scheduler hands the workers. -Instances are broadcast from rank 0 to every worker, so these carry only plain -data: a page's identity (its block hash) and where that page currently lives on -this rank (a layer group and a page slot index). Deliberately no store keys: -each worker prefixes its own rank namespace, so one broadcast serves all shards. +Built by the scheduler adapter, then broadcast from rank 0 to every worker +under tensor parallelism, or bound locally on its owner under attention data +parallelism. These carry only plain data: a page's identity (its block hash) and +where that page currently lives on this rank (a layer group and a page slot +index). Deliberately no store keys: each worker prefixes its own attention-shard +namespace, so one work list serves every shard. """ from dataclasses import dataclass, field @@ -52,11 +54,19 @@ class RequestTransfers: @dataclass class MooncakeStoreMetadata: - """Loads to perform before the next forward pass, saves to start after it.""" + """Loads to perform before the next forward pass, saves to start after it. + + With `async_load` the scheduler hands parked requests' loads to the + workers' background load threads instead, in a third list. + """ loads: List[RequestTransfers] = field(default_factory=list) saves: List[RequestTransfers] = field(default_factory=list) + #: Asynchronous loads to queue when the metadata is bound, after the KV + #: cache manager's allocation fence. Under TP the scheduler exists on rank 0 + #: only and every rank loads its own shard of the parked request's pages. + async_loads: List[RequestTransfers] = field(default_factory=list) def __bool__(self) -> bool: """Whether there is any work at all this iteration.""" - return bool(self.loads or self.saves) + return bool(self.loads or self.saves or self.async_loads) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py index 0bda85e66232..744be81b630f 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py @@ -12,16 +12,32 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Leader side of the Mooncake store KV cache connector. +"""Scheduler adapter side of the Mooncake store KV cache connector. -Runs only on rank 0, deciding what to load and what to save while the workers -do the moving. That needs two pieces of per-request bookkeeping: +Runs on rank 0 under tensor parallelism and on every owner under attention data +parallelism, deciding what to load and what to save for its own requests while +the workers do the moving. Each adapter asks the worker in its own process to +look prefixes up. That needs two pieces of per-request bookkeeping: * a hash chain, since `KVCacheManagerV2` reports `RequestData.block_hashes` empty and a block has no content identity otherwise; * the page slot index per block ordinal, accumulated across iterations. The manager reports only newly allocated indices each step, but a block is allocated before it is full and is savable only once it is full. + +Prompt hashes and pages both exist before the KV they name is computed, so a +block is published only by the pass that computes its last token, and it +counts as stored only once the worker's `wait_for_save` queued that pass. + +Next to a one-model draft on a disaggregated context server, a block's hash +also covers the prompt tokens the draft reads past the block, so a block whose +lookahead runs past the prompt is neither offered nor saved. + +With `async_load`, every hit is offered asynchronously: the runtime parks the +request outside the batch, and the transfer is computed from the request's full +page lists when its allocation is reported, then handed to the workers in the +same iteration's metadata, which they bind after the KV cache manager's +allocation fence. """ from typing import Dict, List, Optional, Tuple @@ -104,6 +120,8 @@ class _RequestState: "saved_upto", "load_first_block", "load_blocks", + "load_async", + "async_loaded_upto", "emitted_saves", ) @@ -113,29 +131,84 @@ def __init__(self, chain: BlockHashChain): self.tokens: List[int] = [] #: Page slot index per block ordinal, per layer group. self.pages: Dict[int, List[int]] = {} - #: First block ordinal not yet considered for saving. + #: First block ordinal not yet handed to the store. Blocks a pass bound + #: for saving count only once the worker queued them; see + #: `MooncakeStoreConnectorScheduler._confirm_submitted_saves`. self.saved_upto = 0 #: The offer made by `get_num_new_matched_tokens`, in block ordinals. self.load_first_block = 0 self.load_blocks = 0 + #: Whether that offer is asynchronous, which decides how + #: `update_state_after_alloc_by_layer_group` settles it. + self.load_async = False + #: End of the offer handed to the load workers, in block ordinals. + #: Those blocks are the store's already, so they are never written + #: back. + self.async_loaded_upto = 0 self.emitted_saves = False class MooncakeStoreConnectorScheduler(KvCacheConnectorScheduler): """Chooses which pages the Mooncake pool serves and which it receives.""" + #: Request state is owner-local and lookups go to the worker in this + #: process, whose attention-DP keys every owner shares. + supports_attention_dp = True + #: Block keys cover a one-model draft's prompt lookahead, offers stop at the + #: last block whose lookahead the prompt holds, and saves publish only the + #: blocks a pass computed. + supports_speculative_prefill = True + #: A query only resets and refills this request's offer state + #: (`get_num_new_matched_tokens`), the pool lookup is a read-only + #: `count_prefix_hit`, and transfers are built only once the allocation is + #: reported (`update_state_after_alloc_by_layer_group`) or the batch is + #: final (`build_connector_meta`), so an admission-time answer has no other + #: effect and a repeated query replaces it. `request_reset` drops the state + #: of a request whose first chunk gave its pages back. + supports_admission_query = True + def __init__(self, llm_args: TorchLlmArgs): super().__init__(llm_args) validate_llm_args(llm_args) self._config = MooncakeStoreConnectorConfig.resolve(llm_args) self._tokens_per_block = int(llm_args.kv_cache_config.tokens_per_block) + # Imported here, as `KVCacheManagerV2` does: the speculative package + # pulls in every drafter, and processes that import this package without + # serving, such as `trtllm-serve mooncake_master`, need none of them. + from tensorrt_llm._torch.speculative import draft_prompt_lookahead + + spec_config = getattr(llm_args, "speculative_config", None) + if spec_config is not None and spec_config.max_draft_len is None: + # Vanilla MTP's lookahead is its draft length, so an unresolved one + # would key draft KV as if the draft read nothing past a block. + raise RuntimeError( + "mooncake-store must be built after the speculative config is " + "resolved from the checkpoint: max_draft_len is still unset, so " + "the prompt lookahead its block keys cover is unknown." + ) + # Prompt tokens past a block that the draft's KV in the block depends + # on: 0 without a draft. The worker refuses every drafter except + # one-model MTP and EAGLE3 unless this server transfers nothing, so the + # `None` of an unestablished span never keys a page. + self._prompt_lookahead = draft_prompt_lookahead(spec_config) or 0 self._requests: Dict[int, _RequestState] = {} self._worker: Optional[MooncakeStoreConnectorWorker] = None + # Saves bound in the last metadata, by request: the block ordinal they + # end at. Settled by `_confirm_submitted_saves` at the next build. + self._unconfirmed_saves: Dict[int, int] = {} + # Asynchronous offers park the request. Their transfers travel in the + # next metadata, which every worker binds after the KV cache manager has + # published the allocation fence for the batch. Under tensor parallelism + # this adapter exists on rank 0 only and every rank loads its own shard; + # under attention DP each owner loads its own pages. + self._async_load = self._config.async_load and self._config.role.loads + self._pending_async_loads: List[RequestTransfers] = [] logger.warning( - f"mooncake-store leader ready (role={self._config.role.value}, " - f"tokens_per_block={self._tokens_per_block})" + f"mooncake-store scheduler adapter ready (role={self._config.role.value}, " + f"tokens_per_block={self._tokens_per_block}, " + f"prompt_lookahead={self._prompt_lookahead}, async_load={self._async_load})" ) def wait_for_initialization(self): @@ -149,12 +222,18 @@ def get_num_new_matched_tokens( ) -> Tuple[int, bool]: """Offer the longest stored prefix beyond what the device already has. + Only whole blocks are offered, starting at the local match and stopping + short of the prompt's last token. Under a one-model draft a block whose + prompt lookahead runs past the prompt has no hash, so it is neither + offered nor, later, saved. + Args: request: The request being scheduled. num_computed_tokens: Tokens already matched in the local KV cache. Returns: - Tokens the store can supply, and `False` for a synchronous load. + Tokens the store can supply, and whether they arrive asynchronously + (every hit when `async_load` is on). A miss is never asynchronous. """ if self._config.capacity_only: # Neither a lookup nor a save will read the hash chain this would @@ -170,8 +249,12 @@ def get_num_new_matched_tokens( tokens = request.get_tokens(0) state = self._state_for(request, tokens, scope) + # Every query replaces the outstanding offer, which is what makes a + # repeated admission-time query harmless. state.load_first_block = 0 state.load_blocks = 0 + state.load_async = False + state.async_loaded_upto = 0 if not self._config.role.loads: return 0, False @@ -197,26 +280,111 @@ def get_num_new_matched_tokens( state.load_first_block = first_block state.load_blocks = hit_blocks + state.load_async = self._async_load logger.debug( f"mooncake-store matched {hit_blocks} blocks " f"({hit_blocks * self._tokens_per_block} tokens) " f"for request {request.request_id}" + f"{' asynchronously' if state.load_async else ''}" ) - return hit_blocks * self._tokens_per_block, False + return hit_blocks * self._tokens_per_block, state.load_async def update_state_after_alloc(self, request: LlmRequest, block_ids: List[int]): - """No-op: page indices are read from the scheduler output instead. + """Settle an offer against a flat page list, as one layer group. - The flat `block_ids` here are a single space, but a V2 page index is - scoped to a layer group, so this connector reads only - `RequestData.new_block_ids_by_layer_group`. + The runtime uses the flat form only when it has no per-group view, so + `block_ids` is the single layer group's list, by block ordinal. The + work happens in `update_state_after_alloc_by_layer_group`; the base + class would otherwise fold that form back into this one. + + Args: + request: The request that was allocated resources. + block_ids: Its page slot indices, by block ordinal. """ + self.update_state_after_alloc_by_layer_group(request, [list(block_ids)]) + + def update_state_after_alloc_by_layer_group( + self, request: LlmRequest, block_ids_by_layer_group: List[List[int]] + ) -> None: + """Trim the outstanding offer to what the runtime honoured, and queue an asynchronous one. + + The runtime calls this once per allocation, after it committed the + honoured count. A synchronous offer is only trimmed: its pages are read + from the scheduler output in `build_connector_meta`, the same + iteration. An asynchronous offer parks the request, which + `build_scheduler_output` then skips until its load lands, so its + transfer is computed here from the lists, which hold every page of the + request by block ordinal. It is handed to the workers in the next + `build_connector_meta`, in the same iteration, and an offer honoured + with no tokens still queues an empty transfer, because the runtime + parks the request regardless and only the worker releases it. + + Args: + request: The request that was allocated resources. Its + `py_num_connector_matched_tokens` is the honoured part of the + offer, which may be shorter than the offer when this + iteration's pages cannot hold all of it. + block_ids_by_layer_group: Page slot indices per layer group, + indexed by layer group id, every page of the request by block + ordinal. + + Raises: + RuntimeError: When the honoured count is not a whole-block part of + the offer, or an honoured block of an asynchronous offer has no + valid destination page in some layer group. The runtime counts + honoured tokens as computed, so an incomplete load must never + be reported as finished. + """ + state = self._requests.get(request.request_id) + if state is None or state.load_blocks == 0: + # A bypassed request, a miss, or a role that does not read. + return + self._settle_offer(request) + if not state.load_async: + return + + transfers = RequestTransfers(request.request_id) + limit = min( + len(state.chain.hashes), + min((len(indices) for indices in block_ids_by_layer_group), default=0), + ) + load_end = state.load_first_block + state.load_blocks + for block in range(state.load_first_block, load_end): + if block >= limit: + raise RuntimeError( + f"mooncake-store: honoured prefix block {block} of request " + f"{request.request_id} has missing destination pages ({limit} " + "addressable); cannot report an incomplete asynchronous load " + "as finished" + ) + block_hash = state.chain.hashes[block] + for layer_group_id, indices in enumerate(block_ids_by_layer_group): + page_index = int(indices[block]) + if page_index == BAD_PAGE_INDEX: + raise RuntimeError( + f"mooncake-store: honoured prefix block {block} of request " + f"{request.request_id} has an invalid destination page in " + f"layer group {layer_group_id}; cannot report an incomplete " + "asynchronous load as finished" + ) + transfers.pages.append(PageTransfer(block_hash, layer_group_id, page_index)) + # Consumed here rather than in `build_connector_meta`: the request is not + # in the scheduler output while it is parked, and must not be loaded a + # second time when it comes back. + state.async_loaded_upto = load_end + state.load_blocks = 0 + self._pending_async_loads.append(transfers) # ---- work lists ---- def build_connector_meta(self, scheduler_output: SchedulerOutput) -> MooncakeStoreMetadata: - """Turn this iteration's scheduled requests into load and save lists.""" + """Turn this iteration's scheduled requests into load and save lists. + + The asynchronous loads queued since the last build travel too, once. + """ + self._confirm_submitted_saves() metadata = MooncakeStoreMetadata() + metadata.async_loads, self._pending_async_loads = self._pending_async_loads, [] for request_data in (*scheduler_output.new_requests, *scheduler_output.cached_requests): state = self._requests.get(request_data.request_id) if state is None: @@ -231,30 +399,46 @@ def build_connector_meta(self, scheduler_output: SchedulerOutput) -> MooncakeSto loads = self._loads_for(state, request_data) if loads.pages: metadata.loads.append(loads) + # The honoured restore, which this iteration's `computed_position` + # leaves out; see `_saves_for`. + num_loaded_tokens = state.load_blocks * self._tokens_per_block # Whatever the store just supplied, and whatever the local cache # matched, is not ours to write back: the store already has the - # former, and the latter was never allocated during this run. - state.saved_upto = max(state.saved_upto, state.load_first_block + state.load_blocks) + # former, and the latter was never allocated during this run. An + # asynchronous offer was consumed when its load was queued, and the + # parked request comes back with all of its pages, so its end is + # remembered separately. + state.saved_upto = max( + state.saved_upto, + state.load_first_block + state.load_blocks, + state.async_loaded_upto, + ) # An offer is consumed once. The load is issued in exactly the # iteration the runtime allocated pages to hold it. state.load_blocks = 0 if self._config.role.saves: - saves = self._saves_for(state, request_data) + saves, saves_end = self._saves_for(state, request_data, num_loaded_tokens) if saves.pages: state.emitted_saves = True metadata.saves.append(saves) + self._unconfirmed_saves[request_data.request_id] = saves_end + else: + # Nothing publishable in the range, so nothing to confirm. + state.saved_upto = max(state.saved_upto, saves_end) return metadata def request_finished(self, request: LlmRequest, cache_block_ids: List[int]) -> bool: """Report whether pages must stay pinned for in-flight saves. - Whether a write has landed is known per rank and this runs only on the - leader, so the answer is "any save was emitted" rather than "a save is - in flight". A request whose writes already retired is reported done by - the next `get_finished`, costing one iteration; the exact answer would - cost a per-rank report back to the leader. + The answer is "any save was emitted" rather than "a save is in + flight". Under tensor parallelism only rank 0's adapter answers while + write completion is known per rank; under attention DP the adapter and + its worker share the owner. Either way a request whose writes already + retired, or whose emitted saves were never queued because their pass + did not run, is reported done by the next `get_finished`, costing one + iteration. Returns: True when this request handed any page to the background save @@ -268,12 +452,74 @@ def request_reset(self, request: LlmRequest) -> None: """Forget a request whose pages were released for it to run again. `pages` and `saved_upto` describe the allocation that was just freed, - and the lookup preceding the next attempt rebuilds the rest. + and the lookup preceding the next attempt rebuilds the rest. Saves still + awaiting confirmation belong to that allocation too, so a report about + them cannot advance the next attempt's save mark. So does an + asynchronous load not yet handed to the workers; the runtime refuses a + reset while a load owns the allocation, so this is defensive. """ self._requests.pop(request.request_id, None) + self._unconfirmed_saves.pop(request.request_id, None) + self._pending_async_loads = [ + transfers + for transfers in self._pending_async_loads + if transfers.request_id != request.request_id + ] # ---- internals ---- + def _settle_offer(self, request: LlmRequest) -> None: + """Trim the request's outstanding offer to the whole blocks the runtime honoured. + + Both the load and the computed-block bound of this pass's saves then + use the honoured count rather than the offer. The runtime sets + `py_num_connector_matched_tokens` when it commits the honoured part of + the answer, before it reports the allocation, and it only ever honours + a whole-block prefix of what was offered. + + Raises: + RuntimeError: When the honoured count is not a whole-block part of + the offer. Loading it anyway would leave the runtime treating + tokens as computed that no page transfer fills. + """ + state = self._requests.get(request.request_id) + if state is None or state.load_blocks == 0: + return + honoured = int(request.py_num_connector_matched_tokens) + blocks, remainder = divmod(honoured, self._tokens_per_block) + if remainder or not 0 <= blocks <= state.load_blocks: + raise RuntimeError( + f"mooncake-store: the runtime honoured {honoured} tokens of request " + f"{request.request_id}, which is not a whole-block part of the " + f"{state.load_blocks}-block offer" + ) + state.load_blocks = blocks + + def _confirm_submitted_saves(self) -> None: + """Count the last bound saves as stored once the worker queued them. + + Bound metadata may never reach `wait_for_save`: the iteration can fail + between `handle_metadata` and the forward pass, and under attention + data parallelism or with asynchronous loads, a fleet vote, a + same-iteration rebuild or a batch the connector emptied can supersede + it. Such blocks were not computed by a pass that ran, so they are + emitted again when the request is next scheduled instead of being lost. + When the executor cannot queue a batch nothing is bound in the first + place, because only `prepare_resources` publishes a scheduler output. + + The worker's report is taken on every build, even with nothing + unconfirmed, so a report about a request reset since cannot confirm + that request's next pass if that pass never runs. + """ + if not self._config.role.saves: + return + submitted = self._require_worker().take_submitted_saves() + for request_id, saves_end in self._unconfirmed_saves.items(): + state = self._requests.get(request_id) + if state is not None and request_id in submitted: + state.saved_upto = max(state.saved_upto, saves_end) + self._unconfirmed_saves.clear() + def _require_worker(self) -> MooncakeStoreConnectorWorker: if self._worker is None: self._worker = resolve_local_worker() @@ -284,7 +530,11 @@ def _state_for( ) -> _RequestState: state = self._requests.get(request.request_id) if state is None: - state = _RequestState(BlockHashChain(self._tokens_per_block, scope=scope)) + state = _RequestState( + BlockHashChain( + self._tokens_per_block, scope=scope, prompt_lookahead=self._prompt_lookahead + ) + ) self._requests[request.request_id] = state # The lookup runs before the request is scheduled, so the chain has to # be ready before any scheduler output exists. @@ -308,25 +558,59 @@ def _addressable_blocks(self, state: _RequestState) -> int: return min(len(state.chain.hashes), min(len(indices) for indices in state.pages.values())) def _loads_for(self, state: _RequestState, request_data: RequestData) -> RequestTransfers: + """The honoured part of the offer, one transfer per block. + + Raises: + RuntimeError: When an honoured block has no destination page. The + runtime counts honoured tokens as computed, so skipping the + block would leave it to attend to KV nothing wrote. + """ transfers = RequestTransfers(request_data.request_id) limit = self._addressable_blocks(state) for offset in range(state.load_blocks): block = state.load_first_block + offset if block >= limit: - # The runtime allocated fewer pages than it accepted tokens - # for, so the tail of the offer has nowhere to land. Dropping - # it here leaves the runtime to recompute it. - break + raise RuntimeError( + f"mooncake-store: honoured prefix block {block} of request " + f"{request_data.request_id} has no destination page " + f"({limit} addressable)" + ) self._append_pages(state, transfers, block) return transfers - def _saves_for(self, state: _RequestState, request_data: RequestData) -> RequestTransfers: + def _saves_for( + self, state: _RequestState, request_data: RequestData, num_loaded_tokens: int + ) -> Tuple[RequestTransfers, int]: + """This pass's newly computed full blocks, and the block ordinal they end at. + + Prompt hashes exist before prefill, and pages can run ahead of + computation: the post-allocation token-budget trim shrinks a context + chunk whose pages were already allocated, and a served prefix grows + capacity past the chunk by any extra KV tokens a one-model draft needs. + The worker saves after this pass, so only the full blocks it completes + are publishable. + + Args: + state: The request's bookkeeping. + request_data: This iteration's scheduler output for the request. + Its `computed_position` excludes this iteration's connector + restore. + num_loaded_tokens: The honoured restore of this iteration, added + back to `computed_position`. + + Returns: + The pages to save, and the block ordinal the save mark may advance + to. The mark itself is left alone; the caller advances it once the + saves are known to have been handed to the store. + """ transfers = RequestTransfers(request_data.request_id) - limit = self._addressable_blocks(state) + computed_end = ( + request_data.computed_position + num_loaded_tokens + request_data.num_scheduled_tokens + ) + limit = min(self._addressable_blocks(state), computed_end // self._tokens_per_block) for block in range(state.saved_upto, limit): self._append_pages(state, transfers, block) - state.saved_upto = max(state.saved_upto, limit) - return transfers + return transfers, max(state.saved_upto, limit) def _append_pages(self, state: _RequestState, transfers: RequestTransfers, block: int) -> None: """Add one block's page from every layer group, or none of them.""" diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/settings.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/settings.py index 68c4244d0aed..93a6d734e103 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/settings.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/settings.py @@ -188,6 +188,20 @@ def _restate_pool_block( ("model_key", effective.model_key, pool.model_key), ("transfer_batch_size", effective.transfer_batch_size, pool.transfer_batch_size), ("stage_through_host", effective.stage_through_host, pool.stage_through_host), + ( + "staging_buffer_bytes", + effective.staging_buffer_bytes, + parse_size(pool.staging_buffer_bytes, strict_units=True), + ), + ("async_load", effective.async_load, pool.async_load), + ("async_load_workers", effective.async_load_workers, pool.async_load_workers), + ("group_blocks", effective.group_blocks, pool.group_blocks), + ) + # Settings where `None` means off rather than "not stated", so a config + # that leaves one out restates the block to off. + nullable = ( + ("soft_pin_ttl_ms", effective.soft_pin_ttl_ms, pool.soft_pin_ttl_ms), + ("stats_interval", effective.stats_interval, pool.stats_interval), ) # Assignment adds to `model_fields_set`, so what the deployment stated for # itself has to be read before any of it is replaced. Only a value it chose @@ -200,6 +214,12 @@ def _restate_pool_block( setattr(pool, setting, value) if setting in stated: overridden[setting] = value + for setting, value, asked in nullable: + if value == asked: + continue + setattr(pool, setting, value) + if setting in stated: + overridden[setting] = value if overridden: logger.warning( f"mooncake-store: {path} states " diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py index 4c1827a6020c..6adccfa37ae3 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py @@ -43,12 +43,7 @@ from tensorrt_llm._utils import CUASSERT from tensorrt_llm.logger import logger -__all__ = ["MAX_STAGING_BUFFER_BYTES", "HostStagingPool", "plan_slot_geometry", "sync_stream"] - -#: Ceiling on the pinned allocation per direction, bounding what a large page -#: times a large batch could demand. The memory is pinned for the process's -#: lifetime and comes out of the DRAM the pool segment is charged against. -MAX_STAGING_BUFFER_BYTES = 1 << 30 +__all__ = ["HostStagingPool", "plan_slot_geometry", "sync_stream"] #: Stated explicitly rather than inferred from the pointers, which would be #: wrong for a host pointer outside the unified address space. @@ -82,19 +77,22 @@ def sync_stream(stream: int) -> None: def plan_slot_geometry( max_bytes_per_page: int, transfer_batch_size: int, - budget_bytes: int = MAX_STAGING_BUFFER_BYTES, + budget_bytes: int, ) -> Tuple[int, int]: """Choose how many pages may be staged at once, and how wide a slot is. - The geometry follows from the layout rather than from configuration: a slot - has to hold the largest page any layer group produces, and there is no - point holding more slots than a transfer batch has pages. A budget below - one page is raised to one rather than refused. + A slot has to hold the largest page any layer group produces, and there is + no point holding more slots than a transfer batch has pages. The budget + bounds the count, never the width: a budget below one page yields one + slot, which pins more than the budget rather than refusing to start. Args: max_bytes_per_page: Largest page payload across layer groups. transfer_batch_size: Pages the connector puts in one store call. - budget_bytes: Ceiling on this pool's pinned allocation. + budget_bytes: Ceiling on this pool's pinned allocation: the configured + `staging_buffer_bytes`, or one asynchronous load worker's share of + it. The memory is pinned for the process's lifetime and comes out + of the DRAM the pool segment is charged against. Returns: Slot width in bytes, and the number of slots. @@ -112,9 +110,10 @@ def plan_slot_geometry( class HostStagingPool: """A registered pinned buffer, sliced into per-page slots. - One pool serves one direction. Loads run on the executor thread and saves - on the connector's background thread, so sharing slots would need a lock on - the transfer path for no benefit. + One pool per user of the slots: the synchronous load path (executor + thread, only when `async_load` is off), the save thread, and each + asynchronous load worker. Slots are never shared between threads, so no + lock sits on the transfer path. """ def __init__( diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py index 0abca8110e6d..9a403c2a0295 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py @@ -16,8 +16,10 @@ Every rejection here is a configuration whose failure mode is a wrong answer rather than a slow one: KV replayed without all of the state it was computed -with. Beam search, attention data parallelism and Mamba caches are rejected for -all connectors in `py_executor`, so they are not repeated here. +with. Beam search and Mamba caches are rejected for all connectors in +`py_executor`, and so is speculative decoding outside the speculative-prefill-only +mode, so they are not repeated here. What is left for speculation is the +drafter: only one whose KV this connector keys and moves may run next to it. Checks run at construction, before any request is admitted, so a bad deployment fails at startup instead of after the first cache hit. @@ -33,7 +35,12 @@ from ..kv_cache_layout import KvCacheLayout from .config import MooncakeStoreConnectorConfig -__all__ = ["validate_layout", "validate_llm_args", "validate_node_budget"] +__all__ = [ + "validate_layout", + "validate_llm_args", + "validate_node_budget", + "validate_speculative_config", +] _GIB = 1 << 30 @@ -75,6 +82,65 @@ def validate_llm_args(llm_args: TorchLlmArgs) -> None: ) +def validate_speculative_config( + llm_args: TorchLlmArgs, config: "MooncakeStoreConnectorConfig" +) -> None: + """Reject a drafter whose KV this connector cannot key or restore. + + Only one-model MTP and EAGLE3 keep their draft KV in the target pool, which + is registered with the connector, with an established prompt lookahead the + block keys cover. Any other drafter's state would be missing or stale after + a restored prefix. Runs before the store is opened, since the cache managers + that would otherwise refuse a separate draft KV cache are built after the + connector. + + Args: + llm_args: The engine arguments. + config: The resolved connector configuration. A capacity-only role + addresses no page, so its drafter does not matter. + + Raises: + NotImplementedError: When a server that transfers KV runs a drafter + other than one-model MTP or EAGLE3, or keeps the draft layers in a + separate KV cache. + """ + spec_config = getattr(llm_args, "speculative_config", None) + if spec_config is None or config.capacity_only: + return + # Imported here: the speculative package pulls in every drafter, which a + # process that imports this package without serving does not need. + from tensorrt_llm._torch.speculative import ( + draft_prompt_lookahead, + should_use_separate_draft_kv_cache, + ) + + mode = spec_config.spec_dec_mode + if not (mode.is_mtp_one_model() or mode.is_eagle3_one_model()): + raise NotImplementedError( + f"The mooncake-store connector does not support speculative decoding " + f"mode {mode.name}. Only one-model MTP and EAGLE3 keep their draft KV " + "in the target KV cache with an established prompt lookahead, which " + "the block keys cover; any other drafter's state would be missing or " + "stale after a restored prefix. Run one-model MTP or EAGLE3 on a " + "disaggregated context server, or use role 'capacity'." + ) + if draft_prompt_lookahead(spec_config) is None: + # Both admitted modes define it; kept so a future change cannot key + # draft KV with a lookahead nobody established. + raise NotImplementedError( + f"The mooncake-store connector needs the prompt lookahead of " + f"speculative decoding mode {mode.name}, which is not established." + ) + if should_use_separate_draft_kv_cache(spec_config): + raise NotImplementedError( + "The mooncake-store connector does not support a separate draft KV " + "cache: the draft layers would live in a cache the connector never " + "sees, so a restored prefix would carry stale draft KV. Run on a " + "disaggregated context server, where one-model draft layers share " + "the target KV cache." + ) + + def _available_host_memory() -> Optional[int]: """Host memory the kernel says is available, or `None` if unknowable. @@ -132,11 +198,14 @@ def validate_node_budget( # DRAM, so it belongs in the sum even when no segment is mounted. staging_claim = 0 if config.stage_through_host: - from .staging import MAX_STAGING_BUFFER_BYTES - - # One pool per direction, and a capacity-only rank opens neither. + # One budget per direction, and a capacity-only rank opens neither. + # `async_load` adds nothing: the load direction is either the + # synchronous load pool or the load workers' pools splitting the same + # budget, never both (see `MooncakeStoreConnectorWorker._open_staging`). + # A pool raised to one page because its share is below one page is + # warned about when it is allocated, not counted here. directions = int(config.role.loads) + int(config.role.saves) - staging_claim = ranks_on_node * directions * MAX_STAGING_BUFFER_BYTES + staging_claim = ranks_on_node * directions * config.staging_buffer_bytes claimed = segment_claim + staging_claim if claimed <= 0: @@ -166,13 +235,20 @@ def validate_node_budget( if affordable > 0: remedy = ( f"Lower segment_size to at most {affordable / _GIB:.1f} GiB, or run " - f"fewer ranks per node." + "fewer ranks per node" + ) + # Naming the staging knobs only where staging is pinned at all. + remedy += ( + ", lower staging_buffer_bytes, or turn off stage_through_host." + if staging_claim + else "." ) else: remedy = ( f"Lending nothing would not be enough: the {staging_claim / _GIB:.1f} " f"GiB of pinned staging buffers alone overruns the budget. Run fewer " - f"ranks per node, or turn off stage_through_host." + f"ranks per node, lower staging_buffer_bytes, or turn off " + f"stage_through_host." ) raise ValueError( f"mooncake-store: this node cannot afford what the pool is configured " diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py index 312dae49c79d..f8c234150a3e 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py @@ -17,31 +17,63 @@ One worker per rank owns a `MooncakeDistributedStore` handle and moves pages between that pool and its own GPU KV cache. It is also the only place that knows how a page is addressed and how a key is spelled, which is why the -leader, colocated with rank 0's worker, asks it to run prefix lookups. - -Loads are synchronous: the runtime has already told the scheduler those tokens -are computed, so a failed load is a wrong answer rather than a slow one. +scheduler adapter in the same process asks it to run prefix lookups: rank 0's +under tensor parallelism, every owner's under attention data parallelism. + +A stored page is one attention shard's KV. Under tensor parallelism each rank +keys its own shard and a prefix counts only when every shard is present. Under +attention data parallelism every owner holds all heads, so all owners key their +pages as shard 0 of 1 and share one key space on the pool, while transfers, +completion and the store handle stay owner-local. + +Loads are synchronous by default, and the runtime has already counted the +loaded tokens as computed, so a page missing at load time must never reach a +forward pass. A rank that holds whole pages and its own batch hands the +affected request back for a restart without the offered prefix +(`reports_failed_loads`): a single-rank server or an attention-DP owner. Under +tensor parallelism the ranks hold shards of the same page and cannot restart a +request alone, so there a missing page fails the executor. + +With `async_load`, the runtime parks a request the pool serves outside the +batch instead, and `async_load_workers` background threads per rank move its +pages: each waits for the KV cache manager's allocation fence on its own +thread, then fetches, and `get_finished` releases the request once its pages +landed. A missing page is handed back or fails the executor as above. Any +exception on a load thread is fatal at the next executor call, and stays so. Saves are asynchronous and gated on a CUDA event, since the pages are complete only once the forward pass that wrote them has retired and blocking the executor loop on an RDMA write is the cost the store exists to avoid. The scheduler reports such a request as saving asynchronously, which keeps its -pages pinned until `get_finished` says the writes landed. A save that fails is -logged and dropped: a write that did not land costs a future cache miss and -nothing else, so the request retires either way. Only the save thread failing -to start is fatal, because that would strand every request behind pages -nothing will release. +pages pinned until `get_finished` says the writes landed. It counts a save as +stored once `wait_for_save` has queued it, so saves bound for a pass that never +ran are emitted again. A save that fails is logged and dropped: a write that +did not land costs a future cache miss and nothing else, so the request retires +either way. Only the save thread failing to start is fatal, because that would +strand every request behind pages nothing will release. With `soft_pin_ttl_ms` +or `group_blocks` set, each put also carries a Mooncake `ReplicateConfig` that +soft-pins its pages or names each page's eviction group. + +Next to one-model MTP or EAGLE3 on a disaggregated context server, the draft +layers share the target pool, so they move with every page. Any other drafter, +or draft layers in a separate KV cache, is refused at construction. A capacity-only worker opens its handle and stops there, with no layout, no buffer registration and no save thread, so a node can lend host memory to the pool without an HCA that can pin GPU pages. + +With `stats_interval` set, a role that reads the pool logs two info-level lines +per interval: lookups that stopped at a hole the pool evicted from a prefix it +otherwise still holds, and the time each phase of a load took. """ import threading +import time import traceback from collections import defaultdict +from contextlib import contextmanager from queue import Queue -from typing import Dict, List, Optional, Sequence, Set, Tuple +from typing import Dict, Iterable, List, Optional, Sequence, Set, Tuple import torch @@ -49,16 +81,21 @@ from tensorrt_llm.llmapi.llm_args import TorchLlmArgs from tensorrt_llm.logger import logger -from ..kv_cache_connector import KvCacheConnectorWorker +from ..kv_cache_connector import AllocationFence, KvCacheConnectorWorker from ..kv_cache_layout import KvCacheLayout from .addressing import PageAddressing -from .config import CONFIG_PATH_ENV, MooncakeStoreConnectorConfig, local_address, pool_config +from .config import ( + CONFIG_PATH_ENV, + MooncakeStoreConnectorConfig, + local_address, + local_reuse_backoff, + pool_config, +) from .gpudirect import PEERMEM_ENV, allocation_range, format_diagnosis, select_registration_path from .keys import KeyNamespace from .ledger import record_segment from .metadata import MooncakeStoreMetadata, RequestTransfers from .staging import ( - MAX_STAGING_BUFFER_BYTES, HostStagingPool, describe_batch_for_get, plan_slot_geometry, @@ -66,24 +103,37 @@ unstage_batch_after_get, ) from .staging import sync_stream as _sync_stream -from .validation import validate_layout, validate_llm_args, validate_node_budget +from .validation import ( + validate_layout, + validate_llm_args, + validate_node_budget, + validate_speculative_config, +) __all__ = ["MooncakeStoreConnectorWorker", "resolve_local_worker"] _GIB = 1 << 30 -#: Set by the worker's constructor so the leader, which the executor builds in -#: the same process on rank 0, can reach the store handle without a second +#: Set by the worker's constructor so the scheduler adapter, which the executor +#: builds in the same process (on rank 0 under tensor parallelism, on every +#: owner under attention DP), can reach the store handle without a second #: connection. `py_executor_creator` constructs the two concurrently, so the -#: leader waits on `_LOCAL_WORKER_READY`. +#: adapter waits on `_LOCAL_WORKER_READY`. _LOCAL_WORKER: Optional["MooncakeStoreConnectorWorker"] = None _LOCAL_WORKER_READY = threading.Event() -#: Seconds `shutdown` gives the save thread to drain its queue and return. -#: Sized for a backlog of pages over a congested fabric rather than a healthy -#: one, since what follows the wait is only safe once the thread has stopped. +#: Seconds `shutdown` gives the save thread and the load threads, against one +#: deadline, to drain and return. Sized for a backlog of pages over a congested +#: fabric rather than a healthy one, since what follows the wait is only safe +#: once every thread has stopped. SAVE_DRAIN_TIMEOUT = 30.0 +#: One parked request's load, as the load threads take it off their queue: the +#: request id, its page transfers, when it was queued, and the allocation fence +#: bound in the iteration that parked it. The fence is never None; see +#: `MooncakeStoreConnectorWorker._start_async_load`. +_QueuedLoad = Tuple[int, RequestTransfers, float, AllocationFence] + def resolve_local_worker(timeout: float = 60.0) -> "MooncakeStoreConnectorWorker": """The worker living in this process, once it has been constructed. @@ -93,9 +143,10 @@ def resolve_local_worker(timeout: float = 60.0) -> "MooncakeStoreConnectorWorker """ if not _LOCAL_WORKER_READY.wait(timeout): raise RuntimeError( - "The mooncake-store leader could not find a worker in its process. " - "The leader only runs on rank 0, where the executor also builds a " - "worker, so this means worker construction failed." + "The mooncake-store scheduler adapter could not find a worker in its " + "process. Each adapter (rank 0 under tensor parallelism, every owner " + "under attention DP) is built next to a worker, so this means worker " + "construction failed." ) assert _LOCAL_WORKER is not None return _LOCAL_WORKER @@ -174,14 +225,306 @@ def _stream_handle(stream: Optional[torch.cuda.Stream]) -> int: return int(stream.cuda_stream) +def _replicate_config(soft_pin_ttl_ms: Optional[int], group_ids: Optional[Sequence[str]] = None): + """The Mooncake `ReplicateConfig` for one put call, or None for the store's defaults. + + Args: + soft_pin_ttl_ms: Soft-pin every page of the call for this long, if set. + group_ids: One eviction group per key of the call, if set. + + Returns: + A `mooncake.store.ReplicateConfig`, or None when neither option is set, + in which case the put keeps its three-argument form. + + Raises: + ImportError: When the installed bindings predate either option. + """ + if soft_pin_ttl_ms is None and group_ids is None: + return None + try: + import mooncake.store as mooncake_store + + config = mooncake_store.ReplicateConfig() + if soft_pin_ttl_ms is not None: + config.soft_pin_action = mooncake_store.SoftPinAction.ENABLE + config.soft_pin_ttl_ms = soft_pin_ttl_ms + if group_ids is not None: + config.group_ids = list(group_ids) + except (ImportError, AttributeError) as exc: + raise ImportError( + "mooncake_store.soft_pin_ttl_ms and group_blocks need " + "mooncake-transfer-engine-cuda13 0.3.13 or newer, whose mooncake.store " + "exposes SoftPinAction and a ReplicateConfig with soft_pin_action, " + "soft_pin_ttl_ms and group_ids" + ) from exc + return config + + +#: The parts of a prefix load timed separately, in the order the stats line +#: lists them. The synchronous path runs them as: resolve keys, wait for the +#: allocation fence (zero copy only), describe each batch's destinations, fetch +#: it, then scatter and sync a staged batch. An asynchronous load waits for the +#: fence first, on its load thread, whether zero copy or staged. +_LOAD_PHASES = ("fence", "resolve", "describe", "get", "unstage", "sync") + + +@contextmanager +def _timed_phase(phases: Optional[Dict[str, List[float]]], name: str): + """Add the block's wall and thread CPU seconds to `phases[name]`. + + Args: + phases: Per-phase `[wall, cpu]` totals of one load, or None to time + nothing, which costs no clock read. + name: One of `_LOAD_PHASES`. + """ + if phases is None: + yield + return + wall, cpu = time.perf_counter(), time.thread_time() + try: + yield + finally: + totals = phases.setdefault(name, [0.0, 0.0]) + totals[0] += time.perf_counter() - wall + totals[1] += time.thread_time() - cpu + + +class _LoadStats: + """Prefix loads, summarised once per `stats_interval`. + + A load is one pass that moved pages: a `start_load_kv` call with pages to + fetch, or one parked request's asynchronous load. Every method takes the + lock, because asynchronous loads record from their load threads into the + same counters the executor thread logs. + """ + + __slots__ = ( + "_lock", + "loads", + "pages", + "batches", + "loaded_bytes", + "failed_pages", + "xfer_sum", + "xfer_max", + "waits", + "wait_sum", + "wait_max", + "phase_wall", + "phase_cpu", + ) + + def __init__(self) -> None: + self._lock = threading.Lock() + self._reset() + + def _reset(self) -> None: + self.loads = 0 + self.pages = 0 + self.batches = 0 + self.loaded_bytes = 0 + self.failed_pages = 0 + self.xfer_sum = 0.0 + self.xfer_max = 0.0 + self.waits = 0 + self.wait_sum = 0.0 + self.wait_max = 0.0 + self.phase_wall = dict.fromkeys(_LOAD_PHASES, 0.0) + self.phase_cpu = dict.fromkeys(_LOAD_PHASES, 0.0) + + def record( + self, + xfer_s: float, + pages: int, + batches: int, + loaded_bytes: int, + phases: Optional[Dict[str, List[float]]] = None, + *, + failed_pages: int = 0, + wait_s: Optional[float] = None, + ) -> None: + """Fold one load into the interval. + + Args: + xfer_s: Seconds from the destinations being ready to the last page + landing. + pages: Pages the load asked for. + batches: Store calls it took. + loaded_bytes: Bytes that landed. + phases: Its per-phase `[wall, cpu]` seconds, from `_timed_phase`. + failed_pages: Pages that did not land. + wait_s: For an asynchronous load, seconds from being queued to its + destinations being ready: the queue plus the fence. + """ + with self._lock: + self.loads += 1 + self.pages += pages + self.batches += batches + self.loaded_bytes += loaded_bytes + self.failed_pages += failed_pages + self.xfer_sum += xfer_s + self.xfer_max = max(self.xfer_max, xfer_s) + if wait_s is not None: + self.waits += 1 + self.wait_sum += wait_s + self.wait_max = max(self.wait_max, wait_s) + for name, (wall, cpu) in (phases or {}).items(): + self.phase_wall[name] = self.phase_wall.get(name, 0.0) + wall + self.phase_cpu[name] = self.phase_cpu.get(name, 0.0) + cpu + + def take_summary(self, *, workers: int = 0, queued: int = 0, active: int = 0) -> Optional[str]: + """The interval's summary, starting a new interval. + + Args: + workers: Asynchronous load threads, or 0 when loads are synchronous, + which leaves the asynchronous fields out. + queued: Asynchronous loads waiting for a load thread right now. + active: Asynchronous loads a load thread is running right now. + + Returns: + The summary, or None when nothing loaded and nothing is waiting. + """ + with self._lock: + if not (self.loads or queued or active): + return None + n = max(1, self.loads) + phases = " ".join( + f"{name}={1000 * self.phase_wall[name] / n:.1f}/" + f"{1000 * self.phase_cpu[name] / n:.1f}" + for name in _LOAD_PHASES + ) + waits = max(1, self.waits) + asynchronous = ( + f"wait_ms mean={1000 * self.wait_sum / waits:.1f} " + f"max={1000 * self.wait_max:.1f} " + f"workers={workers} queue={queued} active={active} " + if workers + else "" + ) + summary = ( + f"loads={self.loads} pages={self.pages} batches={self.batches} " + f"loaded_gib={self.loaded_bytes / _GIB:.2f} " + f"failed_pages={self.failed_pages} " + f"transfer_ms mean={1000 * self.xfer_sum / n:.1f} " + f"max={1000 * self.xfer_max:.1f} " + f"{asynchronous}" + f"phase_ms(wall/cpu per load) {phases}" + ) + self._reset() + return summary + + +class _PrefixHoleStats: + """Prefix lookups that stopped at a hole, summarised once per `stats_interval`. + + A hole is a missing block followed by stored blocks of the same chain: the + pool evicted a block from the middle of a chain it otherwise still holds. + The lookup stops at the hole, so the request recomputes every token behind + it, including the stored ones counted here as holed. + """ + + __slots__ = ( + "lookups", + "lookup_blocks", + "hit_blocks", + "holes", + "hole_blocks", + "holed_tokens", + "max_holed_tokens", + ) + + def __init__(self) -> None: + self.reset() + + def reset(self) -> None: + self.lookups = 0 + self.lookup_blocks = 0 + self.hit_blocks = 0 + self.holes = 0 + self.hole_blocks = 0 + self.holed_tokens = 0 + self.max_holed_tokens = 0 + + def record( + self, + num_blocks: int, + hit_blocks: int, + stored_after_hole: Sequence[int], + tokens_per_block: int, + ) -> None: + """Account for one lookup. + + Args: + num_blocks: Blocks looked up. + hit_blocks: Usable prefix length, which is also the hole's index. + stored_after_hole: Ascending indices of fully stored blocks after + the first missing one. + tokens_per_block: Tokens per block. + """ + self.lookups += 1 + self.lookup_blocks += num_blocks + self.hit_blocks += hit_blocks + if not stored_after_hole: + return + self.holes += 1 + span = stored_after_hole[-1] + 1 - hit_blocks + self.hole_blocks += span - len(stored_after_hole) + holed_tokens = len(stored_after_hole) * tokens_per_block + self.holed_tokens += holed_tokens + self.max_holed_tokens = max(self.max_holed_tokens, holed_tokens) + + def summary(self) -> str: + return ( + f"lookups={self.lookups} lookup_blocks={self.lookup_blocks} " + f"hit_blocks={self.hit_blocks} holes={self.holes} " + f"hole_blocks={self.hole_blocks} holed_tokens={self.holed_tokens} " + f"max_holed_tokens={self.max_holed_tokens}" + ) + + class MooncakeStoreConnectorWorker(KvCacheConnectorWorker): """Moves KV pages between this rank's GPU cache and the Mooncake pool.""" + #: Transfers, completion and the store handle are owner-local, and under + #: attention DP every owner keys its pages as attention shard 0 of 1, so all + #: owners share one key space on the pool. + supports_attention_dp = True + #: A page offered at lookup can be gone by load time, evicted from a full + #: pool or left behind by a put that never completed. On a single rank or an + #: attention-DP owner the affected requests are handed back to the executor + #: to restart (`take_failed_load_requests`). Tensor-parallel ranks hold + #: shards of one page and cannot restart a request alone, so there a failed + #: page stays fatal. + reports_failed_loads = True + #: Loads and saves move whole pages of every layer group in the registered + #: layout, which includes one-model draft layers sharing the target pool; + #: the per-layer hooks are no-ops. `validate_speculative_config` refuses a + #: drafter whose KV lives anywhere else. + supports_speculative_prefill = True + #: Registering pages drops the automatic host tier, and a prefix restore is + #: chosen only after the final batch, so on a disaggregated server the + #: pages each admitted request needs to finish have to stay free. + requires_completion_reservation = True + def __init__(self, llm_args: TorchLlmArgs): super().__init__(llm_args) validate_llm_args(llm_args) self._config = MooncakeStoreConnectorConfig.resolve(llm_args) + # Before the store is opened, so a drafter this connector cannot key + # fails before a segment is mounted. + validate_speculative_config(llm_args, self._config) + loads = self._config.role.loads + #: Read by `KVCacheManagerV2`, which is built after the connector. An + #: offer starts at the local match and covers whole blocks, so a role + #: that reads the pool needs that match to end on a block boundary. + #: Producer and capacity roles never offer and keep the manager's own + #: reuse unchanged. Every rank resolves the same role and arguments, so + #: the ranks of a server agree. + self.requires_block_aligned_reuse = loads + #: See `MooncakeStoreConfig.local_reuse_backoff`: "off" takes local + #: matches whole instead of aligning a backed-off one down. + self.drops_reuse_match_backoff = loads and local_reuse_backoff(llm_args) == "off" # Where this rank records the segment it mounts. `None` when the # deployment named no directory, which only makes the run # unreportable; see `ledger.record_segment`. @@ -189,30 +532,88 @@ def __init__(self, llm_args: TorchLlmArgs): self._run_dir = pool.run_dir if pool is not None else None self._rank = mpi_rank() self._world_size = mpi_world_size() + enable_attention_dp = getattr(llm_args, "enable_attention_dp", False) is True + # A stored page is one attention shard's KV. Every attention-DP owner + # holds all heads, so all owners key shard 0 of 1 on the shared pool; + # under tensor parallelism each rank keeps its own shard. The global + # rank still names this process in the ledger, thread names and logs. + if enable_attention_dp: + self._attention_rank, self._attention_world_size = 0, 1 + else: + self._attention_rank, self._attention_world_size = self._rank, self._world_size self._model_key = self._config.resolve_model_key(llm_args.model) self._addressing: Optional[PageAddressing] = None - # Namespaces for this rank, used for both directions of transfer. + # Namespaces for this rank's attention shard, used for both directions + # of transfer. self._namespaces: Dict[int, KeyNamespace] = {} - # The same namespaces for every rank. A prefix is only reusable when all - # shards of it are present, so a lookup has to ask about all of them. + # The same namespaces for every attention shard. A prefix is only + # reusable when all shards of it are present, so a lookup has to ask + # about all of them: every tensor-parallel rank's, or the single shard + # every attention-DP owner shares. self._peer_namespaces: Dict[int, Tuple[KeyNamespace, ...]] = {} - - # Checked before the segment is mounted, since the kernel answers an - # unaffordable total by killing the process rather than by failing the - # allocation. - validate_node_budget(self._config) - - self._store, self._segment_host = _open_store(self._config) - record_segment( - self._run_dir, - host=self._segment_host, - rank=self._rank, - segment_size=self._config.global_segment_size, - role=self._config.role.value, - model_key=self._model_key, + # A rank that holds whole pages and its own batch can restart a request + # whose page went missing: a single-rank server, or an attention-DP + # owner. Tensor-parallel ranks hold shards of the same page and would + # have to agree on the restart, which they cannot do here, so a failed + # load stays fatal there. + holds_whole_pages = enable_attention_dp or self._world_size == 1 + self._recover_failed_loads = ( + getattr(type(self), "reports_failed_loads", False) is True and holds_whole_pages ) - + # Requests whose last synchronous load came up short, taken by the + # executor right after `start_load_kv`. Executor thread only. + self._failed_load_requests: Set[int] = set() + # Diagnostics, off unless `stats_interval` is set. Only a role that + # reads the pool looks prefixes up or loads them, so the others count + # nothing. + self._stats_interval = self._config.stats_interval + counts = self._stats_interval is not None and loads + self._hole_stats: Optional[_PrefixHoleStats] = _PrefixHoleStats() if counts else None + self._hole_stats_last = time.monotonic() + self._load_stats: Optional[_LoadStats] = _LoadStats() if counts else None + self._load_stats_last = time.monotonic() + + # Asynchronous loads; see `_start_async_load`. Only a role that reads + # the pool loads at all, so producer and capacity roles start no load + # thread and pin no load staging. + self._async_load = self._config.async_load and loads + self._async_workers = self._config.async_load_workers if self._async_load else 0 + # One shared FIFO: an idle load thread takes the next request, so each + # request is handled by exactly one thread, in arrival order. + self._load_queue: "Queue[Optional[_QueuedLoad]]" = Queue() + self._load_threads: List[threading.Thread] = [] + # Why a load thread never reached its loop, read back by + # `_start_load_threads`. + self._load_startup_errors: List[BaseException] = [] + # Set at shutdown, so a load still queued never reaches the store. + self._load_stop = threading.Event() + # Guards the asynchronous bookkeeping below. + self._load_lock = threading.Lock() + # One staging pool per load thread, under `stage_through_host`. + self._async_stagings: List[HostStagingPool] = [] + # Staging pools whose unregistration failed, kept referenced so the + # fabric never reaches memory the allocator has handed out again. + self._unreleased_staging: List[HostStagingPool] = [] + # Queued or loading, by request. + self._async_pending: Set[int] = set() + # Taken by a load thread and loading right now. + self._async_active: Set[int] = set() + # Done but not yet reported: request -> whether a page failed to land. + self._async_done: Dict[int, bool] = {} + # Requests the runtime said it parked (`started_loading_req_ids`). + self._async_announced: Set[int] = set() + # Reported finished with a page missing; see + # `take_failed_async_load_requests`. + self._failed_async_load_requests: Set[int] = set() + # The first exception on a load thread, or the first load this rank + # could not queue. Sticky: every later executor call raises it, until + # shutdown; see `_reraise_background_error`. + self._load_error: Optional[BaseException] = None + + # Transfer state, set before the store is opened so that `shutdown` + # also runs on a worker whose store never opened. + self._store = None self._save_queue: "Queue[Optional[Tuple[torch.cuda.Event, List[RequestTransfers]]]]" = ( Queue() ) @@ -231,26 +632,66 @@ def __init__(self, llm_args: TorchLlmArgs): self._device_index: Optional[int] = None # Pages per store call. Staging narrows this to the slots it can afford. self._batch_size = self._config.transfer_batch_size + # Pages per store call on a load thread; its own staging narrows it. + self._async_batch_size = self._batch_size # Save submissions still in flight, per request. self._outstanding_saves: Dict[int, int] = defaultdict(int) # Requests the runtime has told us are done producing KV. Their pages # stay pinned until we report them back through `get_finished`. self._closed_requests: Set[int] = set() + # Requests whose saves the last `wait_for_save` queued. The scheduler + # adapter takes them at its next `build_connector_meta` and only then + # counts those blocks as stored; see `take_submitted_saves`. Executor + # thread only. + self._submitted_saves: Set[int] = set() # Why the save thread never reached its transfer loop, read back by # `register_kv_cache_layout`. A failure inside the loop is dropped # there instead. self._save_startup_error: Optional[BaseException] = None + # Checked before the segment is mounted, since the kernel answers an + # unaffordable total by killing the process rather than by failing the + # allocation. + validate_node_budget(self._config) + + # Write options exist only for a role that writes, so consumer and + # capacity roles build nothing and need no particular binding. + saves = self._config.role.saves + self._soft_pin_ttl_ms = self._config.soft_pin_ttl_ms if saves else None + self._group_blocks = self._config.group_blocks and saves + # Shared by every put; None keeps the store's defaults. Grouped puts + # need one group per key and build their own per call, and building a + # probe here proves the binding supports them before a segment is + # mounted. + self._put_config = _replicate_config( + self._soft_pin_ttl_ms, ["probe"] if self._group_blocks else None + ) + + self._store, self._segment_host = _open_store(self._config) + record_segment( + self._run_dir, + host=self._segment_host, + rank=self._rank, + segment_size=self._config.global_segment_size, + role=self._config.role.value, + model_key=self._model_key, + ) + global _LOCAL_WORKER _LOCAL_WORKER = self _LOCAL_WORKER_READY.set() logger.warning( f"mooncake-store worker rank {self._rank}/{self._world_size} ready " - f"(role={self._config.role.value}, model_key={self._model_key}, " + f"(attention shard {self._attention_rank}/{self._attention_world_size}, " + f"role={self._config.role.value}, model_key={self._model_key}, " f"master={self._config.master_server_address}, " f"segment={self._config.global_segment_size / (1 << 30):.1f} GiB " - f"on {self._segment_host})" + f"on {self._segment_host}" + f"{f', soft_pin_ttl_ms={self._soft_pin_ttl_ms}' if self._soft_pin_ttl_ms is not None else ''}" + f"{', group_blocks=1' if self._group_blocks else ''}" + f"{f', stats_interval={self._stats_interval:g}s' if self._hole_stats is not None else ''}" + f"{f', async_load workers={self._async_workers}' if self._async_load else ''})" ) if self.capacity_only: # Said plainly, since every other sign of a working connector is @@ -288,7 +729,13 @@ def register_kv_caches(self, kv_cache_tensor: torch.Tensor): ) def register_kv_cache_layout(self, layout: KvCacheLayout) -> None: - """Register the KV pools with Mooncake and start the save thread.""" + """Register the KV pools with Mooncake and start the transfer threads. + + Raises: + RuntimeError: When the save thread or an asynchronous load thread + cannot start, since a request would then wait on a transfer + nothing runs. + """ if self._addressing is not None: raise RuntimeError("KV cache layout already registered") @@ -319,11 +766,11 @@ def register_kv_cache_layout(self, layout: KvCacheLayout) -> None: for layer_group_id in addressing.layer_group_ids: bytes_per_page = addressing.bytes_per_page(layer_group_id) self._namespaces[layer_group_id] = self._namespace( - self._rank, layer_group_id, bytes_per_page + self._attention_rank, layer_group_id, bytes_per_page ) self._peer_namespaces[layer_group_id] = tuple( - self._namespace(rank, layer_group_id, bytes_per_page) - for rank in range(self._world_size) + self._namespace(shard, layer_group_id, bytes_per_page) + for shard in range(self._attention_world_size) ) if self._config.role.saves: @@ -344,11 +791,45 @@ def register_kv_cache_layout(self, layout: KvCacheLayout) -> None: "could not start, so this worker would accept pages it " "could never write to the pool." ) from startup_error + if self._async_load: + self._start_load_threads() logger.warning( f"mooncake-store worker rank {self._rank} registered layout: {addressing.describe()}" + f"{f' ({self._async_workers} asynchronous load workers)' if self._async_load else ''}" ) + def _start_load_threads(self) -> None: + """Start the asynchronous load threads and wait until each has set up. + + Raises: + RuntimeError: When any of them could not bind this rank's device or + open its stream, since a parked request would then wait on a + load nothing runs. + """ + started: List[threading.Event] = [] + for index in range(self._async_workers): + ready = threading.Event() + thread = threading.Thread( + target=self._drain_loads, + args=(index, ready), + name=f"mooncake-store-load-{self._rank}-{index}", + daemon=True, + ) + self._load_threads.append(thread) + started.append(ready) + thread.start() + for ready in started: + ready.wait() + with self._load_lock: + errors, self._load_startup_errors = self._load_startup_errors, [] + if errors: + raise RuntimeError( + f"mooncake-store: {len(errors)} of {self._async_workers} load thread(s) " + f"for rank {self._rank} could not start, so a parked request would " + "wait on a load nothing runs." + ) from errors[0] + def _register_pools(self, addressing: PageAddressing) -> None: """Register the KV pools with Mooncake, one call per allocation. @@ -413,53 +894,117 @@ def _open_staging(self, addressing: PageAddressing) -> None: """Allocate and register the pinned slots pages will pass through. Only the directions this role drives get a pool, since each one costs - a pinned allocation of its own. + a pinned allocation of its own. Each direction gets + `staging_buffer_bytes`. With `async_load` the load direction is the load + threads' pools, which split that budget, and no synchronous load pool + is pinned: the node memory check counts one budget per direction. + + Pools are created loads first, then saves. A pool that fails to + register leaves the ones before it in place for `shutdown` to release. """ max_bytes_per_page = max( addressing.bytes_per_page(layer_group_id) for layer_group_id in addressing.layer_group_ids ) - slot_bytes, num_slots = plan_slot_geometry( + batch = self._config.transfer_batch_size + budget = self._config.staging_buffer_bytes + slot_bytes, num_slots = plan_slot_geometry(max_bytes_per_page, batch, budget) + loads, saves = self._config.role.loads, self._config.role.saves + self._warn_raised_to_one_page( + [ + label + for label, used in (("load", loads and not self._async_load), ("save", saves)) + if used + ], + budget, max_bytes_per_page, - self._config.transfer_batch_size, - MAX_STAGING_BUFFER_BYTES, ) - if self._config.role.loads: + if loads and not self._async_load: self._load_staging = HostStagingPool( slot_bytes=slot_bytes, num_slots=num_slots, store=self._store, label="load", ) - if self._config.role.saves: + if self._async_load: + share = budget // self._async_workers + _, worker_slots = plan_slot_geometry(max_bytes_per_page, batch, share) + self._warn_raised_to_one_page( + [f"async-load-{index}" for index in range(self._async_workers)], + share, + max_bytes_per_page, + split=self._async_workers > 1, + ) + for index in range(self._async_workers): + self._async_stagings.append( + HostStagingPool( + slot_bytes=slot_bytes, + num_slots=worker_slots, + store=self._store, + label=f"async-load-{index}", + ) + ) + self._async_batch_size = min(batch, worker_slots) + if worker_slots < min(64, batch): + logger.warning( + f"mooncake-store rank {self._rank}: each of its {self._async_workers} " + f"load workers stages only {worker_slots} page(s) of " + f"{max_bytes_per_page} B per store call, so per-call overhead " + "dominates its loads. Raise staging_buffer_bytes or lower " + "async_load_workers." + ) + if saves: self._save_staging = HostStagingPool( slot_bytes=slot_bytes, num_slots=num_slots, store=self._store, label="save", ) - self._batch_size = min(self._config.transfer_batch_size, num_slots) - if self._batch_size < self._config.transfer_batch_size: + # Saves, and synchronous loads where they have a pool. + self._batch_size = min(batch, num_slots) + if self._batch_size < batch and (saves or self._load_staging is not None): logger.warning( f"mooncake-store rank {self._rank} reduced its transfer batch from " - f"{self._config.transfer_batch_size} to {self._batch_size} pages: " - f"staging {max_bytes_per_page} B pages within the " - f"{MAX_STAGING_BUFFER_BYTES} B pinned-memory ceiling does not fit " - f"more. Lower transfer_batch_size to make the reduction explicit." + f"{batch} to {self._batch_size} pages: staging {max_bytes_per_page} B " + f"pages within staging_buffer_bytes={budget} does not fit more. Raise " + "staging_buffer_bytes or lower transfer_batch_size." ) - def _namespace(self, rank: int, layer_group_id: int, bytes_per_page: int) -> KeyNamespace: + def _warn_raised_to_one_page( + self, labels: List[str], budget: int, page_bytes: int, *, split: bool = False + ) -> None: + """Say that pools whose budget is below one page pin a page each regardless. + + Args: + labels: The pools sharing this budget. + budget: Pinned bytes each of them was given. + page_bytes: The largest page, which a slot has to hold. + split: Whether `budget` is one load thread's share of + `staging_buffer_bytes`, which fewer load threads would raise. + """ + if not labels or budget >= page_bytes: + return + logger.warning( + f"mooncake-store rank {self._rank}: the {', '.join(labels)} staging " + f"pool(s) get {budget} B each, below one {page_bytes} B page, so each is " + "raised to one page and pins more than the node memory check counted " + "for staging_buffer_bytes. Raise staging_buffer_bytes" + f"{' or lower async_load_workers' if split else ''}." + ) + + def _namespace(self, shard: int, layer_group_id: int, bytes_per_page: int) -> KeyNamespace: + """The key namespace of one attention shard's pages of a layer group.""" return KeyNamespace( namespace=self._config.namespace, model_key=self._model_key, - rank=rank, - world_size=self._world_size, + rank=shard, + world_size=self._attention_world_size, layer_group_id=layer_group_id, tokens_per_block=self._addressing.tokens_per_block, bytes_per_page=bytes_per_page, ) - # ---- leader-facing lookup ---- + # ---- lookup, for the scheduler adapter in this process ---- @property def config(self) -> MooncakeStoreConnectorConfig: @@ -474,9 +1019,9 @@ def is_registered(self) -> bool: def count_prefix_hit(self, block_hashes: Sequence[bytes]) -> int: """How many leading blocks of `block_hashes` are fully present. - A block counts only when every layer group and every rank has its - page, since a prefix is replayed as a whole. The scan stops at the - first incomplete block, a later hit being unusable on its own. + A block counts only when every layer group and every attention shard + has its page, since a prefix is replayed as a whole. The scan stops at + the first incomplete block, a later hit being unusable on its own. Args: block_hashes: Candidate hashes in block order. @@ -517,20 +1062,136 @@ def count_prefix_hit(self, block_hashes: Sequence[bytes]) -> int: if not all(status == 1 for status in window): break hit_blocks += 1 + if self._hole_stats is not None: + # From the same answer, so the diagnostics cost no store call. + self._record_prefix_hole(present, keys_per_block, len(block_hashes), hit_blocks) return hit_blocks + def _record_prefix_hole( + self, present: Sequence[int], keys_per_block: int, num_blocks: int, hit_blocks: int + ) -> None: + """Count the stored blocks behind the first missing one. Never changes the hit.""" + stored_after_hole = [ + index + for index in range(hit_blocks + 1, num_blocks) + if all( + status == 1 + for status in present[index * keys_per_block : (index + 1) * keys_per_block] + ) + ] + self._hole_stats.record( + num_blocks, hit_blocks, stored_after_hole, self._addressing.tokens_per_block + ) + self._maybe_log_hole_stats() + + def _maybe_log_hole_stats(self, force: bool = False) -> Optional[str]: + """Log the prefix-hole summary at most once per `stats_interval`. + + Args: + force: Log what accumulated even if the interval has not elapsed. + + Returns: + The line, when one was logged. Nothing is logged when the stats are + off or no lookup ran since the last line. + """ + # Read defensively: `shutdown` also runs from `__del__` on a worker + # whose constructor may not have got this far. + stats = getattr(self, "_hole_stats", None) + if stats is None: + return None + now = time.monotonic() + due = force or now - self._hole_stats_last >= self._stats_interval + if not (due and stats.lookups): + return None + line = ( + f"mooncake-store prefix-hole stats: rank {self._rank} " + f"window_s={now - self._hole_stats_last:.0f} {stats.summary()}" + ) + stats.reset() + self._hole_stats_last = now + logger.info(line) + return line + + def _maybe_log_load_stats(self, force: bool = False) -> Optional[str]: + """Log the load summary at most once per `stats_interval`. + + Args: + force: Log what accumulated even if the interval has not elapsed. + + Returns: + The line, when one was logged. Nothing is logged when the stats are + off or nothing loaded since the last line. + """ + stats = getattr(self, "_load_stats", None) + if stats is None: + return None + now = time.monotonic() + if not (force or now - self._load_stats_last >= self._stats_interval): + return None + if self._async_load: + with self._load_lock: + active = len(self._async_active) + queued = len(self._async_pending) - active + summary = stats.take_summary(workers=self._async_workers, queued=queued, active=active) + else: + summary = stats.take_summary() + if summary is None: + return None + line = ( + f"mooncake-store load stats: rank {self._rank} " + f"window_s={now - self._load_stats_last:.0f} {summary}" + ) + self._load_stats_last = now + logger.info(line) + return line + # ---- load path ---- def start_load_kv(self, stream: torch.cuda.Stream): - """Pull every scheduled page into its GPU slot before the forward pass.""" + """Pull every scheduled page into its GPU slot before the forward pass. + + The runtime has already counted these tokens as computed. When a page + is missing and this rank can restart a request alone, the request is + recorded for `take_failed_load_requests`, its remaining pages are not + fetched, and its saves leave the bound metadata; the other requests + load as usual. The bound loads are consumed, so calling this again on + the same metadata starts nothing. + + Args: + stream: The executor's stream, on which staged scatters are issued. + + Raises: + RuntimeError: When a page is missing on a tensor-parallel rank, a + zero-copy load has no allocation fence to wait on, a staged + synchronous load arrives on a rank whose loads all go through + the load threads, or a background load failed. + """ + # First, and before the empty-metadata return: under attention DP this + # runs inside the executor's failure vote, and without it the error + # stops the executor loop. + self._reraise_background_error() metadata: Optional[MooncakeStoreMetadata] = self.get_connector_meta() if metadata is None or not metadata.loads: return + loads = list(metadata.loads) + metadata.loads.clear() - keys, addresses, sizes, total_pages = self._resolve(metadata.loads) + # Per-phase times of this load, collected only when stats are on. + phases: Optional[Dict[str, List[float]]] = {} if self._load_stats is not None else None + with _timed_phase(phases, "resolve"): + keys, addresses, sizes, total_pages = self._resolve(loads) if not keys: return - + # `_resolve` emits one key per page in transfer order, so this names the + # request every key is loaded for. + owners = [entry.request_id for entry in loads for _ in entry.pages] + + if self._config.stage_through_host and self._load_staging is None: + raise RuntimeError( + "mooncake-store: a synchronous load arrived on a rank whose loads all " + "go through the async_load workers, so it holds no synchronous " + "staging pool" + ) staging = self._load_staging handle = _stream_handle(stream) if staging is not None else 0 if staging is None: @@ -546,47 +1207,457 @@ def start_load_kv(self, stream: torch.cuda.Stream): "owns. The KV cache manager must offer wait_allocation_fences; " "KVCacheManagerV2 does." ) - fence({entry.request_id for entry in metadata.loads}) - - for batch in zip( - _batched(keys, self._batch_size), - _batched(addresses, self._batch_size), - _batched(sizes, self._batch_size), - ): - batch_keys, batch_addresses, batch_sizes = batch - if staging is None: - target_addresses, target_sizes = list(batch_addresses), list(batch_sizes) - else: - target_addresses, target_sizes = describe_batch_for_get(staging, batch_sizes) - results = self._store.batch_get_into_multi_buffers( - list(batch_keys), target_addresses, target_sizes + with _timed_phase(phases, "fence"): + fence({entry.request_id for entry in loads}) + + transfer_start = time.perf_counter() if phases is not None else 0.0 + failed_requests, num_failed_pages, first_failure, loaded_bytes = self._load_pages( + keys, + addresses, + sizes, + owners, + staging, + handle, + self._batch_size, + recover=self._recover_failed_loads, + phases=phases, + ) + if phases is not None: + self._load_stats.record( + time.perf_counter() - transfer_start, + pages=total_pages, + batches=-(-len(keys) // self._batch_size), + loaded_bytes=loaded_bytes, + phases=phases, + failed_pages=num_failed_pages, + ) + if failed_requests: + # The executor takes these requests out of the batch and restarts + # them; see `take_failed_load_requests`. Nothing of theirs is + # computed this pass, so their saves leave the bound metadata now: + # `wait_for_save` would otherwise read pages the executor has freed, + # or another request has since filled, and publish them under the + # failed request's valid keys. + self.drop_bound_saves(failed_requests) + self._failed_load_requests.update(failed_requests) + logger.warning( + f"mooncake-store rank {self._rank} failed to load {num_failed_pages} page(s) " + f"for {len(failed_requests)} request(s); they restart without the offered " + f"prefix. First failure: {first_failure!r}" ) - failed = [ - key - for key, result in zip(batch_keys, results) - if not isinstance(result, int) or result < 0 + logger.debug(f"mooncake-store rank {self._rank} loaded {total_pages} pages") + + def _load_pages( + self, + keys: List[str], + addresses: List[List[int]], + sizes: List[List[int]], + owners: List[int], + staging: Optional[HostStagingPool], + handle: int, + batch_size: int, + *, + recover: bool, + phases: Optional[Dict[str, List[float]]] = None, + ) -> Tuple[Set[int], int, Optional[str], int]: + """Fetch pages batch by batch, scattering staged ones onto the device. + + Shared by the synchronous path (executor thread and stream) and the + load threads (each with its own stream and staging pool). A store call + is synchronous, and a staged batch's scatter is synced before the next + batch reuses its slots, so no write into a page is in flight on return. + + Args: + keys: Store key of every page, in transfer order. + addresses: Device region addresses of every page. + sizes: Device region sizes of every page. + owners: The request each page is loaded for. + staging: Host slots the pages pass through, or None for zero copy. + handle: Raw CUDA stream the staged scatters are issued on. + batch_size: Pages per store call. + recover: Whether a missing page is handed back rather than raised. + The pages of a request that already lost one are then skipped, + since the request restarts and recomputes all of them. + phases: Per-phase totals the describe, get, unstage and sync + phases add to, or None to time nothing; see `_timed_phase`. + + Returns: + The requests that lost a page, how many pages were lost, the first + missing key, and the bytes that landed. + + Raises: + RuntimeError: On the first batch with a missing page when `recover` + is False. + """ + failed_requests: Set[int] = set() + first_failure: Optional[str] = None + num_failed_pages = 0 + loaded_bytes = 0 + + for start in range(0, len(keys), batch_size): + indices = [ + index + for index in range(start, min(start + batch_size, len(keys))) + if owners[index] not in failed_requests ] - if failed or len(results) != len(batch_keys): - # The runtime already counted these tokens as computed, so a - # partial load would leave the forward pass reading - # uninitialized KV and silently producing wrong tokens. - raise RuntimeError( - f"mooncake-store failed to load {len(failed) or len(batch_keys)} of " - f"{len(batch_keys)} pages; the affected KV slots were already " - f"reported as computed. First failure: {failed[:1]}" + if not indices: + continue + batch_keys = [keys[index] for index in indices] + batch_addresses = [addresses[index] for index in indices] + batch_sizes = [sizes[index] for index in indices] + with _timed_phase(phases, "describe"): + if staging is None: + target_addresses, target_sizes = batch_addresses, batch_sizes + else: + target_addresses, target_sizes = describe_batch_for_get(staging, batch_sizes) + with _timed_phase(phases, "get"): + results = self._store.batch_get_into_multi_buffers( + batch_keys, target_addresses, target_sizes ) - if staging is not None: - # Only reached once every page in the batch landed, so no slot - # holding a failed read is copied over a device page. - unstage_batch_after_get(staging, batch_addresses, batch_sizes, handle) + if len(results) != len(batch_keys): + # Which pages landed is unknowable, so none of them counts. + results = [-1] * len(batch_keys) + landed = [ + position + for position, result in enumerate(results) + if isinstance(result, int) and result >= 0 + ] + loaded_bytes += sum(sum(batch_sizes[position]) for position in landed) + if len(landed) != len(batch_keys): + landed_set = set(landed) + for position, index in enumerate(indices): + if position in landed_set: + continue + failed_requests.add(owners[index]) + num_failed_pages += 1 + if first_failure is None: + first_failure = batch_keys[position] + if not recover: + # The runtime already counted these tokens as computed, so a + # partial load would leave the forward pass reading + # uninitialized KV and silently producing wrong tokens. + raise RuntimeError( + f"mooncake-store failed to load {len(batch_keys) - len(landed)} of " + f"{len(batch_keys)} pages; the affected KV slots were already " + f"reported as computed. First failure: {first_failure!r}" + ) + if staging is not None and landed: + # Only filled slots are scattered, so no slot holding a failed + # read is copied over a device page. + with _timed_phase(phases, "unstage"): + unstage_batch_after_get( + staging, batch_addresses, batch_sizes, handle, only=landed + ) # The next batch reuses the slots and the forward pass reads # these pages, so the scatter has to complete before either. - _sync_stream(handle) + with _timed_phase(phases, "sync"): + _sync_stream(handle) + return failed_requests, num_failed_pages, first_failure, loaded_bytes - logger.debug(f"mooncake-store rank {self._rank} loaded {total_pages} pages") + # ---- asynchronous load path ---- + + def bind_connector_meta(self, metadata: object) -> None: + """Bind this iteration's work and queue the parked requests' loads. + + Every rank binds the same metadata, so every rank queues its own pages + here, including the tensor-parallel ranks whose scheduler adapter never + saw the allocation. The asynchronous loads are consumed, so binding the + same metadata again queues nothing twice. The wait for the destination + pages happens on the load thread; see `_drain_loads`. + + Never raises: it runs outside the executor's attention-DP failure vote, + so what cannot be queued is kept for the next `start_load_kv` or + `get_finished` to raise. + + Args: + metadata: This iteration's `MooncakeStoreMetadata`. + """ + super().bind_connector_meta(metadata) + async_loads = getattr(metadata, "async_loads", None) + if not async_loads: + return + queued = list(async_loads) + metadata.async_loads = [] + if not self._async_load: + self._stash_load_error( + RuntimeError( + "mooncake-store received asynchronous loads, but async_load is off on this rank" + ) + ) + return + if not self._load_threads: + self._stash_load_error( + RuntimeError( + "mooncake-store received asynchronous loads before its KV cache " + "layout was registered, so no load thread runs them" + ) + ) + return + # Bound by `handle_metadata` just before this metadata, for the batch + # that allocated the destination pages. + fence = self.get_allocation_fence() + for transfers in queued: + self._start_async_load(transfers, fence) + + def _start_async_load( + self, transfers: RequestTransfers, fence: Optional[AllocationFence] + ) -> None: + """Queue one parked request's pages for the load threads. + + The fence bound in the iteration that parked the request travels with + it, since a later iteration binds another. A parked request keeps its + allocation, and the fence its pages, until its load is reported, so the + load thread can still wait on it. A transfer with no pages, for an + offer the runtime honoured with no tokens, completes at once, so the + parked request is released at the next `get_finished`. + + Args: + transfers: The parked request's pages, every layer group's. + fence: The KV cache manager's allocation fence for that iteration. + """ + request_id = transfers.request_id + if transfers.pages and fence is None: + # Neither the NIC writing zero copy nor a staged scatter on a load + # thread's stream is ordered behind the allocation stream. + self._stash_load_error( + RuntimeError( + "mooncake-store: an asynchronous load has no allocation fence to " + "wait on, so the NIC could write pages earlier GPU work still " + "owns. The KV cache manager must offer wait_allocation_fences; " + "KVCacheManagerV2 does." + ) + ) + with self._load_lock: + self._async_done[request_id] = True + return + with self._load_lock: + if not transfers.pages: + self._async_done[request_id] = False + return + self._async_pending.add(request_id) + self._load_queue.put((request_id, transfers, time.monotonic(), fence)) + + def _drain_loads(self, index: int, started: threading.Event) -> None: + """Run parked requests' loads, one at a time, until a None arrives. + + Args: + index: This thread's position, which selects its staging pool. + started: Set once the thread has finished its device and stream + setup, whether or not that succeeded. + """ + try: + # A new thread starts on device 0; see `_drain_saves`. + if self._device_index is not None: + torch.cuda.set_device(self._device_index) + staging = self._async_stagings[index] if self._async_stagings else None + # Owned by this thread for its lifetime, so a staged scatter never + # queues behind the executor's work or another load's. + stream = ( + torch.cuda.Stream() if staging is not None and torch.cuda.is_available() else None + ) + handle = _stream_handle(stream) if staging is not None else 0 + except Exception as exc: + # Fatal: a thread that died here would leave the requests parked on + # it waiting for good. + logger.error( + f"mooncake-store load thread {index} failed to start on rank " + f"{self._rank}: {type(exc).__name__}: {exc}" + ) + with self._load_lock: + self._load_startup_errors.append(exc) + return + finally: + started.set() + while True: + item = self._load_queue.get() + if item is None: + return + request_id, transfers, enqueued_at, fence = item + if self._load_stop.is_set(): + # Shutting down: a load still queued never reaches the store. + with self._load_lock: + self._async_pending.discard(request_id) + continue + self._run_async_load(index, request_id, transfers, enqueued_at, fence, staging, handle) + + def _run_async_load( + self, + index: int, + request_id: int, + transfers: RequestTransfers, + enqueued_at: float, + fence: AllocationFence, + staging: Optional[HostStagingPool], + handle: int, + ) -> None: + """Wait for one parked request's pages to be free, then fill them. + + A missing page marks the request failed for `get_finished` to hand + back or raise; any exception is also kept as this worker's background + error. Either way the request is marked done, so it does not stay + parked forever. + """ + phases: Optional[Dict[str, List[float]]] = {} if self._load_stats is not None else None + with self._load_lock: + self._async_active.add(request_id) + started_at = time.monotonic() + failed = False + total_pages = batches = num_failed_pages = loaded_bytes = 0 + try: + # The pages were reserved in the parking iteration, possibly on the + # KV cache manager's allocation stream, and a recycled page can + # wait on a forward pass still in flight. That wait happens here, + # off the executor thread. + with _timed_phase(phases, "fence"): + fence({request_id}) + started_at = time.monotonic() + with _timed_phase(phases, "resolve"): + keys, addresses, sizes, total_pages = self._resolve([transfers]) + batches = -(-len(keys) // self._async_batch_size) + # Recovered here whatever the rank: `get_finished` decides whether + # a failed request is fatal. + failed_requests, num_failed_pages, first_failure, loaded_bytes = self._load_pages( + keys, + addresses, + sizes, + [request_id] * len(keys), + staging, + handle, + self._async_batch_size, + recover=True, + phases=phases, + ) + failed = bool(failed_requests) + if failed: + logger.warning( + f"mooncake-store rank {self._rank} failed to load {num_failed_pages} " + f"page(s) for parked request {request_id}. First failure: " + f"{first_failure!r}" + ) + else: + logger.debug( + f"mooncake-store rank {self._rank} load thread {index} loaded " + f"{total_pages} pages for parked request {request_id}" + ) + except Exception as exc: + # The thread boundary: kept for the executor thread to raise. + failed = True + logger.error( + f"mooncake-store asynchronous load of request {request_id} failed on " + f"rank {self._rank}: {type(exc).__name__}: {exc}\n{traceback.format_exc()}" + ) + self._stash_load_error(exc) + finally: + finished_at = time.monotonic() + # Recorded before the request leaves the pending set, so a summary + # taken once nothing is pending includes it. + if self._load_stats is not None: + self._load_stats.record( + finished_at - started_at, + pages=total_pages, + batches=batches, + loaded_bytes=loaded_bytes, + phases=phases, + failed_pages=num_failed_pages, + wait_s=started_at - enqueued_at, + ) + with self._load_lock: + self._async_active.discard(request_id) + self._async_pending.discard(request_id) + self._async_done[request_id] = failed + + def take_failed_async_load_requests(self) -> Set[int]: + """Parked requests whose last `get_finished` reported a load that came up short. + + The executor drains this in the same `_kv_connector_terminate_requests` + call that ran `get_finished`, before the next scheduling pass, drops + their allocation and lets the scheduler admit them again. Every write + into their pages has completed by then: an id is reported only once its + load thread is done with it. Ids the runtime no longer holds parked, + because the request finished or was cancelled meanwhile, are ignored + there. Cleared on return. + + Returns: + The request IDs to restart. + """ + with self._load_lock: + failed, self._failed_async_load_requests = self._failed_async_load_requests, set() + return failed + + def _stash_load_error(self, error: BaseException) -> None: + """Keep the first background load error for the executor thread to raise.""" + with self._load_lock: + if self._load_error is None: + self._load_error = error + + def _reraise_background_error(self) -> None: + """Raise the kept background load error, if any. + + Sticky rather than raised once, so that no later caller, such as a + second recovery poll or the attention-DP shutdown drain, sees a healthy + worker. It is never cleared; only shutting the worker down ends the + calls that raise it. Save errors keep their log-and-drop policy and are + not raised. + + Raises: + RuntimeError: Chained from the kept error. + """ + with self._load_lock: + error = self._load_error + if error is not None: + raise RuntimeError( + f"mooncake-store rank {self._rank}: a background load failed: " + f"{type(error).__name__}: {error}" + ) from error + + def take_failed_load_requests(self) -> Set[int]: + """Requests whose pages could not all be loaded in the last `start_load_kv`. + + The executor drains this right after `start_load_kv` and before the + forward pass, takes those requests out of the batch, frees their + allocation and lets the scheduler admit them again, so the store is + asked afresh and the lost pages are recomputed. Every write into their + pages has completed by then; see `_load_pages`. Cleared on return. + + Returns: + The request IDs to restart. + """ + failed, self._failed_load_requests = self._failed_load_requests, set() + return failed + + def drop_bound_saves(self, request_ids: Iterable[int]) -> int: + """Remove these requests' saves from the metadata bound for this pass. + + For a request taken out of the batch after `build_connector_meta` ran: + its scheduled tail is never computed, so none of it may be published. + Edits the bound metadata in place, which is the object `wait_for_save` + reads, and a request that is not bound is ignored, so repeating a call + changes nothing. + + Args: + request_ids: The restarted requests. + + Returns: + The number of save entries removed. + """ + metadata: Optional[MooncakeStoreMetadata] = self.get_connector_meta() + if metadata is None or not metadata.saves: + return 0 + drop = set(request_ids) + saves = metadata.saves + kept = [transfers for transfers in saves if transfers.request_id not in drop] + removed = len(saves) - len(kept) + if removed: + dropped_ids = sorted({transfers.request_id for transfers in saves} & drop) + saves[:] = kept + logger.debug( + f"mooncake-store rank {self._rank} dropped the saves of requests " + f"{dropped_ids} ({removed} entr{'y' if removed == 1 else 'ies'}) from this pass" + ) + return removed def wait_for_layer_load(self, layer_idx: int, stream: torch.cuda.Stream): - """No-op: loads complete in `start_load_kv`. + """No-op: loads complete in `start_load_kv`, or before a parked request rejoins. Transfers are whole pages, so every layer of a group lands in one store call and nothing is outstanding when the first layer runs. @@ -617,23 +1688,53 @@ def wait_for_save(self, stream: torch.cuda.Stream): for transfers in metadata.saves: self._outstanding_saves[transfers.request_id] += 1 self._save_queue.put((event, list(metadata.saves))) + # Replaced rather than accumulated: under tensor parallelism only rank + # 0's adapter takes these, so the other ranks must not grow the set + # every pass. + self._submitted_saves = {transfers.request_id for transfers in metadata.saves} + + def take_submitted_saves(self) -> Set[int]: + """Requests whose saves the last `wait_for_save` handed to the save thread. + + A pass that was bound but never ran never reaches `wait_for_save`, so + its saves are not reported and the scheduler adapter emits them again + instead of counting them as stored. Cleared on return. + + Returns: + The request ids of the saves the last `wait_for_save` queued. + """ + submitted, self._submitted_saves = self._submitted_saves, set() + return submitted def get_finished( self, finished_gen_req_ids: List[int], started_loading_req_ids: List[int] ) -> Tuple[List[int], List[int]]: - """Report which requests' saves have landed. + """Report which requests' saves and asynchronous loads have landed. Args: finished_gen_req_ids: Requests that will produce no further KV. - started_loading_req_ids: Requests loading asynchronously. Always - empty here, since `get_num_new_matched_tokens` only offers - synchronous loads; echoed back so the runtime does not wait on - something that already happened. + started_loading_req_ids: Requests the runtime parked for an + asynchronous load. Without `async_load` the scheduler never + offers one, and an id here is echoed back so the runtime does + not wait on something that already happened. With it, an id is + reported only after it was announced here and its load thread + is done with it: its pages landed, or one was found missing. Returns: Requests that have finished saving, and requests that have finished - loading. + loading. A finished load that lost a page is also listed for + `take_failed_async_load_requests`. + + Raises: + RuntimeError: When a background load failed, or a parked request + lost a page on a rank that cannot restart it alone (tensor + parallelism). """ + # Under attention DP this runs inside the executor's failure vote. + self._reraise_background_error() + # Runs every iteration on every rank that transfers, so it is where the + # load summary comes due. + self._maybe_log_load_stats() with self._save_lock: self._closed_requests.update(finished_gen_req_ids) finished_saving = [ @@ -644,7 +1745,37 @@ def get_finished( for request_id in finished_saving: self._closed_requests.discard(request_id) self._outstanding_saves.pop(request_id, None) - return finished_saving, list(started_loading_req_ids) + if not self._async_load: + return finished_saving, list(started_loading_req_ids) + + finished_loading: List[int] = [] + with self._load_lock: + self._async_announced.update(started_loading_req_ids) + for request_id in sorted(self._async_announced): + if request_id in self._async_done: + failed = self._async_done.pop(request_id) + elif request_id in self._async_pending: + continue + else: + # Parked by the runtime but never handed to this worker, so + # its pages hold nothing usable. + logger.warning( + f"mooncake-store rank {self._rank}: request {request_id} was " + "parked for an asynchronous load that was never started" + ) + failed = True + if failed and not self._recover_failed_loads: + raise RuntimeError( + f"mooncake-store rank {self._rank}: the asynchronous load of " + f"request {request_id} lost a page, and this deployment cannot " + "restart the request alone (tensor parallelism, or a connector " + "that does not report failed loads)" + ) + self._async_announced.discard(request_id) + if failed: + self._failed_async_load_requests.add(request_id) + finished_loading.append(request_id) + return finished_saving, finished_loading def _drain_saves(self) -> None: try: @@ -705,13 +1836,15 @@ def _put(self, transfers: Sequence[RequestTransfers]) -> None: staging = self._save_staging handle = _stream_handle(self._save_stream) if staging is not None else 0 + groups = self._block_groups(transfers) if self._group_blocks else [None] * len(keys) for batch in zip( _batched(keys, self._batch_size), _batched(addresses, self._batch_size), _batched(sizes, self._batch_size), + _batched(groups, self._batch_size), ): - batch_keys, batch_addresses, batch_sizes = batch + batch_keys, batch_addresses, batch_sizes, batch_groups = batch # Skip pages another rank or another instance already wrote. The # scheduler holds no store handle, and the answer changes between # the time it builds metadata and now. @@ -733,11 +1866,16 @@ def _put(self, transfers: Sequence[RequestTransfers]) -> None: ) # The store reads the slots on this thread, so fill them first. _sync_stream(handle) - results = self._store.batch_put_from_multi_buffers( - [batch_keys[i] for i in pending], - source_addresses, - source_sizes, - ) + put_args = ([batch_keys[i] for i in pending], source_addresses, source_sizes) + if self._group_blocks: + # One group per key that survived the existence filter. + put_args += ( + _replicate_config(self._soft_pin_ttl_ms, [batch_groups[i] for i in pending]), + ) + elif self._put_config is not None: + put_args += (self._put_config,) + # With neither option set this is the plain three-argument put. + results = self._store.batch_put_from_multi_buffers(*put_args) failures = sum(1 for result in results if not isinstance(result, int) or result < 0) if failures: # A dropped write only costs a future cache miss, which is not @@ -746,6 +1884,28 @@ def _put(self, transfers: Sequence[RequestTransfers]) -> None: f"mooncake-store rank {self._rank} failed to save {failures} of " f"{len(pending)} pages" ) + if self._soft_pin_ttl_ms is not None and failures == len(pending): + logger.warning_once( + f"mooncake-store rank {self._rank}: the pool refused every page " + f"of a put with soft_pin_ttl_ms={self._soft_pin_ttl_ms}. A " + "soft_pin_ttl_ms above the master's --max_kv_soft_pin_ttl (24 h " + "by default) makes the master refuse every put.", + key="mooncake-store-soft-pin-refused", + ) + + def _block_groups(self, transfers: Sequence[RequestTransfers]) -> List[str]: + """The eviction group of every page `_resolve` lists, in the same order. + + One group per block covers every layer group and attention shard of it. + A block is reusable only while all of those pages survive, so the master + should evict them together rather than leave orphans that hold space and + still miss. + """ + return [ + self._namespaces[page.layer_group_id].group(page.block_hash) + for entry in transfers + for page in entry.pages + ] # ---- shared ---- @@ -777,29 +1937,45 @@ def _resolve( return keys, addresses, sizes, pages def shutdown(self) -> None: - """Stop the save thread, then release what it was reading. Idempotent.""" - thread = self._save_thread - if thread is not None: + """Stop the transfer threads, then release what they used. Idempotent.""" + # What accumulated since the last line would otherwise be lost. + self._maybe_log_hole_stats(force=True) + self._maybe_log_load_stats(force=True) + # A load still queued is skipped rather than started. + self._load_stop.set() + save_thread = self._save_thread + threads = ([save_thread] if save_thread is not None else []) + list(self._load_threads) + # Every thread is signalled before any is joined, so they drain + # together against one deadline. A sentinel resent by a retry is + # harmless. + if save_thread is not None: self._save_queue.put(None) - thread.join(timeout=SAVE_DRAIN_TIMEOUT) - if thread.is_alive(): - # The thread reads the KV pools through the store handle and - # the staging slots below, so releasing either takes the memory - # out from under a transfer in flight. Leaking both is the safer - # end: the process is going down anyway, the thread is a daemon - # and does not hold it open, and a later call retries the join. - logger.error( - f"mooncake-store rank {self._rank}: the save thread did not stop " - f"within {SAVE_DRAIN_TIMEOUT:g}s, so the store handle and its " - "staging buffers are left in place rather than freed under a " - "transfer still reading them." - ) - return - self._save_thread = None + for _ in self._load_threads: + self._load_queue.put(None) + deadline = time.monotonic() + SAVE_DRAIN_TIMEOUT + for thread in threads: + thread.join(timeout=max(0.0, deadline - time.monotonic())) + alive = [thread.name for thread in threads if thread.is_alive()] + if alive: + # These threads use the KV pools through the store handle and the + # staging slots, so releasing either takes the memory out from + # under a transfer in flight. Leaking both is the safer end: the + # process is going down anyway, the threads are daemons and do not + # hold it open, and a later call retries the join. + logger.error( + f"mooncake-store rank {self._rank}: {', '.join(alive)} did not stop " + f"within {SAVE_DRAIN_TIMEOUT:g}s, so the store handle and its staging " + "buffers are left in place rather than freed under a transfer still " + "using them." + ) + return + self._save_thread = None + self._load_threads = [] if self._store is not None: # Before the store closes and the cache manager frees the pools, # so the fabric holds no registration over memory handed out again. self._unregister_pools() + self._close_staging() store, self._store = self._store, None if store is not None: try: @@ -809,16 +1985,32 @@ def shutdown(self) -> None: f"mooncake-store close failed: {type(exc).__name__}: {exc}\n" f"{traceback.format_exc()}" ) - # Released only after the store is closed, since it holds registrations - # against this memory. + # Dropped only after the store is closed. A pool whose registration did + # not come back stays referenced in `_unreleased_staging`. self._load_staging = None self._save_staging = None + self._async_stagings = [] self._save_stream = None global _LOCAL_WORKER if _LOCAL_WORKER is self: _LOCAL_WORKER = None _LOCAL_WORKER_READY.clear() + def _close_staging(self) -> None: + """Hand back every staging pool's registration before the store closes. + + A pool whose unregistration fails is kept in `_unreleased_staging`, so + its buffer outlives a registration the fabric may still use. + """ + for pool in (self._load_staging, self._save_staging, *self._async_stagings): + if pool is None: + continue + try: + pool.close() + except RuntimeError as exc: + logger.error(f"mooncake-store rank {self._rank}: {exc}") + self._unreleased_staging.append(pool) + def __del__(self): try: self.shutdown() diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index a963738f9ff8..4ac3dfdfb0ef 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -15,6 +15,7 @@ import math import os import sys +import time from collections import OrderedDict, defaultdict, deque from dataclasses import dataclass, field, replace from enum import IntEnum @@ -934,6 +935,29 @@ def _get_layer_spec_cache_size_components( ) +class ConnectorAdmissionOffer(NamedTuple): + """A KV connector's answer taken while the scheduler admitted the request. + + ``epoch`` is the scheduling pass it was taken in. ``prepare_resources`` + uses the answer only within that pass, and only against the same + ``local_end``. + """ + + # Scheduling pass the answer was taken in. + epoch: int + # The local match the connector was asked from. + local_end: int + # Tokens the connector offered past `local_end`. + num_tokens: int + # Whether the connector loads the offer asynchronously. + load_async: bool + + @property + def parks(self) -> bool: + """Whether the runtime will park the request for an asynchronous load.""" + return self.load_async and self.num_tokens > 0 + + class _MmRunMetadata(NamedTuple): """CPU tensors for exact multimodal run lookup. @@ -1772,6 +1796,20 @@ class KVCacheManagerV2(BaseResourceManager): # _init_mixed_precision assigns fresh per-instance values. low_precision_dtype: Optional[DataType] = None high_layers: dict[int, int] = {} + # Disaggregated completion reservation (see __init__): a logical bound in + # ledger blocks per admitted request, and the GPU slots one ledger block + # takes in each pool group. Declared on the class for instances built + # without __init__; __init__ assigns fresh per-instance dicts, and only a + # non-empty pool map lets anything write the limits. + _disagg_completion_limits: dict[int, int] = {} + _disagg_completion_pool_counts: dict[int, int] = {} + # Scheduling pass of admission-time connector answers (see + # `query_connector_at_admission`), on the class for the same reason: the + # connector prefix path reads it on instances tests build without __init__. + _connector_admission_epoch: int = 0 + # Per-iteration connector offer counters for TLLM_CTX_PARK_STATS, or None + # while disabled; see `enable_connector_offer_stats`. + _connector_offer_stats: Optional[Dict[str, float]] = None # Read by KvCacheCreator: a subclass that overrides the context # commit/history protocol opts out of generic reuse-match backoff. _supports_reuse_match_backoff = True @@ -1856,6 +1894,9 @@ def __init__( self._validate_speculative_config(spec_config) self.is_disagg = is_disagg self.kv_connector_manager = kv_connector_manager + # Scheduling pass of admission-time connector answers; see + # `query_connector_at_admission`. + self._connector_admission_epoch = 0 codec_config = getattr(cold_page_codec_provider, "config", None) self._has_lossy_cold_page_codec = ( getattr(codec_config, "algorithm", None) == "quantization_for_cold_page" @@ -2015,10 +2056,19 @@ def __init__( # check needs this raw value even where a subclass opts out of reuse-match backoff: # a block's reuse key covers no token past its end. self._draft_prompt_lookahead = draft_prompt_lookahead(spec_config) or 0 - # Unsupported adapters keep their main-like reuse endpoint. + # Unsupported adapters keep their main-like reuse endpoint, and so does + # a manager whose KV connector asks for whole local matches instead. self.reuse_match_backoff = self._draft_prompt_lookahead if not self._supports_reuse_match_backoff: self.reuse_match_backoff = 0 + elif self.reuse_match_backoff and self._connector_drops_reuse_match_backoff(): + logger.info( + "[KVCacheManager] The KV connector takes whole local prefix matches: " + f"dropping the {self.reuse_match_backoff}-token one-model draft reuse " + "backoff. Draft KV at the tail of a reused block may belong to another " + "continuation, which affects draft acceptance only." + ) + self.reuse_match_backoff = 0 self.enable_partial_reuse = _effective_partial_reuse( kv_cache_config.enable_partial_reuse, self._supports_partial_reuse, @@ -2479,6 +2529,38 @@ def create_cold_page_codec(cache_config: object) -> Optional[object]: self.max_blocks_per_seq = ((self.max_blocks_per_seq + 3) // 4) * 4 self.enable_block_reuse = kv_cache_config.enable_block_reuse + # With no eviction tier and a connector or the configuration asking for + # it, an admitted full-attention request must keep the pages it needs + # to finish: decode its output growth, a connector prefill its + # remaining prompt, so a final-batch restore is not truncated by other + # requests' chunks. Bounds are logical; no page is allocated early. + # The connector is bound above, so the gate can read its declaration. + self.enable_disagg_completion_reservation = ( + kv_cache_config.enable_disagg_completion_reservation + ) + self._disagg_completion_limits = {} + self._disagg_completion_pool_counts = {} + if self._disagg_completion_reservation_applies(): + # One slot per ledger block per life cycle, as in get_num_free_blocks. + slots_per_block: dict[int, int] = defaultdict(int) + for pool_id in self.impl.get_life_cycle_pool_group_indices(GPU_LEVEL): + slots_per_block[int(pool_id)] += 1 + self._disagg_completion_pool_counts = dict(slots_per_block) + elif ( + self.enable_disagg_completion_reservation + and not self.is_draft + and not self.is_estimating_kv_cache + ): + # Do not keep the existing admission silently. The common miss is a + # server without a connector that keeps the automatic host tier. + # The estimation manager is built without the connector, so only + # the final manager can tell whether the setting applies. + logger.warning( + "kv_cache_config.enable_disagg_completion_reservation is ignored: it " + "applies only to a disaggregated server with no cache tier below GPU " + "(host_cache_size=0 and no disk tier), beam width 1, a single pipeline " + "stage and full attention on every layer." + ) self.disk_prefetch_num_reqs = kv_cache_config.disk_prefetch_num_reqs enable_conversation_manager = ( self.enable_block_reuse @@ -3796,7 +3878,10 @@ def _build_base_config( # Keep the lookahead evidence and its backoff in the same tree match. # A paired scheduler caps the claim by retaining D evidence past the # common usable depth, so this still trims exactly once. + # The alignment is applied in the same tree match after the backoff, + # so the probe, the claim and the paired probes all trim once too. reuse_match_backoff=self.reuse_match_backoff, + reuse_match_alignment=self._connector_reuse_match_alignment(), enable_stats=self.enable_stats, swa_scratch_reuse=scratch_reuse_config, commit_min_snapshot=( @@ -3871,6 +3956,61 @@ def _build_cache_config(self, config: KVCacheManagerConfigPy) -> KVCacheManagerC """Customize the general cache config for a specialized cache manager.""" return config + def _kv_connector_declares(self, capability: str) -> bool: + """Whether the attached KV connector's worker declares ``capability`` exactly True. + + Read through ``getattr``: partially built managers may bind no connector, or a + connector without a worker, and only an explicit True counts, so a mock's truthy + attribute changes nothing. A capacity-only connector serves nothing, so it never + reshapes local reuse or admission. + + Args: + capability: Name of a ``KvCacheConnectorWorker`` declaration, such as + ``"requires_block_aligned_reuse"``. + + Returns: + True only when a serving connector's worker sets ``capability`` to True. + """ + connector = getattr(self, "kv_connector_manager", None) + if connector is None or getattr(connector, "capacity_only", False) is True: + return False + return getattr(getattr(connector, "worker", None), capability, False) is True + + def _connector_drops_reuse_match_backoff(self) -> bool: + """Whether the attached KV connector asks for local matches without backoff. + + A connector that serves from block boundaries can prefer reusing whole local + blocks to aligning a backed-off match down and loading the trimmed block itself; + see `_connector_reuse_match_alignment`. The draft KV at the end of such a block may + then belong to a different next token, which affects draft acceptance only. A + draft manager never drops its own backoff. + + Returns: + True when this target manager should drop the one-model draft backoff. + """ + return not self.is_draft and self._kv_connector_declares("drops_reuse_match_backoff") + + def _connector_reuse_match_alignment(self) -> int: + """Alignment a whole-block KV connector needs of every local prefix match. + + Such a connector serves only from a block boundary, and the local match is where + its offer starts. Disabling partial radix matching is not enough on its own: a + one-model draft's lookahead backoff trims a match of whole blocks to one or more + tokens short of a boundary, which would leave the connector nothing to offer for + any request with a local hit. Aligning trades the trimmed block for the connector + serving it and everything after it. Any other configuration keeps token-granular + reuse. + + Returns: + The ledger ``tokens_per_block`` when a backoff is in effect, partial reuse is + off and the connector requires block-aligned reuse; otherwise 1. + """ + if self.reuse_match_backoff <= 0 or self.enable_partial_reuse: + return 1 + if not self._kv_connector_declares("requires_block_aligned_reuse"): + return 1 + return self._ledger_tokens_per_block + def _remove_zero_size_buffers(self, config: KVCacheManagerConfigPy) -> KVCacheManagerConfigPy: """Exclude empty buffers before creating the runtime storage pools.""" if config.layers and all( @@ -4792,10 +4932,40 @@ def revert_allocate_generation(self, req: LlmRequest) -> None: ) self._after_shrink(request_id, kv_cache) + def drop_context_allocation(self, req: LlmRequest) -> None: + """Release a context request's allocation and rewind it to run from the start. + + For a request the scheduler skipped past a connector prefix whose pages + could not be filled (`KvCacheConnectorWorker.reports_failed_loads`). + Those positions already count as computed, so the request cannot run + this iteration. + + `free_resources` releases unstarted prefix loads and refuses a load that + is still pending, which turns a misreported failure into a RuntimeError + rather than a use-after-free. It resets the connector state and + `py_connector_allocation_reported` / `py_served_prefix_end`, and closes + the cache without committing it: blocks reach the radix tree only in + `try_commit_blocks` after a forward pass, so the never-filled range is + not offered for reuse. The next admission therefore matches local reuse + and asks the connector afresh. This is the same release as the drop + branches of `revert_allocate_context`. + + Args: + req: The context request to restart. + """ + self.free_resources(req) + rewind_context_after_cache_drop(req, self.tokens_per_block) + def revert_allocate_context(self, req: LlmRequest) -> bool: """Undo this iteration's context resize. False means the cache was dropped, not shrunk (history outran pre-resize capacity); the caller drops any draft pool. + + Also returns an unconsumed first chunk's completion bound, for a context + chunk dropped after scheduling or a disaggregated generation-init request + that DisaggCoordinator.revert_deferred_gen_init or the CP admission + alignment did not admit. """ + self._release_unconsumed_context_reservation(req) if self._connector_reservations_enabled(): if self.kv_connector_manager.get_prefix_reservation(req) is not None: self.free_resources(req) @@ -5173,6 +5343,7 @@ def resize_context(self, req: LlmRequest, num_tokens: int) -> bool: rewind_context_after_cache_drop(req, self.tokens_per_block) elif req.is_first_context_chunk: kv_cache.suspend() + self._release_unconsumed_context_reservation(req) return False self._fill_fresh_kv_pages(req.py_request_id) self._log_window_crossing(req, kv_cache, pre_cap, capacity, "context") @@ -5191,13 +5362,135 @@ def try_allocate_draft_context(self, req: LlmRequest, num_tokens: int) -> bool: target = req.context_current_position + num_tokens + self.num_extra_kv_tokens return kv_cache.resize(max(kv_cache.capacity, target)) + def _disagg_completion_reservation_applies(self) -> bool: + """Whether admission reserves each request's completion headroom. + + Opt-in: the attached KV connector's worker must declare + ``requires_completion_reservation``, or + ``kv_cache_config.enable_disagg_completion_reservation`` must be set; + the latter lets a disaggregated generation server without a connector + reserve its decode headroom. A server with neither keeps its existing + admission even without a tier below GPU. Beyond that, only a + disaggregated target manager without a tier below GPU, beam search or + pipeline parallelism, whose layers all use full attention, qualifies; + the configuration does not relax these conditions. Windowed and + recurrent caches keep their existing admission policy. PP is excluded + because each rank decides alone, and a follower refusing a + generation-init request that rank 0 admitted is not reconciled. + + Returns: + True when every condition holds. + """ + return ( + self.is_disagg + and not self.can_evict + and not self.is_draft + and self.max_beam_width == 1 + and self.mapping.pp_size == 1 + and ( + self._kv_connector_declares("requires_completion_reservation") + or self.enable_disagg_completion_reservation + ) + and all( + isinstance(layer, AttentionLayerConfig) + and ( + layer.sliding_window_size is None + or layer.sliding_window_size >= self.max_seq_len + ) + for layer in self.kv_cache_manager_py_config.layers + ) + ) + + def reserve_context_completion(self, req: LlmRequest) -> bool: + """Reserve a connector prefill's completion headroom before chunk allocation. + + Applies on a disaggregated full-attention server without an eviction + tier whose connector declares ``requires_completion_reservation`` or + whose configuration enables the reservation (see + `_disagg_completion_reservation_applies`), to the first chunk of a + request the connector may serve, and only + for connectors without source reservations: those are queried after the + final batch has allocated its chunks (``_run_kv_connector_hooks``), so + every admitted request's full restore must stay allocatable. No page is + allocated here. + + Args: + req: Context request about to allocate its first chunk. + + Returns: + False to defer ``req``: its bound does not fit beside the admitted + requests' bounds. The scheduler skips it instead of preempting, and + its rollback frees the claim it held. True otherwise, including + when no reservation applies. + """ + if not ( + self._disagg_completion_pool_counts + and req.is_first_context_chunk + and self._connector_may_serve(req) + and not self._connector_reservations_enabled() + ): + return True + if req.py_request_id in self._disagg_completion_limits: + return True + limit = self._disagg_completion_limit(req.prompt_len) + if not self._can_reserve_disagg_completion(req.py_request_id, limit): + return False + self._disagg_completion_limits[req.py_request_id] = limit + return True + + def _release_unconsumed_context_reservation(self, req: LlmRequest) -> None: + """Drop the bound of a first chunk that neither a load nor a forward consumed.""" + if ( + req.py_request_id in self._disagg_completion_limits + and req.is_first_context_chunk + and not req.py_connector_allocation_reported + ): + del self._disagg_completion_limits[req.py_request_id] + + def _disagg_generation_limit(self, req: LlmRequest, prompt_len: int) -> int: + """Completion bound of a disaggregated decode request, in ledger blocks.""" + return self._disagg_completion_limit(min(self.max_seq_len, prompt_len + req.max_new_tokens)) + + def _disagg_completion_limit(self, end: int) -> int: + """Ledger blocks to reach ``end`` plus the maximum generation lead.""" + capacity = end + self._generation_kv_capacity_headroom + blocks = (capacity + self._ledger_tokens_per_block - 1) // self._ledger_tokens_per_block + # Matching can share a partial tail page. Reserve its private copy even + # when growing that tail does not change the logical block count. + return blocks + int(self.enable_block_reuse) + + def _can_reserve_disagg_completion(self, request_id: int, limit: int) -> bool: + """Check every full-attention pool without allocating future KV pages. + + ``available`` is free plus evictable slots: it already excludes resident + prompts, generation KV and outgoing transfers, and counts droppable + reusable pages as free. Add only the unallocated part of each admitted + request's bound, including incoming transfers and connector first chunks + awaiting their restore. + """ + needed = max(0, limit - self.kv_cache_map[request_id].num_blocks) + for req_id, req_limit in self._disagg_completion_limits.items(): + if req_id == request_id: + continue + kv_cache = self.kv_cache_map.get(req_id) + if kv_cache is None: + # Freed without its bound: it holds no pages and will not grow. + continue + needed += max(0, req_limit - kv_cache.num_blocks) + stats = self._get_storage_statistics(GPU_LEVEL) + return all( + needed * count <= stats[pool_id].available + for pool_id, count in self._disagg_completion_pool_counts.items() + ) + def prepare_disagg_gen_init(self, req: LlmRequest) -> bool: """Prepare KV cache for a disagg generation init request. Allocates capacity for the full prompt (+ draft) and sets ``kv_cache.history_length`` to ``prompt_len``. Returns True on success, False if preparation or resize failed (cache is suspended - on resize failure). + on resize failure), or if the request's completion bound does not fit + beside the admitted requests' bounds (cache is freed, not suspended). """ reused = self.prepare_context_cache(req) if reused is None: @@ -5225,6 +5518,14 @@ def prepare_disagg_gen_init(self, req: LlmRequest) -> bool: else (req.total_input_len_cp if self._has_cp_helix else req.prompt_len) ) prompt_len = req.total_input_len_cp if self._has_cp_helix else req.prompt_len + limit = None + if self._disagg_completion_pool_counts: + limit = self._disagg_generation_limit(req, prompt_len) + if not self._can_reserve_disagg_completion(req.py_request_id, limit): + # No transfer has started. Release a newly matched prefix too: + # suspension alone cannot reclaim uncommitted KV without a tier. + self.free_resources(req) + return False target = prompt_len + get_draft_token_length(req) + self.num_extra_kv_tokens capacity = max(kv_cache.capacity, target) pre_cap = kv_cache.capacity @@ -5237,6 +5538,8 @@ def prepare_disagg_gen_init(self, req: LlmRequest) -> bool: self._fill_fresh_kv_pages(req.py_request_id) self._log_window_crossing(req, kv_cache, pre_cap, capacity, "disagg_gen_init") req.py_ctx_pre_resize_cap = pre_cap if capacity > pre_cap else None + if limit is not None: + self._disagg_completion_limits[req.py_request_id] = limit # RDMA cannot observe cache-stream ordering. Capture readiness here so # receive publication waits for admission's copies and recycled-slot # dependencies without waiting for later execution-stream work. @@ -5561,9 +5864,12 @@ def _run_kv_connector_hooks(self, scheduled_batch: ScheduledRequests) -> None: # Both forms are offered; the manager picks the one this connector # implements. With several layer groups the flat list is empty, # because a page index is scoped to a group. + started = time.perf_counter() by_group = self.get_page_indices_by_layer_group(request) flat = by_group[0] if len(by_group) == 1 else [] self.kv_connector_manager.update_state_after_alloc(request, flat, by_group) + if self._connector_offer_stats is not None: + self._connector_offer_stats["report_ms"] += (time.perf_counter() - started) * 1e3 if served_any: # A served prefix can carry a request past its last chunk boundary, @@ -5682,6 +5988,200 @@ def _accept_connector_prefix_reservations(self, scheduled_batch: ScheduledReques self.kv_connector_manager.update_state_after_alloc(req, flat, by_group) req.py_connector_allocation_reported = True + def begin_connector_admission_pass(self) -> None: + """Start a scheduling pass: admission-time answers from earlier passes go stale.""" + self._connector_admission_epoch += 1 + + def connector_admission_query_enabled(self) -> bool: + """Whether the scheduler may ask the connector while it admits a request. + + Off without an admission-capable connector (see + ``KvCacheConnectorManager.admission_query_capable``, which is off under + attention DP), on a draft manager, for a capacity-only role, and for a + connector that reserves its prefix during admission instead. + """ + return ( + self.kv_connector_manager is not None + and not self.is_draft + and getattr(self.kv_connector_manager, "admission_query_capable", False) is True + and not self.kv_connector_manager.capacity_only + and not self._connector_reservations_enabled() + ) + + def query_connector_at_admission(self, req: LlmRequest) -> bool: + """Ask the connector now, while the scheduler admits ``req``, and keep the answer. + + Called once the scheduler has sized and allocated ``req``'s chunk, with + the guards of ``_run_kv_connector_hooks``. The answer is the + allocation's one query: ``_apply_connector_matched_prefix`` uses it in + this pass's ``prepare_resources`` instead of asking again. An answer + left from an earlier pass is dropped first. + + Args: + req: The context request being admitted. + + Returns: + Whether the runtime will park ``req`` for an asynchronous load, so + the scheduler need not charge its chunk to the token budget. + """ + req.py_connector_admission_offer = None + if not self.connector_admission_query_enabled(): + return False + if ( + not self._connector_may_serve(req) + or not req.is_first_context_chunk + or req.py_connector_allocation_reported + or not self.kv_connector_manager.should_add_sequence(req) + ): + return False + kv_cache = self.kv_cache_map.get(req.py_request_id) + if kv_cache is None: + return False + local_end = kv_cache.num_committed_tokens + started = time.perf_counter() + num_tokens, load_async = self.kv_connector_manager.query_num_new_matched_tokens( + req, local_end + ) + if self._connector_offer_stats is not None: + self._connector_offer_stats["query_ms"] += (time.perf_counter() - started) * 1e3 + offer = ConnectorAdmissionOffer( + self._connector_admission_epoch, local_end, num_tokens, load_async + ) + req.py_connector_admission_offer = offer + return offer.parks + + def cover_admission_offer(self, req: LlmRequest) -> bool: + """Grow the allocation now to the window this pass's admission answer moves the chunk to. + + An offer that reaches past the scheduled chunk shifts the window and + grows the allocation by the offer. Left to ``prepare_resources``, that + grow competes with every admission after this one in the pass; when it + fails, ``_reserve_connector_prefix`` takes only the whole blocks the + pages already hold and the rest of the offer is recomputed. Growing + here books the pages before the requests behind it allocate. Only the + capacity moves: the history stays at the local match until + ``_reserve_connector_prefix`` places the served prefix, so a + sliding-window layer group holds the offer's pages until then. + + Args: + req: The context request just asked by ``query_connector_at_admission``. + + Returns: + False when the KV cache cannot cover the window; the scheduler then + gives the chunk back and asks again on a later pass. True when there + is nothing to grow (no answer from this pass, a miss, an offer + inside the chunk, a shift ``_reserve_connector_prefix`` would refuse + anyway, or SWA scratch reuse) or the grow succeeded. + """ + offer = req.py_connector_admission_offer + if offer is None or offer.epoch != self._connector_admission_epoch: + return True + if getattr(self, "enable_swa_scratch_reuse", False) is True: + # KVCacheV2Scheduler._pregrown_context_tokens reads the manager-level + # flag and replays any capacity grown past the cursor exactly, even + # for connector-servable caches whose per-cache scratch flag was + # cleared. A grow left on a request dropped before + # prepare_resources would come back as an offer+chunk-sized first + # chunk charged at that size, or be released and re-admitted. Leave + # the shift to prepare_resources, as without admission queries. + return True + kv_cache = self.kv_cache_map.get(req.py_request_id) + if kv_cache is None: + return True + local_end = offer.local_end + offered_end = min(local_end + offer.num_tokens, req.prompt_len - 1) + chunk_end = min(local_end + req.context_chunk_size, req.prompt_len) + if offered_end <= local_end or offered_end < chunk_end: + return True + end = self._shifted_chunk_end(req, offered_end) + if end <= offered_end: + return True + target = self._served_prefix_capacity(req, end) + pre_cap = kv_cache.capacity + if target <= pre_cap: + return True + grown = kv_cache.resize(target) + if self._connector_offer_stats is not None: + self._connector_offer_stats["pregrown" if grown else "pregrow_failed"] += 1 + if not grown: + return False + if req.py_ctx_pre_resize_cap is None: + # As in `_resize_for_served_prefix`: a revert must undo this growth. + req.py_ctx_pre_resize_cap = pre_cap + return True + + def enable_connector_offer_stats(self) -> None: + """Count connector offers per iteration (TLLM_CTX_PARK_STATS); read by `take_connector_offer_stats`.""" + self._connector_offer_stats = self._new_connector_offer_stats() + + @staticmethod + def _new_connector_offer_stats() -> Dict[str, float]: + return { + "queries": 0, + "admission_used": 0, + "offers": 0, + "async_offers": 0, + "offered_tokens": 0, + "honoured_tokens": 0, + "in_chunk": 0, + "shifted": 0, + "capped": 0, + "not_honoured": 0, + # Admission-time grows for a shifting offer (`cover_admission_offer`). + "pregrown": 0, + "pregrow_failed": 0, + # Host time (ms) in the leader query and its broadcast, in fitting the + # served prefix into the allocation, and in reporting the pages. + "query_ms": 0.0, + "reserve_ms": 0.0, + "report_ms": 0.0, + } + + def take_connector_offer_stats(self) -> Optional[Dict[str, float]]: + """This iteration's connector offer counters, reset on read. + + Returns: + The counters since the last call, or None while they are disabled. + """ + if self._connector_offer_stats is None: + return None + stats = self._connector_offer_stats + self._connector_offer_stats = self._new_connector_offer_stats() + return stats + + def _count_connector_offer( + self, + used_admission: bool, + num_tokens: int, + load_async: bool, + local_end: int, + offered_end: int, + chunk_end: int, + position: int, + ) -> None: + """Book one answered query into the TLLM_CTX_PARK_STATS counters. + + The categories follow ``_reserve_connector_prefix``: an offer inside the + scheduled chunk, the shifted window, or the capped fallback. + """ + stats = self._connector_offer_stats + stats["queries"] += 1 + stats["admission_used"] += int(used_admission) + if offered_end <= local_end: + return + stats["offers"] += 1 + stats["async_offers"] += int(load_async) + stats["offered_tokens"] += offered_end - local_end + stats["honoured_tokens"] += position - local_end + if position == local_end: + stats["not_honoured"] += 1 + elif offered_end < chunk_end: + stats["in_chunk"] += 1 + elif position == offered_end: + stats["shifted"] += 1 + else: + stats["capped"] += 1 + def _connector_may_serve(self, req: LlmRequest) -> bool: """Whether the connector is allowed to serve a prefix for ``req``.""" if self.kv_connector_manager is None or self.is_draft: @@ -5717,17 +6217,41 @@ def _apply_connector_matched_prefix(self, req: LlmRequest) -> bool: return False local_end = kv_cache.num_committed_tokens - num_tokens, load_async = self.kv_connector_manager.query_num_new_matched_tokens( - req, local_end - ) + offer = getattr(req, "py_connector_admission_offer", None) + req.py_connector_admission_offer = None + if offer is not None and offer.epoch == self._connector_admission_epoch: + # Asked while the scheduler admitted it, this same pass; asking + # again would break the one-query-per-allocation contract. + if offer.local_end != local_end: + raise RuntimeError( + f"req {req.py_request_id}: admission-time connector answer was " + f"taken at {offer.local_end} but the cache is now at {local_end}" + ) + num_tokens, load_async = offer.num_tokens, offer.load_async + used_admission = True + else: + started = time.perf_counter() + num_tokens, load_async = self.kv_connector_manager.query_num_new_matched_tokens( + req, local_end + ) + if self._connector_offer_stats is not None: + self._connector_offer_stats["query_ms"] += (time.perf_counter() - started) * 1e3 + used_admission = False offered_end = min(local_end + num_tokens, req.prompt_len - 1) + chunk_end = min(local_end + req.context_chunk_size, req.prompt_len) position = local_end + started = time.perf_counter() if offered_end > local_end: fitted = self._reserve_connector_prefix(req, kv_cache, local_end, offered_end) if fitted is not None: position, chunk_size = fitted self._mark_connector_prefix_populated(req, position, chunk_size) + if self._connector_offer_stats is not None: + self._connector_offer_stats["reserve_ms"] += (time.perf_counter() - started) * 1e3 + self._count_connector_offer( + used_admission, num_tokens, load_async, local_end, offered_end, chunk_end, position + ) # The honoured amount, not `num_tokens`: `build_scheduler_output` # reports `context_current_position - recorded` as the range computed @@ -5768,9 +6292,7 @@ def _reserve_connector_prefix( allocation, which is why that arm needs a fallback. """ tokens_per_block = self.tokens_per_block - prompt_len = req.prompt_len - scheduler_chunk = req.context_chunk_size - chunk_end = min(local_end + scheduler_chunk, prompt_len) + chunk_end = min(local_end + req.context_chunk_size, req.prompt_len) candidates = [] if offered_end < chunk_end: @@ -5778,13 +6300,9 @@ def _reserve_connector_prefix( else: # `==` belongs here, not above: an offer that reaches the chunk end # exactly leaves the forward pass nothing to compute unless the - # window moves. - shifted_end = min(offered_end + scheduler_chunk, prompt_len) - if shifted_end < prompt_len: - # A non-final chunk must end on a block boundary or the next - # chunk fragments the cache. - shifted_end = (shifted_end // tokens_per_block) * tokens_per_block - candidates.append((offered_end, shifted_end)) + # window moves. With an admission-time answer, + # `cover_admission_offer` has normally grown the pages already. + candidates.append((offered_end, self._shifted_chunk_end(req, offered_end))) # Fall back to the largest whole-block prefix the pages that already # exist can hold while still leaving a chunk to compute. Whole # blocks because that is the granularity a connector transfers in. @@ -5798,6 +6316,23 @@ def _reserve_connector_prefix( return position, end - position return None + def _shifted_chunk_end(self, req: LlmRequest, offered_end: int) -> int: + """End of the scheduled chunk once a served prefix moves its start to ``offered_end``.""" + prompt_len = req.prompt_len + shifted_end = min(offered_end + req.context_chunk_size, prompt_len) + if shifted_end < prompt_len: + # A non-final chunk must end on a block boundary or the next chunk + # fragments the cache. + shifted_end = (shifted_end // self.tokens_per_block) * self.tokens_per_block + return shifted_end + + def _served_prefix_capacity(self, req: LlmRequest, end: int) -> int: + """The capacity covering ``[0, end)`` for a context chunk ending at ``end``.""" + target = end + self.num_extra_kv_tokens + if end >= req.prompt_len: + target += get_draft_token_length(req) + return target + def _resize_for_served_prefix(self, req: LlmRequest, kv_cache, position: int, end: int) -> bool: """Cover ``[0, end)`` and set the history to ``position``, marking ``[0, position)`` as already valid, for KV served from outside the @@ -5811,9 +6346,7 @@ def _resize_for_served_prefix(self, req: LlmRequest, kv_cache, position: int, en sliding-window layer group. When the capacity is unchanged the call cannot fail, since no new slot is requested. """ - target = end + self.num_extra_kv_tokens - if end >= req.prompt_len: - target += get_draft_token_length(req) + target = self._served_prefix_capacity(req, end) pre_cap = kv_cache.capacity capacity = max(pre_cap, target) if not kv_cache.resize(capacity, position): @@ -7018,6 +7551,10 @@ def free_resources(self, request: LlmRequest, pin_on_release: bool = False): self.kv_connector_manager.release_prefix_reservation(request) # Replay obtains fresh connector metadata for its new destination pages. request.py_connector_allocation_reported = False + # An admission-time answer belongs to the allocation being released. + # `py_connector_offer_deferrals` is not reset: it lasts for the + # request's lifetime. + request.py_connector_admission_offer = None # Same allocation, same lifetime. The pages the served position vouches # for are gone, so a replay recomputes from its own reuse match. request.py_served_prefix_end = 0 @@ -7030,6 +7567,7 @@ def free_resources(self, request: LlmRequest, pin_on_release: bool = False): if self.conversation_manager is not None: self.conversation_manager.finish_request(request) self._allocated_draft_lens.pop(request.py_request_id, None) + self._disagg_completion_limits.pop(request.py_request_id, None) self._inflight_draft_lens.pop(request.py_request_id, None) self._request_stats_enabled_ids.discard(request.py_request_id) # The next owner of these pages fills them again; keeping the set would @@ -7362,6 +7900,7 @@ def shutdown(self): for kv_cache in self.kv_cache_map.values(): kv_cache.close() self.kv_cache_map.clear() + self._disagg_completion_limits.clear() self._disagg_receive_ready.clear() self._on_allocation_stream.clear() self._allocation_ready.clear() diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index a4e42209a2de..303263fdebd9 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -872,6 +872,19 @@ def __init__( # allocation clears this flag so replay can report its new pages. self.py_connector_allocation_reported = False + # The connector's answer for this allocation when the scheduler asked + # at admission (kv_connector_config.admission_query), so + # prepare_resources uses it instead of asking again. Cleared when read + # and when the allocation is released. + self.py_connector_admission_offer = None + # Scheduling iterations this request gave its first chunk back because + # the KV cache could not cover the window its admission-time connector + # offer moves the chunk to. Lasts for the request's lifetime. + self.py_connector_offer_deferrals = 0 + # Set when the scheduler gave up a started request to free pages for + # this one, so its next admission keeps those pages rather than defer. + self.py_preempted_for_pages = False + # End of the prefix served from outside the manager (a KV cache # connector's load, or a staging fetch the request resumes from) retained # by this allocation, or 0. This depth survives async parking; local reuse diff --git a/tensorrt_llm/_torch/pyexecutor/profiling.py b/tensorrt_llm/_torch/pyexecutor/profiling.py index e7bcbd5157fd..e3f70951028a 100644 --- a/tensorrt_llm/_torch/pyexecutor/profiling.py +++ b/tensorrt_llm/_torch/pyexecutor/profiling.py @@ -43,11 +43,13 @@ import functools import os +import threading import time import uuid import weakref +from collections import deque from contextlib import contextmanager -from typing import Dict, List, Optional +from typing import Dict, List, Optional, Tuple import torch @@ -62,6 +64,127 @@ PROFILE_START_STOP_ENV_VAR_NAME = "TLLM_PROFILE_START_STOP" PROFILE_TRACE_ENV_VAR_NAME = "TLLM_TORCH_PROFILE_TRACE" PROFILE_LOG_RANKS_ENV_VAR_NAME = "TLLM_PROFILE_LOG_RANKS" +# Debug-only suffixes of the iteration log line; see `ParkStats` and `GilProbe`. +PARK_STATS_ENV_VAR_NAME = "TLLM_CTX_PARK_STATS" +GIL_PROBE_ENV_VAR_NAME = "TLLM_GIL_PROBE_US" + + +class GilProbe: + """Measure how late a sleeping thread wakes, as a proxy for GIL contention. + + A daemon thread sleeps ``period_us`` in a loop and records how far past the + period each wake-up lands. A thread waiting for the GIL wakes late, so the + overshoot bounds how long other threads held it. Opt-in through + ``TLLM_GIL_PROBE_US`` together with ``TLLM_CTX_PARK_STATS``; the samples + are drained once per executor iteration and kept to the last ~60 s in case + nothing drains them. + """ + + def __init__(self, period_us: int): + """Start the probe thread. + + Args: + period_us: Sleep period of the probe thread, in microseconds. + """ + self._period = period_us / 1e6 + self._samples: deque = deque(maxlen=max(1, int(60.0 / self._period))) + self._stop = threading.Event() + self._thread = threading.Thread(target=self._run, name="gil-probe", daemon=True) + self._thread.start() + + def _run(self) -> None: + period = self._period + while not self._stop.is_set(): + start = time.perf_counter() + time.sleep(period) + self._samples.append((time.perf_counter() - start - period) * 1e3) + + def take(self) -> Optional[Tuple[float, float, int]]: + """Drain the samples taken since the last call. + + Returns: + The maximum and p99 wake-up overshoot in ms and the sample count, + or None when there is no sample. + """ + # popleft never races the probe's append: a late sample stays for the next call. + samples = [self._samples.popleft() for _ in range(len(self._samples))] + if not samples: + return None + samples.sort() + return samples[-1], samples[min(len(samples) - 1, int(0.99 * len(samples)))], len(samples) + + def stop(self) -> None: + """Ask the probe thread to exit after its current sleep.""" + self._stop.set() + + +class ParkStats: + """Per-iteration context-park diagnostics (opt-in through ``TLLM_CTX_PARK_STATS=1``). + + Wall and executor-thread CPU time of each loop phase, the requests parked + for an asynchronous KV connector load and their chunk tokens, the requests + admitted without a token charge, the context tokens left for the forward + pass, and the first-chunk requests the scheduler passed over for lack of + budget. Printed at the end of the iteration log line so existing parsers + keep working. + """ + + PHASES = ("sched", "prep", "meta", "start", "fwd", "upd", "term") + + def __init__(self): + self.reset() + + def reset(self) -> None: + """Zero every counter for the next iteration.""" + self.wall_ms = dict.fromkeys(self.PHASES, 0.0) + self.cpu_ms = dict.fromkeys(self.PHASES, 0.0) + self.parked_reqs = 0 + self.parked_tokens = 0 + self.uncharged = 0 + self.fwd_ctx_tokens = 0 + self.budget_skips = 0 + self.offer: Optional[Dict[str, float]] = None + + @contextmanager + def phase(self, name: str): + """Add the wall and thread CPU time of the block to phase ``name``.""" + wall, cpu = time.perf_counter(), time.thread_time() + try: + yield + finally: + self.wall_ms[name] += (time.perf_counter() - wall) * 1e3 + self.cpu_ms[name] += (time.thread_time() - cpu) * 1e3 + + def format(self, gil: Optional[Tuple[float, float, int]]) -> str: + """Render the ``ctx_park`` suffix of the iteration log line. + + Flat ``key=value`` tokens with no braces, so parsers of the + ``states = {...}`` field before it still match. + + Args: + gil: A ``GilProbe.take`` result, or None. + + Returns: + The suffix, starting with ``", ctx_park = "``. + """ + parts = [ + f"parked={self.parked_reqs}", + f"parked_tok={self.parked_tokens}", + f"unc={self.uncharged}", + f"fwd_tok={self.fwd_ctx_tokens}", + f"bskip={self.budget_skips}", + ] + parts += [ + f"{name}_ms={self.wall_ms[name]:.1f}/{self.cpu_ms[name]:.1f}" for name in self.PHASES + ] + if self.offer is not None: + parts += [ + f"offer_{key}={value:.1f}" if isinstance(value, float) else f"offer_{key}={value}" + for key, value in self.offer.items() + ] + if gil is not None: + parts.append(f"gil={gil[0]:.2f}/{gil[1]:.2f}/{gil[2]}") + return ", ctx_park = " + " ".join(parts) @functools.cache @@ -779,6 +902,7 @@ def profile_step_fn(): f"num_fetched = {num_fetched}, " f"num_dummy = {executor.num_dummy_requests}, " f"num_unscheduled = {executor.num_unscheduled_requests}" + f"{executor._format_park_stats()}" ) # profile_step runs at the start of each executor loop after @@ -786,6 +910,8 @@ def profile_step_fn(): # batch records its ADP dummy work. executor._iter_adp_dummy_ctx_tokens = 0 executor._iter_adp_dummy_gen_tokens = 0 + # Same cadence, including passes whose line is not printed. + executor._reset_park_stats() it += 1 if executor.iter_counter in executor.profile_start_iters and not executor.is_warmup: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index e0502b4cd181..a7c53d85b24d 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -113,7 +113,9 @@ from .perf_metrics_manager import PerfMetricsManager from .pp_utils import (MpiProgressPump, PPCommTag, get_with_mpi_progress, make_mpi_progress_exit_hook, make_mpi_progress_pump) -from .profiling import PROFILE_START_STOP_ENV_VAR_NAME, PyExecutorProfileManager +from .profiling import (GIL_PROBE_ENV_VAR_NAME, PARK_STATS_ENV_VAR_NAME, + PROFILE_START_STOP_ENV_VAR_NAME, GilProbe, ParkStats, + PyExecutorProfileManager) from .profiling import load_iteration_indexes as _load_iteration_indexes from .request_utils import (RequestBroadcaster, attach_py_objects_to_requests, derive_attention_dp_per_rank_request_cap, @@ -162,6 +164,41 @@ def _stats_buffer_is_unbounded(max_stats_len: int) -> bool: _failed_kv_connector_executors: dict[int, "PyExecutor"] = {} +def _validate_connector_admission(connector: Optional[KvCacheConnectorManager], + request: LlmRequest) -> None: + """Refuse work a speculative-prefill-only KV connector cannot serve. + + Such a connector runs next to one-model speculation on context-only + requests alone, and its manager raises inside resource preparation for + anything else, which takes the whole executor down (and, under attention + DP, the connector fail-stop). Refusing at admission fails only this + request, e.g. a /health_generate probe sent straight to a disaggregated + context server. + + A module-level function rather than a method, so the minimal executors + unit tests assemble need no stub for it. + + Args: + connector: The executor's connector manager, or None. + request: The request being admitted. + + Raises: + ValueError: When the connector runs in the speculative-prefill-only + mode, transfers KV, and the request is not context-only. + """ + if connector is None: + return + if getattr(connector, "speculative_prefill_only", False) is not True: + return + if getattr(connector, "capacity_only", False) is True: + return + if getattr(request, "is_context_only_request", False) is not True: + raise ValueError( + "This server's KV cache connector runs next to one-model " + "speculative decoding and accepts context-only " + "(disaggregated prefill) requests only.") + + @dataclasses.dataclass class _KvConnectorTransferDeadline: expires_at: float @@ -439,6 +476,15 @@ class PyExecutor: # If the number of micro batches is too large, the executor will spend too much host memory (No additional GPU memory is required). # 1024 in-flight micro batches can avoid synchronization in most cases and keep host memory usage low. MIN_ASYNC_MICRO_BATCH_NUM = 1024 + # Declared on the class for executors built without __init__, which the + # connector bring-up and shutdown paths still reach. + _kv_connector_config = None + # The scheduler that runs admission-time connector queries, or None; see + # `_resolve_admission_query_scheduler`. + _admission_query_scheduler: Optional[KVCacheV2Scheduler] = None + # Debug-only context-park diagnostics; see `ParkStats` and `GilProbe`. + _park_stats: Optional[ParkStats] = None + _gil_probe: Optional[GilProbe] = None def __init__( self, @@ -593,6 +639,14 @@ def __init__( None) self.guided_decoder = guided_decoder self.disable_overlap_scheduler = disable_overlap_scheduler + # Debug-only context-park diagnostics appended to the iteration log + # line (TLLM_CTX_PARK_STATS=1); the GIL probe reports on that suffix, + # so it runs only with it. + self._park_stats = (ParkStats() if os.environ.get( + PARK_STATS_ENV_VAR_NAME, "0") == "1" else None) + gil_probe_us = int(os.environ.get(GIL_PROBE_ENV_VAR_NAME, "0") or 0) + self._gil_probe = (GilProbe(gil_probe_us) if gil_probe_us > 0 + and self._park_stats is not None else None) self.enable_kv_pool_rebalance = enable_kv_pool_rebalance # Iteration throttle for the KV pool rebalance check. See # _can_pause_for_rebalance / _agreed_need_adjustment. @@ -720,6 +774,9 @@ def __init__( # kv cache events self.kv_cache_manager = self.resource_manager.resource_managers.get( ResourceManagerType.KV_CACHE_MANAGER) + if self._park_stats is not None and hasattr( + self.kv_cache_manager, "enable_connector_offer_stats"): + self.kv_cache_manager.enable_connector_offer_stats() self.draft_kv_cache_manager = self.resource_manager.resource_managers.get( ResourceManagerType.DRAFT_KV_CACHE_MANAGER) self._is_kv_manager_v2 = isinstance(self.kv_cache_manager, @@ -859,6 +916,7 @@ def __init__( kv_cache_manager=self.kv_cache_manager, attention_dp_config=self.llm_args.attention_dp_config, async_transfer_manager=self.async_transfer_manager, + kv_connector_manager=kv_connector_manager, ) # The router publishes the uncached-prefill estimate only when the # scheduler's long-prefill demotion consumes it. @@ -1239,6 +1297,18 @@ def _kv_connector_moves_kv(self) -> bool: connector = getattr(self, "kv_connector_manager", None) return connector is not None and not connector.capacity_only + @property + def _kv_connector_reports_failed_loads(self) -> bool: + """Whether the connector hands failed prefix loads back for a restart. + + See `KvCacheConnectorWorker.reports_failed_loads`. Read through + `getattr` because minimal unit-test executors use + `Mock(spec=KvCacheConnectorManager)`, where the instance attribute is + absent. + """ + connector = getattr(self, "kv_connector_manager", None) + return getattr(connector, "reports_failed_loads", False) is True + def _maybe_init_kv_connector_manager(self) -> None: if self.kv_connector_manager is not None: # A connector that moves no KV cannot contend with the transceiver @@ -1273,13 +1343,34 @@ def _maybe_init_kv_connector_manager(self) -> None: is_kv_cache_manager_v2 = isinstance(self.kv_cache_manager, KVCacheManagerV2) + # Compared against True like `capacity_only`, because a mocked + # connector answers every attribute truthily. + speculative_prefill_only = getattr(self.kv_connector_manager, + "speculative_prefill_only", + False) is True + + if speculative_prefill_only and not is_kv_cache_manager_v2: + # The creator sets the mode for an explicit V2 selection or + # "auto" only, so this is "auto" resolving to V1. + raise NotImplementedError( + "KV connector speculative-prefill-only mode requires " + "KVCacheManagerV2. Set " + "kv_cache_config.use_kv_cache_manager_v2=True.") if self.enable_attention_dp and not is_kv_cache_manager_v2: raise NotImplementedError( "KV connector attention data parallelism requires KVCacheManagerV2. " "Set kv_cache_config.use_kv_cache_manager_v2=True.") - if is_kv_cache_manager_v2 and self.max_draft_len > 0: + if (self._kv_connector_reports_failed_loads + and not is_kv_cache_manager_v2): + raise NotImplementedError( + "A KV connector that reports failed prefix loads requires " + "KVCacheManagerV2, which restarts the affected requests. " + "Set kv_cache_config.use_kv_cache_manager_v2=True.") + + if (is_kv_cache_manager_v2 and self.max_draft_len > 0 + and not speculative_prefill_only): raise NotImplementedError( "KV Cache Connector is not supported with speculative " "decoding. Rejected draft tokens " @@ -1287,7 +1378,10 @@ def _maybe_init_kv_connector_manager(self) -> None: "whichever request allocates next. The connector is only " "told about pages appended since the last report, so it " "would keep addressing a slot another request now owns. " - "Disable speculative decoding to run a connector.") + "One-model MTP or EAGLE3 is supported on a disaggregated " + "context server when the connector scheduler and worker " + "both declare supports_speculative_prefill=True. Otherwise " + "disable speculative decoding to run a connector.") scheduler_config = getattr(self.llm_args, 'scheduler_config', None) capacity_scheduler_policy = getattr(scheduler_config, @@ -1339,10 +1433,24 @@ def _maybe_init_kv_connector_manager(self) -> None: "per-layer load/save hooks have nothing meaningful to " "transfer for those layers.") + # The creator keeps one-model draft layers in the target pool on a + # disaggregated server, so this is a backstop. + if (speculative_prefill_only and getattr( + self, "draft_kv_cache_manager", None) is not None): + raise NotImplementedError( + "KV Cache Connector with one-model speculative prefill " + "requires the draft layers to share the target KV cache. " + "A separate draft KV cache is not registered with the " + "connector, so a restored prefix would carry stale draft " + "KV.") + self.kv_connector_manager.configure_prefix_reservations( is_kv_cache_manager_v2 and (scheduler_config is None or scheduler_config.enable_prefix_aware_scheduling)) + # After the reservation decision, which excludes admission queries. + self._admission_query_scheduler = ( + self._resolve_admission_query_scheduler()) if is_kv_cache_manager_v2: # A registered region is only a valid address while its page is @@ -2064,6 +2172,8 @@ def _shutdown(self, *, release_model_owned_graphs: bool) -> None: # every shutdown path. A no-op when the executor loop already # unwound its own profiling state. self._profile_manager.cleanup() + if self._gil_probe is not None: + self._gil_probe.stop() if self.hang_detector.detected(): # Early return here to avoid waiting for hanging threads. # Since `on_detected` has sent the error message as response, @@ -4585,6 +4695,145 @@ def _release_unused_connector_reservations( } if scheduled_batch is not None else set()) self.kv_cache_manager.release_unused_connector_reservations(accepted) + def _resolve_admission_query_scheduler( + self) -> Optional[KVCacheV2Scheduler]: + """Find the scheduler that runs admission-time connector queries. + + ``kv_connector_config.admission_query`` is ignored, with a warning, + where the connector cannot be asked at admission. Read through + ``getattr`` because executors built without ``__init__`` reach this. + + Returns: + The KVCacheV2Scheduler, under any wrapper, when admission queries + are configured and usable; None otherwise. + + Raises: + NotImplementedError: With the overlap scheduler, or without the + KVCacheV2Scheduler. + """ + config = getattr(getattr(self, "_kv_connector_config", None), + "admission_query", None) + if config is None: + return None + connector = self.kv_connector_manager + reason = None + if getattr(connector, "enable_attention_dp", False) is True: + reason = "attention data parallelism" + elif getattr(connector, "capacity_only", False) is True: + reason = "a capacity-only connector role" + elif getattr(connector, "prefix_reservations_enabled", False) is True: + reason = "a connector that reserves its prefix during admission" + elif getattr(connector, "admission_query_capable", False) is not True: + reason = ("a connector scheduler that does not declare " + "supports_admission_query") + if reason is not None: + logger.warning( + f"kv_connector_config.admission_query has no effect with {reason}." + ) + return None + if not self.disable_overlap_scheduler: + raise NotImplementedError( + "kv_connector_config.admission_query requires " + "disable_overlap_scheduler=True") + scheduler = self._find_kv_cache_v2_scheduler() + if scheduler is None: + raise NotImplementedError( + "kv_connector_config.admission_query requires KVCacheManagerV2 " + "and its KVCacheV2Scheduler, got " + f"{type(getattr(self, 'scheduler', None)).__name__}") + return scheduler + + def _find_kv_cache_v2_scheduler(self) -> Optional[KVCacheV2Scheduler]: + """The KVCacheV2Scheduler itself or under wrappers that keep it as ``.scheduler``.""" + scheduler = getattr(self, "scheduler", None) + # Wrappers nest one or two deep; the bound only stops a cycle. + for _ in range(8): + if scheduler is None or isinstance(scheduler, KVCacheV2Scheduler): + return scheduler + scheduler = getattr(scheduler, "scheduler", None) + return None + + def _check_uncharged_parks_left_batch( + self, scheduled_batch: ScheduledRequests) -> None: + """Fail-stop when a request admitted with a 0-token charge did not park. + + One left in the batch would run a chunk the token budget never counted, + so the forward pass could exceed max_num_tokens. Other parked requests' + asynchronous loads are already registered and may own destination + pages, so this fails through `_kv_connector_fail`, which retains them. + + Args: + scheduled_batch: The batch after the connector took its parked + requests out. + """ + scheduler = self._admission_query_scheduler + if scheduler is None or not scheduler.uncharged_park_ids: + return + unparked = [ + req.py_request_id for req in scheduled_batch.context_requests + if req.py_request_id in scheduler.uncharged_park_ids + ] + if unparked: + self._kv_connector_fail( + f"Context requests {unparked} were admitted with no token " + "charge for an asynchronous connector load but did not park") + + def _park_phase(self, name: str): + """Time one executor-loop phase for TLLM_CTX_PARK_STATS, else a no-op.""" + stats = self._park_stats + return stats.phase(name) if stats is not None else nullcontext() + + def _format_park_stats(self) -> str: + """The ``ctx_park`` suffix of the iteration log line, or ''.""" + stats = self._park_stats + if stats is None: + return "" + scheduler = self._find_kv_cache_v2_scheduler() + if scheduler is not None: + stats.budget_skips = scheduler.last_pass_ctx_budget_skips + take_offers = getattr(self.kv_cache_manager, + "take_connector_offer_stats", None) + stats.offer = take_offers() if take_offers is not None else None + probe = self._gil_probe + return stats.format(probe.take() if probe is not None else None) + + def _reset_park_stats(self) -> None: + """Start the next iteration's ctx_park counters on every rank.""" + stats = self._park_stats + if stats is None: + return + stats.reset() + if self._gil_probe is not None: + self._gil_probe.take() + # Ranks that print no log line drain the offer counters here. + take_offers = getattr(self.kv_cache_manager, + "take_connector_offer_stats", None) + if take_offers is not None: + take_offers() + + def _record_parked_chunks(self, chunks_before: Dict[int, int], + batch: ScheduledRequests) -> None: + """Book the context chunks the connector parked into the ctx_park counters. + + Args: + chunks_before: Chunk size of each context request before the + connector took its parked requests out of the batch. + batch: The batch after the connector took them out. + """ + stats = self._park_stats + parked = dict(chunks_before) + for req in batch.context_requests: + stats.fwd_ctx_tokens += parked.pop(req.py_request_id, 0) + stats.parked_reqs += len(parked) + stats.parked_tokens += sum(parked.values()) + scheduler = self._admission_query_scheduler + if scheduler is not None: + # Only those that reached the connector: a request admitted + # without a charge can be dropped before prepare_resources. + stats.uncharged += sum( + 1 for request_id in scheduler.uncharged_park_ids + if request_id in chunks_before) + def _kv_connector_start_batch( self, scheduled_batch: ScheduledRequests) -> tuple[bool, bool]: """Prepare connector work and return fleet/local compute readiness.""" @@ -4611,6 +4860,20 @@ def _kv_connector_start_batch( if error is None: try: self._kv_connector_terminate_requests(synchronize=False) + if self._kv_connector_reports_failed_loads: + # A parked request whose load failed went back to the + # scheduler without an allocation. Nothing here waits + # for it and it must not rejoin this batch. Its freed + # pages may now let padding succeed where it failed + # before parking. + pending_before = len(loading_context) + loading_context = [ + req for req in loading_context + if not self._is_restarted_connector_load(req) + ] + if pending_before and not loading_context: + self._pad_attention_dp_batch_emptied_by_connector( + scheduled_batch) # Synchronous prefix loads remain in the prepared batch. # Completion can cancel them, so remove them before the # status vote flushes responses and frees their KV pages. @@ -4648,10 +4911,7 @@ def _kv_connector_start_batch( # without allocating surviving contexts' pages again. self.kv_connector_manager.build_scheduler_output( scheduled_batch, self.kv_cache_manager) - self.kv_connector_manager.handle_metadata() - self.kv_connector_manager.mark_prefix_loads_dispatched() - self.kv_connector_manager.worker.start_load_kv( - torch.cuda.current_stream()) + self._kv_connector_start_loads(scheduled_batch) except Exception as exc: error = f"{type(exc).__name__}: {exc}" phase_idle = scheduled_batch.is_attention_dp_phase_idle @@ -4756,28 +5016,167 @@ def _kv_connector_sync_status( def _kv_connector_prepare_batch( self, scheduled_batch: ScheduledRequests) -> list[LlmRequest]: - if self._kv_connector_moves_kv: - loading_context = ([ - req for req in scheduled_batch.context_requests - if self.kv_connector_manager.has_pending_load(req) - ] if self.enable_attention_dp else []) - self.kv_connector_manager.take_scheduled_requests_pending_load( - scheduled_batch) - if self.enable_attention_dp and scheduled_batch.batch_size == 0: - # Async restores can empty a batch after the usual ADP padding - # and resource preparation. Keep ready owners computing while - # this owner's worker makes progress on the restore. - self._pad_empty_attention_dp_batch(scheduled_batch) - if scheduled_batch.batch_size: - self.resource_manager.prepare_resources( - scheduled_batch, publish_connector_output=False) - self.kv_connector_manager.handle_metadata() - self.kv_connector_manager.mark_prefix_loads_dispatched() - self.kv_connector_manager.worker.start_load_kv( - torch.cuda.current_stream()) + if not self._kv_connector_moves_kv: + return [] + loading_context = ([ + req for req in scheduled_batch.context_requests + if self.kv_connector_manager.has_pending_load(req) + ] if self.enable_attention_dp else []) + stats = self._park_stats + if stats is not None: + chunks_before = { + req.py_request_id: req.context_chunk_size + for req in scheduled_batch.context_requests + } + self.kv_connector_manager.take_scheduled_requests_pending_load( + scheduled_batch) + self._check_uncharged_parks_left_batch(scheduled_batch) + if stats is not None: + self._record_parked_chunks(chunks_before, scheduled_batch) + self._pad_attention_dp_batch_emptied_by_connector(scheduled_batch) + self._kv_connector_start_loads(scheduled_batch) + return loading_context + + def _kv_connector_start_loads(self, + scheduled_batch: ScheduledRequests) -> None: + """Bind this pass's connector plan, start its loads and restart failed ones. - return loading_context - return [] + Args: + scheduled_batch: The batch the plan was built for. Context requests + whose synchronous load the connector reports as failed leave it. + """ + self.kv_connector_manager.handle_metadata() + self.kv_connector_manager.mark_prefix_loads_dispatched() + self.kv_connector_manager.worker.start_load_kv( + torch.cuda.current_stream()) + if self._kv_connector_reports_failed_loads: + failed = self.kv_connector_manager.take_failed_load_requests() + if failed: + self._recover_failed_connector_loads(scheduled_batch, failed) + + def _pad_attention_dp_batch_emptied_by_connector( + self, scheduled_batch: ScheduledRequests) -> None: + """Keep an attention-DP owner in the collective forward pass. + + Async restores, and restarts after a failed prefix load, can empty a + batch after the usual ADP padding and resource preparation. Keep ready + owners computing while this owner's worker makes progress. The padding + dummy joins after resource preparation already ran for the batch, so it + is prepared here without reporting the batch to the connector again: a + second report would replace the plan about to be bound. + + Args: + scheduled_batch: The batch to pad when it is empty. + """ + if not self.enable_attention_dp or scheduled_batch.batch_size != 0: + return + self._pad_empty_attention_dp_batch(scheduled_batch) + if scheduled_batch.batch_size: + self.resource_manager.prepare_resources( + scheduled_batch, publish_connector_output=False) + + def _recover_failed_connector_loads(self, + scheduled_batch: ScheduledRequests, + failed_ids: set[int]) -> None: + """Restart the context requests whose connector prefix could not be loaded. + + The connector offered a stored prefix, the scheduler skipped the + request past it, and the pages were then not there to load (evicted + from a full pool, or never completely written). Those positions already + count as computed, so the request cannot take part in this forward + pass. Its allocation and progress are dropped and it stays active: the + scheduler admits it again from local reuse, and the connector is asked + afresh. Restarts are rank-local, so only attention-DP owners and + single-rank servers report failures. + + Args: + scheduled_batch: The batch whose loads were just started. + failed_ids: The request IDs the connector reported. + + Raises: + RuntimeError: When a reported ID is not in this pass's context + batch. + """ + batch_ids = {req.request_id for req in scheduled_batch.context_requests} + unknown = failed_ids - batch_ids + if unknown: + raise RuntimeError( + f"KV connector reported failed prefix loads for requests " + f"{sorted(unknown)} that are not in this pass's context batch") + restarted = [] + for key in ("context_requests_chunking", "context_requests_last_chunk"): + kept = [] + for req in getattr(scheduled_batch, key): + (restarted + if req.request_id in failed_ids else kept).append(req) + setattr(scheduled_batch, key, kept) + for req in restarted: + self._restart_failed_connector_load(req) + logger.warning( + f"kv connector: restarted {len(restarted)} context request(s) " + f"after a failed prefix load: " + f"{[req.request_id for req in restarted]}") + self._pad_attention_dp_batch_emptied_by_connector(scheduled_batch) + + def _restart_failed_connector_load(self, req: LlmRequest) -> None: + """Release a request's allocation and progress so it is admitted afresh. + + The bound saves go first, while the metadata is still bound. Resource + preparation preceded the load, so a speculative-decoding manager + already owns a slot for the request; it is released too, and the next + first-chunk preparation allocates it again. The sequence slot stays + assigned, as for a parked request. + + Args: + req: The context request to restart. + """ + self.kv_connector_manager.forget_request(req) + self.kv_cache_manager.drop_context_allocation(req) + spec_manager = self.resource_manager.get_resource_manager( + ResourceManagerType.SPEC_RESOURCE_MANAGER) + if spec_manager is not None: + spec_manager.free_resources(req) + + def _recover_failed_async_loads(self) -> None: + """Restart the parked requests whose asynchronous prefix load came up short. + + The same outcome as `_recover_failed_connector_loads`, reached from a + different place: the request was parked outside the batch while its + pages were fetched, `get_finished` has just moved it back to context + state, and the bytes it was skipped past are not all there. Its + allocation and progress are dropped before the scheduler can admit it. + """ + requests = self.kv_connector_manager.take_failed_async_load_requests() + for req in requests: + self._restart_failed_connector_load(req) + # Set explicitly so the request stays schedulable whichever path + # moved it out of its parked state. + req.state = LlmRequestState.CONTEXT_INIT + if requests: + logger.warning( + f"kv connector: restarted {len(requests)} parked context " + f"request(s) after a failed asynchronous prefix load: " + f"{[req.request_id for req in requests]}") + + def _is_restarted_connector_load(self, req: LlmRequest) -> bool: + """Whether a parked request was handed back for a restart. + + A parked request whose load completed sits in + `finished_async_loading_requests`, so `should_add_sequence` is False; a + restarted one was popped from there and has no allocation; one still + loading is in DISAGG_GENERATION_TRANS_IN_PROGRESS; a cancelled one is + finished. Valid only when the connector reports failed loads: prefix + reservations are then refused, so a parked request is always a legacy + asynchronous load. + + Args: + req: A request that was parked for an asynchronous load. + + Returns: + True when ``req`` was restarted and must not rejoin the batch. + """ + return (req.state == LlmRequestState.CONTEXT_INIT + and self.kv_connector_manager.should_add_sequence(req)) def _kv_connector_terminate_requests( self, @@ -4797,6 +5196,11 @@ def _kv_connector_terminate_requests( for req in self.kv_connector_manager.take_finished_load_terminations( ): self._finish_connector_load_termination(req) + if self._kv_connector_reports_failed_loads: + # After the terminations: a parked load whose cancellation was + # deferred finished as a termination, so its failure names + # nothing left to restart. + self._recover_failed_async_loads() except Exception as exc: if not self.enable_attention_dp or not synchronize: raise @@ -5225,7 +5629,9 @@ def _executor_loop(self) -> None: self.disagg.handle_errors_synced() - scheduled_batch, iter_stats = self._prepare_and_schedule_batch() + with self._park_phase("sched"): + scheduled_batch, iter_stats = self._prepare_and_schedule_batch( + ) if scheduled_batch is None: # handle_errors_synced() can buffer a non-fatal response @@ -5288,16 +5694,19 @@ def _executor_loop(self) -> None: self._maybe_record_hang_diagnostic_phase( "preparing_resources", scheduled_batch) - self.resource_manager.prepare_resources(scheduled_batch) + with self._park_phase("prep"): + self.resource_manager.prepare_resources(scheduled_batch) if self._kv_connector_moves_kv: - self.kv_connector_manager.handle_metadata() + with self._park_phase("meta"): + self.kv_connector_manager.handle_metadata() if can_queue: # Connector preparation can change the batch. Its status # vote also supplies the updated fleet queue decision. - can_queue, _ = self._kv_connector_start_batch( - scheduled_batch) + with self._park_phase("start"): + can_queue, _ = self._kv_connector_start_batch( + scheduled_batch) elif self._kv_connector_moves_kv: can_queue, _ = self._can_queue(scheduled_batch) @@ -5366,7 +5775,9 @@ def _executor_loop(self) -> None: gpu_forward_end) as fwd_timing: if self.dwdp_manager is not None: self.dwdp_manager.prefetch_first_layers() - batch_outputs = self._forward_step(scheduled_batch) + with self._park_phase("fwd"): + batch_outputs = self._forward_step( + scheduled_batch) # None: _forward_step hit an exception and # _handle_errors already failed the affected requests @@ -5412,8 +5823,9 @@ def _executor_loop(self) -> None: scheduled_batch, spec_metadata) self._update_request_states(scheduled_batch) - self._update_requests(sample_state, - self.resource_manager) + with self._park_phase("upd"): + self._update_requests(sample_state, + self.resource_manager) if self.speculation_gate is not None: self._update_batch_acceptance_rate( scheduled_batch, @@ -5463,7 +5875,8 @@ def _executor_loop(self) -> None: only_with_context_sends=True) if not self.enable_attention_dp: - self._kv_connector_terminate_requests() + with self._park_phase("term"): + self._kv_connector_terminate_requests() # This is the one ADP-synchronized response flush for a # completed executor-loop pass. It is deliberately outside @@ -6737,6 +7150,8 @@ def _configure_csa2_remote_tail(self, request: LlmRequest) -> None: request.py_csa2_remote_tail_split = split def _validate_request(self, request: LlmRequest): + _validate_connector_admission( + getattr(self, "kv_connector_manager", None), request) self._configure_csa2_remote_tail(request) retention = request.kv_cache_retention_config if self._is_kv_manager_v2 and retention is not None: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index f05c2f5504c0..e99d4cbb26d1 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -22,7 +22,8 @@ from tensorrt_llm._torch.pyexecutor.resource_manager import ( BaseResourceManager, ResourceManagerType) from tensorrt_llm._utils import get_sm_version -from tensorrt_llm.llmapi.llm_args import (CapacitySchedulerPolicy, +from tensorrt_llm.llmapi.llm_args import (CacheTransceiverConfig, + CapacitySchedulerPolicy, ContextChunkingPolicy, ExecutorMemoryType, GuidedDecodingConfig, KvCacheConfig, @@ -225,6 +226,47 @@ def _flashinfer_one_engine_spec_supported(attn_backend: str, return not should_use_separate_draft_kv_cache(spec_config) +def _connector_runs_speculative_prefill_only( + spec_config: Optional[SpeculativeConfig], worker_cls: type, + scheduler_cls: Optional[type], + cache_transceiver_config: Optional[CacheTransceiverConfig], + use_kv_cache_manager_v2: bool | str) -> bool: + """Whether a KV connector may run next to this speculative decoding. + + Rejected draft tokens shrink a request's page list underneath the + connector, which is why speculative decoding is refused next to one. A + context-only request on a disaggregated context server never rejects a + draft token, so the append-only page deltas the connector receives stay + valid there. One-model MTP and EAGLE3 keep their draft KV in the target + pool, which is registered with the connector. Generation work is refused + at admission and at the connector boundary + (`KvCacheConnectorManager.speculative_prefill_only`). + + Args: + spec_config: The resolved speculative config, or None. + worker_cls: The connector worker class. + scheduler_cls: The connector scheduler class. + cache_transceiver_config: The cache transceiver config; a backend + marks a disaggregated server. + use_kv_cache_manager_v2: The tri-state + `kv_cache_config.use_kv_cache_manager_v2` (True, False or "auto"). + + Returns: + True when the connector runs in the speculative-prefill-only mode. + """ + if use_kv_cache_manager_v2 is False: + # V1 never refused speculation next to a connector; keep it unchanged. + return False + if spec_config is None or (spec_config.max_draft_len or 0) <= 0: + return False + mode = spec_config.spec_dec_mode + if not (mode.is_mtp_one_model() or mode.is_eagle3_one_model()): + return False + return (is_disagg_enabled(cache_transceiver_config) + and KvCacheConnectorManager.declares_speculative_prefill( + worker_cls, scheduler_cls)) + + ALLOW_UNCORRECTED_OVERLAP_SPEC_ENV_VAR = "TRTLLM_ALLOW_UNCORRECTED_OVERLAP_SPEC" @@ -1022,16 +1064,6 @@ def allocation_scope(current_stage: ExecutorMemoryType): # manager that was actually built. v2_selection = kv_cache_config.use_kv_cache_manager_v2 - # Rejected draft tokens shrink a request's page list, and the freed slot - # goes to whichever request allocates next. The connector is only told - # about pages appended since the last report, so it would keep - # addressing a slot another request now owns. - if (spec_config is not None and spec_config.max_draft_len > 0 - and v2_selection is True): - raise NotImplementedError( - "KV connector is not supported with speculative decoding. " - "Disable speculative decoding to run a connector.") - # V1 registers a single pool. V2 checks its resolved layer groups # during executor initialization. max_attention_window = kv_cache_config.max_attention_window @@ -1063,6 +1095,26 @@ def allocation_scope(current_stage: ExecutorMemoryType): pp_size=mapping.pp_size, cp_size=mapping.cp_size) + # Rejected draft tokens shrink a request's page list, and the freed + # slot goes to whichever request allocates next. The connector is + # only told about pages appended since the last report, so it would + # keep addressing a slot another request now owns. Context-only + # requests never reject a draft token, which is what the + # speculative-prefill-only mode relies on. Checked once the classes + # are resolved, since the mode depends on their declarations, and + # still before any connector instance exists. + speculative_prefill_only = _connector_runs_speculative_prefill_only( + spec_config, worker_cls, scheduler_cls, + cache_transceiver_config, v2_selection) + if (spec_config is not None and (spec_config.max_draft_len or 0) > 0 + and v2_selection is True and not speculative_prefill_only): + raise NotImplementedError( + "KV connector is not supported with speculative decoding. " + "One-model MTP or EAGLE3 is supported on a disaggregated " + "context server when the connector scheduler and worker " + "both declare supports_speculative_prefill=True. Otherwise " + "disable speculative decoding to run a connector.") + rank = tensorrt_llm.mpi_rank() # Some connector API implementations may need to establish out-of-band communication between the scheduler and workers. # In this case, the worker may be dependent on the scheduler, or vice-versa. @@ -1089,7 +1141,23 @@ def allocation_scope(current_stage: ExecutorMemoryType): kv_connector_manager = KvCacheConnectorManager( connector_worker, connector_scheduler, - enable_attention_dp=mapping.enable_attention_dp) + enable_attention_dp=mapping.enable_attention_dp, + speculative_prefill_only=speculative_prefill_only, + # Read from the class, which every rank resolved above: under + # tensor parallelism only rank 0 holds a scheduler instance. + admission_query=getattr(scheduler_cls, + "supports_admission_query", False) + is True) + # Logged once the worker says whether it moves KV: a capacity-only + # role is never refused, which is the recommended generation-server + # setup and should not be told otherwise. + if (kv_connector_manager.speculative_prefill_only + and not kv_connector_manager.capacity_only): + logger.info( + "KV connector runs in speculative-prefill-only mode: " + "requests that are not context-only are refused at " + "admission. A disaggregated generation server should use " + "a connector role that transfers no KV (capacity_only).") except Exception as e: logger.error(f"Error instantiating connector: {e}") diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py index ece654d11bc5..685170fd3c7a 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py @@ -28,6 +28,10 @@ - DefaultADPRouter: load-balanced min-heap routing - KVCacheAwareADPRouter: cache-aware routing that factors in prefix match length from the KV cache radix tree + - ConversationAwareADPRouter: explicit conversation-to-rank affinity + - WorkAwareADPRouter: balances ranks by estimated uncached prefill on + servers whose KV cache connector restores prefixes from a pool shared + by all ranks """ from __future__ import annotations @@ -48,6 +52,7 @@ from tensorrt_llm._torch.distributed.communicator import Distributed from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest + from ..connectors.kv_cache_connector import KvCacheConnectorManager from ..executor_request_queue import RequestQueueItem HeapVal = namedtuple("HeapVal", ["num_tokens", "num_requests", "rank", "request_list"]) @@ -64,6 +69,23 @@ def _num_input_tokens(request) -> int: return len(getattr(request, "input_token_ids", [])) +def _lru_put(table: OrderedDict, key, value, cap: int) -> None: + """Insert or refresh ``key`` as most recent, then evict the oldest entries above ``cap``.""" + table[key] = value + table.move_to_end(key) + while len(table) > cap: + table.popitem(last=False) + + +def _attention_dp_relax(req_item) -> bool: + """Sort key that puts requests with a strict ``attention_dp_rank`` first.""" + scheduling_params = getattr(req_item.request, "py_scheduling_params", None) + if scheduling_params is None or scheduling_params.attention_dp_rank is None: + return True + relax = scheduling_params.attention_dp_relax + return True if relax is None else relax + + def is_retiring_request(request) -> bool: """True if ``request`` has produced its final token and is being torn down.""" return request.state == LlmRequestState.GENERATION_TO_COMPLETE @@ -110,6 +132,7 @@ def create( kv_cache_manager=None, attention_dp_config=None, async_transfer_manager=None, + kv_connector_manager: Optional["KvCacheConnectorManager"] = None, ) -> "ADPRouter": """Factory method to create the appropriate ADP router. @@ -122,12 +145,20 @@ def create( async_transfer_manager: PyExecutor's AsyncTransferManager, used by KVCacheAwareADPRouter to account for KV-transfer-in-progress requests in per-rank load. Ignored by DefaultADPRouter. + kv_connector_manager: The executor's KV cache connector manager (may + be None). WorkAwareADPRouter needs one that moves KV (not + ``capacity_only``). Under attention DP the manager has already + validated that such a connector shares one storage pool across + the ranks (``KvCacheConnectorManager.validate_attention_dp``), so + a conversation's earlier prompt is restorable on any rank. Returns: A ConversationAwareADPRouter if kv_cache_routing_conversation_affinity is set; otherwise a KVCacheAwareADPRouter if enable_kv_cache_aware_routing is set, the kv_cache_manager has block reuse enabled and helix CP is - off; otherwise a DefaultADPRouter. + off; otherwise a WorkAwareADPRouter if enable_work_aware_routing is + set and a KV-moving connector is attached; otherwise a + DefaultADPRouter. """ if ( attention_dp_config is not None @@ -176,6 +207,22 @@ def create( account_for_in_transfer=attention_dp_config.kv_cache_routing_account_for_in_transfer, ) + # `is True`: unvalidated (Mock) configs carry truthy auto-attributes. + if ( + attention_dp_config is not None + and attention_dp_config.enable_work_aware_routing is True + ): + if kv_connector_manager is not None and not kv_connector_manager.capacity_only: + return WorkAwareADPRouter(dist=dist, has_seq_slot_headroom=has_seq_slot_headroom) + # The KV-cache estimation executor is built without its connector + # (create_py_executor); the final executor decides, so stay quiet here. + if getattr(kv_cache_manager, "is_estimating_kv_cache", False) is not True: + logger.warning( + "attention_dp_config.enable_work_aware_routing requires a KV cache " + "connector that restores prefixes from a pool shared by the " + "attention-DP ranks; falling back to load-based attention-DP routing." + ) + return DefaultADPRouter(dist=dist, has_seq_slot_headroom=has_seq_slot_headroom) @abstractmethod @@ -1080,3 +1127,157 @@ def _least_active_tokens(soft_cap: int) -> int: ) return all_ranks_new_requests, expected_num_active_requests + + +class WorkAwareADPRouter(ADPRouter): + """Balances attention-DP ranks by estimated prefill work, not prompt length. + + With prefix reuse (and especially a shared KV store), most of a prompt is + restored rather than computed, so prompt length says little about the work a + request brings: two prompts of the same length can differ by 50x in the + tokens they actually prefill. The default router balances by prompt length, + which leaves one rank working through a large uncached prefill while the + others run padding iterations in lockstep. + + The estimate for a request is how much its prompt grew since the previous + request of the same conversation (the earlier prompt is the reusable part); + a request with no known predecessor, or whose prompt shrank, counts its + whole prompt. A rank's load is the summed estimate of its requests still in + the context phase, and each new request goes, largest first, to the + least-loaded rank with room for it. + + Every rank routes the same new requests in the same order, so the + conversation and per-request bookkeeping stay identical across ranks. + + Intended for context servers whose KV cache connector restores prefixes + from a pool shared by all attention-DP ranks: a conversation's earlier + prompt is restorable on whichever rank a later request lands, so the + estimate does not depend on the rank. ``ADPRouter.create`` builds it only + when such a connector is attached. + + Only ``route_requests`` changes the two tables, including their LRU order, + and it runs on every rank with the same gathered states and broadcast + requests, so the tables stay identical across ranks. ``create_rank_state`` + only reads them. The estimate is also published as + ``py_uncached_prefill_tokens`` for the V2 scheduler's long-prefill demotion + (``TLLM_V2_CTX_DEMOTE_UNCACHED_TOKENS``) while that is enabled (see + ``set_publish_uncached_estimate``). + """ + + # LRU cap on each table (entries are tens of bytes each), bounding memory on + # long-running servers as conversations and requests churn. + DEFAULT_MAX_ENTRIES = 1 << 16 + + def __init__( + self, + dist: "Distributed", + has_seq_slot_headroom: bool = False, + max_entries: int = DEFAULT_MAX_ENTRIES, + ): + super().__init__(dist, has_seq_slot_headroom=has_seq_slot_headroom) + self._max_entries = max(1, int(max_entries)) + # conversation_id -> prompt length of its latest request. + self._conv_prompt_len: "OrderedDict[str, int]" = OrderedDict() + # request id (and n>1 child ids) -> estimated prefill tokens. + self._estimated_work: "OrderedDict[int, int]" = OrderedDict() + self._publish_uncached_estimate_enabled = False + logger.info( + "[adp_router] work-aware attention-DP routing: ranks balanced by " + "estimated uncached prefill tokens" + ) + + @staticmethod + def _conversation_id(req_item) -> Optional[str]: + # The request's own conversation_id, never subagent_affinity_id: sub-agents + # have prompts of their own (same rule as ConversationAwareADPRouter's estimate). + params = getattr(req_item.request, "py_conversation_params", None) + return None if params is None else params.conversation_id + + def set_publish_uncached_estimate(self, enabled: bool) -> None: + """Publish py_uncached_prefill_tokens only while the scheduler uses it. + + Args: + enabled: Whether to set ``py_uncached_prefill_tokens`` on routed + requests. + """ + self._publish_uncached_estimate_enabled = enabled + + def _estimate_prefill_tokens(self, req_item: RequestQueueItem) -> int: + """Estimate, record and (optionally) publish a new request's uncached prefill. + + Args: + req_item: A new request being routed. + + Returns: + The estimated prefill tokens, at least 1. + """ + prompt_len = _num_input_tokens(req_item.request) + conv_id = self._conversation_id(req_item) + work = prompt_len + if conv_id is not None: + previous = self._conv_prompt_len.get(conv_id) + if previous is not None and prompt_len >= previous: + work = prompt_len - previous + _lru_put(self._conv_prompt_len, conv_id, prompt_len, self._max_entries) + work = max(1, work) + _lru_put(self._estimated_work, req_item.id, work, self._max_entries) + for child_id in req_item.child_req_ids or (): + _lru_put(self._estimated_work, child_id, work, self._max_entries) + if self._publish_uncached_estimate_enabled and req_item.request is not None: + req_item.request.py_uncached_prefill_tokens = work + return work + + def create_rank_state( + self, + active_requests: list[LlmRequest], + new_requests: list[RequestQueueItem], + ) -> RankState: + # Read-only on the tables (dict.get, no LRU touch): active requests are + # rank-local, and the tables must stay identical across ranks. + load = 0 + for req in active_requests: + if not req.is_context_init_state: + continue + load += self._estimated_work.get(req.py_request_id, req.py_orig_prompt_len) + return RankState( + rank=self.dist.tp_rank, + num_active_requests=len(active_requests), + num_active_tokens=load, + ) + + def route_requests( + self, + all_rank_states: list[RankState], + new_requests: list[RequestQueueItem], + max_num_active_requests: int, + ) -> Tuple[Dict[int, List[RequestQueueItem]], int]: + tp_size = len(all_rank_states) + all_ranks_new_requests: Dict[int, List[RequestQueueItem]] = { + s.rank: [] for s in all_rank_states + } + counts = [s.num_active_requests for s in all_rank_states] + loads = [s.num_active_tokens for s in all_rank_states] + ordered = sorted(new_requests, key=_attention_dp_relax) + # Every request is estimated once, in this order, identically on every rank. + work = {req_item.id: self._estimate_prefill_tokens(req_item) for req_item in ordered} + remaining = self._assign_explicit_dp_ranks( + ordered, all_ranks_new_requests, counts, max_num_active_requests + ) + for rank, assigned in all_ranks_new_requests.items(): + loads[rank] += sum(work[req_item.id] for req_item in assigned) + # Largest first, so big prefills spread before small ones fill in (stable sort). + for req_item in sorted(remaining, key=lambda item: work[item.id], reverse=True): + candidates = [r for r in range(tp_size) if counts[r] < max_num_active_requests] + rank = min(candidates or range(tp_size), key=lambda r: (loads[r], counts[r], r)) + all_ranks_new_requests[rank].append(req_item) + loads[rank] += work[req_item.id] + counts[rank] += 1 + # Capped like the other routers; a rank already above the cap is + # tolerated by PyExecutor's padding check. + expected_num_active_requests = min(max(counts), max_num_active_requests) + logger.debug( + f"[adp_router][work] new_reqs_per_rank=" + f"{[len(all_ranks_new_requests[r]) for r in range(tp_size)]} " + f"est_load_per_rank={loads}" + ) + return all_ranks_new_requests, expected_num_active_requests diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 9c17e72d3fb3..8d9dd3e17de7 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -75,6 +75,13 @@ def __init__( self.max_num_requests = max_num_requests self.num_tokens = 0 self.num_requests = 0 + # Requests admitted with a 0-token charge because the KV connector will + # park them for an asynchronous load; they still count in num_requests. + self.num_uncharged_parks = 0 + # First-chunk context requests passed over because the token budget was + # spent: the pass was limited by budget, not by queued work. Read only + # by the TLLM_CTX_PARK_STATS diagnostics. + self.num_ctx_budget_skips = 0 # PEFT accounting self._peft_cache_manager = peft_cache_manager @@ -155,6 +162,14 @@ class KVCacheV2Scheduler(RequestScheduler): assignment into one unified loop. """ + # Admission-time connector queries are off unless __init__ turns them on; + # declared on the class so an instance built without __init__ schedules as + # before and reports no uncharged request. + _connector_admission_query: bool = False + _admission_pass_active: bool = False + uncharged_park_ids: frozenset = frozenset() + last_pass_ctx_budget_skips: int = 0 + def __init__( self, max_batch_size: int, @@ -170,7 +185,38 @@ def __init__( cross_kv_cache_manager: object | None = None, # KVCacheManagerV2 for enc-dec cross-attn enable_prefix_aware_scheduling: bool = True, enable_recompute_pause: bool = True, + connector_admission_query: bool = False, + max_uncharged_connector_parks: int = 0, + max_connector_offer_deferrals: int = 16, ) -> None: + """Build the scheduler. + + Args: + max_batch_size: Maximum number of requests per batch. + max_num_tokens: Token budget per iteration, or None for unlimited. + kv_cache_manager: The KVCacheManagerV2 that allocates inline. + scheduler_policy: Capacity policy; only MAX_UTILIZATION is supported. + ctx_chunk_config: ``(policy, chunk_unit_size)`` to chunk context + requests, or None. + peft_cache_manager: PEFT cache manager, if any. + scheduler_capacity: Request capacity; defaults to max_batch_size. + no_schedule_until_state: First schedulable request state. + no_schedule_after_state: First request state past the schedulable range. + draft_kv_cache_manager: KVCacheManagerV2 for the draft layers, if any. + cross_kv_cache_manager: KVCacheManagerV2 for cross attention, if any. + enable_prefix_aware_scheduling: Defer requests behind a contributor + of the same uncached block. + enable_recompute_pause: Allow preempting a started request to free pages. + connector_admission_query: Ask an admission-capable KV connector for + a first chunk's prefix during admission and book the pages a + shifting offer needs (``kv_connector_config.admission_query``). + Needs context chunking. + max_uncharged_connector_parks: Per pass, at most this many first + chunks that the connector parks for an asynchronous load are + admitted with a 0-token charge. + max_connector_offer_deferrals: Passes a first chunk may give its + pages back while the KV cache cannot cover its offer's window. + """ self.remote_tail_phase: RemoteTailPhase | None = None self._remote_tail_context_ids: frozenset[int] = frozenset() self._remote_tail_generation_ids: frozenset[int] = frozenset() @@ -287,6 +333,43 @@ def __init__( f"(max iterations demoted: {self._ctx_demote_max_iters or 'unbounded'})" ) + # Admission-time connector queries (kv_connector_config.admission_query). + # A connector that declares `supports_admission_query` is asked for a + # first chunk's prefix once the chunk is sized and allocated + # (`_parks_for_async_load`), and the pages a shifting offer needs are + # booked right away. A first chunk the connector will load + # asynchronously is parked by the executor before the forward pass, so + # up to `max_uncharged_connector_parks` of them per pass are admitted + # with a 0-token charge; they still take a request slot. Chunks under + # two blocks keep the charge: they free little budget. + self._connector_admission_query = connector_admission_query and self.chunking_enabled + self._max_uncharged_parks = max(0, max_uncharged_connector_parks) + self._max_offer_deferrals = max(0, max_connector_offer_deferrals) + self._uncharged_park_min_tokens = 2 * self.tokens_per_block + self._admission_pass_active = False + self._uncharged_park_ids_pass: set[int] = set() + self._offer_page_deferrals_pass = 0 + self._uncharged_parks_logged = 0 + #: Request ids admitted with a 0-token charge in the last pass. The + #: executor checks that each one parked. + self.uncharged_park_ids = frozenset() + #: `BudgetTracker.num_ctx_budget_skips` of the last pass, for the + #: TLLM_CTX_PARK_STATS diagnostics. + self.last_pass_ctx_budget_skips = 0 + if connector_admission_query: + if ctx_chunk_config is None: + logger.warning( + "KVCacheV2Scheduler: kv_connector_config.admission_query has no " + "effect without context chunking" + ) + else: + logger.info( + "KVCacheV2Scheduler: admission-time connector queries, up to " + f"{self._max_uncharged_parks} uncharged parks per pass (chunks >= " + f"{self._uncharged_park_min_tokens} tokens), up to " + f"{self._max_offer_deferrals} deferrals for offer pages" + ) + # Registered by PyExecutor; see set_async_transfer_manager. self._async_transfer_manager = None @@ -340,6 +423,18 @@ def schedule_request( self, active_requests: RequestList, inflight_request_ids: set[int] ) -> SchedulerOutput: active_requests = drop_decoder_context_requests_waiting_for_encoder_output(active_requests) + # Checked per pass: the manager answers False under attention DP, for a + # capacity-only role, with prefix reservations, or without an + # admission-capable connector. `is True` because a mocked manager + # answers truthily. + self._admission_pass_active = ( + self._connector_admission_query + and self.kv_cache_manager.connector_admission_query_enabled() is True + ) + self._uncharged_park_ids_pass = set() + self._offer_page_deferrals_pass = 0 + if self._admission_pass_active: + self.kv_cache_manager.begin_connector_admission_pass() # Main scheduling loop ( scheduled_encoder, @@ -351,6 +446,8 @@ def schedule_request( has_chunking, ) = self._schedule_loop(active_requests, inflight_request_ids) + self.uncharged_park_ids = frozenset(self._uncharged_park_ids_pass) + # Sort by LoRA task ID scheduled_encoder.sort(key=_get_lora_task_id) self._sort_requests(scheduled_ctx, scheduled_gen, has_chunking) @@ -616,9 +713,13 @@ def preempt_for_pages(req: LlmRequest) -> bool: break if other.py_request_id in demoted_ids: protected.add(other.py_request_id) - return self._try_preempt_for_pages( + freed = self._try_preempt_for_pages( requests_list, protected, inflight_request_ids, recompute_paused, preempted_ids ) + if freed and self._admission_pass_active: + # Its next admission keeps these pages; see _defer_for_offer_pages. + req.py_preempted_for_pages = True + return freed # --- Phase 2: schedule deferred context / encoder requests --- # Generation PEFT pages are now fully committed in the budget. @@ -660,8 +761,11 @@ def preempt_for_pages(req: LlmRequest) -> bool: and req.state_value == self._context_init_state_value and not self._has_context_chunk_budget(budget) ): + if req.is_first_context_chunk: + budget.num_ctx_budget_skips += 1 continue if scheduled_ctx and self._first_chunk_waits_for_budget(req, budget): + budget.num_ctx_budget_skips += 1 continue if req.py_request_id in preempted_ids: continue @@ -743,6 +847,9 @@ def preempt_for_pages(req: LlmRequest) -> bool: or evicted or recompute_paused or deferred_behind_contributor + # A deferral for offer pages gives its pages back, and + # max_connector_offer_deferrals bounds the wait per request. + or self._offer_page_deferrals_pass > 0 ) if ( self.remote_tail_phase is None @@ -762,6 +869,8 @@ def preempt_for_pages(req: LlmRequest) -> bool: rid: self._ctx_demote_skips.get(rid, 0) + 1 for rid in demoted_ids } + self.last_pass_ctx_budget_skips = budget.num_ctx_budget_skips + # Phase and frozen-membership deferrals are not KV capacity stalls. self._detect_deadlock( requests_list, @@ -997,6 +1106,12 @@ def _try_schedule_context( Source reservations advance the prefix before budget checks. Destination pages are allocated here; loads are accepted after the final batch trims. A loading request remains outside the schedulable state range. + + On a disaggregated server without an eviction tier whose connector + declares `requires_completion_reservation`, or whose + `kv_cache_config.enable_disagg_completion_reservation` is set, a + connector first chunk can be deferred for completion headroom; the + rollback below also drops its bound. """ first_chunk = req.is_first_context_chunk if self.chunking_enabled: @@ -1071,6 +1186,14 @@ def _try_schedule_context_full( # V2 resizes KV cache directly in the scheduler (no separate # prepareResources for main cache), so include draft tokens. req.context_chunk_size = context_tokens + if not self.kv_cache_manager.reserve_context_completion(req): + # Admitted requests own the completion headroom. Defer this prompt + # instead of preempting one of them to make room for it. + logger.debug( + "Deferring context request %s: admitted requests hold the completion headroom", + req.py_request_id, + ) + return ScheduleAction.SKIP, 0, False if not self._try_allocate_context(req, context_tokens + draft_len): # Out of pages. Give up one started request so this one can # proceed, and retry next iteration: a failed resize leaves a @@ -1238,6 +1361,13 @@ def _try_schedule_context_chunked( chunk_tokens += draft_len resize_tokens += draft_len + if not self.kv_cache_manager.reserve_context_completion(req): + # As in _try_schedule_context_full: wait for headroom, do not preempt. + logger.debug( + "Deferring context request %s: admitted requests hold the completion headroom", + req.py_request_id, + ) + return ScheduleAction.SKIP, 0, False # V2 resizes KV cache directly in the scheduler, so include # draft tokens for last chunk. if not self._try_allocate_context(req, resize_tokens): @@ -1253,8 +1383,107 @@ def _try_schedule_context_chunked( chunking_flag = req.context_chunk_size < req.context_remaining_length + # After the pre-grown clamp, MM alignment, allocation and cross context, + # so a pre-grown replay is asked with its exact chunk. + if self._admission_pass_active: + uncharged = self._parks_for_async_load(req, budget, chunk_tokens) + if uncharged is None: + return ScheduleAction.SKIP, 0, False + if uncharged: + return ScheduleAction.SCHEDULED, 0, chunking_flag return ScheduleAction.SCHEDULED, chunk_tokens, chunking_flag + def _parks_for_async_load( + self, req: LlmRequest, budget: BudgetTracker, chunk_tokens: int + ) -> Optional[bool]: + """Ask the connector at admission and decide the chunk's token charge. + + Every scheduled first chunk is asked here, once its chunk is sized and + its pages allocated; a request dropped before ``prepare_resources`` is + asked again on its next pass, as ``supports_admission_query`` allows. + An offer that moves the chunk past its sized end gets its pages now + (``KVCacheManagerV2.cover_admission_offer``). When the KV cache cannot + cover them the request gives its chunk back and waits for a later pass, + up to ``max_connector_offer_deferrals`` passes, rather than load only + the blocks the chunk's pages hold and recompute the rest. Past that + bound it is admitted as before and ``prepare_resources`` falls back to + those blocks. So is a request whose last pass preempted a started + request for pages: deferring it would hand those pages to the requests + behind it and make it preempt again. + + Only the token charge is waived, since the executor removes a parked + request from the batch before the forward pass. At most + ``max_uncharged_connector_parks`` requests per pass are admitted + uncharged, and only chunks of at least two blocks; a synchronous or + miss answer does not count. Past the cap a request that parks keeps the + normal charge. + + Args: + req: The context request just allocated. + budget: This pass's budget tracker. + chunk_tokens: The tokens the chunk would be charged. + + Returns: + None to skip ``req`` this pass (deferred for offer pages), True to + admit it with a 0-token charge, False to charge ``chunk_tokens``. + """ + preempted_for_pages = req.py_preempted_for_pages + req.py_preempted_for_pages = False + parks = self.kv_cache_manager.query_connector_at_admission(req) + if ( + not self.kv_cache_manager.cover_admission_offer(req) + and not preempted_for_pages + and self._defer_for_offer_pages(req) + ): + return None + if not ( + parks + and budget.num_uncharged_parks < self._max_uncharged_parks + and chunk_tokens >= self._uncharged_park_min_tokens + ): + return False + budget.num_uncharged_parks += 1 + self._uncharged_park_ids_pass.add(req.py_request_id) + if self._uncharged_parks_logged < 3: + self._uncharged_parks_logged += 1 + logger.info( + f"[ctx-park] request {req.py_request_id} admitted uncharged: its " + f"{chunk_tokens}-token chunk waits for an asynchronous load" + ) + return True + + def _defer_for_offer_pages(self, req: LlmRequest) -> bool: + """Let ``req`` wait for the pages its connector offer needs. + + Only a first chunk can hold an admission-time offer. The caller returns + SKIP, and ``_try_schedule_context`` then runs + ``release_context_admission``, which frees every pool, resets the + connector's state for the request and rewinds its cursor, so the next + pass matches and asks afresh. Counts as progress for the deadlock + detector, since the bound keeps the wait finite. + + Args: + req: The first chunk whose offer the KV cache could not cover. + + Returns: + True to defer ``req`` this pass. False once it has waited + ``max_connector_offer_deferrals`` passes; it is then admitted and + takes the capped prefix. + """ + deferrals = req.py_connector_offer_deferrals + if deferrals >= self._max_offer_deferrals: + if deferrals == self._max_offer_deferrals: + req.py_connector_offer_deferrals = deferrals + 1 + logger.debug( + f"[ctx-park] request {req.py_request_id}: no pages for its connector " + f"offer after {deferrals} passes; admitting it with the capped prefix" + ) + return False + req.py_connector_offer_deferrals = deferrals + 1 + req.py_connector_admission_offer = None + self._offer_page_deferrals_pass += 1 + return True + def _reuse_claim_limit(self, req: LlmRequest) -> int | None: """Pair the two pools on one depth before either claims. diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index e02c42ae24f5..61dd318f63ba 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -33,7 +33,8 @@ from pydantic import AliasChoices, BaseModel, ConfigDict from pydantic import Field as PydanticField from pydantic import (NonNegativeFloat, NonNegativeInt, PositiveInt, - PrivateAttr, StrictInt, field_validator, model_validator) + PrivateAttr, StrictInt, ValidationInfo, field_validator, + model_validator) from strenum import StrEnum from transformers import PreTrainedTokenizerBase @@ -2113,6 +2114,23 @@ class AttentionDpConfig(StrictBaseModel): "Active prompt tokens estimate work, not cached KV residency. Existing " "conversation->rank pins are unaffected. Only used when " "kv_cache_routing_conversation_affinity is True.") + enable_work_aware_routing: bool = Field( + default=False, + description= + "Route attention-DP requests by estimated uncached prefill instead of " + "prompt length. A request's estimate is its prompt growth since the " + "previous request with the same conversation_params.conversation_id " + "(its whole prompt when there is none or the prompt shrank); a rank's " + "load is the summed estimate of its requests still in the context " + "phase, and new requests go, largest first, to the least-loaded rank " + "below max_num_active_requests. Intended for context servers whose KV " + "cache connector restores prefixes from a pool shared by all " + "attention-DP ranks, where a conversation's earlier prompt is reusable " + "on any rank. Without a KV cache connector that moves KV, load-based " + "routing is used and a warning is logged. Cannot be combined with " + "enable_kv_cache_aware_routing or " + "kv_cache_routing_conversation_affinity.", + status="prototype") @model_validator(mode='after') def validate_attention_dp_config(self) -> 'AttentionDpConfig': @@ -2125,6 +2143,13 @@ def validate_attention_dp_config(self) -> 'AttentionDpConfig': raise ValueError( "attention_dp_config.timeout_iters must be greater or equal to 0 when enable_balance is true" ) + if self.enable_work_aware_routing and ( + self.enable_kv_cache_aware_routing + or self.kv_cache_routing_conversation_affinity): + raise ValueError( + "attention_dp_config.enable_work_aware_routing cannot be " + "combined with enable_kv_cache_aware_routing or " + "kv_cache_routing_conversation_affinity") return self @@ -2654,7 +2679,65 @@ class MooncakeStoreConfig(StrictBaseModel): description="Copy pages through a pinned host buffer instead of " "registering the KV pools with Mooncake. For a host whose HCA cannot " "pin GPU pages; costs a copy each way. Ignored when role is " - "'capacity'.") + "'capacity'. Pinned memory per direction is staging_buffer_bytes.") + staging_buffer_bytes: Union[int, str] = Field( + "1GiB", + telemetry=False, + description="Pinned host memory per transfer direction when " + "stage_through_host is set, as a binary size ('512MiB') or a byte " + "count; ambiguous units such as 'GB' are refused. It bounds how many " + "pages one store call stages, not how large a page may be. With " + "async_load the load direction's share is split evenly across " + "async_load_workers. Counted per rank in the node memory check. " + "Ignored without stage_through_host.") + async_load: bool = Field( + False, + telemetry=False, + status="prototype", + description="Load a prefix the pool serves on background threads " + "while the request waits outside the batch, instead of inside the " + "executor iteration, so other requests keep computing during the " + "transfer. A page that is gone by load time restarts only that " + "request on a single rank or an attention-DP owner, and fails the " + "server under tensor parallelism. Validated with " + "disable_overlap_scheduler. Ignored unless role reads the pool " + "('both' or 'consumer').") + async_load_workers: PositiveInt = Field( + 1, + telemetry=False, + status="prototype", + description="Background load threads per rank when async_load is " + "set. Each has its own CUDA stream and, under stage_through_host, an " + "equal share of staging_buffer_bytes.") + soft_pin_ttl_ms: Optional[PositiveInt] = Field( + None, + telemetry=False, + description="Soft-pin every page this server writes for this many " + "milliseconds, so the master evicts the pool by age instead of " + "refilling an eviction shortfall in shard order, which punches holes " + "into prefixes still in use. Must not exceed the master's " + "--max_kv_soft_pin_ttl (24 h by default), or every put is refused, " + "and the master must allow evicting soft-pinned objects (Mooncake's " + "default). Ignored unless the role writes; needs Mooncake 0.3.13 or " + "newer.") + group_blocks: bool = Field( + False, + telemetry=False, + description="Put every page of a block (all layer groups and " + "tensor-parallel shards) into one Mooncake eviction group, so the " + "master evicts a block whole rather than leaving pages that hold " + "space but can never hit. Mostly matters under tensor parallelism. " + "Ignored unless the role writes; needs Mooncake 0.3.13 or newer.") + stats_interval: Optional[float] = Field( + None, + gt=0, + allow_inf_nan=False, + telemetry=False, + description="Seconds between per-rank info-level log lines counting " + "prefix lookups stopped by a hole the pool evicted from a prefix it " + "otherwise still holds, and timing each phase of a prefix load. " + "Only a role that reads the pool counts anything. None disables the " + "counters at no cost.") local_hostname: Optional[str] = Field( None, telemetry=False, @@ -2677,23 +2760,67 @@ class MooncakeStoreConfig(StrictBaseModel): "master to accept connections. Raise it when the wait spans a " "container start on another node. Expiring fails the server at " "startup.") - - @field_validator("segment_size", mode="after") + local_reuse_backoff: Literal["off", "align"] = Field( + "off", + telemetry=False, + description="How a local prefix match is handled under one-model " + "speculative decoding, whose draft reads prompt tokens past each " + "position and so backs every local reuse match off by that lookahead. " + "'off' drops the backoff: local hits are used whole, and the draft KV " + "at the tail of a reused block may belong to a different next token, " + "which can lower draft acceptance but does not change the output. " + "'align' keeps the backoff and aligns the match down to a block " + "boundary, so the pool serves the trimmed block (a transfer) along " + "with the rest of the prefix. Has no effect without one-model " + "speculative decoding or for a role that does not read the pool.") + + @field_validator("segment_size", "staging_buffer_bytes", mode="after") @classmethod - def _check_segment_size(cls, value): + def _check_sizes(cls, value, info: ValidationInfo): """Reject a bad size here rather than in every rank after bringup.""" from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.config import \ parse_size + name = info.field_name try: parsed = parse_size(value, strict_units=True) except ValueError as exc: - raise ValueError(f"mooncake_store.segment_size: {exc}") - if parsed < 0: + raise ValueError(f"mooncake_store.{name}: {exc}") + if name == "segment_size" and parsed < 0: raise ValueError(f"mooncake_store.segment_size: {value!r} is " f"{parsed} bytes; it cannot be negative.") + if name == "staging_buffer_bytes" and parsed <= 0: + raise ValueError(f"mooncake_store.staging_buffer_bytes: {value!r} " + f"is {parsed} bytes; it must be positive.") return value +class KvCacheConnectorAdmissionConfig(StrictBaseModel): + """Admission-time connector queries for the KVCacheManagerV2 scheduler. + + The scheduler asks a connector that declares + ``KvCacheConnectorScheduler.supports_admission_query`` for its prefix while + it admits a first context chunk, once the chunk is sized and its pages are + allocated, and books the pages a shifting offer needs at that point. + """ + max_uncharged_parks: NonNegativeInt = Field( + default=0, + status="prototype", + description="Per scheduling iteration, at most this many first context " + "chunks of at least two KV blocks, whose prefix the connector will load " + "asynchronously, are admitted without charging their tokens to " + "max_num_tokens. They still take a batch slot and keep their pages. " + "0 charges every chunk.") + max_offer_deferrals: NonNegativeInt = Field( + default=16, + status="prototype", + description="How many scheduling iterations a first context chunk may " + "give its pages back while the KV cache cannot cover the window a " + "connector offer moves the chunk to. Past this bound, or right after " + "the request preempted another one for pages, it is admitted and the " + "offer is capped to the whole blocks its chunk's pages hold. 0 never " + "defers.") + + class KvCacheConnectorConfig(StrictBaseModel): """Configuration for the KV Cache Connector. @@ -2746,6 +2873,18 @@ class KvCacheConnectorConfig(StrictBaseModel): "transfers after cancellation or shutdown. Can only shorten the " "existing transfer deadline. Expiration fails all ADP owners and requires process restart. " "Unused without attention DP.") + admission_query: Optional[KvCacheConnectorAdmissionConfig] = Field( + default=None, + status="prototype", + description="Ask the connector for a first context chunk's prefix while " + "the scheduler admits it, instead of after the batch is final. None " + "keeps final-batch queries; {} enables admission queries and page " + "booking only; max_uncharged_parks > 0 also admits chunks that park for " + "an asynchronous load without a token charge. Requires KVCacheManagerV2 " + "with context chunking and disable_overlap_scheduler=True. Ignored, with " + "a warning, under attention DP, for a capacity-only role, for " + "prefix-reservation connectors, and for connectors that do not declare " + "supports_admission_query.") @model_validator(mode="after") def _resolve_preset(self) -> "KvCacheConnectorConfig": @@ -5177,6 +5316,24 @@ class KvCacheConfig(StrictBaseModel, PybindMirror): description="KV cache manager v2 configuration for block reuse policies." ) + # This is a pure python field, not a pybind field. It is only for the Pytorch backend. + enable_disagg_completion_reservation: bool = Field( + default=False, + status="prototype", + description= + "Reserve the GPU KV cache pages each admitted request needs to finish " + "on a disaggregated server. Applies only to KV cache manager v2 on a " + "disaggregated server with no cache tier below GPU (host_cache_size=0 " + "and no disk tier), beam width 1, a single pipeline stage and full " + "attention on every layer; elsewhere it is ignored, and KV cache " + "manager v2 logs a warning. A request is then admitted only while the " + "completion bound of every admitted request still fits: a generation " + "request through its output limit and, under a KV connector without " + "source reservations, a context request the connector may serve " + "through the end of its prompt. The bounds are logical; no page is " + "allocated early. A KV connector can request the same through its " + "worker's requires_completion_reservation declaration.") + def _to_pybind(self): config = _KvCacheConfig( enable_block_reuse=self.enable_block_reuse, diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py index 3eb903af9973..86a7cef76fdd 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py @@ -136,6 +136,7 @@ class _KVCacheManagerConfigFieldSpec: enable_partial_commit: bool = True pool_rebalance: object = None cold_tier_offload_min_priority: int = DEFAULT_COLD_TIER_OFFLOAD_MIN_PRIORITY + reuse_match_alignment: int = 1 KVCacheManagerConfig.__dataclass_fields__ = _KVCacheManagerConfigFieldSpec.__dataclass_fields__ diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi index e418a261853e..6f4b4bc94d6a 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi @@ -333,6 +333,9 @@ class KVCacheManagerConfig: enable_partial_commit: bool = True pool_rebalance: PoolRebalanceConfig = ... cold_tier_offload_min_priority: int = DEFAULT_COLD_TIER_OFFLOAD_MIN_PRIORITY + # Required alignment of the final prefix match, after backoff and coverage + # pruning; a positive divisor of tokens_per_block. 1 keeps token-granular reuse. + reuse_match_alignment: int = 1 @property def enable_swa_scratch_reuse(self) -> bool: ... diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index de355f70ce33..bab056050b54 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -70,6 +70,13 @@ "title": "Enable Kv Cache Aware Routing", "type": "boolean" }, + "enable_work_aware_routing": { + "default": false, + "description": "Route attention-DP requests by estimated uncached prefill instead of prompt length. A request's estimate is its prompt growth since the previous request with the same conversation_params.conversation_id (its whole prompt when there is none or the prompt shrank); a rank's load is the summed estimate of its requests still in the context phase, and new requests go, largest first, to the least-loaded rank below max_num_active_requests. Intended for context servers whose KV cache connector restores prefixes from a pool shared by all attention-DP ranks, where a conversation's earlier prompt is reusable on any rank. Without a KV cache connector that moves KV, load-based routing is used and a warning is logged. Cannot be combined with enable_kv_cache_aware_routing or kv_cache_routing_conversation_affinity.", + "status": "prototype", + "title": "Enable Work Aware Routing", + "type": "boolean" + }, "kv_cache_routing_account_for_in_transfer": { "default": false, "description": "In-transfer load accounting in KV cache-aware routing. When True, requests still streaming KV to the GEN worker (tracked by the PyExecutor AsyncTransferManager but no longer in active_requests) are folded back into the per-rank load reported via RankState. This can improve balance under heavy disagg traffic but inflates num_active_requests reported upstream, which lets the inference loop's idle-fetch wait expire even when no requests are runnable and causes empty fetch cycles. Default False preserves the prior behaviour (fetch blocks when truly idle). Only used when enable_kv_cache_aware_routing is True.", @@ -2868,6 +2875,13 @@ "title": "Enable Block Reuse", "type": "boolean" }, + "enable_disagg_completion_reservation": { + "default": false, + "description": "Reserve the GPU KV cache pages each admitted request needs to finish on a disaggregated server. Applies only to KV cache manager v2 on a disaggregated server with no cache tier below GPU (host_cache_size=0 and no disk tier), beam width 1, a single pipeline stage and full attention on every layer; elsewhere it is ignored, and KV cache manager v2 logs a warning. A request is then admitted only while the completion bound of every admitted request still fits: a generation request through its output limit and, under a KV connector without source reservations, a context request the connector may serve through the end of its prompt. The bounds are logical; no page is allocated early. A KV connector can request the same through its worker's requires_completion_reservation declaration.", + "status": "prototype", + "title": "Enable Disagg Completion Reservation", + "type": "boolean" + }, "enable_kv_pool_rebalance": { "default": false, "description": "Opt in to the KVCacheManagerV2 auto-tuner (``adjust()``) for rebalancing pool-group ratios between iterations. When True the PyExecutor calls ``adjust()`` opportunistically; the auto-tuner itself remains gated by the ``kv_pool_rebalance_config`` thresholds (2000 samples / 120s by default). When False (default) the rebalance hook is skipped entirely and pool ratios remain at their warmup-derived values. Beta: enable at your own risk. Only used when using KV cache manager v2 (experimental). This option is incompatible with dtype='fp8_ds_mla'.", @@ -3164,10 +3178,47 @@ "title": "KvCacheConfig", "type": "object" }, + "KvCacheConnectorAdmissionConfig": { + "additionalProperties": false, + "description": "Admission-time connector queries for the KVCacheManagerV2 scheduler.\n\nThe scheduler asks a connector that declares\n``KvCacheConnectorScheduler.supports_admission_query`` for its prefix while\nit admits a first context chunk, once the chunk is sized and its pages are\nallocated, and books the pages a shifting offer needs at that point.", + "properties": { + "max_offer_deferrals": { + "default": 16, + "description": "How many scheduling iterations a first context chunk may give its pages back while the KV cache cannot cover the window a connector offer moves the chunk to. Past this bound, or right after the request preempted another one for pages, it is admitted and the offer is capped to the whole blocks its chunk's pages hold. 0 never defers.", + "minimum": 0, + "status": "prototype", + "title": "Max Offer Deferrals", + "type": "integer" + }, + "max_uncharged_parks": { + "default": 0, + "description": "Per scheduling iteration, at most this many first context chunks of at least two KV blocks, whose prefix the connector will load asynchronously, are admitted without charging their tokens to max_num_tokens. They still take a batch slot and keep their pages. 0 charges every chunk.", + "minimum": 0, + "status": "prototype", + "title": "Max Uncharged Parks", + "type": "integer" + } + }, + "title": "KvCacheConnectorAdmissionConfig", + "type": "object" + }, "KvCacheConnectorConfig": { "additionalProperties": false, "description": "Configuration for the KV Cache Connector.\n\nCan be configured either by specifying a named preset via ``connector``\n(e.g. ``\"lmcache\"``), or by providing explicit ``connector_module``,\n``connector_scheduler_class``, and ``connector_worker_class`` fields.\nWhen ``connector`` is set, the module/class fields are auto-populated\nfrom the preset registry and can be omitted.", "properties": { + "admission_query": { + "anyOf": [ + { + "$ref": "#/$defs/KvCacheConnectorAdmissionConfig" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Ask the connector for a first context chunk's prefix while the scheduler admits it, instead of after the batch is final. None keeps final-batch queries; {} enables admission queries and page booking only; max_uncharged_parks > 0 also admits chunks that park for an asynchronous load without a token charge. Requires KVCacheManagerV2 with context chunking and disable_overlap_scheduler=True. Ignored, with a warning, under attention DP, for a capacity-only role, for prefix-reservation connectors, and for connectors that do not declare supports_admission_query.", + "status": "prototype" + }, "connector": { "anyOf": [ { @@ -4121,6 +4172,36 @@ "additionalProperties": false, "description": "How this server joins a Mooncake store pool.\n\nThe settings every participant shares come from the master named by\n`pool`; what is left here is per-server. Setting this makes\n`trtllm-serve` render the Mooncake client config and export\n`MOONCAKE_CONFIG_PATH`; an inherited `MOONCAKE_CONFIG_PATH` wins.", "properties": { + "async_load": { + "default": false, + "description": "Load a prefix the pool serves on background threads while the request waits outside the batch, instead of inside the executor iteration, so other requests keep computing during the transfer. A page that is gone by load time restarts only that request on a single rank or an attention-DP owner, and fails the server under tensor parallelism. Validated with disable_overlap_scheduler. Ignored unless role reads the pool ('both' or 'consumer').", + "status": "prototype", + "telemetry": { + "exclude": true + }, + "title": "Async Load", + "type": "boolean" + }, + "async_load_workers": { + "default": 1, + "description": "Background load threads per rank when async_load is set. Each has its own CUDA stream and, under stage_through_host, an equal share of staging_buffer_bytes.", + "exclusiveMinimum": 0, + "status": "prototype", + "telemetry": { + "exclude": true + }, + "title": "Async Load Workers", + "type": "integer" + }, + "group_blocks": { + "default": false, + "description": "Put every page of a block (all layer groups and tensor-parallel shards) into one Mooncake eviction group, so the master evicts a block whole rather than leaving pages that hold space but can never hit. Mostly matters under tensor parallelism. Ignored unless the role writes; needs Mooncake 0.3.13 or newer.", + "telemetry": { + "exclude": true + }, + "title": "Group Blocks", + "type": "boolean" + }, "local_hostname": { "anyOf": [ { @@ -4137,6 +4218,19 @@ }, "title": "Local Hostname" }, + "local_reuse_backoff": { + "default": "off", + "description": "How a local prefix match is handled under one-model speculative decoding, whose draft reads prompt tokens past each position and so backs every local reuse match off by that lookahead. 'off' drops the backoff: local hits are used whole, and the draft KV at the tail of a reused block may belong to a different next token, which can lower draft acceptance but does not change the output. 'align' keeps the backoff and aligns the match down to a block boundary, so the pool serves the trimmed block (a transfer) along with the rest of the prefix. Has no effect without one-model speculative decoding or for a role that does not read the pool.", + "enum": [ + "off", + "align" + ], + "telemetry": { + "exclude": true + }, + "title": "Local Reuse Backoff", + "type": "string" + }, "master_timeout": { "default": 60.0, "description": "Seconds to wait for the pool manifest to appear and the master to accept connections. Raise it when the wait spans a container start on another node. Expiring fails the server at startup.", @@ -4213,15 +4307,65 @@ "description": "Host memory each of this server's ranks contributes to the pool, as a binary size ('16GiB') or a byte count; ambiguous units such as 'GB' are refused. A node's demand is ranks_on_node x segment_size, checked against available memory at startup. Zero lends nothing, leaving the server to use only capacity its peers hold.", "title": "Segment Size" }, + "soft_pin_ttl_ms": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Soft-pin every page this server writes for this many milliseconds, so the master evicts the pool by age instead of refilling an eviction shortfall in shard order, which punches holes into prefixes still in use. Must not exceed the master's --max_kv_soft_pin_ttl (24 h by default), or every put is refused, and the master must allow evicting soft-pinned objects (Mooncake's default). Ignored unless the role writes; needs Mooncake 0.3.13 or newer.", + "telemetry": { + "exclude": true + }, + "title": "Soft Pin Ttl Ms" + }, "stage_through_host": { "default": false, - "description": "Copy pages through a pinned host buffer instead of registering the KV pools with Mooncake. For a host whose HCA cannot pin GPU pages; costs a copy each way. Ignored when role is 'capacity'.", + "description": "Copy pages through a pinned host buffer instead of registering the KV pools with Mooncake. For a host whose HCA cannot pin GPU pages; costs a copy each way. Ignored when role is 'capacity'. Pinned memory per direction is staging_buffer_bytes.", "telemetry": { "exclude": true }, "title": "Stage Through Host", "type": "boolean" }, + "staging_buffer_bytes": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "string" + } + ], + "default": "1GiB", + "description": "Pinned host memory per transfer direction when stage_through_host is set, as a binary size ('512MiB') or a byte count; ambiguous units such as 'GB' are refused. It bounds how many pages one store call stages, not how large a page may be. With async_load the load direction's share is split evenly across async_load_workers. Counted per rank in the node memory check. Ignored without stage_through_host.", + "telemetry": { + "exclude": true + }, + "title": "Staging Buffer Bytes" + }, + "stats_interval": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "number" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Seconds between per-rank info-level log lines counting prefix lookups stopped by a hole the pool evicted from a prefix it otherwise still holds, and timing each phase of a prefix load. Only a role that reads the pool counts anything. None disables the counters at no cost.", + "telemetry": { + "exclude": true + }, + "title": "Stats Interval" + }, "transfer_batch_size": { "default": 64, "description": "Page keys per store call.", diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index 2bb5e4c910b0..f8a254ff9a9d 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -70,6 +70,13 @@ "title": "Enable Kv Cache Aware Routing", "type": "boolean" }, + "enable_work_aware_routing": { + "default": false, + "description": "Route attention-DP requests by estimated uncached prefill instead of prompt length. A request's estimate is its prompt growth since the previous request with the same conversation_params.conversation_id (its whole prompt when there is none or the prompt shrank); a rank's load is the summed estimate of its requests still in the context phase, and new requests go, largest first, to the least-loaded rank below max_num_active_requests. Intended for context servers whose KV cache connector restores prefixes from a pool shared by all attention-DP ranks, where a conversation's earlier prompt is reusable on any rank. Without a KV cache connector that moves KV, load-based routing is used and a warning is logged. Cannot be combined with enable_kv_cache_aware_routing or kv_cache_routing_conversation_affinity.", + "status": "prototype", + "title": "Enable Work Aware Routing", + "type": "boolean" + }, "kv_cache_routing_account_for_in_transfer": { "default": false, "description": "In-transfer load accounting in KV cache-aware routing. When True, requests still streaming KV to the GEN worker (tracked by the PyExecutor AsyncTransferManager but no longer in active_requests) are folded back into the per-rank load reported via RankState. This can improve balance under heavy disagg traffic but inflates num_active_requests reported upstream, which lets the inference loop's idle-fetch wait expire even when no requests are runnable and causes empty fetch cycles. Default False preserves the prior behaviour (fetch blocks when truly idle). Only used when enable_kv_cache_aware_routing is True.", @@ -4432,6 +4439,13 @@ "title": "Enable Block Reuse", "type": "boolean" }, + "enable_disagg_completion_reservation": { + "default": false, + "description": "Reserve the GPU KV cache pages each admitted request needs to finish on a disaggregated server. Applies only to KV cache manager v2 on a disaggregated server with no cache tier below GPU (host_cache_size=0 and no disk tier), beam width 1, a single pipeline stage and full attention on every layer; elsewhere it is ignored, and KV cache manager v2 logs a warning. A request is then admitted only while the completion bound of every admitted request still fits: a generation request through its output limit and, under a KV connector without source reservations, a context request the connector may serve through the end of its prompt. The bounds are logical; no page is allocated early. A KV connector can request the same through its worker's requires_completion_reservation declaration.", + "status": "prototype", + "title": "Enable Disagg Completion Reservation", + "type": "boolean" + }, "enable_kv_pool_rebalance": { "default": false, "description": "Opt in to the KVCacheManagerV2 auto-tuner (``adjust()``) for rebalancing pool-group ratios between iterations. When True the PyExecutor calls ``adjust()`` opportunistically; the auto-tuner itself remains gated by the ``kv_pool_rebalance_config`` thresholds (2000 samples / 120s by default). When False (default) the rebalance hook is skipped entirely and pool ratios remain at their warmup-derived values. Beta: enable at your own risk. Only used when using KV cache manager v2 (experimental). This option is incompatible with dtype='fp8_ds_mla'.", @@ -4728,10 +4742,47 @@ "title": "KvCacheConfig", "type": "object" }, + "KvCacheConnectorAdmissionConfig": { + "additionalProperties": false, + "description": "Admission-time connector queries for the KVCacheManagerV2 scheduler.\n\nThe scheduler asks a connector that declares\n``KvCacheConnectorScheduler.supports_admission_query`` for its prefix while\nit admits a first context chunk, once the chunk is sized and its pages are\nallocated, and books the pages a shifting offer needs at that point.", + "properties": { + "max_offer_deferrals": { + "default": 16, + "description": "How many scheduling iterations a first context chunk may give its pages back while the KV cache cannot cover the window a connector offer moves the chunk to. Past this bound, or right after the request preempted another one for pages, it is admitted and the offer is capped to the whole blocks its chunk's pages hold. 0 never defers.", + "minimum": 0, + "status": "prototype", + "title": "Max Offer Deferrals", + "type": "integer" + }, + "max_uncharged_parks": { + "default": 0, + "description": "Per scheduling iteration, at most this many first context chunks of at least two KV blocks, whose prefix the connector will load asynchronously, are admitted without charging their tokens to max_num_tokens. They still take a batch slot and keep their pages. 0 charges every chunk.", + "minimum": 0, + "status": "prototype", + "title": "Max Uncharged Parks", + "type": "integer" + } + }, + "title": "KvCacheConnectorAdmissionConfig", + "type": "object" + }, "KvCacheConnectorConfig": { "additionalProperties": false, "description": "Configuration for the KV Cache Connector.\n\nCan be configured either by specifying a named preset via ``connector``\n(e.g. ``\"lmcache\"``), or by providing explicit ``connector_module``,\n``connector_scheduler_class``, and ``connector_worker_class`` fields.\nWhen ``connector`` is set, the module/class fields are auto-populated\nfrom the preset registry and can be omitted.", "properties": { + "admission_query": { + "anyOf": [ + { + "$ref": "#/$defs/KvCacheConnectorAdmissionConfig" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Ask the connector for a first context chunk's prefix while the scheduler admits it, instead of after the batch is final. None keeps final-batch queries; {} enables admission queries and page booking only; max_uncharged_parks > 0 also admits chunks that park for an asynchronous load without a token charge. Requires KVCacheManagerV2 with context chunking and disable_overlap_scheduler=True. Ignored, with a warning, under attention DP, for a capacity-only role, for prefix-reservation connectors, and for connectors that do not declare supports_admission_query.", + "status": "prototype" + }, "connector": { "anyOf": [ { @@ -5685,6 +5736,36 @@ "additionalProperties": false, "description": "How this server joins a Mooncake store pool.\n\nThe settings every participant shares come from the master named by\n`pool`; what is left here is per-server. Setting this makes\n`trtllm-serve` render the Mooncake client config and export\n`MOONCAKE_CONFIG_PATH`; an inherited `MOONCAKE_CONFIG_PATH` wins.", "properties": { + "async_load": { + "default": false, + "description": "Load a prefix the pool serves on background threads while the request waits outside the batch, instead of inside the executor iteration, so other requests keep computing during the transfer. A page that is gone by load time restarts only that request on a single rank or an attention-DP owner, and fails the server under tensor parallelism. Validated with disable_overlap_scheduler. Ignored unless role reads the pool ('both' or 'consumer').", + "status": "prototype", + "telemetry": { + "exclude": true + }, + "title": "Async Load", + "type": "boolean" + }, + "async_load_workers": { + "default": 1, + "description": "Background load threads per rank when async_load is set. Each has its own CUDA stream and, under stage_through_host, an equal share of staging_buffer_bytes.", + "exclusiveMinimum": 0, + "status": "prototype", + "telemetry": { + "exclude": true + }, + "title": "Async Load Workers", + "type": "integer" + }, + "group_blocks": { + "default": false, + "description": "Put every page of a block (all layer groups and tensor-parallel shards) into one Mooncake eviction group, so the master evicts a block whole rather than leaving pages that hold space but can never hit. Mostly matters under tensor parallelism. Ignored unless the role writes; needs Mooncake 0.3.13 or newer.", + "telemetry": { + "exclude": true + }, + "title": "Group Blocks", + "type": "boolean" + }, "local_hostname": { "anyOf": [ { @@ -5701,6 +5782,19 @@ }, "title": "Local Hostname" }, + "local_reuse_backoff": { + "default": "off", + "description": "How a local prefix match is handled under one-model speculative decoding, whose draft reads prompt tokens past each position and so backs every local reuse match off by that lookahead. 'off' drops the backoff: local hits are used whole, and the draft KV at the tail of a reused block may belong to a different next token, which can lower draft acceptance but does not change the output. 'align' keeps the backoff and aligns the match down to a block boundary, so the pool serves the trimmed block (a transfer) along with the rest of the prefix. Has no effect without one-model speculative decoding or for a role that does not read the pool.", + "enum": [ + "off", + "align" + ], + "telemetry": { + "exclude": true + }, + "title": "Local Reuse Backoff", + "type": "string" + }, "master_timeout": { "default": 60.0, "description": "Seconds to wait for the pool manifest to appear and the master to accept connections. Raise it when the wait spans a container start on another node. Expiring fails the server at startup.", @@ -5777,15 +5871,65 @@ "description": "Host memory each of this server's ranks contributes to the pool, as a binary size ('16GiB') or a byte count; ambiguous units such as 'GB' are refused. A node's demand is ranks_on_node x segment_size, checked against available memory at startup. Zero lends nothing, leaving the server to use only capacity its peers hold.", "title": "Segment Size" }, + "soft_pin_ttl_ms": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Soft-pin every page this server writes for this many milliseconds, so the master evicts the pool by age instead of refilling an eviction shortfall in shard order, which punches holes into prefixes still in use. Must not exceed the master's --max_kv_soft_pin_ttl (24 h by default), or every put is refused, and the master must allow evicting soft-pinned objects (Mooncake's default). Ignored unless the role writes; needs Mooncake 0.3.13 or newer.", + "telemetry": { + "exclude": true + }, + "title": "Soft Pin Ttl Ms" + }, "stage_through_host": { "default": false, - "description": "Copy pages through a pinned host buffer instead of registering the KV pools with Mooncake. For a host whose HCA cannot pin GPU pages; costs a copy each way. Ignored when role is 'capacity'.", + "description": "Copy pages through a pinned host buffer instead of registering the KV pools with Mooncake. For a host whose HCA cannot pin GPU pages; costs a copy each way. Ignored when role is 'capacity'. Pinned memory per direction is staging_buffer_bytes.", "telemetry": { "exclude": true }, "title": "Stage Through Host", "type": "boolean" }, + "staging_buffer_bytes": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "string" + } + ], + "default": "1GiB", + "description": "Pinned host memory per transfer direction when stage_through_host is set, as a binary size ('512MiB') or a byte count; ambiguous units such as 'GB' are refused. It bounds how many pages one store call stages, not how large a page may be. With async_load the load direction's share is split evenly across async_load_workers. Counted per rank in the node memory check. Ignored without stage_through_host.", + "telemetry": { + "exclude": true + }, + "title": "Staging Buffer Bytes" + }, + "stats_interval": { + "anyOf": [ + { + "exclusiveMinimum": 0, + "type": "number" + }, + { + "type": "null" + } + ], + "default": null, + "description": "Seconds between per-rank info-level log lines counting prefix lookups stopped by a hole the pool evicted from a prefix it otherwise still holds, and timing each phase of a prefix load. Only a role that reads the pool counts anything. None disables the counters at no cost.", + "telemetry": { + "exclude": true + }, + "title": "Stats Interval" + }, "transfer_batch_size": { "default": 64, "description": "Page keys per store call.", diff --git a/tensorrt_llm/serve/openai_client.py b/tensorrt_llm/serve/openai_client.py index 5a58e4a7c293..881ded0a4ff4 100644 --- a/tensorrt_llm/serve/openai_client.py +++ b/tensorrt_llm/serve/openai_client.py @@ -120,6 +120,19 @@ def _metrics_phase(role: ServerRole) -> str: return "ctx" if role is ServerRole.CONTEXT else "gen" +# Failures that leave a generation leg undelivered: the connection could not be +# opened (ClientConnectorError), or the connection closed before the whole +# request was written (aiohttp's ClientConnectionResetError, a +# ConnectionResetError). Any other failure may happen after the generation +# server accepted the request. Once the connection is open, a server-side close +# or reset before the response surfaces as ServerDisconnectedError or as +# ClientOSError (ECONNRESET) whether or not the server read the request first. +_GENERATION_LEG_RETRYABLE_ERRORS = ( + aiohttp.ClientConnectorError, + ConnectionResetError, +) + + class OpenAIClient(ABC): async def send_request( self, @@ -439,11 +452,20 @@ async def _post_with_retry( # attempt that reaches self._max_retries. _TRANSIENT_TCP_BUDGET = 0 if self._no_retry or is_remote_tail_handoff else 5 loop_max = max(max_retries, _TRANSIENT_TCP_BUDGET) + 1 + # A generation leg takes its prompt KV cache from the context server's + # send session named by disagg_request_id. + is_generation_leg = ( + self._role == ServerRole.GENERATION + and disagg_params is not None + and disagg_params.request_type == "generation_only" + ) for attempt in range(loop_max): if attempt > 0: await self._router.renew_request(request, req_id=req_id) - # Regenerate disagg_request_id on retry to avoid ID collision on workers - if attempt > 0 and self._disagg_id_generator is not None: + # A retried request gets a new disagg_request_id so it cannot collide + # with the abandoned attempt on the worker. A generation leg keeps its + # id: a new one could never be paired with the context KV cache. + if attempt > 0 and self._disagg_id_generator is not None and not is_generation_leg: dp = getattr(request, "disaggregated_params", None) if dp is not None and getattr(dp, "disagg_request_id", None) is not None: dp.disagg_request_id = await self._disagg_id_generator() @@ -457,6 +479,7 @@ async def _post_with_retry( if self._request_perf_metrics: headers[RETURN_METRICS_HEADER] = "1" headers.update(self._get_request_headers(request)) + response_started = False try: lines_yielded = 0 start_time = get_steady_clock_now_in_seconds() @@ -465,6 +488,7 @@ async def _post_with_retry( data=body, headers=headers, ) as http_response: + response_started = True response_clock = AdjustedSteadyClock() if self._request_perf_metrics: destination_time = get_steady_clock_now_in_seconds() @@ -494,6 +518,18 @@ async def _post_with_retry( role, response_metrics, ) + # An error status is checked before the stream branch too, so a + # streaming leg answered with an error body raises a + # ClientResponseError that carries the worker's status. + if http_response.status >= 400: + error_body = await http_response.text() + raise aiohttp.ClientResponseError( + http_response.request_info, + http_response.history, + status=http_response.status, + message=f"{http_response.reason}: {error_body[:2048]}", + headers=http_response.headers, + ) if not is_stream and "text/event-stream" in content_type: raise ValueError( "Received an event-stream although request stream was False" @@ -516,15 +552,6 @@ async def _post_with_retry( lines_yielded += 1 yield line else: - if http_response.status >= 400: - error_body = await http_response.text() - raise aiohttp.ClientResponseError( - http_response.request_info, - http_response.history, - status=http_response.status, - message=f"{http_response.reason}: {error_body[:2048]}", - headers=http_response.headers, - ) response_dict = await http_response.json() # yield here since python forbids return statements in async generators yield response_dict @@ -539,6 +566,20 @@ async def _post_with_retry( traceback.format_exc(), ) raise + # Once the generation server may have accepted a generation leg, + # that attempt owns the context KV cache: the transfer either went + # to it or failed and was released on the context server. A re-post + # would reuse its disagg_request_id and collide with the original + # on the worker, or wait for KV cache that is never sent. + if is_generation_leg and ( + response_started or not isinstance(e, _GENERATION_LEG_RETRYABLE_ERRORS) + ): + logger.error( + f"Client error to {url}: {e} - not retrying a generation request " + "the generation server may have accepted", + traceback.format_exc(), + ) + raise # ServerDisconnectedError and connection resets can result from # transient TCP races, such as client/server keepalive races at # burst start. Give them an extended retry budget to recover diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index 686612b2aa81..6b66e6f582c7 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -20,8 +20,8 @@ import socket import traceback import uuid -from contextlib import asynccontextmanager -from typing import Callable, Optional +from contextlib import aclosing, asynccontextmanager +from typing import Any, AsyncGenerator, Callable, Optional import aiohttp import uvicorn @@ -439,17 +439,23 @@ async def wrapper(req: request_type, raw_req: Request) -> Response: raw_req, self._perf_metrics_collector.queue_latency_seconds, self._collect_perf_metrics) response_or_generator = await entry_point(req, hooks) - self._perf_metrics_collector.total_responses.inc() if req.stream: stream = response_or_generator if isinstance(req, ResponsesRequest): - # A relayed Responses stream that fails early ends with - # an `error` event; other protocols end on a sentinel. + # A relayed Responses stream that fails ends with an + # `error` event, even before its first frame. stream = guard_responses_stream(stream, stream_error_event) + # Hold the headers until the first chunk: a chat or + # completion stream that fails before it (e.g. a generation + # server's KV-cache transfer error) becomes an HTTP error, + # not a 200 whose body is cut off. + stream = await self._start_stream(stream) + self._perf_metrics_collector.total_responses.inc() return StreamingResponse( content=stream, media_type="text/event-stream") + self._perf_metrics_collector.total_responses.inc() # by_alias: e.g. text.format.schema goes out under its wire name. return JSONResponse( content=response_or_generator.model_dump(by_alias=True)) @@ -459,6 +465,39 @@ async def wrapper(req: request_type, raw_req: Request) -> Response: return self._handle_exception(e) return wrapper + @staticmethod + async def _start_stream( + stream: AsyncGenerator[Any, None]) -> AsyncGenerator[Any, None]: + """Pull a relayed stream's first chunk before the response headers go out. + + A worker that fails before its first chunk then raises here and becomes an + HTTP error response instead of a 200 whose body is cut off. The returned + relay is already started, so closing it before its first chunk still closes + ``stream`` and releases the request's routing state. + + Args: + stream: The relayed response stream from the disaggregated service. + + Returns: + ``stream`` itself when it is already exhausted; otherwise a started + relay that yields the first chunk and then the rest of ``stream``. + """ + try: + first_chunk = await anext(stream) + except StopAsyncIteration: + return stream + + async def relay() -> AsyncGenerator[Any, None]: + async with aclosing(stream): + yield None # Consumed below. + yield first_chunk + async for chunk in stream: + yield chunk + + relayed = relay() + await anext(relayed) + return relayed + async def anthropic_messages(self, request: AnthropicMessagesRequest, raw_request: Request) -> Response: """Serve Anthropic Messages through the disaggregated chat pipeline.""" diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index f940f2a80937..37b44c6714df 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -31,6 +31,11 @@ "kind": "value", "path": "attention_dp_config.enable_kv_cache_aware_routing" }, + { + "capture_policy": "bool", + "kind": "value", + "path": "attention_dp_config.enable_work_aware_routing" + }, { "capture_policy": "bool", "kind": "value", @@ -673,6 +678,11 @@ "kind": "value", "path": "kv_cache_config.enable_block_reuse" }, + { + "capture_policy": "bool", + "kind": "value", + "path": "kv_cache_config.enable_disagg_completion_reservation" + }, { "capture_policy": "bool", "kind": "value", @@ -867,6 +877,16 @@ "kind": "value", "path": "kv_cache_config.use_uvm" }, + { + "capture_policy": "int", + "kind": "value", + "path": "kv_connector_config.admission_query.max_offer_deferrals" + }, + { + "capture_policy": "int", + "kind": "value", + "path": "kv_connector_config.admission_query.max_uncharged_parks" + }, { "allowed_values": [ "lmcache", diff --git a/tests/integration/defs/llmapi/test_llm_api_connector_mooncake.py b/tests/integration/defs/llmapi/test_llm_api_connector_mooncake.py index 097978fb3361..33108ccdca78 100644 --- a/tests/integration/defs/llmapi/test_llm_api_connector_mooncake.py +++ b/tests/integration/defs/llmapi/test_llm_api_connector_mooncake.py @@ -18,15 +18,17 @@ engine publishes its KV cache into a real Mooncake pool, a second engine replays it, and the generated tokens have to agree. That covers the seams no narrower test reaches, namely the KV cache layout a `KVCacheManagerV2` actually -reports, per-rank namespacing under tensor parallelism, and the scheduler -interaction when a request is preempted while the connector is still reading -its pages. +reports, per-rank namespacing under tensor parallelism, the key space attention +DP owners share, the draft layers a one-model drafter keeps in the target pool on +a disaggregated context server, and the scheduler interaction when a request is +preempted while the connector is still reading its pages. One master and pool serve the whole module and nothing else on the machine. Each test gets a key namespace of its own; see `mooncake_pool`. """ import contextlib +import glob import json import math import os @@ -36,12 +38,15 @@ import pytest from test_common.mooncake_utils import running_master_on_free_ports -from tensorrt_llm import LLM, SamplingParams +from tensorrt_llm import LLM, DisaggregatedParams, SamplingParams from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.config import ( MooncakeStoreConnectorConfig, parse_size, ) -from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.ledger import read_segments +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.ledger import ( + read_segments, + segments_dir, +) from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.master import ( POOL_MANIFEST_NAME, provision_pool, @@ -49,7 +54,14 @@ from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.scheduler import ( MooncakeStoreConnectorScheduler, ) -from tensorrt_llm.llmapi.llm_args import KvCacheConfig, KvCacheConnectorConfig, MooncakeStoreConfig +from tensorrt_llm.llmapi import RequestError, SchedulingParams +from tensorrt_llm.llmapi.llm_args import ( + CacheTransceiverConfig, + Eagle3DecodingConfig, + KvCacheConfig, + KvCacheConnectorConfig, + MooncakeStoreConfig, +) from ..conftest import llm_models_root from .test_llm_api_connector import E2E_MIN_TOKEN_AGREEMENT @@ -141,17 +153,22 @@ def pool_capacity(config_path: str, segment_size): store.close() -def _set_namespace(config_path: str, namespace: str) -> None: - """Point the rendered client config at `namespace`, renaming into place.""" +def _set_client_settings(config_path: str, **settings) -> None: + """Rewrite keys of the rendered client config, renaming into place.""" with open(config_path) as handle: config = json.load(handle) - config["namespace"] = namespace + config.update(settings) staging = f"{config_path}.partial" with open(staging, "w") as handle: json.dump(config, handle, indent=2) os.replace(staging, config_path) +def _set_namespace(config_path: str, namespace: str) -> None: + """Point the rendered client config at `namespace`.""" + _set_client_settings(config_path, namespace=namespace) + + @pytest.fixture(scope="module") def _mooncake_pool_for_module(tmp_path_factory): """One master, one rendered client config and one lent segment. @@ -248,6 +265,25 @@ def recording(self, request, num_computed_tokens): return counts +@pytest.fixture +def connector_answers(monkeypatch): + """Every answer the leader gave, as (matched tokens, whether asynchronous). + + Only usable when the executor runs in this process, like + `matched_token_counts`. + """ + answers = [] + original = MooncakeStoreConnectorScheduler.get_num_new_matched_tokens + + def recording(self, request, num_computed_tokens): + result = original(self, request, num_computed_tokens) + answers.append(tuple(result)) + return result + + monkeypatch.setattr(MooncakeStoreConnectorScheduler, "get_num_new_matched_tokens", recording) + return answers + + @pytest.fixture def single_process_worker(monkeypatch): monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") @@ -296,6 +332,59 @@ def test_mooncake_e2e_cross_engine_prefix_reuse( assert_tokens_agree(cold_token_ids, warm_token_ids, context="Cross-engine prefix reuse") +@pytest.mark.threadleak(enabled=False) +def test_mooncake_e2e_async_load_matches_sync_tokens( + single_process_worker, mooncake_pool, connector_answers +): + """A prefix loaded on background threads replays like one loaded in the iteration. + + The cold engine computes and publishes the prompt synchronously. The warm + engine reads the same pool with `async_load` and two load threads, so the + request is parked while a load thread fills its pages through that thread's + own staging slots, and only rejoins the batch once they landed. Pages reported before they + landed, or scattered into the wrong slots, diverge early. + """ + sampling_params = SamplingParams(max_tokens=32, ignore_eos=True) + + cold = LLM(**llm_kwargs(mooncake_pool)) + try: + cold_output = cold.generate([PROMPT], sampling_params) + finally: + cold.shutdown() + assert connector_answers and not any(tokens for tokens, _ in connector_answers), ( + f"The first run already found tokens in the pool (answers {connector_answers}), " + "so it was not a cold start and the comparison below is meaningless." + ) + connector_answers.clear() + + # The module's tests share this config, so it is restored whatever happens. + _set_client_settings(mooncake_pool.config_path, async_load=True, async_load_workers=2) + try: + warm = LLM(**llm_kwargs(mooncake_pool)) + try: + warm_output = warm.generate([PROMPT], sampling_params) + finally: + warm.shutdown() + finally: + _set_client_settings(mooncake_pool.config_path, async_load=False, async_load_workers=1) + + hits = [(tokens, asynchronous) for tokens, asynchronous in connector_answers if tokens] + assert hits, ( + f"The second engine read nothing back from the pool (answers {connector_answers}), " + "so it recomputed the prefix and this test would pass with the connector disabled." + ) + assert all(asynchronous for _, asynchronous in hits), ( + f"The second engine offered a prefix synchronously (answers {connector_answers}), " + "so async_load did not reach its scheduler adapter." + ) + + assert_tokens_agree( + list(cold_output[0].outputs[0].token_ids), + list(warm_output[0].outputs[0].token_ids), + context="Prefix reuse through asynchronous loads", + ) + + @pytest.mark.skip_less_device(2) @pytest.mark.threadleak(enabled=False) def test_mooncake_e2e_prefix_reuse_across_tensor_parallel_ranks(mooncake_pool): @@ -343,6 +432,143 @@ def test_mooncake_e2e_prefix_reuse_across_tensor_parallel_ranks(mooncake_pool): ) +def _segment_record_stamps(run_dir: str) -> dict: + """Each segment record under `run_dir`, with the inode and time it was last written.""" + stamps = {} + for path in glob.glob(os.path.join(segments_dir(run_dir), "*.json")): + with contextlib.suppress(FileNotFoundError): + info = os.stat(path) + stamps[path] = (info.st_ino, info.st_mtime_ns) + return stamps + + +def _ranks_recorded_since(run_dir: str, stamps: dict) -> set: + """Ranks whose segment record is new or rewritten since `stamps` was taken. + + Not a count of records: a record is named by host and pid, and the test + session reuses MPI worker processes across engines, so an engine can + rewrite an earlier engine's record instead of adding one. A rewrite renames + a new file into place, which changes the inode. + """ + ranks = set() + for path, stamp in _segment_record_stamps(run_dir).items(): + if stamps.get(path) != stamp: + with open(path) as handle: + ranks.add(int(json.load(handle)["rank"])) + return ranks + + +def _rank0_rows(stats) -> list: + """Rank 0's iteration statistics. + + Under attention DP rank 0 reports one row per owner, each tagged with + `attentionDpRank`. Connector and KV cache counters ride on rank 0's rows + only. + """ + return [entry for entry in stats if entry.get("attentionDpRank", 0) == 0] + + +def _connector_cached_tokens(stats) -> int: + """Prompt tokens rank 0's connector served, summed over iteration statistics.""" + return sum(entry.get("iterConnectorCachedTokens", 0) for entry in _rank0_rows(stats)) + + +@pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("disable_overlap_scheduler", [False, True]) +@pytest.mark.skip_less_device(2) +def test_mooncake_e2e_prefix_reuse_across_attention_dp_owners( + mooncake_pool, disable_overlap_scheduler +): + """Save and reuse at TP 2 with attention DP, where each owner holds every head. + + Every owner keys its pages as attention shard 0 of 1, so a prefix one owner + saved is a hit for the other. Under the per-rank keys tensor parallelism + uses, an owner would ask for both ranks' shards of a page only one of them + could have written, and nothing would ever be reused. + + The cold engine serves one request pinned to owner 1, so rank 0 saves + nothing and every page of the prompt in the pool is one owner 1 saved. The + warm engine starts with empty GPU caches and serves four copies with one + request per owner at a time, so both owners serve the prompt and each one's + first copy finds nothing in its own GPU cache. The assertions read rank 0's + rows of the iteration statistics. Rank 0 serves at least one of the copies + and saved none of the prompt's pages before its first copy looked them up, + so a hit there is a page owner 1 saved: a read across owners, not a + same-owner hit. + + This runs aggregated serving, which the connector documents as not + supported yet under attention DP because of a skipped-pass hazard on + generation pages; one request per owner and short generations do not + provoke it, and the key space under test is the same on a disaggregated + context server. + """ + sampling_params = SamplingParams(max_tokens=32, ignore_eos=True) + kwargs = llm_kwargs( + mooncake_pool, + tensor_parallel_size=2, + enable_attention_dp=True, + max_batch_size=1, + disable_overlap_scheduler=disable_overlap_scheduler, + enable_iter_perf_stats=True, + # Keep every iteration rather than the most recent 1000, so the sums + # below cover the whole run. + iter_stats_max_iterations=-1, + ) + records_before = _segment_record_stamps(mooncake_pool.run_dir) + + cold = LLM(**kwargs) + try: + # Strict, so the router cannot move it to rank 0 for balance. + cold_output = cold.generate( + [PROMPT], + sampling_params, + scheduling_params=SchedulingParams(attention_dp_rank=1, attention_dp_relax=False), + ) + cold_stats = cold.get_stats(timeout=30) + finally: + cold.shutdown() + cold_token_ids = list(cold_output[0].outputs[0].token_ids) + + joined = _ranks_recorded_since(mooncake_pool.run_dir, records_before) + assert joined == {0, 1}, ( + "Both attention-DP owners should have joined the pool, but this run recorded " + f"ranks {sorted(joined)}. An owner that never opened a store handle saves " + "nothing and reads every prefix as a miss." + ) + # The test's own key namespace keeps the pool cold. Rank 0's rows cannot + # see owner 1's load, but they do show where the request ran: attention-DP + # padding requests are left out of them. + cold_context_tokens = sum( + entry.get("inflightBatchingStats", {}).get("numCtxTokens", 0) + for entry in _rank0_rows(cold_stats) + ) + assert cold_context_tokens == 0, ( + f"Rank 0 prefilled {cold_context_tokens} tokens in the cold engine, so the cold " + "request did not stay on owner 1 and a hit on rank 0 below could be a page rank 0 " + "saved itself." + ) + + warm = LLM(**kwargs) + try: + warm_outputs = warm.generate([PROMPT] * 4, sampling_params) + warm_stats = warm.get_stats(timeout=30) + finally: + warm.shutdown() + + assert warm_stats, "The warm engine returned no iteration statistics to assert on." + assert _connector_cached_tokens(warm_stats) > 0, ( + f"Rank 0 served no prompt tokens from the pool across {len(warm_stats)} " + "iterations, so it recomputed a prefix owner 1 stored and no owner read " + "another owner's pages." + ) + for index, output in enumerate(warm_outputs): + assert_tokens_agree( + cold_token_ids, + list(output.outputs[0].token_ids), + context=f"Copy {index} replayed under attention DP", + ) + + @pytest.mark.threadleak(enabled=False) def test_mooncake_e2e_chunked_prefill_survives_a_small_cache(mooncake_pool): """Chunked prefill against a cache too small to hold every request. @@ -457,3 +683,139 @@ def test_mooncake_e2e_chunked_prefill_survives_a_small_cache(mooncake_pool): f"prompt tokens, past the {budget}-token budget. Requests are being " "preempted and recomputed repeatedly rather than making progress." ) + + +#: Target and one-model draft checkpoints under `LLM_MODELS_ROOT`, per drafter: +#: the pair the single-GPU disaggregated EAGLE3 test serves. +SPECULATIVE_PREFILL_MODELS = { + "eagle3": ("Qwen3/Qwen3-8B", "Qwen3/qwen3_8b_eagle3"), +} + + +# Two 8B engines, a context and a generation server, share the device. +@pytest.mark.skip_less_device_memory(80000) +@pytest.mark.threadleak(enabled=False) +@pytest.mark.parametrize("drafter", ["eagle3"]) +def test_mooncake_e2e_speculative_prefill_reuse( + single_process_worker, mooncake_pool, matched_token_counts, drafter +): + """A disaggregated context server with a one-model draft replays a prefix it stored. + + One-model EAGLE3 keeps its draft layers in the target KV cache, so a page + the pool serves carries draft KV as well, keyed by the prompt token the + draft reads past each block. The draft token a context pass hands to the + generation server is computed from that draft KV: were the draft layers + missing from the layout registered with the connector, or keyed without + their lookahead, the warm engine's drafts would degenerate rather than tie. + So the first token and the first draft token of every prompt must agree + between a cold and a warm context engine, and so must the generation the + generation server continues from each. + + Chunked prefill below the prompt length exercises the KV slots one-model + speculation allocates past every chunk, which must not be published before + a pass computes them. A request that is not context-only is refused at + admission, and the engine keeps serving afterwards. + + Both engines run in this process, so the generation server pulls each + context request's KV and the context engine shuts down with nothing in + flight. The pool keys these pages under the module's model key, inside a + namespace of this test's own. + """ + target, draft = SPECULATIVE_PREFILL_MODELS[drafter] + + def engine(*, context: bool) -> dict: + """One engine's settings; each gets a speculative config of its own to resolve.""" + settings = dict( + model=f"{llm_models_root()}/{target}", + speculative_config=Eagle3DecodingConfig( + max_draft_len=3, speculative_model=f"{llm_models_root()}/{draft}" + ), + cache_transceiver_config=CacheTransceiverConfig( + backend="NIXL", transceiver_runtime="PYTHON" + ), + max_batch_size=4, + # Bounds the warmup allocation, which would otherwise size itself + # against the checkpoint's position limit. + max_seq_len=1024, + ) + if context: + # Two blocks per chunk, shorter than every prompt. + settings.update(enable_chunked_prefill=True, max_num_tokens=64) + return llm_kwargs(mooncake_pool, **settings) + return llm_kwargs(mooncake_pool, kv_connector_config=None, **settings) + + # Distinct from their first token, so no prompt reuses another's blocks. + prompts = [f"{index}. {PROMPT}" for index in range(3)] + sampling_params = SamplingParams(max_tokens=16, ignore_eos=True) + + def serve(context, generation): + """Prefill on `context`, continue on `generation`, and return what each produced.""" + context_outputs = context.generate( + prompts, + sampling_params, + disaggregated_params=[ + DisaggregatedParams(request_type="context_only") for _ in prompts + ], + ) + handoffs = [] + for output in context_outputs: + params = output.disaggregated_params + assert params is not None and params.request_type == "context_only" + assert params.first_gen_tokens, "The context pass handed over no first token." + assert params.draft_tokens, ( + "The context pass handed over no draft tokens, so the draft layers " + "never ran and this test cannot see the draft KV it is about." + ) + handoffs.append((list(params.first_gen_tokens), list(params.draft_tokens))) + params.request_type = "generation_only" + generation_outputs = generation.generate( + prompts, + sampling_params, + disaggregated_params=[output.disaggregated_params for output in context_outputs], + ) + return handoffs, [list(output.outputs[0].token_ids) for output in generation_outputs] + + generation = LLM(**engine(context=False)) + try: + cold = LLM(**engine(context=True)) + try: + cold_handoffs, cold_generations = serve(cold, generation) + finally: + cold.shutdown() + + assert matched_token_counts and not any(matched_token_counts), ( + f"The cold context engine already found tokens in the pool (matched " + f"counts {matched_token_counts}), so the comparison below is meaningless." + ) + matched_token_counts.clear() + + warm = LLM(**engine(context=True)) + try: + with pytest.raises(RequestError, match="context-only"): + warm.generate([PROMPT], SamplingParams(max_tokens=4)) + warm_handoffs, warm_generations = serve(warm, generation) + finally: + warm.shutdown() + finally: + generation.shutdown() + + assert matched_token_counts and max(matched_token_counts) > 0, ( + f"The warm context engine read nothing back from the pool (matched counts " + f"{matched_token_counts}), so it recomputed every prefix and this test " + "would pass with the connector disabled." + ) + for index, (cold_handoff, warm_handoff) in enumerate(zip(cold_handoffs, warm_handoffs)): + (cold_first, cold_drafts), (warm_first, warm_drafts) = cold_handoff, warm_handoff + assert warm_first == cold_first, ( + f"Prompt {index}: the first token changed from {cold_first} to {warm_first} " + "once the prefix came from the pool." + ) + assert warm_drafts[0] == cold_drafts[0], ( + f"Prompt {index}: the first draft token changed from {cold_drafts} to " + f"{warm_drafts} once the prefix came from the pool, so the draft read " + "draft KV that was missing or belonged to other tokens." + ) + for index, (cold_tokens, warm_tokens) in enumerate(zip(cold_generations, warm_generations)): + assert_tokens_agree( + cold_tokens, warm_tokens, context=f"Prompt {index} continued after a warm prefill" + ) diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 1bce3cdda018..06515b5e2a18 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -246,6 +246,7 @@ l0_a10: - llmapi/test_llm_api_connector.py::test_connector_transfers_only_in_window_blocks_to_the_sliding_group[qwen3_forced_windows] - llmapi/test_llm_api_connector_mooncake.py::test_mooncake_e2e_cross_engine_prefix_reuse - llmapi/test_llm_api_connector_mooncake.py::test_mooncake_e2e_chunked_prefill_survives_a_small_cache + - llmapi/test_llm_api_connector_mooncake.py::test_mooncake_e2e_async_load_matches_sync_tokens # third-party policy checks CPU-only - thirdparty/test_cmake_third_party.py::test_cmake_listfiles - thirdparty/test_git_modules.py::test_gitmodules diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index 8354a7290642..348750c14ba8 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -35,6 +35,8 @@ l0_dgx_h100: - llmapi/test_llm_api_pytorch_bart.py::test_bart_pytorch_continuous_admission_replays_encoder_and_mixed_cuda_graphs[tp2] - llmapi/test_llm_api_pytorch_whisper.py::test_whisper_pytorch_feature_combinations[fp32-kv-v2-graphs-off-greedy-tp2] - llmapi/test_llm_api_connector_mooncake.py::test_mooncake_e2e_prefix_reuse_across_tensor_parallel_ranks + - llmapi/test_llm_api_connector_mooncake.py::test_mooncake_e2e_prefix_reuse_across_attention_dp_owners[False] + - llmapi/test_llm_api_connector_mooncake.py::test_mooncake_e2e_prefix_reuse_across_attention_dp_owners[True] # ------------- Disaggregated serving tests --------------- - accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_auto_dtype[False-True] - accuracy/test_disaggregated_serving.py::TestQwen3_8B::test_auto_dtype[True-True] diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 209c866aa838..e382d37c948b 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -207,6 +207,7 @@ l0_h100: - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_simple_qwen3[True-True-Qwen3-8B-FP8] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_llama_context_capacity[False-False-DeepSeek-V3-Lite-fp8/fp8] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_spec_dec_batch_slot_limit[False-Qwen3-8B-eagle3-Qwen3-8B] + - llmapi/test_llm_api_connector_mooncake.py::test_mooncake_e2e_speculative_prefill_reuse[eagle3] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_cancel_gen_requests[Qwen3-0.6B] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_logits[False-Qwen3-0.6B] - disaggregated/test_disaggregated_single_gpu.py::test_disaggregated_logits[True-Qwen3-0.6B] diff --git a/tests/unittest/_torch/disaggregation/conftest.py b/tests/unittest/_torch/disaggregation/conftest.py index 769c91c0c81a..b7716abac812 100644 --- a/tests/unittest/_torch/disaggregation/conftest.py +++ b/tests/unittest/_torch/disaggregation/conftest.py @@ -12,6 +12,12 @@ def inflight_cancel(monkeypatch): monkeypatch.setattr(coordinator_module, "is_disagg_inflight_cancel_enabled", lambda: True) +@pytest.fixture +def no_inflight_cancel(monkeypatch): + """Pin the in-flight cancellation opt-in off; the env read is process-cached.""" + monkeypatch.setattr(coordinator_module, "is_disagg_inflight_cancel_enabled", lambda: False) + + @pytest.fixture def clock(monkeypatch): """Freeze the coordinator's ``time.monotonic``; tests advance ``clock["t"]``.""" diff --git a/tests/unittest/_torch/disaggregation/test_benchmark_disagg.py b/tests/unittest/_torch/disaggregation/test_benchmark_disagg.py index 9c7eeeb18722..e37e2246cc65 100644 --- a/tests/unittest/_torch/disaggregation/test_benchmark_disagg.py +++ b/tests/unittest/_torch/disaggregation/test_benchmark_disagg.py @@ -1102,6 +1102,10 @@ def test_warmup_skips_throttle(self, mock_get_from_waiting_queue): "tensorrt_llm._torch.pyexecutor.scheduler.adp_router.KVCacheAwareADPRouter", id="KVCacheAwareADPRouter", ), + pytest.param( + "tensorrt_llm._torch.pyexecutor.scheduler.adp_router.WorkAwareADPRouter", + id="WorkAwareADPRouter", + ), ] @@ -1135,6 +1139,9 @@ def _make_unscheduled_request_item(req_id=None): req_item.request.py_scheduling_params = None req_item.request.input_token_ids = [1, 2, 3] req_item.id = req_id if req_id is not None else id(req_item) + # RequestQueueItem.child_req_ids is Optional[list]; a bare Mock attribute + # is truthy and not iterable. + req_item.child_req_ids = None return req_item @@ -1234,7 +1241,8 @@ def test_cap_not_applied_when_below_max(self, router_cls_path): # ceil(2.0 * ceil(8/4)) = ceil(2.0 * 2) = 4. expected_natural = 4 else: - # DefaultADPRouter natural value is ceil(8/4) = 2. + # DefaultADPRouter natural value is ceil(8/4) = 2; WorkAwareADPRouter + # spreads eight equal requests two per rank. expected_natural = 2 assert expected == expected_natural, ( f"{cls_name}: natural expected was {expected_natural}, got {expected}" diff --git a/tests/unittest/_torch/disaggregation/test_disagg_coordinator_errors.py b/tests/unittest/_torch/disaggregation/test_disagg_coordinator_errors.py index 944271cf6cba..4864f29f5b5a 100644 --- a/tests/unittest/_torch/disaggregation/test_disagg_coordinator_errors.py +++ b/tests/unittest/_torch/disaggregation/test_disagg_coordinator_errors.py @@ -576,3 +576,96 @@ def test_cancelling_part_of_a_saturated_receive_batch_settles_only_those_request assert all(req.state == LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE for req in kept) assert transceiver.check_gen_transfer_complete() assert len(executor.resource_manager.free_resources.call_args_list) == 8 + + +# -- generation transfer timeouts through the real executor error path ------- +# +# The coordinator-level timeout tests cannot see what happens after +# fail_requests: FakeExecutorEffects never runs PyExecutor._handle_errors. +# These follow a timed-out receive through it, where a request still in its +# receive state would be parked for a transfer that no longer exists and then +# dropped without a response or a free. + + +def _cancel_only_marks_pending_receives(transceiver: FakeKvCacheTransceiver) -> list: + """Model the Python transceiver with a receive in flight: cancelling it only + marks the session cancelled and returns False, and the next status poll + reports the cancelled session. Returns the ids of the deferred cancels.""" + deferred = [] + cancel_finished_transfer = transceiver.cancel_request + + def cancel_request(req) -> bool: + if req.py_request_id not in transceiver._pending_recvs: + return cancel_finished_transfer(req) + deferred.append(req.py_request_id) + transceiver.cancel_recv_remotely(req) + return False + + transceiver.cancel_request = cancel_request + return deferred + + +@pytest.mark.parametrize( + "cancel_finishes_at_once, error_msg", + [ + (True, "Request 1 timed out"), + (False, "Error in kv cache transfer for generation requests"), + ], +) +def test_idle_generation_timeout_is_answered_and_freed_through_the_real_error_path( + monkeypatch, clock, no_inflight_cancel, cancel_finishes_at_once: bool, error_msg: str +) -> None: + """Idle loop heads in the executor's order (transfer poll, timeout check, + held-failure sweep) cancel a receive whose KV never arrived. The real + error path answers the request once, frees its resources once and parks + nothing, whether the cancel finishes at once or is reported by the next + status poll.""" + executor, coordinator, transceiver = _executor_with_real_error_path(monkeypatch) + transceiver.kv_transfer_timeout_ms = 1000 + req = _gen_request(1) + transceiver.request_and_receive_async(req) + executor.active_requests.append(req) + req.py_kv_transfer_start_time = clock["t"] + clock["t"] += 2.0 + deferred = [] if cancel_finishes_at_once else _cancel_only_marks_pending_receives(transceiver) + + for _ in range(3): + coordinator.poll_gen_transfers() + coordinator.check_transfer_timeouts() + executor._sweep_transfer_held_failures() + + (response,) = executor.responses[1] + assert response.error_msg == error_msg + assert req.state == LlmRequestState.GENERATION_COMPLETE + assert executor.active_requests == [] + executor.resource_manager.free_resources.assert_called_once_with(req) + assert executor._held_transfer_failures() == {} + if cancel_finishes_at_once: + assert transceiver.call_log.count("cancel_request:1") == 1 + else: + assert deferred == [1] + + +def test_timed_out_receive_cancelled_by_the_response_pass_is_answered_not_parked( + monkeypatch, no_inflight_cancel +) -> None: + """The response pass of a busy server cancels a timed-out receive and + fails it through the same helper; the request must be answered and + freed, not parked and then dropped.""" + executor, _, transceiver = _executor_with_real_error_path(monkeypatch) + transceiver.kv_transfer_timeout_ms = 1000 + req = _gen_request(1) + transceiver.request_and_receive_async(req) + executor.active_requests.append(req) + req.py_kv_transfer_timed_out = True + + executor._handle_responses() + executor._sweep_transfer_held_failures() + + assert transceiver.call_log.count("cancel_request:1") == 1 + (response,) = executor.responses[1] + assert response.error_msg == "Request 1 timed out" + assert req.state == LlmRequestState.GENERATION_COMPLETE + assert executor.active_requests == [] + executor.resource_manager.free_resources.assert_called_once_with(req) + assert executor._held_transfer_failures() == {} diff --git a/tests/unittest/_torch/disaggregation/test_disagg_coordinator_progress.py b/tests/unittest/_torch/disaggregation/test_disagg_coordinator_progress.py index e1290f4e7c97..45f9bdc0296a 100644 --- a/tests/unittest/_torch/disaggregation/test_disagg_coordinator_progress.py +++ b/tests/unittest/_torch/disaggregation/test_disagg_coordinator_progress.py @@ -16,6 +16,7 @@ from coordinator_harness import CoordinatorHarness, TransferRequest from fake_dist import FakeDistGroup +from tensorrt_llm._torch.disaggregation.orchestration import coordinator as coordinator_module from tensorrt_llm.bindings import LlmRequestState from tensorrt_llm.disaggregated_params import DisaggScheduleStyle @@ -254,3 +255,157 @@ def test_context_gate_is_entered_once_per_rank_whatever_arrived_locally() -> Non ["prepare_context_requests:[1]"], ] assert [h.dist.calls for h in ranks] == [[], []] + + +# -- generation transfer timeout on an idle server --------------------------- + +_GEN_ERROR = "Error in kv cache transfer for generation requests" + + +def _cancel_only_marks_inflight_receives(transceiver) -> None: + """Model the Python transceiver with a receive in flight: cancelling it only + marks the session cancelled and returns False; the next status poll + reports the cancelled session.""" + cancel_finished_transfer = transceiver.cancel_request + + def cancel_request(req) -> bool: + if req.py_request_id not in transceiver._pending_recvs: + return cancel_finished_transfer(req) + transceiver.call_log.append(f"cancel_request:{req.py_request_id}") + transceiver.cancel_recv_remotely(req) + return False + + transceiver.cancel_request = cancel_request + + +def _executor_drops_failed_requests(h: CoordinatorHarness) -> None: + """The executor's error path removes the requests it fails from active_requests.""" + record_failure = h.effects.fail_requests + + def fail_requests(error_msg, requests, *, charge_budget) -> None: + record_failure(error_msg, requests, charge_budget=charge_budget) + for req in requests: + if req in h.active: + h.active.remove(req) + + h.effects.fail_requests = fail_requests + + +def _idle_loop_heads(h: CoordinatorHarness, iterations: int) -> None: + """Loop heads of an executor with nothing scheduled: the transfer poll, + then the timeout check, and no response pass.""" + for _ in range(iterations): + h.coordinator.poll_gen_transfers() + h.coordinator.check_transfer_timeouts() + + +@pytest.mark.parametrize("inflight_requested", [False, True]) +@pytest.mark.parametrize( + "cancel_finishes_at_once, error_msg", + [(True, "Request 1 timed out"), (False, _GEN_ERROR)], +) +def test_idle_generation_server_fails_a_receive_whose_kv_never_arrives( + monkeypatch, clock, cancel_finishes_at_once: bool, error_msg: str, inflight_requested: bool +) -> None: + """Nothing is scheduled, so the response pass never runs. The loop head + alone must cancel the receive once kv_transfer_timeout_ms passes and fail + the request, whether the cancel finishes at once or is reported by the + next status poll. Requesting in-flight cancellation from a transceiver + that does not support it takes the same path.""" + monkeypatch.setattr( + coordinator_module, "is_disagg_inflight_cancel_enabled", lambda: inflight_requested + ) + h = CoordinatorHarness(kv_transfer_timeout_ms=1000) + if not cancel_finishes_at_once: + _cancel_only_marks_inflight_receives(h.transceiver) + _executor_drops_failed_requests(h) + req = _receiving(h, 1, started_at=clock["t"]) + clock["t"] += 2.0 + + _idle_loop_heads(h, 3) + + assert h.transceiver.call_log.count("cancel_request:1") == 1 + assert h.effects.failed == [(error_msg, [req], False)] + assert h.active == [] + assert req.state == LlmRequestState.DISAGG_TRANS_ERROR + assert h.registry.removed == ([req] if cancel_finishes_at_once else []) + + +def test_idle_generation_timeout_asks_a_refused_cancel_only_once(no_inflight_cancel, clock) -> None: + """C++ semantics without in-flight cancellation: a receive already in + flight cannot be cancelled, and nothing is reported for the refusal. + Asking again every loop head would only repeat the transceiver's warning, + so the cancel is asked once and the request waits until the transceiver + settles the receive.""" + h = CoordinatorHarness(kv_transfer_timeout_ms=1000) + _executor_drops_failed_requests(h) + req = _receiving(h, 1, started_at=clock["t"]) + clock["t"] += 2.0 + + def refuse(request) -> bool: + h.transceiver.call_log.append(f"cancel_request:{request.py_request_id}") + return False + + h.transceiver.cancel_request = refuse + + _idle_loop_heads(h, 5) + + assert h.transceiver.call_log.count("cancel_request:1") == 1 + assert h.effects.failed == [] + assert h.active == [req] + assert req.state == LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS + + h.transceiver.finish_recv(req, outcome="error") + h.coordinator.poll_gen_transfers() + + assert h.effects.failed == [(_GEN_ERROR, [req], False)] + assert h.active == [] + assert h.transceiver.call_log.count("cancel_request:1") == 1 + + +def test_idle_generation_timeout_under_adp_waits_for_the_synced_drain( + no_inflight_cancel, clock +) -> None: + """Under multi-rank ADP the loop head cancels the receive and ends its + receive state at once, but the error response waits for the synced + drain. The loop head itself enters no collective.""" + group = FakeDistGroup(world_size=2, tp_size=2) + ranks = _ranks(group, kv_transfer_timeout_ms=1000, enable_attention_dp=True) + req = _receiving(ranks[0], 1, started_at=clock["t"]) + clock["t"] += 2.0 + ranks[0].coordinator.check_transfer_timeouts() + + group.run(lambda rank: ranks[rank].coordinator.poll_gen_transfers()) + + assert ranks[0].transceiver.call_log.count("cancel_request:1") == 1 + assert ranks[0].active == [] + assert req.state == LlmRequestState.DISAGG_TRANS_ERROR + assert [h.effects.failed for h in ranks] == [[], []] + assert [h.dist.calls for h in ranks] == [[], []] + + group.run(lambda rank: ranks[rank].coordinator.handle_timeouts_synced()) + + assert ranks[0].dist.calls == [("tp_allgather_int64", [True])] + assert ranks[1].dist.calls == [("tp_allgather_int64", [False])] + assert ranks[0].effects.failed == [("Request timed out (KV transfer)", [req], False)] + assert ranks[1].effects.failed == [("Request timed out (KV transfer)", [], False)] + + +def test_idle_generation_timeout_leaves_fresh_and_user_canceled_receives( + no_inflight_cancel, clock +) -> None: + """A receive inside its timeout is not touched, and a timed-out receive + the user cancelled is left to the executor's cancel path.""" + h = CoordinatorHarness(kv_transfer_timeout_ms=1000) + fresh = _receiving(h, 1, started_at=clock["t"]) + user_canceled = _receiving(h, 2, started_at=clock["t"] - 2.0) + h.registry.canceled.append(2) + h.coordinator.check_transfer_timeouts() + + h.coordinator.poll_gen_transfers() + + assert not fresh.py_kv_transfer_timed_out and user_canceled.py_kv_transfer_timed_out + assert "cancel_request:1" not in h.transceiver.call_log + assert "cancel_request:2" not in h.transceiver.call_log + assert h.active == [fresh, user_canceled] + assert h.effects.failed == [] diff --git a/tests/unittest/_torch/disaggregation/test_disagg_coordinator_transfers.py b/tests/unittest/_torch/disaggregation/test_disagg_coordinator_transfers.py index 2e9fb0275eb3..0d74a736884e 100644 --- a/tests/unittest/_torch/disaggregation/test_disagg_coordinator_transfers.py +++ b/tests/unittest/_torch/disaggregation/test_disagg_coordinator_transfers.py @@ -401,6 +401,50 @@ def test_timed_out_requests_under_adp_fail_only_at_the_synced_drain() -> None: assert len(h.effects.failed) == 1 +def test_fail_timed_out_ends_the_receive_state_of_cancelled_requests() -> None: + """cancel_request leaves a receive in progress. Failing the request must + end that state, or the executor's error path would park the failure for a + transfer that no longer exists and then drop it without a response. A + receive that completed after its deadline keeps its state.""" + h = _Harness() + cancelled = _Request( + 1, + is_context_only_request=False, + state=LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS, + ) + completed_late = _Request( + 2, + is_context_only_request=False, + state=LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE, + ) + + h.coordinator.fail_timed_out([cancelled, completed_late]) + + assert cancelled.state == LlmRequestState.DISAGG_TRANS_ERROR + assert completed_late.state == LlmRequestState.DISAGG_GENERATION_TRANS_COMPLETE + assert h.effects.failed == [ + ("Request 1 timed out", [cancelled], False), + ("Request 2 timed out", [completed_late], False), + ] + + +def test_fail_timed_out_ends_the_receive_state_before_the_adp_buffer() -> None: + """Under multi-rank ADP the state ends when the request is buffered, not + at the synced drain, because the request has already left the active + list the error vote reads.""" + h = _Harness(enable_attention_dp=True, world_size=2) + req = _Request( + 1, + is_context_only_request=False, + state=LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS, + ) + + h.coordinator.fail_timed_out([req]) + + assert req.state == LlmRequestState.DISAGG_TRANS_ERROR + assert h.effects.failed == [] + + def test_fatal_error_during_the_synced_drain_leaves_no_stale_pending_state() -> None: """If the executor's error path raises (fatal), the drained requests must not be re-failed by the next drain.""" diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_connector_alignment.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_connector_alignment.py new file mode 100644 index 000000000000..7be52f6d2936 --- /dev/null +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_connector_alignment.py @@ -0,0 +1,148 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""How KVCacheManagerV2 shapes reuse matches for a whole-block connector. + +A one-model draft's lookahead backoff trims a local match of whole blocks to a +few tokens short of a boundary, and a connector that serves only from a block +boundary could then offer nothing. The manager asks the core to align the +match only when that combination is actually present, or drops the backoff +when the connector asks for whole local matches instead. +""" + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import KvCacheConnectorWorker +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 + +pytestmark = pytest.mark.cpu_only + +TOKENS_PER_BLOCK = 128 + + +def _manager( + *, backoff: int, partial_reuse: bool, connector, is_draft: bool = False +) -> KVCacheManagerV2: + manager = object.__new__(KVCacheManagerV2) + manager.reuse_match_backoff = backoff + manager.enable_partial_reuse = partial_reuse + manager.is_draft = is_draft + manager._ledger_tokens_per_block = TOKENS_PER_BLOCK + if connector is not None: + manager.kv_connector_manager = connector + return manager + + +def _connector(whole_blocks, *, capacity_only: bool = False) -> SimpleNamespace: + return SimpleNamespace( + capacity_only=capacity_only, + worker=SimpleNamespace(requires_block_aligned_reuse=whole_blocks), + ) + + +def _dropping_connector(drops, *, capacity_only: bool = False) -> SimpleNamespace: + return SimpleNamespace( + capacity_only=capacity_only, + worker=SimpleNamespace(drops_reuse_match_backoff=drops), + ) + + +@pytest.mark.parametrize( + "backoff, partial_reuse, connector, is_draft, expected", + [ + (1, False, _connector(True), False, TOKENS_PER_BLOCK), + # No backoff: matches of whole blocks already end on a boundary. + (0, False, _connector(True), False, 1), + # Partial reuse wants token-granular matches by definition. + (1, True, _connector(True), False, 1), + # A connector that serves from any position needs nothing. + (1, False, _connector(False), False, 1), + (1, False, SimpleNamespace(worker=object()), False, 1), + # Only an explicit declaration counts, not any truthy attribute such + # as the one a MagicMock worker answers with. + (1, False, _connector("yes"), False, 1), + (1, False, None, False, 1), + # Connector doubles that carry no worker at all must not crash __init__. + (1, False, SimpleNamespace(capacity_only=False), False, 1), + # A capacity-only connector serves nothing, so it never reshapes reuse. + (1, False, _connector(True, capacity_only=True), False, 1), + # A paired draft pool aligns too, so the joint claim stays on a boundary. + (1, False, _connector(True), True, TOKENS_PER_BLOCK), + ], + ids=[ + "whole_block_connector_with_backoff", + "no_backoff", + "partial_reuse", + "token_granular_connector", + "connector_without_the_attribute", + "truthy_non_bool_declaration", + "no_connector", + "connector_without_a_worker", + "capacity_only_connector", + "draft_manager", + ], +) +def test_reuse_match_alignment(backoff, partial_reuse, connector, is_draft, expected): + manager = _manager( + backoff=backoff, partial_reuse=partial_reuse, connector=connector, is_draft=is_draft + ) + assert manager._connector_reuse_match_alignment() == expected + + +def test_a_manager_without_a_bound_connector_attribute_keeps_token_granularity(): + """Tests and specialized managers build partial instances that never bind one.""" + manager = object.__new__(KVCacheManagerV2) + manager.reuse_match_backoff = 1 + manager.enable_partial_reuse = False + manager._ledger_tokens_per_block = TOKENS_PER_BLOCK + assert manager._connector_reuse_match_alignment() == 1 + + +@pytest.mark.parametrize( + "is_draft, connector, expected", + [ + (False, _dropping_connector(True), True), + (False, _dropping_connector(False), False), + # A mock answers truthily; only a declared True counts. + (False, _dropping_connector("yes"), False), + (False, SimpleNamespace(worker=object()), False), + (False, None, False), + (False, SimpleNamespace(capacity_only=False), False), + (False, _dropping_connector(True, capacity_only=True), False), + (True, _dropping_connector(True), False), + ], + ids=[ + "asked", + "not_asked", + "truthy_but_not_true", + "no_attribute", + "no_connector", + "connector_without_a_worker", + "capacity_only_connector", + "draft", + ], +) +def test_connector_drops_reuse_match_backoff(is_draft, connector, expected): + manager = object.__new__(KVCacheManagerV2) + manager.is_draft = is_draft + manager.kv_connector_manager = connector + assert manager._connector_drops_reuse_match_backoff() is expected + + +def test_the_worker_base_class_declares_neither(): + """Third-party connectors keep upstream reuse unless they opt in.""" + assert KvCacheConnectorWorker.requires_block_aligned_reuse is False + assert KvCacheConnectorWorker.drops_reuse_match_backoff is False diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_disagg_completion_reservation.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_disagg_completion_reservation.py new file mode 100644 index 000000000000..5897532a87e8 --- /dev/null +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_disagg_completion_reservation.py @@ -0,0 +1,785 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Disaggregated completion reservation in KVCacheManagerV2. + +With no eviction tier and a KV connector whose worker declares +``requires_completion_reservation``, or ``kv_cache_config.enable_disagg_completion_reservation`` +set, a full-attention disaggregated server reserves the pages an admitted request still needs to +finish: a decode request its output growth, a connector prefill its remaining prompt. A new +request is admitted only if the reservation of every admitted request still fits. Without either +opt-in the server keeps its existing admission. + +Tests that build a manager without ``__init__`` assign both reservation dicts whenever they +assign either, so the class-level defaults are never written. +""" + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState +from tensorrt_llm.bindings import DataType, SamplingConfig +from tensorrt_llm.bindings.internal.batch_manager import CacheType +from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.logger import logger +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.runtime.kv_cache_manager_v2 import ( + AttentionLayerConfig, + BufferConfig, + SsmLayerConfig, +) + + +class _ConnectorDouble: + """The KV connector surface the manager touches on the paths these tests take. + + It never offers a prefix and holds no load, so freeing a request only resets state. + ``declares`` is what the worker sets ``requires_completion_reservation`` to. + """ + + def __init__(self, declares: object = True, capacity_only: bool = False) -> None: + self.prefix_reservations_enabled = False + self.capacity_only = capacity_only + self.admission_query_capable = False + self.worker = SimpleNamespace(requires_completion_reservation=declares) + + def release_unstarted_prefix_loads(self, request: LlmRequest) -> None: + pass + + def has_pending_load(self, request: LlmRequest) -> bool: + return False + + def release_prefix_reservation(self, request: LlmRequest) -> None: + pass + + def reset_request_state(self, request: LlmRequest) -> None: + pass + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "available, capacities, limits, candidate_limit, expected", + [ + ([7], [8, 0], {1: 6}, 3, True), + ([6], [8, 0], {1: 6}, 3, False), + ([4], [20, 0], {1: 6}, 3, True), + ([3], [8, 8], {1: 6, 2: 3}, 3, False), + ([5], [8, 8], {1: 6, 2: 3}, 3, True), + ([100, 6], [8, 0], {1: 6}, 3, False), + ], +) +def test_completion_reservation_counts_unallocated_pages_in_every_pool( + available: list[int], + capacities: list[int], + limits: dict[int, int], + candidate_limit: int, + expected: bool, +) -> None: + manager = object.__new__(KVCacheManagerV2) + manager._disagg_completion_limits = dict(limits) + manager._disagg_completion_pool_counts = {i: 1 for i in range(len(available))} + manager.kv_cache_map = { + i + 1: Mock(capacity=capacity, num_blocks=(capacity + 3) // 4) + for i, capacity in enumerate(capacities) + } + manager._get_storage_statistics = Mock( + return_value=[SimpleNamespace(available=count) for count in available] + ) + + assert manager._can_reserve_disagg_completion(2, candidate_limit) is expected + # Checking admission neither consumes a reservation nor touches a cache. + assert manager._disagg_completion_limits == limits + for cache in manager.kv_cache_map.values(): + cache.resize.assert_not_called() + + +@pytest.mark.cpu_only +def test_completion_reservation_skips_a_bound_whose_cache_is_gone() -> None: + manager = object.__new__(KVCacheManagerV2) + manager._disagg_completion_limits = {1: 100, 2: 3} + manager._disagg_completion_pool_counts = {0: 1} + manager.kv_cache_map = {2: Mock(num_blocks=0)} + manager._get_storage_statistics = Mock(return_value=[SimpleNamespace(available=3)]) + + assert manager._can_reserve_disagg_completion(2, 3) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("enable_block_reuse", [False, True]) +@pytest.mark.parametrize("ledger_tokens_per_block, expected_blocks", [(4, 10), (16, 3)]) +def test_generation_limit_counts_pool_variants_and_speculative_padding( + enable_block_reuse: bool, ledger_tokens_per_block: int, expected_blocks: int +) -> None: + manager = object.__new__(KVCacheManagerV2) + manager._ledger_tokens_per_block = ledger_tokens_per_block + manager.max_seq_len = 32 + manager._generation_kv_capacity_headroom = 7 # 1 extra KV + 5 draft + 1 base token + manager.enable_block_reuse = enable_block_reuse + request = Mock(max_new_tokens=20) + + # Clamp to the model sequence length, then add the maximum generation lead: + # 32 + 7 = 39 tokens. + limit = manager._disagg_generation_limit(request, 24) + assert limit == expected_blocks + int(enable_block_reuse) + + manager._disagg_completion_limits = {} + manager._disagg_completion_pool_counts = {0: 2} + manager.kv_cache_map = {1: Mock(capacity=0, num_blocks=0)} + stats = SimpleNamespace(available=2 * limit - 1) + manager._get_storage_statistics = Mock(return_value=[stats]) + assert not manager._can_reserve_disagg_completion(1, limit) + stats.available = 2 * limit + assert manager._can_reserve_disagg_completion(1, limit) + + +def _context_reservation_manager(available: int, **connector_overrides) -> KVCacheManagerV2: + manager = object.__new__(KVCacheManagerV2) + manager._ledger_tokens_per_block = 4 + manager._generation_kv_capacity_headroom = 1 + manager.enable_block_reuse = False + manager.is_draft = False + manager._disagg_completion_limits = {} + manager._disagg_completion_pool_counts = {0: 1} + manager.kv_cache_map = {7: Mock(num_blocks=0)} + manager._get_storage_statistics = Mock(return_value=[SimpleNamespace(available=available)]) + connector = dict( + prefix_reservations_enabled=False, + capacity_only=False, + worker=SimpleNamespace(requires_completion_reservation=True), + ) + connector.update(connector_overrides) + manager.kv_connector_manager = SimpleNamespace(**connector) + return manager + + +def _first_chunk(**overrides) -> SimpleNamespace: + fields = dict( + py_request_id=7, + prompt_len=16, + is_first_context_chunk=True, + is_dummy=False, + is_generation_only_request=False, + is_disagg_generation_init_state=False, + py_connector_allocation_reported=False, + ) + fields.update(overrides) + return SimpleNamespace(**fields) + + +@pytest.mark.cpu_only +def test_context_reservation_records_a_bound_once_it_fits() -> None: + # 16 prompt tokens + 1 base token in 4-token blocks: 5 blocks. + manager = _context_reservation_manager(available=5) + request = _first_chunk() + + assert manager.reserve_context_completion(request) + assert manager._disagg_completion_limits == {7: 5} + # A recorded bound is not checked again on later passes. + manager._get_storage_statistics.return_value = [SimpleNamespace(available=0)] + assert manager.reserve_context_completion(request) + + +@pytest.mark.cpu_only +def test_context_reservation_defers_without_recording() -> None: + manager = _context_reservation_manager(available=4) + + assert not manager.reserve_context_completion(_first_chunk()) + assert manager._disagg_completion_limits == {} + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "request_overrides, connector_overrides, has_pool_counts", + [ + (dict(is_first_context_chunk=False), {}, True), + ({}, {}, False), + ({}, dict(prefix_reservations_enabled=True), True), + ({}, dict(capacity_only=True), True), + (dict(is_dummy=True), {}, True), + (dict(is_generation_only_request=True), {}, True), + (dict(is_disagg_generation_init_state=True), {}, True), + ], + ids=[ + "later_chunk", + "eviction_tier", + "source_reservations", + "capacity_only", + "dummy", + "generation_only", + "gen_init", + ], +) +def test_context_reservation_is_a_noop_outside_conservative_connector_prefill( + request_overrides: dict, connector_overrides: dict, has_pool_counts: bool +) -> None: + manager = _context_reservation_manager(available=0, **connector_overrides) + if not has_pool_counts: + manager._disagg_completion_pool_counts = {} + + assert manager.reserve_context_completion(_first_chunk(**request_overrides)) + assert manager._disagg_completion_limits == {} + manager._get_storage_statistics.assert_not_called() + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "request_overrides, released", + [ + ({}, True), + (dict(py_connector_allocation_reported=True), False), + (dict(is_first_context_chunk=False), False), + # A reverted disaggregated generation-init admission, before its transfer. + (dict(is_disagg_generation_init_state=True), True), + ], + ids=["unconsumed", "reported", "later_chunk", "gen_init"], +) +def test_unconsumed_context_reservation_is_returned( + request_overrides: dict, released: bool +) -> None: + manager = object.__new__(KVCacheManagerV2) + manager._disagg_completion_limits = {7: 5} + manager._disagg_completion_pool_counts = {0: 1} + + manager._release_unconsumed_context_reservation(_first_chunk(**request_overrides)) + + assert (7 not in manager._disagg_completion_limits) is released + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("first_chunk", [True, False]) +def test_failed_first_chunk_resize_releases_its_bound(first_chunk: bool) -> None: + """A first chunk that cannot grow is suspended and returns its bound with it.""" + manager = object.__new__(KVCacheManagerV2) + manager._disagg_completion_limits = {7: 5} + manager._disagg_completion_pool_counts = {0: 1} + kv_cache = Mock(capacity=0, resize=Mock(return_value=False)) + manager.kv_cache_map = {7: kv_cache} + manager._has_cp_helix = False + manager.is_disagg = True + manager.num_extra_kv_tokens = 0 + manager.is_draft = False + manager.kv_connector_manager = SimpleNamespace( + prefix_reservations_enabled=False, + capacity_only=False, + worker=SimpleNamespace(requires_completion_reservation=True), + ) + request = SimpleNamespace( + py_request_id=7, + is_disagg_generation_init_state=False, + is_dummy_request=False, + context_current_position=0, + is_first_context_chunk=first_chunk, + py_connector_allocation_reported=False, + ) + + assert not KVCacheManagerV2.resize_context(manager, request, 4) + + kv_cache.resize.assert_called_once_with(4) + if first_chunk: + kv_cache.suspend.assert_called_once_with() + assert 7 not in manager._disagg_completion_limits + else: + kv_cache.suspend.assert_not_called() + assert manager._disagg_completion_limits == {7: 5} + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("reported", [False, True]) +def test_revert_without_capacity_growth_returns_an_unconsumed_bound(reported: bool) -> None: + """The release comes before every early return of ``revert_allocate_context``. + + A first chunk whose resize grew no capacity keeps its cache, but an unconsumed bound must + still go back so that the next admission checks it again. + """ + manager = object.__new__(KVCacheManagerV2) + manager._disagg_completion_limits = {7: 5} + manager._disagg_completion_pool_counts = {0: 1} + manager.is_draft = False + manager.kv_connector_manager = _ConnectorDouble() + request = _first_chunk(py_connector_allocation_reported=reported, py_ctx_pre_resize_cap=None) + + assert manager.revert_allocate_context(request) + + assert (7 in manager._disagg_completion_limits) is reported + + +def _attention_layer(sliding_window_size: int | None) -> AttentionLayerConfig: + return AttentionLayerConfig( + layer_id=0, + buffers=[BufferConfig(role="key", size=256)], + sliding_window_size=sliding_window_size, + ) + + +def _gate_manager(**overrides) -> KVCacheManagerV2: + manager = object.__new__(KVCacheManagerV2) + fields = dict( + is_disagg=True, + can_evict=False, + is_draft=False, + max_beam_width=1, + mapping=SimpleNamespace(pp_size=1), + max_seq_len=64, + kv_cache_manager_py_config=SimpleNamespace(layers=[_attention_layer(None)]), + kv_connector_manager=_ConnectorDouble(), + enable_disagg_completion_reservation=False, + ) + fields.update(overrides) + for name, value in fields.items(): + setattr(manager, name, value) + return manager + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + "overrides, applies", + [ + ({}, True), + # A window that covers max_seq_len behaves as full attention. + ( + dict(kv_cache_manager_py_config=SimpleNamespace(layers=[_attention_layer(64)])), + True, + ), + (dict(is_disagg=False), False), + (dict(can_evict=True), False), + # With the configuration off, only a serving connector that declares it turns it on. + (dict(kv_connector_manager=None), False), + (dict(kv_connector_manager=_ConnectorDouble(declares=False)), False), + # Only an explicit True counts, not a truthy value such as a mock's attribute. + (dict(kv_connector_manager=_ConnectorDouble(declares="yes")), False), + (dict(kv_connector_manager=_ConnectorDouble(capacity_only=True)), False), + (dict(is_draft=True), False), + (dict(max_beam_width=2), False), + (dict(mapping=SimpleNamespace(pp_size=2)), False), + ( + dict(kv_cache_manager_py_config=SimpleNamespace(layers=[_attention_layer(16)])), + False, + ), + ( + dict( + kv_cache_manager_py_config=SimpleNamespace( + layers=[ + _attention_layer(None), + SsmLayerConfig( + layer_id=1, buffers=[BufferConfig(role="ssm_state", size=256)] + ), + ] + ) + ), + False, + ), + # The configuration opts in without a declaring connector, under the same conditions. + (dict(kv_connector_manager=None, enable_disagg_completion_reservation=True), True), + ( + dict( + kv_connector_manager=_ConnectorDouble(declares=False), + enable_disagg_completion_reservation=True, + ), + True, + ), + ( + dict( + kv_connector_manager=None, + enable_disagg_completion_reservation=True, + is_disagg=False, + ), + False, + ), + ( + dict( + kv_connector_manager=None, + enable_disagg_completion_reservation=True, + can_evict=True, + ), + False, + ), + ( + dict( + kv_connector_manager=None, + enable_disagg_completion_reservation=True, + mapping=SimpleNamespace(pp_size=2), + ), + False, + ), + ( + dict( + kv_connector_manager=None, + enable_disagg_completion_reservation=True, + kv_cache_manager_py_config=SimpleNamespace(layers=[_attention_layer(16)]), + ), + False, + ), + ], + ids=[ + "full_attention", + "window_covers_max_seq_len", + "not_disagg", + "eviction_tier", + "no_connector", + "undeclared_connector", + "truthy_non_bool_declaration", + "capacity_only_connector", + "draft", + "beam_search", + "pipeline_parallel", + "sliding_window", + "recurrent_layer", + "config_without_connector", + "config_with_undeclared_connector", + "config_not_disagg", + "config_eviction_tier", + "config_pipeline_parallel", + "config_sliding_window", + ], +) +def test_completion_reservation_gate(overrides: dict, applies: bool) -> None: + assert _gate_manager(**overrides)._disagg_completion_reservation_applies() is applies + + +def _real_manager( + kv_cache_config: KvCacheConfig | None = None, **overrides: object +) -> KVCacheManagerV2: + """A one-layer manager whose tiny GPU pool makes page pressure reachable. + + Its connector declares the reservation unless ``kv_connector_manager`` is overridden. + """ + if kv_cache_config is None: + kv_cache_config = KvCacheConfig( + enable_block_reuse=False, + host_cache_size=0, + max_gpu_total_bytes=128 << 10, + max_util_for_resume=1.0, + ) + kwargs = dict( + num_layers=1, + num_kv_heads=8, + head_dim=64, + tokens_per_block=4, + max_seq_len=65536, + max_batch_size=4, + max_num_tokens=4, + mapping=Mapping(world_size=1, rank=0, tp_size=1, pp_size=1), + dtype=DataType.HALF, + vocab_size=4096, + is_disagg=True, + kv_connector_manager=_ConnectorDouble(), + ) + kwargs.update(overrides) + torch.cuda.init() + return KVCacheManagerV2(kv_cache_config, CacheType.SELF, **kwargs) + + +def _gen_init_request(request_id: int, prompt: int, output: int) -> LlmRequest: + request = LlmRequest( + request_id=request_id, + max_new_tokens=output, + input_tokens=[request_id] * prompt, + sampling_config=SamplingConfig(1), + is_streaming=False, + ) + request.state = LlmRequestState.DISAGG_GENERATION_INIT + return request + + +@pytest.mark.parametrize("reserve_growth", [False, True]) +@pytest.mark.parametrize("enable_block_reuse", [False, True]) +def test_disagg_generation_progress_under_real_gpu_page_pressure( + reserve_growth: bool, enable_block_reuse: bool +) -> None: + """Prompt-only admission can fill the pool before either request finishes. + + Use the real allocator and admission path; the control disables only the reservation. + The guarded path must defer the second long prompt, admit a short request alongside the + first, and release reservations on free. + """ + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + manager = _real_manager( + KvCacheConfig( + enable_block_reuse=enable_block_reuse, + host_cache_size=0, + max_gpu_total_bytes=128 << 10, + max_util_for_resume=1.0, + ) + ) + try: + assert manager._disagg_completion_pool_counts == {0: 1} + if not reserve_growth: + manager._disagg_completion_pool_counts.clear() + total_blocks = manager.get_kv_cache_stats().max_num_blocks + prompt_len = total_blocks // 3 * manager.tokens_per_block + assert prompt_len >= 8 + # Allocation granularity can make the physical pool larger than the tiny + # requested byte quota. Keep this pressure probe within the actual table. + assert 2 * prompt_len < manager.max_seq_len + + first = _gen_init_request(1, prompt_len, prompt_len) + second = _gen_init_request(2, prompt_len, prompt_len) + short = _gen_init_request(3, 4, 1) + + assert manager.prepare_disagg_gen_init(first) + assert manager.kv_cache_map[1].capacity == prompt_len + assert manager.prepare_disagg_gen_init(second) is not reserve_growth + first.state = LlmRequestState.GENERATION_IN_PROGRESS + if reserve_growth: + assert 2 not in manager.kv_cache_map + assert 2 not in manager._disagg_completion_limits + assert manager.prepare_disagg_gen_init(short) + for _ in range(prompt_len): + assert manager.try_allocate_generation(first) + manager.free_resources(first) + assert 1 not in manager._disagg_completion_limits + assert manager.prepare_disagg_gen_init(second) + # Cancellation must return both the physical pages and the logical bound. + manager.free_resources(second) + manager.free_resources(short) + assert not manager._disagg_completion_limits + assert manager.get_kv_cache_stats().free_num_blocks == total_blocks + else: + second.state = LlmRequestState.GENERATION_IN_PROGRESS + for _ in range(prompt_len): + first_fits = manager.try_allocate_generation(first) + second_fits = manager.try_allocate_generation(second) + if not first_fits and not second_fits: + break + else: + pytest.fail("The unguarded control did not exhaust its GPU KV pool") + finally: + manager.shutdown() + + +_GATE_DEFAULT_CONFIG = dict( + enable_block_reuse=False, + host_cache_size=0, + max_gpu_total_bytes=128 << 10, + max_util_for_resume=1.0, +) + + +@pytest.mark.parametrize( + "config_overrides, manager_overrides, expected", + [ + ({}, {}, {0: 1}), + ({}, dict(is_disagg=False), {}), + # An explicit host tier; the automatic one depends on the host's memory. + (dict(host_cache_size=16 << 20), {}, {}), + # Without either opt-in, a server with no tier below GPU keeps main's admission. + ({}, dict(kv_connector_manager=None), {}), + ({}, dict(kv_connector_manager=_ConnectorDouble(declares=False)), {}), + ({}, dict(kv_connector_manager=_ConnectorDouble(capacity_only=True)), {}), + ({}, dict(max_beam_width=2), {}), + ({}, dict(is_draft=True), {}), + (dict(max_attention_window=[16]), {}, {}), + # The configuration opts a server with no connector in, under the same conditions. + (dict(enable_disagg_completion_reservation=True), dict(kv_connector_manager=None), {0: 1}), + ( + dict(enable_disagg_completion_reservation=True), + dict(kv_connector_manager=None, is_disagg=False), + {}, + ), + ( + dict(enable_disagg_completion_reservation=True, host_cache_size=16 << 20), + dict(kv_connector_manager=None), + {}, + ), + ( + dict(enable_disagg_completion_reservation=True), + dict(kv_connector_manager=None, is_draft=True), + {}, + ), + ( + dict(enable_disagg_completion_reservation=True, host_cache_size=16 << 20), + dict(kv_connector_manager=None, is_estimating_kv_cache=True), + {}, + ), + ], + ids=[ + "baseline", + "not_disagg", + "host_tier", + "no_connector", + "undeclared_connector", + "capacity_only_connector", + "beam_search", + "draft", + "sliding_window", + "config_without_connector", + "config_not_disagg", + "config_host_tier", + "config_draft", + "config_estimating", + ], +) +def test_completion_reservation_gate_on_a_real_manager( + config_overrides: dict, + manager_overrides: dict, + expected: dict, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The constructor turns the reservation on only where the gate allows it. + + A target manager whose configuration asks for the reservation where it cannot apply says so; + a draft manager, which never reserves, stays quiet, and so does the estimation manager, which + is built without the connector. Pipeline parallelism needs one process per rank; + ``test_completion_reservation_gate`` covers that clause. + """ + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + warned: list[str] = [] + monkeypatch.setattr(logger, "warning", lambda *msg: warned.append(" ".join(map(str, msg)))) + config = dict(_GATE_DEFAULT_CONFIG) + config.update(config_overrides) + manager = _real_manager(KvCacheConfig(**config), **manager_overrides) + try: + assert manager._disagg_completion_pool_counts == expected + assert manager._disagg_completion_limits == {} + ignored = ( + config.get("enable_disagg_completion_reservation", False) + and not expected + and not manager_overrides.get("is_draft", False) + and not manager_overrides.get("is_estimating_kv_cache", False) + ) + assert any("enable_disagg_completion_reservation" in msg for msg in warned) is ignored + finally: + manager.shutdown() + + +@pytest.mark.parametrize( + "connector", + [None, _ConnectorDouble(declares=False)], + ids=["no_connector", "undeclared_connector"], +) +def test_admission_is_unchanged_without_the_opt_in(connector: _ConnectorDouble | None) -> None: + """A disaggregated server with no tier below GPU keeps main's admission unless asked. + + Two generation-init requests whose completion bounds together exceed the pool are both + admitted, and a first context chunk is neither deferred nor bounded. + """ + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + manager = _real_manager(kv_connector_manager=connector) + try: + assert not manager.can_evict + assert manager._disagg_completion_pool_counts == {} + total_blocks = manager.get_kv_cache_stats().max_num_blocks + prompt_len = total_blocks // 3 * manager.tokens_per_block + assert prompt_len >= 8 + first = _gen_init_request(1, prompt_len, prompt_len) + second = _gen_init_request(2, prompt_len, prompt_len) + # Each bound covers twice the prompt, so the two cannot both fit the pool. + assert 2 * manager._disagg_generation_limit(first, prompt_len) > total_blocks + + assert manager.prepare_disagg_gen_init(first) + assert manager.prepare_disagg_gen_init(second) + assert manager._disagg_completion_limits == {} + + context = LlmRequest( + request_id=3, + max_new_tokens=1, + input_tokens=[3] * prompt_len, + sampling_config=SamplingConfig(1), + is_streaming=False, + ) + assert manager.reserve_context_completion(context) + assert manager._disagg_completion_limits == {} + + manager.free_resources(first) + manager.free_resources(second) + assert manager.get_kv_cache_stats().free_num_blocks == total_blocks + finally: + manager.shutdown() + + +def test_config_opt_in_reserves_decode_headroom_without_a_connector() -> None: + """A generation server with no connector reserves decode headroom when configured to. + + The same pair of generation-init requests that main admits together is split: the second + waits until the first frees its bound. A context request is never bounded, since only a + connector prefill reserves its prompt. + """ + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + manager = _real_manager( + KvCacheConfig(**_GATE_DEFAULT_CONFIG, enable_disagg_completion_reservation=True), + kv_connector_manager=None, + ) + try: + assert manager._disagg_completion_pool_counts == {0: 1} + total_blocks = manager.get_kv_cache_stats().max_num_blocks + prompt_len = total_blocks // 3 * manager.tokens_per_block + assert prompt_len >= 8 + first = _gen_init_request(1, prompt_len, prompt_len) + second = _gen_init_request(2, prompt_len, prompt_len) + + assert manager.prepare_disagg_gen_init(first) + assert 1 in manager._disagg_completion_limits + assert not manager.prepare_disagg_gen_init(second) + assert 2 not in manager.kv_cache_map + assert 2 not in manager._disagg_completion_limits + + context = LlmRequest( + request_id=3, + max_new_tokens=1, + input_tokens=[3] * prompt_len, + sampling_config=SamplingConfig(1), + is_streaming=False, + ) + assert manager.reserve_context_completion(context) + assert 3 not in manager._disagg_completion_limits + + manager.free_resources(first) + assert manager.prepare_disagg_gen_init(second) + manager.free_resources(second) + assert not manager._disagg_completion_limits + assert manager.get_kv_cache_stats().free_num_blocks == total_blocks + finally: + manager.shutdown() + + +def test_reverted_gen_init_releases_its_completion_bound() -> None: + """A generation-init admission that the coordinator reverts leaves no bound behind. + + This mirrors DisaggCoordinator.revert_deferred_gen_init, which reaches + revert_allocate_context through PyExecutor._revert_ctx_alloc. A fresh gen-init holds + history past its pre-resize capacity of zero, so revert frees its cache outright; the + test pins that no bound leaks and that the request is admitted again afterwards. + """ + if not torch.cuda.is_available(): + pytest.skip("requires CUDA") + manager = _real_manager() + try: + assert manager._disagg_completion_pool_counts == {0: 1} + total_blocks = manager.get_kv_cache_stats().max_num_blocks + prompt_len = total_blocks // 3 * manager.tokens_per_block + assert prompt_len >= 8 + request = _gen_init_request(1, prompt_len, prompt_len) + + assert manager.prepare_disagg_gen_init(request) + assert 1 in manager._disagg_completion_limits + + assert not manager.revert_allocate_context(request) + assert 1 not in manager._disagg_completion_limits + assert 1 not in manager.kv_cache_map + assert manager.get_kv_cache_stats().free_num_blocks == total_blocks + + assert manager.prepare_disagg_gen_init(request) + assert 1 in manager._disagg_completion_limits + manager.free_resources(request) + assert not manager._disagg_completion_limits + finally: + manager.shutdown() diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py index 951be07f199b..18d11ae2c439 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py @@ -392,6 +392,9 @@ def make_ctx_request( req.encoder_output_len = encoder_output_len req.py_encoder_output_ready_event = None req.py_skip_cross_kv_projection = False + # Real values: the admission-time connector path compares and clears them. + req.py_connector_offer_deferrals = 0 + req.py_preempted_for_pages = False return req @@ -682,6 +685,8 @@ def preempt_request(req): mgr.preempt_request.side_effect = preempt_request # A bare Mock would answer truthily and suppress every stall check. mgr.has_pending_preemption.return_value = False + # Completion headroom is plentiful unless a test says otherwise. + mgr.reserve_context_completion.return_value = True return mgr @@ -724,6 +729,9 @@ def make_scheduler( cross_kv_cache_manager: Mock | None = None, enable_prefix_aware_scheduling: bool = True, enable_recompute_pause: bool = True, + connector_admission_query: bool = False, + max_uncharged_connector_parks: int = 0, + max_connector_offer_deferrals: int = 16, ) -> object: """Create KVCacheV2Scheduler, patching isinstance check for mock mgr.""" from tensorrt_llm._torch.pyexecutor.scheduler.scheduler_v2 import KVCacheV2Scheduler @@ -751,6 +759,9 @@ def make_scheduler( scheduler_capacity=scheduler_capacity, enable_prefix_aware_scheduling=enable_prefix_aware_scheduling, enable_recompute_pause=enable_recompute_pause, + connector_admission_query=connector_admission_query, + max_uncharged_connector_parks=max_uncharged_connector_parks, + max_connector_offer_deferrals=max_connector_offer_deferrals, **kwargs, ) @@ -2333,6 +2344,75 @@ def test_a_deferred_preemption_is_not_a_stall(self): sched.schedule_request(reqs, set()) +@pytest.mark.parametrize("chunked", [False, True]) +def test_completion_reservation_deferral_preserves_started_context(chunked: bool) -> None: + """A new prompt waiting for completion headroom must let admitted prefill finish.""" + manager = make_kv_cache_manager(has_cache_tier_below_gpu=False) + # Only the new prompt is refused, whatever order the pass consults them in. + manager.reserve_context_completion.side_effect = lambda req: req.py_request_id != 0 + scheduler = make_scheduler( + manager, + max_num_tokens=1000, + ctx_chunk_config=(ContextChunkingPolicy.FIRST_COME_FIRST_SERVED, 64) if chunked else None, + ) + waiting = make_ctx_request(0, 500) + started = make_ctx_request(99, 100, is_first_context_chunk=False) + + batch = scheduler.schedule_request([waiting, started], set()) + + assert ids(batch.context_requests) == [99] + assert batch.paused_requests == [] + manager.preempt_request.assert_not_called() + manager.resize_context.assert_called_once_with(started, 100) + started.pause.assert_not_called() + + # Once headroom returns, admission accepts the queued prompt. + manager.reserve_context_completion.side_effect = None + manager.reserve_context_completion.return_value = True + resumed = scheduler.schedule_request([waiting], set()) + assert ids(resumed.context_requests) == [0] + manager.preempt_request.assert_not_called() + + +@pytest.mark.parametrize( + "chunked, max_num_tokens, align_to_nothing, reaches_prepare", + [ + # Full path: the prompt does not fit the token budget once prepared. + (False, 100, False, True), + # Chunked path: the pass has no chunk budget left at all. + (True, 32, False, False), + # Chunked path: no multimodal-safe chunk boundary this pass. + (True, 1000, True, True), + ], + ids=["full_token_budget", "no_chunk_budget", "mm_alignment"], +) +def test_completion_reservation_not_consulted_before_budget_skips( + chunked: bool, max_num_tokens: int, align_to_nothing: bool, reaches_prepare: bool +) -> None: + """Budget and alignment refusals come first, so they never record a completion bound.""" + manager = make_kv_cache_manager(has_cache_tier_below_gpu=False) + scheduler = make_scheduler( + manager, + max_num_tokens=max_num_tokens, + ctx_chunk_config=(ContextChunkingPolicy.FIRST_COME_FIRST_SERVED, 64) if chunked else None, + ) + if align_to_nothing: + scheduler._align_chunk_to_mm_block = Mock(return_value=0) + request = make_ctx_request(0, 200) + + batch = scheduler.schedule_request([request], set()) + + assert batch.context_requests == [] + if reaches_prepare: + manager.prepare_context.assert_called_once_with(request) + else: + manager.prepare_context.assert_not_called() + if align_to_nothing: + scheduler._align_chunk_to_mm_block.assert_called_once() + manager.reserve_context_completion.assert_not_called() + manager.resize_context.assert_not_called() + + class TestDeadlockDetection: """The scheduler must fail loudly rather than spin scheduling nothing. @@ -3456,6 +3536,20 @@ def test_force_chunk_first_chunk_waits_on_leftover_budget(self): assert mgr.prepare_context.call_args_list == [call(started)] mgr.free_resources.assert_not_called() + def test_force_chunk_wait_counts_as_a_budget_skip(self): + """Both first chunks waiting on the 20-token leftover were limited by budget.""" + mgr = make_kv_cache_manager(tokens_per_block=32) + sched = make_scheduler( + mgr, max_num_tokens=150, ctx_chunk_config=(ContextChunkingPolicy.FORCE_CHUNK, 32) + ) + started = self._started_ctx_ending_at_snapshot(0, 1000, 1130) + waiting = [make_ctx_request(1, 4000), make_ctx_request(2, 4000)] + + out = sched.schedule_request([started, *waiting], set()) + + assert ids(out.context_requests) == [0] + assert sched.last_pass_ctx_budget_skips == 2 + def test_force_chunk_first_chunk_that_fits_leftover_still_runs(self): mgr = make_kv_cache_manager(tokens_per_block=32) sched = make_scheduler( @@ -5306,3 +5400,337 @@ def test_preemption_keeps_arrival_order(self, long_arrived_first): assert ids(out.context_requests) == [long_id] else: assert ids(out.recompute_paused_requests) == [long_id] + + +# =========================================================================== +# Admission-time connector queries +# =========================================================================== + + +class TestConnectorAdmission: + """kv_connector_config.admission_query: the connector is asked while a first chunk is admitted. + + A chunk the connector will park for an asynchronous load is admitted with no + token charge (up to max_uncharged_connector_parks per pass), and a chunk + whose offer needs pages the KV cache cannot give waits for a later pass. + """ + + @staticmethod + def _parking_manager(parked_ids, calls=None): + """A manager whose connector parks exactly ``parked_ids`` when asked at admission.""" + mgr = make_kv_cache_manager(tokens_per_block=64) + + def query(req): + if calls is not None: + calls.append(req.py_request_id) + return req.py_request_id in parked_ids + + mgr.query_connector_at_admission.side_effect = query + mgr.connector_admission_query_enabled.return_value = True + mgr.cover_admission_offer.return_value = True + return mgr + + @staticmethod + def _scheduler(mgr, parks=2, deferrals=16, max_batch_size=100, admission=True): + return make_scheduler( + mgr, + max_batch_size=max_batch_size, + max_num_tokens=8192, + ctx_chunk_config=(None, 64), + connector_admission_query=admission, + max_uncharged_connector_parks=parks, + max_connector_offer_deferrals=deferrals, + ) + + @staticmethod + def _parked_head_and_two_more(): + return ( + make_ctx_request(0, context_remaining_length=20000), + make_ctx_request(1, context_remaining_length=6000), + make_ctx_request(2, context_remaining_length=1000), + ) + + def test_off_by_default(self): + mgr = self._parking_manager({0}) + sched = self._scheduler(mgr, admission=False) + head = make_ctx_request(0, context_remaining_length=20000) + out = sched.schedule_request([head, make_ctx_request(1, 1000)], set()) + assert ids(out.context_requests) == [0] + assert head.context_chunk_size == 8192 + mgr.query_connector_at_admission.assert_not_called() + mgr.begin_connector_admission_pass.assert_not_called() + assert sched.uncharged_park_ids == frozenset() + + def test_parked_head_is_uncharged_and_the_budget_goes_to_the_rest(self): + sched = self._scheduler(self._parking_manager({0})) + head, long_second, short = self._parked_head_and_two_more() + out = sched.schedule_request([head, long_second, short], set()) + assert sorted(ids(out.context_requests)) == [0, 1, 2] + # The parked head keeps its sized chunk and pages; only the charge is waived. + assert head.context_chunk_size == 8192 + assert sched.uncharged_park_ids == frozenset({0}) + assert long_second.context_chunk_size == 6000 + assert short.context_chunk_size == 1000 + + def test_without_the_knob_the_park_burns_its_chunk(self): + sched = self._scheduler(self._parking_manager({0}), admission=False) + head, long_second, short = self._parked_head_and_two_more() + out = sched.schedule_request([head, long_second, short], set()) + assert ids(out.context_requests) == [0] + assert long_second.context_chunk_size == 0 + assert short.context_chunk_size == 0 + # Both were passed over for budget, which the park diagnostics report. + assert sched.last_pass_ctx_budget_skips == 2 + + def test_a_follower_gets_the_whole_budget(self): + sched = self._scheduler(self._parking_manager({0}), parks=1) + head = make_ctx_request(0, context_remaining_length=20000) + follower = make_ctx_request(1, context_remaining_length=20000) + out = sched.schedule_request([head, follower], set()) + assert sorted(ids(out.context_requests)) == [0, 1] + assert head.context_chunk_size == 8192 + assert follower.context_chunk_size == 8192 + assert sched.uncharged_park_ids == frozenset({0}) + + def test_cap_limits_the_uncharged_parks_per_pass(self): + calls = [] + sched = self._scheduler(self._parking_manager({0, 1, 2}, calls), parks=1) + reqs = [make_ctx_request(i, context_remaining_length=6000) for i in range(3)] + sched.schedule_request(reqs, set()) + # Every first chunk is asked, so a pool hit past the cap also gets its pages booked. + assert calls == [0, 1, 2] + assert sched.uncharged_park_ids == frozenset({0}) + # Past the cap a parking request is charged: 6000, then the 2192 left + # rounded down to the 64-token chunk unit. + assert [r.context_chunk_size for r in reqs] == [6000, 6000, 2176] + # The cap is per pass: the head is admitted uncharged again. + sched.schedule_request(reqs, set()) + assert sched.uncharged_park_ids == frozenset({0}) + + def test_a_miss_does_not_use_up_the_cap(self): + """The cap counts parks, not questions: a miss leaves the slot to the next request.""" + calls = [] + sched = self._scheduler(self._parking_manager({1}, calls), parks=1) + sched.schedule_request( + [make_ctx_request(0, context_remaining_length=4000), make_ctx_request(1, 20000)], set() + ) + assert calls == [0, 1] + assert sched.uncharged_park_ids == frozenset({1}) + + def test_parked_request_still_takes_a_request_slot(self): + sched = self._scheduler(self._parking_manager({0}), max_batch_size=1) + out = sched.schedule_request( + [make_ctx_request(0, context_remaining_length=20000), make_ctx_request(1, 1000)], set() + ) + assert ids(out.context_requests) == [0] + assert sched.uncharged_park_ids == frozenset({0}) + + def test_a_request_that_fails_allocation_is_never_asked(self): + calls = [] + mgr = self._parking_manager({0, 1}, calls) + mgr.resize_context.side_effect = lambda req, n: req.py_request_id != 0 + sched = self._scheduler(mgr) + sched.schedule_request( + [make_ctx_request(0, context_remaining_length=20000), make_ctx_request(1, 1000)], set() + ) + assert calls == [1] + + def test_unchunked_scheduling_never_asks(self): + mgr = self._parking_manager({0}) + sched = make_scheduler(mgr, max_num_tokens=8192, connector_admission_query=True) + sched.schedule_request([make_ctx_request(0, context_remaining_length=4000)], set()) + mgr.query_connector_at_admission.assert_not_called() + mgr.begin_connector_admission_pass.assert_not_called() + + def test_a_small_chunk_keeps_its_charge(self): + """Under two blocks a park frees too little budget to use up an uncharged slot.""" + calls = [] + sched = self._scheduler(self._parking_manager({0}, calls)) + out = sched.schedule_request([make_ctx_request(0, context_remaining_length=100)], set()) + assert ids(out.context_requests) == [0] + assert calls == [0] + assert sched.uncharged_park_ids == frozenset() + + def test_uncharged_ids_are_per_pass(self): + parked = {0} + mgr = self._parking_manager(parked) + sched = self._scheduler(mgr) + head = make_ctx_request(0, context_remaining_length=20000) + sched.schedule_request([head], set()) + assert sched.uncharged_park_ids == frozenset({0}) + mgr.begin_connector_admission_pass.assert_called_once() + parked.clear() + sched.schedule_request([head], set()) + assert sched.uncharged_park_ids == frozenset() + assert mgr.begin_connector_admission_pass.call_count == 2 + + def test_an_offer_without_pages_gives_its_chunk_back(self): + """Waiting a pass beats loading only the blocks the chunk's pages hold and recomputing the rest.""" + mgr = self._parking_manager({0}) + mgr.cover_admission_offer.side_effect = lambda req: req.py_request_id != 0 + sched = self._scheduler(mgr) + head = make_ctx_request(0, context_remaining_length=20000) + other = make_ctx_request(1, context_remaining_length=6000) + with patch.object( + sched, "release_context_admission", wraps=sched.release_context_admission + ) as release: + out = sched.schedule_request([head, other], set()) + assert ids(out.context_requests) == [1] + assert other.context_chunk_size == 6000 + assert head.py_connector_offer_deferrals == 1 + assert head.py_connector_admission_offer is None + # The first-chunk release frees the pools and rewinds; nothing is shrunk first. + release.assert_called_once_with(head) + mgr.revert_allocate_context.assert_not_called() + assert sched.uncharged_park_ids == frozenset() + + def test_a_charged_pool_hit_past_the_cap_also_waits_for_pages(self): + mgr = self._parking_manager({0, 1}) + mgr.cover_admission_offer.side_effect = lambda req: req.py_request_id != 1 + sched = self._scheduler(mgr, parks=1) + reqs = [make_ctx_request(i, context_remaining_length=3000) for i in range(3)] + out = sched.schedule_request(reqs, set()) + assert sorted(ids(out.context_requests)) == [0, 2] + assert sched.uncharged_park_ids == frozenset({0}) + assert reqs[1].py_connector_offer_deferrals == 1 + + def test_the_wait_is_bounded(self): + mgr = self._parking_manager({0}) + mgr.cover_admission_offer.return_value = False + sched = self._scheduler(mgr, deferrals=2) + head = make_ctx_request(0, context_remaining_length=20000) + with patch.object( + sched, "release_context_admission", wraps=sched.release_context_admission + ) as release: + scheduled = [ + ids(sched.schedule_request([head], set()).context_requests) for _ in range(4) + ] + # Two passes waiting, then admitted (uncharged, since it parks) to take the capped prefix. + assert scheduled == [[], [], [0], [0]] + assert sched.uncharged_park_ids == frozenset({0}) + assert release.call_count == 2 + + def test_a_wait_for_offer_pages_is_not_a_stall(self): + mgr = self._parking_manager({0}) + mgr.cover_admission_offer.return_value = False + sched = self._scheduler(mgr) + head = make_ctx_request(0, context_remaining_length=20000) + for _ in range(5): + assert ids(sched.schedule_request([head], set()).context_requests) == [] + assert sched._stalled_schedules == 0 + assert head.py_connector_offer_deferrals == 5 + + def test_a_request_that_preempted_for_pages_keeps_them(self): + """Waiting would hand the pages a re-prefill paid for to the requests behind it.""" + mgr = self._parking_manager({0}) + mgr.cover_admission_offer.return_value = False + sched = self._scheduler(mgr) + head = make_ctx_request(0, context_remaining_length=20000) + head.py_preempted_for_pages = True + assert ids(sched.schedule_request([head], set()).context_requests) == [0] + assert head.py_connector_offer_deferrals == 0 + assert head.py_preempted_for_pages is False + # The flag is spent: the next admission waits for pages again. + assert ids(sched.schedule_request([head], set()).context_requests) == [] + assert head.py_connector_offer_deferrals == 1 + + def test_no_wait_without_the_knob(self): + mgr = self._parking_manager({0}) + mgr.cover_admission_offer.return_value = False + sched = self._scheduler(mgr, admission=False) + out = sched.schedule_request([make_ctx_request(0, context_remaining_length=20000)], set()) + assert ids(out.context_requests) == [0] + mgr.cover_admission_offer.assert_not_called() + + def test_zero_parks_still_books_pages(self): + """``admission_query: {}``: the connector is asked and the offer's pages are booked, + but every chunk keeps its charge.""" + mgr = self._parking_manager({0}) + sched = self._scheduler(mgr, parks=0) + head = make_ctx_request(0, context_remaining_length=20000) + out = sched.schedule_request([head, make_ctx_request(1, 1000)], set()) + assert ids(out.context_requests) == [0] + assert head.context_chunk_size == 8192 + mgr.query_connector_at_admission.assert_called_once_with(head) + mgr.cover_admission_offer.assert_called_once_with(head) + assert sched.uncharged_park_ids == frozenset() + + @pytest.mark.parametrize("admission", [False, True]) + def test_preemption_marks_the_requester_only_with_admission_queries(self, admission): + mgr = make_kv_cache_manager( + resize_context_fn=_out_of_pages_for(0), + has_cache_tier_below_gpu=False, + ) + mgr.connector_admission_query_enabled.return_value = True + mgr.query_connector_at_admission.return_value = False + mgr.cover_admission_offer.return_value = True + sched = make_scheduler( + mgr, + max_num_tokens=1000, + ctx_chunk_config=(ContextChunkingPolicy.FIRST_COME_FIRST_SERVED, 64), + connector_admission_query=admission, + ) + requester = make_ctx_request(0, 500) + victim = make_ctx_request(99, 100, is_first_context_chunk=False) + + out = sched.schedule_request([requester, victim], set()) + + mgr.preempt_request.assert_called_once_with(victim) + assert ids(out.recompute_paused_requests) == [99] + assert requester.py_preempted_for_pages is admission + + @pytest.mark.parametrize("answer", [False, "mock"]) + def test_inert_when_the_manager_reports_no_capability(self, answer): + """Attention DP, a capacity-only role, prefix reservations or a connector without + the capability all read as no capability; a mocked truthy answer does too.""" + mgr = self._parking_manager({0}) + mgr.connector_admission_query_enabled.return_value = Mock() if answer == "mock" else False + sched = self._scheduler(mgr) + head = make_ctx_request(0, context_remaining_length=20000) + out = sched.schedule_request([head, make_ctx_request(1, 1000)], set()) + assert ids(out.context_requests) == [0] + assert head.context_chunk_size == 8192 + mgr.begin_connector_admission_pass.assert_not_called() + mgr.query_connector_at_admission.assert_not_called() + assert sched.uncharged_park_ids == frozenset() + + def test_pregrown_replay_is_asked_with_its_exact_chunk(self): + """Under SWA scratch reuse a pre-grown chunk replays exactly; the connector sees that chunk.""" + sizes = [] + mgr = self._parking_manager({0}) + mgr.query_connector_at_admission.side_effect = lambda req: ( + sizes.append(req.context_chunk_size), + True, + )[1] + mgr.enable_swa_scratch_reuse = True + sched = make_scheduler( + mgr, + max_num_tokens=1024, + ctx_chunk_config=(None, 64), + connector_admission_query=True, + max_uncharged_connector_parks=1, + ) + req = make_ctx_request(0, context_remaining_length=10_000) + req.context_current_position = 1_024 + mgr.kv_cache_map[0].capacity = 1_024 + 512 + mgr.num_extra_kv_tokens + out = sched.schedule_request([req], set()) + assert ids(out.context_requests) == [0] + assert sizes == [512] + assert req.context_chunk_size == 512 + mgr.resize_context.assert_called_once_with(req, 512) + assert sched.uncharged_park_ids == frozenset({0}) + + def test_budget_skips_are_reported(self): + sched = self._scheduler(make_kv_cache_manager(tokens_per_block=64), admission=False) + sched.schedule_request( + [make_ctx_request(0, context_remaining_length=20000), make_ctx_request(1, 1000)], set() + ) + assert sched.last_pass_ctx_budget_skips == 1 + + def test_budget_skips_count_only_first_chunks(self): + """A continuation passed over for budget is not a new request starved of it.""" + sched = self._scheduler(make_kv_cache_manager(tokens_per_block=64), admission=False) + head = make_ctx_request(0, context_remaining_length=20000) + continuation = make_ctx_request(1, 12000, is_first_context_chunk=False) + sched.schedule_request([head, continuation, make_ctx_request(2, 1000)], set()) + assert sched.last_pass_ctx_budget_skips == 1 diff --git a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py index addb326f0a91..5fba4654f9b1 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -385,6 +385,7 @@ def _make_manager_for_cache_tier_test( vocab_size: int | None = 16, dtype: DataType = DataType.HALF, layer_specs: list[KVCacheLayerSpec] | None = None, + kv_connector_manager: object | None = None, ) -> tuple[KVCacheManagerV2, Mock]: impl_constructor = Mock(side_effect=impl_side_effect) if mapping is None: @@ -456,6 +457,7 @@ def build_cache_config( vocab_size=vocab_size, execution_stream=Mock(), cold_page_codec_provider=cold_page_codec_provider, + kv_connector_manager=kv_connector_manager, ) return manager, impl_constructor @@ -1228,6 +1230,98 @@ def test_one_model_prompt_lookahead_configures_reuse_backoff( assert replace(core_config, commit_min_snapshot=True).reuse_match_backoff == 1 +@pytest.mark.parametrize( + ( + "requires", + "drops", + "capacity_only", + "is_draft", + "expected_backoff", + "expected_alignment", + "expected_windows", + ), + [ + (False, False, False, False, 1, 1, [6]), + (True, False, False, False, 1, TOKENS_PER_BLOCK, [6]), + # A dropped backoff has no rewind, so the SWA window is not widened. + (True, True, False, False, 0, 1, [5]), + (True, True, True, False, 1, 1, [6]), + # A paired draft pool keeps its backoff but aligns with the target. + (True, True, False, True, 1, TOKENS_PER_BLOCK, [6]), + ], + ids=["no_declaration", "align", "drop", "capacity_only", "joint_draft"], +) +@pytest.mark.parametrize( + "spec_config", + [ + Eagle3DecodingConfig( + max_draft_len=1, + speculative_model="draft-model", + ), + MTPDecodingConfig(max_draft_len=1), + ], + ids=["eagle3_one_model", "mtp_eagle_one_model"], +) +def test_kv_connector_reuse_declarations_shape_the_core_match( + spec_config: Eagle3DecodingConfig | MTPDecodingConfig, + requires: bool, + drops: bool, + capacity_only: bool, + is_draft: bool, + expected_backoff: int, + expected_alignment: int, + expected_windows: list[int | None], +) -> None: + """A whole-block connector's declarations reach the real core config.""" + kv_cache_config = KvCacheConfig( + enable_block_reuse=True, + enable_partial_reuse=False, + max_gpu_total_bytes=16 << 20, + max_attention_window=[5, MAX_SEQ_LEN, MAX_SEQ_LEN - 1], + ) + connector = SimpleNamespace( + capacity_only=capacity_only, + worker=SimpleNamespace( + requires_block_aligned_reuse=requires, + drops_reuse_match_backoff=drops, + ), + ) + manager, _ = _make_manager_for_cache_tier_test( + kv_cache_config, + [Mock()], + spec_config=spec_config, + is_draft=is_draft, + joint_reuse=is_draft, + kv_connector_manager=connector, + ) + + assert manager.reuse_match_backoff == expected_backoff + # The raw lookahead stays for the lender and the public property. + assert manager.draft_prompt_lookahead == 1 + assert manager.max_attention_window_vec == expected_windows + if expected_backoff == 0: + # Without the backoff the lookup stops one token short of the prompt, + # exactly as it does without a draft. + prompt = list(range(65)) + request = SimpleNamespace( + multimodal_hashes=None, + multimodal_positions=None, + multimodal_lengths=None, + ) + manager._reuse_token_source = lambda _: prompt + assert list(manager._context_reuse_tokens(request)) == prompt[:-1] + + core_config = manager._build_base_config( + kv_cache_config, + tokens_per_block=TOKENS_PER_BLOCK, + cache_tiers=[GpuCacheTierConfig(quota=1 << 30)], + ) + assert core_config.reuse_match_backoff == expected_backoff + assert core_config.reuse_match_alignment == expected_alignment + rebuilt = replace(core_config, commit_min_snapshot=True) + assert rebuilt.reuse_match_alignment == expected_alignment + + @pytest.mark.parametrize( ("fresh_cache", "expected_lookup_tokens"), [(True, 3), (False, None)], @@ -1769,6 +1863,24 @@ def test_context_revert_shrinks_in_place_when_history_fits() -> None: assert request.context_chunk_size == 192 +@pytest.mark.cpu_only +def test_drop_context_allocation_frees_and_rewinds() -> None: + """A failed connector load restarts the request from the start of its prompt.""" + request = _revert_context_request(93) + manager = object.__new__(KVCacheManagerV2) + manager.tokens_per_block = 64 + manager.free_resources = Mock() + + manager.drop_context_allocation(request) + + manager.free_resources.assert_called_once_with(request) + request.set_prepopulated_prompt_len.assert_called_once_with(0, 64) + assert request.context_current_position == 0 + assert request.context_chunk_size == 512 + assert request.estimated_reusable_tokens == 0 + assert request.py_ctx_pre_resize_cap is None + + def test_draft_manager_keeps_shared_progress_across_context_and_generation() -> None: request = LlmRequest( request_id=39, diff --git a/tests/unittest/_torch/executor/test_adp_router.py b/tests/unittest/_torch/executor/test_adp_router.py index fbb33674022a..77388bb86480 100644 --- a/tests/unittest/_torch/executor/test_adp_router.py +++ b/tests/unittest/_torch/executor/test_adp_router.py @@ -23,11 +23,13 @@ """ import inspect -from unittest.mock import MagicMock, Mock +from types import SimpleNamespace +from unittest.mock import MagicMock, Mock, patch import numpy as np import pytest +import tensorrt_llm._torch.pyexecutor.scheduler.adp_router as adp_router_module from tensorrt_llm._torch.pyexecutor.executor_request_queue import RequestQueueItem from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState from tensorrt_llm._torch.pyexecutor.request_utils import get_from_waiting_queue @@ -39,6 +41,7 @@ KVCacheAwareADPRouter, RankIterStatsPayload, RankState, + WorkAwareADPRouter, _num_input_tokens, build_active_requests_for_overlap, count_retiring_requests, @@ -445,7 +448,12 @@ def test_router_factory_requires_the_headroom_flag(self): assert params["has_seq_slot_headroom"].default is inspect.Parameter.empty def test_router_constructors_default_to_no_headroom(self): - for cls in (DefaultADPRouter, ConversationAwareADPRouter, KVCacheAwareADPRouter): + for cls in ( + DefaultADPRouter, + ConversationAwareADPRouter, + KVCacheAwareADPRouter, + WorkAwareADPRouter, + ): default = inspect.signature(cls.__init__).parameters["has_seq_slot_headroom"].default assert default is False, cls.__name__ @@ -470,6 +478,11 @@ def test_router_factory_propagates_the_headroom_flag(self, has_seq_slot_headroom kv_cache_routing_fair_share_multiplier=2.0, kv_cache_routing_new_conv_placement="round_robin", ), + Mock( + kv_cache_routing_conversation_affinity=False, + enable_kv_cache_aware_routing=False, + enable_work_aware_routing=True, + ), ] built = set() for attention_dp_config in configs: @@ -478,14 +491,16 @@ def test_router_factory_propagates_the_headroom_flag(self, has_seq_slot_headroom has_seq_slot_headroom=has_seq_slot_headroom, kv_cache_manager=mgr, attention_dp_config=attention_dp_config, + kv_connector_manager=Mock(capacity_only=False), ) built.add(type(router).__name__) assert router.exclude_retiring_requests is has_seq_slot_headroom - # Anti-vacuity: the three configs must actually reach three branches. + # Anti-vacuity: the four configs must actually reach four branches. assert built == { "DefaultADPRouter", "KVCacheAwareADPRouter", "ConversationAwareADPRouter", + "WorkAwareADPRouter", } def test_router_factory_uses_default_router_for_kv_cache_aware_routing_under_helix(self): @@ -1920,6 +1935,316 @@ def test_new_conversations_never_exceed_slot_capacity(self, placement: str) -> N assert sum(len(v) for v in assign.values()) == 8 # nothing dropped +class TestWorkAwareADPRouter: + """Work-aware routing: ranks balanced by estimated uncached prefill.""" + + @staticmethod + def _item( + req_id, + prompt_len, + conversation_id=None, + *, + target_dp_rank=None, + attention_dp_relax=True, + child_req_ids=None, + subagent_affinity_id=None, + ): + scheduling_params = None + if target_dp_rank is not None or subagent_affinity_id is not None: + scheduling_params = SchedulingParams( + attention_dp_rank=target_dp_rank, + attention_dp_relax=attention_dp_relax, + subagent_affinity_id=subagent_affinity_id, + ) + request = SimpleNamespace( + num_input_tokens=prompt_len, + py_scheduling_params=scheduling_params, + py_conversation_params=None + if conversation_id is None + else ConversationParams(conversation_id=conversation_id), + ) + return RequestQueueItem(req_id, request, child_req_ids=child_req_ids) + + @staticmethod + def _states(*rank_loads): + """One RankState per (num_active_requests, num_active_tokens) pair.""" + return [ + RankState(rank=r, num_active_requests=a, num_active_tokens=t) + for r, (a, t) in enumerate(rank_loads) + ] + + @classmethod + def _route(cls, router, items, states=None, max_num_active_requests=8): + if states is None: + states = cls._states((0, 0), (0, 0)) + return router.route_requests(states, items, max_num_active_requests) + + @staticmethod + def _ranks(assigned): + return {item.id: rank for rank, items in assigned.items() for item in items} + + @staticmethod + def _enabled_config(): + cfg = AttentionDpConfig(enable_work_aware_routing=True) + return AttentionDpConfig.model_validate(cfg.model_dump()) + + @staticmethod + def _create(attention_dp_config, kv_connector_manager, is_estimating_kv_cache=False): + return ADPRouter.create( + dist=_mock_dist(tp_size=2), + has_seq_slot_headroom=True, + kv_cache_manager=SimpleNamespace( + enable_block_reuse=True, is_estimating_kv_cache=is_estimating_kv_cache + ), + attention_dp_config=attention_dp_config, + kv_connector_manager=kv_connector_manager, + ) + + def test_interface_compliance(self): + assert isinstance(WorkAwareADPRouter(_mock_dist(tp_size=2)), ADPRouter) + + def test_factory_selects_work_aware_router_from_config(self): + router = self._create(self._enabled_config(), Mock(capacity_only=False)) + assert type(router) is WorkAwareADPRouter + assert router.exclude_retiring_requests is True + + @pytest.mark.parametrize("capacity_only", [None, True], ids=["no_connector", "capacity_only"]) + def test_factory_falls_back_without_a_kv_moving_connector(self, capacity_only): + kv_connector_manager = None if capacity_only is None else Mock(capacity_only=True) + with patch.object(adp_router_module.logger, "warning") as warning: + router = self._create(self._enabled_config(), kv_connector_manager) + assert type(router) is DefaultADPRouter + warning.assert_called_once() + assert "enable_work_aware_routing" in warning.call_args.args[0] + + def test_factory_default_when_work_aware_routing_is_off(self): + router = self._create(AttentionDpConfig(), Mock(capacity_only=False)) + assert type(router) is DefaultADPRouter + + def test_estimation_executor_falls_back_silently(self): + """The estimation executor has no connector; the final executor decides.""" + with patch.object(adp_router_module.logger, "warning") as warning: + router = self._create(self._enabled_config(), None, is_estimating_kv_cache=True) + assert type(router) is DefaultADPRouter + warning.assert_not_called() + + def test_estimate_is_prompt_growth_within_a_conversation(self): + router = WorkAwareADPRouter(_mock_dist(tp_size=2)) + router.set_publish_uncached_estimate(True) + steps = [ + (1, 100_000, "conv", 100_000), + # The next turn reuses the previous prompt; only the growth is prefilled. + (2, 102_000, "conv", 2_000), + # A shrinking prompt (non-monotonic turn) has no usable predecessor. + (3, 50_000, "conv", 50_000), + # No conversation id: the whole prompt. + (4, 7_000, None, 7_000), + # An identical resubmission still counts one token. + (5, 50_000, "conv", 1), + ] + for req_id, prompt_len, conversation_id, expected_work in steps: + item = self._item(req_id, prompt_len, conversation_id) + self._route(router, [item]) + assert item.request.py_uncached_prefill_tokens == expected_work, req_id + assert router._estimated_work[req_id] == expected_work, req_id + + def test_large_prefill_and_small_requests_land_on_different_ranks(self): + router = WorkAwareADPRouter(_mock_dist(tp_size=2)) + seeds = [ + self._item(1000 + i, 130_000, conv) for i, conv in enumerate(("a", "b", "c", "whale")) + ] + self._route(router, seeds) + # Same prompt length everywhere, but "whale" brings 100K new tokens and the + # three others 2K each. Prompt-length balancing would split them 2/2. + items = [ + self._item(10, 230_000, "whale"), + self._item(11, 132_000, "a"), + self._item(12, 132_000, "b"), + self._item(13, 132_000, "c"), + ] + assigned, expected = self._route(router, items) + ranks = self._ranks(assigned) + assert ranks[10] != ranks[11] == ranks[12] == ranks[13] + assert expected == 3 + + def test_new_requests_go_to_the_rank_with_less_outstanding_work(self): + router = WorkAwareADPRouter(_mock_dist(tp_size=2)) + assigned, _ = self._route( + router, [self._item(20, 3_000)], states=self._states((0, 50_000), (0, 1_000)) + ) + assert [item.id for item in assigned[1]] == [20] + assert assigned[0] == [] + + def test_expected_count_covers_every_rank_and_respects_capacity(self): + router = WorkAwareADPRouter(_mock_dist(tp_size=2)) + items = [self._item(30 + i, 1_000) for i in range(3)] + assigned, expected = self._route( + router, items, states=self._states((2, 0), (0, 10**9)), max_num_active_requests=3 + ) + # Rank 0 is least loaded but full after one more request; the rest spill over. + assert len(assigned[0]) == 1 + assert len(assigned[1]) == 2 + assert expected == 3 + + def test_expected_is_capped_when_a_rank_is_already_over_capacity(self): + router = WorkAwareADPRouter(_mock_dist(tp_size=2)) + items = [self._item(30 + i, 1_000) for i in range(3)] + assigned, expected = self._route( + router, items, states=self._states((6, 0), (0, 0)), max_num_active_requests=4 + ) + assert assigned[0] == [] + assert len(assigned[1]) == 3 + assert expected == 4 + + def test_rank_state_counts_only_requests_still_in_context(self): + router = WorkAwareADPRouter(_mock_dist(tp_rank=1, tp_size=2)) + self._route(router, [self._item(1, 6_000, "x")]) + # Request 40 is estimated at 3_000 (growth over request 1), 41 at 5_000. + self._route(router, [self._item(40, 9_000, "x"), self._item(41, 5_000)]) + table_before = list(router._estimated_work.items()) + prefilling = SimpleNamespace( + py_request_id=40, is_context_init_state=True, py_orig_prompt_len=9_000 + ) + transferring = SimpleNamespace( + py_request_id=41, is_context_init_state=False, py_orig_prompt_len=5_000 + ) + unknown = SimpleNamespace( + py_request_id=99, is_context_init_state=True, py_orig_prompt_len=700 + ) + state = router.create_rank_state([prefilling, transferring, unknown], []) + assert (state.rank, state.num_active_requests, state.num_active_tokens) == (1, 3, 3_700) + # Read-only: no insertion and no LRU touch. + assert list(router._estimated_work.items()) == table_before + + def test_gather_all_rank_states_keeps_connector_fields(self): + dist = _mock_dist(tp_rank=0, tp_size=1) + dist.tp_allgather.side_effect = lambda payload: [payload] + router = WorkAwareADPRouter(dist) + single_rank = self._states((0, 0)) + self._route(router, [self._item(1, 6_000, "x")], states=single_rank) + self._route(router, [self._item(40, 9_000, "x")], states=single_rank) + states = router.gather_all_rank_states( + [ + SimpleNamespace( + py_request_id=40, is_context_init_state=True, py_orig_prompt_len=9_000 + ) + ], + num_connector_transfers=2, + connector_poll_failed=True, + ) + assert len(states) == 1 + assert states[0].num_active_requests == 1 + assert states[0].num_active_tokens == 3_000 + assert states[0].num_connector_transfers == 2 + assert states[0].connector_poll_failed is True + + def test_explicit_rank_requests_are_charged(self): + router = WorkAwareADPRouter(_mock_dist(tp_size=2)) + strict = self._item(1, 5_000, target_dp_rank=0, attention_dp_relax=False) + relaxed = self._item(2, 100) + assigned, _ = self._route( + router, [relaxed, strict], states=self._states((0, 0), (2, 1_000)) + ) + # Charged, the loads are [5_000, 1_000]; uncharged, [0, 1_000] would pull + # the relaxed request onto rank 0. + assert self._ranks(assigned) == {1: 0, 2: 1} + assert router._estimated_work[1] == 5_000 + + def test_child_ids_inherit_the_estimate(self): + router = WorkAwareADPRouter(_mock_dist(tp_size=2)) + self._route(router, [self._item(6, 1_000, "c")]) + self._route(router, [self._item(7, 1_300, "c", child_req_ids=[501, 502])]) + assert router._estimated_work[7] == 300 + assert router._estimated_work[501] == 300 + assert router._estimated_work[502] == 300 + + def test_estimate_published_only_when_enabled(self): + router = WorkAwareADPRouter(_mock_dist(tp_size=2)) + first = self._item(1, 1_000, "a") + self._route(router, [first]) + assert not hasattr(first.request, "py_uncached_prefill_tokens") + router.set_publish_uncached_estimate(True) + second = self._item(2, 1_500, "a") + self._route(router, [second]) + assert second.request.py_uncached_prefill_tokens == 500 + + def test_estimate_keyed_by_own_conversation_id(self): + """Sub-agents share their parent's affinity id but not its prompt history.""" + router = WorkAwareADPRouter(_mock_dist(tp_size=2)) + router.set_publish_uncached_estimate(True) + self._route(router, [self._item(0, 3_000, "parent")]) + # Keyed by the affinity id, this would count as 1_000 of growth over the parent. + child = self._item(1, 4_000, "child", subagent_affinity_id="parent") + self._route(router, [child]) + assert child.request.py_uncached_prefill_tokens == 4_000 + child_turn2 = self._item(2, 4_500, "child", subagent_affinity_id="parent") + self._route(router, [child_turn2]) + assert child_turn2.request.py_uncached_prefill_tokens == 500 + # The parent's history is its own too: growth over its 3_000 prompt. + parent_turn2 = self._item(3, 3_200, "parent") + self._route(router, [parent_turn2]) + assert parent_turn2.request.py_uncached_prefill_tokens == 200 + assert [router._estimated_work[i] for i in range(4)] == [3_000, 4_000, 500, 200] + + def test_routing_is_deterministic_across_ranks(self): + """Every rank routes the same inputs, so placements and tables must agree.""" + + def batches(): + return [ + [self._item(1, 8_000, "A"), self._item(2, 3_000, "B"), self._item(3, 500)], + [ + self._item(4, 9_000, "A", child_req_ids=[40, 41]), + self._item(5, 2_000, "C", target_dp_rank=1, attention_dp_relax=False), + self._item(6, 3_500, "B"), + ], + [ + self._item(7, 12_000, "A"), + self._item(8, 1_000, "C"), + self._item(9, 700, target_dp_rank=0), + ], + ] + + def run(tp_rank): + router = WorkAwareADPRouter(_mock_dist(tp_rank=tp_rank, tp_size=2)) + results = [] + for batch in batches(): + assigned, expected = self._route( + router, batch, states=self._states((1, 2_000), (0, 0)) + ) + results.append( + ( + {rank: [item.id for item in items] for rank, items in assigned.items()}, + expected, + ) + ) + return router, results + + r0, results0 = run(0) + r1, results1 = run(1) + assert results0 == results1 + assert list(r0._conv_prompt_len.items()) == list(r1._conv_prompt_len.items()) + assert list(r0._estimated_work.items()) == list(r1._estimated_work.items()) + + def test_lru_tables_are_bounded(self): + router = WorkAwareADPRouter(_mock_dist(tp_size=2), max_entries=2) + for i in range(5): + self._route(router, [self._item(i, 1_000 + i, f"c{i}")]) + assert list(router._conv_prompt_len) == ["c3", "c4"] + assert list(router._estimated_work) == [3, 4] + + @pytest.mark.parametrize( + "other_mode", ["kv_cache_routing_conversation_affinity", "enable_kv_cache_aware_routing"] + ) + def test_work_aware_rejects_other_routing_modes(self, other_mode): + with pytest.raises(ValueError, match="enable_work_aware_routing"): + AttentionDpConfig(enable_work_aware_routing=True, **{other_mode: True}) + + def test_work_aware_config_alone_validates(self): + cfg = AttentionDpConfig(enable_work_aware_routing=True) + assert cfg.enable_work_aware_routing is True + assert AttentionDpConfig().enable_work_aware_routing is False + + class TestNumInputTokens: """The cheap token-count accessor used by the routers (avoids copying the full input_token_ids list just to count it).""" diff --git a/tests/unittest/_torch/executor/test_disagg_receive_ordering.py b/tests/unittest/_torch/executor/test_disagg_receive_ordering.py index 89d9e727ff11..f9f6bd299605 100644 --- a/tests/unittest/_torch/executor/test_disagg_receive_ordering.py +++ b/tests/unittest/_torch/executor/test_disagg_receive_ordering.py @@ -58,6 +58,9 @@ def _manager() -> Mock: manager._on_allocation_stream = set() manager._allocation_ready = {} manager._pending_preemption = {} + # Real completion-reservation state: off (empty pool map), no bounds. + manager._disagg_completion_pool_counts = {} + manager._disagg_completion_limits = {} manager.is_draft = False manager.kv_connector_manager = None manager.enable_block_reuse = False diff --git a/tests/unittest/_torch/executor/test_kv_connector_executor_lifetime.py b/tests/unittest/_torch/executor/test_kv_connector_executor_lifetime.py index 36766f504881..810fb7ea730d 100644 --- a/tests/unittest/_torch/executor/test_kv_connector_executor_lifetime.py +++ b/tests/unittest/_torch/executor/test_kv_connector_executor_lifetime.py @@ -797,3 +797,214 @@ def test_shutdown_drains_the_connector_before_the_pools_it_reads_are_freed() -> # The worker thread first: it is the one still handing the connector work. executor.worker_thread.join.assert_called_once_with() assert order == ["connector", "kv pools"] + + +# ---- admission-time connector queries ---- + + +def _prepare_batch_executor(monkeypatch, parked_ids) -> PyExecutor: + """An executor whose connector parks ``parked_ids``, as take_scheduled_requests_pending_load does.""" + executor = _executor() + connector = executor.kv_connector_manager + + def park(batch): + batch.context_requests = [ + r for r in batch.context_requests if r.py_request_id not in parked_ids + ] + + connector.take_scheduled_requests_pending_load.side_effect = park + # Not on the Mock spec: it is an instance attribute of the real manager. + connector.worker = Mock() + monkeypatch.setattr("torch.cuda.current_stream", lambda: None) + return executor + + +def _ctx_request(request_id: int, chunk: int) -> SimpleNamespace: + return SimpleNamespace( + py_request_id=request_id, request_id=request_id, context_chunk_size=chunk + ) + + +def _admission_config(max_uncharged_parks: int = 2) -> SimpleNamespace: + return SimpleNamespace( + admission_query=SimpleNamespace( + max_uncharged_parks=max_uncharged_parks, max_offer_deferrals=16 + ) + ) + + +def _admission_executor(scheduler) -> PyExecutor: + """An executor whose connector can be asked at admission, with ``scheduler`` installed.""" + executor = _executor() + executor._kv_connector_config = _admission_config() + executor.kv_connector_manager.admission_query_capable = True + executor.kv_connector_manager.prefix_reservations_enabled = False + executor.scheduler = scheduler + executor.disable_overlap_scheduler = True + return executor + + +def test_an_uncharged_request_that_did_not_park_fails_stop(monkeypatch) -> None: + """A 0-charged chunk left in the batch would push the forward past max_num_tokens. + + Other parked requests' loads may already own destination pages, so the + failure retains DMA resources instead of taking the ordinary crash path. + """ + executor = _prepare_batch_executor(monkeypatch, parked_ids=set()) + executor._admission_query_scheduler = SimpleNamespace(uncharged_park_ids=frozenset({1})) + batch = SimpleNamespace(context_requests=[_ctx_request(1, 8192)]) + with pytest.raises(RuntimeError, match="did not park"): + executor._kv_connector_prepare_batch(batch) + assert executor._kv_connector_failed is True + executor.kv_connector_manager.worker.start_load_kv.assert_not_called() + + +def test_an_uncharged_request_that_parked_passes(monkeypatch) -> None: + executor = _prepare_batch_executor(monkeypatch, parked_ids={1}) + executor._admission_query_scheduler = SimpleNamespace(uncharged_park_ids=frozenset({1})) + batch = SimpleNamespace(context_requests=[_ctx_request(1, 8192), _ctx_request(2, 4096)]) + executor._kv_connector_prepare_batch(batch) + assert [r.py_request_id for r in batch.context_requests] == [2] + assert executor._kv_connector_failed is False + executor.kv_connector_manager.worker.start_load_kv.assert_called_once() + + +def test_admission_query_resolves_the_v2_scheduler_through_a_wrapper() -> None: + from tensorrt_llm._torch.pyexecutor.scheduler.scheduler_v2 import KVCacheV2Scheduler + + inner = object.__new__(KVCacheV2Scheduler) + executor = _admission_executor(SimpleNamespace(scheduler=inner)) + assert executor._resolve_admission_query_scheduler() is inner + # Off unless configured, including on an executor built without __init__. + executor._kv_connector_config = SimpleNamespace(admission_query=None) + assert executor._resolve_admission_query_scheduler() is None + del executor._kv_connector_config + assert executor._resolve_admission_query_scheduler() is None + + +@pytest.mark.parametrize("problem", ["no_v2_scheduler", "overlap_scheduler"]) +def test_admission_query_is_refused_where_unchecked(problem) -> None: + from tensorrt_llm._torch.pyexecutor.scheduler.scheduler_v2 import KVCacheV2Scheduler + + executor = _admission_executor( + SimpleNamespace() if problem == "no_v2_scheduler" else object.__new__(KVCacheV2Scheduler) + ) + executor.disable_overlap_scheduler = problem != "overlap_scheduler" + with pytest.raises(NotImplementedError, match="admission_query"): + executor._resolve_admission_query_scheduler() + + +@pytest.mark.parametrize( + "setting, reason", + [ + ("not_capable", "supports_admission_query"), + ("prefix_reservations", "reserves its prefix"), + ("capacity_only", "capacity-only"), + ], +) +def test_admission_query_without_capability_only_warns(monkeypatch, setting, reason) -> None: + from tensorrt_llm._torch.pyexecutor.scheduler.scheduler_v2 import KVCacheV2Scheduler + + executor = _admission_executor(object.__new__(KVCacheV2Scheduler)) + # Neither would be accepted if the configuration were not ignored first. + executor.disable_overlap_scheduler = False + connector = executor.kv_connector_manager + if setting == "not_capable": + connector.admission_query_capable = False + elif setting == "prefix_reservations": + connector.prefix_reservations_enabled = True + else: + connector.capacity_only = True + warning = Mock() + monkeypatch.setattr("tensorrt_llm._torch.pyexecutor.py_executor.logger.warning", warning) + assert executor._resolve_admission_query_scheduler() is None + warning.assert_called_once() + assert reason in warning.call_args.args[0] + + +# ---- context-park diagnostics (TLLM_CTX_PARK_STATS) ---- + + +def test_park_stats_count_the_parked_chunks(monkeypatch) -> None: + from tensorrt_llm._torch.pyexecutor.profiling import ParkStats + + executor = _prepare_batch_executor(monkeypatch, parked_ids={1}) + # Request 9 was admitted uncharged but dropped before the connector saw the batch. + executor._admission_query_scheduler = SimpleNamespace(uncharged_park_ids=frozenset({1, 9})) + executor._park_stats = ParkStats() + batch = SimpleNamespace(context_requests=[_ctx_request(1, 8192), _ctx_request(2, 4096)]) + executor._kv_connector_prepare_batch(batch) + stats = executor._park_stats + assert (stats.parked_reqs, stats.parked_tokens, stats.fwd_ctx_tokens, stats.uncharged) == ( + 1, + 8192, + 4096, + 1, + ) + + +def test_park_stats_suffix_is_empty_when_off_and_parseable_when_on() -> None: + from tensorrt_llm._torch.pyexecutor.profiling import ParkStats + + executor = _executor() + assert executor._format_park_stats() == "" + executor._park_stats = ParkStats() + executor.scheduler = None + executor.kv_cache_manager = Mock() + executor.kv_cache_manager.take_connector_offer_stats.return_value = { + "queries": 2, + "query_ms": 1.234, + } + with executor._park_phase("fwd"): + pass + suffix = executor._format_park_stats() + assert suffix.startswith(", ctx_park = parked=0 parked_tok=0 unc=0 fwd_tok=0 bskip=0 sched_ms=") + assert "offer_queries=2 offer_query_ms=1.2" in suffix + # No braces: parsers of the `states = {...}` field just before it keep matching. + assert "{" not in suffix and "}" not in suffix + executor._reset_park_stats() + assert executor._park_stats.wall_ms["fwd"] == 0.0 + + +def test_park_stats_report_the_schedulers_budget_skips() -> None: + from tensorrt_llm._torch.pyexecutor.profiling import ParkStats + from tensorrt_llm._torch.pyexecutor.scheduler.scheduler_v2 import KVCacheV2Scheduler + + executor = _executor() + executor._park_stats = ParkStats() + inner = object.__new__(KVCacheV2Scheduler) + inner.last_pass_ctx_budget_skips = 3 + executor.scheduler = SimpleNamespace(scheduler=inner) + executor.kv_cache_manager = SimpleNamespace() + assert " bskip=3 " in executor._format_park_stats() + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_reset_drains_offer_stats_only_when_enabled(enabled) -> None: + """Ranks that print no log line still drain the counters; a disabled executor never asks.""" + from tensorrt_llm._torch.pyexecutor.profiling import ParkStats + + executor = _executor() + executor._park_stats = ParkStats() if enabled else None + executor._reset_park_stats() + assert executor.kv_cache_manager.take_connector_offer_stats.call_count == int(enabled) + + +def test_gil_probe_reports_and_stops() -> None: + import time + + from tensorrt_llm._torch.pyexecutor.profiling import GilProbe + + probe = GilProbe(500) + try: + deadline = time.monotonic() + 2.0 + while not probe._samples and time.monotonic() < deadline: + time.sleep(0.01) + reading = probe.take() + assert reading is not None + worst, p99, count = reading + assert worst >= p99 and count > 0 + finally: + probe.stop() + probe._thread.join(timeout=1.0) + assert not probe._thread.is_alive() diff --git a/tests/unittest/_torch/executor/test_kv_connector_initialization.py b/tests/unittest/_torch/executor/test_kv_connector_initialization.py index c8c6d0db713a..7e3d95bf1aa4 100644 --- a/tests/unittest/_torch/executor/test_kv_connector_initialization.py +++ b/tests/unittest/_torch/executor/test_kv_connector_initialization.py @@ -63,6 +63,146 @@ def test_v2_connector_registers_layer_groups( windows[0] if len(windows) == 1 else None ) connector.wait_for_initialization.assert_called_once() + # No kv_connector_config on an executor built without __init__: no admission queries. + assert executor._admission_query_scheduler is None + + +def test_admission_query_is_resolved_during_connector_init( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from tensorrt_llm._torch.pyexecutor.scheduler.scheduler_v2 import KVCacheV2Scheduler + + executor = object.__new__(PyExecutor) + connector = Mock(capacity_only=False, enable_attention_dp=False) + # Real values: a Mock auto-attribute would answer truthily. + connector.admission_query_capable = True + connector.prefix_reservations_enabled = False + connector.speculative_prefill_only = False + executor.kv_connector_manager = connector + executor.kv_cache_transceiver = None + executor.dist = SimpleNamespace(pp_size=1, cp_size=1) + executor.max_beam_width = 1 + executor.max_draft_len = 0 + executor.enable_attention_dp = False + executor.llm_args = SimpleNamespace(kv_cache_config=None, scheduler_config=None) + executor._kv_connector_config = SimpleNamespace( + admission_query=SimpleNamespace(max_uncharged_parks=0, max_offer_deferrals=16) + ) + executor.disable_overlap_scheduler = True + scheduler = object.__new__(KVCacheV2Scheduler) + executor.scheduler = scheduler + executor.kv_cache_manager = Mock(spec=KVCacheManagerV2, enable_block_reuse=True) + executor.kv_cache_manager.impl = SimpleNamespace( + init_config=SimpleNamespace(cache_tiers=[SimpleNamespace(tier=CacheTier.GPU_MEM)]) + ) + executor.model_engine = Mock() + executor.model_engine.model.named_modules.return_value = [] + layout = SimpleNamespace(groups=[SimpleNamespace(window_size=None)]) + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.build_kv_cache_layout_v2", + lambda manager: layout, + ) + + executor._maybe_init_kv_connector_manager() + + assert executor._admission_query_scheduler is scheduler + connector.wait_for_initialization.assert_called_once() + + +def _speculative_connector_executor( + monkeypatch: pytest.MonkeyPatch, + connector: Mock, + *, + max_draft_len: int, + draft_kv_cache_manager: Mock | None = None, + kv_cache_manager: Mock | None = None, +) -> tuple[PyExecutor, SimpleNamespace]: + """The executor shell of `test_v2_connector_registers_layer_groups`, with a draft.""" + executor = object.__new__(PyExecutor) + executor.kv_connector_manager = connector + executor.kv_cache_transceiver = None + executor.dist = SimpleNamespace(pp_size=1, cp_size=1) + executor.max_beam_width = 1 + executor.max_draft_len = max_draft_len + executor.draft_kv_cache_manager = draft_kv_cache_manager + executor.enable_attention_dp = False + executor.llm_args = SimpleNamespace(kv_cache_config=None, scheduler_config=None) + if kv_cache_manager is None: + kv_cache_manager = Mock(spec=KVCacheManagerV2, enable_block_reuse=True) + kv_cache_manager.impl = SimpleNamespace( + init_config=SimpleNamespace(cache_tiers=[SimpleNamespace(tier=CacheTier.GPU_MEM)]) + ) + executor.kv_cache_manager = kv_cache_manager + executor.model_engine = Mock() + executor.model_engine.model.named_modules.return_value = [] + layout = SimpleNamespace(groups=[SimpleNamespace(window_size=None)]) + monkeypatch.setattr( + "tensorrt_llm._torch.pyexecutor.py_executor.build_kv_cache_layout_v2", + lambda manager: layout, + ) + return executor, layout + + +@pytest.mark.parametrize( + "max_draft_len,speculative_prefill_only,separate_draft_cache,error", + [ + pytest.param(1, False, False, "speculative decoding", id="refused"), + pytest.param(1, True, False, None, id="speculative-prefill"), + pytest.param(1, True, True, "share the target KV cache", id="separate-draft-cache"), + pytest.param(0, False, False, None, id="no-speculation"), + ], +) +def test_v2_connector_speculative_decoding_guard( + monkeypatch: pytest.MonkeyPatch, + max_draft_len: int, + speculative_prefill_only: bool, + separate_draft_cache: bool, + error: str | None, +) -> None: + connector = Mock( + capacity_only=False, + enable_attention_dp=False, + speculative_prefill_only=speculative_prefill_only, + ) + executor, layout = _speculative_connector_executor( + monkeypatch, + connector, + max_draft_len=max_draft_len, + draft_kv_cache_manager=Mock() if separate_draft_cache else None, + ) + + if error is not None: + with pytest.raises(NotImplementedError, match=error): + executor._maybe_init_kv_connector_manager() + connector.worker.register_kv_cache_layout.assert_not_called() + else: + executor._maybe_init_kv_connector_manager() + connector.worker.register_kv_cache_layout.assert_called_once_with(layout) + connector.wait_for_initialization.assert_called_once() + + +def test_speculative_prefill_requires_kv_cache_manager_v2( + monkeypatch: pytest.MonkeyPatch, +) -> None: + connector = Mock(capacity_only=False, enable_attention_dp=False, speculative_prefill_only=True) + executor, _ = _speculative_connector_executor( + monkeypatch, connector, max_draft_len=1, kv_cache_manager=Mock() + ) + + with pytest.raises(NotImplementedError, match="KVCacheManagerV2"): + executor._maybe_init_kv_connector_manager() + connector.worker.register_kv_caches.assert_not_called() + + +def test_a_mock_mode_attribute_does_not_lift_the_speculative_refusal( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # `speculative_prefill_only` is a truthy Mock auto-attribute here, not True. + connector = Mock(capacity_only=False, enable_attention_dp=False) + executor, _ = _speculative_connector_executor(monkeypatch, connector, max_draft_len=1) + + with pytest.raises(NotImplementedError, match="speculative decoding"): + executor._maybe_init_kv_connector_manager() @pytest.mark.parametrize( diff --git a/tests/unittest/_torch/executor/test_kv_connector_speculative_prefill.py b/tests/unittest/_torch/executor/test_kv_connector_speculative_prefill.py new file mode 100644 index 000000000000..d8f96e891b85 --- /dev/null +++ b/tests/unittest/_torch/executor/test_kv_connector_speculative_prefill.py @@ -0,0 +1,311 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The KV connector manager's speculative-prefill-only mode. + +The mode lets one-model speculative decoding run next to a connector by +refusing every request that is not context-only work: at admission, where only +the request fails, and again at the connector boundary for anything that +bypasses admission. Dummy requests and a connector that transfers no KV are +never refused. +""" + +from collections.abc import Iterator +from contextlib import ExitStack +from types import SimpleNamespace +from typing import Optional +from unittest.mock import MagicMock, Mock, patch + +import pytest + +from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import KvCacheConnectorManager +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor, _validate_connector_admission +from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests + +pytestmark = pytest.mark.cpu_only +CONNECTOR_MODULE = "tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector" + + +class SpecMock(MagicMock): + supports_attention_dp = True + supports_speculative_prefill = True + + +def _request( + request_id: int, + *, + context_only: bool = True, + generation_only: bool = False, + dummy: bool = False, +) -> MagicMock: + """A request whose phase flags are explicit bools. + + The manager compares `is_context_only_request` against True, and a bare + MagicMock attribute is merely truthy, so every flag is set here. + """ + req = MagicMock( + is_dummy_request=dummy, + is_attention_dp_dummy=dummy, + is_context_only_request=context_only, + is_generation_only_request=generation_only, + ) + req.request_id = request_id + req.py_request_id = request_id + req.is_child = False + req.is_finished = False + req.state = ( + LlmRequestState.CONTEXT_INIT + if context_only or not generation_only + else LlmRequestState.GENERATION_IN_PROGRESS + ) + req.context_current_position = 0 + req.context_remaining_length = 8 + req.context_chunk_size = 8 + req.get_tokens.return_value = list(range(8)) + req.py_draft_tokens = [] + req.kv_cache_retention_config = None + req.cache_salt = None + return req + + +@pytest.fixture +def exit_stack() -> Iterator[ExitStack]: + with ExitStack() as stack: + yield stack + + +def _manager( + attention_dp: bool, + *, + stack: Optional[ExitStack] = None, + capacity_only: bool = False, + speculative_prefill_only: bool = True, +) -> KvCacheConnectorManager: + """A manager in the speculative-prefill-only mode by default. + + A tensor-parallel manager reaches the leader broadcast on every query, so + it needs its MPI patches for the whole test: they are entered into + ``stack``, and the broadcast is an identity that records its calls. + """ + worker = SpecMock(capacity_only=capacity_only) + worker.get_finished.return_value = ([], []) + scheduler = SpecMock() + scheduler.get_num_new_matched_tokens.return_value = (4, False) + scheduler.reserve_prefix.return_value = (4, False) + if attention_dp: + with patch(f"{CONNECTOR_MODULE}.mpi_rank", return_value=1): + return KvCacheConnectorManager( + worker, + scheduler, + enable_attention_dp=True, + speculative_prefill_only=speculative_prefill_only, + ) + assert stack is not None, "a tensor-parallel manager needs MPI patches for the whole test" + stack.enter_context(patch(f"{CONNECTOR_MODULE}.mpi_rank", return_value=0)) + stack.enter_context( + patch(f"{CONNECTOR_MODULE}.mpi_broadcast", side_effect=lambda result, root: result) + ) + return KvCacheConnectorManager( + worker, scheduler, speculative_prefill_only=speculative_prefill_only + ) + + +def _broadcast(manager: KvCacheConnectorManager) -> Mock: + from tensorrt_llm._torch.pyexecutor.connectors import kv_cache_connector + + assert not manager.enable_attention_dp + return kv_cache_connector.mpi_broadcast + + +def test_modes_default_off() -> None: + with ( + patch(f"{CONNECTOR_MODULE}.mpi_rank", return_value=0), + patch(f"{CONNECTOR_MODULE}.mpi_broadcast", side_effect=lambda result, root: result), + ): + manager = KvCacheConnectorManager(MagicMock(), MagicMock()) + assert manager.speculative_prefill_only is False + assert manager.enable_attention_dp is False + + +@pytest.mark.parametrize("undeclared", ["worker", "scheduler"]) +def test_constructor_refuses_undeclared_connector(undeclared: str) -> None: + worker = MagicMock() if undeclared == "worker" else SpecMock() + scheduler = MagicMock() if undeclared == "scheduler" else SpecMock() + with patch(f"{CONNECTOR_MODULE}.mpi_rank", return_value=0): + with pytest.raises(NotImplementedError, match="supports_speculative_prefill"): + KvCacheConnectorManager(worker, scheduler, speculative_prefill_only=True) + + +def test_tp_follower_without_scheduler_checks_worker_only() -> None: + with patch(f"{CONNECTOR_MODULE}.mpi_rank", return_value=1): + manager = KvCacheConnectorManager(SpecMock(), None, speculative_prefill_only=True) + assert manager.speculative_prefill_only is True + with pytest.raises(NotImplementedError, match="supports_speculative_prefill"): + KvCacheConnectorManager(MagicMock(), None, speculative_prefill_only=True) + + +@pytest.mark.parametrize("context_only,dummy", [(True, False), (False, True)]) +def test_speculative_prefill_allows_context_work_and_adp_dummies( + context_only: bool, dummy: bool +) -> None: + manager = _manager(attention_dp=True) + request = _request(7, context_only=context_only, dummy=dummy) + + num_tokens, _ = manager.query_num_new_matched_tokens(request, 0) + + assert num_tokens == (0 if dummy else 4) + assert manager.scheduler.get_num_new_matched_tokens.call_count == (0 if dummy else 1) + + +@pytest.mark.parametrize("attention_dp", [True, False], ids=["adp", "tp"]) +@pytest.mark.parametrize("generation_only", [False, True]) +def test_speculative_prefill_rejects_generation_before_lookup( + exit_stack: ExitStack, attention_dp: bool, generation_only: bool +) -> None: + manager = _manager(attention_dp, stack=exit_stack) + request = _request(7, context_only=False, generation_only=generation_only) + + with pytest.raises(NotImplementedError, match="context-only"): + manager.query_num_new_matched_tokens(request, 0) + + manager.scheduler.get_num_new_matched_tokens.assert_not_called() + if not attention_dp: + # Every rank raises for the same request; none waits in the broadcast. + _broadcast(manager).assert_not_called() + + +@pytest.mark.parametrize("attention_dp", [True, False], ids=["adp", "tp"]) +def test_speculative_prefill_rejects_before_prefix_reservation( + exit_stack: ExitStack, attention_dp: bool +) -> None: + manager = _manager(attention_dp, stack=exit_stack) + manager.prefix_reservations_enabled = True + + with pytest.raises(NotImplementedError, match="context-only"): + manager.reserve_prefix(_request(7, context_only=False), 0) + manager.scheduler.reserve_prefix.assert_not_called() + if not attention_dp: + _broadcast(manager).assert_not_called() + + reservation = manager.reserve_prefix(_request(8), 0) + assert (reservation.start, reservation.end) == (0, 4) + manager.scheduler.reserve_prefix.assert_called_once() + + +@pytest.mark.parametrize("generation_only", [False, True]) +def test_speculative_prefill_rejects_generation_before_transfer_scheduling( + generation_only: bool, +) -> None: + manager = _manager(attention_dp=True) + batch = ScheduledRequests() + batch.context_requests_last_chunk.append(_request(6)) + batch.generation_requests.append( + _request(7, context_only=False, generation_only=generation_only) + ) + + with patch(f"{CONNECTOR_MODULE}.mpi_rank", return_value=1): + with pytest.raises(NotImplementedError, match="context-only"): + manager.build_scheduler_output(batch, MagicMock(spec=KVCacheManagerV2)) + + assert manager.scheduler_output_manager.requests == {} + manager.handle_metadata() + manager.scheduler.build_connector_meta.assert_not_called() + + +def test_context_batch_with_adp_dummy_is_reported() -> None: + manager = _manager(attention_dp=True) + context = _request(6) + batch = ScheduledRequests() + batch.context_requests_last_chunk = [context, _request(900, context_only=False, dummy=True)] + batch.generation_requests = [_request(901, context_only=False, dummy=True)] + cache = MagicMock(spec=KVCacheManagerV2) + cache.get_page_indices_by_layer_group.return_value = [[5]] + + with patch(f"{CONNECTOR_MODULE}.mpi_rank", return_value=1): + manager.build_scheduler_output(batch, cache) + manager.handle_metadata() + + output = manager.scheduler.build_connector_meta.call_args.args[0] + assert [req.request_id for req in output.new_requests] == [6] + assert output.cached_requests == [] + + +def test_capacity_only_connector_never_refuses() -> None: + manager = _manager(attention_dp=True, capacity_only=True) + assert manager.capacity_only is True + assert manager.speculative_prefill_only is True + request = _request(7, context_only=False, generation_only=True) + + assert manager.query_num_new_matched_tokens(request, 0) == (0, False) + assert manager.get_num_new_matched_tokens(request, 0) == 0 + assert manager.reserve_prefix(request, 0) is None + batch = ScheduledRequests() + batch.generation_requests.append(request) + assert manager.build_scheduler_output(batch, MagicMock(spec=KVCacheManagerV2)) is None + _validate_connector_admission(manager, request) + manager.scheduler.get_num_new_matched_tokens.assert_not_called() + manager.scheduler.reserve_prefix.assert_not_called() + + +@pytest.mark.parametrize("capacity_only", [False, True]) +@pytest.mark.parametrize("context_only", [False, True]) +@pytest.mark.parametrize("prefill_only", [False, True]) +def test_admission_refuses_only_non_context_requests( + prefill_only: bool, context_only: bool, capacity_only: bool +) -> None: + """Only the refused request fails; the executor never reaches resource preparation.""" + manager = _manager( + attention_dp=True, capacity_only=capacity_only, speculative_prefill_only=prefill_only + ) + executor = object.__new__(PyExecutor) + executor.kv_connector_manager = manager + request = _request(7, context_only=context_only, generation_only=not context_only) + + if prefill_only and not context_only and not capacity_only: + # Nothing else is set on the executor, so this raising proves the + # admission check is the first thing `_validate_request` does. + with pytest.raises(ValueError, match="context-only"): + PyExecutor._validate_request(executor, request) + else: + _validate_connector_admission(executor.kv_connector_manager, request) + + +def test_only_an_explicit_context_only_flag_is_admitted() -> None: + """A merely truthy `is_context_only_request` is not context-only work.""" + manager = _manager(attention_dp=True) + request = _request(7) + request.is_context_only_request = "yes" + + with pytest.raises(ValueError, match="context-only"): + _validate_connector_admission(manager, request) + with pytest.raises(NotImplementedError, match="context-only"): + manager.query_num_new_matched_tokens(request, 0) + manager.scheduler.get_num_new_matched_tokens.assert_not_called() + + +@pytest.mark.parametrize( + "shell", + [ + SimpleNamespace(), + SimpleNamespace(kv_connector_manager=None), + SimpleNamespace(kv_connector_manager=Mock(capacity_only=False)), + SimpleNamespace(kv_connector_manager=Mock(spec=KvCacheConnectorManager)), + ], + ids=["no_attribute", "no_connector", "mock_connector", "spec_mock_connector"], +) +def test_admission_ignores_executor_shells_without_the_mode(shell: SimpleNamespace) -> None: + request = _request(7, context_only=False, generation_only=True) + _validate_connector_admission(getattr(shell, "kv_connector_manager", None), request) diff --git a/tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py b/tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py index 3688bc0ea054..d686634be560 100644 --- a/tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py +++ b/tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py @@ -100,6 +100,7 @@ def __init__(self, request_id=0, prompt_len=PROMPT_LEN): self.is_disagg_generation_init_state = False self.py_num_connector_matched_tokens = 0 self.py_connector_allocation_reported = False + self.py_connector_admission_offer = None self.py_served_prefix_end = 0 self.py_ctx_pre_resize_cap = None self.py_draft_tokens = [] @@ -155,6 +156,7 @@ def __init__(self, num_matched=0, load_async=False, add_sequence=True): # Read by the manager to decide whether the prefix path runs at all. self.capacity_only = False self.prefix_reservations_enabled = False + self.admission_query_capable = False self.reservations = {} self.reservation_requests = {} self.next_reservation_id = 1 @@ -1095,6 +1097,262 @@ def test_final_batch_releases_only_rejected_requests(self) -> None: assert other.context_current_position == 0 +class TestAdmissionQuery: + """kv_connector_config.admission_query: the scheduler asks at admission, and prepare reuses that answer.""" + + @staticmethod + def _admitted(num_matched=64, load_async=True, capable=True, committed=0): + connector = FakeConnectorManager(num_matched=num_matched, load_async=load_async) + connector.admission_query_capable = capable + manager = make_manager(connector) + req = FakeRequest() + kv_cache = FakeKvCache(committed=committed, capacity=PROMPT_LEN) + manager.kv_cache_map[req.py_request_id] = kv_cache + req.context_current_position = committed + req.set_prepopulated_prompt_len(committed, TOKENS_PER_BLOCK) + manager.begin_connector_admission_pass() + return connector, manager, req, kv_cache + + def test_an_async_answer_parks_and_is_not_asked_again(self): + connector, manager, req, _ = self._admitted() + assert manager.query_connector_at_admission(req) + assert connector.queries == [(0, 0)] + + manager._run_kv_connector_hooks(scheduled(req)) + + assert connector.queries == [(0, 0)] + assert connector.commits == [(0, 64, True)] + assert len(connector.allocs) == 1 + assert req.py_connector_admission_offer is None + + @pytest.mark.parametrize("num_matched, load_async", [(64, False), (0, False)]) + def test_sync_and_miss_answers_do_not_park_and_are_used_once(self, num_matched, load_async): + connector, manager, req, _ = self._admitted(num_matched, load_async) + assert not manager.query_connector_at_admission(req) + + manager._run_kv_connector_hooks(scheduled(req)) + + assert connector.queries == [(0, 0)] + assert connector.commits == [(0, num_matched, False)] + + def test_no_capability_means_no_admission_query(self): + connector, manager, req, _ = self._admitted(capable=False) + assert not manager.query_connector_at_admission(req) + assert connector.queries == [] + assert req.py_connector_admission_offer is None + + @pytest.mark.parametrize( + "guard", ["reservations", "draft", "reported", "no_add_sequence", "capacity_only"] + ) + def test_every_guard_blocks_the_query(self, guard): + connector, manager, req, _ = self._admitted() + if guard == "reservations": + connector.prefix_reservations_enabled = True + elif guard == "draft": + manager.is_draft = True + elif guard == "reported": + req.py_connector_allocation_reported = True + elif guard == "no_add_sequence": + connector.add_sequence = False + else: + connector.capacity_only = True + assert not manager.query_connector_at_admission(req) + assert connector.queries == [] + + def test_a_blocked_query_drops_an_old_answer(self): + _, manager, req, _ = self._admitted() + assert manager.query_connector_at_admission(req) + assert req.py_connector_admission_offer is not None + req.py_connector_allocation_reported = True + assert not manager.query_connector_at_admission(req) + assert req.py_connector_admission_offer is None + + def test_an_answer_from_an_earlier_pass_is_asked_again(self): + """A request dropped before prepare (e.g. by the total-KV cap) is asked afresh next pass, + and its pages are still reported once.""" + connector, manager, req, _ = self._admitted() + assert manager.query_connector_at_admission(req) + manager.begin_connector_admission_pass() + + manager._run_kv_connector_hooks(scheduled(req)) + + assert connector.queries == [(0, 0), (0, 0)] + assert len(connector.allocs) == 1 + assert req.py_connector_admission_offer is None + + def test_a_moved_cache_is_refused(self): + _, manager, req, kv_cache = self._admitted() + assert manager.query_connector_at_admission(req) + kv_cache.num_committed_tokens = TOKENS_PER_BLOCK + with pytest.raises(RuntimeError, match="admission-time connector answer"): + manager._run_kv_connector_hooks(scheduled(req)) + + def test_releasing_the_allocation_drops_the_answer(self): + _, manager, req, _ = self._admitted() + assert manager.query_connector_at_admission(req) + manager.kv_cache_map.pop(req.py_request_id) + manager.free_resources(req) + assert req.py_connector_admission_offer is None + + @pytest.mark.parametrize( + "change", + [None, "no_connector", "draft", "not_capable", "capacity_only", "reservations"], + ) + def test_connector_admission_query_enabled_gates(self, change): + connector, manager, _, _ = self._admitted() + if change == "no_connector": + manager.kv_connector_manager = None + elif change == "draft": + manager.is_draft = True + elif change == "not_capable": + connector.admission_query_capable = False + elif change == "capacity_only": + connector.capacity_only = True + elif change == "reservations": + connector.prefix_reservations_enabled = True + assert manager.connector_admission_query_enabled() is (change is None) + + def test_offer_stats_classify_the_served_prefix(self): + connector, manager, req, _ = self._admitted() + manager.enable_connector_offer_stats() + assert manager.query_connector_at_admission(req) + + manager._run_kv_connector_hooks(scheduled(req)) + stats = manager.take_connector_offer_stats() + + assert stats["queries"] == 1 and stats["admission_used"] == 1 + assert stats["offers"] == 1 and stats["async_offers"] == 1 + assert stats["offered_tokens"] == 64 and stats["honoured_tokens"] == 64 + assert stats["in_chunk"] == 1 + assert stats["shifted"] == stats["capped"] == stats["not_honoured"] == 0 + assert manager.take_connector_offer_stats()["queries"] == 0 + + def test_a_final_batch_query_is_counted_too(self): + connector = FakeConnectorManager(num_matched=64) + manager = make_manager(connector) + manager.enable_connector_offer_stats() + req = FakeRequest() + manager.kv_cache_map[req.py_request_id] = FakeKvCache(committed=0, capacity=PROMPT_LEN) + + manager._run_kv_connector_hooks(scheduled(req)) + stats = manager.take_connector_offer_stats() + + assert stats["queries"] == 1 and stats["admission_used"] == 0 + assert stats["offers"] == 1 and stats["async_offers"] == 0 + + def test_offer_stats_are_off_by_default(self): + _, manager, req, _ = self._admitted() + assert manager.query_connector_at_admission(req) + manager._run_kv_connector_hooks(scheduled(req)) + assert manager.take_connector_offer_stats() is None + + +class TestAdmissionGrow: + """``cover_admission_offer``: an admission-time offer past the chunk gets its pages before later admissions.""" + + @staticmethod + def _asked(num_matched=128, grow_ok=True, capacity=64, chunk=64, scratch_reuse=False): + connector = FakeConnectorManager(num_matched=num_matched, load_async=True) + connector.admission_query_capable = True + manager = make_manager(connector) + manager.enable_swa_scratch_reuse = scratch_reuse + manager.enable_connector_offer_stats() + req = FakeRequest() + req.context_chunk_size = chunk + kv_cache = FakeKvCache(committed=0, capacity=capacity, grow_ok=grow_ok) + manager.kv_cache_map[req.py_request_id] = kv_cache + manager.begin_connector_admission_pass() + assert manager.query_connector_at_admission(req) == (num_matched > 0) + return connector, manager, req, kv_cache + + def test_a_shifting_offer_is_covered_at_admission_and_served_in_full(self): + connector, manager, req, kv_cache = self._asked() + + assert manager.cover_admission_offer(req) + # The shifted window [128, 192): capacity only, the history stays at the local match. + assert kv_cache.resize_calls == [(192, None)] + assert kv_cache.history_length == 0 + assert req.py_ctx_pre_resize_cap == 64 + + manager._run_kv_connector_hooks(scheduled(req)) + stats = manager.take_connector_offer_stats() + + assert req.context_current_position == 128 and req.context_chunk_size == 64 + assert kv_cache.capacity == 192 + assert connector.commits == [(0, 128, True)] + assert stats["pregrown"] == 1 and stats["shifted"] == 1 and stats["capped"] == 0 + + def test_without_pages_the_cover_fails_and_leaves_the_allocation(self): + _, manager, req, kv_cache = self._asked(grow_ok=False) + + assert not manager.cover_admission_offer(req) + + assert kv_cache.capacity == 64 and kv_cache.history_length == 0 + assert req.py_ctx_pre_resize_cap is None + assert manager.take_connector_offer_stats()["pregrow_failed"] == 1 + + @pytest.mark.parametrize("num_matched", [0, 32]) + def test_a_miss_or_an_offer_inside_the_chunk_needs_no_pages(self, num_matched): + _, manager, req, kv_cache = self._asked(num_matched=num_matched, grow_ok=False) + assert manager.cover_admission_offer(req) + assert kv_cache.resize_calls == [] + + def test_an_answer_from_an_earlier_pass_is_not_covered(self): + _, manager, req, kv_cache = self._asked(grow_ok=False) + manager.begin_connector_admission_pass() + assert manager.cover_admission_offer(req) + assert kv_cache.resize_calls == [] + + def test_the_target_matches_the_prepare_time_grow(self): + """Extra KV tokens and, on a final chunk, draft tokens: the prepare-time resize must not grow again.""" + connector = FakeConnectorManager(num_matched=192, load_async=True) + connector.admission_query_capable = True + manager = make_manager(connector, num_extra_kv_tokens=3) + req = FakeRequest() + req.py_draft_tokens = [0, 0] + req.context_chunk_size = 64 + kv_cache = FakeKvCache(committed=0, capacity=67) + manager.kv_cache_map[req.py_request_id] = kv_cache + manager.begin_connector_admission_pass() + assert manager.query_connector_at_admission(req) + + assert manager.cover_admission_offer(req) + # Offer 192 shifts the 64-token chunk to the prompt end: 256 + 3 extra + 2 draft. + assert kv_cache.capacity == PROMPT_LEN + 5 + + manager._run_kv_connector_hooks(scheduled(req)) + assert [capacity for capacity, _ in kv_cache.resize_calls] == [PROMPT_LEN + 5] * 2 + assert req.context_current_position == 192 and req.context_chunk_size == 64 + + def test_a_revert_undoes_the_admission_grow(self): + _, manager, req, kv_cache = self._asked() + assert manager.cover_admission_offer(req) + assert manager.revert_allocate_context(req) + assert kv_cache.capacity == 64 + + @pytest.mark.parametrize("grow_ok", [True, False]) + def test_swa_scratch_reuse_leaves_the_shift_to_prepare(self, grow_ok): + """A grow past the cursor would be replayed exactly as a pre-grown chunk, so it is not made.""" + connector, manager, req, kv_cache = self._asked(grow_ok=grow_ok, scratch_reuse=True) + + assert manager.cover_admission_offer(req) + assert kv_cache.resize_calls == [] + assert req.py_ctx_pre_resize_cap is None + + manager._run_kv_connector_hooks(scheduled(req)) + + if grow_ok: + # As in TestChunkedBeyondTheChunk: the window shifts at prepare time. + assert (req.context_current_position, req.context_chunk_size) == (128, 64) + assert kv_cache.capacity == 192 + assert connector.commits == [(0, 128, True)] + else: + # The grow fails at prepare time and the capped fallback applies. + assert (req.context_current_position, req.context_chunk_size) == (32, 32) + assert kv_cache.capacity == 64 + assert connector.commits == [(0, 32, True)] + + def test_disagg_metadata_build_keeps_pending_context_reservation() -> None: manager, connector, req, cache = TestPrefixReservations.prepare() assert manager.resize_context(req, 32) diff --git a/tests/unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py b/tests/unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py index 59746a8a969f..4ed7dba83390 100644 --- a/tests/unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py +++ b/tests/unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py @@ -29,11 +29,15 @@ import tensorrt_llm.bindings from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_layout import valid_page_slots from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 -from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, SamplingConfig +from tensorrt_llm._torch.pyexecutor.llm_request import ( + LlmRequest, + SamplingConfig, + rewind_context_after_cache_drop, +) from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.llmapi.llm_args import KvCacheConfig from tensorrt_llm.mapping import Mapping -from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX, AttnLifeCycle +from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX, AttnLifeCycle, ReuseScope DataType = tensorrt_llm.bindings.DataType CacheType = tensorrt_llm.bindings.internal.batch_manager.CacheType @@ -56,6 +60,8 @@ class FakeConnectorManager: def __init__(self, num_matched=OFFER_TOKENS, load_async=False): self.prefix_reservations_enabled = False self.capacity_only = False + # Each admission-time query test turns it on. + self.admission_query_capable = False self.reservations = {} self.reservation_requests = {} self.next_reservation_id = 1 @@ -191,6 +197,21 @@ def schedule(manager, request, num_tokens=None): return manager.resize_context(request, num_tokens) +def schedule_reserving(manager, request, num_tokens=None): + """One scheduling pass that first reserves completion headroom, as the V2 scheduler does. + + A deferred first chunk is rolled back like ``release_context_admission``. + """ + assert manager.prepare_context(request) + if not manager.reserve_context_completion(request): + manager.free_resources(request) + rewind_context_after_cache_drop(request, manager.tokens_per_block) + return False + if num_tokens is None: + num_tokens = request.context_remaining_length + return manager.resize_context(request, num_tokens) + + def run(manager, *requests): """One ``prepare_resources``, i.e. the requests reached the final batch.""" batch = ScheduledRequests() @@ -378,6 +399,113 @@ def test_freeing_the_allocation_makes_the_request_askable_again(manager, connect assert len(connector.queries) == 2 +def test_a_dropped_context_allocation_commits_nothing_and_is_asked_again(manager, connector): + """A restart after a failed connector load leaves no trace of the lost prefix. + + The served range was skipped, never computed. Dropping the allocation must + neither publish it to the radix tree, where every later request would reuse + unfilled pages, nor keep the ask memo, which would leave the restarted + request without a connector query. + """ + request = make_request() + assert schedule(manager, request) + run(manager, request) + assert request.context_current_position == OFFER_TOKENS + + manager.drop_context_allocation(request) + + assert request.py_request_id not in manager.kv_cache_map + assert connector.forgotten == [request.py_request_id] + assert request.context_current_position == 0 + assert request.context_chunk_size == PROMPT_LEN + assert not request.py_connector_allocation_reported + assert request.py_served_prefix_end == 0 + + # The same prompt finds nothing to reuse locally. + twin = make_request(request_id=2) + assert manager.probe_context_reuse(twin) == 0 + assert manager.probe_context_reuse(request) == 0 + + assert schedule(manager, request) + run(manager, request) + + assert request.context_current_position == OFFER_TOKENS + assert [query for query in connector.queries if query[0] == request.py_request_id] == [ + (request.py_request_id, 0), + (request.py_request_id, 0), + ] + assert connector.commits == [ + (request.py_request_id, OFFER_TOKENS, False), + (request.py_request_id, OFFER_TOKENS, False), + ] + + +@pytest.mark.parametrize( + ("requires", "drops", "expected_local_end", "expected_backoff"), + [ + # A one-token backoff leaves a match of four whole blocks one short. + (False, False, 4 * TOKENS_PER_BLOCK - 1, 1), + # Aligned down: the connector serves the trimmed block and the rest. + (True, False, 3 * TOKENS_PER_BLOCK, 1), + # The backoff is dropped, so the whole local match is used. + (True, True, 4 * TOKENS_PER_BLOCK, 0), + ], + ids=["no_declaration", "align", "off"], +) +def test_whole_block_connector_is_queried_from_a_block_boundary( + monkeypatch, requires, drops, expected_local_end, expected_backoff +): + """The local match a whole-block connector is asked from, against the real tree. + + A one-model draft's lookahead backoff is stood in for by patching the + lookahead the manager reads at construction; the connector offers nothing, + so the query records the local match alone. + """ + monkeypatch.setattr( + "tensorrt_llm._torch.speculative.draft_prompt_lookahead", lambda _spec_config: 1 + ) + torch.cuda.init() + gc.collect() + torch.cuda.empty_cache() + connector = FakeConnectorManager(num_matched=0) + connector.worker = SimpleNamespace( + requires_block_aligned_reuse=requires, drops_reuse_match_backoff=drops + ) + manager = make_manager( + connector, + # Partial reuse keeps matches token-granular, so alignment needs it off. + kv_cache_config=KvCacheConfig( + max_tokens=2048, enable_block_reuse=True, enable_partial_reuse=False + ), + ) + try: + assert manager.reuse_match_backoff == expected_backoff + + # Seed four whole blocks under the scope the manager derives for + # `make_request`, which carries no LoRA task and no cache salt. + seed = manager.impl.create_kv_cache( + ReuseScope(lora_id=None, salt=manager._derive_reuse_salt(None)) + ) + assert seed.resume(manager._stream.cuda_stream) + assert seed.resize(4 * TOKENS_PER_BLOCK) + seed.commit(list(range(4 * TOKENS_PER_BLOCK))) + seed.close() + + request = make_request(prompt_len=6 * TOKENS_PER_BLOCK) + assert manager.probe_context_reuse(request) == expected_local_end + + # The connector is asked in `prepare_resources`, not while scheduling. + assert schedule(manager, request) + assert connector.queries == [] + run(manager, request) + assert connector.queries[-1] == (request.py_request_id, expected_local_end) + finally: + manager.shutdown() + del manager + gc.collect() + torch.cuda.empty_cache() + + # --------------------------------------------------------------------------- # Variable sliding-window attention. # @@ -783,8 +911,6 @@ def test_token_budget_rejection_releases_real_cache( def test_reserved_prefix_tracks_range_and_allocation( manager: KVCacheManagerV2, connector: FakeConnectorManager ) -> None: - from tensorrt_llm._torch.pyexecutor.llm_request import rewind_context_after_cache_drop - connector.prefix_reservations_enabled = True connector.num_matched = PROMPT_LEN request = make_request() @@ -890,3 +1016,232 @@ def test_real_kv_pressure_rejects_reserved_prefix_without_transmission() -> None if blocker is not None: blocker.close() manager.shutdown() + + +# -- disaggregated completion reservation ----------------------------------------- + + +@pytest.mark.parametrize("enable_block_reuse", [False, True]) +@pytest.mark.parametrize("reserve_completion", [False, True]) +def test_connector_prefill_admission_preserves_full_restore_under_pressure( + reserve_completion: bool, enable_block_reuse: bool +) -> None: + """Two first chunks fit, but both stored prefixes do not fit at once. + + The control, whose worker does not declare ``requires_completion_reservation``, keeps + main's admission and shows a truncated restore against real device pools. The reservation + must defer the second request before the connector is asked, keep room for a short + request, and honor the full offer once the first request releases its KV. + """ + connector = FakeConnectorManager() + connector.worker = SimpleNamespace(requires_completion_reservation=reserve_completion) + manager = make_manager( + connector, + kv_cache_config=KvCacheConfig( + enable_block_reuse=enable_block_reuse, + host_cache_size=0, + max_gpu_total_bytes=4 << 20, + max_util_for_resume=1.0, + ), + max_seq_len=65536, + max_num_tokens=TOKENS_PER_BLOCK, + is_disagg=True, + ) + try: + assert bool(manager._disagg_completion_pool_counts) is reserve_completion + total_blocks = manager.get_kv_cache_stats().max_num_blocks + prompt_len = (total_blocks // 2 + 1) * TOKENS_PER_BLOCK + assert 3 * TOKENS_PER_BLOCK < prompt_len < manager.max_seq_len + connector.num_matched = prompt_len - TOKENS_PER_BLOCK + first = make_request(1001, prompt_len) + second = make_request(1002, prompt_len) + short = make_request(1003, TOKENS_PER_BLOCK) + + assert schedule_reserving(manager, first, TOKENS_PER_BLOCK) + # Reserving future growth must not allocate it. + assert manager.kv_cache_map[first.py_request_id].capacity < prompt_len + assert connector.queries == [] + assert schedule_reserving(manager, second, TOKENS_PER_BLOCK) is not reserve_completion + if reserve_completion: + assert second.py_request_id not in manager.kv_cache_map + assert second.py_request_id not in manager._disagg_completion_limits + assert schedule_reserving(manager, short, TOKENS_PER_BLOCK) + run(manager, first) + else: + run(manager, first, second) + assert connector.commits[0] == (first.py_request_id, connector.num_matched, False) + if not reserve_completion: + assert connector.commits[1][1] < connector.num_matched + return + + assert connector.queries == [(first.py_request_id, 0)] + manager.free_resources(first) + assert first.py_request_id not in manager._disagg_completion_limits + assert schedule_reserving(manager, second, TOKENS_PER_BLOCK) + run(manager, second) + assert connector.commits[-1] == (second.py_request_id, connector.num_matched, False) + assert second.context_current_position == connector.num_matched + slots = manager.get_page_indices_by_layer_group(second) + assert all(index != BAD_PAGE_INDEX for group in slots for index in group) + manager.free_resources(second) + manager.free_resources(short) + assert not manager._disagg_completion_limits + assert manager.get_kv_cache_stats().free_num_blocks == total_blocks + finally: + manager.shutdown() + + +def test_refuted_connector_chunk_releases_unused_completion_reservation() -> None: + """Returning a first chunk to the queue must not retain a hidden quota.""" + connector = FakeConnectorManager() + connector.worker = SimpleNamespace(requires_completion_reservation=True) + manager = make_manager( + connector, + kv_cache_config=KvCacheConfig(max_tokens=2048, enable_block_reuse=True, host_cache_size=0), + is_disagg=True, + ) + request = make_request(1004) + try: + assert schedule_reserving(manager, request, TOKENS_PER_BLOCK) + assert request.py_request_id in manager._disagg_completion_limits + assert manager.revert_allocate_context(request) + assert request.py_request_id not in manager._disagg_completion_limits + assert connector.queries == [] + assert schedule_reserving(manager, request, TOKENS_PER_BLOCK) + run(manager, request) + assert connector.commits == [(request.py_request_id, OFFER_TOKENS, False)] + manager.free_resources(request) + assert not manager._disagg_completion_limits + finally: + manager.shutdown() + + +# -- admission-time connector queries ---------------------------------------------- + + +def _admission_scheduler(manager: KVCacheManagerV2, **kwargs): + """A real KVCacheV2Scheduler that admits one-block context chunks.""" + from tensorrt_llm._torch.pyexecutor.scheduler.scheduler_v2 import KVCacheV2Scheduler + from tensorrt_llm.llmapi.llm_args import CapacitySchedulerPolicy, ContextChunkingPolicy + + return KVCacheV2Scheduler( + max_batch_size=4, + max_num_tokens=TOKENS_PER_BLOCK, + kv_cache_manager=manager, + scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION, + ctx_chunk_config=(ContextChunkingPolicy.FIRST_COME_FIRST_SERVED, TOKENS_PER_BLOCK), + **kwargs, + ) + + +def test_admission_books_a_shifting_offer_on_real_pages( + manager: KVCacheManagerV2, connector: FakeConnectorManager +) -> None: + """The window an offer moves the chunk to is allocated while the request is admitted.""" + connector.num_matched = 64 + connector.admission_query_capable = True + request = make_request() + scheduler = _admission_scheduler(manager, connector_admission_query=True) + + output = scheduler.schedule_request([request], set()) + + assert [req.py_request_id for req in output.context_requests] == [request.py_request_id] + kv_cache = manager.kv_cache_map[request.py_request_id] + assert kv_cache.capacity >= 64 + TOKENS_PER_BLOCK + # Only the capacity moved: the served prefix is placed in prepare_resources. + assert kv_cache.history_length == 0 + assert connector.queries == [(request.py_request_id, 0)] + assert scheduler.uncharged_park_ids == frozenset() + + run(manager, request) + + # The admission-time answer was used, not asked again. + assert connector.queries == [(request.py_request_id, 0)] + assert request.context_current_position == 64 + assert request.context_chunk_size == TOKENS_PER_BLOCK + assert kv_cache.history_length == 64 + assert connector.commits == [(request.py_request_id, 64, False)] + + +def test_admission_deferral_releases_real_pages() -> None: + """A first chunk whose offer cannot be covered gives its real pages back and is asked afresh.""" + connector = FakeConnectorManager(num_matched=64) + connector.admission_query_capable = True + torch.cuda.init() + gc.collect() + torch.cuda.empty_cache() + manager = make_manager( + connector, + # The request resumes into a nearly full pool. + kv_cache_config=KvCacheConfig( + max_tokens=2048, enable_block_reuse=True, max_util_for_resume=1.0 + ), + ) + blocker = None + try: + # Leave exactly one free block: the one-block chunk fits, the window + # its offer moves it to does not. + available_blocks = manager.get_num_free_blocks() + blocker = manager.impl.create_kv_cache() + assert blocker.resume(manager._stream.cuda_stream) + for num_blocks in range(1, available_blocks + 2): + if not blocker.resize(num_blocks * TOKENS_PER_BLOCK): + break + else: + pytest.fail("The blocker did not exhaust the resolved KV pool") + held_blocks = num_blocks - 1 + assert held_blocks > 1 + assert blocker.resize((held_blocks - 1) * TOKENS_PER_BLOCK) + + request = make_request() + scheduler = _admission_scheduler( + manager, connector_admission_query=True, max_connector_offer_deferrals=1 + ) + + output = scheduler.schedule_request([request], set()) + + assert output.context_requests == [] + assert request.py_request_id not in manager.kv_cache_map + assert connector.forgotten == [request.py_request_id] + assert request.py_connector_offer_deferrals == 1 + assert request.context_current_position == 0 + assert connector.queries == [(request.py_request_id, 0)] + + # At the bound the request is admitted with the pages its chunk holds. + output = scheduler.schedule_request([request], set()) + + assert [req.py_request_id for req in output.context_requests] == [request.py_request_id] + assert request.py_request_id in manager.kv_cache_map + run(manager, request) + assert connector.queries == [(request.py_request_id, 0)] * 2 + # The shifted window cannot grow and the capped fallback cannot advance. + assert connector.commits == [(request.py_request_id, 0, False)] + assert request.context_current_position == 0 + manager.free_resources(request) + finally: + if blocker is not None: + blocker.close() + manager.shutdown() + del manager + gc.collect() + torch.cuda.empty_cache() + + +def test_admission_off_keeps_the_final_batch_query( + manager: KVCacheManagerV2, connector: FakeConnectorManager +) -> None: + """Without kv_connector_config.admission_query a capable connector is asked in prepare_resources only.""" + connector.num_matched = 64 + connector.admission_query_capable = True + request = make_request() + scheduler = _admission_scheduler(manager) + + output = scheduler.schedule_request([request], set()) + + assert [req.py_request_id for req in output.context_requests] == [request.py_request_id] + assert connector.queries == [] + + run(manager, request) + + assert connector.queries == [(request.py_request_id, 0)] + assert connector.commits == [(request.py_request_id, 64, False)] diff --git a/tests/unittest/_torch/executor/test_mooncake_store_common.py b/tests/unittest/_torch/executor/test_mooncake_store_common.py index a97103f80032..5a32a07bd489 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_common.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_common.py @@ -22,17 +22,21 @@ in-process fake that records what it was handed. """ +import hashlib import json from types import SimpleNamespace import pytest from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import config as config_module +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import master as master_module from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import staging as staging_module from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.config import ( + DEFAULT_STAGING_BUFFER_BYTES, MooncakeStoreConnectorConfig, StoreRole, local_address, + parse_size, split_address, ) from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.keys import ( @@ -40,6 +44,7 @@ KeyNamespace, ReuseScope, ) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.master import PoolManifest from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.staging import ( HostStagingPool, describe_batch_for_get, @@ -47,6 +52,11 @@ stage_batch_for_put, unstage_batch_after_get, ) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.validation import ( + NODE_BUDGET_RESERVE_BYTES, + validate_node_budget, +) +from tensorrt_llm.llmapi.llm_args import MooncakeStoreConfig pytestmark = pytest.mark.cpu_only @@ -203,6 +213,233 @@ def test_hash_chain_rejects_shrinking_token_list(): chain.extend(list(range(TOKENS_PER_BLOCK))) +def _reference_chain(scope: ReuseScope, tokens, tokens_per_block: int): + """A frozen copy of the per-token encoding stored keys were first written with. + + Each block is blake2b-128 over its parent and its tokens as signed 8-byte + little-endian ids, starting from the scope's seed. Kept here, rather than + read from the module, so that a change to how the chain encodes tokens + cannot also change what it is checked against. + """ + + def digest(*parts: bytes) -> bytes: + hasher = hashlib.blake2b(digest_size=16) + for part in parts: + hasher.update(part) + return hasher.digest() + + parent = scope.seed() + hashes = [] + for start in range(0, len(tokens) - tokens_per_block + 1, tokens_per_block): + block = tokens[start : start + tokens_per_block] + parent = digest(parent, b"".join(int(t).to_bytes(8, "little", signed=True) for t in block)) + hashes.append(parent) + return hashes + + +@pytest.mark.parametrize( + "scope", + [ + ReuseScope(), + ReuseScope(cache_salt="tenant-a"), + ReuseScope(multimodal_digests=(b"\x01" * 32,)), + ], + ids=["unscoped", "cache_salt", "multimodal"], +) +def test_hash_chain_keeps_the_stored_key_encoding(scope): + """Pages already in a pool stay addressable however the chain encodes tokens. + + The token list covers both ends of int64 and -1, and is fed in pieces that + end mid-block, so incremental extension is checked against the same + reference as a single pass. + """ + tokens = [7, -1, 2**63 - 1, -(2**63), 0, 151_000] * 5 + list(range(23)) + expected = _reference_chain(scope, tokens, TOKENS_PER_BLOCK) + assert len(expected) == len(tokens) // TOKENS_PER_BLOCK + + chain = BlockHashChain(TOKENS_PER_BLOCK, scope=scope) + for end in (5, 6, 17, len(tokens)): + chain.extend(tokens[:end]) + assert list(chain.hashes) == expected + # Extending again with no new block leaves the chain as it is. + assert list(chain.extend(tuple(tokens))) == expected + # A tuple encodes like a list; a fresh chain makes it encode every block. + from_tuple = BlockHashChain(TOKENS_PER_BLOCK, scope=scope) + assert list(from_tuple.extend(tuple(tokens))) == expected + + +def test_hash_chain_golden_digest(): + """Pinned keys, captured from the per-token encoding before it was replaced. + + A change here changes every key a deployment has stored, so pages written + by an older build stop being found. Regenerating these from the code under + test would prove nothing. + """ + assert [h.hex() for h in BlockHashChain(4).extend(list(range(8)))] == [ + "2e8c598e5152642e1a12a2345cedf3e3", + "4c62122a866d675f3f436e12723d692a", + ] + salted = BlockHashChain(2, scope=ReuseScope(cache_salt="tenant-a")) + assert [h.hex() for h in salted.extend([-1, 2**63 - 1, -(2**63), 0])] == [ + "4308a783baf926ba62be9f96df5db867", + "3205a431e8c69e7077da3e22ce5b56da", + ] + multimodal = BlockHashChain(3, scope=ReuseScope(multimodal_digests=(b"\x01" * 32,))) + assert [h.hex() for h in multimodal.extend([151_000, 7, -5, 2**40, 0, 1, 99])] == [ + "9e3d925110ea5d2fc057a7874a546f11", + "29d8726fadccb94fc9beacdde9042f0a", + ] + + +@pytest.mark.parametrize("token", [2**63, -(2**63) - 1], ids=["above", "below"]) +def test_hash_chain_rejects_tokens_outside_int64(token): + """A token id no 8-byte key can encode is refused, and leaves the chain as it was.""" + chain = BlockHashChain(TOKENS_PER_BLOCK) + chain.extend(list(range(TOKENS_PER_BLOCK))) + before = list(chain.hashes) + tokens = list(range(TOKENS_PER_BLOCK)) + [0] * (TOKENS_PER_BLOCK - 1) + [token] + with pytest.raises(OverflowError): + chain.extend(tokens) + assert list(chain.hashes) == before + + +@pytest.mark.parametrize("lookahead", [1, TOKENS_PER_BLOCK + 1]) +def test_hash_chain_requires_lookahead_and_matches_incremental_extension(lookahead): + """A block is hashed once the draft's lookahead past it is in the prompt too.""" + tokens = list(range(4 * TOKENS_PER_BLOCK + lookahead)) + chain = BlockHashChain(TOKENS_PER_BLOCK, prompt_lookahead=lookahead) + assert chain.prompt_lookahead == lookahead + assert list(chain.extend(tokens[: TOKENS_PER_BLOCK + lookahead - 1])) == [] + assert len(chain.extend(tokens[: TOKENS_PER_BLOCK + lookahead])) == 1 + for end in range(TOKENS_PER_BLOCK + lookahead + 1, len(tokens) + 1): + chain.extend(tokens[:end]) + assert len(chain.hashes) == 4 + one_shot = BlockHashChain(TOKENS_PER_BLOCK, prompt_lookahead=lookahead).extend(tokens) + assert list(chain.hashes) == list(one_shot) + + # The last lookahead token of block 0 is part of its identity; the token + # after it is not. + covered = list(tokens) + covered[TOKENS_PER_BLOCK + lookahead - 1] += 1000 + assert ( + chain.hashes[0] + != BlockHashChain(TOKENS_PER_BLOCK, prompt_lookahead=lookahead).extend(covered)[0] + ) + beyond = list(tokens) + beyond[TOKENS_PER_BLOCK + lookahead] += 1000 + assert ( + chain.hashes[0] + == BlockHashChain(TOKENS_PER_BLOCK, prompt_lookahead=lookahead).extend(beyond)[0] + ) + + +def test_hash_chain_rejects_negative_lookahead(): + with pytest.raises(ValueError, match="prompt_lookahead"): + BlockHashChain(TOKENS_PER_BLOCK, prompt_lookahead=-1) + + +def test_hash_chain_lookahead_separates_keys(): + """Servers whose drafts read different lookaheads never share a page. + + Nor does a speculative server share one with a server without a draft, and + the request scope still separates keys under a lookahead. + """ + tokens = list(range(3 * TOKENS_PER_BLOCK)) + first_blocks = { + BlockHashChain(TOKENS_PER_BLOCK, prompt_lookahead=lookahead).extend(tokens)[0] + for lookahead in (0, 1, 2) + } + assert len(first_blocks) == 3 + + scopes = [ + ReuseScope(), + ReuseScope(cache_salt="tenant-a"), + ReuseScope(cache_salt="tenant-b"), + ReuseScope(multimodal_digests=(b"\x01" * 32,)), + ReuseScope(multimodal_digests=(b"\x02" * 32,)), + ] + scoped = { + BlockHashChain(TOKENS_PER_BLOCK, scope=scope, prompt_lookahead=1).extend(tokens)[0] + for scope in scopes + } + assert len(scoped) == len(scopes) + + +def _hand_built_chain(lookahead, cache_salt, multimodal_digests, tokens, tokens_per_block): + """Expected key bytes, built from the specification rather than from `ReuseScope`. + + The seed is blake2b-128 over the length-framed, presence-tagged salt and + the multimodal digests. A non-zero lookahead re-digests it with + `b"prompt-lookahead"` and the lookahead as 8 little-endian bytes. Each block + is then blake2b-128 over its parent and its tokens plus the lookahead tokens + after it, as signed 8-byte little-endian ids, at a one-block stride. + """ + + def digest(*parts: bytes) -> bytes: + hasher = hashlib.blake2b(digest_size=16) + for part in parts: + hasher.update(part) + return hasher.digest() + + salt_bytes = b"\x00" if cache_salt is None else b"\x01" + cache_salt.encode("utf-8") + framed = b"".join( + len(part).to_bytes(4, "little") + part for part in (salt_bytes, *multimodal_digests) + ) + parent = digest(framed) + if lookahead: + parent = digest(parent, b"prompt-lookahead", lookahead.to_bytes(8, "little")) + hashes = [] + for start in range(0, len(tokens) - lookahead - tokens_per_block + 1, tokens_per_block): + window = tokens[start : start + tokens_per_block + lookahead] + parent = digest(parent, b"".join(t.to_bytes(8, "little", signed=True) for t in window)) + hashes.append(parent) + return hashes + + +@pytest.mark.parametrize( + "lookahead, cache_salt, multimodal_digests", + [ + (0, None, ()), + (0, "tenant-a", ()), + (0, None, (b"\x01" * 32,)), + (1, None, ()), + (3, None, ()), + (1, "tenant-a", ()), + (3, None, (b"\x01" * 32,)), + ], + ids=[ + "no_lookahead", + "no_lookahead_salted", + "no_lookahead_multimodal", + "lookahead_1", + "lookahead_3", + "lookahead_1_salted", + "lookahead_3_multimodal", + ], +) +def test_hash_chain_key_bytes_follow_the_framed_scope_and_lookahead( + lookahead, cache_salt, multimodal_digests +): + """Pins the bytes of every stored key, with and without a draft lookahead. + + Lookahead 0 pins the keys pools already hold. The token list covers both + ends of int64 and -1, and is fed in pieces that end mid-block. + """ + tokens = [7, -1, 2**63 - 1, -(2**63), 0, 151_000] * 5 + list(range(23)) + expected = _hand_built_chain( + lookahead, cache_salt, multimodal_digests, tokens, TOKENS_PER_BLOCK + ) + assert len(expected) == (len(tokens) - lookahead) // TOKENS_PER_BLOCK + + scope = ReuseScope(cache_salt=cache_salt, multimodal_digests=multimodal_digests) + chain = BlockHashChain(TOKENS_PER_BLOCK, scope=scope, prompt_lookahead=lookahead) + for end in (5, 6, 17, len(tokens)): + chain.extend(tokens[:end]) + assert list(chain.hashes) == expected + one_shot = BlockHashChain(TOKENS_PER_BLOCK, scope=scope, prompt_lookahead=lookahead) + assert list(one_shot.extend(tuple(tokens))) == expected + + def test_key_namespace_separates_every_dimension(): base = dict( namespace="trtllm", @@ -227,6 +464,35 @@ def test_key_namespace_separates_every_dimension(): assert KeyNamespace(**{**base, field: value}).key(block_hash) != reference +def test_key_namespace_group_ignores_rank_and_layer_group(): + """Every page a hit on a block needs shares one eviction group.""" + base = dict( + namespace="trtllm", + model_key="m", + rank=0, + world_size=2, + layer_group_id=0, + tokens_per_block=4, + bytes_per_page=64, + ) + block = b"\x01" * 16 + group = KeyNamespace(**base).group(block) + + for field, value in [ + ("rank", 1), + ("layer_group_id", 3), + ("tokens_per_block", 8), + ("bytes_per_page", 128), + ]: + assert KeyNamespace(**{**base, field: value}).group(block) == group + for field, value in [("world_size", 4), ("namespace", "other"), ("model_key", "n")]: + assert KeyNamespace(**{**base, field: value}).group(block) != group + assert KeyNamespace(**base).group(b"\x02" * 16) != group + # The key still names one shard's page. + assert KeyNamespace(**base).key(block) != KeyNamespace(**{**base, "rank": 1}).key(block) + assert group == f"trtllm/m/w2/{block.hex()}" + + # ---- config ---- @@ -366,6 +632,207 @@ def test_config_requires_an_explicit_model_key(store_config, tmp_path, monkeypat config.resolve_model_key("/models/MiniMax-M3/") +# ---- pool eviction and diagnostics settings ---- + + +def test_config_eviction_settings_default_off_and_read_the_json(store_config): + config = resolve_config() + assert config.soft_pin_ttl_ms is None + assert config.group_blocks is False + + set_pool_setting(store_config, soft_pin_ttl_ms=86400000, group_blocks=True) + config = resolve_config() + assert config.soft_pin_ttl_ms == 86400000 + assert config.group_blocks is True + + # An explicit null is the same as leaving the TTL out. + set_pool_setting(store_config, soft_pin_ttl_ms=None) + assert resolve_config().soft_pin_ttl_ms is None + + +@pytest.mark.parametrize("ttl", [0, -1]) +def test_config_rejects_a_soft_pin_ttl_that_is_not_positive(store_config, ttl): + set_pool_setting(store_config, soft_pin_ttl_ms=ttl) + with pytest.raises(ValueError, match="soft_pin_ttl_ms"): + resolve_config() + + +@pytest.mark.parametrize("value", ["yes", "false", 1]) +def test_config_group_blocks_must_be_a_json_boolean(store_config, value): + """`bool("false")` is True, so anything but a JSON boolean is refused.""" + set_pool_setting(store_config, group_blocks=value) + with pytest.raises(ValueError, match="group_blocks"): + resolve_config() + + +def test_config_stats_interval_defaults_off_and_reads_the_json(store_config): + assert resolve_config().stats_interval is None + + set_pool_setting(store_config, stats_interval=30) + assert resolve_config().stats_interval == 30.0 + + set_pool_setting(store_config, stats_interval=None) + assert resolve_config().stats_interval is None + + +@pytest.mark.parametrize("interval", [0, -1, float("inf")]) +def test_config_rejects_a_stats_interval_that_is_not_a_positive_number(store_config, interval): + set_pool_setting(store_config, stats_interval=interval) + with pytest.raises(ValueError, match="stats_interval"): + resolve_config() + + +def test_config_has_no_environment_override_for_eviction_settings(store_config, monkeypatch): + """The source branch read these from the environment; this connector reads the JSON only.""" + monkeypatch.setenv("TRTLLM_MOONCAKE_STORE_SOFT_PIN_TTL_MS", "86400000") + monkeypatch.setenv("TRTLLM_MOONCAKE_STORE_GROUP_BLOCKS", "1") + monkeypatch.setenv("TRTLLM_MOONCAKE_HOLE_STATS", "1") + monkeypatch.setenv("TLLM_M3_LOAD_PHASE_TIMING", "1") + monkeypatch.setenv("TLLM_MOONCAKE_LOAD_PHASE_TIMING", "1") + + config = resolve_config() + + assert config.soft_pin_ttl_ms is None + assert config.group_blocks is False + assert config.stats_interval is None + + +def test_config_async_load_defaults_off_and_reads_the_json(store_config): + assert resolve_config().async_load is False + set_pool_setting(store_config, async_load=True) + assert resolve_config().async_load is True + # `bool("false")` is True, so only a JSON boolean is read. + set_pool_setting(store_config, async_load="false") + with pytest.raises(ValueError, match="async_load"): + resolve_config() + + +def test_config_async_load_workers_defaults_to_one_and_rejects_zero(store_config): + assert resolve_config().async_load_workers == 1 + set_pool_setting(store_config, async_load_workers=3) + assert resolve_config().async_load_workers == 3 + set_pool_setting(store_config, async_load_workers=0) + with pytest.raises(ValueError, match="async_load_workers"): + resolve_config() + + +def test_config_staging_buffer_bytes_defaults_and_parses_binary_units(store_config): + """A config without the key, such as vLLM's, stages what an unset field asks.""" + default = MooncakeStoreConfig.model_fields["staging_buffer_bytes"].default + assert resolve_config().staging_buffer_bytes == DEFAULT_STAGING_BUFFER_BYTES + assert DEFAULT_STAGING_BUFFER_BYTES == parse_size(default, strict_units=True) + + set_pool_setting(store_config, staging_buffer_bytes="512MiB") + assert resolve_config().staging_buffer_bytes == 512 * 1024**2 + + set_pool_setting(store_config, staging_buffer_bytes="512MB") + with pytest.raises(ValueError, match="ambiguous unit"): + resolve_config() + set_pool_setting(store_config, staging_buffer_bytes=0) + with pytest.raises(ValueError, match="staging_buffer_bytes"): + resolve_config() + + +_GIB = 1 << 30 + + +def _budget_config(**overrides) -> MooncakeStoreConnectorConfig: + settings = dict( + master_server_address="10.0.0.7:50051", + global_segment_size=4 * _GIB, + stage_through_host=True, + staging_buffer_bytes=_GIB, + ) + settings.update(overrides) + return MooncakeStoreConnectorConfig(**settings) + + +def _assert_node_claim(config, *, ranks, claimed): + """`claimed` bytes fit exactly, and one byte less fails with the remedy text.""" + validate_node_budget( + config, ranks_on_node=ranks, available_bytes=claimed + NODE_BUDGET_RESERVE_BYTES + ) + with pytest.raises(ValueError, match="cannot afford") as excinfo: + validate_node_budget( + config, ranks_on_node=ranks, available_bytes=claimed + NODE_BUDGET_RESERVE_BYTES - 1 + ) + return str(excinfo.value) + + +def test_validate_node_budget_counts_staging_once_per_direction(): + """Each direction a role drives pins `staging_buffer_bytes` per rank, and no more.""" + # Two ranks, each lending 4 GiB and staging 1 GiB for loads and 1 GiB for saves. + message = _assert_node_claim(_budget_config(), ranks=2, claimed=2 * 4 * _GIB + 2 * 2 * _GIB) + assert "lower staging_buffer_bytes" in message + # The load threads split the load direction's budget rather than adding to it. + _assert_node_claim( + _budget_config(async_load=True, async_load_workers=4), + ranks=2, + claimed=2 * 4 * _GIB + 2 * 2 * _GIB, + ) + _assert_node_claim( + _budget_config(role=StoreRole.CONSUMER), ranks=2, claimed=2 * 4 * _GIB + 2 * _GIB + ) + _assert_node_claim( + _budget_config(staging_buffer_bytes=512 * 1024**2), ranks=1, claimed=4 * _GIB + _GIB + ) + # A capacity role transfers nothing, so it stages nothing either. + _assert_node_claim(_budget_config(role=StoreRole.CAPACITY), ranks=2, claimed=2 * 4 * _GIB) + message = _assert_node_claim( + _budget_config(stage_through_host=False), ranks=2, claimed=2 * 4 * _GIB + ) + assert "staging_buffer_bytes" not in message + + +def test_client_config_renders_the_async_and_staging_settings(tmp_path): + rendered, config = _rendered_client_config( + tmp_path, async_load=True, async_load_workers=3, staging_buffer_bytes="256MiB" + ) + assert rendered["async_load"] is True + assert rendered["async_load_workers"] == 3 + # Resolved to a byte count, as the segment size is. + assert rendered["staging_buffer_bytes"] == 256 * 1024**2 + assert config.async_load is True + assert config.async_load_workers == 3 + assert config.staging_buffer_bytes == 256 * 1024**2 + + rendered, config = _rendered_client_config(tmp_path) + assert rendered["async_load"] is False + assert rendered["async_load_workers"] == 1 + assert rendered["staging_buffer_bytes"] == DEFAULT_STAGING_BUFFER_BYTES + assert config.async_load is False + assert config.staging_buffer_bytes == DEFAULT_STAGING_BUFFER_BYTES + + +def _rendered_client_config(tmp_path, **settings): + """What `provision_pool` writes for a block with `settings`, and that file read back.""" + pool = MooncakeStoreConfig(pool="h:1", model_key="m", **settings) + rendered = master_module._client_config(pool, PoolManifest(master_server_address="h:1"), "") + path = tmp_path / "rendered.json" + path.write_text(json.dumps(rendered)) + return rendered, MooncakeStoreConnectorConfig.from_file(str(path)) + + +def test_client_config_renders_eviction_and_stats_settings(tmp_path): + rendered, config = _rendered_client_config( + tmp_path, soft_pin_ttl_ms=86400000, group_blocks=True, stats_interval=30 + ) + assert rendered["soft_pin_ttl_ms"] == 86400000 + assert rendered["group_blocks"] is True + assert rendered["stats_interval"] == 30 + assert config.soft_pin_ttl_ms == 86400000 + assert config.group_blocks is True + assert config.stats_interval == 30.0 + + rendered, config = _rendered_client_config(tmp_path) + assert "soft_pin_ttl_ms" not in rendered + assert "stats_interval" not in rendered + assert rendered["group_blocks"] is False + assert config.soft_pin_ttl_ms is None + assert config.group_blocks is False + assert config.stats_interval is None + + # The connector itself, the registry preset that selects it, and the startup # gates on parallelism and index-V caching land with the KV cache manager V2 # support they depend on. @@ -493,6 +960,12 @@ def test_plan_slot_geometry_rejects_degenerate_inputs(bad): plan_slot_geometry(*bad) +def test_plan_slot_geometry_needs_the_configured_budget(): + """The budget is `staging_buffer_bytes` or a load thread's share, never a default.""" + with pytest.raises(TypeError): + plan_slot_geometry(1024, 8) + + def test_staging_pool_registers_its_whole_buffer_once(): """One registration covering every slot, since the store addresses bytes.""" store = FakeStore() diff --git a/tests/unittest/_torch/executor/test_mooncake_store_connector.py b/tests/unittest/_torch/executor/test_mooncake_store_connector.py index 95cdf73dee6a..f4450555081f 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_connector.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_connector.py @@ -27,15 +27,18 @@ import contextlib import json import os +import sys import threading import time from collections.abc import Iterator -from types import SimpleNamespace +from enum import Enum +from types import ModuleType, SimpleNamespace import pytest import torch from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import ( + KvCacheConnectorManager, RequestData, SchedulerOutput, ) @@ -61,6 +64,7 @@ ) from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.keys import BlockHashChain from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.metadata import ( + MooncakeStoreMetadata, PageTransfer, RequestTransfers, ) @@ -70,6 +74,12 @@ from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.worker import ( MooncakeStoreConnectorWorker, ) +from tensorrt_llm.llmapi.llm_args import ( + Eagle3DecodingConfig, + MTPDecodingConfig, + NGramDecodingConfig, + SADecodingConfig, +) from tensorrt_llm.runtime.kv_cache_manager_v2 import BAD_PAGE_INDEX pytestmark = pytest.mark.cpu_only @@ -92,6 +102,8 @@ def __init__(self): #: Index of the `register_buffer` call to fail, counting from 0. self.fail_register_at = None self.put_calls = [] + #: The arguments of every put past the three buffers, as `(args, kwargs)`. + self.put_options = [] self.get_calls = [] self.exist_calls = [] self.closed = False @@ -117,8 +129,9 @@ def batch_is_exist(self, keys): self.exist_calls.append(list(keys)) return [1 if key in self.objects else 0 for key in keys] - def batch_put_from_multi_buffers(self, keys, addresses, sizes, *_args, **_kwargs): + def batch_put_from_multi_buffers(self, keys, addresses, sizes, *args, **kwargs): self.put_calls.append((list(keys), [list(a) for a in addresses], [list(s) for s in sizes])) + self.put_options.append((args, kwargs)) self.objects.update(keys) return [sum(size) for size in sizes] @@ -286,19 +299,29 @@ def set_pool_setting(store_config, **settings): store_config.write_text(json.dumps(raw)) -def make_llm_args(*, run_dir: str | None = None) -> SimpleNamespace: +def make_llm_args( + *, + run_dir: str | None = None, + enable_attention_dp: bool = False, + local_reuse_backoff: str = "off", + speculative_config=None, +) -> SimpleNamespace: return SimpleNamespace( model="/models/test-model", kv_cache_config=SimpleNamespace(tokens_per_block=TOKENS_PER_BLOCK), tensor_parallel_size=1, + enable_attention_dp=enable_attention_dp, pipeline_parallel_size=1, context_parallel_size=1, sparse_attention_config=None, + speculative_config=speculative_config, # How a rank reaches settings the process that rendered the client # config could not hand it; see config.pool_config. kv_connector_config=SimpleNamespace( connector="mooncake-store", - mooncake_store=SimpleNamespace(run_dir=run_dir), + mooncake_store=SimpleNamespace( + run_dir=run_dir, local_reuse_backoff=local_reuse_backoff + ), ), ) @@ -324,6 +347,9 @@ def make_worker( *, layout: KvCacheLayout | None = None, run_dir: str | None = None, + enable_attention_dp: bool = False, + local_reuse_backoff: str = "off", + speculative_config=None, ) -> Iterator[MooncakeStoreConnectorWorker]: """Build a worker and shut it down before the test call phase ends. @@ -331,7 +357,14 @@ def make_worker( pytest-threadleak snapshots threads around the call phase only, so fixture teardown would run too late to keep it quiet. """ - worker = MooncakeStoreConnectorWorker(make_llm_args(run_dir=run_dir)) + worker = MooncakeStoreConnectorWorker( + make_llm_args( + run_dir=run_dir, + enable_attention_dp=enable_attention_dp, + local_reuse_backoff=local_reuse_backoff, + speculative_config=speculative_config, + ) + ) fake_store.workers.append(worker) if layout is not None: worker.register_kv_cache_layout(layout) @@ -369,6 +402,9 @@ def make_request( multimodal_positions=multimodal_positions, multimodal_lengths=multimodal_lengths, get_tokens=lambda _beam=0, _tokens=tuple(tokens): list(_tokens), + # What the runtime honoured of the connector's offer, set before it + # reports the allocation; `settle_offer` assigns it. + py_num_connector_matched_tokens=0, ) @@ -789,6 +825,14 @@ def test_format_diagnosis_reports_the_value_mooncake_read( assert "nvidia_peermem is not loaded, so every range fails" in report +# ---- worker declarations ---- + + +def test_worker_declares_completion_reservation(): + """A disaggregated server keeps completion headroom only for a connector that asks.""" + assert MooncakeStoreConnectorWorker.requires_completion_reservation is True + + # ---- worker registration ---- @@ -923,13 +967,162 @@ def test_worker_prefix_hit_stops_at_the_first_gap(store_config, fake_store): assert worker.count_prefix_hit(hashes) == 1 -def test_worker_load_raises_when_a_page_is_missing(store_config, fake_store): +def test_worker_declares_failed_load_reporting(): + """The runtime restarts a request whose page went missing only if the worker says so.""" + assert MooncakeStoreConnectorWorker.reports_failed_loads is True + # The hooks the manager requires of a reporting worker are implemented. + KvCacheConnectorManager._validate_failed_load_reporting(MooncakeStoreConnectorWorker) + + +def test_worker_load_reports_the_request_when_a_page_is_missing(store_config, fake_store): + """A single-rank worker hands a request whose page is gone back to the executor.""" + with make_worker(fake_store, layout=make_layout()) as worker: + assert worker._recover_failed_loads + transfers = RequestTransfers(7, [PageTransfer(b"\x00" * 16, 0, 1)]) + worker.bind_allocation_fence(AllocationFence(fake_store)) + worker.bind_connector_meta(SimpleNamespace(loads=[transfers], saves=[])) + worker.start_load_kv(None) + assert worker.take_failed_load_requests() == {7} + # Cleared on return: the next iteration starts clean. + assert worker.take_failed_load_requests() == set() + + +def test_worker_load_raises_when_a_page_is_missing_under_tensor_parallelism( + store_config, fake_store +): + """Tensor-parallel ranks hold shards of one page and cannot restart a request alone. + + Set on the built worker rather than through the world size, which would + also change every key namespace. + """ with make_worker(fake_store, layout=make_layout()) as worker: + worker._recover_failed_loads = False transfers = RequestTransfers(7, [PageTransfer(b"\x00" * 16, 0, 1)]) worker.bind_allocation_fence(AllocationFence(fake_store)) worker.bind_connector_meta(SimpleNamespace(loads=[transfers], saves=[])) with pytest.raises(RuntimeError, match="already"): worker.start_load_kv(None) + assert worker.take_failed_load_requests() == set() + + +def test_worker_load_skips_the_rest_of_a_failed_request_and_keeps_the_others( + store_config, fake_store +): + """Only the request that lost a page restarts; its later pages are not fetched.""" + with make_worker(fake_store, layout=make_layout()) as worker: + present, missing, later = (bytes([1]) * 16, bytes([2]) * 16, bytes([3]) * 16) + for block_hash in (present, later): + fake_store.objects.add(worker._namespaces[0].key(block_hash)) + # The first store call holds request 7's page and request 8's missing + # one, so request 8's second page would be the second call. + worker._batch_size = 2 + loads = [ + RequestTransfers(7, [PageTransfer(present, 0, 1)]), + RequestTransfers(8, [PageTransfer(missing, 0, 2), PageTransfer(later, 0, 3)]), + ] + worker.bind_allocation_fence(AllocationFence(fake_store)) + worker.bind_connector_meta(SimpleNamespace(loads=loads, saves=[])) + worker.start_load_kv(None) + + assert worker.take_failed_load_requests() == {8} + assert len(fake_store.get_calls) == 1 + fetched = [key for keys, _addresses, _sizes in fake_store.get_calls for key in keys] + assert worker._namespaces[0].key(present) in fetched + assert worker._namespaces[0].key(later) not in fetched + + +def test_worker_load_counts_a_short_answer_as_every_page_failed( + store_config, fake_store, monkeypatch +): + """A store answer of the wrong length cannot say which pages landed. + + Both pages exist, but the store answers for the first one only. Reading + that as one landed page would let request 7 compute over a page nothing + confirmed, so every request in the call restarts. + """ + with make_worker(fake_store, layout=make_layout()) as worker: + first, second = bytes([1]) * 16, bytes([2]) * 16 + for block_hash in (first, second): + fake_store.objects.add(worker._namespaces[0].key(block_hash)) + answer = fake_store.batch_get_into_multi_buffers + monkeypatch.setattr( + fake_store, + "batch_get_into_multi_buffers", + lambda keys, addresses, sizes: answer(keys, addresses, sizes)[:1], + ) + worker._batch_size = 2 + loads = [ + RequestTransfers(7, [PageTransfer(first, 0, 1)]), + RequestTransfers(8, [PageTransfer(second, 0, 2)]), + ] + worker.bind_allocation_fence(AllocationFence(fake_store)) + worker.bind_connector_meta(SimpleNamespace(loads=loads, saves=[])) + worker.start_load_kv(None) + + assert len(fake_store.get_calls) == 1 + assert worker.take_failed_load_requests() == {7, 8} + + +def test_worker_failed_load_drops_that_requests_saves_from_the_pass(store_config, fake_store): + """The request that lost a page must not publish its scheduled tail. + + `build_connector_meta` bound the saves of every scheduled request before + the load ran. The failed request leaves the batch without computing + anything, so only the surviving request's saves may reach `wait_for_save`. + """ + with make_worker(fake_store, layout=make_layout()) as worker: + present, missing = bytes([1]) * 16, bytes([2]) * 16 + tail_ok, tail_failed = bytes([3]) * 16, bytes([4]) * 16 + fake_store.objects.add(worker._namespaces[0].key(present)) + metadata = MooncakeStoreMetadata( + loads=[ + RequestTransfers(7, [PageTransfer(present, 0, 1)]), + RequestTransfers(9, [PageTransfer(missing, 0, 2)]), + ], + saves=[ + RequestTransfers(7, [PageTransfer(tail_ok, 0, 3)]), + RequestTransfers(9, [PageTransfer(tail_failed, 0, 4)]), + ], + ) + saves = metadata.saves + worker.bind_allocation_fence(AllocationFence(fake_store)) + worker.bind_connector_meta(metadata) + worker.start_load_kv(None) + + assert worker.take_failed_load_requests() == {9} + # What `wait_for_save` will hand to the save thread after the pass. + assert [transfers.request_id for transfers in worker.get_connector_meta().saves] == [7] + # Edited in place, since that is the object `wait_for_save` reads. + assert worker.get_connector_meta().saves is saves + assert 9 not in worker._outstanding_saves + # The manager drops the same request again before the restart. + assert worker.drop_bound_saves([9]) == 0 + assert [transfers.request_id for transfers in saves] == [7] + + +def test_worker_drops_bound_saves_of_a_restarted_request_even_when_it_was_the_whole_batch( + store_config, fake_store +): + """A pass whose only request restarted publishes nothing.""" + with make_worker(fake_store, layout=make_layout()) as worker: + missing, tail = bytes([2]) * 16, bytes([4]) * 16 + metadata = MooncakeStoreMetadata( + loads=[RequestTransfers(8, [PageTransfer(missing, 0, 2)])], + saves=[RequestTransfers(8, [PageTransfer(tail, 0, 4)])], + ) + worker.bind_allocation_fence(AllocationFence(fake_store)) + worker.bind_connector_meta(metadata) + worker.start_load_kv(None) + + assert worker.take_failed_load_requests() == {8} + assert metadata.saves == [] + # With nothing bound to save, the post-pass handoff queues nothing. + worker.wait_for_save(None) + assert worker._save_queue.empty() + assert not worker._outstanding_saves + assert worker.take_submitted_saves() == set() + # A second drop is harmless, as is dropping an unknown request. + assert worker.drop_bound_saves([8, 9]) == 0 def test_worker_zero_copy_load_refuses_to_run_unordered(store_config, fake_store): @@ -1019,220 +1212,934 @@ def refuse(_keys): assert worker.get_finished([42], []) == ([42], []) -# ---- host staging ---- +def test_worker_reports_the_saves_it_queued(store_config, fake_store, monkeypatch): + """Only saves `wait_for_save` handed to the save thread are reported, once.""" + class FakeEvent: + def record(self, stream): + pass -@contextlib.contextmanager -def make_staged_worker(fake_store, store_config, *, layout, budget=None): - """A worker configured to pass pages through pinned host slots. + def synchronize(self): + pass - `budget` patches the pinned-memory ceiling rather than setting a field: - the allocation follows from the layout and the transfer batch now, so the - ceiling is the only thing left that can bind. - """ - set_pool_setting(store_config, stage_through_host=True) - with contextlib.ExitStack() as stack: - if budget is not None: - patch = stack.enter_context(pytest.MonkeyPatch.context()) - patch.setattr(worker_module, "MAX_STAGING_BUFFER_BYTES", budget) - worker = stack.enter_context(make_worker(fake_store, layout=layout)) - yield worker + monkeypatch.setattr(worker_module.torch.cuda, "Event", FakeEvent) + with make_worker(fake_store, layout=make_layout()) as worker: + # Bound but never run: nothing is reported. + worker.bind_connector_meta( + SimpleNamespace( + loads=[], saves=[RequestTransfers(5, [PageTransfer(b"\x05" * 16, 0, 1)])] + ) + ) + assert worker.take_submitted_saves() == set() + saves = [ + RequestTransfers(7, [PageTransfer(b"\x07" * 16, 0, 2)]), + RequestTransfers(8, [PageTransfer(b"\x08" * 16, 0, 3)]), + ] + metadata = SimpleNamespace(loads=[], saves=list(saves)) + worker.bind_connector_meta(metadata) + # A request dropped from the bound saves before the pass runs is not + # queued, so it is not reported either. + metadata.saves[:] = saves[:1] + worker.wait_for_save(None) + assert worker.take_submitted_saves() == {7} + # Cleared on return. + assert worker.take_submitted_saves() == set() -@pytest.fixture -def staged_copies(monkeypatch): - """Record the copies staging would issue, instead of running them.""" - copies = [] - monkeypatch.setattr( - staging_module, - "_memcpy_async", - lambda dst, src, size, kind, stream: copies.append((int(dst), int(src), int(size))), - ) - # Imported into the worker by name, so replace the worker's binding. - monkeypatch.setattr(worker_module, "_sync_stream", lambda _stream: None) - return copies +# ---- soft-pinned and grouped puts ---- -class StubStagingPool: - """Stands in for `HostStagingPool` where no device backs its slots. - The real pool page-locks its buffer whenever a device is reported, which - is the one claim `fake_cuda` cannot make true. - """ +def _save_two_pages(worker): + hashes = [bytes([index]) * 16 for index in range(2)] + worker._put( + [RequestTransfers(1, [PageTransfer(hashes[0], 0, 0), PageTransfer(hashes[1], 0, 1)])] + ) - def __init__(self, *, slot_bytes, num_slots, store, label): - self.slot_bytes = slot_bytes - self.num_slots = num_slots - self.label = label - def close(self) -> None: - pass +def _record_put_configs(monkeypatch): + """Replace the `ReplicateConfig` factory; each call returns a namespace of its arguments.""" + built = [] + def fake_config(soft_pin_ttl_ms, group_ids=None): + if soft_pin_ttl_ms is None and group_ids is None: + return None + config = SimpleNamespace( + soft_pin_ttl_ms=soft_pin_ttl_ms, + group_ids=None if group_ids is None else list(group_ids), + ) + built.append(config) + return config -@pytest.fixture -def fake_cuda(monkeypatch): - """Present a CUDA device on a host that has none, recording set_device calls. + monkeypatch.setattr(worker_module, "_replicate_config", fake_config) + return built - The device bookkeeping around the save thread is control flow, so testing - it should not need a GPU. The two things that path does need one for, the - staging slots registration pins and the stream the thread opens over them, - are stood in for, so claiming a device cannot reach a real allocation. - """ - recorded = [] - monkeypatch.setattr(torch.cuda, "is_available", lambda: True) - monkeypatch.setattr(torch.cuda, "current_device", lambda: 3) - monkeypatch.setattr(torch.cuda, "set_device", lambda index: recorded.append(index)) - monkeypatch.setattr(torch.cuda, "Stream", lambda *_args, **_kwargs: SimpleNamespace()) - monkeypatch.setattr(worker_module, "HostStagingPool", StubStagingPool) - return recorded +def test_worker_save_passes_no_config_by_default(store_config, fake_store): + """Without either option the put keeps its plain three-argument form.""" + with make_worker(fake_store, layout=make_layout()) as worker: + assert worker._put_config is None + _save_two_pages(worker) + assert fake_store.put_options == [((), {})] -def test_staging_put_hands_the_store_one_host_buffer_per_page( - store_config, fake_store, staged_copies -): - layout = make_layout(regions_per_group=3) - with make_staged_worker(fake_store, store_config, layout=layout) as worker: - block_hash = b"\x07" * 16 - worker._put([RequestTransfers(1, [PageTransfer(block_hash, 0, 2)])]) - device_addresses, device_sizes = PageAddressing(layout).buffers(0, 2) - slot = worker._save_staging.slot_address(0) - (keys, addresses, sizes) = fake_store.put_calls[0] +def test_worker_save_soft_pins_every_page_when_configured(store_config, fake_store, monkeypatch): + """The one config is built once and handed to every put of every batch.""" + built = _record_put_configs(monkeypatch) + set_pool_setting(store_config, soft_pin_ttl_ms=86400000, transfer_batch_size=1) + with make_worker(fake_store, layout=make_layout()) as worker: + _save_two_pages(worker) + assert len(built) == 1 + assert built[0].soft_pin_ttl_ms == 86400000 + assert built[0].group_ids is None + assert len(fake_store.put_calls) == 2 + assert fake_store.put_options == [((built[0],), {}), ((built[0],), {})] - assert keys == [worker._namespaces[0].key(block_hash)] - # The store sees one contiguous host buffer whose length is the sum of - # the device regions. That equality is what keeps a staged write - # byte-identical to a zero-copy one. - assert addresses == [[slot]] - assert sizes == [[sum(device_sizes)]] - # Every region was gathered, in order, into its place in the slot. - expected = [] - offset = 0 - for address, size in zip(device_addresses, device_sizes): - expected.append((slot + offset, address, size)) - offset += size - assert staged_copies == expected +@pytest.mark.parametrize("role", ["consumer", "capacity"]) +def test_worker_builds_no_put_config_for_a_role_that_does_not_write( + store_config, fake_store, monkeypatch, role +): + """Write options cost nothing, and need no particular binding, where nothing is written.""" + built = _record_put_configs(monkeypatch) + set_pool_setting(store_config, role=role, soft_pin_ttl_ms=86400000, group_blocks=True) + with make_worker(fake_store, layout=make_layout()) as worker: + assert worker._put_config is None + assert worker._soft_pin_ttl_ms is None + assert worker._group_blocks is False + assert built == [] -def test_staging_get_scatters_back_to_the_device_regions(store_config, fake_store, staged_copies): - layout = make_layout(regions_per_group=3) - with make_staged_worker(fake_store, store_config, layout=layout) as worker: - block_hash = b"\x03" * 16 - fake_store.objects.add(worker._namespaces[0].key(block_hash)) +def test_worker_groups_every_page_of_a_block(store_config, fake_store, monkeypatch): + """Each put names one group per key. - worker.bind_connector_meta( - SimpleNamespace(loads=[RequestTransfers(7, [PageTransfer(block_hash, 0, 5)])], saves=[]) + The group is shared across a block's layer groups, differs across blocks, + and lines up with the pages left after the existence filter. + """ + built = _record_put_configs(monkeypatch) + set_pool_setting(store_config, soft_pin_ttl_ms=5000, group_blocks=True) + with make_worker(fake_store, layout=make_layout(num_groups=2)) as worker: + hashes = [bytes([index]) * 16 for index in range(2)] + # Block 0's layer-group-1 page is already stored, so only three are put. + fake_store.objects.add(worker._namespaces[1].key(hashes[0])) + worker._put( + [ + RequestTransfers( + 1, + [ + PageTransfer(hashes[0], 0, 0), + PageTransfer(hashes[0], 1, 0), + PageTransfer(hashes[1], 0, 1), + PageTransfer(hashes[1], 1, 1), + ], + ) + ] ) - worker.start_load_kv(None) - - device_addresses, device_sizes = PageAddressing(layout).buffers(0, 5) - slot = worker._load_staging.slot_address(0) - (keys, addresses, sizes) = fake_store.get_calls[0] - assert keys == [worker._namespaces[0].key(block_hash)] - assert addresses == [[slot]] - assert sizes == [[sum(device_sizes)]] + group0 = worker._namespaces[0].group(hashes[0]) + group1 = worker._namespaces[0].group(hashes[1]) + assert worker._namespaces[1].group(hashes[0]) == group0 + assert group0 != group1 + expected_keys = [ + worker._namespaces[0].key(hashes[0]), + worker._namespaces[0].key(hashes[1]), + worker._namespaces[1].key(hashes[1]), + ] + ((keys, _, _),) = fake_store.put_calls + assert keys == expected_keys + ((args, kwargs),) = fake_store.put_options + assert kwargs == {} + (config,) = args + assert config.group_ids == [group0, group1, group1] + assert config.soft_pin_ttl_ms == 5000 + # The construction-time probe, then the one put. + assert [entry.group_ids for entry in built] == [["probe"], [group0, group1, group1]] + + +def test_worker_warns_once_when_the_pool_refuses_every_soft_pinned_put( + store_config, fake_store, monkeypatch +): + """A TTL above the master's ceiling fails every put, which is worth naming.""" + _record_put_configs(monkeypatch) + set_pool_setting(store_config, soft_pin_ttl_ms=90000000) + monkeypatch.setattr( + fake_store, + "batch_put_from_multi_buffers", + lambda keys, addresses, sizes, *args: [-1] * len(keys), + ) + refused = [] + monkeypatch.setattr( + worker_module.logger, + "warning_once", + lambda *message, key: refused.append((" ".join(message), key)), + ) + with make_worker(fake_store, layout=make_layout()) as worker: + _save_two_pages(worker) + ((message, key),) = refused + assert "--max_kv_soft_pin_ttl" in message + assert key == "mooncake-store-soft-pin-refused" - # The scatter is the mirror of the gather: same split, opposite direction. - expected = [] - offset = 0 - for address, size in zip(device_addresses, device_sizes): - expected.append((address, slot + offset, size)) - offset += size - assert staged_copies == expected +def test_worker_does_not_blame_soft_pin_for_a_partial_failure( + store_config, fake_store, monkeypatch +): + _record_put_configs(monkeypatch) + set_pool_setting(store_config, soft_pin_ttl_ms=5000) + monkeypatch.setattr( + fake_store, + "batch_put_from_multi_buffers", + lambda keys, addresses, sizes, *args: [-1] + [0] * (len(keys) - 1), + ) + refused = [] + monkeypatch.setattr( + worker_module.logger, "warning_once", lambda *message, key: refused.append(key) + ) + with make_worker(fake_store, layout=make_layout()) as worker: + _save_two_pages(worker) + assert refused == [] -def test_staging_does_not_scatter_a_failed_load(store_config, fake_store, staged_copies): - layout = make_layout() - with make_staged_worker(fake_store, store_config, layout=layout) as worker: - block_hash = b"\x09" * 16 - key = worker._namespaces[0].key(block_hash) - fake_store.objects.add(key) - fake_store.fail_gets_for.add(key) - worker.bind_connector_meta( - SimpleNamespace(loads=[RequestTransfers(7, [PageTransfer(block_hash, 0, 1)])], saves=[]) - ) - with pytest.raises(RuntimeError, match="failed to load"): - worker.start_load_kv(None) +def _fake_mooncake_store(monkeypatch, **attributes): + """Install a `mooncake.store` module carrying only `attributes`.""" + package = ModuleType("mooncake") + store = ModuleType("mooncake.store") + for name, value in attributes.items(): + setattr(store, name, value) + package.store = store + monkeypatch.setitem(sys.modules, "mooncake", package) + monkeypatch.setitem(sys.modules, "mooncake.store", store) - # A failed read leaves the slot holding whatever it held before, and - # copying that onto the page would put unrelated bytes where the - # runtime already promised computed KV. - assert staged_copies == [] +class _SoftPinAction(Enum): + PRESERVE = 0 + ENABLE = 1 + DISABLE = 2 -def test_staging_load_leaves_the_ordering_to_the_stream(store_config, fake_store, staged_copies): - """The scatter is queued on the stream, behind the previous pass.""" - with make_staged_worker(fake_store, store_config, layout=make_layout()) as worker: - fence = AllocationFence(fake_store) - worker.bind_allocation_fence(fence) - bind_one_load(worker, fake_store, b"\x05" * 16) - worker.start_load_kv(None) - assert fence.reads_before_wait == [] - assert len(staged_copies) == 1 +class _ReplicateConfig: + def __init__(self): + self.soft_pin_action = _SoftPinAction.PRESERVE + self.soft_pin_ttl_ms = None + self.group_ids = None -def test_the_ranks_device_is_captured_and_adopted_by_the_save_thread( - store_config, fake_store, fake_cuda -): - """The save thread must not run on torch's default device. - It issues staging copies against pointers owned by the rank's device. A - stream created on device 0 instead fails every copy with - cudaErrorInvalidValue, and only on ranks other than 0. - """ - with make_worker(fake_store, layout=make_layout()) as worker: - assert worker._device_index == 3 +def test_replicate_config_is_none_without_options(): + assert worker_module._replicate_config(None) is None + assert worker_module._replicate_config(None, None) is None - deadline = time.monotonic() + 5.0 - while 3 not in fake_cuda and time.monotonic() < deadline: - time.sleep(0.01) - assert 3 in fake_cuda, f"save thread set devices {fake_cuda}, expected the rank's 3" - # Never the thread-local default. - assert 0 not in fake_cuda +def test_replicate_config_sets_soft_pin_and_groups(monkeypatch): + _fake_mooncake_store( + monkeypatch, ReplicateConfig=_ReplicateConfig, SoftPinAction=_SoftPinAction + ) + config = worker_module._replicate_config(86400000) + assert isinstance(config, _ReplicateConfig) + assert config.soft_pin_action is _SoftPinAction.ENABLE + assert config.soft_pin_ttl_ms == 86400000 + assert config.group_ids is None -def test_a_save_thread_that_cannot_start_fails_registration( - store_config, fake_store, fake_cuda, monkeypatch -): - """A worker whose save thread died must not go on to report itself ready. + config = worker_module._replicate_config(None, ("a", "b")) + assert config.soft_pin_action is _SoftPinAction.PRESERVE + assert config.group_ids == ["a", "b"] - The thread's setup binds this rank's device. If that raises and the thread - exits, every later `wait_for_save` counts a save that nothing will ever - consume, and the requests holding those pages stay pinned for good. - """ - def refuse(_index): - raise RuntimeError("cudaSetDevice failed") +def test_replicate_config_names_the_binding_it_needs(monkeypatch): + _fake_mooncake_store(monkeypatch, ReplicateConfig=_ReplicateConfig) + with pytest.raises(ImportError, match="SoftPinAction"): + worker_module._replicate_config(1000) - monkeypatch.setattr(torch.cuda, "set_device", refuse) + class OldReplicateConfig: + __slots__ = ("soft_pin_action", "soft_pin_ttl_ms") - with pytest.raises(RuntimeError, match="save thread"): - with make_worker(fake_store, layout=make_layout()): - pytest.fail("registration should not have completed") + _fake_mooncake_store( + monkeypatch, ReplicateConfig=OldReplicateConfig, SoftPinAction=_SoftPinAction + ) + with pytest.raises(ImportError, match="group_ids"): + worker_module._replicate_config(None, ["a"]) -def test_shutdown_leaves_the_store_open_under_a_save_still_reading( - store_config, fake_store, staged_copies, monkeypatch -): - """A timed join is not evidence that the thread stopped. +# ---- diagnostics ---- - The save thread reads the KV pools through the store handle and the staging - slots, so closing the handle or dropping those buffers mid-transfer takes - the memory out from under it. - """ - monkeypatch.setattr(worker_module, "SAVE_DRAIN_TIMEOUT", 0.1) - reading = threading.Event() - release = threading.Event() - store_put = fake_store.batch_put_from_multi_buffers - def blocking_put(*args, **kwargs): - reading.set() - assert release.wait(30.0), "the test never released the save thread" - return store_put(*args, **kwargs) +def _store_blocks(fake_store, worker, hashes, blocks): + for block in blocks: + fake_store.objects.add(worker._namespaces[0].key(hashes[block])) - fake_store.batch_put_from_multi_buffers = blocking_put + +def test_worker_counts_a_single_block_hole(store_config, fake_store): + set_pool_setting(store_config, stats_interval=60) + with make_worker(fake_store, layout=make_layout()) as worker: + hashes = [bytes([index]) * 16 for index in range(5)] + # Block 2 was evicted from the middle of a chain the pool still holds. + _store_blocks(fake_store, worker, hashes, [0, 1, 3, 4]) + + # The hit is unchanged: it stops at the hole's index. + assert worker.count_prefix_hit(hashes) == 2 + + stats = worker._hole_stats + assert (stats.lookups, stats.lookup_blocks, stats.hit_blocks) == (1, 5, 2) + assert (stats.holes, stats.hole_blocks) == (1, 1) + assert stats.holed_tokens == 2 * TOKENS_PER_BLOCK + assert stats.max_holed_tokens == 2 * TOKENS_PER_BLOCK + # No extra store traffic: holes come from the one existence probe. + assert len(fake_store.exist_calls) == 1 + + # Not due yet, so nothing is logged and nothing is reset. + assert worker._maybe_log_hole_stats() is None + line = worker._maybe_log_hole_stats(force=True) + assert line is not None + assert "holes=1" in line + assert f"holed_tokens={2 * TOKENS_PER_BLOCK}" in line + assert stats.lookups == 0 + + +def test_worker_hole_stats_ignore_a_clean_prefix(store_config, fake_store): + set_pool_setting(store_config, stats_interval=60) + with make_worker(fake_store, layout=make_layout()) as worker: + hashes = [bytes([index]) * 16 for index in range(4)] + _store_blocks(fake_store, worker, hashes, [0, 1]) + + assert worker.count_prefix_hit(hashes) == 2 + assert worker._hole_stats.lookups == 1 + assert worker._hole_stats.holes == 0 + assert worker._hole_stats.holed_tokens == 0 + + +def test_worker_hole_stats_log_once_per_interval(store_config, fake_store, monkeypatch): + set_pool_setting(store_config, stats_interval=60) + with make_worker(fake_store, layout=make_layout()) as worker: + hashes = [bytes([index]) * 16 for index in range(3)] + _store_blocks(fake_store, worker, hashes, [0, 2]) + lines = [] + monkeypatch.setattr(worker_module.logger, "info", lines.append) + + worker.count_prefix_hit(hashes) + assert lines == [] + worker._hole_stats_last -= worker._stats_interval + worker.count_prefix_hit(hashes) + assert len(lines) == 1 + assert "prefix-hole stats" in lines[0] + assert "lookups=2" in lines[0] and "holes=2" in lines[0] + + # Shutdown flushes what accumulated since the last line. + worker.count_prefix_hit(hashes) + worker.shutdown() + assert len(lines) == 2 + assert "lookups=1" in lines[1] + + +def test_worker_hole_stats_off_by_default(store_config, fake_store): + with make_worker(fake_store, layout=make_layout()) as worker: + hashes = [bytes([index]) * 16 for index in range(5)] + _store_blocks(fake_store, worker, hashes, [0, 1, 3, 4]) + + assert worker.count_prefix_hit(hashes) == 2 + assert worker._hole_stats is None + assert worker._load_stats is None + assert worker._maybe_log_hole_stats(force=True) is None + assert worker._maybe_log_load_stats(force=True) is None + assert len(fake_store.exist_calls) == 1 + + +def test_worker_that_never_reads_the_pool_counts_nothing(store_config, fake_store): + """Only a role that loads looks prefixes up, so a producer has nothing to count.""" + set_pool_setting(store_config, role="producer", stats_interval=60) + with make_worker(fake_store, layout=make_layout()) as worker: + assert worker._hole_stats is None + assert worker._load_stats is None + + +def _load_stats_lines(lines): + return [line for line in lines if "mooncake-store load stats" in line] + + +def test_worker_load_stats_time_each_phase(store_config, fake_store, monkeypatch): + """A zero-copy load waits on the fence and has nothing to scatter.""" + set_pool_setting(store_config, stats_interval=60) + lines = [] + monkeypatch.setattr(worker_module.logger, "info", lines.append) + with make_worker(fake_store, layout=make_layout()) as worker: + # Nothing loaded yet, so even a forced line has nothing to say. + assert worker._maybe_log_load_stats(force=True) is None + + worker.bind_allocation_fence(AllocationFence(fake_store)) + bind_one_load(worker, fake_store, b"\x04" * 16) + worker.start_load_kv(None) + + stats = worker._load_stats + for phase in ("resolve", "fence", "describe", "get"): + assert stats.phase_wall[phase] > 0, phase + assert stats.phase_wall["unstage"] == stats.phase_wall["sync"] == 0 + + # Not due yet. + assert worker.get_finished([], []) == ([], []) + assert _load_stats_lines(lines) == [] + + worker._load_stats_last -= worker._stats_interval + worker.get_finished([], []) + (line,) = _load_stats_lines(lines) + assert "loads=1 pages=1 batches=1 " in line + for phase in worker_module._LOAD_PHASES: + assert f" {phase}=" in line + # The counters start over, so a quiet interval logs nothing. + assert stats.loads == 0 + worker._load_stats_last -= worker._stats_interval + worker.get_finished([], []) + assert len(_load_stats_lines(lines)) == 1 + + +def test_worker_load_stats_time_the_staged_scatter(store_config, fake_store, staged_copies): + """A staged load has no fence to wait on, and scatters and syncs every batch.""" + set_pool_setting(store_config, stats_interval=60, transfer_batch_size=1) + with make_staged_worker(fake_store, store_config, layout=make_layout()) as worker: + hashes = [bytes([index]) * 16 for index in range(2)] + _store_blocks(fake_store, worker, hashes, [0, 1]) + worker.bind_connector_meta( + SimpleNamespace( + loads=[ + RequestTransfers( + 7, [PageTransfer(hashes[0], 0, 1), PageTransfer(hashes[1], 0, 2)] + ) + ], + saves=[], + ) + ) + worker.start_load_kv(None) + + stats = worker._load_stats + assert (stats.loads, stats.pages, stats.batches) == (1, 2, 2) + assert stats.loaded_bytes == 2 * 64 + for phase in ("resolve", "describe", "get", "unstage", "sync"): + assert stats.phase_wall[phase] > 0, phase + assert stats.phase_wall["fence"] == 0 + + line = worker._maybe_log_load_stats(force=True) + assert "mooncake-store load stats" in line + assert "loads=1 pages=2 batches=2 " in line + assert worker._maybe_log_load_stats(force=True) is None + + +def test_worker_load_stats_are_flushed_at_shutdown(store_config, fake_store, monkeypatch): + set_pool_setting(store_config, stats_interval=60) + lines = [] + monkeypatch.setattr(worker_module.logger, "info", lines.append) + with make_worker(fake_store, layout=make_layout()) as worker: + worker.bind_allocation_fence(AllocationFence(fake_store)) + bind_one_load(worker, fake_store, b"\x05" * 16) + worker.start_load_kv(None) + worker.shutdown() + assert len(_load_stats_lines(lines)) == 1 + # Shut down again on the way out, with nothing left to flush. + assert len(_load_stats_lines(lines)) == 1 + + +# ---- attention data parallelism ---- + + +def test_mooncake_declares_adp_support(): + assert MooncakeStoreConnectorScheduler.supports_attention_dp is True + assert MooncakeStoreConnectorWorker.supports_attention_dp is True + # The gate the executor creator runs before constructing either half. + KvCacheConnectorManager.validate_attention_dp( + MooncakeStoreConnectorWorker, MooncakeStoreConnectorScheduler, pp_size=1, cp_size=1 + ) + + +@pytest.mark.parametrize("producer_rank,consumer_rank", [(0, 3), (3, 1), (1, 0)]) +def test_adp_reuses_another_owners_pages( + store_config, fake_store, monkeypatch, producer_rank, consumer_rank +): + """A full prefix needs one copy per layer group, whichever owner saved it.""" + monkeypatch.setattr(worker_module, "mpi_world_size", lambda: 4) + monkeypatch.setattr(worker_module, "mpi_rank", lambda: producer_rank) + layout = make_layout(num_groups=3, regions_per_group=2) + hashes = BlockHashChain(TOKENS_PER_BLOCK).extend(list(range(2 * TOKENS_PER_BLOCK))) + with make_worker(fake_store, layout=layout, enable_attention_dp=True) as producer: + producer._put( + [RequestTransfers(7, [PageTransfer(hashes[0], group, 1) for group in range(3)])] + ) + # An incompletely published next block must not extend the offer. + producer._put([RequestTransfers(7, [PageTransfer(hashes[1], 0, 2)])]) + monkeypatch.setattr(worker_module, "mpi_rank", lambda: consumer_rank) + with make_worker(fake_store, layout=layout, enable_attention_dp=True) as consumer: + assert consumer._namespaces == producer._namespaces + assert all("/w1r0/" in namespace.prefix for namespace in consumer._namespaces.values()) + + assert consumer.count_prefix_hit(hashes) == 1 + # Two blocks over three layer groups: one key per page, with no + # fan-out over the other owners. + assert len(fake_store.exist_calls[-1]) == 2 * 3 + + transfers = RequestTransfers( + 19, [PageTransfer(hashes[0], group, 5) for group in range(3)] + ) + consumer.bind_allocation_fence(AllocationFence(fake_store)) + consumer.bind_connector_meta(SimpleNamespace(loads=[transfers], saves=[])) + consumer.start_load_kv(None) + assert consumer.take_failed_load_requests() == set() + # The consumer's own pages, filled under the producer's keys. + keys, addresses, sizes = fake_store.get_calls[-1] + assert keys == [producer._namespaces[group].key(hashes[0]) for group in range(3)] + addressing = PageAddressing(layout) + assert addresses == [addressing.buffers(group, 5)[0] for group in range(3)] + assert sizes == [addressing.buffers(group, 5)[1] for group in range(3)] + + # Reuse is shared, while request completion belongs to each owner. + producer._outstanding_saves[7] = 1 + assert producer.get_finished([7], []) == ([], []) + assert consumer.get_finished([7], []) == ([7], []) + producer._outstanding_saves.clear() + assert producer.get_finished([], []) == ([7], []) + + +def test_tp_lookup_still_requires_all_attention_shards(store_config, fake_store, monkeypatch): + monkeypatch.setattr(worker_module, "mpi_world_size", lambda: 4) + monkeypatch.setattr(worker_module, "mpi_rank", lambda: 0) + block_hash = BlockHashChain(TOKENS_PER_BLOCK).extend(list(range(TOKENS_PER_BLOCK)))[0] + with make_worker(fake_store, layout=make_layout(num_groups=2)) as worker: + worker._put( + [RequestTransfers(7, [PageTransfer(block_hash, group, 1) for group in range(2)])] + ) + # Rank 0's shard alone is not a prefix the other ranks can load. + assert worker.count_prefix_hit([block_hash]) == 0 + for namespaces in worker._peer_namespaces.values(): + for namespace in namespaces: + fake_store.objects.add(namespace.key(block_hash)) + assert worker.count_prefix_hit([block_hash]) == 1 + assert len(fake_store.exist_calls[-1]) == 4 * 2 + # Tensor-parallel shards are not what an attention-DP owner stores. + with make_worker( + fake_store, layout=make_layout(num_groups=2), enable_attention_dp=True + ) as owner: + assert owner.count_prefix_hit([block_hash]) == 0 + + +def test_tp_keys_are_unchanged(store_config, fake_store, monkeypatch): + """Pages written before attention-DP support must still be found.""" + block_hash = bytes(range(16)) + suffix = f"/lg0/t{TOKENS_PER_BLOCK}b64/{block_hash.hex()}" + with make_worker(fake_store, layout=make_layout()) as single: + assert single._namespaces[0].key(block_hash) == f"trtllm/test-model/w1r0{suffix}" + + monkeypatch.setattr(worker_module, "mpi_world_size", lambda: 4) + monkeypatch.setattr(worker_module, "mpi_rank", lambda: 2) + with make_worker(fake_store, layout=make_layout()) as rank_two: + assert rank_two._namespaces[0].key(block_hash) == f"trtllm/test-model/w4r2{suffix}" + assert [namespace.key(block_hash) for namespace in rank_two._peer_namespaces[0]] == [ + f"trtllm/test-model/w4r{rank}{suffix}" for rank in range(4) + ] + # The global rank still names the process in the ledger and logs. + assert (rank_two._rank, rank_two._world_size) == (2, 4) + + +def test_attention_dp_owner_restarts_a_failed_load_where_tensor_parallel_raises( + store_config, fake_store, monkeypatch +): + """An owner holds whole pages and its own batch, so it can restart a request alone.""" + monkeypatch.setattr(worker_module, "mpi_world_size", lambda: 4) + monkeypatch.setattr(worker_module, "mpi_rank", lambda: 1) + transfers = RequestTransfers(7, [PageTransfer(b"\x00" * 16, 0, 1)]) + with make_worker(fake_store, layout=make_layout(), enable_attention_dp=True) as owner: + owner.bind_allocation_fence(AllocationFence(fake_store)) + owner.bind_connector_meta(SimpleNamespace(loads=[transfers], saves=[])) + owner.start_load_kv(None) + assert owner.take_failed_load_requests() == {7} + + with make_worker(fake_store, layout=make_layout()) as tp_rank: + assert not tp_rank._recover_failed_loads + tp_rank.bind_allocation_fence(AllocationFence(fake_store)) + tp_rank.bind_connector_meta(SimpleNamespace(loads=[transfers], saves=[])) + with pytest.raises(RuntimeError, match="already"): + tp_rank.start_load_kv(None) + + +def test_capacity_only_worker_builds_under_attention_dp(store_config, fake_store, monkeypatch): + """A generation server that only lends memory may run with attention DP.""" + set_pool_setting(store_config, role="capacity") + monkeypatch.setattr(worker_module, "mpi_world_size", lambda: 2) + monkeypatch.setattr(worker_module, "mpi_rank", lambda: 1) + with make_worker(fake_store, layout=make_layout(), enable_attention_dp=True) as worker: + assert worker.capacity_only + assert (worker._attention_rank, worker._attention_world_size) == (0, 1) + # Nothing registered, no save thread, and nothing to look up. + assert worker._addressing is None + assert worker._save_thread is None + assert fake_store.registered == [] + assert worker.count_prefix_hit([b"\x01" * 16]) == 0 + assert fake_store.exist_calls == [] + + +# ---- local reuse under a one-model draft ---- + + +@pytest.mark.parametrize("mode, drops", [("off", True), ("align", False)]) +def test_worker_local_reuse_backoff_comes_from_the_config(store_config, fake_store, mode, drops): + """`mooncake_store.local_reuse_backoff` picks whole local matches or aligned ones.""" + with make_worker(fake_store, local_reuse_backoff=mode) as worker: + assert worker.drops_reuse_match_backoff is drops + # Either way an offer starts on a block boundary. + assert worker.requires_block_aligned_reuse is True + + +def test_worker_local_reuse_backoff_defaults_off_for_an_inherited_config(store_config, fake_store): + """A deployment that provisions the pool itself may set no `mooncake_store` block.""" + llm_args = make_llm_args() + llm_args.kv_connector_config.mooncake_store = None + worker = MooncakeStoreConnectorWorker(llm_args) + fake_store.workers.append(worker) + try: + assert worker.drops_reuse_match_backoff is True + assert worker.requires_block_aligned_reuse is True + finally: + worker.shutdown() + + +@pytest.mark.parametrize( + "role, declares", + [("both", True), ("consumer", True), ("producer", False), ("capacity", False)], +) +def test_mooncake_asks_for_block_aligned_local_reuse_only_where_it_loads( + store_config, fake_store, role, declares +): + """Only a role that reads the pool reshapes the local match its offer starts from. + + A producer or capacity rank never offers, so it keeps the cache manager's own + reuse, draft backoff included. + """ + set_pool_setting(store_config, role=role) + with make_worker(fake_store) as worker: + assert worker.requires_block_aligned_reuse is declares + assert worker.drops_reuse_match_backoff is declares + # Class defaults stay off, so a role that loads is the only thing that turns them on. + assert MooncakeStoreConnectorWorker.requires_block_aligned_reuse is False + assert MooncakeStoreConnectorWorker.drops_reuse_match_backoff is False + + +# ---- speculative prefill ---- + + +def _eagle3(*, separate_draft_cache: bool) -> Eagle3DecodingConfig: + config = Eagle3DecodingConfig(max_draft_len=3, speculative_model="draft") + # The executor creator turns this off on a disaggregated server, before the + # connector is built. + config._allow_separate_draft_kv_cache = separate_draft_cache + return config + + +def _mtp(*, separate_draft_cache: bool, vanilla: bool = False) -> MTPDecodingConfig: + config = MTPDecodingConfig(max_draft_len=2 if vanilla else 1, use_mtp_vanilla=vanilla) + config._allow_separate_draft_kv_cache = separate_draft_cache + return config + + +def test_mooncake_declares_speculative_prefill_support(): + assert MooncakeStoreConnectorScheduler.supports_speculative_prefill is True + assert MooncakeStoreConnectorWorker.supports_speculative_prefill is True + # The check the executor creator and the connector manager run. + assert KvCacheConnectorManager.declares_speculative_prefill( + MooncakeStoreConnectorWorker, MooncakeStoreConnectorScheduler + ) + + +@pytest.mark.parametrize( + "make_spec_config, match", + [ + (lambda: _eagle3(separate_draft_cache=True), "separate draft KV cache"), + (lambda: _mtp(separate_draft_cache=True), "separate draft KV cache"), + (lambda: SADecodingConfig(max_draft_len=4), "mode SA"), + ( + lambda: NGramDecodingConfig(max_draft_len=3, max_matching_ngram_size=2), + "mode NGRAM", + ), + ], + ids=["eagle3_separate_cache", "mtp_separate_cache", "suffix_automaton", "ngram"], +) +def test_worker_refuses_a_drafter_it_cannot_key( + store_config, fake_store, monkeypatch, make_spec_config, match +): + """Refused before a segment is mounted, since the store would serve stale draft KV.""" + opened = [] + monkeypatch.setattr( + worker_module, "_open_store", lambda config: opened.append(config) or (fake_store, "") + ) + with pytest.raises(NotImplementedError, match=match): + MooncakeStoreConnectorWorker(make_llm_args(speculative_config=make_spec_config())) + assert opened == [] + + +@pytest.mark.parametrize( + "role, make_spec_config", + [ + ("both", lambda: _eagle3(separate_draft_cache=False)), + ("both", lambda: _mtp(separate_draft_cache=False)), + ("consumer", lambda: _mtp(separate_draft_cache=False, vanilla=True)), + # A capacity role addresses no page, so its drafter does not matter. + ("capacity", lambda: _eagle3(separate_draft_cache=True)), + ("capacity", lambda: NGramDecodingConfig(max_draft_len=3, max_matching_ngram_size=2)), + ], + ids=[ + "eagle3", + "mtp_eagle", + "mtp_vanilla_consumer", + "capacity_separate_cache", + "capacity_ngram", + ], +) +def test_worker_accepts_shared_one_model_draft_and_capacity_role( + store_config, fake_store, role, make_spec_config +): + set_pool_setting(store_config, role=role) + with make_worker(fake_store, speculative_config=make_spec_config()) as worker: + assert worker.capacity_only is (role == "capacity") + assert len(fake_store.workers) == 1 + + +# ---- host staging ---- + + +@contextlib.contextmanager +def make_staged_worker(fake_store, store_config, *, layout, budget=None): + """A worker configured to pass pages through pinned host slots. + + `budget` sets `staging_buffer_bytes`, the pinned memory per direction: the + slot width follows from the layout, so the budget only bounds how many + slots a pool holds. + """ + settings = {"stage_through_host": True} + if budget is not None: + settings["staging_buffer_bytes"] = budget + set_pool_setting(store_config, **settings) + with make_worker(fake_store, layout=layout) as worker: + yield worker + + +@pytest.fixture +def staged_copies(monkeypatch): + """Record the copies staging would issue, instead of running them.""" + copies = [] + monkeypatch.setattr( + staging_module, + "_memcpy_async", + lambda dst, src, size, kind, stream: copies.append((int(dst), int(src), int(size))), + ) + # Imported into the worker by name, so replace the worker's binding. + monkeypatch.setattr(worker_module, "_sync_stream", lambda _stream: None) + return copies + + +class StubStagingPool: + """Stands in for `HostStagingPool` where no device backs its slots. + + The real pool page-locks its buffer whenever a device is reported, which + is the one claim `fake_cuda` cannot make true. + """ + + def __init__(self, *, slot_bytes, num_slots, store, label): + self.slot_bytes = slot_bytes + self.num_slots = num_slots + self.label = label + + def close(self) -> None: + pass + + +@pytest.fixture +def fake_cuda(monkeypatch): + """Present a CUDA device on a host that has none, recording set_device calls. + + The device bookkeeping around the save thread is control flow, so testing + it should not need a GPU. The two things that path does need one for, the + staging slots registration pins and the stream the thread opens over them, + are stood in for, so claiming a device cannot reach a real allocation. + """ + recorded = [] + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 3) + monkeypatch.setattr(torch.cuda, "set_device", lambda index: recorded.append(index)) + monkeypatch.setattr(torch.cuda, "Stream", lambda *_args, **_kwargs: SimpleNamespace()) + monkeypatch.setattr(worker_module, "HostStagingPool", StubStagingPool) + return recorded + + +def test_staging_put_hands_the_store_one_host_buffer_per_page( + store_config, fake_store, staged_copies +): + layout = make_layout(regions_per_group=3) + with make_staged_worker(fake_store, store_config, layout=layout) as worker: + block_hash = b"\x07" * 16 + worker._put([RequestTransfers(1, [PageTransfer(block_hash, 0, 2)])]) + + device_addresses, device_sizes = PageAddressing(layout).buffers(0, 2) + slot = worker._save_staging.slot_address(0) + (keys, addresses, sizes) = fake_store.put_calls[0] + + assert keys == [worker._namespaces[0].key(block_hash)] + # The store sees one contiguous host buffer whose length is the sum of + # the device regions. That equality is what keeps a staged write + # byte-identical to a zero-copy one. + assert addresses == [[slot]] + assert sizes == [[sum(device_sizes)]] + + # Every region was gathered, in order, into its place in the slot. + expected = [] + offset = 0 + for address, size in zip(device_addresses, device_sizes): + expected.append((slot + offset, address, size)) + offset += size + assert staged_copies == expected + + +def test_staging_get_scatters_back_to_the_device_regions(store_config, fake_store, staged_copies): + layout = make_layout(regions_per_group=3) + with make_staged_worker(fake_store, store_config, layout=layout) as worker: + block_hash = b"\x03" * 16 + fake_store.objects.add(worker._namespaces[0].key(block_hash)) + + worker.bind_connector_meta( + SimpleNamespace(loads=[RequestTransfers(7, [PageTransfer(block_hash, 0, 5)])], saves=[]) + ) + worker.start_load_kv(None) + + device_addresses, device_sizes = PageAddressing(layout).buffers(0, 5) + slot = worker._load_staging.slot_address(0) + (keys, addresses, sizes) = fake_store.get_calls[0] + assert keys == [worker._namespaces[0].key(block_hash)] + assert addresses == [[slot]] + assert sizes == [[sum(device_sizes)]] + + # The scatter is the mirror of the gather: same split, opposite direction. + expected = [] + offset = 0 + for address, size in zip(device_addresses, device_sizes): + expected.append((address, slot + offset, size)) + offset += size + assert staged_copies == expected + + +@pytest.mark.parametrize("recover", [True, False], ids=["restart", "tensor_parallel"]) +def test_staging_does_not_scatter_a_failed_load(store_config, fake_store, staged_copies, recover): + layout = make_layout() + with make_staged_worker(fake_store, store_config, layout=layout) as worker: + worker._recover_failed_loads = recover + block_hash = b"\x09" * 16 + key = worker._namespaces[0].key(block_hash) + fake_store.objects.add(key) + fake_store.fail_gets_for.add(key) + + worker.bind_connector_meta( + SimpleNamespace(loads=[RequestTransfers(7, [PageTransfer(block_hash, 0, 1)])], saves=[]) + ) + if recover: + worker.start_load_kv(None) + assert worker.take_failed_load_requests() == {7} + else: + with pytest.raises(RuntimeError, match="failed to load"): + worker.start_load_kv(None) + + # A failed read leaves the slot holding whatever it held before, and + # copying that onto the page would put unrelated bytes where the + # runtime already promised computed KV. + assert staged_copies == [] + + +def test_staging_scatters_only_the_pages_that_landed(store_config, fake_store, staged_copies): + """In a mixed batch the filled slot reaches its page and the failed one does not.""" + layout = make_layout(regions_per_group=2) + with make_staged_worker(fake_store, store_config, layout=layout) as worker: + good, bad = bytes([0x0A]) * 16, bytes([0x0B]) * 16 + fake_store.objects.add(worker._namespaces[0].key(good)) + fake_store.objects.add(worker._namespaces[0].key(bad)) + fake_store.fail_gets_for.add(worker._namespaces[0].key(bad)) + loads = [ + RequestTransfers(7, [PageTransfer(good, 0, 1)]), + RequestTransfers(8, [PageTransfer(bad, 0, 2)]), + ] + worker.bind_connector_meta(SimpleNamespace(loads=loads, saves=[])) + worker.start_load_kv(None) + + assert worker.take_failed_load_requests() == {8} + # Both pages shared one store call, so slot 0 holds the landed page. + assert len(fake_store.get_calls) == 1 + good_addresses, good_sizes = PageAddressing(layout).buffers(0, 1) + slot = worker._load_staging.slot_address(0) + expected = [] + offset = 0 + for address, size in zip(good_addresses, good_sizes): + expected.append((address, slot + offset, size)) + offset += size + assert staged_copies == expected + + +def test_staging_load_leaves_the_ordering_to_the_stream(store_config, fake_store, staged_copies): + """The scatter is queued on the stream, behind the previous pass.""" + with make_staged_worker(fake_store, store_config, layout=make_layout()) as worker: + fence = AllocationFence(fake_store) + worker.bind_allocation_fence(fence) + bind_one_load(worker, fake_store, b"\x05" * 16) + worker.start_load_kv(None) + assert fence.reads_before_wait == [] + assert len(staged_copies) == 1 + + +def test_the_ranks_device_is_captured_and_adopted_by_the_save_thread( + store_config, fake_store, fake_cuda +): + """The save thread must not run on torch's default device. + + It issues staging copies against pointers owned by the rank's device. A + stream created on device 0 instead fails every copy with + cudaErrorInvalidValue, and only on ranks other than 0. + """ + with make_worker(fake_store, layout=make_layout()) as worker: + assert worker._device_index == 3 + + deadline = time.monotonic() + 5.0 + while 3 not in fake_cuda and time.monotonic() < deadline: + time.sleep(0.01) + assert 3 in fake_cuda, f"save thread set devices {fake_cuda}, expected the rank's 3" + # Never the thread-local default. + assert 0 not in fake_cuda + + +def test_a_save_thread_that_cannot_start_fails_registration( + store_config, fake_store, fake_cuda, monkeypatch +): + """A worker whose save thread died must not go on to report itself ready. + + The thread's setup binds this rank's device. If that raises and the thread + exits, every later `wait_for_save` counts a save that nothing will ever + consume, and the requests holding those pages stay pinned for good. + """ + + def refuse(_index): + raise RuntimeError("cudaSetDevice failed") + + monkeypatch.setattr(torch.cuda, "set_device", refuse) + + with pytest.raises(RuntimeError, match="save thread"): + with make_worker(fake_store, layout=make_layout()): + pytest.fail("registration should not have completed") + + +def test_shutdown_leaves_the_store_open_under_a_save_still_reading( + store_config, fake_store, staged_copies, monkeypatch +): + """A timed join is not evidence that the thread stopped. + + The save thread reads the KV pools through the store handle and the staging + slots, so closing the handle or dropping those buffers mid-transfer takes + the memory out from under it. + """ + monkeypatch.setattr(worker_module, "SAVE_DRAIN_TIMEOUT", 0.1) + reading = threading.Event() + release = threading.Event() + store_put = fake_store.batch_put_from_multi_buffers + + def blocking_put(*args, **kwargs): + reading.set() + assert release.wait(30.0), "the test never released the save thread" + return store_put(*args, **kwargs) + + fake_store.batch_put_from_multi_buffers = blocking_put with make_worker(fake_store, layout=make_layout()) as worker: # Queued directly: what matters is a thread inside the store call, not @@ -1269,52 +2176,851 @@ def test_staging_narrows_the_batch_to_the_budget(store_config, fake_store, stage assert worker._save_staging.num_slots == 2 assert worker._batch_size == 2 - hashes = [bytes([index]) * 16 for index in range(5)] - worker._put( - [ - RequestTransfers( - 1, - [PageTransfer(block_hash, 0, index) for index, block_hash in enumerate(hashes)], - ) - ] + hashes = [bytes([index]) * 16 for index in range(5)] + worker._put( + [ + RequestTransfers( + 1, + [PageTransfer(block_hash, 0, index) for index, block_hash in enumerate(hashes)], + ) + ] + ) + # Five pages through two slots, and no call wider than the slot count. + assert [len(keys) for keys, _a, _s in fake_store.put_calls] == [2, 2, 1] + + +# ---- asynchronous loads ---- + + +def enable_async_load(store_config, workers=None): + """Turn on `async_load` in the rendered client config.""" + settings = {"async_load": True} + if workers is not None: + settings["async_load_workers"] = workers + set_pool_setting(store_config, **settings) + + +def queue_async(worker, *transfers, fence=None): + """Bind metadata carrying asynchronous loads, as `handle_metadata` does. + + The allocation fence is bound first, as the manager binds it just before + the metadata of the batch that allocated the pages. + + Returns: + The fence the loads were queued with. + """ + fence = fence if fence is not None else AllocationFence(worker._store) + worker.bind_allocation_fence(fence) + worker.bind_connector_meta(MooncakeStoreMetadata(async_loads=list(transfers))) + return fence + + +def wait_async_idle(worker, timeout=5.0): + """Block until the load threads have finished every queued load.""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + with worker._load_lock: + if not worker._async_pending: + return + time.sleep(0.005) + raise AssertionError("the asynchronous load did not finish in time") + + +def test_worker_has_no_load_thread_by_default(store_config, fake_store): + with make_worker(fake_store, layout=make_layout()) as worker: + assert worker._async_load is False + assert worker._load_threads == [] + assert worker._async_stagings == [] + # Without asynchronous offers a parked id is echoed straight back. + assert worker.get_finished([], [5]) == ([], [5]) + + +def test_worker_starts_async_loads_from_metadata(store_config, fake_store): + """A rank queues its parked loads when it binds the metadata, behind the fence.""" + enable_async_load(store_config) + with make_worker(fake_store, layout=make_layout()) as worker: + assert worker._recover_failed_loads + block_hash = b"\x0a" * 16 + fake_store.objects.add(worker._namespaces[0].key(block_hash)) + fence = queue_async( + worker, + RequestTransfers(9, [PageTransfer(block_hash, 0, 2)]), + # Parked for an offer honoured with no tokens: nothing to wait for. + RequestTransfers(10, []), + ) + wait_async_idle(worker) + + # Only the request with pages to write waits for its allocation. + assert fence.waited_for == [{9}] + _saving, finished = worker.get_finished([], [9, 10]) + assert sorted(finished) == [9, 10] + assert worker.take_failed_async_load_requests() == set() + + +def test_worker_starts_async_loads_from_metadata_under_tensor_parallelism(store_config, fake_store): + """Every tensor-parallel rank queues its own shard's load from the broadcast metadata. + + Tensor parallelism is set on the built worker, as a rank that cannot restart a + request alone, rather than through the world size, which would also change + every key namespace. + """ + enable_async_load(store_config) + with make_worker(fake_store, layout=make_layout()) as worker: + worker._recover_failed_loads = False + block_hash = b"\x07" * 16 + fake_store.objects.add(worker._namespaces[0].key(block_hash)) + + queue_async(worker, RequestTransfers(9, [PageTransfer(block_hash, 0, 2)])) + wait_async_idle(worker) + + assert worker.get_finished([], [9]) == ([], [9]) + assert [keys for keys, _addresses, _sizes in fake_store.get_calls] == [ + [worker._namespaces[0].key(block_hash)] + ] + + +def test_worker_async_load_failure_is_fatal_under_tensor_parallelism(store_config, fake_store): + """Tensor-parallel ranks cannot restart a request alone, so a lost page fails the rank.""" + enable_async_load(store_config) + with make_worker(fake_store, layout=make_layout()) as worker: + worker._recover_failed_loads = False + queue_async(worker, RequestTransfers(4, [PageTransfer(b"\x08" * 16, 0, 1)])) + wait_async_idle(worker) + with pytest.raises(RuntimeError, match="cannot restart the request alone"): + worker.get_finished([], [4]) + + +def test_worker_rejects_async_loads_in_metadata_when_async_load_is_off(store_config, fake_store): + """Binding never raises; the next executor calls do.""" + with make_worker(fake_store, layout=make_layout()) as worker: + worker.bind_allocation_fence(AllocationFence(fake_store)) + worker.bind_connector_meta(MooncakeStoreMetadata(async_loads=[RequestTransfers(1, [])])) + with pytest.raises(RuntimeError, match="async_load is off"): + worker.start_load_kv(None) + with pytest.raises(RuntimeError, match="async_load is off"): + worker.get_finished([], [1]) + + +def test_worker_async_load_waits_on_the_fence_off_the_executor_thread(store_config, fake_store): + """Binding never blocks on the fence; the load thread waits on it before fetching.""" + enable_async_load(store_config) + with make_worker(fake_store, layout=make_layout()) as worker: + block_hash = b"\x0c" * 16 + fake_store.objects.add(worker._namespaces[0].key(block_hash)) + calls = [] + store_get = fake_store.batch_get_into_multi_buffers + + def recording_get(*args, **kwargs): + calls.append("get") + return store_get(*args, **kwargs) + + fake_store.batch_get_into_multi_buffers = recording_get + released = threading.Event() + fence_threads = [] + + def blocking_fence(request_ids): + fence_threads.append(threading.current_thread().name) + assert released.wait(5.0), "the test never released the fence" + calls.append(("fence", set(request_ids))) + + try: + queue_async( + worker, RequestTransfers(12, [PageTransfer(block_hash, 0, 2)]), fence=blocking_fence + ) + # The fence of the parking iteration travels with the request, so a + # later iteration binding another changes nothing it waits on. + worker.bind_allocation_fence(None) + assert worker.get_finished([], [12]) == ([], []) + finally: + released.set() + wait_async_idle(worker) + + assert calls == [("fence", {12}), "get"] + assert len(fence_threads) == 1 + assert fence_threads[0].startswith("mooncake-store-load-") + assert worker.get_finished([], []) == ([], [12]) + + +def test_worker_async_load_is_reported_once_the_runtime_announces_it(store_config, fake_store): + """A landed load waits for the runtime to say it parked the request.""" + enable_async_load(store_config) + layout = make_layout() + with make_worker(fake_store, layout=layout) as worker: + assert len(worker._load_threads) == 1 + block_hash = b"\x05" * 16 + fake_store.objects.add(worker._namespaces[0].key(block_hash)) + + queue_async(worker, RequestTransfers(7, [PageTransfer(block_hash, 0, 2)])) + wait_async_idle(worker) + + # Landed, but not announced: nothing to report yet. + assert worker.get_finished([], []) == ([], []) + assert worker.get_finished([], [7]) == ([], [7]) + assert worker.take_failed_async_load_requests() == set() + # Reported once only. + assert worker.get_finished([], []) == ([], []) + + (keys, addresses, sizes) = fake_store.get_calls[0] + expected_addresses, expected_sizes = PageAddressing(layout).buffers(0, 2) + assert keys == [worker._namespaces[0].key(block_hash)] + assert addresses == [expected_addresses] + assert sizes == [expected_sizes] + + +def test_worker_async_load_announced_first_waits_for_the_pages(store_config, fake_store): + """The runtime is never told a load finished while its bytes are in flight.""" + enable_async_load(store_config) + with make_worker(fake_store, layout=make_layout()) as worker: + block_hash = b"\x06" * 16 + fake_store.objects.add(worker._namespaces[0].key(block_hash)) + gate = threading.Event() + store_get = fake_store.batch_get_into_multi_buffers + + def gated(*args, **kwargs): + assert gate.wait(5.0), "the test never opened the gate" + return store_get(*args, **kwargs) + + fake_store.batch_get_into_multi_buffers = gated + try: + queue_async(worker, RequestTransfers(7, [PageTransfer(block_hash, 0, 1)])) + assert worker.get_finished([], [7]) == ([], []) + finally: + gate.set() + wait_async_idle(worker) + assert worker.get_finished([], []) == ([], [7]) + assert worker.take_failed_async_load_requests() == set() + + +def test_worker_async_load_failure_is_handed_to_the_executor(store_config, fake_store): + """A missing page still releases the parked request, marked for a restart.""" + enable_async_load(store_config) + with make_worker(fake_store, layout=make_layout()) as worker: + present, missing, later = (bytes([1]) * 16, bytes([2]) * 16, bytes([3]) * 16) + for block_hash in (present, later): + fake_store.objects.add(worker._namespaces[0].key(block_hash)) + # One page per store call, so the page after the missing one is a + # separate call the load thread can skip. + worker._async_batch_size = 1 + queue_async( + worker, + RequestTransfers( + 8, + [ + PageTransfer(present, 0, 1), + PageTransfer(missing, 0, 2), + PageTransfer(later, 0, 3), + ], + ), + ) + wait_async_idle(worker) + + assert worker.get_finished([], [8]) == ([], [8]) + assert worker.take_failed_async_load_requests() == {8} + # Cleared on return. + assert worker.take_failed_async_load_requests() == set() + # The page after the missing one was not fetched: the request restarts. + fetched = [key for keys, _addresses, _sizes in fake_store.get_calls for key in keys] + assert worker._namespaces[0].key(later) not in fetched + + +def test_worker_async_load_with_nothing_to_load_completes_at_once(store_config, fake_store): + enable_async_load(store_config) + with make_worker(fake_store, layout=make_layout()) as worker: + fence = queue_async(worker, RequestTransfers(9, [])) + assert worker.get_finished([], [9]) == ([], [9]) + assert worker.take_failed_async_load_requests() == set() + assert fence.waited_for == [] + assert fake_store.get_calls == [] + + +def test_worker_restarts_a_parked_request_it_was_never_handed(store_config, fake_store): + """An announced request with no load behind it must not stay parked.""" + enable_async_load(store_config) + with make_worker(fake_store, layout=make_layout()) as worker: + assert worker.get_finished([], [11]) == ([], [11]) + assert worker.take_failed_async_load_requests() == {11} + + +def test_worker_async_load_refuses_to_queue_without_a_fence(store_config, fake_store): + """Without the allocation fence nothing orders a load thread behind the allocation.""" + enable_async_load(store_config) + with make_worker(fake_store, layout=make_layout()) as worker: + block_hash = b"\x0b" * 16 + fake_store.objects.add(worker._namespaces[0].key(block_hash)) + worker.bind_allocation_fence(None) + worker.bind_connector_meta( + MooncakeStoreMetadata( + async_loads=[RequestTransfers(3, [PageTransfer(block_hash, 0, 1)])] + ) + ) + with pytest.raises(RuntimeError, match="no allocation fence"): + worker.start_load_kv(None) + assert fake_store.get_calls == [] + + +def test_worker_async_load_exception_is_fatal_at_the_next_executor_call(store_config, fake_store): + """A store call that raises fails the rank, even where a missing page would restart. + + The error is sticky, so no later call, such as a second recovery poll, + sees a healthy worker. + """ + enable_async_load(store_config) + with make_worker(fake_store, layout=make_layout()) as worker: + assert worker._recover_failed_loads + + def broken(*_args, **_kwargs): + raise RuntimeError("transport lost") + + fake_store.batch_get_into_multi_buffers = broken + queue_async(worker, RequestTransfers(5, [PageTransfer(b"\x0d" * 16, 0, 1)])) + wait_async_idle(worker) + + with pytest.raises(RuntimeError, match="background load failed") as excinfo: + worker.start_load_kv(None) + assert "transport lost" in str(excinfo.value) + with pytest.raises(RuntimeError, match="background load failed"): + worker.get_finished([], [5]) + with pytest.raises(RuntimeError, match="background load failed"): + worker.start_load_kv(None) + + +def test_bind_connector_meta_never_raises(store_config, fake_store): + """Binding runs outside the attention-DP failure vote, so the next calls raise instead.""" + block_hash = b"\x1a" * 16 + + def async_metadata(): + return MooncakeStoreMetadata( + async_loads=[RequestTransfers(1, [PageTransfer(block_hash, 0, 1)])] + ) + + with make_worker(fake_store, layout=make_layout()) as worker: + worker.bind_allocation_fence(AllocationFence(fake_store)) + worker.bind_connector_meta(async_metadata()) + with pytest.raises(RuntimeError, match="async_load is off"): + worker.get_finished([], []) + + enable_async_load(store_config) + # No layout registered yet, so no load thread exists to run the load. + with make_worker(fake_store) as worker: + worker.bind_allocation_fence(AllocationFence(fake_store)) + worker.bind_connector_meta(async_metadata()) + with pytest.raises(RuntimeError, match="layout was registered"): + worker.get_finished([], []) + + with make_worker(fake_store, layout=make_layout()) as worker: + worker.bind_allocation_fence(None) + worker.bind_connector_meta(async_metadata()) + with pytest.raises(RuntimeError, match="no allocation fence"): + worker.get_finished([], [1]) + assert fake_store.get_calls == [] + + +def test_worker_start_load_kv_starts_only_newly_bound_loads(store_config, fake_store): + """Recovery may call `start_load_kv` twice per pass and rebind the same metadata.""" + block_hash = b"\x0e" * 16 + with make_worker(fake_store, layout=make_layout()) as worker: + worker.bind_allocation_fence(AllocationFence(fake_store)) + bind_one_load(worker, fake_store, block_hash) + worker.start_load_kv(None) + worker.start_load_kv(None) + assert len(fake_store.get_calls) == 1 + + fake_store.get_calls.clear() + enable_async_load(store_config) + with make_worker(fake_store, layout=make_layout()) as worker: + fake_store.objects.add(worker._namespaces[0].key(block_hash)) + metadata = MooncakeStoreMetadata( + async_loads=[RequestTransfers(4, [PageTransfer(block_hash, 0, 1)])] + ) + worker.bind_allocation_fence(AllocationFence(fake_store)) + worker.bind_connector_meta(metadata) + assert metadata.async_loads == [] + worker.bind_connector_meta(metadata) + wait_async_idle(worker) + assert len(fake_store.get_calls) == 1 + assert worker.get_finished([], [4]) == ([], [4]) + + +def test_staged_async_worker_refuses_a_synchronous_load(store_config, fake_store, staged_copies): + """Its loads all go through the load threads, so no synchronous load pool is pinned.""" + enable_async_load(store_config) + with make_staged_worker(fake_store, store_config, layout=make_layout()) as worker: + assert worker._load_staging is None + bind_one_load(worker, fake_store, b"\x0f" * 16) + with pytest.raises(RuntimeError, match="synchronous staging pool"): + worker.start_load_kv(None) + assert fake_store.get_calls == [] + assert staged_copies == [] + + +@pytest.mark.parametrize("role", ["producer", "capacity"]) +def test_async_load_is_off_for_roles_that_do_not_read( + store_config, fake_store, staged_copies, role +): + """A role that never loads starts no load thread and pins no load staging.""" + enable_async_load(store_config, workers=2) + set_pool_setting(store_config, role=role, stage_through_host=True) + with make_worker(fake_store, layout=make_layout()) as worker: + assert worker._async_load is False + assert worker._load_threads == [] + assert worker._async_stagings == [] + assert worker._load_staging is None + + +def test_async_staging_registers_one_pool_per_worker_and_no_sync_load_pool( + store_config, fake_store, staged_copies +): + enable_async_load(store_config, workers=3) + with make_staged_worker(fake_store, store_config, layout=make_pool_layout()) as worker: + assert worker._load_staging is None + assert len(worker._async_stagings) == 3 + pools = [worker._save_staging, *worker._async_stagings] + # The save pool and one pool per load thread, and never the KV pools. + assert sorted(address for address, _ in fake_store.registered) == sorted( + pool.slot_address(0) for pool in pools + ) + assert POOL_BASE not in [address for address, _ in fake_store.registered] + + +def test_async_staged_load_scatters_only_landed_pages_on_the_loader_thread( + store_config, fake_store, staged_copies, monkeypatch +): + """The load thread scatters through its own slots, and never a failed one.""" + copy_threads = [] + + def record_copy(dst, src, size, kind, stream): + staged_copies.append((int(dst), int(src), int(size))) + copy_threads.append(threading.current_thread().name) + + monkeypatch.setattr(staging_module, "_memcpy_async", record_copy) + enable_async_load(store_config) + layout = make_layout(regions_per_group=2) + with make_staged_worker(fake_store, store_config, layout=layout) as worker: + good, bad = bytes([0x1A]) * 16, bytes([0x1B]) * 16 + for block_hash in (good, bad): + fake_store.objects.add(worker._namespaces[0].key(block_hash)) + fake_store.fail_gets_for.add(worker._namespaces[0].key(bad)) + queue_async( + worker, RequestTransfers(7, [PageTransfer(good, 0, 1), PageTransfer(bad, 0, 2)]) + ) + wait_async_idle(worker) + + assert worker.get_finished([], [7]) == ([], [7]) + assert worker.take_failed_async_load_requests() == {7} + # Both pages shared one store call, so slot 0 holds the landed page. + assert len(fake_store.get_calls) == 1 + slot = worker._async_stagings[0].slot_address(0) + good_addresses, good_sizes = PageAddressing(layout).buffers(0, 1) + expected = [] + offset = 0 + for address, size in zip(good_addresses, good_sizes): + expected.append((address, slot + offset, size)) + offset += size + assert staged_copies == expected + assert copy_threads and all( + name.startswith("mooncake-store-load-") for name in copy_threads + ) + + +def test_a_pool_below_one_page_warns_that_it_exceeds_the_budget( + store_config, fake_store, staged_copies, monkeypatch +): + """A budget below one page still pins a page, which the node check did not count.""" + layout = make_layout(regions_per_group=2) + page_bytes = PageAddressing(layout).bytes_per_page(0) + warnings = [] + monkeypatch.setattr(worker_module.logger, "warning", warnings.append) + + with make_staged_worker( + fake_store, store_config, layout=layout, budget=page_bytes - 1 + ) as worker: + assert worker._load_staging.num_slots == 1 + assert worker._save_staging.num_slots == 1 + raised = [message for message in warnings if "raised to one page" in message] + assert len(raised) == 1 + assert "the load, save staging pool(s)" in raised[0] + + warnings.clear() + enable_async_load(store_config, workers=2) + # A whole page for the save pool, half of one for each load thread. + with make_staged_worker(fake_store, store_config, layout=layout, budget=page_bytes) as worker: + assert [pool.num_slots for pool in worker._async_stagings] == [1, 1] + assert worker._save_staging.num_slots == 1 + raised = [message for message in warnings if "raised to one page" in message] + assert len(raised) == 1 + assert "async-load-0, async-load-1" in raised[0] + assert "async_load_workers" in raised[0] + + warnings.clear() + # Below one page for the save pool too, whose budget no load thread shares, + # so only the load threads' warning suggests fewer of them. + with make_staged_worker( + fake_store, store_config, layout=layout, budget=page_bytes - 1 + ) as worker: + assert worker._save_staging.num_slots == 1 + raised = [message for message in warnings if "raised to one page" in message] + assert len(raised) == 2 + (save_warning,) = [message for message in raised if "the save staging pool(s)" in message] + assert "async_load_workers" not in save_warning + + +def test_worker_pool_default_is_one_worker_with_the_full_staging( + store_config, fake_store, staged_copies +): + """One load thread gets the load direction's whole budget, and no other load pool is pinned.""" + enable_async_load(store_config) + layout = make_layout(regions_per_group=2) + page_bytes = PageAddressing(layout).bytes_per_page(0) + with make_staged_worker( + fake_store, store_config, layout=layout, budget=8 * page_bytes + ) as worker: + assert len(worker._load_threads) == 1 + assert worker._load_staging is None + assert [pool.num_slots for pool in worker._async_stagings] == [8] + # The save pool is the budget's too. + assert worker._save_staging.num_slots == 8 + assert worker._async_batch_size == worker._batch_size == 8 + + +def test_worker_pool_splits_the_staging_budget( + store_config, fake_store, staged_copies, monkeypatch +): + layout = make_layout(regions_per_group=2) + page_bytes = PageAddressing(layout).bytes_per_page(0) + # 512 pages of budget across 4 load threads: 128 slots each. + enable_async_load(store_config, workers=4) + set_pool_setting(store_config, transfer_batch_size=256) + with make_staged_worker( + fake_store, store_config, layout=layout, budget=512 * page_bytes + ) as worker: + assert [pool.num_slots for pool in worker._async_stagings] == [128] * 4 + assert worker._async_batch_size == 128 + # Saves keep the whole direction's budget. + assert worker._batch_size == 256 + + # 8 pages across 2 threads: 4 each, with no floor, and a warning saying so. + warnings = [] + monkeypatch.setattr(worker_module.logger, "warning", warnings.append) + enable_async_load(store_config, workers=2) + set_pool_setting(store_config, transfer_batch_size=16) + with make_staged_worker( + fake_store, store_config, layout=layout, budget=8 * page_bytes + ) as worker: + assert [pool.num_slots for pool in worker._async_stagings] == [4, 4] + assert worker._async_batch_size == 4 + assert worker._batch_size == 8 + assert any("stages only 4 page(s)" in message for message in warnings) + # The save pool narrows the batch it uses, and says so. + assert any("reduced its transfer batch" in message for message in warnings) + + # A consumer pins no pool the narrowed batch applies to, so only the load + # threads' narrowing is reported. + warnings.clear() + set_pool_setting(store_config, role="consumer") + with make_staged_worker( + fake_store, store_config, layout=layout, budget=8 * page_bytes + ) as worker: + assert worker._save_staging is None and worker._load_staging is None + assert worker._async_batch_size == 4 + assert any("stages only 4 page(s)" in message for message in warnings) + assert not any("reduced its transfer batch" in message for message in warnings) + + +def test_worker_pool_loads_two_requests_concurrently(store_config, fake_store): + """Two load threads transfer at the same time; each request is handled once.""" + enable_async_load(store_config, workers=2) + with make_worker(fake_store, layout=make_layout()) as worker: + assert len(worker._load_threads) == 2 + hashes = [bytes([0x21]) * 16, bytes([0x22]) * 16] + for block_hash in hashes: + fake_store.objects.add(worker._namespaces[0].key(block_hash)) + # Both threads must be inside the store call at once, or this times out. + rendezvous = threading.Barrier(2, timeout=5.0) + store_get = fake_store.batch_get_into_multi_buffers + + def meet_then_get(*args, **kwargs): + rendezvous.wait() + return store_get(*args, **kwargs) + + fake_store.batch_get_into_multi_buffers = meet_then_get + queue_async( + worker, + RequestTransfers(7, [PageTransfer(hashes[0], 0, 1)]), + RequestTransfers(8, [PageTransfer(hashes[1], 0, 2)]), + ) + wait_async_idle(worker) + + _saving, finished = worker.get_finished([], [7, 8]) + assert sorted(finished) == [7, 8] + assert worker.take_failed_async_load_requests() == set() + fetched = [key for keys, _addresses, _sizes in fake_store.get_calls for key in keys] + assert sorted(fetched) == sorted(worker._namespaces[0].key(h) for h in hashes) + + +def test_worker_pool_keeps_failure_and_success_apart(store_config, fake_store): + enable_async_load(store_config, workers=2) + with make_worker(fake_store, layout=make_layout()) as worker: + good, bad = bytes([0x31]) * 16, bytes([0x32]) * 16 + fake_store.objects.add(worker._namespaces[0].key(good)) + queue_async( + worker, + *( + RequestTransfers(request_id, [PageTransfer(block_hash, 0, page)]) + for request_id, block_hash, page in ((7, good, 1), (8, bad, 2), (9, good, 3)) + ), + ) + wait_async_idle(worker) + _saving, finished = worker.get_finished([], [7, 8, 9]) + assert sorted(finished) == [7, 8, 9] + assert worker.take_failed_async_load_requests() == {8} + + +def test_worker_pool_stats_line_summarises_the_interval(store_config, fake_store, monkeypatch): + """Asynchronous loads fold into the `stats_interval` load line, with their queue.""" + enable_async_load(store_config, workers=2) + set_pool_setting(store_config, stats_interval=60) + with make_worker(fake_store, layout=make_layout()) as worker: + # Idle and nothing loaded: no line, even when forced. + assert worker._maybe_log_load_stats(force=True) is None + + good, bad = bytes([0x41]) * 16, bytes([0x42]) * 16 + fake_store.objects.add(worker._namespaces[0].key(good)) + queue_async( + worker, + RequestTransfers(7, [PageTransfer(good, 0, 1)]), + RequestTransfers(8, [PageTransfer(bad, 0, 2)]), + ) + wait_async_idle(worker) + + lines = [] + monkeypatch.setattr(worker_module.logger, "info", lines.append) + # Not due yet: the interval gates the line. + assert worker._maybe_log_load_stats() is None + line = worker._maybe_log_load_stats(force=True) + assert line is not None and lines == [line] + assert line.startswith("mooncake-store load stats: rank 0 ") + assert "loads=2 pages=2 batches=2 " in line + assert "failed_pages=1 " in line + assert "wait_ms mean=" in line + assert "workers=2 queue=0 active=0 " in line + # The load thread waits on the fence, so that phase is timed too. + assert "phase_ms(wall/cpu per load) fence=" in line + # The counters start over, so a quiet interval logs nothing. + assert worker._maybe_log_load_stats(force=True) is None + + +def test_the_ranks_device_is_adopted_by_the_load_threads(store_config, fake_store, fake_cuda): + """Like the save thread, a load thread must not run on torch's default device.""" + enable_async_load(store_config, workers=2) + # No save thread, so every device switch below is a load thread's. + set_pool_setting(store_config, role="consumer") + with make_worker(fake_store, layout=make_layout()) as worker: + assert worker._device_index == 3 + assert len(worker._load_threads) == 2 + # Each thread sets the device before registration stops waiting for it. + assert fake_cuda == [3, 3] + + +def test_a_load_thread_that_cannot_start_fails_registration( + store_config, fake_store, fake_cuda, monkeypatch +): + """A parked request must not wait on a load thread that died at startup.""" + + def refuse(_index): + raise RuntimeError("cudaSetDevice failed") + + monkeypatch.setattr(torch.cuda, "set_device", refuse) + enable_async_load(store_config) + set_pool_setting(store_config, role="consumer") + + with pytest.raises(RuntimeError, match="load thread"): + with make_worker(fake_store, layout=make_layout()): + pytest.fail("registration should not have completed") + + +def _blocking_store_get(fake_store): + """Make the store's get block until released, and say when it is reading.""" + reading = threading.Event() + release = threading.Event() + store_get = fake_store.batch_get_into_multi_buffers + + def blocking_get(*args, **kwargs): + reading.set() + assert release.wait(30.0), "the test never released the load thread" + return store_get(*args, **kwargs) + + fake_store.batch_get_into_multi_buffers = blocking_get + return reading, release + + +def test_shutdown_leaves_the_store_open_under_a_load_still_reading( + store_config, fake_store, monkeypatch +): + """A load thread inside a zero-copy get still writes the registered KV pools.""" + monkeypatch.setattr(worker_module, "SAVE_DRAIN_TIMEOUT", 0.1) + enable_async_load(store_config) + reading, release = _blocking_store_get(fake_store) + + with make_worker(fake_store, layout=make_layout()) as worker: + try: + queue_async(worker, RequestTransfers(1, [PageTransfer(b"\x05" * 16, 0, 0)])) + assert reading.wait(30.0), "the load thread never reached the store" + + worker.shutdown() + + assert not fake_store.closed + assert worker._store is fake_store + assert fake_store.unregistered == [] + # Kept rather than dropped, so a later call retries the join. + assert worker._load_threads + finally: + release.set() + worker.shutdown() + + assert fake_store.closed + assert worker._store is None + assert worker._load_threads == [] + + +def test_shutdown_skips_loads_still_queued(store_config, fake_store, monkeypatch): + """A load queued behind the one in flight never reaches the store once shutting down.""" + monkeypatch.setattr(worker_module, "SAVE_DRAIN_TIMEOUT", 0.1) + enable_async_load(store_config) + reading, release = _blocking_store_get(fake_store) + + with make_worker(fake_store, layout=make_layout()) as worker: + first, second = bytes([0x51]) * 16, bytes([0x52]) * 16 + first_key, second_key = (worker._namespaces[0].key(h) for h in (first, second)) + fake_store.objects.update((first_key, second_key)) + try: + queue_async( + worker, + RequestTransfers(1, [PageTransfer(first, 0, 1)]), + RequestTransfers(2, [PageTransfer(second, 0, 2)]), + ) + assert reading.wait(30.0), "the load thread never reached the store" + worker.shutdown() + finally: + release.set() + worker.shutdown() + + assert [key for keys, _addresses, _sizes in fake_store.get_calls for key in keys] == [ + first_key + ] + assert worker._async_pending == set() + + +@pytest.mark.parametrize("async_load", [False, True], ids=["sync", "async"]) +def test_shutdown_closes_every_staging_pool_before_closing_the_store( + store_config, fake_store, staged_copies, async_load +): + if async_load: + enable_async_load(store_config, workers=2) + with make_staged_worker(fake_store, store_config, layout=make_layout()) as worker: + pools = [worker._load_staging, worker._save_staging, *worker._async_stagings] + bases = [pool.slot_address(0) for pool in pools if pool is not None] + assert len(bases) == (3 if async_load else 2) + close_at = fake_store.events.index(("close",)) + unregistered = [event[1] for event in fake_store.events[:close_at] if event[0] == "unregister"] + assert sorted(unregistered) == sorted(bases) + assert worker._unreleased_staging == [] + + +def test_shutdown_keeps_a_staging_pool_whose_unregistration_failed( + store_config, fake_store, staged_copies, monkeypatch +): + """The fabric may still reach a pool whose registration did not come back.""" + with make_staged_worker(fake_store, store_config, layout=make_layout()) as worker: + save_pool = worker._save_staging + save_base = save_pool.slot_address(0) + unregister = fake_store.unregister_buffer + monkeypatch.setattr( + fake_store, + "unregister_buffer", + lambda address: -1 if address == save_base else unregister(address), ) - # Five pages through two slots, and no call wider than the slot count. - assert [len(keys) for keys, _a, _s in fake_store.put_calls] == [2, 2, 1] + worker.shutdown() + + assert worker._unreleased_staging == [save_pool] + assert worker._save_staging is None + assert fake_store.closed # ---- scheduler ---- class FakeWorker: - """Stands in for the process-local worker's lookup service.""" + """Stands in for the process-local worker's lookup service and save handoff.""" def __init__(self, hit_blocks=0): self.hit_blocks = hit_blocks self.queries = [] + #: What the last `wait_for_save` queued; see `submit_saves`. + self.submitted = set() def count_prefix_hit(self, block_hashes): self.queries.append(list(block_hashes)) return min(self.hit_blocks, len(block_hashes)) + def take_submitted_saves(self): + submitted, self.submitted = self.submitted, set() + return submitted + + +def submit_saves(scheduler, metadata): + """What `wait_for_save` reports once the pass's forward has been enqueued.""" + scheduler._worker.submitted = {transfers.request_id for transfers in metadata.saves} -def make_scheduler(store_config, hit_blocks=0): - scheduler = MooncakeStoreConnectorScheduler(make_llm_args()) + +def make_scheduler(store_config, hit_blocks=0, speculative_config=None): + scheduler = MooncakeStoreConnectorScheduler( + make_llm_args(speculative_config=speculative_config) + ) scheduler._worker = FakeWorker(hit_blocks) return scheduler -def request_data(request_id, new_tokens, page_indices): +def request_data( + request_id, + new_tokens, + page_indices, + *, + computed_position=0, + num_scheduled_tokens=None, +): + """One request's scheduler output. + + By default the pass computes exactly `new_tokens` from position 0: a first + context chunk covering the whole prompt. + """ return RequestData( request_id=request_id, new_tokens=list(new_tokens), new_block_ids=list(page_indices), - computed_position=0, - num_scheduled_tokens=len(new_tokens), + computed_position=computed_position, + num_scheduled_tokens=( + len(new_tokens) if num_scheduled_tokens is None else num_scheduled_tokens + ), # One entry per layer group, indexed by layer group id. new_block_ids_by_layer_group=[list(page_indices)], ) +def settle_offer(scheduler, request, honoured_tokens, page_indices): + """Report an allocation the way `KVCacheManagerV2` does, after it honoured part of an offer. + + The manager always calls the per-layer-group form, and the base class folds + a single layer group into the flat `update_state_after_alloc`. + """ + request.py_num_connector_matched_tokens = honoured_tokens + scheduler.update_state_after_alloc_by_layer_group(request, [list(page_indices)]) + + +def saved_pages(metadata): + """Every page index a build bound for saving, in order.""" + return [page.page_index for transfers in metadata.saves for page in transfers.pages] + + def test_scheduler_never_offers_the_whole_prompt(store_config): scheduler = make_scheduler(store_config, hit_blocks=99) # Exactly three full blocks: the last one is withheld so the runtime still @@ -1343,15 +3049,22 @@ def test_scheduler_skips_local_prefix_when_looking_up(store_config): def test_scheduler_builds_loads_for_the_offered_blocks(store_config): scheduler = make_scheduler(store_config, hit_blocks=2) - tokens = list(range(5 * TOKENS_PER_BLOCK)) + tokens = list(range(6 * TOKENS_PER_BLOCK)) request = make_request(1, tokens) - scheduler.get_num_new_matched_tokens(request, 0) - - output = SchedulerOutput(new_requests=[request_data(1, tokens, [10, 11, 12, 13, 14])]) - metadata = scheduler.build_connector_meta(output) + assert scheduler.get_num_new_matched_tokens(request, 0) == (2 * TOKENS_PER_BLOCK, False) + pages = [10, 11, 12, 13, 14, 15] + settle_offer(scheduler, request, 2 * TOKENS_PER_BLOCK, pages) + + # The runtime leaves the restore out of `computed_position`, and this chunk + # computes the three blocks after it. + data = request_data( + 1, tokens, pages, computed_position=0, num_scheduled_tokens=3 * TOKENS_PER_BLOCK + ) + metadata = scheduler.build_connector_meta(SchedulerOutput(new_requests=[data])) assert [page.page_index for page in metadata.loads[0].pages] == [10, 11] # Blocks 0 and 1 came from the store, so only blocks 2..4 are written back. + # Block 5 has a hash and a page, but this pass does not compute it. assert [page.page_index for page in metadata.saves[0].pages] == [12, 13, 14] @@ -1365,15 +3078,116 @@ def test_scheduler_does_not_resave_blocks_across_iterations(store_config): SchedulerOutput(new_requests=[request_data(1, tokens, [4, 5])]) ) assert [page.page_index for page in first.saves[0].pages] == [4, 5] + submit_saves(scheduler, first) # A generation step completes one more block; only that block is saved. more_tokens = list(range(2 * TOKENS_PER_BLOCK, 3 * TOKENS_PER_BLOCK)) second = scheduler.build_connector_meta( - SchedulerOutput(cached_requests=[request_data(1, more_tokens, [6])]) + SchedulerOutput( + cached_requests=[ + request_data(1, more_tokens, [6], computed_position=2 * TOKENS_PER_BLOCK) + ] + ) ) assert [page.page_index for page in second.saves[0].pages] == [6] +def test_scheduler_resaves_blocks_of_a_pass_that_never_ran(store_config): + """Saves bound for a pass that never reached `wait_for_save` are emitted again.""" + scheduler = make_scheduler(store_config, hit_blocks=0) + tokens = list(range(2 * TOKENS_PER_BLOCK)) + scheduler.get_num_new_matched_tokens(make_request(1, tokens), 0) + first = scheduler.build_connector_meta( + SchedulerOutput(new_requests=[request_data(1, tokens, [4, 5])]) + ) + assert [page.page_index for page in first.saves[0].pages] == [4, 5] + + # The pass was abandoned after binding, so `wait_for_save` never queued it. + # The request is scheduled again from the same position. + retry = request_data(1, [], [], num_scheduled_tokens=2 * TOKENS_PER_BLOCK) + again = scheduler.build_connector_meta(SchedulerOutput(cached_requests=[retry])) + assert [page.page_index for page in again.saves[0].pages] == [4, 5] + submit_saves(scheduler, again) + + more_tokens = list(range(2 * TOKENS_PER_BLOCK, 3 * TOKENS_PER_BLOCK)) + after = scheduler.build_connector_meta( + SchedulerOutput( + cached_requests=[ + request_data(1, more_tokens, [6], computed_position=2 * TOKENS_PER_BLOCK) + ] + ) + ) + assert [page.page_index for page in after.saves[0].pages] == [6] + + +def test_scheduler_confirms_only_the_saves_the_worker_queued(store_config): + """A request dropped from the bound saves keeps its blocks unsaved.""" + scheduler = make_scheduler(store_config, hit_blocks=0) + tokens = list(range(TOKENS_PER_BLOCK)) + for request_id in (1, 2): + scheduler.get_num_new_matched_tokens(make_request(request_id, tokens), 0) + first = scheduler.build_connector_meta( + SchedulerOutput(new_requests=[request_data(1, tokens, [4]), request_data(2, tokens, [8])]) + ) + assert [transfers.request_id for transfers in first.saves] == [1, 2] + scheduler._worker.submitted = {1} + + retry = [ + request_data(request_id, [], [], num_scheduled_tokens=TOKENS_PER_BLOCK) + for request_id in (1, 2) + ] + again = scheduler.build_connector_meta(SchedulerOutput(cached_requests=retry)) + assert [transfers.request_id for transfers in again.saves] == [2] + assert [page.page_index for page in again.saves[0].pages] == [8] + + +def test_scheduler_forgets_unconfirmed_saves_of_a_reset_request(store_config): + scheduler = make_scheduler(store_config, hit_blocks=0) + tokens = list(range(TOKENS_PER_BLOCK)) + request = make_request(1, tokens) + scheduler.get_num_new_matched_tokens(request, 0) + scheduler.build_connector_meta(SchedulerOutput(new_requests=[request_data(1, tokens, [4])])) + assert scheduler._unconfirmed_saves == {1: 1} + scheduler.request_reset(request) + assert scheduler._unconfirmed_saves == {} + + +def test_scheduler_ignores_a_stale_report_after_a_reset(store_config): + """A report about a reset request's last pass cannot confirm its next one. + + The worker's report is taken on every build, including one with nothing + awaiting confirmation. Were it left behind, the build after a retry whose + pass never ran would take it, count the retry's saves as stored, and never + publish those blocks. + """ + scheduler = make_scheduler(store_config, hit_blocks=0) + tokens = list(range(2 * TOKENS_PER_BLOCK)) + request = make_request(1, tokens) + scheduler.get_num_new_matched_tokens(request, 0) + first = scheduler.build_connector_meta( + SchedulerOutput(new_requests=[request_data(1, tokens, [4, 5])]) + ) + assert saved_pages(first) == [4, 5] + submit_saves(scheduler, first) + scheduler.request_reset(request) + + # Readmitted to different pages. + scheduler.get_num_new_matched_tokens(request, 0) + retry = scheduler.build_connector_meta( + SchedulerOutput(new_requests=[request_data(1, tokens, [9, 10])]) + ) + assert saved_pages(retry) == [9, 10] + + # That pass never ran, so nothing reports it and the next build emits its + # saves again. + again = scheduler.build_connector_meta( + SchedulerOutput( + cached_requests=[request_data(1, [], [], num_scheduled_tokens=2 * TOKENS_PER_BLOCK)] + ) + ) + assert saved_pages(again) == [9, 10] + + def test_scheduler_waits_for_a_block_to_fill_before_saving(store_config): scheduler = make_scheduler(store_config, hit_blocks=0) tokens = list(range(TOKENS_PER_BLOCK + 1)) @@ -1387,6 +3201,238 @@ def test_scheduler_waits_for_a_block_to_fill_before_saving(store_config): assert [page.page_index for page in metadata.saves[0].pages] == [4] +def test_scheduler_publishes_only_computed_full_blocks(store_config): + scheduler = make_scheduler(store_config) + tokens = list(range(3 * TOKENS_PER_BLOCK + 1)) + scheduler.get_num_new_matched_tokens(make_request(1, tokens), 0) + data = request_data(1, tokens, [4, 5, 6, 7], num_scheduled_tokens=TOKENS_PER_BLOCK - 1) + first = scheduler.build_connector_meta(SchedulerOutput(new_requests=[data])) + assert first.saves == [] + + # Every prompt hash and capacity page was reported up front, but no page + # may be published until the pass computing its final token runs. + for ordinal in range(3): + computed = TOKENS_PER_BLOCK - 1 if ordinal == 0 else ordinal * TOKENS_PER_BLOCK + scheduled = 1 if ordinal == 0 else TOKENS_PER_BLOCK + data = request_data(1, [], [], computed_position=computed, num_scheduled_tokens=scheduled) + metadata = scheduler.build_connector_meta(SchedulerOutput(cached_requests=[data])) + assert [page.page_index for page in metadata.saves[0].pages] == [4 + ordinal] + submit_saves(scheduler, metadata) + + +@pytest.mark.parametrize("local_blocks", [0, 1]) +@pytest.mark.parametrize("use_mtp", [False, True], ids=["no_draft", "mtp"]) +def test_scheduler_saves_computed_tail_after_restore(store_config, local_blocks, use_mtp): + """The same saves with and without a one-model draft's prompt lookahead. + + The prompt runs one token past its last full block, so a lookahead of one + leaves every block hashable and the bound is the computed position alone. + """ + config = MTPDecodingConfig(max_draft_len=1) if use_mtp else None + scheduler = make_scheduler(store_config, hit_blocks=2, speculative_config=config) + assert scheduler._prompt_lookahead == (1 if use_mtp else 0) + tokens = list(range(6 * TOKENS_PER_BLOCK + 1)) + request = make_request(1, tokens) + assert scheduler.get_num_new_matched_tokens(request, local_blocks * TOKENS_PER_BLOCK) == ( + 2 * TOKENS_PER_BLOCK, + False, + ) + pages = list(range(10, 17)) + settle_offer(scheduler, request, 2 * TOKENS_PER_BLOCK, pages) + # The runtime subtracts the connector's restore from `computed_position`, + # while every prompt hash and capacity page is available ahead of + # computation. + data = request_data( + 1, + tokens, + pages, + computed_position=local_blocks * TOKENS_PER_BLOCK, + num_scheduled_tokens=TOKENS_PER_BLOCK + 1, + ) + metadata = scheduler.build_connector_meta(SchedulerOutput(new_requests=[data])) + tail_block = local_blocks + 2 + assert [page.page_index for page in metadata.loads[0].pages] == [ + 10 + local_blocks, + 11 + local_blocks, + ] + assert [page.page_index for page in metadata.saves[0].pages] == [10 + tail_block] + submit_saves(scheduler, metadata) + + # A later chunk includes the restored prefix in its position. The restore + # is counted only once, and the partial page waits until it is complete. + data = request_data( + 1, + [], + [], + computed_position=(tail_block + 1) * TOKENS_PER_BLOCK + 1, + num_scheduled_tokens=TOKENS_PER_BLOCK - 1, + ) + metadata = scheduler.build_connector_meta(SchedulerOutput(cached_requests=[data])) + assert metadata.loads == [] + assert [page.page_index for page in metadata.saves[0].pages] == [11 + tail_block] + + +def test_scheduler_mtp_keys_require_matching_lookahead(store_config): + """Block keys cover the prompt token a one-model draft reads past each block.""" + tokens = list(range(3 * TOKENS_PER_BLOCK)) + other_tokens = list(tokens) + # The first token of block 1: the draft's lookahead past block 0. + other_tokens[TOKENS_PER_BLOCK] += 1000 + scheduler = make_scheduler( + store_config, hit_blocks=99, speculative_config=MTPDecodingConfig(max_draft_len=1) + ) + assert scheduler._prompt_lookahead == 1 + + scheduler.get_num_new_matched_tokens(make_request(1, tokens), 0) + other = make_request(2, other_tokens) + assert scheduler.get_num_new_matched_tokens(other, 0) == (2 * TOKENS_PER_BLOCK, False) + first, second = scheduler._worker.queries + assert first == list(BlockHashChain(TOKENS_PER_BLOCK, prompt_lookahead=1).extend(tokens)) + assert first[0] != second[0] + # Pages a server without the draft wrote are a different key space. + assert first[0] != BlockHashChain(TOKENS_PER_BLOCK).extend(tokens)[0] + + pages = [4, 5, 6] + settle_offer(scheduler, other, 2 * TOKENS_PER_BLOCK, pages) + data = request_data( + 2, other_tokens, pages, computed_position=0, num_scheduled_tokens=TOKENS_PER_BLOCK + ) + metadata = scheduler.build_connector_meta(SchedulerOutput(new_requests=[data])) + assert [page.page_index for page in metadata.loads[0].pages] == [4, 5] + # The pass computes the final block, but the prompt holds no token past it + # for the draft to have read, so it has no key and is not saved. + assert metadata.saves == [] + + +def test_scheduler_mtp_vanilla_lookahead_is_the_draft_length(store_config): + """Vanilla MTP chains its draft layers, so it reads `max_draft_len` tokens ahead.""" + config = MTPDecodingConfig(max_draft_len=2, use_mtp_vanilla=True) + assert config.spec_dec_mode.is_mtp_vanilla() + scheduler = make_scheduler(store_config, hit_blocks=99, speculative_config=config) + assert scheduler._prompt_lookahead == 2 + + # Block 2's lookahead runs one token past this prompt, so it is not offered, + # where a server without a draft offers it. + tokens = list(range(3 * TOKENS_PER_BLOCK + 1)) + assert scheduler.get_num_new_matched_tokens(make_request(1, tokens), 0) == ( + 2 * TOKENS_PER_BLOCK, + False, + ) + assert scheduler._worker.queries[0] == list( + BlockHashChain(TOKENS_PER_BLOCK, prompt_lookahead=2).extend(tokens) + ) + without_draft = make_scheduler(store_config, hit_blocks=99) + assert without_draft._prompt_lookahead == 0 + assert without_draft.get_num_new_matched_tokens(make_request(1, tokens), 0) == ( + 3 * TOKENS_PER_BLOCK, + False, + ) + + +def test_scheduler_eagle3_lookahead_is_one(store_config): + config = Eagle3DecodingConfig(max_draft_len=3, speculative_model="draft") + scheduler = make_scheduler(store_config, speculative_config=config) + assert scheduler._prompt_lookahead == 1 + + +def test_scheduler_requires_resolved_spec_config(store_config): + """An unresolved draft length would key draft KV as if it read nothing ahead.""" + config = MTPDecodingConfig() + assert config.max_draft_len is None + with pytest.raises(RuntimeError, match="resolved"): + make_scheduler(store_config, speculative_config=config) + + +def test_scheduler_does_not_save_capacity_pages_past_the_chunk(store_config): + """One-model speculation allocates pages past every context chunk. + + `num_extra_kv_tokens` (`max_draft_len - 1`) gives the block after a chunk + that ends on a block boundary a page before any pass has computed it, so + each block is published only by the chunk that computes it. + """ + scheduler = make_scheduler( + store_config, hit_blocks=0, speculative_config=MTPDecodingConfig(max_draft_len=3) + ) + tokens = list(range(4 * TOKENS_PER_BLOCK + 1)) + assert scheduler.get_num_new_matched_tokens(make_request(1, tokens), 0) == (0, False) + + for ordinal in range(4): + # The first chunk computes one block and is allocated two pages. + new_pages = [10, 11] if ordinal == 0 else [11 + ordinal] + data = request_data( + 1, + tokens if ordinal == 0 else [], + new_pages, + computed_position=ordinal * TOKENS_PER_BLOCK, + num_scheduled_tokens=TOKENS_PER_BLOCK + (1 if ordinal == 3 else 0), + ) + output = ( + SchedulerOutput(new_requests=[data]) + if ordinal == 0 + else SchedulerOutput(cached_requests=[data]) + ) + metadata = scheduler.build_connector_meta(output) + assert saved_pages(metadata) == [10 + ordinal] + submit_saves(scheduler, metadata) + + +def test_scheduler_rebuild_in_the_same_iteration_emits_the_same_saves(store_config): + """A second build of one iteration bounds its saves exactly like the first. + + By then the runtime no longer subtracts the restore from + `computed_position`, and the offer is already consumed, so the two halves + of the computed end trade places and the bound comes out the same. + """ + scheduler = make_scheduler(store_config, hit_blocks=2) + tokens = list(range(5 * TOKENS_PER_BLOCK + 1)) + request = make_request(1, tokens) + scheduler.get_num_new_matched_tokens(request, 0) + pages = list(range(10, 16)) + settle_offer(scheduler, request, 2 * TOKENS_PER_BLOCK, pages) + + first = scheduler.build_connector_meta( + SchedulerOutput( + new_requests=[ + request_data( + 1, tokens, pages, computed_position=0, num_scheduled_tokens=3 * TOKENS_PER_BLOCK + ) + ] + ) + ) + assert [page.page_index for page in first.loads[0].pages] == [10, 11] + assert saved_pages(first) == [12, 13, 14] + + # Rebuilt before the pass runs: no report in between. + rebuilt = scheduler.build_connector_meta( + SchedulerOutput( + cached_requests=[ + request_data( + 1, + [], + [], + computed_position=2 * TOKENS_PER_BLOCK, + num_scheduled_tokens=3 * TOKENS_PER_BLOCK, + ) + ] + ) + ) + assert rebuilt.loads == [] + assert saved_pages(rebuilt) == [12, 13, 14] + submit_saves(scheduler, rebuilt) + + # The final token's pass completes no new block, and nothing is re-emitted. + last = scheduler.build_connector_meta( + SchedulerOutput( + cached_requests=[ + request_data( + 1, [], [], computed_position=5 * TOKENS_PER_BLOCK, num_scheduled_tokens=1 + ) + ] + ) + ) + assert last.saves == [] + + def test_scheduler_skips_blocks_without_a_page_in_every_group(store_config): scheduler = make_scheduler(store_config, hit_blocks=0) tokens = list(range(2 * TOKENS_PER_BLOCK)) @@ -1405,6 +3451,66 @@ def test_scheduler_skips_blocks_without_a_page_in_every_group(store_config): ] +@pytest.mark.parametrize( + "honoured_blocks, loaded_pages, saved", + [(3, [10, 11, 12], [13]), (2, [10, 11], [12]), (0, [], [10])], + ids=["whole_offer", "short", "none"], +) +def test_scheduler_sync_load_stops_at_the_honoured_prefix( + store_config, honoured_blocks, loaded_pages, saved +): + """A short honour trims both the load and the restored count the save bound uses.""" + scheduler = make_scheduler(store_config, hit_blocks=3) + tokens = list(range(6 * TOKENS_PER_BLOCK)) + request = make_request(1, tokens) + assert scheduler.get_num_new_matched_tokens(request, 0) == (3 * TOKENS_PER_BLOCK, False) + pages = list(range(10, 16)) + # The runtime's pages may cover less than the offer. + settle_offer(scheduler, request, honoured_blocks * TOKENS_PER_BLOCK, pages) + data = request_data( + 1, tokens, pages, computed_position=0, num_scheduled_tokens=TOKENS_PER_BLOCK + ) + metadata = scheduler.build_connector_meta(SchedulerOutput(new_requests=[data])) + assert [page.page_index for transfers in metadata.loads for page in transfers.pages] == ( + loaded_pages + ) + # The block after the honoured prefix is computed by this pass, so it is + # the only block published. + assert saved_pages(metadata) == saved + + +@pytest.mark.parametrize("async_load", [False, True], ids=["sync", "async"]) +def test_scheduler_rejects_an_honour_beyond_the_offer(store_config, async_load): + if async_load: + enable_async_load(store_config) + scheduler = make_scheduler(store_config, hit_blocks=1) + request = make_request(1, list(range(4 * TOKENS_PER_BLOCK))) + assert scheduler.get_num_new_matched_tokens(request, 0) == (TOKENS_PER_BLOCK, async_load) + with pytest.raises(RuntimeError, match="whole-block part of"): + settle_offer(scheduler, request, 2 * TOKENS_PER_BLOCK, [10, 11, 12]) + assert scheduler.build_connector_meta(SchedulerOutput()).async_loads == [] + + +def test_scheduler_rejects_an_honour_off_a_block_boundary(store_config): + scheduler = make_scheduler(store_config, hit_blocks=2) + request = make_request(1, list(range(4 * TOKENS_PER_BLOCK))) + assert scheduler.get_num_new_matched_tokens(request, 0) == (2 * TOKENS_PER_BLOCK, False) + with pytest.raises(RuntimeError, match="whole-block part of"): + settle_offer(scheduler, request, TOKENS_PER_BLOCK + 1, [10, 11, 12]) + + +def test_scheduler_refuses_an_honoured_block_without_a_page(store_config): + """The runtime counts an honoured block as computed, so it cannot be skipped.""" + scheduler = make_scheduler(store_config, hit_blocks=2) + tokens = list(range(4 * TOKENS_PER_BLOCK)) + request = make_request(1, tokens) + scheduler.get_num_new_matched_tokens(request, 0) + settle_offer(scheduler, request, 2 * TOKENS_PER_BLOCK, [10]) + data = request_data(1, tokens, [10], computed_position=0, num_scheduled_tokens=TOKENS_PER_BLOCK) + with pytest.raises(RuntimeError, match="has no destination page"): + scheduler.build_connector_meta(SchedulerOutput(new_requests=[data])) + + def test_scheduler_request_finished_pins_pages_only_when_saving(store_config): scheduler = make_scheduler(store_config, hit_blocks=0) tokens = list(range(2 * TOKENS_PER_BLOCK)) @@ -1551,3 +3657,266 @@ def test_scheduler_bypasses_multimodal_requests_it_cannot_identify(store_config) assert metadata.saves == [] # Nothing was handed to a save thread, so nothing pins its pages either. assert scheduler.request_finished(request, [7, 8, 9]) is False + + +# ---- admission-time queries and asynchronous offers ---- + + +def make_async_scheduler(store_config, hit_blocks): + enable_async_load(store_config) + return make_scheduler(store_config, hit_blocks=hit_blocks) + + +def take_async_loads(scheduler): + """The parked loads the next metadata hands to the workers.""" + return scheduler.build_connector_meta(SchedulerOutput()).async_loads + + +def test_scheduler_declares_admission_queries(): + """A query has no effect beyond its answer until the pages are reported. + + So the KV cache manager may ask while the scheduler is still admitting the + request (`kv_connector_config.admission_query`). + """ + assert MooncakeStoreConnectorScheduler.supports_admission_query is True + + +def test_scheduler_repeated_query_replaces_the_answer(store_config): + """A request asked at admission and asked again is served by the last answer only.""" + scheduler = make_scheduler(store_config, hit_blocks=2) + tokens = list(range(5 * TOKENS_PER_BLOCK)) + request = make_request(1, tokens) + assert scheduler.get_num_new_matched_tokens(request, 0) == (2 * TOKENS_PER_BLOCK, False) + scheduler._worker.hit_blocks = 3 + assert scheduler.get_num_new_matched_tokens(request, 0) == (3 * TOKENS_PER_BLOCK, False) + assert len(scheduler._worker.queries) == 2 + # The chain is kept and extended idempotently, not rebuilt. + assert len(scheduler._requests[1].chain.hashes) == 5 + + pages = [10, 11, 12, 13, 14] + settle_offer(scheduler, request, 3 * TOKENS_PER_BLOCK, pages) + data = request_data( + 1, tokens, pages, computed_position=0, num_scheduled_tokens=2 * TOKENS_PER_BLOCK + ) + metadata = scheduler.build_connector_meta(SchedulerOutput(new_requests=[data])) + assert [page.page_index for transfers in metadata.loads for page in transfers.pages] == [ + 10, + 11, + 12, + ] + assert saved_pages(metadata) == [13, 14] + + # The offer is consumed once. The follow-up carries no new tokens, since the + # adapter accumulates the per-step token deltas. + again = scheduler.build_connector_meta( + SchedulerOutput( + cached_requests=[ + request_data( + 1, [], [], computed_position=5 * TOKENS_PER_BLOCK, num_scheduled_tokens=0 + ) + ] + ) + ) + assert again.loads == [] + + +def test_scheduler_query_reset_query_is_clean(store_config): + """A deferred first chunk gives its pages back between two queries.""" + scheduler = make_scheduler(store_config, hit_blocks=2) + tokens = list(range(5 * TOKENS_PER_BLOCK)) + request = make_request(1, tokens) + scheduler.get_num_new_matched_tokens(request, 0) + first_state = scheduler._requests[1] + + scheduler.request_reset(request) + scheduler._worker.hit_blocks = 1 + assert scheduler.get_num_new_matched_tokens(request, 0) == (TOKENS_PER_BLOCK, False) + assert scheduler._requests[1] is not first_state + + pages = [20, 21, 22, 23, 24] + settle_offer(scheduler, request, TOKENS_PER_BLOCK, pages) + data = request_data( + 1, tokens, pages, computed_position=0, num_scheduled_tokens=4 * TOKENS_PER_BLOCK + ) + metadata = scheduler.build_connector_meta(SchedulerOutput(new_requests=[data])) + assert [page.page_index for transfers in metadata.loads for page in transfers.pages] == [20] + assert saved_pages(metadata) == [21, 22, 23, 24] + + +def test_scheduler_offers_asynchronously_when_enabled(store_config): + scheduler = make_async_scheduler(store_config, hit_blocks=2) + request = make_request(1, list(range(5 * TOKENS_PER_BLOCK))) + assert scheduler.get_num_new_matched_tokens(request, 0) == (2 * TOKENS_PER_BLOCK, True) + # A miss is never asynchronous: there would be nothing to wait for. + scheduler._worker.hit_blocks = 0 + request = make_request(2, list(range(5 * TOKENS_PER_BLOCK))) + assert scheduler.get_num_new_matched_tokens(request, 0) == (0, False) + + +def test_scheduler_sends_async_loads_through_metadata(store_config): + """The offer reaches every rank in the metadata, not only the adapter's own worker.""" + scheduler = make_async_scheduler(store_config, hit_blocks=2) + request = make_request(1, list(range(5 * TOKENS_PER_BLOCK))) + assert scheduler.get_num_new_matched_tokens(request, 0) == (2 * TOKENS_PER_BLOCK, True) + + settle_offer(scheduler, request, 2 * TOKENS_PER_BLOCK, [10, 11, 12, 13, 14]) + + metadata = scheduler.build_connector_meta(SchedulerOutput()) + (transfers,) = metadata.async_loads + assert transfers.request_id == 1 + assert [page.page_index for page in transfers.pages] == [10, 11] + assert metadata.loads == [] and metadata.saves == [] + assert metadata + # Sent once only. + assert take_async_loads(scheduler) == [] + + +def test_scheduler_repeated_query_replaces_the_answer_and_queues_one_load(store_config): + """Only the last answer counts, and exactly one asynchronous load is queued.""" + scheduler = make_async_scheduler(store_config, hit_blocks=2) + request = make_request(1, list(range(5 * TOKENS_PER_BLOCK))) + assert scheduler.get_num_new_matched_tokens(request, 0) == (2 * TOKENS_PER_BLOCK, True) + assert take_async_loads(scheduler) == [] + scheduler._worker.hit_blocks = 3 + assert scheduler.get_num_new_matched_tokens(request, 0) == (3 * TOKENS_PER_BLOCK, True) + + settle_offer(scheduler, request, 3 * TOKENS_PER_BLOCK, [10, 11, 12, 13, 14]) + + (transfers,) = take_async_loads(scheduler) + assert [page.page_index for page in transfers.pages] == [10, 11, 12] + assert take_async_loads(scheduler) == [] + + +def test_scheduler_queues_the_async_load_when_the_pages_are_reported(store_config): + """The offer reaches the workers, with the pages the runtime reserved, in the next metadata.""" + scheduler = make_async_scheduler(store_config, hit_blocks=2) + tokens = list(range(5 * TOKENS_PER_BLOCK)) + request = make_request(1, tokens) + assert scheduler.get_num_new_matched_tokens(request, 0) == (2 * TOKENS_PER_BLOCK, True) + pages = [10, 11, 12, 13, 14] + settle_offer(scheduler, request, 2 * TOKENS_PER_BLOCK, pages) + + (transfers,) = take_async_loads(scheduler) + state = scheduler._requests[1] + assert transfers.request_id == 1 + assert [(page.layer_group_id, page.page_index) for page in transfers.pages] == [ + (0, 10), + (0, 11), + ] + assert [page.block_hash for page in transfers.pages] == list(state.chain.hashes[:2]) + + # Once its load landed, the parked request comes back as a new request with + # all of its pages. It loads nothing more and does not write the stored + # blocks back. + data = request_data( + 1, + tokens, + pages, + computed_position=2 * TOKENS_PER_BLOCK, + num_scheduled_tokens=3 * TOKENS_PER_BLOCK, + ) + metadata = scheduler.build_connector_meta(SchedulerOutput(new_requests=[data])) + assert metadata.loads == [] + assert metadata.async_loads == [] + assert saved_pages(metadata) == [12, 13, 14] + + +def test_scheduler_async_load_stops_at_the_honoured_prefix(store_config): + scheduler = make_async_scheduler(store_config, hit_blocks=3) + request = make_request(1, list(range(6 * TOKENS_PER_BLOCK))) + scheduler.get_num_new_matched_tokens(request, 0) + # Three pages exist, but only two offered blocks were honoured. The third + # page holds the locally computed tail and must not be overwritten. + settle_offer(scheduler, request, 2 * TOKENS_PER_BLOCK, [20, 21, 22]) + (transfers,) = take_async_loads(scheduler) + assert [page.page_index for page in transfers.pages] == [20, 21] + assert scheduler._requests[1].async_loaded_upto == 2 + + +@pytest.mark.parametrize("pages", [[], [[20]], [[20, BAD_PAGE_INDEX]]]) +def test_scheduler_async_load_rejects_incomplete_honoured_prefix(store_config, pages): + scheduler = make_async_scheduler(store_config, hit_blocks=2) + request = make_request(1, list(range(6 * TOKENS_PER_BLOCK))) + scheduler.get_num_new_matched_tokens(request, 0) + request.py_num_connector_matched_tokens = 2 * TOKENS_PER_BLOCK + with pytest.raises(RuntimeError, match="honoured prefix"): + scheduler.update_state_after_alloc_by_layer_group(request, pages) + assert take_async_loads(scheduler) == [] + + +def test_scheduler_async_load_with_no_honoured_tokens_loads_nothing(store_config): + """The runtime parks the request anyway, so an empty load releases it.""" + scheduler = make_async_scheduler(store_config, hit_blocks=2) + request = make_request(1, list(range(6 * TOKENS_PER_BLOCK))) + scheduler.get_num_new_matched_tokens(request, 0) + settle_offer(scheduler, request, 0, [20, 21]) + (transfers,) = take_async_loads(scheduler) + assert transfers.request_id == 1 + assert transfers.pages == [] + assert scheduler._requests[1].async_loaded_upto == 0 + + +def test_scheduler_async_flat_alloc_is_a_single_group(store_config): + scheduler = make_async_scheduler(store_config, hit_blocks=1) + request = make_request(1, list(range(4 * TOKENS_PER_BLOCK))) + scheduler.get_num_new_matched_tokens(request, 0) + request.py_num_connector_matched_tokens = TOKENS_PER_BLOCK + scheduler.update_state_after_alloc(request, [30, 31, 32, 33]) + (transfers,) = take_async_loads(scheduler) + assert [(page.layer_group_id, page.page_index) for page in transfers.pages] == [(0, 30)] + + +def test_scheduler_alloc_hooks_queue_nothing_for_synchronous_loads(store_config): + scheduler = make_scheduler(store_config, hit_blocks=2) + tokens = list(range(5 * TOKENS_PER_BLOCK)) + request = make_request(1, tokens) + scheduler.get_num_new_matched_tokens(request, 0) + pages = [10, 11, 12, 13, 14] + request.py_num_connector_matched_tokens = 2 * TOKENS_PER_BLOCK + scheduler.update_state_after_alloc_by_layer_group(request, [list(pages)]) + scheduler.update_state_after_alloc(request, list(pages)) + data = request_data( + 1, tokens, pages, computed_position=0, num_scheduled_tokens=3 * TOKENS_PER_BLOCK + ) + metadata = scheduler.build_connector_meta(SchedulerOutput(new_requests=[data])) + assert metadata.async_loads == [] + assert [page.page_index for page in metadata.loads[0].pages] == [10, 11] + + +def test_scheduler_request_reset_drops_a_pending_async_load(store_config): + """A load for an allocation that was released must never reach the workers.""" + scheduler = make_async_scheduler(store_config, hit_blocks=2) + kept, reset = (make_request(i, list(range(5 * TOKENS_PER_BLOCK))) for i in (1, 2)) + for request, pages in ((kept, [10, 11, 12]), (reset, [20, 21, 22])): + scheduler.get_num_new_matched_tokens(request, 0) + settle_offer(scheduler, request, 2 * TOKENS_PER_BLOCK, pages) + + scheduler.request_reset(reset) + + assert [transfers.request_id for transfers in take_async_loads(scheduler)] == [1] + + +def test_scheduler_never_offers_async_to_a_bypassed_request(store_config): + """A request the store cannot key gets neither a lookup nor a load.""" + scheduler = make_async_scheduler(store_config, hit_blocks=2) + tokens = list(range(3 * TOKENS_PER_BLOCK)) + for request in ( + make_request(1, tokens, lora_task_id=0), + make_request(2, tokens, multimodal_positions=[4], multimodal_lengths=[4]), + ): + assert scheduler.get_num_new_matched_tokens(request, 0) == (0, False) + settle_offer(scheduler, request, 0, [7, 8, 9]) + assert scheduler._worker.queries == [] + assert take_async_loads(scheduler) == [] + + +@pytest.mark.parametrize("role", ["producer", "capacity"]) +def test_scheduler_capacity_and_producer_roles_never_offer_async(store_config, role): + enable_async_load(store_config) + set_pool_setting(store_config, role=role) + scheduler = make_scheduler(store_config, hit_blocks=2) + assert scheduler._async_load is False + request = make_request(1, list(range(3 * TOKENS_PER_BLOCK))) + assert scheduler.get_num_new_matched_tokens(request, 0) == (0, False) + settle_offer(scheduler, request, 0, [7, 8, 9]) + assert take_async_loads(scheduler) == [] diff --git a/tests/unittest/_torch/executor/test_mooncake_store_real_pool.py b/tests/unittest/_torch/executor/test_mooncake_store_real_pool.py index 7aacef3ce28c..54be4358e974 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_real_pool.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_real_pool.py @@ -178,12 +178,15 @@ def _layout(buffer: torch.Tensor) -> KvCacheLayout: return KvCacheLayout(tokens_per_block=TOKENS_PER_BLOCK, groups=tuple(groups)) -def _llm_args(store_config: MooncakeStoreConfig) -> SimpleNamespace: +def _llm_args( + store_config: MooncakeStoreConfig, *, enable_attention_dp: bool = False +) -> SimpleNamespace: """The fields the worker reads off `TorchLlmArgs`, and nothing else.""" return SimpleNamespace( model=f"/models/{store_config.model_key}", kv_cache_config=SimpleNamespace(tokens_per_block=TOKENS_PER_BLOCK), tensor_parallel_size=1, + enable_attention_dp=enable_attention_dp, pipeline_parallel_size=1, context_parallel_size=1, sparse_attention_config=None, @@ -274,6 +277,10 @@ def open_worker( layout: KvCacheLayout | None = None, stage_through_host: bool, model_key: str = MODEL_KEY, + enable_attention_dp: bool = False, + soft_pin_ttl_ms: int | None = None, + group_blocks: bool = False, + async_load: bool = False, ) -> Iterator[MooncakeStoreConnectorWorker]: """A worker holding a real store handle, over `buffer` or `layout` as its KV cache. @@ -287,9 +294,14 @@ def open_worker( segment_size=WORKER_SEGMENT_SIZE, stage_through_host=stage_through_host, run_dir=pool.run_dir, + soft_pin_ttl_ms=soft_pin_ttl_ms, + group_blocks=group_blocks, + async_load=async_load, ) with provision_pool(store_config, run_dir=pool.run_dir): - worker = MooncakeStoreConnectorWorker(_llm_args(store_config)) + worker = MooncakeStoreConnectorWorker( + _llm_args(store_config, enable_attention_dp=enable_attention_dp) + ) try: worker.register_kv_cache_layout(layout if layout is not None else _layout(buffer)) yield worker @@ -385,6 +397,52 @@ def test_a_saved_page_comes_back_byte_identical(mooncake_pool, stage_through_hos ) +@pytest.mark.parametrize("stage_through_host", [False, True], ids=["zero_copy", "host_staged"]) +def test_an_async_load_comes_back_byte_identical(mooncake_pool, stage_through_host): + """A load thread restores a page exactly as the synchronous path does. + + It runs on its own thread and, when staged, through its own pinned slots + and CUDA stream, and the page counts as loaded only once it is reported. + """ + buffer = _device_buffer(filled=True) + page_index = 3 + expected = _snapshot(buffer, page_index) + + with open_worker( + mooncake_pool, buffer=buffer, stage_through_host=stage_through_host, async_load=True + ) as worker: + save_page(worker, request_id=1, page_index=page_index) + buffer.zero_() + # The zeroing has to retire before the load thread writes the page, + # which is what the KV cache manager's allocation fence promises. + worker.bind_allocation_fence(lambda _request_ids: torch.cuda.synchronize()) + worker.bind_connector_meta( + MooncakeStoreMetadata( + async_loads=_page_transfers(2, page_index, tuple(range(NUM_GROUPS))) + ) + ) + + deadline = time.monotonic() + SAVE_TIMEOUT_SECONDS + while True: + _, finished_loading = worker.get_finished([], [2]) + if 2 in finished_loading: + break + assert time.monotonic() < deadline, ( + f"The asynchronous load of request 2 was not reported within " + f"{SAVE_TIMEOUT_SECONDS}s." + ) + time.sleep(0.05) + assert worker.take_failed_async_load_requests() == set() + + assert_page_matches(buffer, page_index, expected, context="Asynchronous load from the pool") + for group_id, region_id in _every_region(): + neighbour = _page_slice(buffer, group_id, region_id, page_index + 1) + assert not neighbour.any(), ( + f"The load wrote into page {page_index + 1} of layer group {group_id} " + f"({REGION_ROLES[region_id]}), which it was not asked for." + ) + + def test_a_page_written_by_one_worker_is_read_back_by_another(mooncake_pool): """The pool outlives the worker that filled it. @@ -445,6 +503,46 @@ def test_a_prefix_counts_only_once_every_rank_has_its_shard(mooncake_pool, monke ) +def test_attention_dp_owners_share_one_shard(mooncake_pool, monkeypatch): + """A page one attention-DP owner saved is a hit for another, and loads back intact. + + Every owner holds all heads, so all of them key their pages as attention + shard 0 of 1. Owner 1 saves nothing here, so its hit can only be owner 0's + page. A tensor-parallel rank of the same world size keys per-rank shards + and misses the same block, which shows the hit came from the shared + attention-DP key space rather than from a lookup that matches anything. + """ + monkeypatch.setattr(worker_module, "mpi_world_size", lambda: 2) + page_index, reader_page = 1, 6 + + writer_buffer = _device_buffer(filled=True) + expected = _snapshot(writer_buffer, page_index) + with open_worker( + mooncake_pool, buffer=writer_buffer, stage_through_host=True, enable_attention_dp=True + ) as owner_zero: + save_page(owner_zero, request_id=1, page_index=page_index) + + monkeypatch.setattr(worker_module, "mpi_rank", lambda: 1) + reader_buffer = _device_buffer(filled=False) + with open_worker( + mooncake_pool, buffer=reader_buffer, stage_through_host=True, enable_attention_dp=True + ) as owner_one: + assert owner_one.count_prefix_hit([BLOCK_HASH]) == 1, ( + "Attention-DP owner 1 did not find the block owner 0 published, so " + "nothing one owner computes would be reused by another." + ) + load_page(owner_one, request_id=2, page_index=reader_page) + assert_page_matches( + reader_buffer, reader_page, expected, context="Replay by another attention-DP owner" + ) + + with open_worker(mooncake_pool, buffer=reader_buffer, stage_through_host=True) as tp_rank_one: + assert tp_rank_one.count_prefix_hit([BLOCK_HASH]) == 0, ( + "A tensor-parallel rank found a block only attention-DP owners saved. " + "Its shard holds a slice of the heads, so it would load the wrong bytes." + ) + + def test_a_different_model_key_shares_nothing(mooncake_pool): """Two models' pages stay apart even when their tokens agree. @@ -478,6 +576,46 @@ def test_a_different_model_key_shares_nothing(mooncake_pool): ) +@pytest.mark.parametrize( + "soft_pin_ttl_ms, group_blocks", + [(60_000, False), (None, True), (60_000, True)], + ids=["soft_pin", "group", "both"], +) +def test_soft_pinned_grouped_pages_round_trip(mooncake_pool, soft_pin_ttl_ms, group_blocks): + """Pages put with a `ReplicateConfig` are stored and served like any other. + + Proves the pinned bindings accept the soft-pin and group fields all the way + from the `mooncake_store` block, through the rendered client config, to the + put. Which page the master evicts first is left to cluster validation. + """ + writer_page, reader_page = 4, 1 + writer_buffer = _device_buffer(filled=True) + expected = _snapshot(writer_buffer, writer_page) + + with open_worker( + mooncake_pool, + buffer=writer_buffer, + stage_through_host=True, + soft_pin_ttl_ms=soft_pin_ttl_ms, + group_blocks=group_blocks, + ) as writer: + assert writer._soft_pin_ttl_ms == soft_pin_ttl_ms + assert writer._group_blocks is group_blocks + save_page(writer, request_id=1, page_index=writer_page) + + reader_buffer = _device_buffer(filled=False) + with open_worker(mooncake_pool, buffer=reader_buffer, stage_through_host=True) as reader: + assert reader.count_prefix_hit([BLOCK_HASH]) == 1, ( + "A block put with a ReplicateConfig was not found, so the pool refused " + "or dropped the soft-pinned or grouped put." + ) + load_page(reader, request_id=2, page_index=reader_page) + + assert_page_matches( + reader_buffer, reader_page, expected, context="Replay of a soft-pinned or grouped put" + ) + + # ---- RDMA over a real KVCacheManagerV2 ---- diff --git a/tests/unittest/_torch/executor/test_mooncake_store_settings.py b/tests/unittest/_torch/executor/test_mooncake_store_settings.py index a740afcc8a4b..35249a9920b4 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_settings.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_settings.py @@ -28,12 +28,14 @@ import json import pytest +from pydantic import ValidationError import tensorrt_llm.usage as usage from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store import settings from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.config import ( CLIENT_CONFIG_NAME, CONFIG_PATH_ENV, + local_reuse_backoff, ) from tensorrt_llm.llmapi.llm import LLM, _TorchLLM from tensorrt_llm.llmapi.llm_args import ( @@ -90,6 +92,9 @@ def test_an_inherited_config_supersedes_the_block_it_ignores(monkeypatch, tmp_pa model_key="their-checkpoint", transfer_batch_size=128, stage_through_host=True, + staging_buffer_bytes="256MiB", + async_load=True, + async_load_workers=2, ) monkeypatch.setenv(CONFIG_PATH_ENV, str(inherited)) pool = mooncake_store_config(role="both", segment_size="16GiB") @@ -102,6 +107,9 @@ def test_an_inherited_config_supersedes_the_block_it_ignores(monkeypatch, tmp_pa assert pool.model_key == "their-checkpoint" assert pool.transfer_batch_size == 128 assert pool.stage_through_host is True + assert pool.staging_buffer_bytes == 256 * (1 << 20) + assert pool.async_load is True + assert pool.async_load_workers == 2 # How to reach the pool and where this run's files go are the server's own, # so an inherited config has nothing to say about them. @@ -156,6 +164,76 @@ def test_only_the_values_the_deployment_chose_are_reported_as_overridden(monkeyp assert "namespace" not in overrides[0] +def test_an_inherited_config_restates_the_eviction_and_stats_settings(monkeypatch, tmp_path): + """Soft pin, grouping and the stats interval are what the ranks read from the client config.""" + inherited = write_client_config( + tmp_path / "external.json", soft_pin_ttl_ms=5000, group_blocks=True, stats_interval=10 + ) + monkeypatch.setenv(CONFIG_PATH_ENV, str(inherited)) + pool = mooncake_store_config() + + settings.apply_effective_settings(llm_args(pool=pool)) + + assert pool.soft_pin_ttl_ms == 5000 + assert pool.group_blocks is True + assert pool.stats_interval == 10.0 + + +def test_an_inherited_config_without_soft_pin_or_stats_turns_them_off(monkeypatch, tmp_path): + """A TTL or stats interval the inherited config leaves out is off, not unstated.""" + inherited = write_client_config(tmp_path / "external.json") + monkeypatch.setenv(CONFIG_PATH_ENV, str(inherited)) + pool = mooncake_store_config(soft_pin_ttl_ms=86400000, group_blocks=True, stats_interval=30) + + warnings = [] + monkeypatch.setattr(settings.logger, "warning", warnings.append) + settings.apply_effective_settings(llm_args(pool=pool)) + + assert pool.soft_pin_ttl_ms is None + assert pool.group_blocks is False + assert pool.stats_interval is None + overrides = [line for line in warnings if "joins the pool as" in line] + assert len(overrides) == 1 + assert "soft_pin_ttl_ms=None" in overrides[0] + assert "group_blocks=False" in overrides[0] + assert "stats_interval=None" in overrides[0] + + +@pytest.mark.parametrize("ttl", [0, -1]) +def test_a_soft_pin_ttl_must_be_positive(ttl): + with pytest.raises(ValidationError, match="soft_pin_ttl_ms"): + mooncake_store_config(soft_pin_ttl_ms=ttl) + + +@pytest.mark.parametrize("interval", [0, -1, float("inf"), float("nan")]) +def test_a_stats_interval_must_be_a_positive_number(interval): + with pytest.raises(ValidationError, match="stats_interval"): + mooncake_store_config(stats_interval=interval) + + +def test_the_staging_and_async_load_settings_are_checked_up_front(): + """A bad value names the field it came from rather than failing every rank later.""" + pool = mooncake_store_config( + staging_buffer_bytes="512MiB", async_load=True, async_load_workers=4 + ) + assert pool.staging_buffer_bytes == "512MiB" + assert mooncake_store_config().staging_buffer_bytes == "1GiB" + assert mooncake_store_config().async_load is False + assert mooncake_store_config().async_load_workers == 1 + + with pytest.raises(ValidationError, match="async_load_workers"): + mooncake_store_config(async_load_workers=0) + with pytest.raises(ValidationError, match="mooncake_store.staging_buffer_bytes"): + mooncake_store_config(staging_buffer_bytes="1GB") + with pytest.raises(ValidationError, match="mooncake_store.staging_buffer_bytes"): + mooncake_store_config(staging_buffer_bytes=0) + # The size check is shared, and still names the segment size for its own. + with pytest.raises(ValidationError, match="mooncake_store.segment_size"): + mooncake_store_config(segment_size="1GB") + with pytest.raises(ValidationError, match="mooncake_store.segment_size"): + mooncake_store_config(segment_size=-1) + + def test_the_pool_replaces_this_engines_own_offload_tiers(monkeypatch, tmp_path): """A native tier beside the pool claims a second share of the node's DRAM.""" monkeypatch.setenv( @@ -246,6 +324,36 @@ def test_a_settled_deployment_is_restated_without_a_word(monkeypatch, tmp_path): assert args.kv_cache_config.host_cache_size == 0 +def test_local_reuse_backoff_takes_only_the_modes_it_has(): + assert mooncake_store_config().local_reuse_backoff == "off" + assert mooncake_store_config(local_reuse_backoff="align").local_reuse_backoff == "align" + with pytest.raises(ValidationError): + mooncake_store_config(local_reuse_backoff="partial") + + +def test_local_reuse_backoff_is_read_from_the_block_with_off_as_its_default(): + """Every rank reads it off its own arguments, so no block means the default.""" + pool = mooncake_store_config(local_reuse_backoff="align") + assert local_reuse_backoff(llm_args(pool=pool)) == "align" + assert local_reuse_backoff(llm_args(pool=mooncake_store_config())) == "off" + # An inherited `MOONCAKE_CONFIG_PATH` needs no block at all. + assert local_reuse_backoff(llm_args()) == "off" + + +def test_an_inherited_config_leaves_local_reuse_backoff_alone(monkeypatch, tmp_path): + """It is an engine setting the client config does not carry, so nothing restates it.""" + inherited = write_client_config(tmp_path / "external.json", role="consumer") + monkeypatch.setenv(CONFIG_PATH_ENV, str(inherited)) + pool = mooncake_store_config(local_reuse_backoff="align") + args = llm_args(pool=pool) + + settings.apply_effective_settings(args) + + assert pool.role == "consumer" + assert pool.local_reuse_backoff == "align" + assert local_reuse_backoff(args) == "align" + + def test_another_connector_keeps_its_own_cache_tiers(): """Only this connector claims to be the deployment's offload tier.""" args = llm_args(connector="lmcache", host_cache_size=4 * GIB) diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 067c082067a6..3a4ed27f0d21 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/tests/unittest/_torch/executor/test_py_executor.py @@ -1849,6 +1849,7 @@ def _make_disagg_transfer_request( req.py_prompt_len = prompt_len req.total_input_len_cp = prompt_len if total_input_len_cp is None else total_input_len_cp req.is_disagg_generation_transmission_in_progress = in_progress + req.py_kv_transfer_timed_out = False return req diff --git a/tests/unittest/_torch/executor/test_py_executor_creator_connector_spec_prefill.py b/tests/unittest/_torch/executor/test_py_executor_creator_connector_spec_prefill.py new file mode 100644 index 000000000000..89393fd776c9 --- /dev/null +++ b/tests/unittest/_torch/executor/test_py_executor_creator_connector_spec_prefill.py @@ -0,0 +1,223 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Which speculative configs may run next to a KV connector. + +A disaggregated context server with a one-model drafter (MTP or EAGLE3) serves +context-only requests, which never reject a draft token, so a connector whose +two halves declare `supports_speculative_prefill` stays valid there. Every other +combination keeps the speculative-decoding refusal, and an explicit +KVCacheManagerV1 selection keeps V1's unchanged behaviour. `mooncake-store` is +the connector shipped with such declarations. +""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest + +from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import ( + KvCacheConnectorManager, + KvCacheConnectorScheduler, + KvCacheConnectorWorker, +) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.config import ( + MooncakeStoreConnectorConfig, + StoreRole, +) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.scheduler import ( + MooncakeStoreConnectorScheduler, +) +from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.worker import ( + MooncakeStoreConnectorWorker, +) +from tensorrt_llm._torch.pyexecutor.py_executor import _validate_connector_admission +from tensorrt_llm._torch.pyexecutor.py_executor_creator import ( + _connector_runs_speculative_prefill_only, +) +from tensorrt_llm.llmapi.llm_args import ( + CacheTransceiverConfig, + Eagle3DecodingConfig, + MTPDecodingConfig, + NGramDecodingConfig, + SADecodingConfig, +) + +pytestmark = pytest.mark.cpu_only + +DISAGG = CacheTransceiverConfig(backend="DEFAULT") + + +class _Declaring: + supports_speculative_prefill = True + + +class _Plain: + pass + + +def _mtp_eagle() -> MTPDecodingConfig: + config = MTPDecodingConfig(max_draft_len=1) + assert config.spec_dec_mode.is_mtp_eagle_one_model() + return config + + +def _mtp_vanilla() -> MTPDecodingConfig: + # Not `num_nextn_predict_layers=2`: that deprecated input is remapped to + # max_draft_len and leaves the mode at MTP-Eagle. + config = MTPDecodingConfig(max_draft_len=2, use_mtp_vanilla=True) + assert config.spec_dec_mode.is_mtp_vanilla() + return config + + +def _eagle3() -> Eagle3DecodingConfig: + config = Eagle3DecodingConfig(max_draft_len=3, speculative_model="draft") + assert config.spec_dec_mode.is_eagle3_one_model() + return config + + +MOONCAKE = (MooncakeStoreConnectorWorker, MooncakeStoreConnectorScheduler) + + +@pytest.mark.parametrize("use_v2", [True, "auto"], ids=["v2", "auto"]) +@pytest.mark.parametrize( + "connector_classes", + [(_Declaring, _Declaring), MOONCAKE], + ids=["declaring", "mooncake_store"], +) +@pytest.mark.parametrize( + "make_spec_config", + [_mtp_eagle, _mtp_vanilla, _eagle3], + ids=["mtp_eagle", "mtp_vanilla", "eagle3"], +) +def test_one_model_drafter_runs_prefill_only_on_a_declaring_disagg_connector( + make_spec_config, connector_classes, use_v2 +) -> None: + worker_cls, scheduler_cls = connector_classes + assert _connector_runs_speculative_prefill_only( + make_spec_config(), worker_cls, scheduler_cls, DISAGG, use_v2 + ) + + +@pytest.mark.parametrize( + "spec_config, transceiver", + [ + (_eagle3(), None), + (_mtp_eagle(), None), + (_eagle3(), CacheTransceiverConfig()), + (NGramDecodingConfig(max_draft_len=3, max_matching_ngram_size=2), DISAGG), + ], + ids=["eagle3_aggregated", "mtp_aggregated", "transceiver_without_backend", "ngram"], +) +def test_mooncake_store_keeps_the_refusal_off_a_disaggregated_context_server( + spec_config, transceiver +) -> None: + assert not _connector_runs_speculative_prefill_only(spec_config, *MOONCAKE, transceiver, True) + + +def test_a_capacity_role_mooncake_server_runs_mtp_with_attention_dp() -> None: + """A disaggregated generation server lends memory and serves generation work. + + Its classes declare the mode, so one-model MTP on a disaggregated server + selects it, and attention DP is accepted too. A role that transfers no KV is + exempt from the context-only refusal, so generation requests are admitted. + """ + assert _connector_runs_speculative_prefill_only(_mtp_eagle(), *MOONCAKE, DISAGG, True) + KvCacheConnectorManager.validate_attention_dp(*MOONCAKE, pp_size=1, cp_size=1) + + # The halves as an attention-DP owner holds them, without opening a pool. + worker = object.__new__(MooncakeStoreConnectorWorker) + worker._config = MooncakeStoreConnectorConfig( + master_server_address="127.0.0.1:50051", role=StoreRole.CAPACITY + ) + scheduler = object.__new__(MooncakeStoreConnectorScheduler) + with patch( + "tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector.mpi_rank", return_value=1 + ): + manager = KvCacheConnectorManager( + worker, scheduler, enable_attention_dp=True, speculative_prefill_only=True + ) + assert manager.capacity_only is True + assert manager.speculative_prefill_only is True + generation_request = SimpleNamespace( + is_context_only_request=False, + is_generation_only_request=True, + is_dummy_request=False, + is_attention_dp_dummy=False, + ) + _validate_connector_admission(manager, generation_request) + + +@pytest.mark.parametrize( + "spec_config, worker_cls, scheduler_cls, transceiver, use_v2", + [ + (None, _Declaring, _Declaring, DISAGG, True), + (MTPDecodingConfig(), _Declaring, _Declaring, DISAGG, True), + ( + NGramDecodingConfig(max_draft_len=3, max_matching_ngram_size=2), + _Declaring, + _Declaring, + DISAGG, + True, + ), + (SADecodingConfig(max_draft_len=4), _Declaring, _Declaring, DISAGG, True), + (_eagle3(), _Declaring, _Plain, DISAGG, True), + (_eagle3(), _Plain, _Declaring, DISAGG, True), + (_eagle3(), KvCacheConnectorWorker, KvCacheConnectorScheduler, DISAGG, True), + (_eagle3(), _Declaring, _Declaring, None, True), + (_eagle3(), _Declaring, _Declaring, CacheTransceiverConfig(), True), + (_mtp_eagle(), _Declaring, _Declaring, None, "auto"), + (_mtp_eagle(), _Declaring, _Declaring, DISAGG, False), + (_eagle3(), _Declaring, _Declaring, DISAGG, False), + ], + ids=[ + "no_spec", + "unresolved_mtp_draft_len", + "ngram", + "suffix_automaton", + "worker_only_declares", + "scheduler_only_declares", + "base_classes", + "aggregated", + "transceiver_without_backend", + "mtp_aggregated", + "mtp_on_v1", + "eagle3_on_v1", + ], +) +def test_other_combinations_keep_the_speculative_refusal( + spec_config, worker_cls, scheduler_cls, transceiver, use_v2 +) -> None: + assert not _connector_runs_speculative_prefill_only( + spec_config, worker_cls, scheduler_cls, transceiver, use_v2 + ) + + +def test_declares_speculative_prefill() -> None: + assert KvCacheConnectorManager.declares_speculative_prefill(_Declaring, _Declaring) + # A TP follower rank holds no scheduler, so only its worker is checked. + assert KvCacheConnectorManager.declares_speculative_prefill(_Declaring, None) + assert not KvCacheConnectorManager.declares_speculative_prefill(_Declaring, _Plain) + assert not KvCacheConnectorManager.declares_speculative_prefill(_Plain, _Declaring) + assert not KvCacheConnectorManager.declares_speculative_prefill(_Plain, None) + assert not KvCacheConnectorManager.declares_speculative_prefill( + KvCacheConnectorWorker, KvCacheConnectorScheduler + ) + + +def test_only_an_explicit_true_declares() -> None: + class Truthy: + supports_speculative_prefill = "yes" + + assert not KvCacheConnectorManager.declares_speculative_prefill(Truthy, Truthy) diff --git a/tests/unittest/_torch/executor/test_remote_tail_adp_phase.py b/tests/unittest/_torch/executor/test_remote_tail_adp_phase.py index e403e1e1533f..2f4cc575b29e 100644 --- a/tests/unittest/_torch/executor/test_remote_tail_adp_phase.py +++ b/tests/unittest/_torch/executor/test_remote_tail_adp_phase.py @@ -24,7 +24,7 @@ from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState, SamplingConfig from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor from tensorrt_llm._torch.pyexecutor.resource_manager import ResourceManagerType -from tensorrt_llm._torch.pyexecutor.scheduler.adp_router import RankState +from tensorrt_llm._torch.pyexecutor.scheduler.adp_router import ADPRouter, RankState from tensorrt_llm._torch.pyexecutor.scheduler.scheduler import RemoteTailPhase, ScheduledRequests pytestmark = pytest.mark.cpu_only @@ -656,6 +656,16 @@ def start_worker(executor): events.append("start_worker") monkeypatch.setattr(PyExecutor, "start_worker", start_worker) + # The router factory must receive the executor's connector manager; a + # missing keyword is invisible to every router test. + real_create = ADPRouter.create + seen_create_kwargs = {} + + def recording_create(**kwargs): + seen_create_kwargs.update(kwargs) + return real_create(**kwargs) + + monkeypatch.setattr(ADPRouter, "create", recording_create) executor = PyExecutor( resource_manager, scheduler, @@ -665,6 +675,8 @@ def start_worker(executor): max_num_sequences=8, kv_cache_transceiver=transceiver, ) + assert "kv_connector_manager" in seen_create_kwargs + assert seen_create_kwargs["kv_connector_manager"] is None assert executor.kv_cache_transceiver is transceiver assert events == ( ["warmup", "reserve", "start_worker"] if should_enable else ["warmup", "start_worker"] diff --git a/tests/unittest/_torch/test_connector_adp.py b/tests/unittest/_torch/test_connector_adp.py index ad1e6829630e..8a116ba0fb81 100644 --- a/tests/unittest/_torch/test_connector_adp.py +++ b/tests/unittest/_torch/test_connector_adp.py @@ -6,7 +6,7 @@ import weakref from collections.abc import Iterator from concurrent.futures import ThreadPoolExecutor -from contextlib import nullcontext +from contextlib import ExitStack, nullcontext from functools import partial from threading import Barrier, Event from types import SimpleNamespace @@ -77,6 +77,46 @@ def _manager(rank: int = 1) -> KvCacheConnectorManager: return KvCacheConnectorManager(worker, scheduler, enable_attention_dp=True) +def _failed_load_manager( + rank: int = 1, + *, + adp: bool = True, + offer: tuple[int, bool] = (4, False), + stack: ExitStack | None = None, + world_size: int = 1, +) -> KvCacheConnectorManager: + """A manager whose worker hands failed prefix loads back for a restart. + + The query, update and finish callbacks of a tensor-parallel manager go + through the leader broadcast, and the restart guard reads the world size, + so a non-ADP manager needs its MPI patches for the whole test: they are + entered into ``stack``. + """ + worker = ADPAdapterMock(reports_failed_loads=True) + worker.get_finished.return_value = ([], []) + worker.take_failed_load_requests.return_value = set() + worker.take_failed_async_load_requests.return_value = set() + scheduler = ADPAdapterMock() + scheduler.get_num_new_matched_tokens.return_value = offer + scheduler.request_finished.return_value = True + if adp: + with patch(f"{CONNECTOR_MODULE}.mpi_rank", return_value=rank): + return KvCacheConnectorManager(worker, scheduler, enable_attention_dp=True) + assert stack is not None, "a tensor-parallel manager needs MPI patches for the whole test" + stack.enter_context(patch(f"{CONNECTOR_MODULE}.mpi_rank", return_value=0)) + stack.enter_context( + patch(f"{CONNECTOR_MODULE}.mpi_broadcast", side_effect=lambda result, root: result) + ) + stack.enter_context(patch(f"{CONNECTOR_MODULE}.mpi_world_size", return_value=world_size)) + return KvCacheConnectorManager(worker, scheduler) + + +@pytest.fixture +def exit_stack() -> Iterator[ExitStack]: + with ExitStack() as stack: + yield stack + + def _configure_recovery_controls(executor: "PyExecutor") -> None: from tensorrt_llm._torch.pyexecutor.executor_request_queue import ExecutorRequestQueue @@ -147,6 +187,23 @@ def test_adp_requires_scheduler_on_every_owner(rank: int) -> None: KvCacheConnectorManager(ADPAdapterMock(), None, enable_attention_dp=True) +def test_admission_query_is_off_under_attention_dp() -> None: + """Owners schedule independently, so no owner is asked while it admits a request.""" + with patch(f"{CONNECTOR_MODULE}.mpi_rank", return_value=1): + manager = KvCacheConnectorManager( + ADPAdapterMock(), ADPAdapterMock(), enable_attention_dp=True, admission_query=True + ) + assert manager.admission_query_capable is False + + +def test_admission_query_capability_under_tp() -> None: + with patch(f"{CONNECTOR_MODULE}.mpi_rank", return_value=0): + capable = KvCacheConnectorManager(ADPAdapterMock(), ADPAdapterMock(), admission_query=True) + default = KvCacheConnectorManager(ADPAdapterMock(), ADPAdapterMock()) + assert capable.admission_query_capable is True + assert default.admission_query_capable is False + + @pytest.mark.parametrize("request_counts", [(2, 2), (3, 0), (0, 4)]) def test_disjoint_owner_callbacks_and_metadata( forbid_connector_collectives: None, request_counts: tuple[int, int] @@ -228,7 +285,80 @@ def test_dummy_requests_never_reach_storage(forbid_connector_collectives: None) batch.context_requests_last_chunk = [dummy] batch.generation_requests = [dummy] cache = MagicMock(spec=KVCacheManagerV2) - output_manager = KvCacheConnectorSchedulerOutputManager(enable_attention_dp=True) + output_manager = KvCacheConnectorSchedulerOutputManager() + output = output_manager.build_scheduler_output(batch, AsyncRequests({}, {}), cache) + assert not output.new_requests and not output.cached_requests + assert not output_manager.requests + cache.get_page_indices_by_layer_group.assert_not_called() + + +def test_mock_requests_are_not_mistaken_for_dummies(forbid_connector_collectives: None) -> None: + """Only an explicit True marks a dummy; a mock answers every attribute truthily.""" + manager = _manager() + request = MagicMock() + request.request_id = 5 + manager.scheduler.request_finished.return_value = True + assert manager.request_finished(request, [1]) + manager.scheduler.request_finished.assert_called_once_with(request, [1]) + + +@pytest.mark.parametrize( + "is_dummy_request,is_attention_dp_dummy", + [(True, False), (False, True)], + ids=["is_dummy_request", "is_attention_dp_dummy"], +) +def test_tp_manager_skips_dummy_requests( + is_dummy_request: bool, is_attention_dp_dummy: bool +) -> None: + """Without attention DP, either dummy flag keeps a request from the connector. + + The early returns sit before the leader broadcast, and every TP rank sees + the same flags, so no rank enters a collective the others skip. + """ + worker, scheduler = MagicMock(), MagicMock() + dummy = _request(999) + dummy.is_dummy_request = is_dummy_request + dummy.is_attention_dp_dummy = is_attention_dp_dummy + real = _request(7) + real.is_attention_dp_dummy = False + with ( + patch(f"{CONNECTOR_MODULE}.mpi_rank", return_value=0), + patch( + f"{CONNECTOR_MODULE}.mpi_broadcast", side_effect=lambda result, root: result + ) as broadcast, + ): + manager = KvCacheConnectorManager(worker, scheduler) + assert not manager.enable_attention_dp + assert manager.query_num_new_matched_tokens(dummy, 0) == (0, False) + assert manager.get_num_new_matched_tokens(dummy, 0) == 0 + manager.commit_new_matched_tokens(dummy, 4, True) + manager.update_state_after_alloc(dummy, [5]) + manager.update_state_after_alloc(dummy, [], [[3], [5]]) + assert not manager.request_finished(dummy, [5]) + manager.prefix_reservations_enabled = True + assert manager.reserve_prefix(dummy, 0) is None + manager.prefix_reservations_enabled = False + broadcast.assert_not_called() + + # A request with neither flag set still reaches the leader. + scheduler.get_num_new_matched_tokens.return_value = (0, False) + assert manager.query_num_new_matched_tokens(real, 0) == (0, False) + scheduler.get_num_new_matched_tokens.assert_called_once_with(real, 0) + broadcast.assert_called_once() + + scheduler.update_state_after_alloc.assert_not_called() + scheduler.update_state_after_alloc_by_layer_group.assert_not_called() + scheduler.request_finished.assert_not_called() + scheduler.request_finished_by_layer_group.assert_not_called() + scheduler.reserve_prefix.assert_not_called() + assert not manager.scheduler_output_manager.external_loads + assert not manager.has_pending_transfers(999) + + batch = ScheduledRequests() + batch.context_requests_last_chunk = [dummy] + batch.generation_requests = [dummy] + cache = MagicMock(spec=KVCacheManagerV2) + output_manager = KvCacheConnectorSchedulerOutputManager() output = output_manager.build_scheduler_output(batch, AsyncRequests({}, {}), cache) assert not output.new_requests and not output.cached_requests assert not output_manager.requests @@ -1629,3 +1759,461 @@ def test_phase_idle_connector_owner_joins_ready_peer( assert results == [(peer_has_work, peer_has_work), (peer_has_work, False)] for executor in executors: executor.dist.tp_allgather_int64.assert_called_once() + + +# --------------------------------------------------------------------------- +# Failed prefix loads handed back for a restart (reports_failed_loads). +# --------------------------------------------------------------------------- + + +def test_failed_loads_and_restarts_reach_the_worker(forbid_connector_collectives: None) -> None: + manager = _failed_load_manager(offer=(4, True)) + parked = _request(8) + assert manager.get_num_new_matched_tokens(parked, 0) == 4 + batch = ScheduledRequests() + batch.context_requests_last_chunk = [parked] + manager.take_scheduled_requests_pending_load(batch) + manager.worker.get_finished.return_value = ([], [8]) + assert manager.get_finished() == [] + assert 8 in manager.finished_async_loading_requests + manager.worker.take_failed_load_requests.return_value = {7} + manager.worker.take_failed_async_load_requests.return_value = {8} + + assert manager.take_failed_load_requests() == {7} + assert manager.take_failed_async_load_requests() == [parked] + assert 8 not in manager.finished_async_loading_requests + + request = _request(7) + manager.forget_request(request) + manager.worker.drop_bound_saves.assert_called_once_with([7]) + # The allocation release that follows resets the connector state. + manager.scheduler.request_reset.assert_not_called() + + +def test_failed_loads_are_reported_only_by_a_connector_that_declares_it() -> None: + assert _failed_load_manager().reports_failed_loads + with ( + patch(f"{CONNECTOR_MODULE}.mpi_rank", return_value=0), + patch(f"{CONNECTOR_MODULE}.mpi_broadcast", side_effect=lambda result, root: result), + ): + # A mock answers every attribute truthily; only a declared True counts. + assert not KvCacheConnectorManager(MagicMock(), MagicMock()).reports_failed_loads + # A role that loads nothing has nothing to restart. + capacity_only = MagicMock(reports_failed_loads=True, capacity_only=True) + assert not KvCacheConnectorManager(capacity_only, MagicMock()).reports_failed_loads + + +def test_failed_load_sets_are_empty_for_a_connector_without_them() -> None: + class PlainWorker: + capacity_only = False + + with ( + patch(f"{CONNECTOR_MODULE}.mpi_rank", return_value=0), + patch(f"{CONNECTOR_MODULE}.mpi_broadcast", side_effect=lambda result, root: result), + ): + manager = KvCacheConnectorManager(PlainWorker(), MagicMock()) + assert manager.take_failed_load_requests() == set() + assert manager.take_failed_async_load_requests() == [] + + +@pytest.mark.parametrize("adp", [True, False], ids=["adp", "single_rank"]) +def test_failed_sync_load_restarts_the_request_and_pads_the_batch( + exit_stack: ExitStack, adp: bool +) -> None: + from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor + + manager = _failed_load_manager(adp=adp, stack=exit_stack) + req = _request(7) + assert manager.get_num_new_matched_tokens(req, 0) == 4 + batch = ScheduledRequests() + batch.context_requests_last_chunk = [req] + manager.set_scheduler_output(SchedulerOutput()) + manager.worker.take_failed_load_requests.return_value = {7} + + executor = object.__new__(PyExecutor) + _configure_recovery_controls(executor) + executor.kv_connector_manager = manager + executor.kv_cache_manager = MagicMock(spec=KVCacheManagerV2) + executor.resource_manager = MagicMock() + executor.resource_manager.get_resource_manager.return_value = None + padded = [] + executor._pad_empty_attention_dp_batch = lambda scheduled: padded.append(scheduled.batch_size) + if not adp: + executor.enable_attention_dp = False + executor._can_queue = MagicMock(return_value=(False, False)) + + with patch("torch.cuda.current_stream"): + assert executor._kv_connector_start_batch(batch) == (False, False) + + # The request left the batch, its saves and allocation were dropped, and it + # is admitted again rather than skipped as an already-loaded request. + assert batch.context_requests == [] + executor.kv_cache_manager.drop_context_allocation.assert_called_once_with(req) + manager.worker.drop_bound_saves.assert_called_with([7]) + manager.worker.start_load_kv.assert_called_once() + assert manager.should_add_sequence(req) + assert not executor._kv_connector_failed + # Under attention DP the emptied batch is padded so other owners keep running. + assert padded == ([0] if adp else []) + + +def test_failed_async_load_restarts_the_parked_request( + forbid_connector_collectives: None, +) -> None: + from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor + + manager = _failed_load_manager(offer=(4, True)) + req = _request(8) + assert manager.get_num_new_matched_tokens(req, 0) == 4 + batch = ScheduledRequests() + batch.context_requests_last_chunk = [req] + manager.take_scheduled_requests_pending_load(batch) + assert batch.batch_size == 0 + assert req.state == LlmRequestState.DISAGG_GENERATION_TRANS_IN_PROGRESS + manager.worker.get_finished.return_value = ([], [8]) + manager.worker.take_failed_async_load_requests.return_value = {8} + + executor = object.__new__(PyExecutor) + _configure_recovery_controls(executor) + executor.kv_connector_manager = manager + executor.kv_cache_manager = MagicMock(spec=KVCacheManagerV2) + spec_manager = MagicMock() + executor.resource_manager = MagicMock() + executor.resource_manager.get_resource_manager.return_value = spec_manager + + executor._kv_connector_terminate_requests(synchronize=False) + + executor.kv_cache_manager.drop_context_allocation.assert_called_once_with(req) + spec_manager.free_resources.assert_called_once_with(req) + manager.worker.drop_bound_saves.assert_called_once_with([8]) + assert req.state == LlmRequestState.CONTEXT_INIT + # It is asked about again rather than skipped as an already-loaded request. + assert manager.should_add_sequence(req) + assert not manager.has_pending_transfers(8) + + +def test_failed_async_load_of_a_finished_request_restarts_nothing_silently( + forbid_connector_collectives: None, +) -> None: + """A failure reported for a request that already left is not a restart.""" + from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor + from tensorrt_llm.logger import logger + + manager = _failed_load_manager() + manager.worker.take_failed_async_load_requests.return_value = {9} + executor = object.__new__(PyExecutor) + _configure_recovery_controls(executor) + executor.kv_connector_manager = manager + executor.kv_cache_manager = MagicMock(spec=KVCacheManagerV2) + executor.resource_manager = MagicMock() + + with patch.object(logger, "warning") as warning: + executor._kv_connector_terminate_requests(synchronize=False) + + executor.kv_cache_manager.drop_context_allocation.assert_not_called() + manager.worker.drop_bound_saves.assert_not_called() + warning.assert_not_called() + + +def test_failed_async_load_in_adp_recovery_loop_is_not_resumed( + forbid_connector_collectives: None, +) -> None: + """A parked request restarted mid-recovery stays out of the batch it left.""" + from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor + + manager = _failed_load_manager(offer=(4, True)) + failed, ready = _request(10), _request(11) + batch = ScheduledRequests() + batch.context_requests_last_chunk = [failed, ready] + cache = MagicMock(spec=KVCacheManagerV2) + cache.get_page_indices_by_layer_group.return_value = [[5, 6]] + for req in batch.context_requests: + manager.get_num_new_matched_tokens(req, 0) + req.context_current_position = 4 + req.context_remaining_length = 4 + req.context_chunk_size = 4 + manager.update_state_after_alloc(req, [5, 6]) + + # Mirror the executor's already-prepared batch and initial load plan. + manager.build_scheduler_output(batch, cache) + manager.handle_metadata() + executor = object.__new__(PyExecutor) + _configure_recovery_controls(executor) + executor.kv_connector_manager = manager + executor.kv_cache_manager = cache + executor.kv_cache_transceiver = None + executor.active_requests = [failed, ready] + executor.expected_num_active_requests = 2 + # No sequence slot is left for a padding dummy. + executor.max_num_active_requests = 2 + executor._count_schedulable_active_requests = MagicMock(return_value=0) + executor._should_skip_dummy_for_benchmark_disagg = MagicMock(return_value=False) + executor._has_adp_dummy_kv_capacity = MagicMock(return_value=False) + executor.resource_manager = MagicMock() + executor.resource_manager.get_resource_manager.return_value = None + executor.resource_manager.prepare_resources.side_effect = AssertionError("double allocation") + manager.worker.get_finished.side_effect = [([], []), ([], [10]), ([], [11])] + manager.worker.take_failed_async_load_requests.side_effect = [set(), {10}, set()] + + with ( + patch("torch.cuda.current_stream"), + patch("tensorrt_llm._torch.pyexecutor.py_executor.time.sleep") as sleep, + ): + executor._kv_connector_start_batch(batch) + + assert batch.context_requests == [ready] + assert failed.state == LlmRequestState.CONTEXT_INIT + assert manager.should_add_sequence(failed) + assert not manager.should_add_sequence(ready) + cache.drop_context_allocation.assert_called_once_with(failed) + manager.worker.drop_bound_saves.assert_called_once_with([10]) + assert manager.worker.start_load_kv.call_count == 2 + executor.resource_manager.prepare_resources.assert_not_called() + assert sleep.call_count == 2 + assert not executor._kv_connector_failed + compute_plan = manager.scheduler.build_connector_meta.call_args.args[0] + assert [req.request_id for req in compute_plan.new_requests] == [11] + + +@pytest.mark.parametrize("pad_ok", [False, True]) +def test_restart_of_the_only_parked_request_ends_the_wait( + forbid_connector_collectives: None, pad_ok: bool +) -> None: + """Nothing is left to wait for once the only parked request restarts.""" + from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor + + manager = _failed_load_manager(offer=(4, True)) + req = _request(10) + manager.get_num_new_matched_tokens(req, 0) + manager.update_state_after_alloc(req, [5, 6]) + batch = ScheduledRequests() + batch.context_requests_last_chunk = [req] + cache = MagicMock(spec=KVCacheManagerV2) + cache.get_page_indices_by_layer_group.return_value = [[5, 6]] + manager.build_scheduler_output(batch, cache) + + executor = object.__new__(PyExecutor) + _configure_recovery_controls(executor) + executor.kv_connector_manager = manager + executor.kv_cache_manager = cache + executor.resource_manager = MagicMock() + executor.resource_manager.get_resource_manager.return_value = None + dummy = _request(999, dummy=True) + pad_calls = 0 + + def pad(scheduled: ScheduledRequests) -> None: + # Padding fails while the parked request holds its pages, and succeeds + # afterwards only if the freed pages are enough. + nonlocal pad_calls + pad_calls += 1 + assert scheduled.batch_size == 0 + if pad_calls > 1 and pad_ok: + scheduled.generation_requests.append(dummy) + + executor._pad_empty_attention_dp_batch = pad + manager.worker.get_finished.return_value = ([], [10]) + manager.worker.take_failed_async_load_requests.side_effect = [{10}] + + with ( + patch("torch.cuda.current_stream"), + patch("tensorrt_llm._torch.pyexecutor.py_executor.time.sleep") as sleep, + ): + result = executor._kv_connector_start_batch(batch) + + assert result == (pad_ok, pad_ok) + sleep.assert_not_called() + assert pad_calls == 2 + assert batch.context_requests == [] + assert batch.generation_requests == ([dummy] if pad_ok else []) + cache.drop_context_allocation.assert_called_once_with(req) + assert req.state == LlmRequestState.CONTEXT_INIT + assert manager.should_add_sequence(req) + assert not executor._kv_connector_failed + if pad_ok: + executor.resource_manager.prepare_resources.assert_called_once_with( + batch, publish_connector_output=False + ) + else: + executor.resource_manager.prepare_resources.assert_not_called() + + +def test_failed_load_report_under_tp_is_fatal(exit_stack: ExitStack) -> None: + manager = _failed_load_manager(adp=False, stack=exit_stack, world_size=2) + manager.worker.take_failed_load_requests.return_value = {7} + with pytest.raises(RuntimeError, match="tensor-parallel"): + manager.take_failed_load_requests() + manager.worker.take_failed_async_load_requests.return_value = {8} + with pytest.raises(RuntimeError, match="tensor-parallel"): + manager.take_failed_async_load_requests() + + +class _ReportingWorker(KvCacheConnectorWorker): + """Every abstract method, and `take_failed_load_requests`, but no `drop_bound_saves`.""" + + reports_failed_loads = True + + def register_kv_caches(self, kv_cache_tensor: object) -> None: + pass + + def start_load_kv(self, stream: object) -> None: + pass + + def wait_for_layer_load(self, layer_idx: int, stream: object) -> None: + pass + + def save_kv_layer(self, layer_idx: int, stream: object) -> None: + pass + + def wait_for_save(self, stream: object) -> None: + pass + + def get_finished( + self, finished_gen_req_ids: list[int], started_loading_req_ids: list[int] + ) -> tuple[list[int], list[int]]: + return [], [] + + def take_failed_load_requests(self) -> set[int]: + return set() + + +def test_reports_failed_loads_requires_its_hooks() -> None: + class CompleteWorker(_ReportingWorker): + def drop_bound_saves(self, request_ids: object) -> int: + return 0 + + with ( + patch(f"{CONNECTOR_MODULE}.mpi_rank", return_value=0), + patch(f"{CONNECTOR_MODULE}.mpi_broadcast", side_effect=lambda result, root: result), + ): + with pytest.raises(ValueError, match="does not implement drop_bound_saves;"): + KvCacheConnectorManager(_ReportingWorker(MagicMock()), MagicMock()) + assert KvCacheConnectorManager( + CompleteWorker(MagicMock()), MagicMock() + ).reports_failed_loads + with pytest.raises(NotImplementedError, match="drop_bound_saves"): + _ReportingWorker(MagicMock()).drop_bound_saves([7]) + + +def test_reports_failed_loads_rejects_prefix_reservations( + forbid_connector_collectives: None, +) -> None: + class ReservingScheduler(ADPAdapterMock): + def reserve_prefix( + self, request: object, num_computed_tokens: int, reservation_id: int + ) -> tuple[int, bool]: + return 0, False + + def release_prefix_reservation( + self, request: object, reservation_id: int, start: int, end: int + ) -> None: + pass + + class ReservingWorker(ADPAdapterMock): + reports_failed_loads = True + + def get_finished_prefix_loads(self) -> list[int]: + return [] + + with patch(f"{CONNECTOR_MODULE}.mpi_rank", return_value=1): + manager = KvCacheConnectorManager( + ReservingWorker(), ReservingScheduler(), enable_attention_dp=True + ) + assert manager.reports_failed_loads + with pytest.raises(ValueError, match="reports_failed_loads"): + manager.configure_prefix_reservations(True) + # Without the protocol asked for, the final-batch query path stays usable. + manager.configure_prefix_reservations(False) + assert not manager.prefix_reservations_enabled + + +def test_reports_failed_loads_requires_kv_cache_manager_v2(exit_stack: ExitStack) -> None: + """Restarts are a KVCacheManagerV2 feature; a resolved V1 manager is refused.""" + from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor + + executor = object.__new__(PyExecutor) + executor.kv_connector_manager = _failed_load_manager(adp=False, stack=exit_stack) + executor.kv_cache_transceiver = None + executor.dist = SimpleNamespace(pp_size=1, cp_size=1) + executor.max_beam_width = 1 + executor.enable_attention_dp = False + executor.kv_cache_manager = object() + + with pytest.raises( + NotImplementedError, match="reports failed prefix loads requires KVCacheManagerV2" + ): + executor._maybe_init_kv_connector_manager() + + +def test_failed_sync_load_for_a_request_outside_the_batch_is_fatal( + forbid_connector_collectives: None, +) -> None: + """A misreported ID is a connector failure, not a restart.""" + from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor + + manager = _failed_load_manager() + req = _request(7) + manager.get_num_new_matched_tokens(req, 0) + batch = ScheduledRequests() + batch.context_requests_last_chunk = [req] + manager.set_scheduler_output(SchedulerOutput()) + manager.worker.take_failed_load_requests.return_value = {99} + + executor = object.__new__(PyExecutor) + _configure_recovery_controls(executor) + executor.kv_connector_manager = manager + executor.kv_cache_manager = MagicMock(spec=KVCacheManagerV2) + executor.resource_manager = MagicMock() + executor._pad_empty_attention_dp_batch = MagicMock() + + with ( + patch("torch.cuda.current_stream"), + pytest.raises(RuntimeError, match="KV connector failed.*not in this pass's context batch"), + ): + executor._kv_connector_start_batch(batch) + + assert executor._kv_connector_failed + assert batch.context_requests == [req] + executor.kv_cache_manager.drop_context_allocation.assert_not_called() + manager.worker.drop_bound_saves.assert_not_called() + + +def test_failed_async_load_still_pending_is_fatal(forbid_connector_collectives: None) -> None: + """The pages of a load the worker has not finished may still be written.""" + manager = _failed_load_manager(offer=(4, True)) + req = _request(8) + manager.get_num_new_matched_tokens(req, 0) + batch = ScheduledRequests() + batch.context_requests_last_chunk = [req] + manager.take_scheduled_requests_pending_load(batch) + assert manager.get_finished() == [] + assert manager.has_pending_transfers(8) + manager.worker.take_failed_async_load_requests.return_value = {8} + + with pytest.raises(RuntimeError, match="not reported finished"): + manager.take_failed_async_load_requests() + assert manager.has_pending_transfers(8) + + +def test_failed_async_load_reported_after_its_finish_is_fatal( + forbid_connector_collectives: None, +) -> None: + """A report after the finish comes too late: the request may already be scheduled.""" + manager = _failed_load_manager(offer=(4, True)) + req = _request(8) + manager.get_num_new_matched_tokens(req, 0) + batch = ScheduledRequests() + batch.context_requests_last_chunk = [req] + manager.take_scheduled_requests_pending_load(batch) + manager.worker.get_finished.return_value = ([], [8]) + assert manager.get_finished() == [] + assert manager.take_failed_async_load_requests() == [] + assert not manager.should_add_sequence(req) + + # The scheduler can now admit the request over the loaded pages; the next + # poll names it anyway. + manager.worker.get_finished.return_value = ([], []) + assert manager.get_finished() == [] + manager.worker.take_failed_async_load_requests.return_value = {8} + with pytest.raises(RuntimeError, match="earlier get_finished returned them as finished"): + manager.take_failed_async_load_requests() + assert 8 in manager.finished_async_loading_requests diff --git a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py index b2c33ae83ce3..c893054ee376 100644 --- a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py +++ b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py @@ -175,6 +175,9 @@ class KvCacheConfigV2: # tensorrt_llm/llmapi/llm_args.py; KVCacheManagerV2._build_base_config() # reads it unconditionally. kv_pool_rebalance_config: KvPoolRebalanceConfig = field(default_factory=KvPoolRebalanceConfig) + # Mirrors KvCacheConfig.enable_disagg_completion_reservation in + # tensorrt_llm/llmapi/llm_args.py; KVCacheManagerV2.__init__() reads it. + enable_disagg_completion_reservation: bool = False # --------------------------------------------------------------------------- diff --git a/tests/unittest/disaggregated/test_disagg_openai_client.py b/tests/unittest/disaggregated/test_disagg_openai_client.py index a1bf5ff66e27..a6a2fae63a2a 100644 --- a/tests/unittest/disaggregated/test_disagg_openai_client.py +++ b/tests/unittest/disaggregated/test_disagg_openai_client.py @@ -15,6 +15,7 @@ import asyncio import errno import json +import re from unittest.mock import AsyncMock, MagicMock, Mock, patch import aiohttp @@ -33,6 +34,7 @@ ) from tensorrt_llm.serve.openai_client import OpenAIHttpClient, _encode_request_body from tensorrt_llm.serve.openai_protocol import ( + ChatCompletionRequest, CompletionRequest, CompletionResponse, CompletionResponseChoice, @@ -909,8 +911,9 @@ async def test_retry_affinity_signature_matches_wire_request( internal_disagg_auth_key="secret", disagg_id_generator=AsyncMock(return_value=1000) if regenerate_id else None, ) + # A failure that leaves the request undelivered, so both roles retry. session.post.side_effect = [ - aiohttp.ClientError("transient"), + ConnectionResetError(), self._mock_http_ok(self._ok_response()), ] request = CompletionRequest( @@ -928,7 +931,10 @@ async def test_retry_affinity_signature_matches_wire_request( ) await client.send_request(request) - expected_ids = [42, 1000 if regenerate_id else 42] + # Only a context retry regenerates the request ID; a generation leg keeps + # the ID that names its context KV cache. + regenerated = regenerate_id and role == ServerRole.CONTEXT + expected_ids = [42, 1000 if regenerated else 42] # Serialized attempts retain their own IDs when a context retry # regenerates the request ID. @@ -954,7 +960,7 @@ async def test_retry_affinity_signature_matches_wire_request( assert ( headers[0][INTERNAL_DISAGG_AUTH_HEADER] == headers[1][INTERNAL_DISAGG_AUTH_HEADER] ) - if regenerate_id: + if regenerated: assert ( headers[0][SUBAGENT_AFFINITY_AUTH_HEADER] != headers[1][SUBAGENT_AFFINITY_AUTH_HEADER] @@ -1598,9 +1604,6 @@ async def test_remote_tail_handoff_does_not_retry_or_change_id( "role,request_type,schedule_style,bounded_replay", [ (ServerRole.GENERATION, None, None, True), - (ServerRole.GENERATION, "generation_only", None, False), - (ServerRole.GENERATION, "generation_only", DisaggScheduleStyle.CONTEXT_FIRST, False), - (ServerRole.GENERATION, "generation_only", DisaggScheduleStyle.GENERATION_FIRST, False), (ServerRole.CONTEXT, "context_only", DisaggScheduleStyle.GENERATION_FIRST, False), (ServerRole.CONTEXT, "context_only", DisaggScheduleStyle.CONTEXT_FIRST, True), ], @@ -1655,3 +1658,315 @@ async def test_other_requests_retain_retry( mock_router.finish_request.assert_awaited_once_with( completion_request, mock_session, success=True, req_id=42 ) + + +class TestGenerationNoReplayAfterAccept: + """A generation leg is never replayed once the generation server may hold it. + + The context server's KV-cache handoff is one-shot and is named by + ``disagg_request_id``. When the transfer fails, the generation server raises + in its stream generator after the response headers and drops the + connection. A replay under a new id then waits for KV cache that no context + server holds, and one under the same id can collide with the original on the + worker, so the request fails instead. A leg that cannot have reached the + server (connection refused, or reset while it was being written) is still + re-sent, under the same id. + """ + + _GEN_ID = 42 + + def _make_client(self, role, session, max_retries=1, disagg_id_generator=None): + _reset_prometheus_registry() + router = AsyncMock(spec=Router) + router.servers = ["localhost:8000"] + router.get_next_server = AsyncMock(return_value=("localhost:8000", None)) + router.finish_request = AsyncMock() + client = OpenAIHttpClient( + router=router, + role=role, + timeout_secs=30, + max_retries=max_retries, + retry_interval_sec=0, + session=session, + disagg_id_generator=disagg_id_generator, + ) + client._metrics_collector = MagicMock() + return client + + def _gen_request(self, kind, stream, schedule_style=None): + disaggregated_params = DisaggregatedParams( + request_type="generation_only", + first_gen_tokens=[7], + ctx_request_id=self._GEN_ID, + disagg_request_id=self._GEN_ID, + schedule_style=schedule_style, + ) + if kind == "completion": + return CompletionRequest( + model="m", prompt="hi", stream=stream, disaggregated_params=disaggregated_params + ) + if kind == "chat": + return ChatCompletionRequest( + model="m", + messages=[{"role": "user", "content": "hi"}], + stream=stream, + disaggregated_params=disaggregated_params, + ) + assert kind == "responses" + return ResponsesRequest( + model="m", input="hi", stream=stream, disaggregated_params=disaggregated_params + ) + + def _mock_http_response(self, status, body): + r = AsyncMock() + r.status = status + r.reason = "OK" if status == 200 else "Internal Server Error" + r.headers = {"Content-Type": "application/json"} + r.json = AsyncMock(return_value=body) + r.text = AsyncMock(return_value=json.dumps(body)) + r.request_info = MagicMock() + r.history = () + r.__aenter__ = AsyncMock(return_value=r) + r.__aexit__ = AsyncMock(return_value=False) + return r + + @staticmethod + def _mock_ok_stream(): + async def iter_any(): + yield b'data: "ok"\n\n' + + r = AsyncMock() + r.status = 200 + r.headers = {"Content-Type": "text/event-stream"} + r.content = AsyncMock() + r.content.iter_any = iter_any + r.__aenter__ = AsyncMock(return_value=r) + r.__aexit__ = AsyncMock(return_value=False) + return r + + def _ok_body(self): + return CompletionResponse( + model="m", + usage=UsageInfo(prompt_tokens=1, completion_tokens=1), + choices=[CompletionResponseChoice(index=0, text="ok")], + ).model_dump() + + @staticmethod + def _wire_ids(session): + return [ + msgspec.msgpack.decode(call.kwargs["data"])["disaggregated_params"]["disagg_request_id"] + for call in session.post.call_args_list + ] + + @pytest.mark.asyncio + @pytest.mark.parametrize("kind", ["completion", "chat", "responses"]) + async def test_gen_kv_transfer_failure_fails_request_without_replay(self, kind): + attempts = [] + replay_blocker = asyncio.Event() + + async def generation_server(reader, writer): + head = await reader.readuntil(b"\r\n\r\n") + length = int(re.search(rb"(?i)content-length: *(\d+)", head).group(1)) + raw_body = await reader.readexactly(length) + if re.search(rb"(?i)x-trtllm-msgpack: *1", head): + body = msgspec.msgpack.decode(raw_body) + else: + body = json.loads(raw_body) + attempts.append(body["disaggregated_params"]["disagg_request_id"]) + if len(attempts) > 1: + # No context server holds KV for a replayed request. + await replay_blocker.wait() + # What uvicorn does when the stream generator raises + # "Error in kv cache transfer for generation requests" after the headers. + writer.write( + b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\n" + b"Transfer-Encoding: chunked\r\n\r\n" + ) + await writer.drain() + writer.close() + + server = await asyncio.start_server(generation_server, "127.0.0.1", 0) + port = server.sockets[0].getsockname()[1] + ids = iter(range(1000, 2000)) + + async def next_id(): + return next(ids) + + session = aiohttp.ClientSession() + try: + client = self._make_client(ServerRole.GENERATION, session, disagg_id_generator=next_id) + request = self._gen_request(kind, stream=True) + response_generator = await client.send_request(request, server=f"127.0.0.1:{port}") + + async def consume(): + return [chunk async for chunk in response_generator] + + with pytest.raises(aiohttp.ClientPayloadError): + await asyncio.wait_for(consume(), timeout=10) + finally: + replay_blocker.set() + await session.close() + server.close() + await server.wait_closed() + + assert attempts == [self._GEN_ID] + assert request.disaggregated_params.disagg_request_id == self._GEN_ID + client._router.finish_request.assert_awaited_once_with(request, session, success=False) + client._metrics_collector.retry_requests.inc.assert_not_called() + + @pytest.mark.asyncio + @pytest.mark.parametrize("stream", [False, True]) + async def test_gen_error_status_is_not_replayed(self, stream): + session = AsyncMock(spec=aiohttp.ClientSession) + client = self._make_client(ServerRole.GENERATION, session) + session.post.side_effect = [ + self._mock_http_response(500, {"error": "Error in kv cache transfer"}), + self._mock_ok_stream() if stream else self._mock_http_response(200, self._ok_body()), + ] + request = self._gen_request("completion", stream=stream) + + with patch("asyncio.sleep", new_callable=AsyncMock): + with pytest.raises(aiohttp.ClientResponseError) as exc_info: + response = await client.send_request(request) + if stream: + # A streaming leg answered with an error status raises the + # worker's status, not "Response is not streaming". + _ = [chunk async for chunk in response] + + assert exc_info.value.status == 500 + assert session.post.call_count == 1 + client._router.finish_request.assert_awaited_once_with(request, session, success=False) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "schedule_style", + [None, DisaggScheduleStyle.CONTEXT_FIRST, DisaggScheduleStyle.GENERATION_FIRST], + ids=["no_style", "context_first", "generation_first"], + ) + @pytest.mark.parametrize( + "error", + [ + ConnectionResetError("Cannot write to closing transport"), + aiohttp.ClientConnectorError( + Mock(), ConnectionRefusedError(errno.ECONNREFUSED, "Connection refused") + ), + ], + ids=["reset_while_writing", "connection_refused"], + ) + async def test_gen_undelivered_request_is_retried_with_same_id(self, error, schedule_style): + session = AsyncMock(spec=aiohttp.ClientSession) + client = self._make_client( + ServerRole.GENERATION, session, disagg_id_generator=AsyncMock(return_value=999) + ) + session.post.side_effect = [error, self._mock_http_response(200, self._ok_body())] + request = self._gen_request("completion", stream=False, schedule_style=schedule_style) + hooks = MagicMock(spec=ResponseHooks) + + with patch("asyncio.sleep", new_callable=AsyncMock): + response = await client.send_request(request, hooks=hooks) + + assert isinstance(response, CompletionResponse) + assert session.post.call_count == 2 + assert self._wire_ids(session) == [self._GEN_ID, self._GEN_ID] + client._disagg_id_generator.assert_not_awaited() + hooks.on_disagg_request_id.assert_not_called() + assert client._router.renew_request.await_count == 2 + client._router.finish_request.assert_awaited_once_with(request, session, success=True) + + @pytest.mark.asyncio + @pytest.mark.parametrize("stream", [False, True]) + @pytest.mark.parametrize( + "error", + [ + asyncio.TimeoutError(), + aiohttp.ServerTimeoutError("t"), + aiohttp.ServerDisconnectedError(), + aiohttp.ClientOSError(errno.ECONNRESET, "reset"), + aiohttp.ClientConnectionError("x"), + aiohttp.ClientPayloadError("p"), + ], + ids=[ + "timeout", + "server_timeout", + "server_disconnected", + "connection_reset", + "connection_error", + "payload_error", + ], + ) + async def test_gen_possibly_delivered_request_is_not_retried(self, error, stream): + session = AsyncMock(spec=aiohttp.ClientSession) + client = self._make_client(ServerRole.GENERATION, session, max_retries=1) + session.post.side_effect = error + request = self._gen_request("completion", stream=stream) + + with patch("asyncio.sleep", new_callable=AsyncMock): + with pytest.raises(type(error)): + response = await client.send_request(request) + if stream: + _ = [chunk async for chunk in response] + + # Neither max_retries nor the transient-TCP budget re-posts the leg. + assert session.post.call_count == 1 + client._metrics_collector.retry_requests.inc.assert_not_called() + client._metrics_collector.error_requests.inc.assert_called_once_with() + client._router.finish_request.assert_awaited_once_with(request, session, success=False) + + @pytest.mark.asyncio + async def test_gen_undelivered_reset_keeps_transient_budget(self): + session = AsyncMock(spec=aiohttp.ClientSession) + client = self._make_client( + ServerRole.GENERATION, + session, + max_retries=0, + disagg_id_generator=AsyncMock(return_value=999), + ) + session.post.side_effect = [ + ConnectionResetError(), + ConnectionResetError(), + self._mock_http_response(200, self._ok_body()), + ] + request = self._gen_request("completion", stream=False) + + with patch("asyncio.sleep", new_callable=AsyncMock): + response = await client.send_request(request) + + assert isinstance(response, CompletionResponse) + assert session.post.call_count == 3 + assert self._wire_ids(session) == [self._GEN_ID] * 3 + + @pytest.mark.asyncio + async def test_ctx_error_status_is_still_retried(self): + session = AsyncMock(spec=aiohttp.ClientSession) + client = self._make_client(ServerRole.CONTEXT, session) + session.post.side_effect = [ + self._mock_http_response(500, {"error": "transient"}), + self._mock_http_response(200, self._ok_body()), + ] + request = CompletionRequest( + model="m", + prompt="hi", + stream=False, + disaggregated_params=DisaggregatedParams( + request_type="context_only", disagg_request_id=1 + ), + ) + + with patch("asyncio.sleep", new_callable=AsyncMock): + response = await client.send_request(request) + + assert isinstance(response, CompletionResponse) + assert session.post.call_count == 2 + + @pytest.mark.asyncio + async def test_no_retry_env_wins_for_undelivered_generation_leg(self, monkeypatch): + # The client reads the flag when it is constructed. + monkeypatch.setenv("TRTLLM_DISAGG_NO_RETRY", "1") + session = AsyncMock(spec=aiohttp.ClientSession) + client = self._make_client(ServerRole.GENERATION, session, max_retries=5) + session.post.side_effect = ConnectionResetError() + + with pytest.raises(ConnectionResetError): + await client.send_request(self._gen_request("completion", stream=False)) + + assert session.post.call_count == 1 diff --git a/tests/unittest/disaggregated/test_kv_transfer.py b/tests/unittest/disaggregated/test_kv_transfer.py index 03b706bae59b..6a27d35e5851 100644 --- a/tests/unittest/disaggregated/test_kv_transfer.py +++ b/tests/unittest/disaggregated/test_kv_transfer.py @@ -112,6 +112,9 @@ class KvCacheConfigV2: # tensorrt_llm/llmapi/llm_args.py; KVCacheManagerV2._build_base_config() # reads it unconditionally. kv_pool_rebalance_config: KvPoolRebalanceConfig = field(default_factory=KvPoolRebalanceConfig) + # Mirrors KvCacheConfig.enable_disagg_completion_reservation in + # tensorrt_llm/llmapi/llm_args.py; KVCacheManagerV2.__init__() reads it. + enable_disagg_completion_reservation: bool = False _THIS_DIR = Path(__file__).resolve().parent diff --git a/tests/unittest/disaggregated/test_openai_disagg_server.py b/tests/unittest/disaggregated/test_openai_disagg_server.py index dbc2575c3ae7..bfc8c83de11a 100644 --- a/tests/unittest/disaggregated/test_openai_disagg_server.py +++ b/tests/unittest/disaggregated/test_openai_disagg_server.py @@ -11,20 +11,26 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +import json from types import SimpleNamespace -from unittest.mock import AsyncMock, Mock, patch +from unittest.mock import AsyncMock, MagicMock, Mock, patch +import aiohttp import pytest -from fastapi import Request +from fastapi import HTTPException, Request +from fastapi.responses import StreamingResponse from starlette.datastructures import Headers from tensorrt_llm.llmapi.disagg_utils import ServerRole, extract_disagg_cfg from tensorrt_llm.serve import openai_disagg_server +from tensorrt_llm.serve.anthropic_protocol import AnthropicMessagesRequest from tensorrt_llm.serve.openai_disagg_server import OpenAIDisaggServer from tensorrt_llm.serve.openai_protocol import ( + ChatCompletionRequest, CompletionRequest, ConversationParams, DisaggregatedParams, + ResponsesRequest, ) pytestmark = pytest.mark.cpu_only @@ -322,3 +328,227 @@ def test_disagg_config_rejects_non_bool_request_chat_template_opt_in(value): generation_servers={"num_instances": 0}, allow_request_chat_template=value, ) + + +def _streaming_server(): + server = OpenAIDisaggServer.__new__(OpenAIDisaggServer) + server._perf_metrics_collector = MagicMock() + server._allow_request_chat_template = False + server._collect_perf_metrics = False + server._config = SimpleNamespace(conversation_affinity_header_for_subagents=None) + return server + + +def _streaming_raw_request(): + return SimpleNamespace( + headers=Headers(headers={}), state=SimpleNamespace(server_arrival_time=0.0) + ) + + +def _stream_request(kind): + if kind == "completion": + return CompletionRequest(model="test-model", prompt="hello", stream=True) + assert kind == "chat" + return ChatCompletionRequest( + model="test-model", messages=[{"role": "user", "content": "hello"}], stream=True + ) + + +async def _failing_generation_stream(): + raise aiohttp.ClientPayloadError("Response payload is not completed") + yield b"unreachable" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["completion", "chat"]) +async def test_stream_failing_before_first_chunk_returns_http_error(kind): + """A generation server that fails its KV transfer yields no chunk; answer 500, not a cut 200.""" + + async def entry_point(request, hooks): + return _failing_generation_stream() + + request = _stream_request(kind) + wrapper = _streaming_server()._wrap_entry_point(entry_point, type(request)) + + with pytest.raises(HTTPException) as exc_info: + await wrapper(request, _streaming_raw_request()) + + assert exc_info.value.status_code == 500 + assert "Response payload is not completed" in exc_info.value.detail + + +@pytest.mark.asyncio +@pytest.mark.parametrize("chunks", [[b"data: 1\n\n", b"data: 2\n\n", b"data: [DONE]\n\n"], []]) +async def test_stream_forwards_every_chunk(chunks): + closed = [] + + async def generation_stream(): + try: + for chunk in chunks: + yield chunk + finally: + closed.append(True) + + async def entry_point(request, hooks): + return generation_stream() + + wrapper = _streaming_server()._wrap_entry_point(entry_point, CompletionRequest) + request = CompletionRequest(model="test-model", prompt="hello", stream=True) + + response = await wrapper(request, _streaming_raw_request()) + + assert isinstance(response, StreamingResponse) + assert [chunk async for chunk in response.body_iterator] == chunks + assert closed == [True] + + +# Responses streams are left out: guard_responses_stream does not close its +# source when it is itself closed before the end. +@pytest.mark.asyncio +@pytest.mark.parametrize("kind", ["completion", "chat"]) +async def test_stream_closed_before_iteration_closes_upstream(kind): + """A client that leaves before the body is sent still releases the upstream stream.""" + closed = [] + + async def generation_stream(): + try: + yield b"data: 1\n\n" + yield b"data: 2\n\n" + finally: + closed.append(True) + + async def entry_point(request, hooks): + return generation_stream() + + request = _stream_request(kind) + wrapper = _streaming_server()._wrap_entry_point(entry_point, type(request)) + + response = await wrapper(request, _streaming_raw_request()) + assert isinstance(response, StreamingResponse) + await response.body_iterator.aclose() + + assert closed == [True] + + +@pytest.mark.asyncio +async def test_stream_error_after_first_chunk_propagates(): + closed = [] + + async def generation_stream(): + try: + yield b"data: 1\n\n" + raise RuntimeError("mid-stream") + finally: + closed.append(True) + + async def entry_point(request, hooks): + return generation_stream() + + wrapper = _streaming_server()._wrap_entry_point(entry_point, CompletionRequest) + + response = await wrapper(_stream_request("completion"), _streaming_raw_request()) + + assert isinstance(response, StreamingResponse) + received = [] + with pytest.raises(RuntimeError, match="mid-stream"): + async for chunk in response.body_iterator: + received.append(chunk) + assert received == [b"data: 1\n\n"] + assert closed == [True] + + +@pytest.mark.asyncio +async def test_stream_upstream_error_status_keeps_worker_status(): + async def rejected_stream(): + raise aiohttp.ClientResponseError( + MagicMock(), + (), + status=400, + message='Bad Request: {"error": {"message": "prompt too long"}}', + ) + yield b"unreachable" + + async def entry_point(request, hooks): + return rejected_stream() + + wrapper = _streaming_server()._wrap_entry_point(entry_point, CompletionRequest) + + with pytest.raises(HTTPException) as exc_info: + await wrapper(_stream_request("completion"), _streaming_raw_request()) + + assert exc_info.value.status_code == 400 + assert exc_info.value.detail == "prompt too long" + + +@pytest.mark.asyncio +async def test_responses_stream_failing_before_first_frame_ends_with_error_event(): + """A Responses relay keeps its 200 and ends with an `error` event.""" + + async def failing_relay(): + raise RuntimeError("no generation worker") + yield b"unreachable" + + async def entry_point(request, hooks): + return failing_relay() + + wrapper = _streaming_server()._wrap_entry_point(entry_point, ResponsesRequest) + request = ResponsesRequest.model_validate({"model": "m", "input": "hi", "stream": True}) + + response = await wrapper(request, _streaming_raw_request()) + + assert isinstance(response, StreamingResponse) + frames = [] + with pytest.raises(RuntimeError, match="no generation worker"): + async for frame in response.body_iterator: + frames.append(frame) + assert len(frames) == 1 + assert frames[0].startswith(b"event: error") + assert json.loads(frames[0].decode().split("data: ", 1)[1])["sequence_number"] == 0 + + +@pytest.mark.asyncio +async def test_failed_stream_is_not_counted_as_response(): + async def ok_stream(): + yield b"data: 1\n\n" + + streams = iter([_failing_generation_stream(), ok_stream()]) + + async def entry_point(request, hooks): + return next(streams) + + server = _streaming_server() + metrics = server._perf_metrics_collector + wrapper = server._wrap_entry_point(entry_point, CompletionRequest) + + with pytest.raises(HTTPException): + await wrapper(_stream_request("completion"), _streaming_raw_request()) + + metrics.total_responses.inc.assert_not_called() + metrics.internal_errors.inc.assert_called_once_with() + + response = await wrapper(_stream_request("completion"), _streaming_raw_request()) + + assert [chunk async for chunk in response.body_iterator] == [b"data: 1\n\n"] + metrics.total_responses.inc.assert_called_once_with() + + +@pytest.mark.asyncio +async def test_anthropic_stream_failing_before_first_chunk_is_api_error(): + """/v1/messages maps a chat stream that fails before its first chunk to an error envelope.""" + server = _streaming_server() + server._service = SimpleNamespace( + openai_chat_completion=AsyncMock(return_value=_failing_generation_stream()) + ) + request = AnthropicMessagesRequest.model_validate( + { + "model": "test-model", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hello"}], + "stream": True, + } + ) + + response = await server.anthropic_messages(request, _streaming_raw_request()) + + assert response.status_code == 500 + assert json.loads(response.body)["error"]["type"] == "api_error" diff --git a/tests/unittest/executor/test_forward_error_survival.py b/tests/unittest/executor/test_forward_error_survival.py index b37aa5edc235..d9fd6d9f6082 100644 --- a/tests/unittest/executor/test_forward_error_survival.py +++ b/tests/unittest/executor/test_forward_error_survival.py @@ -25,6 +25,8 @@ import pytest +from tensorrt_llm._torch.disaggregation.kv_cache_transceiver import GenTransferStatus +from tensorrt_llm._torch.disaggregation.orchestration import coordinator as coordinator_module from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState from tensorrt_llm._torch.pyexecutor.py_executor import BatchState, PyExecutor @@ -582,6 +584,54 @@ def test_fatal_error_flushes_parked_failures(): executor._terminate_request.assert_called_once_with(receiving) +def test_timed_out_receive_cancelled_at_loop_head_is_buffered_not_parked_under_adp(monkeypatch): + """A receive cancelled at the loop head is buffered, not parked, under ADP. + + The loop head cancelled it after its KV-transfer timeout, and multi-rank + ADP fails it at the synced drain. The cancel leaves no transfer behind, + so the real _handle_errors must buffer the error response and the + termination for the ADP flush instead of parking the failure: the + request has already left active_requests, so the held-failure sweep + would drop it without a response. + """ + monkeypatch.setattr(coordinator_module, "is_disagg_inflight_cancel_enabled", lambda: False) + monkeypatch.delenv("TRTLLM_DISAGG_BENCHMARK_GEN_ONLY", raising=False) + monkeypatch.delenv("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP", raising=False) + executor = _transfer_executor() + # Set before the first executor.disagg access, which builds the coordinator. + executor.enable_attention_dp = True + executor.dist = Mock(rank=0, world_size=2, tp_size=2, pp_size=1, cp_size=1) + executor.dist.tp_allgather_int64.return_value = Mock(any=lambda: True) + executor.canceled_req_ids = [] + transceiver = executor.kv_cache_transceiver + transceiver.kv_transfer_timeout_ms = 1000 + transceiver.supports_inflight_request_cancellation.return_value = False + transceiver.check_gen_transfer_status.return_value = GenTransferStatus([], [], []) + transceiver.cancel_request.return_value = True + receiving = _gen_receive_request(1) + receiving.is_child = False + receiving.parent_request_id = None + receiving.py_kv_transfer_timed_out = True + receiving.is_disagg_generation_transmission_in_progress = True + executor.active_requests = [receiving] + + executor.disagg.poll_gen_transfers() + + transceiver.cancel_request.assert_called_once_with(receiving) + assert executor.active_requests == [] + assert receiving.state == LlmRequestState.DISAGG_TRANS_ERROR + + executor.disagg.handle_timeouts_synced() + + executor.dist.tp_allgather_int64.assert_called_once_with([True]) + assert executor._transfer_held_failures == {} + assert executor._pending_response_terminations == [receiving] + assert [req_id for req_id, _ in executor._pending_transfer_responses] == [1] + assert receiving.state == LlmRequestState.GENERATION_COMPLETE + executor._terminate_request.assert_not_called() + executor._enqueue_responses.assert_not_called() + + def _terminating_executor(): """Executor shell exposing only what _terminate_request touches.""" executor = object.__new__(PyExecutor) diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py index 923f4767eda2..1cf4408a73f4 100755 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py @@ -1589,7 +1589,10 @@ def test_commit_min_snapshot_reuses_swa_post_commit_prefix(self) -> None: self.assertEqual(kv2.num_committed_tokens, len(prompt)) kv2.close() - def test_reuse_match_backoff_trims_the_tail_of_every_match(self) -> None: + @parameterized.expand([("tokens", 1, True), ("blocks", 8, False)]) + def test_reuse_match_backoff_trims_the_tail_of_every_match( + self, _case: str, alignment: int, partial_reuse: bool + ) -> None: """A pool holding state that reads D tokens ahead cannot use the last D. Set for a one-model speculative decoding draft pool, whose KV at position @@ -1597,6 +1600,9 @@ def test_reuse_match_backoff_trims_the_tail_of_every_match(self) -> None: the first m - D positions. Applied inside match(), so probe_reuse and the claim agree and neither needs a second tree walk. Commit is untouched -- the blocks are stored in full; only what a later match may use shrinks. + + A whole-block KV connector also needs the trimmed match aligned to a + block boundary, which reuse_match_alignment applies after the backoff. """ tokens_per_block = 8 backoff = 3 @@ -1615,6 +1621,8 @@ def test_reuse_match_backoff_trims_the_tail_of_every_match(self) -> None: ) ], reuse_match_backoff=backoff, + reuse_match_alignment=alignment, + enable_partial_reuse=partial_reuse, ) ) @@ -1630,17 +1638,67 @@ def test_reuse_match_backoff_trims_the_tail_of_every_match(self) -> None: # Every other test in this file runs with the default backoff of 0 and # matches the full prompt, so the delta here is the backoff alone. The - # trim crosses a block boundary rather than stopping at one. - self.assertEqual(self.manager.probe_reuse(None, prompt), len(prompt) - backoff) + # trim crosses a block boundary rather than stopping at one. Alignment + # is applied after the backoff, never in its place. + expected = (len(prompt) - backoff) // alignment * alignment + self.assertEqual(self.manager.probe_reuse(None, prompt), expected) # A match shorter than the backoff collapses to nothing, not negative. self.assertEqual(self.manager.probe_reuse(None, prompt[: backoff - 1]), 0) + # The first block a connector would be asked about is the one holding the + # first token past the trimmed and aligned match. + keys = list(sequence_to_blockchain_keys(tokens_per_block, ReuseScope(), prompt)) + self.assertEqual( + self.manager.probe_first_new_block_key(None, prompt), + keys[expected // tokens_per_block + 1][1], + ) + # The claim agrees with the probe: same trim, no second tree walk. kv2 = self.manager.create_kv_cache(input_tokens=prompt) - self.assertEqual(kv2.num_committed_tokens, len(prompt) - backoff) + self.assertEqual(kv2.num_committed_tokens, expected) kv2.close() + def test_reuse_match_alignment_must_divide_tokens_per_block(self) -> None: + """An alignment the block size is not a multiple of is refused at construction.""" + for alignment in (3, 0, -8): + with self.subTest(alignment=alignment), self.assertRaises(AssertionError): + KVCacheManagerConfig( + tokens_per_block=8, + cache_tiers=[GpuCacheTierConfig(quota=16 << 20)], + layers=[ + AttentionLayerConfig( + layer_id=LayerId(0), + buffers=[BufferConfig(role=Role.KEY, size=1024)], + ) + ], + reuse_match_alignment=alignment, + ) + + def test_replace_preserves_reuse_match_alignment(self) -> None: + """dataclasses.replace() on a config must carry reuse_match_alignment over. + + KVCacheManagerV2 rebuilds its config with replace(config, layers=...). The + C++ binding replaces the Python dataclass, so replace() is keyed on the + field spec in tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py: a field + missing from it is silently reset to its constructor default. + """ + layers = [ + AttentionLayerConfig( + layer_id=LayerId(0), + buffers=[BufferConfig(role=Role.KEY, size=1024)], + ) + ] + config = KVCacheManagerConfig( + tokens_per_block=8, + cache_tiers=[GpuCacheTierConfig(quota=16 << 20)], + layers=layers, + reuse_match_alignment=8, + ) + self.assertEqual(config.reuse_match_alignment, 8) + self.assertEqual(replace(config, layers=config.layers).reuse_match_alignment, 8) + self.assertEqual(replace(config).reuse_match_alignment, 8) + @parameterized.expand([("backoff_0", 0), ("backoff_1", 1)]) def test_planned_drop_handle(self, _case: str, reuse_match_backoff: int) -> None: window_size = 8 @@ -3451,6 +3509,51 @@ def test_first_new_block_probe_uses_snapshot_after_backoff(self) -> None: self.assertTrue(claimed.resume(stream)) claimed.close() + def test_reuse_match_alignment_applies_after_snapshot_pruning(self) -> None: + """Alignment trims the final, snapshot-pruned match and the attention-only diagnostic. + + The tree holds a full block and a partial one ending at token 40, with an SSM + snapshot at each end, so both the attention prefix and the recurrent state + reach 40. Aligning to the block size cuts the match back to the snapshot at 32. + KvCacheManager copies its config, so every alignment needs its own manager. + The retry after an SSM snapshot truncation lands off the alignment is not + reached here. + """ + tokens_per_block = 32 + prompt = [self.next_token() for _ in range(64)] + for alignment, expected in ((1, 40), (tokens_per_block, 32)): + cfg = self._make_ssm_config( + tokens_per_block=tokens_per_block, enable_partial_reuse=True + ) + cfg.reuse_match_alignment = alignment + self.manager = KVCacheManager(cfg) + stream_holder = CachedCudaStream() + stream = cast(CudaStream, stream_holder.handle) + + seed = self.manager.create_kv_cache() + self.assertTrue(seed.resume(stream)) + seed.capacity = 32 + seed.commit(prompt[:32]) + seed.capacity = 40 + seed.commit(prompt[32:40]) + seed.close() + + kv = self.manager.create_kv_cache(input_tokens=prompt[:48]) + self.assertEqual(kv.num_committed_tokens, expected, f"alignment {alignment}") + self.assertEqual( + kv._get_num_reusable_tokens_before_hybrid_pruning(), + expected, + f"alignment {alignment}", + ) + self.assertTrue(kv.resume(stream)) + kv.close() + + del seed, kv + gc.collect() + stream_holder.synchronize() + self.manager.shutdown() + del self.manager + def test_reuse_match_reports_content_divergence_depth(self) -> None: """The raw walk depth locates the fork, even where pruning shortens reuse.""" cfg = self._make_ssm_config(tokens_per_block=32, enable_partial_reuse=True) diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index ede5add24e59..11025e6c0e41 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -1308,6 +1308,37 @@ def test_KvCacheConfig_declaration(offload_field): BlockReuseConfig(max_num_turns=0) +@pytest.mark.cpu_only +def test_KvCacheConfig_disagg_completion_reservation(): + assert KvCacheConfig().enable_disagg_completion_reservation is False + field = KvCacheConfig.model_fields["enable_disagg_completion_reservation"] + assert field.json_schema_extra["status"] == "prototype" + + config = KvCacheConfig(host_cache_size=0, + enable_disagg_completion_reservation=True) + assert config.enable_disagg_completion_reservation is True + assert KvCacheConfig.model_validate( + config.model_dump()).enable_disagg_completion_reservation is True + assert KvCacheConfig.model_validate_json( + config.model_dump_json()).enable_disagg_completion_reservation is True + # The C++ config has no such field, and the conversion leaves it out. + pybind_config = config._to_pybind() + assert pybind_config.host_cache_size == 0 + assert not hasattr(pybind_config, "enable_disagg_completion_reservation") + + # The deployment path: a YAML extra-options file merged into the LLM args. + merged = update_llm_args_with_extra_dict( + {"model": "dummy"}, + yaml.safe_load("kv_cache_config:\n" + " host_cache_size: 0\n" + " enable_disagg_completion_reservation: true\n")) + assert merged[ + "kv_cache_config"].enable_disagg_completion_reservation is True + with pytest.raises(ValidationError, + match="enable_disagg_completion_reservation"): + KvCacheConfig(enable_disagg_completion_reservation="sometimes") + + @pytest.mark.cpu_only def test_KvCacheConfig_offload_priority_compatibility(): assert KvCacheConfig().cold_tier_offload_min_priority is None @@ -5469,3 +5500,46 @@ def test_deepseek_resolves_auto_to_python_on_nixl(self) -> None: cfg = self._pretrained_config(["DeepseekV3ForCausalLM"], "deepseek_v3") _resolve_transceiver_runtime_auto(args, DeepseekV3ForCausalLM, cfg) assert args.cache_transceiver_config.transceiver_runtime == "PYTHON" + + +@pytest.mark.cpu_only +class TestKvCacheConnectorAdmissionQuery: + """kv_connector_config.admission_query: admission-time connector queries.""" + + @staticmethod + def _config(**kwargs): + from tensorrt_llm.llmapi.llm_args import KvCacheConnectorConfig + return KvCacheConnectorConfig(connector="mooncake-store", **kwargs) + + def test_off_by_default(self) -> None: + assert self._config().admission_query is None + + def test_empty_mapping_enables_page_booking_only(self) -> None: + admission = self._config(admission_query={}).admission_query + assert admission.max_uncharged_parks == 0 + assert admission.max_offer_deferrals == 16 + + def test_round_trip(self) -> None: + from tensorrt_llm.llmapi.llm_args import KvCacheConnectorConfig + config = self._config(admission_query={ + "max_uncharged_parks": 4, + "max_offer_deferrals": 0 + }) + dumped = config.model_dump() + assert dumped["admission_query"] == { + "max_uncharged_parks": 4, + "max_offer_deferrals": 0 + } + restored = KvCacheConnectorConfig.model_validate_json( + config.model_dump_json()) + assert restored.admission_query == config.admission_query + + @pytest.mark.parametrize("field", + ["max_uncharged_parks", "max_offer_deferrals"]) + def test_negative_values_are_rejected(self, field: str) -> None: + with pytest.raises(ValidationError, match=field): + self._config(admission_query={field: -1}) + + def test_unknown_key_is_rejected(self) -> None: + with pytest.raises(ValidationError, match="max_parks"): + self._config(admission_query={"max_parks": 1})