From e99294139a9bc66d66cded600fb2b6504ffb74c0 Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Tue, 27 Oct 2026 15:33:07 -0700 Subject: [PATCH 01/20] [None][feat] KVCacheManagerV2: align radix-tree reuse matches to a block multiple (reuse_match_alignment) A one-model EAGLE3 or MTP draft sets reuse_match_backoff, so the radix tree trims every local match by the draft lookahead and a match of K whole blocks ends a few tokens short of a block boundary. A KV connector that stores and offers whole blocks is queried from that endpoint and has nothing to offer for any request with a local hit. Add KVCacheManagerConfig::reuseMatchAlignment. After the backoff, BlockRadixTree::match trims the final match down to a multiple of the alignment, and pruneMatch re-aligns before validating pages and again whenever SSM snapshot truncation moves the endpoint off the alignment, so the probe, the claim and the first-new-block probe all trim once in the same tree walk. The attention-only diagnostic pass is aligned too; the content-divergence depth stays raw. The default of 1 keeps token-granular reuse, and validate() refuses an alignment that is not a positive divisor of tokens_per_block. The nanobind constructor gains a trailing reuse_match_alignment=1 kwarg and a read-write field, and the runtime field spec lists it so that dataclasses.replace() rebuilds of the config keep it. The struct layout changes, so the bindings and the kvCacheManagerV2 gtests must be rebuilt together with the library. Signed-off-by: Simeng Liu --- .../kv_cache_manager_v2/blockRadixTree.cpp | 69 +++++++---- .../kv_cache_manager_v2/blockRadixTree.h | 7 +- .../kv_cache_manager_v2/config.cpp | 5 + .../kv_cache_manager_v2/config.h | 4 + .../kv_cache_manager_v2/kvCacheManager.cpp | 3 +- .../batch_manager/kvCacheManagerV2.cpp | 7 +- .../kvCacheManagerV2ConfigTest.cpp | 18 +++ .../runtime/kv_cache_manager_v2/__init__.py | 1 + .../runtime/kv_cache_manager_v2/__init__.pyi | 3 + .../test_kv_cache_manager_v2.py | 111 +++++++++++++++++- 10 files changed, 193 insertions(+), 35 deletions(-) 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/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/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) From d6b6ae7eaf808249f42dff0c2bc8ab2ff23edb45 Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Tue, 27 Oct 2026 15:38:58 -0700 Subject: [PATCH 02/20] [None][feat] KV connector: whole-block local reuse declarations read by KVCacheManagerV2 A connector that stores and offers KV by whole blocks is queried from the local prefix match. Under a one-model draft that match is backed off by the draft's prompt lookahead and ends a few tokens short of a block boundary, so such a connector has nothing to offer for any request with a local hit. KvCacheConnectorWorker gains two opt-in declarations next to supports_attention_dp, both False by default and honoured only when exactly True: - requires_block_aligned_reuse: when a backoff is in effect and partial reuse is off, KVCacheManagerV2 passes its ledger tokens_per_block as the core reuse_match_alignment, so every local match is aligned down to a block boundary and the connector serves the trimmed block with the rest of the prefix. A paired draft pool aligns too, keeping the joint claim on a boundary. - drops_reuse_match_backoff: the target manager drops the one-model draft backoff instead and uses whole local matches. Draft KV at the tail of a reused block may then come from another continuation, which affects draft acceptance only. A draft manager never drops its own backoff, and the raw draft_prompt_lookahead stays for the sharing lender. The manager reads both through one helper that ignores a capacity-only connector and tolerates partially built managers and connector doubles without a worker. Nothing declares them yet, and speculative decoding is still refused with a connector, so production behaviour is unchanged. The connector guide documents the declarations in the worker interface, and the shutdown and capacity_only worker bullets that had drifted under the attention data parallelism section move back there. Signed-off-by: Simeng Liu --- docs/source/features/kv-cache-connector.md | 21 ++- .../connectors/kv_cache_connector.py | 15 ++ .../kv_cache/kv_cache_manager_v2.py | 69 +++++++- ...kv_cache_manager_v2_connector_alignment.py | 148 ++++++++++++++++++ .../kv_cache/test_kvcm2_integration.py | 94 +++++++++++ ...est_kv_connector_v2_prefix_real_manager.py | 68 +++++++- 6 files changed, 407 insertions(+), 8 deletions(-) create mode 100644 tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_connector_alignment.py diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index 53918b857687..5fc9d74f76f9 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -345,6 +345,21 @@ 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`. + +* **`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. + +* **`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. + +Both declarations 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 @@ -470,12 +485,6 @@ 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. - -* **`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. - ## Example Implementation The file `examples/llm-api/llm_kv_cache_connector.py` provides a reference implementation of a **Persistent KV Cache**. diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py index 8c4c3815493a..6cc92171a0bc 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py @@ -180,6 +180,21 @@ 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 + #: 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 def __init__(self, llm_args: TorchLlmArgs): self._llm_args = llm_args 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..5237f9fed138 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 @@ -2015,10 +2015,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, @@ -3796,7 +3805,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 +3883,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 leaves reuse unchanged. A capacity-only connector serves nothing, so it + never reshapes local reuse. + + 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( 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_kvcm2_integration.py b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py index addb326f0a91..9089091fc9d1 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)], 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..d0818d26b078 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 @@ -33,7 +33,7 @@ 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 @@ -378,6 +378,72 @@ def test_freeing_the_allocation_makes_the_request_askable_again(manager, connect assert len(connector.queries) == 2 +@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. # From 35330dac81a26faea099dc1580ccf64035d1931e Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Tue, 27 Oct 2026 15:55:37 -0700 Subject: [PATCH 03/20] [None][fix] KV connector: skip dummy requests without attention DP KvCacheConnectorManager and its scheduler-output manager skipped placeholder requests only when attention DP was on, by testing `enable_attention_dp and req.is_dummy_request`. A dummy request carries no prompt a connector could load or save, whatever the parallel mode, so the check now goes through one module helper, `_is_dummy_request`, used by the scheduler output, prefix reservation, query, commit, request_finished and update_state_after_alloc paths. The helper reads both `is_dummy_request` (bound by the C++ LlmRequest, so readable on the base class C++ callers hand over) and the Python-only `is_attention_dp_dummy`, and compares each against True. The explicit comparison keeps a MagicMock request, which answers every attribute truthily, from being mistaken for a dummy and silently skipped. Under tensor parallelism this is defensive rather than a behaviour change: C++ addSequence already skips dummies, KVCacheManagerV2 does not ask the connector for them, and padding dummies exist only under attention DP. Every TP rank sees the same flags, so the early returns before the leader broadcast stay collective-consistent. The output manager's `enable_attention_dp` keyword is now dead and is removed together with the one call site that passed it. Signed-off-by: Simeng Liu --- .../connectors/kv_cache_connector.py | 48 +++++++----- tests/unittest/_torch/test_connector_adp.py | 75 ++++++++++++++++++- 2 files changed, 104 insertions(+), 19 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py index 6cc92171a0bc..1b5f81e8d1ae 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: @@ -698,8 +719,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() @@ -712,7 +732,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 @@ -733,7 +753,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 @@ -813,9 +833,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 @@ -920,9 +938,7 @@ 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 if local_end < 0: raise ValueError("A prefix reservation cannot start before the prompt") @@ -1123,7 +1139,7 @@ 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 @@ -1165,7 +1181,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 @@ -1440,7 +1456,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: @@ -1616,11 +1632,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/tests/unittest/_torch/test_connector_adp.py b/tests/unittest/_torch/test_connector_adp.py index ad1e6829630e..3cebbec1e5a0 100644 --- a/tests/unittest/_torch/test_connector_adp.py +++ b/tests/unittest/_torch/test_connector_adp.py @@ -228,7 +228,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 From 06b56d5ede499c75ac8434776005bd50f4e7326e Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Tue, 27 Oct 2026 16:05:45 -0700 Subject: [PATCH 04/20] [None][feat] KV connector: restart only the requests whose prefix load failed (reports_failed_loads) A page a connector offered at lookup can be gone by the time it is loaded: evicted from a full pool, or left behind by a put that never completed. The runtime has already counted those positions as computed, so today the only safe answer is for the connector to raise, which stops the executor for one request's cache miss. This adds an opt-in protocol that hands such a request back for a restart instead. A worker declares `reports_failed_loads = True` and implements `take_failed_load_requests()` (polled after every `start_load_kv`), optionally `take_failed_async_load_requests()` (polled after every `get_finished`), and `drop_bound_saves(request_ids)`. The manager masks the declaration for a capacity-only role, refuses a declaring worker that leaves its hooks at the base defaults, and refuses the combination with prefix reservations, whose completion protocol has no failure outcome. The executor refuses it on KVCacheManagerV1, which cannot restart a request. For each reported request the executor drops its bound saves, then calls the new `KVCacheManagerV2.drop_context_allocation`, which is `free_resources` followed by a context rewind, and releases any speculative-decoding slot. The cache is closed without committing, so the never-filled range does not reach the radix tree, and the scheduler readmits the request and asks the connector afresh. `_kv_connector_prepare_batch` now parks, pads an attention-DP batch the connector emptied (`_pad_attention_dp_batch_emptied_by_connector`), then binds the plan and starts loads (`_kv_connector_start_loads`), which is where synchronous failures are taken. Asynchronous failures are taken right after the terminations in `_kv_connector_terminate_requests`, so a parked load whose cancellation was deferred is not restarted. The attention-DP recovery loop drops restarted parked requests from the requests it waits for, so a restarted request (CONTEXT_INIT with no cache) is never put back into the batch, and a loop whose only parked request restarted stops waiting and retries padding. A rebuilt plan goes through `_kv_connector_start_loads` too. A restart is owner-local, so it applies on a single rank or an attention-DP owner; a tensor-parallel report is fatal. Connector exceptions, deadline misses and misreported IDs remain fail-stop through `_kv_connector_fail`, with no new collective or vote field. A misreported ID is a synchronous ID outside the batch, an asynchronous ID whose load has not been reported finished, or one an earlier `get_finished` already returned as finished: a finished load stays in `finished_async_loading_requests` after the scheduler admits the request, so only the `get_finished` call the poll follows can still name a request that is parked. Signed-off-by: Simeng Liu --- docs/source/features/kv-cache-connector.md | 25 +- .../connectors/kv_cache_connector.py | 248 ++++++++- .../kv_cache/kv_cache_manager_v2.py | 24 + tensorrt_llm/_torch/pyexecutor/py_executor.py | 215 +++++++- .../kv_cache/test_kvcm2_integration.py | 18 + ...est_kv_connector_v2_prefix_real_manager.py | 41 ++ tests/unittest/_torch/test_connector_adp.py | 500 +++++++++++++++++- 7 files changed, 1040 insertions(+), 31 deletions(-) diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index 5fc9d74f76f9..7b1c08fa27a3 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`** @@ -310,6 +310,8 @@ request in the same iteration. With chunked prefill, allocation growth can limit 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. @@ -346,6 +348,20 @@ These methods run on all workers (GPU processes) and interact with the actual GP * **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. @@ -452,6 +468,13 @@ 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. diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py index 1b5f81e8d1ae..a03706c90a0a 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py @@ -216,6 +216,16 @@ class KvCacheConnectorWorker(ABC): #: 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 def __init__(self, llm_args: TorchLlmArgs): self._llm_args = llm_args @@ -396,6 +406,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. @@ -525,10 +594,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. @@ -818,6 +889,15 @@ 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 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)) # Requests that haven't yet been passed into get_finished. self.new_async_requests = AsyncRequests(dict(), dict()) @@ -848,6 +928,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( @@ -866,6 +949,32 @@ def validate_attention_dp( "and implement owner-local scheduling/transfers against a shared pool." ) + @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. @@ -930,6 +1039,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 @@ -1322,6 +1443,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( @@ -1569,6 +1808,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 = ( 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 5237f9fed138..86c33602c3dc 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 @@ -4859,6 +4859,30 @@ 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. diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index e0502b4cd181..b5faf5676d6b 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -1239,6 +1239,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 @@ -1279,6 +1291,13 @@ def _maybe_init_kv_connector_manager(self) -> None: "KV connector attention data parallelism requires KVCacheManagerV2. " "Set kv_cache_config.use_kv_cache_manager_v2=True.") + 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: raise NotImplementedError( "KV Cache Connector is not supported with speculative " @@ -4611,6 +4630,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 +4681,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 +4786,158 @@ 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 []) + self.kv_connector_manager.take_scheduled_requests_pending_load( + 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 +4957,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 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 9089091fc9d1..5fba4654f9b1 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kvcm2_integration.py @@ -1863,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_kv_connector_v2_prefix_real_manager.py b/tests/unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py index d0818d26b078..54bf447bc365 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 @@ -378,6 +378,47 @@ 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"), [ diff --git a/tests/unittest/_torch/test_connector_adp.py b/tests/unittest/_torch/test_connector_adp.py index 3cebbec1e5a0..fb98484530e6 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 @@ -1702,3 +1742,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 From 46e447ae00c1991b915bce718d0277c2950e95f2 Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Tue, 27 Oct 2026 16:31:53 -0700 Subject: [PATCH 05/20] [None][feat] KV connector: speculative-prefill-only mode for one-model MTP/EAGLE3 on disaggregated context servers KVCacheManagerV2 refuses speculative decoding next to a KV connector because a rejected draft token shrinks a request's page list underneath the connector, which only hears about appended pages. A context-only request on a disaggregated context server never rejects a draft token, so its page list only grows, and one-model MTP and EAGLE3 keep their draft KV in the target pool the connector has registered. This lets such a server run speculative decoding next to a connector that declares it can. Both connector halves gain a `supports_speculative_prefill` class attribute, modelled on `supports_attention_dp`. The scheduler's declaration promises that its keys cover the draft's prompt lookahead past each block and that it saves only computed blocks; the worker's promises whole-page transfers of every layer group, draft layers included. `KvCacheConnectorManager` takes a keyword-only `speculative_prefill_only` and refuses it, after the attention-DP and failed-load checks, unless both present halves declare support (`declares_speculative_prefill`; a TP follower has no scheduler). The executor creator computes the mode once the connector classes are resolved: one-model MTP (MTP-Eagle or vanilla) or EAGLE3 with a positive draft length, a cache transceiver backend, and both declarations. It is never set for an explicit `use_kv_cache_manager_v2=False`, so V1 keeps constructing exactly as before. The speculative-decoding refusal moves after class resolution, keeps its wording, and is lifted only in this mode. The executor accepts the mode on KVCacheManagerV2 only (an "auto" that resolves to V1 is refused), lifts its own speculative refusal for it, and refuses a separate draft KV cache, which the connector never sees. A request that is not context-only is refused at admission by the module-level `_validate_connector_admission`, the first check in `_validate_request`, so only that request fails, for example a /health_generate probe sent straight to the context server. Being a module-level helper read through getattr, it needs no stub in the minimal executors unit tests build. As a backstop, the manager refuses such a request again in `query_num_new_matched_tokens` and `reserve_prefix` before the leader broadcast, so every TP rank raises together, and in `build_scheduler_output` before any output state is recorded. Dummy requests and a capacity-only connector are never refused, which also lets a capacity-role generation server whose classes declare the mode run one-model speculation. No connector declares the mode yet, so behaviour is unchanged until one opts in. Signed-off-by: Simeng Liu --- docs/source/features/kv-cache-connector.md | 70 +++- .../connectors/kv_cache_connector.py | 123 +++++++ tensorrt_llm/_torch/pyexecutor/py_executor.py | 69 +++- .../_torch/pyexecutor/py_executor_creator.py | 87 ++++- .../test_kv_connector_initialization.py | 96 ++++++ .../test_kv_connector_speculative_prefill.py | 311 ++++++++++++++++++ ...executor_creator_connector_spec_prefill.py | 150 +++++++++ 7 files changed, 887 insertions(+), 19 deletions(-) create mode 100644 tests/unittest/_torch/executor/test_kv_connector_speculative_prefill.py create mode 100644 tests/unittest/_torch/executor/test_py_executor_creator_connector_spec_prefill.py diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index 7b1c08fa27a3..0a5c0f24ebdb 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -82,6 +82,9 @@ 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). + ##### 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 +201,64 @@ 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`. ##### `RequestData` fields that may not be populated @@ -368,13 +422,16 @@ These methods run on all workers (GPU processes) and interact with the actual GP * **`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. -Both declarations 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. +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 @@ -481,7 +538,10 @@ 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. +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. This change adds no Mooncake backend, deployment configuration or capacity tuning. Select a connector through `connector_module`, `connector_scheduler_class`, diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py index a03706c90a0a..de21fc1aaa77 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py @@ -201,6 +201,18 @@ 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 @@ -477,6 +489,21 @@ 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 def __init__(self, llm_args: TorchLlmArgs): self._llm_args = llm_args @@ -861,6 +888,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. @@ -872,7 +902,26 @@ def __init__( scheduler: Optional[KvCacheConnectorScheduler], *, enable_attention_dp: bool = False, + speculative_prefill_only: 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`. + + 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." ) @@ -898,6 +947,18 @@ def __init__( ) 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()) @@ -949,6 +1010,28 @@ 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. @@ -1061,6 +1144,9 @@ def reserve_prefix(self, request: LlmRequest, local_end: int) -> Optional[Prefix """Query once per pending attempt and protect the promised source range.""" 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) @@ -1267,6 +1353,9 @@ def query_num_new_matched_tokens( # 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) @@ -1352,6 +1441,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", @@ -1593,6 +1709,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( {}, { diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index b5faf5676d6b..e247c93a389d 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -162,6 +162,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 @@ -1285,6 +1320,19 @@ 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( @@ -1298,7 +1346,8 @@ def _maybe_init_kv_connector_manager(self) -> None: "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: + 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 " @@ -1306,7 +1355,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, @@ -1358,6 +1410,17 @@ 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 @@ -6902,6 +6965,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..c2f374b8e4e1 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,18 @@ 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) + # 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/tests/unittest/_torch/executor/test_kv_connector_initialization.py b/tests/unittest/_torch/executor/test_kv_connector_initialization.py index c8c6d0db713a..6a88589c9a84 100644 --- a/tests/unittest/_torch/executor/test_kv_connector_initialization.py +++ b/tests/unittest/_torch/executor/test_kv_connector_initialization.py @@ -65,6 +65,102 @@ def test_v2_connector_registers_layer_groups( 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( "use_v2,connector_mode,num_groups,with_retention,rejected", [ 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_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..9f49fd6bdbbd --- /dev/null +++ b/tests/unittest/_torch/executor/test_py_executor_creator_connector_spec_prefill.py @@ -0,0 +1,150 @@ +# 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. +""" + +import pytest + +from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import ( + KvCacheConnectorManager, + KvCacheConnectorScheduler, + KvCacheConnectorWorker, +) +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 + + +@pytest.mark.parametrize("use_v2", [True, "auto"], ids=["v2", "auto"]) +@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, use_v2 +) -> None: + assert _connector_runs_speculative_prefill_only( + make_spec_config(), _Declaring, _Declaring, DISAGG, use_v2 + ) + + +@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) From e3ad6bab7cf29ac151e7f3b6e8521d96c3dd8917 Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Tue, 27 Oct 2026 16:57:56 -0700 Subject: [PATCH 06/20] [None][feat] KVCacheManagerV2: reserve completion headroom on disaggregated servers without an eviction tier With no eviction tier, a full-attention disaggregated server admitted requests against their prompt-sized allocations only. Incoming transfers and new prompts could then take pages that admitted requests still needed: a decode request ran out of room for its output, and a connector prefill, whose prefix is restored only after the final batch is chosen, found the pages for that restore taken by other requests' chunks. KVCacheManagerV2 can now keep a logical completion bound, in ledger blocks, per admitted request and admit a new one only while every bound still fits each GPU pool's free and evictable slots. No page is allocated early. Disaggregated generation init reserves through the output limit plus the maximum generation lead, and frees a refused request's cache instead of suspending it. A connector first chunk reserves through the end of its prompt when the connector has no source reservations; the V2 scheduler consults the reservation after every budget, chunking and multimodal-alignment refusal and defers the request instead of preempting an admitted one. The reservation is opt-in. A connector asks for it with a new KvCacheConnectorWorker declaration, requires_completion_reservation, which the manager reads from the attached connector's worker like requires_block_aligned_reuse, so every rank sees it. Only an explicit True counts, and a capacity-only connector's is ignored. The mooncake_store worker declares it: registering its pages drops the automatic host tier, and its prefix restores are chosen only after the final batch. A server opts in without such a connector through a new prototype KvCacheConfig field, enable_disagg_completion_reservation, False by default. It lets a disaggregated generation server with no connector and host_cache_size: 0 reserve its decode headroom; the connector prefill reservation still needs a connector that may serve the prompt. A target manager that cannot apply the field logs a warning rather than ignore it silently: without a connector the automatic host tier stays unless host_cache_size is 0. The field is Python-only, so the C++ KvCacheConfig conversion leaves it out, and usage telemetry captures it like the other KvCacheConfig booleans. A server with neither opt-in keeps its existing admission; that includes host_cache_size: 0 and page-registering connectors such as LMCache and KVBM, whose automatic host tier the manager drops. Beyond either opt-in, the reservation applies only to a disaggregated target manager with no tier below GPU, beam width 1, a single pipeline stage and full attention on every layer. Pipeline parallelism is excluded because each rank decides alone and a follower refusing a generation-init request that rank 0 admitted would not be reconciled. A bound is released by free_resources, by revert_allocate_context and a failed first-chunk resize while the chunk is unconsumed, and on shutdown. The receive-ordering tests bind the real prepare_disagg_gen_init onto a spec'd Mock, which would otherwise answer the new class-level state with truthy child mocks, so their fixture sets both reservation maps explicitly. Signed-off-by: Simeng Liu --- docs/source/features/kv-cache-connector.md | 45 + .../run_cache_transceiver_test.py | 3 + .../connectors/kv_cache_connector.py | 11 + .../connectors/mooncake_store/worker.py | 5 + .../kv_cache/kv_cache_manager_v2.py | 186 ++++- .../pyexecutor/scheduler/scheduler_v2.py | 21 + tensorrt_llm/llmapi/llm_args.py | 18 + .../schemas/trtllm-serve-config.schema.json | 7 + .../trtllm-serve-disagg-config.schema.json | 7 + .../usage/llm_args_golden_manifest.json | 5 + ..._cache_v2_disagg_completion_reservation.py | 785 ++++++++++++++++++ .../kv_cache/test_kv_cache_v2_scheduler.py | 71 ++ .../executor/test_disagg_receive_ordering.py | 3 + ...est_kv_connector_v2_prefix_real_manager.py | 121 ++- .../executor/test_mooncake_store_connector.py | 8 + .../test_cache_transceiver_single_process.py | 3 + .../disaggregated/test_kv_transfer.py | 3 + tests/unittest/llmapi/test_llm_args.py | 31 + 18 files changed, 1327 insertions(+), 6 deletions(-) create mode 100644 tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_disagg_completion_reservation.py diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index 0a5c0f24ebdb..2b30f88a3e22 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -369,6 +369,48 @@ 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. +##### 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. @@ -431,6 +473,9 @@ These methods run on all workers (GPU processes) and interact with the actual GP * **`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 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/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py index de21fc1aaa77..6b7265a6fa62 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py @@ -238,6 +238,17 @@ class KvCacheConnectorWorker(ABC): #: 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 diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py index 312dae49c79d..e411ce22419b 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py @@ -177,6 +177,11 @@ def _stream_handle(stream: Optional[torch.cuda.Stream]) -> int: class MooncakeStoreConnectorWorker(KvCacheConnectorWorker): """Moves KV pages between this rank's GPU cache and the Mooncake pool.""" + #: 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) 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 86c33602c3dc..0f919bdb2099 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 @@ -1772,6 +1772,13 @@ 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] = {} # Read by KvCacheCreator: a subclass that overrides the context # commit/history protocol opts out of generic reuse-match backoff. _supports_reuse_match_backoff = True @@ -2488,6 +2495,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 @@ -3888,8 +3927,8 @@ def _kv_connector_declares(self, capability: str) -> bool: 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 leaves reuse unchanged. A capacity-only connector serves nothing, so it - never reshapes local reuse. + 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 @@ -4886,7 +4925,13 @@ def drop_context_allocation(self, req: LlmRequest) -> None: 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) @@ -5264,6 +5309,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") @@ -5282,13 +5328,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: @@ -5316,6 +5484,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 @@ -5328,6 +5504,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. @@ -7121,6 +7299,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 @@ -7453,6 +7632,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/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 9c17e72d3fb3..ce9f5cef0a2b 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -997,6 +997,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 +1077,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 +1252,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): diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index e02c42ae24f5..b6e8bfe453ef 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -5177,6 +5177,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/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index de355f70ce33..80f1c313b130 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -2868,6 +2868,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'.", diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index 2bb5e4c910b0..b1bde7f4d66a 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -4432,6 +4432,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'.", diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index f940f2a80937..6bacef4ce177 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -673,6 +673,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", 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..1e9b3d1c82cf 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 @@ -682,6 +682,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 @@ -2333,6 +2335,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. 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_v2_prefix_real_manager.py b/tests/unittest/_torch/executor/test_kv_connector_v2_prefix_real_manager.py index 54bf447bc365..4aee698eaf12 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,7 +29,11 @@ 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 @@ -191,6 +195,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() @@ -890,8 +909,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() @@ -997,3 +1014,101 @@ 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() diff --git a/tests/unittest/_torch/executor/test_mooncake_store_connector.py b/tests/unittest/_torch/executor/test_mooncake_store_connector.py index 95cdf73dee6a..143e0d2efbb9 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_connector.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_connector.py @@ -789,6 +789,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 ---- 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_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/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index ede5add24e59..f0d9c1e90e0f 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 From 41130b6087c17bb91eff7c63c0dfcc08e1b8549f Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Tue, 27 Oct 2026 17:51:45 -0700 Subject: [PATCH 07/20] [None][feat] KV connector: admission-time prefix queries, offer page booking and uncharged async parks The KVCacheManagerV2 scheduler sizes and charges a first context chunk before the connector is asked for its prefix, which only happens in prepare_resources once the batch is final. A chunk the connector then loads asynchronously is parked before the forward pass, so its charge is lost budget, and an offer that moves the chunk window has to grow its pages after every later admission of the pass took theirs; when that grow fails, only the whole blocks the chunk's pages hold are loaded and the rest of the offer is recomputed. kv_connector_config.admission_query (KvCacheConnectorAdmissionConfig, prototype) lets the scheduler ask a connector that declares KvCacheConnectorScheduler.supports_admission_query once a first chunk is sized and allocated. The answer is stored on the request as an epoch-checked ConnectorAdmissionOffer and used by prepare_resources, so each allocation is still queried and reported once. KVCacheManagerV2.cover_admission_offer grows the allocation to the shifted window right away. When the pool cannot cover it, the chunk is skipped and the existing first-chunk release returns its pages, up to max_offer_deferrals passes or until the request has just preempted another one for pages; deferrals count as scheduler progress. With max_uncharged_parks > 0, up to that many chunks per pass that will park are admitted with a 0-token charge; they still take a request slot. The manager reports the capability from the scheduler class and turns it off under attention DP. The executor ignores the config with a warning for a capacity-only role, prefix-reservation connectors and connectors without the capability, and refuses it with the overlap scheduler or without the KVCacheV2Scheduler. A request admitted without a charge that does not park fails the executor through _kv_connector_fail, because other parked loads may already own destination pages. The cover is a no-op under SWA scratch reuse, whose pre-grown replay would otherwise read the grow as a skipped forward. Signed-off-by: Simeng Liu --- docs/source/features/kv-cache-connector.md | 69 +++- tensorrt_llm/_torch/pyexecutor/_util.py | 11 + .../connectors/kv_cache_connector.py | 34 +- .../kv_cache/kv_cache_manager_v2.py | 204 ++++++++++- tensorrt_llm/_torch/pyexecutor/llm_request.py | 13 + tensorrt_llm/_torch/pyexecutor/py_executor.py | 93 +++++ .../_torch/pyexecutor/py_executor_creator.py | 7 +- .../pyexecutor/scheduler/scheduler_v2.py | 197 ++++++++++- tensorrt_llm/llmapi/llm_args.py | 39 +++ .../schemas/trtllm-serve-config.schema.json | 37 ++ .../trtllm-serve-disagg-config.schema.json | 37 ++ .../usage/llm_args_golden_manifest.json | 10 + .../kv_cache/test_kv_cache_v2_scheduler.py | 326 ++++++++++++++++++ .../test_kv_connector_executor_lifetime.py | 123 +++++++ .../test_kv_connector_initialization.py | 44 +++ .../executor/test_kv_connector_v2_prefix.py | 220 ++++++++++++ ...est_kv_connector_v2_prefix_real_manager.py | 133 +++++++ tests/unittest/_torch/test_connector_adp.py | 17 + tests/unittest/llmapi/test_llm_args.py | 43 +++ 19 files changed, 1634 insertions(+), 23 deletions(-) diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index 2b30f88a3e22..853b2a91c817 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -85,6 +85,9 @@ on the leader process; with attention DP, they run on each owning rank. * **`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. @@ -356,10 +359,12 @@ 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 @@ -369,6 +374,62 @@ 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. + +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. + ##### Completion headroom on disaggregated servers A disaggregated server with KV cache manager V2 can reserve logical GPU-page headroom for each 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 6b7265a6fa62..e5c62345bfce 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/kv_cache_connector.py @@ -515,6 +515,22 @@ class KvCacheConnectorScheduler(ABC): #: 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 @@ -914,6 +930,7 @@ def __init__( *, 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. @@ -928,6 +945,10 @@ def __init__( 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 @@ -949,6 +970,13 @@ 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 @@ -1346,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``. 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 0f919bdb2099..902273d771cd 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 @@ -934,6 +934,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. @@ -1779,6 +1802,10 @@ class KVCacheManagerV2(BaseResourceManager): # 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 # Read by KvCacheCreator: a subclass that overrides the context # commit/history protocol opts out of generic reuse-match backoff. _supports_reuse_match_backoff = True @@ -1863,6 +1890,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" @@ -5951,6 +5981,123 @@ 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 + num_tokens, load_async = self.kv_connector_manager.query_num_new_matched_tokens( + req, local_end + ) + 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 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 _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: @@ -5986,9 +6133,21 @@ 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 + else: + num_tokens, load_async = self.kv_connector_manager.query_num_new_matched_tokens( + req, local_end + ) offered_end = min(local_end + num_tokens, req.prompt_len - 1) position = local_end @@ -6037,9 +6196,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: @@ -6047,13 +6204,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. @@ -6067,6 +6220,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 @@ -6080,9 +6250,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): @@ -7287,6 +7455,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 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/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index e247c93a389d..0a3dd9e0a97f 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -474,6 +474,12 @@ 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 def __init__( self, @@ -1425,6 +1431,9 @@ def _maybe_init_kv_connector_manager(self) -> None: 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 @@ -4667,6 +4676,89 @@ 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 _kv_connector_start_batch( self, scheduled_batch: ScheduledRequests) -> tuple[bool, bool]: """Prepare connector work and return fleet/local compute readiness.""" @@ -4857,6 +4949,7 @@ def _kv_connector_prepare_batch( ] if self.enable_attention_dp else []) self.kv_connector_manager.take_scheduled_requests_pending_load( scheduled_batch) + self._check_uncharged_parks_left_batch(scheduled_batch) self._pad_attention_dp_batch_emptied_by_connector(scheduled_batch) self._kv_connector_start_loads(scheduled_batch) return loading_context diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index c2f374b8e4e1..e99d4cbb26d1 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -1142,7 +1142,12 @@ def allocation_scope(current_stage: ExecutorMemoryType): connector_worker, connector_scheduler, enable_attention_dp=mapping.enable_attention_dp, - speculative_prefill_only=speculative_prefill_only) + 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. diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index ce9f5cef0a2b..4fdd3390d64a 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -75,6 +75,9 @@ 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 # PEFT accounting self._peft_cache_manager = peft_cache_manager @@ -155,6 +158,13 @@ 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() + def __init__( self, max_batch_size: int, @@ -170,7 +180,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 +328,40 @@ 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() + 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 +415,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 +438,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 +705,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. @@ -743,6 +836,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 @@ -1274,8 +1370,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 b6e8bfe453ef..11e1b21b8271 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2694,6 +2694,33 @@ def _check_segment_size(cls, value): 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 +2773,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": diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index 80f1c313b130..ec7150f2324d 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -3171,10 +3171,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": [ { diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index b1bde7f4d66a..b13158905fa5 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -4735,10 +4735,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": [ { diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 6bacef4ce177..ff7cb7b26268 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -872,6 +872,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/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 1e9b3d1c82cf..d0d0d693500c 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 @@ -726,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 @@ -753,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, ) @@ -5377,3 +5386,320 @@ 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 + + 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}) 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..79f925b495b3 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,126 @@ 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] diff --git a/tests/unittest/_torch/executor/test_kv_connector_initialization.py b/tests/unittest/_torch/executor/test_kv_connector_initialization.py index 6a88589c9a84..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,50 @@ 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( 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..65e24b50e600 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,224 @@ 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) + + +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 + 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)) + + assert req.context_current_position == 128 and req.context_chunk_size == 64 + assert kv_cache.capacity == 192 + assert connector.commits == [(0, 128, True)] + + 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 + + @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 4aee698eaf12..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 @@ -60,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 @@ -1112,3 +1114,134 @@ def test_refuted_connector_chunk_releases_unused_completion_reservation() -> Non 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/test_connector_adp.py b/tests/unittest/_torch/test_connector_adp.py index fb98484530e6..8a116ba0fb81 100644 --- a/tests/unittest/_torch/test_connector_adp.py +++ b/tests/unittest/_torch/test_connector_adp.py @@ -187,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] diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index f0d9c1e90e0f..11025e6c0e41 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -5500,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}) From 934702bac215fd46c6f0e59fb51a85334544b1a9 Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Tue, 27 Oct 2026 18:02:12 -0700 Subject: [PATCH 08/20] [None][perf] Add CTX park diagnostics to the iteration log (TLLM_CTX_PARK_STATS) Tuning admission-time connector queries on a disaggregated context server needs to see, per iteration, how much of the context budget went to requests the connector parked, how many first chunks were passed over for budget, how connector offers were placed, and where the executor thread spent its time. TLLM_CTX_PARK_STATS=1 appends a flat ", ctx_park = key=value ..." suffix, with no braces so existing parsers keep matching, to every print_iter_log line of the non-overlap executor loop. ParkStats records the wall and executor-thread CPU time of the sched/prep/meta/start/fwd/upd/term phases, the parked requests and tokens, the requests admitted without a token charge, the context tokens left for the forward pass, and the scheduler's budget-limited first chunks (BudgetTracker.num_ctx_budget_skips, including FORCE_CHUNK first chunks that wait for a fresh budget). KVCacheManagerV2 counts each connector query and offer, how the offer was placed, the admission-time grows, and the host time in the query, the placement and the page report, enabled by the executor and drained on every rank each iteration. TLLM_GIL_PROBE_US adds the wake-up overshoot of a sleeping probe thread as a GIL contention proxy; the executor stops it on shutdown. Both switches are debug-only environment variables, like the TLLM_PROFILE_* ones next to them in profiling.py, and add no configuration field. Signed-off-by: Simeng Liu --- docs/source/features/kv-cache-connector.md | 13 ++ .../kv_cache/kv_cache_manager_v2.py | 96 +++++++++++++ tensorrt_llm/_torch/pyexecutor/profiling.py | 128 +++++++++++++++++- tensorrt_llm/_torch/pyexecutor/py_executor.py | 111 +++++++++++++-- .../pyexecutor/scheduler/scheduler_v2.py | 13 ++ .../kv_cache/test_kv_cache_v2_scheduler.py | 31 +++++ .../test_kv_connector_executor_lifetime.py | 88 ++++++++++++ .../executor/test_kv_connector_v2_prefix.py | 38 ++++++ 8 files changed, 507 insertions(+), 11 deletions(-) diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index 853b2a91c817..ddae31575775 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -430,6 +430,19 @@ not booked at admission; `prepare_resources` grows them as on the final-batch pa 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 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 902273d771cd..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 @@ -1806,6 +1807,9 @@ class KVCacheManagerV2(BaseResourceManager): # `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 @@ -5860,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, @@ -6031,9 +6038,12 @@ def query_connector_at_admission(self, req: LlmRequest) -> bool: 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 ) @@ -6091,6 +6101,8 @@ def cover_admission_offer(self, req: LlmRequest) -> bool: 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: @@ -6098,6 +6110,78 @@ def cover_admission_offer(self, req: LlmRequest) -> bool: 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: @@ -6144,18 +6228,30 @@ def _apply_connector_matched_prefix(self, req: LlmRequest) -> bool: 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 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 0a3dd9e0a97f..96854daedd72 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, @@ -480,6 +482,9 @@ class PyExecutor: # 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, @@ -634,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. @@ -761,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, @@ -2155,6 +2171,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, @@ -4759,6 +4777,62 @@ def _check_uncharged_parks_left_batch( 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.""" @@ -4947,9 +5021,17 @@ def _kv_connector_prepare_batch( 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 @@ -5546,7 +5628,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 @@ -5609,16 +5693,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) @@ -5687,7 +5774,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 @@ -5733,8 +5822,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, @@ -5784,7 +5874,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 diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 4fdd3390d64a..8d9dd3e17de7 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -78,6 +78,10 @@ def __init__( # 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 @@ -164,6 +168,7 @@ class KVCacheV2Scheduler(RequestScheduler): _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, @@ -348,6 +353,9 @@ def __init__( #: 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( @@ -753,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 @@ -858,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, 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 d0d0d693500c..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 @@ -3536,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( @@ -5465,6 +5479,8 @@ def test_without_the_knob_the_park_burns_its_chunk(self): 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) @@ -5703,3 +5719,18 @@ def test_pregrown_replay_is_asked_with_its_exact_chunk(self): 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/test_kv_connector_executor_lifetime.py b/tests/unittest/_torch/executor/test_kv_connector_executor_lifetime.py index 79f925b495b3..810fb7ea730d 100644 --- a/tests/unittest/_torch/executor/test_kv_connector_executor_lifetime.py +++ b/tests/unittest/_torch/executor/test_kv_connector_executor_lifetime.py @@ -920,3 +920,91 @@ def test_admission_query_without_capability_only_warns(monkeypatch, setting, rea 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_v2_prefix.py b/tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py index 65e24b50e600..d686634be560 100644 --- a/tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py +++ b/tests/unittest/_torch/executor/test_kv_connector_v2_prefix.py @@ -1212,6 +1212,40 @@ def test_connector_admission_query_enabled_gates(self, change): 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.""" @@ -1222,6 +1256,7 @@ def _asked(num_matched=128, grow_ok=True, capacity=64, chunk=64, scratch_reuse=F 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) @@ -1240,10 +1275,12 @@ def test_a_shifting_offer_is_covered_at_admission_and_served_in_full(self): 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) @@ -1252,6 +1289,7 @@ def test_without_pages_the_cover_fails_and_leaves_the_allocation(self): 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): From a3972bd46c1836dadaf8f8b62878f6735ec28a74 Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Tue, 27 Oct 2026 18:28:19 -0700 Subject: [PATCH 09/20] [None][perf] mooncake_store: encode block tokens in one numpy pass BlockHashChain.extend encoded every token with int.to_bytes and joined the pieces per block, which dominated the lookup of a long prompt. It now encodes all newly completed blocks with a single np.asarray(dtype=" --- .../connectors/mooncake_store/keys.py | 37 +++++--- .../executor/test_mooncake_store_common.py | 91 +++++++++++++++++++ 2 files changed, 117 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py index b6ec205a85e8..7d87003cd7a6 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py @@ -29,6 +29,8 @@ from dataclasses import dataclass, field from typing import List, Optional, Sequence, Tuple +import numpy as np + __all__ = [ "BlockHashChain", "KeyNamespace", @@ -118,20 +120,33 @@ 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 + num_full_blocks = len(tokens) // 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. ' 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 + + def test_key_namespace_separates_every_dimension(): base = dict( namespace="trtllm", From acaebab996c04b4ee3f2e6311fdba2a31591def4 Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Tue, 27 Oct 2026 18:35:23 -0700 Subject: [PATCH 10/20] [None][fix] mooncake_store: publish only computed blocks and count a save once it is queued The scheduler adapter published every block that had a hash and a page. The whole prompt is hashed at lookup, 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 by the extra KV tokens a one-model draft needs. Such a block was saved with KV that had not been computed, and since the save mark moved past it and _put never overwrites an existing key, a later prefix hit loaded it. Saves are now bounded by the blocks this pass completes: computed_position, plus this iteration's honoured restore (which the runtime leaves out of computed_position), plus num_scheduled_tokens. The offer is trimmed to what the runtime honoured. update_state_after_alloc settles the outstanding offer to request.py_num_connector_matched_tokens, which KVCacheManagerV2 reaches through the base by-layer-group fold, and raises when that count is not a whole-block part of the offer. _loads_for now raises when an honoured block has no destination page instead of silently dropping the tail, because the runtime already counts those tokens as computed. A save counts as stored only once the worker queued it. _saves_for returns the block ordinal its saves end at without moving the save mark; build_connector_meta records it as unconfirmed, and the next build promotes it only for requests the worker's wait_for_save reported through the new take_submitted_saves. Saves bound for a pass that never reached wait_for_save are therefore emitted again rather than lost. The 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; request_reset also drops the request's unconfirmed entry. The scheduler test helpers gain FakeWorker.take_submitted_saves, submit_saves, settle_offer, saved_pages and request_data's computed_position / num_scheduled_tokens keywords, and make_request carries a settable py_num_connector_matched_tokens. Signed-off-by: Simeng Liu --- docs/source/features/kv-cache-connector.md | 4 + .../connectors/mooncake_store/scheduler.py | 147 ++++++- .../connectors/mooncake_store/worker.py | 33 +- .../executor/test_mooncake_store_connector.py | 382 +++++++++++++++++- 4 files changed, 536 insertions(+), 30 deletions(-) diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index ddae31575775..80f58f1be044 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -786,6 +786,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 diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py index 0bda85e66232..e007e6f0d1a3 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py @@ -22,6 +22,10 @@ * 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. """ from typing import Dict, List, Optional, Tuple @@ -113,7 +117,9 @@ 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 @@ -132,6 +138,9 @@ def __init__(self, llm_args: TorchLlmArgs): self._tokens_per_block = int(llm_args.kv_cache_config.tokens_per_block) 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] = {} logger.warning( f"mooncake-store leader ready (role={self._config.role.value}, " @@ -205,17 +214,29 @@ def get_num_new_matched_tokens( return hit_blocks * self._tokens_per_block, False def update_state_after_alloc(self, request: LlmRequest, block_ids: List[int]): - """No-op: page indices are read from the scheduler output instead. + """Trim the outstanding offer to what the runtime honoured. + + Page indices still come from the scheduler output: 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`. `KVCacheManagerV2` reports + through `update_state_after_alloc_by_layer_group`, whose base form + folds the single layer group this connector accepts into this call. - 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`. + 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: Unused; see above. """ + self._settle_offer(request) # ---- work lists ---- def build_connector_meta(self, scheduler_output: SchedulerOutput) -> MooncakeStoreMetadata: """Turn this iteration's scheduled requests into load and save lists.""" + self._confirm_submitted_saves() metadata = MooncakeStoreMetadata() for request_data in (*scheduler_output.new_requests, *scheduler_output.cached_requests): state = self._requests.get(request_data.request_id) @@ -231,6 +252,9 @@ 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 @@ -241,10 +265,14 @@ def build_connector_meta(self, scheduler_output: SchedulerOutput) -> MooncakeSto 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: @@ -268,12 +296,67 @@ 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. """ self._requests.pop(request.request_id, None) + self._unconfirmed_saves.pop(request.request_id, None) # ---- 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 once attention data + parallelism or asynchronous loads are supported, 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() @@ -308,25 +391,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/worker.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py index e411ce22419b..059a4afd5d29 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py @@ -26,11 +26,12 @@ 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. 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 @@ -241,6 +242,11 @@ def __init__(self, llm_args: TorchLlmArgs): # 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. @@ -622,6 +628,23 @@ 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: only the adapter colocated with + # rank 0's worker 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] diff --git a/tests/unittest/_torch/executor/test_mooncake_store_connector.py b/tests/unittest/_torch/executor/test_mooncake_store_connector.py index 143e0d2efbb9..bf72e126697c 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_connector.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_connector.py @@ -369,6 +369,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, ) @@ -1027,6 +1030,41 @@ def refuse(_keys): assert worker.get_finished([42], []) == ([42], []) +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 + + def synchronize(self): + pass + + 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() + + # ---- host staging ---- @@ -1294,16 +1332,27 @@ def test_staging_narrows_the_batch_to_the_budget(store_config, fake_store, stage 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()) @@ -1311,18 +1360,47 @@ def make_scheduler(store_config, hit_blocks=0): 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 @@ -1351,15 +1429,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] @@ -1373,15 +1458,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)) @@ -1395,6 +1581,126 @@ 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]) +def test_scheduler_saves_computed_tail_after_restore(store_config, local_blocks): + scheduler = make_scheduler(store_config, hit_blocks=2) + 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_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)) @@ -1413,6 +1719,62 @@ 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 + + +def test_scheduler_rejects_an_honour_beyond_the_offer(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, False) + with pytest.raises(RuntimeError, match="whole-block part of"): + settle_offer(scheduler, request, 2 * TOKENS_PER_BLOCK, [10, 11, 12]) + + +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)) From 6015401a4207c68a115f80228bb5b5c3691c1b0f Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Tue, 27 Oct 2026 18:57:13 -0700 Subject: [PATCH 11/20] [None][feat] mooncake_store: hand failed prefix loads back for a restart on a single rank A page the pool offered at lookup can be gone by the time it is loaded: evicted from a full pool, or left behind by a put that never completed. The runtime has already counted those positions as computed, so the worker used to fail the server on any missing page. The worker now declares reports_failed_loads. On a rank that holds whole pages and its own batch (a single-rank server, or an attention-DP owner once this connector supports attention DP) a request that lost a page is recorded for take_failed_load_requests instead, and the executor restarts only that request without the offered prefix. Under tensor parallelism the ranks hold shards of the same page and cannot restart a request alone, so there the load still raises with the existing "failed to load ... already reported as computed" error. Loads go through a new _load_pages(keys, addresses, sizes, owners, staging, handle, batch_size, *, recover), which later work reuses for asynchronous loads. It skips the remaining pages of a request that already lost one, treats a result list of the wrong length as all pages failed, scatters only the staged slots that landed and syncs the stream per batch, so no write into a restarted request's pages is in flight when it is reported. It returns the failed requests, the failed page count, the first missing key and the bytes that landed. A failed request's saves leave the bound metadata at once through drop_bound_saves, which edits the saves list wait_for_save reads in place and is idempotent with the manager's forget_request; otherwise the pass would publish pages the request never computed under its valid keys. The fake-store tests cover the single-rank report and its TP raise, the skip of a failed request's later pages, the save drop for a mixed batch and for a batch the failed request was all of, and staged loads that scatter only the landed pages. TP variants set _recover_failed_loads on the built worker rather than changing the world size, which would change every key namespace. Signed-off-by: Simeng Liu --- docs/source/features/kv-cache-connector.md | 8 + .../connectors/mooncake_store/worker.py | 235 +++++++++++++++--- .../executor/test_mooncake_store_connector.py | 191 +++++++++++++- 3 files changed, 401 insertions(+), 33 deletions(-) diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index 80f58f1be044..6b3d46cb42a0 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -838,6 +838,14 @@ sliding-window attention, and `sparse_attention_config` with an index-V cache un `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. +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. The worker declares +[`reports_failed_loads`](#2-worker-interface-kvcacheconnectorworker), so on a single-rank server +(and on attention-DP owners once `mooncake-store` supports attention data parallelism) 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. + ### GPUDirect RDMA By default the connector registers the KV pools with Mooncake so the NIC reads and writes pages in diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py index 059a4afd5d29..319743afcdb1 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py @@ -19,8 +19,13 @@ 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. +Loads are synchronous, 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, and an attention-DP owner once this connector supports attention data +parallelism. 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. 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 @@ -42,7 +47,7 @@ import traceback from collections import defaultdict 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 @@ -178,6 +183,13 @@ def _stream_handle(stream: Optional[torch.cuda.Stream]) -> int: class MooncakeStoreConnectorWorker(KvCacheConnectorWorker): """Moves KV pages between this rank's GPU cache and the Mooncake pool.""" + #: 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 #: 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. @@ -203,6 +215,20 @@ def __init__(self, llm_args: TorchLlmArgs): # 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. self._peer_namespaces: Dict[int, Tuple[KeyNamespace, ...]] = {} + # 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 = ( + getattr(llm_args, "enable_attention_dp", False) is True 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() # Checked before the segment is mounted, since the kernel answers an # unaffordable total by killing the process rather than by failing the @@ -533,7 +559,21 @@ def count_prefix_hit(self, block_hashes: Sequence[bytes]) -> int: # ---- 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. + + Args: + stream: The executor's stream, on which staged scatters are issued. + + Raises: + RuntimeError: When a page is missing on a tensor-parallel rank, or + a zero-copy load has no allocation fence to wait on. + """ metadata: Optional[MooncakeStoreMetadata] = self.get_connector_meta() if metadata is None or not metadata.loads: return @@ -541,6 +581,9 @@ def start_load_kv(self, stream: torch.cuda.Stream): keys, addresses, sizes, total_pages = self._resolve(metadata.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 metadata.loads for _ in entry.pages] staging = self._load_staging handle = _stream_handle(stream) if staging is not None else 0 @@ -559,42 +602,174 @@ def start_load_kv(self, stream: torch.cuda.Stream): ) 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 + failed_requests, num_failed_pages, first_failure, _ = self._load_pages( + keys, + addresses, + sizes, + owners, + staging, + handle, + self._batch_size, + recover=self._recover_failed_loads, + ) + 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}" + ) + 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, + ) -> Tuple[Set[int], int, Optional[str], int]: + """Fetch pages batch by batch, scattering staged ones onto the device. + + 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. + + 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 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] if staging is None: - target_addresses, target_sizes = list(batch_addresses), list(batch_sizes) + target_addresses, target_sizes = batch_addresses, 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 + batch_keys, target_addresses, target_sizes ) - failed = [ - key - for key, result in zip(batch_keys, results) - if not isinstance(result, int) or result < 0 + 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 ] - 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 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) + 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. + 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) + return failed_requests, num_failed_pages, first_failure, loaded_bytes - logger.debug(f"mooncake-store rank {self._rank} loaded {total_pages} pages") + 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`. diff --git a/tests/unittest/_torch/executor/test_mooncake_store_connector.py b/tests/unittest/_torch/executor/test_mooncake_store_connector.py index bf72e126697c..aebfb81ab732 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_connector.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_connector.py @@ -36,6 +36,7 @@ import torch from tensorrt_llm._torch.pyexecutor.connectors.kv_cache_connector import ( + KvCacheConnectorManager, RequestData, SchedulerOutput, ) @@ -61,6 +62,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, ) @@ -934,13 +936,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): @@ -1188,9 +1339,11 @@ def test_staging_get_scatters_back_to_the_device_regions(store_config, fake_stor assert staged_copies == expected -def test_staging_does_not_scatter_a_failed_load(store_config, fake_store, staged_copies): +@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) @@ -1199,8 +1352,12 @@ def test_staging_does_not_scatter_a_failed_load(store_config, fake_store, staged worker.bind_connector_meta( SimpleNamespace(loads=[RequestTransfers(7, [PageTransfer(block_hash, 0, 1)])], saves=[]) ) - with pytest.raises(RuntimeError, match="failed to load"): + 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 @@ -1208,6 +1365,34 @@ def test_staging_does_not_scatter_a_failed_load(store_config, fake_store, staged 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: From 5602a84e0123650a807d7159ff195b7a1002ae5b Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Tue, 27 Oct 2026 19:03:00 -0700 Subject: [PATCH 12/20] [None][feat] mooncake_store: support attention data parallelism The attention-DP framework creates a scheduler adapter and a worker on every owner and asks a connector to opt in, but mooncake-store declared nothing, so the executor creator refused it under enable_attention_dp, including for a capacity-only server. Its keys also assumed tensor parallelism: each rank keyed pages w{N}r{rank} and a lookup required every rank's shard, so under attention DP no owner could ever find a prefix, because no other owner writes the same page under a different rank. Both connector halves now declare supports_attention_dp. The worker keys pages by attention shard rather than by global rank: 0 of 1 under attention DP, where every owner holds all heads, and the global rank and world size under tensor parallelism. All owners therefore share one key space (w1r0), a prefix one owner saved is a hit for every other owner and for a single-rank server with the same model key and page layout, and a lookup asks one key per layer group per block. Tensor-parallel keys are byte-identical to before. The global rank still names the process in the segment ledger, thread names and logs. Request state, transfers, completion and the store handle were already owner-local. With the failed-load restart in place, a page lost between lookup and load restarts only the affected request on its owner instead of fail-stopping every owner; the docs and the worker docstring now say so without a hedge. Docstrings that described a single rank-0 leader now describe the scheduler adapter, which runs on rank 0 under tensor parallelism and on every owner under attention DP. The connector docs gain an attention-DP subsection for Mooncake, drop attention DP from the unsupported list, replace the line that said no Mooncake backend exists, and state that aggregated serving under attention DP is not supported yet: when the attention-DP gate skips a prepared pass, a generation page already reported to the connector can be reassigned without the connector hearing of it. Disaggregated context servers and roles that save nothing are unaffected. Tests: the fake-store suite covers owners sharing pages across rank pairs, tensor-parallel lookups still needing every shard, pinned key strings for a single rank and for rank 2 of 4, the class declarations against validate_attention_dp, the restart formula on an owner versus a TP rank, and a capacity-only worker under attention DP. The real-pool suite saves a page from one owner and loads it byte-identical on another while a TP rank misses it. A new TP2 attention-DP end-to-end test, parametrized over the overlap scheduler and listed in l0_dgx_h100, checks that both owners join the pool, that rank 0 serves prompt tokens from it, and that every replayed copy agrees with the cold output. Signed-off-by: Simeng Liu --- docs/source/features/kv-cache-connector.md | 44 ++++- .../connectors/mooncake_store/keys.py | 14 +- .../connectors/mooncake_store/metadata.py | 10 +- .../connectors/mooncake_store/scheduler.py | 26 +-- .../connectors/mooncake_store/validation.py | 4 +- .../connectors/mooncake_store/worker.py | 85 ++++++---- .../llmapi/test_llm_api_connector_mooncake.py | 150 +++++++++++++++++- .../test_lists/test-db/l0_dgx_h100.yml | 2 + .../executor/test_mooncake_store_connector.py | 146 ++++++++++++++++- .../executor/test_mooncake_store_real_pool.py | 50 +++++- 10 files changed, 464 insertions(+), 67 deletions(-) diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index 6b3d46cb42a0..5323e5ccbe01 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -661,7 +661,8 @@ 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. -This change adds no Mooncake backend, deployment configuration or capacity tuning. +The [`mooncake-store`](#mooncake-store) connector declares `supports_attention_dp`; see +[its notes on attention data parallelism](#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 @@ -826,13 +827,34 @@ 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. + +`mooncake-store` does not declare `supports_speculative_prefill`, so the +[speculative-prefill-only mode](#speculative-decoding-on-a-disaggregated-context-server) is refused +for it at startup. + ### 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, speculative decoding and Mamba/hybrid models, 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 @@ -840,11 +862,17 @@ 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. The worker declares -[`reports_failed_loads`](#2-worker-interface-kvcacheconnectorworker), so on a single-rank server -(and on attention-DP owners once `mooncake-store` supports attention data parallelism) 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. +[`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 diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py index 7d87003cd7a6..611a32d998d7 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py @@ -20,9 +20,10 @@ chain is seeded with a `ReuseScope` covering what the tokens do not record. 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. """ import hashlib @@ -156,9 +157,10 @@ class KeyNamespace: namespace: str model_key: str - #: Global rank of the shard whose KV these bytes are, and the world size it - #: was produced under. Both are needed: rank 3 of 8 holds different heads - #: than rank 3 of 4. + #: The attention shard whose KV these bytes are, and the shard count it was + #: produced under: the global rank and world size under tensor parallelism, + #: 0 of 1 under attention DP, where every owner holds all heads. Both are + #: needed: rank 3 of 8 holds different heads than rank 3 of 4. rank: int world_size: int layer_group_id: int diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py index f1eb2db7ee14..3736bf1b753c 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 diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py index e007e6f0d1a3..b30f66454d79 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py @@ -12,10 +12,12 @@ # 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; @@ -130,6 +132,10 @@ def __init__(self, chain: BlockHashChain): 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 + def __init__(self, llm_args: TorchLlmArgs): super().__init__(llm_args) @@ -143,7 +149,7 @@ def __init__(self, llm_args: TorchLlmArgs): self._unconfirmed_saves: Dict[int, int] = {} logger.warning( - f"mooncake-store leader ready (role={self._config.role.value}, " + f"mooncake-store scheduler adapter ready (role={self._config.role.value}, " f"tokens_per_block={self._tokens_per_block})" ) @@ -278,11 +284,13 @@ def build_connector_meta(self, scheduler_output: SchedulerOutput) -> MooncakeSto 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 diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py index 0abca8110e6d..649092c6f24f 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py @@ -16,8 +16,8 @@ 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`, so they are not repeated here. Checks run at construction, before any request is admitted, so a bad deployment fails at startup instead of after the first cache hit. diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py index 319743afcdb1..50bf5e9afd5a 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py @@ -17,15 +17,22 @@ 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. +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, 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, and an attention-DP owner once this connector supports attention data -parallelism. 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. +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. 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 @@ -78,10 +85,11 @@ _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() @@ -99,9 +107,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 @@ -183,6 +192,10 @@ def _stream_handle(stream: Optional[torch.cuda.Stream]) -> int: 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 @@ -207,22 +220,32 @@ 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, ...]] = {} # 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 = ( - getattr(llm_args, "enable_attention_dp", False) is True or self._world_size == 1 - ) + 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 ) @@ -284,7 +307,8 @@ def __init__(self, llm_args: TorchLlmArgs): 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})" @@ -356,11 +380,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: @@ -485,18 +509,19 @@ def _open_staging(self, addressing: PageAddressing) -> None: f"more. Lower transfer_batch_size to make the reduction explicit." ) - def _namespace(self, rank: int, layer_group_id: int, bytes_per_page: int) -> KeyNamespace: + 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: @@ -511,9 +536,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. @@ -803,9 +828,9 @@ 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: only the adapter colocated with - # rank 0's worker takes these, so the other ranks must not grow the - # set every pass. + # 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]: 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..836bbd443088 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,16 @@ 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, 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 @@ -41,7 +42,10 @@ 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,6 +53,7 @@ from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.scheduler import ( MooncakeStoreConnectorScheduler, ) +from tensorrt_llm.llmapi import SchedulingParams from tensorrt_llm.llmapi.llm_args import KvCacheConfig, KvCacheConnectorConfig, MooncakeStoreConfig from ..conftest import llm_models_root @@ -343,6 +348,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. 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/unittest/_torch/executor/test_mooncake_store_connector.py b/tests/unittest/_torch/executor/test_mooncake_store_connector.py index aebfb81ab732..2f47b93c8be2 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_connector.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_connector.py @@ -288,11 +288,14 @@ 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 +) -> 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, @@ -326,6 +329,7 @@ def make_worker( *, layout: KvCacheLayout | None = None, run_dir: str | None = None, + enable_attention_dp: bool = False, ) -> Iterator[MooncakeStoreConnectorWorker]: """Build a worker and shut it down before the test call phase ends. @@ -333,7 +337,9 @@ 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) + ) fake_store.workers.append(worker) if layout is not None: worker.register_kv_cache_layout(layout) @@ -1216,6 +1222,142 @@ def synchronize(self): assert worker.take_submitted_saves() == set() +# ---- 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 == [] + + # ---- host staging ---- 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..2d8dcd49867e 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,7 @@ def open_worker( layout: KvCacheLayout | None = None, stage_through_host: bool, model_key: str = MODEL_KEY, + enable_attention_dp: bool = False, ) -> Iterator[MooncakeStoreConnectorWorker]: """A worker holding a real store handle, over `buffer` or `layout` as its KV cache. @@ -289,7 +293,9 @@ def open_worker( run_dir=pool.run_dir, ) 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 @@ -445,6 +451,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. From 9661d622d9de8a9d8abac1ed809b24e1368735fd Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Tue, 27 Oct 2026 19:24:15 -0700 Subject: [PATCH 13/20] [None][feat] mooncake_store: block-aligned local reuse under one-model drafts (local_reuse_backoff) A one-model draft (MTP or EAGLE3) makes KVCacheManagerV2 back every local reuse match off by its prompt lookahead, so a local match of whole blocks ends one or more tokens short of a block boundary. mooncake-store offers only from a block boundary, so after any local hit it would offer nothing. The worker now fills in the two whole-block reuse declarations the cache manager reads once it is built after the connector. A role that reads the pool (both, consumer) always declares requires_block_aligned_reuse, and declares drops_reuse_match_backoff according to the new kv_connector_config.mooncake_store.local_reuse_backoff: - "off", the default, drops the backoff and takes local matches whole. The draft KV at the tail of a reused block may belong to another continuation, which can lower draft acceptance but never the output. - "align" keeps the backoff and has the manager align local matches down to a block, so the pool serves the trimmed block with the rest of the prefix at the cost of that block's transfer. Producer and capacity roles never offer, so they declare neither and keep the manager's own reuse. Every rank resolves the same role and arguments, so the declarations agree across the ranks of a server. The mode is an engine setting read from llm_args through the new config.local_reuse_backoff helper on every rank. It is not rendered into the client config or restated from an inherited one, so a deployment that inherits MOONCAKE_CONFIG_PATH without a mooncake_store block gets "off". The field is telemetry=False. The setting takes effect once the connector may run next to one-model speculative decoding. The single-block synchronous offer that the source branch made under "align" with asynchronous loads is not ported; this connector loads synchronously. Signed-off-by: Simeng Liu --- docs/source/features/kv-cache-connector.md | 18 ++++++ .../connectors/mooncake_store/config.py | 22 ++++++- .../connectors/mooncake_store/worker.py | 19 +++++- tensorrt_llm/llmapi/llm_args.py | 13 ++++ .../schemas/trtllm-serve-config.schema.json | 13 ++++ .../trtllm-serve-disagg-config.schema.json | 13 ++++ .../executor/test_mooncake_store_connector.py | 62 ++++++++++++++++++- .../executor/test_mooncake_store_settings.py | 32 ++++++++++ 8 files changed, 187 insertions(+), 5 deletions(-) diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index 5323e5ccbe01..d9a12aff469d 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -848,6 +848,24 @@ adapter and a worker on every owner, each deciding and moving pages for its own [speculative-prefill-only mode](#speculative-decoding-on-a-disaggregated-context-server) is refused for it at startup. +`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. The setting +has no effect while this connector refuses speculative decoding. + ### Restrictions Checked at startup, before any request is admitted, because each one's failure mode is KV replayed diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py index 4cbd8243bd11..66836f18afa4 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py @@ -29,7 +29,7 @@ import socket from dataclasses import dataclass from enum import Enum -from typing import TYPE_CHECKING, Any, Optional, Tuple +from typing import TYPE_CHECKING, Any, Literal, Optional, Tuple if TYPE_CHECKING: from tensorrt_llm.llmapi.llm_args import MooncakeStoreConfig, TorchLlmArgs @@ -41,6 +41,7 @@ "SEGMENTS_DIR_NAME", "StoreRole", "local_address", + "local_reuse_backoff", "parse_size", "pool_config", "provisioned_config_path", @@ -249,6 +250,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. diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py index 50bf5e9afd5a..51a47d57bf9d 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py @@ -65,7 +65,13 @@ from ..kv_cache_connector import 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 @@ -213,6 +219,17 @@ def __init__(self, llm_args: TorchLlmArgs): validate_llm_args(llm_args) self._config = MooncakeStoreConnectorConfig.resolve(llm_args) + 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`. diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 11e1b21b8271..2411b986d7d8 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2677,6 +2677,19 @@ 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.") + 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", mode="after") @classmethod diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index ec7150f2324d..820c1fa96bea 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -4181,6 +4181,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.", diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index b13158905fa5..c4bff194a140 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -5745,6 +5745,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.", diff --git a/tests/unittest/_torch/executor/test_mooncake_store_connector.py b/tests/unittest/_torch/executor/test_mooncake_store_connector.py index 2f47b93c8be2..43e8bcc72f17 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_connector.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_connector.py @@ -289,7 +289,10 @@ def set_pool_setting(store_config, **settings): def make_llm_args( - *, run_dir: str | None = None, enable_attention_dp: bool = False + *, + run_dir: str | None = None, + enable_attention_dp: bool = False, + local_reuse_backoff: str = "off", ) -> SimpleNamespace: return SimpleNamespace( model="/models/test-model", @@ -303,7 +306,9 @@ def make_llm_args( # 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 + ), ), ) @@ -330,6 +335,7 @@ def make_worker( layout: KvCacheLayout | None = None, run_dir: str | None = None, enable_attention_dp: bool = False, + local_reuse_backoff: str = "off", ) -> Iterator[MooncakeStoreConnectorWorker]: """Build a worker and shut it down before the test call phase ends. @@ -338,7 +344,11 @@ def make_worker( fixture teardown would run too late to keep it quiet. """ worker = MooncakeStoreConnectorWorker( - make_llm_args(run_dir=run_dir, enable_attention_dp=enable_attention_dp) + make_llm_args( + run_dir=run_dir, + enable_attention_dp=enable_attention_dp, + local_reuse_backoff=local_reuse_backoff, + ) ) fake_store.workers.append(worker) if layout is not None: @@ -1358,6 +1368,52 @@ def test_capacity_only_worker_builds_under_attention_dp(store_config, fake_store 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 + + # ---- host staging ---- diff --git a/tests/unittest/_torch/executor/test_mooncake_store_settings.py b/tests/unittest/_torch/executor/test_mooncake_store_settings.py index a740afcc8a4b..06bf6d22d1dc 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 ( @@ -246,6 +248,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) From e41b5c75792ebc551f869f89edf77260da47f981 Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Tue, 27 Oct 2026 19:38:27 -0700 Subject: [PATCH 14/20] [None][feat] mooncake_store: speculative prefill on disaggregated context servers The speculative-prefill-only mode lets one-model MTP or EAGLE3 run next to a KV connector on a disaggregated context server, but only for a connector whose two halves declare supports_speculative_prefill. mooncake-store now declares it on both and makes the promises behind the declaration true. A one-model draft writes the KV at a position from the prompt tokens after it, so a stored block's draft-layer KV depends on tokens past the block. BlockHashChain takes a keyword-only prompt_lookahead: with a lookahead of L, block b is hashed over tokens[b*T : (b+1)*T + L] at a one-block stride over the same int64 payload, so a block exists only once its lookahead does, and a non-zero L re-digests the scope seed with b"prompt-lookahead" and L. With L = 0 the seed and every key are byte-identical to before, so pools already populated stay addressable, while a speculative server never shares pages with a server without a draft or with a different lookahead. The scheduler adapter reads the lookahead with draft_prompt_lookahead (one for MTP-Eagle and EAGLE3, max_draft_len for vanilla MTP), imported inside its constructor so importing the connector package does not pull in every drafter. It raises when max_draft_len is still unresolved, since vanilla MTP's lookahead is its draft length and lookahead-0 keys would silently name draft KV that reads ahead. Offers are already bounded by the chain, so a block whose lookahead runs past the prompt is neither offered nor saved. Saves were already capped at computed blocks, which this mode needs because one-model speculation allocates KV slots past every context chunk. The worker moves whole pages of every layer group, draft layers included, and its per-layer hooks are no-ops. Before the store is opened it runs the new validate_speculative_config, an allow-list of one-model MTP and EAGLE3 whose draft layers share the target KV cache: any other drafter, or a separate draft KV cache, would leave missing or stale draft KV behind a restored prefix and is refused. A capacity role addresses no page and skips the check, so a disaggregated generation server may run one-model MTP or EAGLE3, with or without attention DP. Tests pin key bytes built by hand from the framed scope and the lookahead digest for lookahead 0, 1 and 3 with salt and multimodal scopes, and cover lookahead keys, the unresolved draft length, chunk-ahead capacity pages, the drafter allow-list and the creator's acceptance of the real classes. An end-to-end test runs Qwen3-8B with EAGLE3 as a cold and a warm disaggregated context server against a live pool, with a generation server pulling each request, and requires reuse, matching first and draft tokens, and the refusal of a request that is not context-only; it is listed in l0_h100. Signed-off-by: Simeng Liu --- docs/source/features/kv-cache-connector.md | 46 +++- .../connectors/mooncake_store/keys.py | 62 ++++- .../connectors/mooncake_store/scheduler.py | 41 +++- .../connectors/mooncake_store/validation.py | 70 +++++- .../connectors/mooncake_store/worker.py | 19 +- .../llmapi/test_llm_api_connector_mooncake.py | 153 ++++++++++++- .../test_lists/test-db/l0_h100.yml | 1 + .../executor/test_mooncake_store_common.py | 137 ++++++++++++ .../executor/test_mooncake_store_connector.py | 211 +++++++++++++++++- ...executor_creator_connector_spec_prefill.py | 79 ++++++- 10 files changed, 786 insertions(+), 33 deletions(-) diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index d9a12aff469d..009ebdd9dc7c 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -263,6 +263,9 @@ At runtime: * 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 Two `RequestData` fields can be reported empty depending on the cache configuration. A connector must tolerate both. @@ -661,8 +664,9 @@ 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`; see -[its notes on attention data parallelism](#attention-data-parallelism-and-speculative-prefill). +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 @@ -844,9 +848,25 @@ adapter and a worker on every owner, each deciding and moving pages for its own as on a single-rank server; see [Restrictions](#restrictions). * A `role: capacity` server may run with attention DP; it transfers nothing either way. -`mooncake-store` does not declare `supports_speculative_prefill`, so the -[speculative-prefill-only mode](#speculative-decoding-on-a-disaggregated-context-server) is refused -for it at startup. +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 @@ -863,20 +883,22 @@ block boundary the pool serves from. 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. The setting -has no effect while this connector refuses speculative decoding. +`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. ### 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, 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. The worker declares diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py index 611a32d998d7..56623793aaf7 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py @@ -19,6 +19,13 @@ 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 attention shard that produced them, the layer group inside that shard, the tokens each page holds and @@ -91,13 +98,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 @@ -105,6 +146,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.""" @@ -113,6 +159,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 @@ -122,7 +170,8 @@ def extend(self, tokens: Sequence[int]) -> Sequence[bytes]: The full chain, indexed by block ordinal. """ tokens_per_block = self._tokens_per_block - num_full_blocks = len(tokens) // 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( @@ -137,16 +186,21 @@ def extend(self, tokens: Sequence[int]) -> Sequence[bytes]: # keys are unchanged; it is done in one call for every new block because # per-token encoding dominated the lookup of a long prompt. A token # outside int64 raises OverflowError before any block is appended. + # Windows of a block plus its lookahead start one block apart, so the + # payload runs `lookahead` tokens past the last new block. payload = memoryview( np.asarray( - tokens[first_new * tokens_per_block : num_full_blocks * tokens_per_block], + tokens[ + first_new * tokens_per_block : num_full_blocks * tokens_per_block + lookahead + ], dtype=" 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. @@ -375,7 +408,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. diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py index 649092c6f24f..4a4c8af35456 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py @@ -17,7 +17,9 @@ 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 and Mamba caches are rejected for all connectors in -`py_executor`, so they are not repeated here. +`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. diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py index 51a47d57bf9d..a39bc2806f80 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py @@ -45,6 +45,10 @@ either way. Only the save thread failing to start is fatal, because that would strand every request behind pages nothing will release. +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. @@ -85,7 +89,12 @@ 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"] @@ -209,6 +218,11 @@ class MooncakeStoreConnectorWorker(KvCacheConnectorWorker): #: 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. @@ -219,6 +233,9 @@ def __init__(self, llm_args: TorchLlmArgs): 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 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 836bbd443088..3d70350f3c68 100644 --- a/tests/integration/defs/llmapi/test_llm_api_connector_mooncake.py +++ b/tests/integration/defs/llmapi/test_llm_api_connector_mooncake.py @@ -19,8 +19,9 @@ 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, the key space attention -DP owners share, and the scheduler interaction when a request is preempted while -the connector is still reading its pages. +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`. @@ -37,7 +38,7 @@ 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, @@ -53,8 +54,14 @@ from tensorrt_llm._torch.pyexecutor.connectors.mooncake_store.scheduler import ( MooncakeStoreConnectorScheduler, ) -from tensorrt_llm.llmapi import SchedulingParams -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 @@ -599,3 +606,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_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/executor/test_mooncake_store_common.py b/tests/unittest/_torch/executor/test_mooncake_store_common.py index f0e059dd200c..26b73847fded 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_common.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_common.py @@ -294,6 +294,143 @@ def test_hash_chain_rejects_tokens_outside_int64(token): 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", diff --git a/tests/unittest/_torch/executor/test_mooncake_store_connector.py b/tests/unittest/_torch/executor/test_mooncake_store_connector.py index 43e8bcc72f17..77d72b8d2674 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_connector.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_connector.py @@ -72,6 +72,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 @@ -293,6 +299,7 @@ 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", @@ -302,6 +309,7 @@ def make_llm_args( 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( @@ -336,6 +344,7 @@ def make_worker( 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. @@ -348,6 +357,7 @@ def make_worker( 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) @@ -1414,6 +1424,85 @@ def test_mooncake_asks_for_block_aligned_local_reuse_only_where_it_loads( 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 ---- @@ -1737,8 +1826,10 @@ def submit_saves(scheduler, metadata): 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 @@ -1984,8 +2075,16 @@ def test_scheduler_publishes_only_computed_full_blocks(store_config): @pytest.mark.parametrize("local_blocks", [0, 1]) -def test_scheduler_saves_computed_tail_after_restore(store_config, local_blocks): - scheduler = make_scheduler(store_config, hit_blocks=2) +@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) == ( @@ -2027,6 +2126,110 @@ def test_scheduler_saves_computed_tail_after_restore(store_config, local_blocks) 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. 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 index 9f49fd6bdbbd..89393fd776c9 100644 --- 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 @@ -18,9 +18,13 @@ 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. +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 ( @@ -28,6 +32,17 @@ 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, ) @@ -72,18 +87,76 @@ def _eagle3() -> Eagle3DecodingConfig: 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, use_v2 + make_spec_config, connector_classes, use_v2 ) -> None: + worker_cls, scheduler_cls = connector_classes assert _connector_runs_speculative_prefill_only( - make_spec_config(), _Declaring, _Declaring, DISAGG, use_v2 + 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( From 0e5b3ff4fb8c2e67995d93f330f5c241468593d8 Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Tue, 27 Oct 2026 20:02:57 -0700 Subject: [PATCH 15/20] [None][feat] mooncake_store: soft-pinned and per-block grouped puts A full Mooncake pool evicts on the master, and two of its behaviours cost hits that the connector can avoid with per-put options. 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. That evicts blocks from the middle of prefixes still in use, and a prefix is replayed only up to its first missing block, so every stored block behind such a hole is recomputed. Soft-pinned objects are evicted only by the second, age-ordered pass. The new mooncake_store.soft_pin_ttl_ms soft-pins every page a server writes for that many milliseconds. A hit on a block needs every page of it: all layer groups and, under tensor parallelism, every rank's shard. The master evicts pages independently, so it can keep a shard that holds space but can never hit again. With the new mooncake_store.group_blocks, every page of a block is put into one eviction group, named by the new KeyNamespace.group as //w/, which leaves out the shard, the layer group, the tokens per block and the page size but keeps the attention world size. Both fields are off by default, telemetry=False, and apply only to a role that writes. They are rendered into the client config, read back by MooncakeStoreConnectorConfig.from_file, and restated from an inherited MOONCAKE_CONFIG_PATH; a config without a TTL restates the block to off. The JSON keys match the source branch's, so an existing client config keeps working, while its TRTLLM_MOONCAKE_STORE_SOFT_PIN_TTL_MS and TRTLLM_MOONCAKE_STORE_GROUP_BLOCKS environment overrides are not read. group_blocks must be a JSON boolean, since bool("false") is True. The worker builds a Mooncake ReplicateConfig through _replicate_config (the source branch's _put_config, renamed so it no longer shadows the worker's _put_config attribute). A soft pin alone shares one config across every put; grouped puts build one per call with a group per key that survived the existence filter. The constructor builds a probe config before the segment is mounted, so bindings older than 0.3.13 fail at startup with an ImportError naming SoftPinAction and the ReplicateConfig fields. With neither option the put keeps its three-argument form. A put whose every page is refused under a soft pin logs once that a TTL above the master's --max_kv_soft_pin_ttl (24 h by default) makes the master refuse every put. Signed-off-by: Simeng Liu --- docs/source/features/kv-cache-connector.md | 33 +++ ...trtllm_mooncake_store_connector_extra.yaml | 5 + .../connectors/mooncake_store/config.py | 33 ++- .../connectors/mooncake_store/keys.py | 20 +- .../connectors/mooncake_store/master.py | 4 + .../connectors/mooncake_store/settings.py | 10 + .../connectors/mooncake_store/worker.py | 97 +++++++- tensorrt_llm/llmapi/llm_args.py | 19 ++ .../schemas/trtllm-serve-config.schema.json | 26 +++ .../trtllm-serve-disagg-config.schema.json | 26 +++ .../executor/test_mooncake_store_common.py | 101 ++++++++ .../executor/test_mooncake_store_connector.py | 219 +++++++++++++++++- .../executor/test_mooncake_store_real_pool.py | 44 ++++ .../executor/test_mooncake_store_settings.py | 38 +++ 14 files changed, 663 insertions(+), 12 deletions(-) diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index 009ebdd9dc7c..304737524898 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -815,6 +815,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 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..2a5a8150e024 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,8 @@ 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 diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py index 66836f18afa4..68af0173e233 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py @@ -29,7 +29,7 @@ import socket from dataclasses import dataclass from enum import Enum -from typing import TYPE_CHECKING, Any, Literal, Optional, Tuple +from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple if TYPE_CHECKING: from tensorrt_llm.llmapi.llm_args import MooncakeStoreConfig, TorchLlmArgs @@ -240,6 +240,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. @@ -310,6 +322,16 @@ 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 + #: 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 def __post_init__(self) -> None: """Reject settings that would fail later, inside a transfer.""" @@ -321,6 +343,11 @@ 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.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}" + ) @property def capacity_only(self) -> bool: @@ -356,6 +383,10 @@ 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)), + 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), ) @staticmethod diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py index 56623793aaf7..f7dbb534b174 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/keys.py @@ -30,7 +30,8 @@ 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. +as garbage. With `group_blocks`, the pages of one block also share an eviction +group named by `KeyNamespace.group`. """ import hashlib @@ -234,3 +235,20 @@ def prefix(self) -> 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..002174e2db36 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py @@ -423,6 +423,10 @@ def _client_config( "role": StoreRole(pool.role).value, "transfer_batch_size": pool.transfer_batch_size, "stage_through_host": pool.stage_through_host, + "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 {}), } return config diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/settings.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/settings.py index 68c4244d0aed..9539f11c6ba1 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/settings.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/settings.py @@ -188,7 +188,11 @@ 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), + ("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),) # 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 # is worth reporting; the rest are defaults being resolved. @@ -200,6 +204,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/worker.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py index a39bc2806f80..20c00dd6b93c 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py @@ -43,7 +43,9 @@ 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. +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, @@ -204,6 +206,41 @@ 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 + + class MooncakeStoreConnectorWorker(KvCacheConnectorWorker): """Moves KV pages between this rank's GPU cache and the Mooncake pool.""" @@ -292,6 +329,19 @@ def __init__(self, llm_args: TorchLlmArgs): # 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, @@ -345,7 +395,9 @@ def __init__(self, llm_args: TorchLlmArgs): 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 ''})" ) if self.capacity_only: # Said plainly, since every other sign of a working connector is @@ -967,13 +1019,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. @@ -995,11 +1049,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 @@ -1008,6 +1067,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 ---- diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 2411b986d7d8..87291d2f0b54 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2655,6 +2655,25 @@ class MooncakeStoreConfig(StrictBaseModel): "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'.") + 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.") local_hostname: Optional[str] = Field( None, telemetry=False, diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index 820c1fa96bea..4b533ce753d9 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -4165,6 +4165,15 @@ "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": { + "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": [ { @@ -4270,6 +4279,23 @@ "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'.", diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index c4bff194a140..83d72edb6029 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -5729,6 +5729,15 @@ "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": { + "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": [ { @@ -5834,6 +5843,23 @@ "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'.", diff --git a/tests/unittest/_torch/executor/test_mooncake_store_common.py b/tests/unittest/_torch/executor/test_mooncake_store_common.py index 26b73847fded..fcfd94ee9283 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_common.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_common.py @@ -29,6 +29,7 @@ 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 ( MooncakeStoreConnectorConfig, @@ -41,6 +42,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, @@ -48,6 +50,7 @@ stage_batch_for_put, unstage_batch_after_get, ) +from tensorrt_llm.llmapi.llm_args import MooncakeStoreConfig pytestmark = pytest.mark.cpu_only @@ -455,6 +458,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 ---- @@ -594,6 +626,75 @@ def test_config_requires_an_explicit_model_key(store_config, tmp_path, monkeypat config.resolve_model_key("/models/MiniMax-M3/") +# ---- pool eviction 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_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") + + config = resolve_config() + + assert config.soft_pin_ttl_ms is None + assert config.group_blocks is False + + +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_settings(tmp_path): + rendered, config = _rendered_client_config( + tmp_path, soft_pin_ttl_ms=86400000, group_blocks=True + ) + assert rendered["soft_pin_ttl_ms"] == 86400000 + assert rendered["group_blocks"] is True + assert config.soft_pin_ttl_ms == 86400000 + assert config.group_blocks is True + + rendered, config = _rendered_client_config(tmp_path) + assert "soft_pin_ttl_ms" not in rendered + assert rendered["group_blocks"] is False + assert config.soft_pin_ttl_ms is None + assert config.group_blocks is False + + # 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. diff --git a/tests/unittest/_torch/executor/test_mooncake_store_connector.py b/tests/unittest/_torch/executor/test_mooncake_store_connector.py index 77d72b8d2674..a33802aa2e79 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_connector.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_connector.py @@ -27,10 +27,12 @@ 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 @@ -100,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 @@ -125,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] @@ -1242,6 +1247,216 @@ def synchronize(self): assert worker.take_submitted_saves() == set() +# ---- soft-pinned and grouped puts ---- + + +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 _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 + + monkeypatch.setattr(worker_module, "_replicate_config", fake_config) + return built + + +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_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],), {})] + + +@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_worker_groups_every_page_of_a_block(store_config, fake_store, monkeypatch): + """Each put names one group per key. + + 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), + ], + ) + ] + ) + 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" + + +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 _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) + + +class _SoftPinAction(Enum): + PRESERVE = 0 + ENABLE = 1 + DISABLE = 2 + + +class _ReplicateConfig: + def __init__(self): + self.soft_pin_action = _SoftPinAction.PRESERVE + self.soft_pin_ttl_ms = None + self.group_ids = None + + +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 + + +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 + + config = worker_module._replicate_config(None, ("a", "b")) + assert config.soft_pin_action is _SoftPinAction.PRESERVE + assert config.group_ids == ["a", "b"] + + +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) + + class OldReplicateConfig: + __slots__ = ("soft_pin_action", "soft_pin_ttl_ms") + + _fake_mooncake_store( + monkeypatch, ReplicateConfig=OldReplicateConfig, SoftPinAction=_SoftPinAction + ) + with pytest.raises(ImportError, match="group_ids"): + worker_module._replicate_config(None, ["a"]) + + # ---- attention data parallelism ---- 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 2d8dcd49867e..3e1ff619f82e 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_real_pool.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_real_pool.py @@ -278,6 +278,8 @@ def open_worker( 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, ) -> Iterator[MooncakeStoreConnectorWorker]: """A worker holding a real store handle, over `buffer` or `layout` as its KV cache. @@ -291,6 +293,8 @@ 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, ) with provision_pool(store_config, run_dir=pool.run_dir): worker = MooncakeStoreConnectorWorker( @@ -524,6 +528,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 06bf6d22d1dc..fb0319779f00 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_settings.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_settings.py @@ -158,6 +158,44 @@ 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_settings(monkeypatch, tmp_path): + """Soft pin and grouping 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 + ) + 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 + + +def test_an_inherited_config_without_soft_pin_turns_it_off(monkeypatch, tmp_path): + """A TTL 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) + + 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 + 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] + + +@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) + + 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( From a8f47f66413534a4b5b46f0a5d5587000b351386 Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Tue, 27 Oct 2026 20:07:32 -0700 Subject: [PATCH 16/20] [None][feat] mooncake_store: prefix-hole and load-phase statistics (stats_interval) Tuning a shared pool needs two numbers the connector did not report: how often a prefix lookup stops at a block the pool evicted from the middle of a chain it otherwise still holds, and where the time of a prefix load goes. The source branch read them behind three environment variables and a fixed interval; here one setting, mooncake_store.stats_interval (seconds, gt=0, finite, telemetry=False), turns on both. It is rendered into the client config, read back by MooncakeStoreConnectorConfig.from_file and restated from an inherited MOONCAKE_CONFIG_PATH like the eviction settings; a config without it restates the block to off. TRTLLM_MOONCAKE_HOLE_STATS, TLLM_M3_LOAD_PHASE_TIMING and TLLM_MOONCAKE_LOAD_PHASE_TIMING are not read. Only a role that reads the pool counts anything, since only it looks prefixes up and loads them. count_prefix_hit then folds its existing batch_is_exist answer into _PrefixHoleStats: lookups, looked-up and hit blocks, holes, the missing blocks inside them, and the stored tokens behind a hole that are recomputed anyway. The hit itself is unchanged and no store call is added. A "prefix-hole stats" line is logged at info level at most once per interval. The synchronous load times its phases with _timed_phase, which sums wall and thread CPU seconds per phase and reads no clock when stats are off: resolve, the allocation fence, and, inside _load_pages through its new phases argument, describe, get, unstage and sync. Each start_load_kv pass that moved pages is folded into _LoadStats with its pages, store calls, landed bytes and transfer time, and get_finished logs a "load stats" line at most once per interval. _LoadStats records and drains under a lock, so loads that later run on background threads can record into the same counters. Both lines are flushed when the worker shuts down. Signed-off-by: Simeng Liu --- docs/source/features/kv-cache-connector.md | 23 ++ .../connectors/mooncake_store/config.py | 12 + .../connectors/mooncake_store/master.py | 1 + .../connectors/mooncake_store/settings.py | 5 +- .../connectors/mooncake_store/worker.py | 327 +++++++++++++++++- tensorrt_llm/llmapi/llm_args.py | 10 + .../schemas/trtllm-serve-config.schema.json | 17 + .../trtllm-serve-disagg-config.schema.json | 17 + .../executor/test_mooncake_store_common.py | 31 +- .../executor/test_mooncake_store_connector.py | 175 ++++++++++ .../executor/test_mooncake_store_settings.py | 21 +- 11 files changed, 616 insertions(+), 23 deletions(-) diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index 304737524898..2477c5e8d742 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -919,6 +919,29 @@ rendered client config, so a deployment that inherits `MOONCAKE_CONFIG_PATH` wit `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), their + `pages`, store calls (`batches`) and `loaded_gib`, the transfer time from the destinations being + ready to the last page landing (`transfer_ms`), and each pass's mean wall and CPU milliseconds + for `fence`, `resolve`, `describe`, `get`, `unstage` and `sync`. + +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 diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py index 68af0173e233..0dd7680ff131 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py @@ -24,6 +24,7 @@ """ import json +import math import os import re import socket @@ -332,6 +333,9 @@ class MooncakeStoreConnectorConfig: #: 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.""" @@ -348,6 +352,11 @@ def __post_init__(self) -> None: 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: @@ -387,6 +396,9 @@ def from_file(path: str) -> "MooncakeStoreConnectorConfig": 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/master.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py index 002174e2db36..4b76192ae93e 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py @@ -427,6 +427,7 @@ def _client_config( # 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/settings.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/settings.py index 9539f11c6ba1..ba877acf296e 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/settings.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/settings.py @@ -192,7 +192,10 @@ def _restate_pool_block( ) # 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),) + 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 # is worth reporting; the rest are defaults being resolved. diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py index 20c00dd6b93c..a208231d5bbc 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py @@ -54,11 +54,17 @@ 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, Iterable, List, Optional, Sequence, Set, Tuple @@ -241,6 +247,187 @@ def _replicate_config(soft_pin_ttl_ms: Optional[int], group_ids: Optional[Sequen 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. +_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. Every method takes the lock, so a load that runs off the executor + thread can record into the same counters the executor thread logs. + """ + + __slots__ = ( + "_lock", + "loads", + "pages", + "batches", + "loaded_bytes", + "xfer_sum", + "xfer_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.xfer_sum = 0.0 + self.xfer_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, + ) -> 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`. + """ + with self._lock: + self.loads += 1 + self.pages += pages + self.batches += batches + self.loaded_bytes += loaded_bytes + self.xfer_sum += xfer_s + self.xfer_max = max(self.xfer_max, xfer_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) -> Optional[str]: + """The interval's summary, starting a new interval; None when nothing loaded.""" + with self._lock: + if not self.loads: + return None + n = 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 + ) + summary = ( + f"loads={n} pages={self.pages} batches={self.batches} " + f"loaded_gib={self.loaded_bytes / _GIB:.2f} " + f"transfer_ms mean={1000 * self.xfer_sum / n:.1f} " + f"max={1000 * self.xfer_max:.1f} " + 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.""" @@ -323,6 +510,15 @@ def __init__(self, llm_args: TorchLlmArgs): # 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() # Checked before the segment is mounted, since the kernel answers an # unaffordable total by killing the process rather than by failing the @@ -397,7 +593,8 @@ def __init__(self, llm_args: TorchLlmArgs): f"segment={self._config.global_segment_size / (1 << 30):.1f} GiB " 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"{', 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 ''})" ) if self.capacity_only: # Said plainly, since every other sign of a working connector is @@ -665,8 +862,83 @@ 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 + 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): @@ -689,7 +961,10 @@ def start_load_kv(self, stream: torch.cuda.Stream): if metadata is None or not metadata.loads: return - 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(metadata.loads) if not keys: return # `_resolve` emits one key per page in transfer order, so this names the @@ -711,9 +986,11 @@ 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}) + with _timed_phase(phases, "fence"): + fence({entry.request_id for entry in metadata.loads}) - failed_requests, num_failed_pages, first_failure, _ = self._load_pages( + 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, @@ -722,7 +999,16 @@ def start_load_kv(self, stream: torch.cuda.Stream): 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, + ) if failed_requests: # The executor takes these requests out of the batch and restarts # them; see `take_failed_load_requests`. Nothing of theirs is @@ -750,6 +1036,7 @@ def _load_pages( 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. @@ -768,6 +1055,8 @@ def _load_pages( 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 @@ -793,13 +1082,15 @@ def _load_pages( batch_keys = [keys[index] for index in indices] batch_addresses = [addresses[index] for index in indices] batch_sizes = [sizes[index] for index in indices] - if staging is None: - target_addresses, target_sizes = batch_addresses, batch_sizes - else: - target_addresses, target_sizes = describe_batch_for_get(staging, batch_sizes) - results = self._store.batch_get_into_multi_buffers( - batch_keys, target_addresses, target_sizes - ) + 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 len(results) != len(batch_keys): # Which pages landed is unknowable, so none of them counts. results = [-1] * len(batch_keys) @@ -830,10 +1121,14 @@ def _load_pages( 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. - unstage_batch_after_get(staging, batch_addresses, batch_sizes, handle, only=landed) + 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 def take_failed_load_requests(self) -> Set[int]: @@ -948,6 +1243,9 @@ def get_finished( Requests that have finished saving, and requests that have finished loading. """ + # 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 = [ @@ -1121,6 +1419,9 @@ def _resolve( def shutdown(self) -> None: """Stop the save thread, then release what it was reading. 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) thread = self._save_thread if thread is not None: self._save_queue.put(None) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 87291d2f0b54..36b707aea86e 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2674,6 +2674,16 @@ class MooncakeStoreConfig(StrictBaseModel): "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, diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index 4b533ce753d9..d8cc625e9263 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -4305,6 +4305,23 @@ "title": "Stage Through Host", "type": "boolean" }, + "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 83d72edb6029..e5be00381f5c 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -5869,6 +5869,23 @@ "title": "Stage Through Host", "type": "boolean" }, + "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/tests/unittest/_torch/executor/test_mooncake_store_common.py b/tests/unittest/_torch/executor/test_mooncake_store_common.py index fcfd94ee9283..5285770a29b0 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_common.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_common.py @@ -626,7 +626,7 @@ def test_config_requires_an_explicit_model_key(store_config, tmp_path, monkeypat config.resolve_model_key("/models/MiniMax-M3/") -# ---- pool eviction settings ---- +# ---- pool eviction and diagnostics settings ---- def test_config_eviction_settings_default_off_and_read_the_json(store_config): @@ -659,15 +659,36 @@ def test_config_group_blocks_must_be_a_json_boolean(store_config, value): 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 _rendered_client_config(tmp_path, **settings): @@ -679,20 +700,24 @@ def _rendered_client_config(tmp_path, **settings): return rendered, MooncakeStoreConnectorConfig.from_file(str(path)) -def test_client_config_renders_eviction_settings(tmp_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 + 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 diff --git a/tests/unittest/_torch/executor/test_mooncake_store_connector.py b/tests/unittest/_torch/executor/test_mooncake_store_connector.py index a33802aa2e79..191671a32d18 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_connector.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_connector.py @@ -1457,6 +1457,181 @@ class OldReplicateConfig: worker_module._replicate_config(None, ["a"]) +# ---- diagnostics ---- + + +def _store_blocks(fake_store, worker, hashes, blocks): + for block in blocks: + fake_store.objects.add(worker._namespaces[0].key(hashes[block])) + + +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 ---- diff --git a/tests/unittest/_torch/executor/test_mooncake_store_settings.py b/tests/unittest/_torch/executor/test_mooncake_store_settings.py index fb0319779f00..7ad6d42a73e5 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_settings.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_settings.py @@ -158,10 +158,10 @@ 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_settings(monkeypatch, tmp_path): - """Soft pin and grouping are what the ranks read from the client config.""" +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 + 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() @@ -170,13 +170,14 @@ def test_an_inherited_config_restates_the_eviction_settings(monkeypatch, tmp_pat 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_turns_it_off(monkeypatch, tmp_path): - """A TTL the inherited config leaves out is off, not unstated.""" +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) + pool = mooncake_store_config(soft_pin_ttl_ms=86400000, group_blocks=True, stats_interval=30) warnings = [] monkeypatch.setattr(settings.logger, "warning", warnings.append) @@ -184,10 +185,12 @@ def test_an_inherited_config_without_soft_pin_turns_it_off(monkeypatch, tmp_path 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]) @@ -196,6 +199,12 @@ def test_a_soft_pin_ttl_must_be_positive(ttl): 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_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( From 5450f5ef1c231bd85e78172c90cf67ec3e32d825 Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Tue, 27 Oct 2026 20:45:52 -0700 Subject: [PATCH 17/20] [None][feat] mooncake_store: asynchronous prefix loads, configurable host staging and admission-query support A prefix the pool serves is loaded inside the executor iteration that schedules the request, so every other request in the batch waits on the transfer. With the new mooncake_store.async_load, a role that reads the pool offers every hit asynchronously instead: the runtime parks the request outside the batch and async_load_workers background threads per rank load its pages while other requests keep computing. The scheduler adapter computes a parked request's transfer when its allocation is reported. update_state_after_alloc now delegates to the new update_state_after_alloc_by_layer_group, which settles the honoured count first, as before, and for an asynchronous offer builds the transfer from the request's full page lists, refusing a missing or invalid destination page. An offer honoured with no tokens still queues an empty transfer, because the runtime parks the request regardless. The transfer travels once, in MooncakeStoreMetadata.async_loads of the same iteration's metadata; the blocks it covers are never written back, and request_reset drops one not yet handed out. The adapter also declares supports_admission_query: a query only refills the request's offer state and transfers are built only once the allocation is reported, so a repeated admission-time answer replaces the earlier one. Every rank queues its own pages when it binds the metadata, so under tensor parallelism each rank loads its own shard. Binding never raises, since it runs outside the attention-DP failure vote; what cannot be queued (async_load off, no load thread, no allocation fence) is kept for the next start_load_kv or get_finished to raise. A load thread waits for the KV cache manager's allocation fence on its own thread, fetches through the shared _load_pages, and get_finished reports a request only after the runtime announced it and its pages landed or one was found missing. A missing page hands the request back through take_failed_async_load_requests on a single rank or an attention-DP owner, and fails the rank under tensor parallelism. Any exception on a load thread is kept and raised by every later start_load_kv and get_finished until shutdown. start_load_kv now consumes the bound synchronous loads, so a second call on the same metadata starts nothing. The pinned host memory per direction under stage_through_host becomes the configurable staging_buffer_bytes (default 1GiB, the previous fixed ceiling), which replaces staging.MAX_STAGING_BUFFER_BYTES and is what the node memory check counts. With async_load each load thread gets its own staging pool and CUDA stream from an even share of the load direction's budget, and no synchronous load pool is pinned, so the check still counts one budget per direction; a pool whose share is below one page is warned about. Shutdown signals the save and load threads together, joins them against one deadline, skips loads still queued, and unregisters every staging pool before closing the store, keeping any pool whose unregistration failed. The three settings are telemetry-free MooncakeStoreConfig fields, rendered into the client config, read back by MooncakeStoreConnectorConfig.from_file (async_load must be a JSON boolean) and restated from an inherited MOONCAKE_CONFIG_PATH. Asynchronous loads record into the stats_interval load line, which gains failed_pages and, with async_load, wait_ms and the load threads' queue. The integration list runs a new end-to-end test comparing an asynchronously loaded prefix with a synchronous run. Signed-off-by: Simeng Liu --- docs/source/features/kv-cache-connector.md | 63 +- ...trtllm_mooncake_store_connector_extra.yaml | 4 + .../connectors/mooncake_store/config.py | 29 + .../connectors/mooncake_store/master.py | 3 + .../connectors/mooncake_store/metadata.py | 12 +- .../connectors/mooncake_store/scheduler.py | 156 ++- .../connectors/mooncake_store/settings.py | 7 + .../connectors/mooncake_store/staging.py | 29 +- .../connectors/mooncake_store/validation.py | 22 +- .../connectors/mooncake_store/worker.py | 748 ++++++++++-- tensorrt_llm/llmapi/llm_args.py | 46 +- .../schemas/trtllm-serve-config.schema.json | 39 +- .../trtllm-serve-disagg-config.schema.json | 39 +- .../llmapi/test_llm_api_connector_mooncake.py | 83 +- .../integration/test_lists/test-db/l0_a10.yml | 1 + .../executor/test_mooncake_store_common.py | 119 ++ .../executor/test_mooncake_store_connector.py | 1045 ++++++++++++++++- .../executor/test_mooncake_store_real_pool.py | 48 + .../executor/test_mooncake_store_settings.py | 29 + 19 files changed, 2355 insertions(+), 167 deletions(-) diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index 2477c5e8d742..13420dcba028 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -392,6 +392,7 @@ A connector that starts a transfer inside `get_num_new_matched_tokens` must not 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`: @@ -933,10 +934,14 @@ its last line and skipped when there is nothing to report: 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), their - `pages`, store calls (`batches`) and `loaded_gib`, the transfer time from the destinations being - ready to the last page landing (`transfer_ms`), and each pass's mean wall and CPU milliseconds - for `fence`, `resolve`, `describe`, `get`, `unstage` and `sync`. +* `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 @@ -957,7 +962,8 @@ one-model MTP or EAGLE3 sharing the target KV cache; it bypasses individual requ 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. The worker declares +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 @@ -1015,3 +1021,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/llm-api/configs/trtllm_mooncake_store_connector_extra.yaml b/examples/llm-api/configs/trtllm_mooncake_store_connector_extra.yaml index 2a5a8150e024..22c7ad9f5015 100644 --- a/examples/llm-api/configs/trtllm_mooncake_store_connector_extra.yaml +++ b/examples/llm-api/configs/trtllm_mooncake_store_connector_extra.yaml @@ -30,3 +30,7 @@ kv_connector_config: # 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/pyexecutor/connectors/mooncake_store/config.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py index 0dd7680ff131..c7213f223354 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/config.py @@ -38,6 +38,7 @@ __all__ = [ "CLIENT_CONFIG_NAME", "CONFIG_PATH_ENV", + "DEFAULT_STAGING_BUFFER_BYTES", "MooncakeStoreConnectorConfig", "SEGMENTS_DIR_NAME", "StoreRole", @@ -66,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` @@ -323,6 +328,19 @@ 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 @@ -347,6 +365,10 @@ 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; " @@ -392,6 +414,13 @@ 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 ), diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/master.py index 4b76192ae93e..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,9 @@ 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. diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py index 3736bf1b753c..5ad3121c0c1c 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/metadata.py @@ -54,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 701fb2109c11..744be81b630f 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/scheduler.py @@ -32,6 +32,12 @@ 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 @@ -114,6 +120,8 @@ class _RequestState: "saved_upto", "load_first_block", "load_blocks", + "load_async", + "async_loaded_upto", "emitted_saves", ) @@ -130,6 +138,13 @@ def __init__(self, chain: BlockHashChain): #: 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 @@ -143,6 +158,14 @@ class MooncakeStoreConnectorScheduler(KvCacheConnectorScheduler): #: 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) @@ -174,11 +197,18 @@ def __init__(self, llm_args: TorchLlmArgs): # 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 scheduler adapter ready (role={self._config.role.value}, " f"tokens_per_block={self._tokens_per_block}, " - f"prompt_lookahead={self._prompt_lookahead})" + f"prompt_lookahead={self._prompt_lookahead}, async_load={self._async_load})" ) def wait_for_initialization(self): @@ -202,7 +232,8 @@ def get_num_new_matched_tokens( 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 @@ -218,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 @@ -245,38 +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]): - """Trim the outstanding offer to what the runtime honoured. + """Settle an offer against a flat page list, as one 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. - Page indices still come from the scheduler output: 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`. `KVCacheManagerV2` reports - through `update_state_after_alloc_by_layer_group`, whose base form - folds the single layer group this connector accepts into this call. + 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: Unused; see above. + 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: @@ -297,8 +405,15 @@ def build_connector_meta(self, scheduler_output: SchedulerOutput) -> MooncakeSto # 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 @@ -339,10 +454,17 @@ def request_reset(self, request: LlmRequest) -> None: `pages` and `saved_upto` describe the allocation that was just freed, 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. + 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 ---- @@ -377,8 +499,8 @@ 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 once attention data - parallelism or asynchronous loads are supported, a fleet vote, a + 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. diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/settings.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/settings.py index ba877acf296e..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,13 @@ 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 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 4a4c8af35456..9a403c2a0295 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/validation.py @@ -198,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: @@ -232,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 a208231d5bbc..f8c234150a3e 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/worker.py @@ -26,13 +26,20 @@ 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, 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. +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 @@ -74,7 +81,7 @@ 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 ( @@ -89,7 +96,6 @@ 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, @@ -116,11 +122,18 @@ _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. @@ -250,7 +263,8 @@ def _replicate_config(soft_pin_ttl_ms: Optional[int], group_ids: Optional[Sequen #: 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. +#: 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") @@ -279,8 +293,9 @@ 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. Every method takes the lock, so a load that runs off the executor - thread can record into the same counters the executor thread logs. + 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__ = ( @@ -289,8 +304,12 @@ class _LoadStats: "pages", "batches", "loaded_bytes", + "failed_pages", "xfer_sum", "xfer_max", + "waits", + "wait_sum", + "wait_max", "phase_wall", "phase_cpu", ) @@ -304,8 +323,12 @@ def _reset(self) -> None: 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) @@ -316,6 +339,9 @@ def record( 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. @@ -326,34 +352,62 @@ def record( 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) -> Optional[str]: - """The interval's summary, starting a new interval; None when nothing loaded.""" + 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: + if not (self.loads or queued or active): return None - n = self.loads + 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={n} pages={self.pages} batches={self.batches} " + 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() @@ -520,34 +574,46 @@ def __init__(self, llm_args: TorchLlmArgs): self._load_stats: Optional[_LoadStats] = _LoadStats() if counts else None self._load_stats_last = time.monotonic() - # 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, - ) - + # 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() ) @@ -566,6 +632,8 @@ 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 @@ -581,6 +649,34 @@ def __init__(self, llm_args: TorchLlmArgs): # 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() @@ -594,7 +690,8 @@ def __init__(self, llm_args: TorchLlmArgs): 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', 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 @@ -632,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") @@ -688,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. @@ -757,41 +894,104 @@ 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 _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( @@ -928,7 +1128,13 @@ def _maybe_log_load_stats(self, force: bool = False) -> Optional[str]: now = time.monotonic() if not (force or now - self._load_stats_last >= self._stats_interval): return None - summary = stats.take_summary() + 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 = ( @@ -948,29 +1154,44 @@ def start_load_kv(self, stream: torch.cuda.Stream): 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. + 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, or - a zero-copy load has no allocation fence to wait on. + 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() # 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(metadata.loads) + 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 metadata.loads for _ in entry.pages] + 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: @@ -987,7 +1208,7 @@ def start_load_kv(self, stream: torch.cuda.Stream): "KVCacheManagerV2 does." ) with _timed_phase(phases, "fence"): - fence({entry.request_id for entry in metadata.loads}) + 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( @@ -1008,6 +1229,7 @@ def start_load_kv(self, stream: torch.cuda.Stream): 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 @@ -1040,9 +1262,10 @@ def _load_pages( ) -> Tuple[Set[int], int, Optional[str], int]: """Fetch pages batch by batch, scattering staged ones onto the device. - 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. + 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. @@ -1131,6 +1354,262 @@ def _load_pages( _sync_stream(handle) return failed_requests, num_failed_pages, first_failure, loaded_bytes + # ---- 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`. @@ -1178,7 +1657,7 @@ def drop_bound_saves(self, request_ids: Iterable[int]) -> int: 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. @@ -1230,19 +1709,29 @@ def take_submitted_saves(self) -> Set[int]: 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() @@ -1256,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: @@ -1418,32 +1937,45 @@ def _resolve( return keys, addresses, sizes, pages def shutdown(self) -> None: - """Stop the save thread, then release what it was reading. Idempotent.""" + """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) - thread = self._save_thread - if thread is not None: + # 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: @@ -1453,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/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 36b707aea86e..011a14154f94 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 @@ -2654,7 +2655,36 @@ 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, @@ -2720,19 +2750,23 @@ class MooncakeStoreConfig(StrictBaseModel): "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", mode="after") + @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 diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index d8cc625e9263..2239efc4797f 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -4165,6 +4165,27 @@ "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.", @@ -4298,13 +4319,29 @@ }, "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": [ { diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index e5be00381f5c..9d15424d20ca 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -5729,6 +5729,27 @@ "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.", @@ -5862,13 +5883,29 @@ }, "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": [ { 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 3d70350f3c68..33108ccdca78 100644 --- a/tests/integration/defs/llmapi/test_llm_api_connector_mooncake.py +++ b/tests/integration/defs/llmapi/test_llm_api_connector_mooncake.py @@ -153,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. @@ -260,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") @@ -308,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): 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/unittest/_torch/executor/test_mooncake_store_common.py b/tests/unittest/_torch/executor/test_mooncake_store_common.py index 5285770a29b0..5a32a07bd489 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_common.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_common.py @@ -32,9 +32,11 @@ 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 ( @@ -50,6 +52,10 @@ 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 @@ -691,6 +697,113 @@ def test_config_has_no_environment_override_for_eviction_settings(store_config, 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) @@ -847,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 191671a32d18..f4450555081f 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_connector.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_connector.py @@ -1900,16 +1900,15 @@ def test_worker_accepts_shared_one_model_draft_and_capacity_role( def make_staged_worker(fake_store, store_config, *, layout, budget=None): """A worker configured to pass pages through pinned host slots. - `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. + `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. """ - 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)) + 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 @@ -2190,6 +2189,763 @@ def test_staging_narrows_the_batch_to_the_budget(store_config, fake_store, stage 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), + ) + worker.shutdown() + + assert worker._unreleased_staging == [save_pool] + assert worker._save_staging is None + assert fake_store.closed + + # ---- scheduler ---- @@ -2723,12 +3479,16 @@ def test_scheduler_sync_load_stops_at_the_honoured_prefix( assert saved_pages(metadata) == saved -def test_scheduler_rejects_an_honour_beyond_the_offer(store_config): +@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, False) + 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): @@ -2897,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 3e1ff619f82e..54be4358e974 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_real_pool.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_real_pool.py @@ -280,6 +280,7 @@ def open_worker( 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. @@ -295,6 +296,7 @@ def open_worker( 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( @@ -395,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. diff --git a/tests/unittest/_torch/executor/test_mooncake_store_settings.py b/tests/unittest/_torch/executor/test_mooncake_store_settings.py index 7ad6d42a73e5..35249a9920b4 100644 --- a/tests/unittest/_torch/executor/test_mooncake_store_settings.py +++ b/tests/unittest/_torch/executor/test_mooncake_store_settings.py @@ -92,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") @@ -104,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. @@ -205,6 +211,29 @@ def test_a_stats_interval_must_be_a_positive_number(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( From 20391192edc1a9a378fb0c772dde292fdb75122d Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Tue, 27 Oct 2026 15:39:55 -0700 Subject: [PATCH 18/20] [None][fix] Disagg proxy: never replay a generation leg the generation server may hold; fail a stream before its first chunk with an HTTP error A generation leg takes its prompt KV cache from the context server's one-shot send session named by disagg_request_id. When that transfer fails, the generation server raises in its stream generator after the 200 headers, so the proxy saw a ClientPayloadError with no lines yielded and re-posted the leg under a freshly minted disagg_request_id. No context server holds KV cache for that id, so the replay waited forever and the client request never completed. The transient-TCP budget made this wider: a server disconnect or connection reset can arrive after the generation server already read the request. The client now uses the generation-leg contract of open PR #20938, with its text kept verbatim so a later merge of that PR is a no-op. A generation leg keeps its id on retry, and it is re-sent only when it cannot have reached the server: the connection was refused (ClientConnectorError) or reset while the request was being written (ConnectionResetError). Any failure after the response headers, and any other error, fails the request. Context requests are unchanged. The error-status check in _post_with_retry now runs before the stream branch as well, so a streaming leg answered with a 4xx/5xx JSON body raises a ClientResponseError carrying the worker's status instead of tripping the "Response is not streaming" assertion. On the server, _wrap_entry_point pulls a relayed stream's first chunk before it sends the headers. A chat or completion stream that fails before its first chunk now reaches _handle_exception and becomes an HTTP error (and an Anthropic error envelope on /v1/messages) instead of a 200 whose body is cut off. The returned relay is already started, so closing it before it is iterated still closes the upstream stream and releases the router reservation. Responses streams keep their 200 plus an `error` event, because guard_responses_stream stays inside the first-chunk pull. total_responses now counts a stream only once its first chunk arrived. Signed-off-by: Simeng Liu --- docs/source/features/disagg-serving.md | 2 + tensorrt_llm/serve/openai_client.py | 63 +++- tensorrt_llm/serve/openai_disagg_server.py | 49 ++- .../test_disagg_openai_client.py | 327 +++++++++++++++++- .../test_openai_disagg_server.py | 234 ++++++++++++- 5 files changed, 651 insertions(+), 24 deletions(-) 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/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/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_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" From c861c0dee5c210cd77b1d5582f393d5c58bbbf5c Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Tue, 27 Oct 2026 16:01:48 -0700 Subject: [PATCH 19/20] [None][fix] Disagg: cancel and fail timed-out KV receives on an idle generation server Without in-flight cancellation, the generation-side KV-transfer timeout only flags the request in check_transfer_timeouts. The cancel and the error response happen in the response pass, which runs only after a forward step. A generation server with nothing else scheduled therefore kept a request whose KV never arrived until the client gave up, even after logging "Observed timeout on generation request". DisaggTransferCoordinator.poll_gen_transfers, which runs at every loop head, now cancels and fails the receives that check_transfer_timeouts flagged whenever in-flight cancellation is inactive. This covers any transceiver without that capability, including KvCacheTransceiverV2 and TRTLLM_DISAGG_ENABLE_INFLIGHT_CANCEL=1 on an unsupporting transceiver. 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, while the C++ transceiver cannot cancel a receive already in flight and would only repeat its "Cannot cancel request" warning at the idle loop rate. Requests the user cancelled are left to the cancel path. No collective is added. The in-flight cancellation path is unchanged. fail_timed_out now moves a request still in DISAGG_GENERATION_TRANS_IN_PROGRESS to DISAGG_TRANS_ERROR before failing it. cancel_request does not change the state, so PyExecutor._handle_errors parked the failure for a transfer that no longer existed, and _sweep_transfer_held_failures then dropped it because the request had already left active_requests: no error response was sent and the KV was never freed. This also fixes the same loss on the existing response-pass timeout path. Signed-off-by: Simeng Liu --- .../orchestration/coordinator.py | 80 ++++++++- .../_torch/disaggregation/conftest.py | 6 + .../test_disagg_coordinator_errors.py | 93 +++++++++++ .../test_disagg_coordinator_progress.py | 155 ++++++++++++++++++ .../test_disagg_coordinator_transfers.py | 44 +++++ .../_torch/executor/test_py_executor.py | 1 + .../executor/test_forward_error_survival.py | 50 ++++++ 7 files changed, 426 insertions(+), 3 deletions(-) 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/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_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/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/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) From 2cebc8360860a6e79cccea70f542a09120a14cd6 Mon Sep 17 00:00:00 2001 From: Simeng Liu Date: Tue, 27 Oct 2026 21:10:52 -0700 Subject: [PATCH 20/20] [None][feat] Attention DP: work-aware router that balances ranks by estimated uncached prefill The default attention-DP router balances ranks by prompt length. When a KV cache connector restores prefixes from a pool shared by all attention-DP ranks, most of a later turn's prompt is restored rather than computed, so equal-length prompts can bring very different prefill work and one rank can end up working through a large prefill while its partners pad in lockstep. WorkAwareADPRouter estimates a new request's prefill as its prompt growth since the previous request with the same conversation_id (the whole prompt for a first turn, a request without an id, or a shrinking prompt; at least one token). 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. Every request, including explicitly placed ones, is estimated once in routing order and charged to its rank; n>1 child ids inherit the parent's estimate, and the estimate is published as py_uncached_prefill_tokens through the existing hook while the V2 scheduler's long-prefill demotion consumes it. The conversation and per-request tables are LRU-bounded and change only inside route_requests, which every rank runs on identical inputs; create_rank_state only reads them, so placements stay consistent across ranks. The #19369 connector fields in the gathered rank state are untouched. The router is selected by the new prototype field attention_dp_config.enable_work_aware_routing, replacing the fork's TRTLLM_ADP_WORK_AWARE_ROUTING environment variable so the choice is validated and identical on every rank. ADPRouter.create gains a kv_connector_manager argument, which PyExecutor passes through, and builds the router only with a connector that moves KV; otherwise it falls back to the default router with a warning, staying silent for the KV-cache estimation executor, which is built without its connector. The field cannot be combined with enable_kv_cache_aware_routing or kv_cache_routing_conversation_affinity. Signed-off-by: Simeng Liu --- docs/source/features/kv-cache-connector.md | 15 + tensorrt_llm/_torch/pyexecutor/py_executor.py | 1 + .../_torch/pyexecutor/scheduler/adp_router.py | 203 ++++++++++- tensorrt_llm/llmapi/llm_args.py | 24 ++ .../schemas/trtllm-serve-config.schema.json | 7 + .../trtllm-serve-disagg-config.schema.json | 7 + .../usage/llm_args_golden_manifest.json | 5 + .../disaggregation/test_benchmark_disagg.py | 10 +- .../_torch/executor/test_adp_router.py | 331 +++++++++++++++++- .../executor/test_remote_tail_adp_phase.py | 14 +- 10 files changed, 611 insertions(+), 6 deletions(-) diff --git a/docs/source/features/kv-cache-connector.md b/docs/source/features/kv-cache-connector.md index 13420dcba028..b011bf9a17dd 100644 --- a/docs/source/features/kv-cache-connector.md +++ b/docs/source/features/kv-cache-connector.md @@ -693,6 +693,21 @@ 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. +**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 + +```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 The file `examples/llm-api/llm_kv_cache_connector.py` provides a reference implementation of a **Persistent KV Cache**. diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 96854daedd72..a7c53d85b24d 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -916,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. 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/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 011a14154f94..61dd318f63ba 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -2114,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': @@ -2126,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 diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index 2239efc4797f..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.", diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index 9d15424d20ca..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.", diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index ff7cb7b26268..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", 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/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_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"]