[None][feat] Close the Mooncake store feature gap: async loads, attention DP, speculative prefill and KV connector support in KVCacheManagerV2

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp

@@ -849,8 +849,29 @@ std::vector<BlockRadixTree::MatchResult> 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<BlockRadixTree::MatchResult>& 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::MatchResult> BlockRadixTree::pruneMatch(
-    std::vector<MatchResult> matched, std::optional<LifeCycleId> ssmLcId) const
+    std::vector<MatchResult> matched, std::optional<LifeCycleId> 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::MatchResult> 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::MatchResult> 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::MatchResult> 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<BlockRadixTree::MatchResult>& 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<int> 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<int>(tokens.size());

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<MatchResult> pruneMatch(std::vector<MatchResult> matched, std::optional<LifeCycleId> ssmLcId) const;
+    std::vector<MatchResult> pruneMatch(
+        std::vector<MatchResult> matched, std::optional<LifeCycleId> ssmLcId, int alignment = 1) const;
 
     // Erase any pending empty root blocks from mRoots.
     // Const-qualified: deferred cleanup is not a logical mutation.

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();

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<BatchDesc> constraints;   // batches that must always be supportable
     std::optional<BatchDesc> typicalStep; // typical step for initial ratio computation

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

cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp

@@ -1753,7 +1753,7 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m)
                 std::optional<std::vector<float>> initialPoolRatio,
                 std::optional<kv::SwaScratchReuseConfig> 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)

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();

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