[TRTLLM-11875][feat] Support fine-grained context chunk management

cpp/include/tensorrt_llm/batch_manager/llmRequest.h

@@ -1234,6 +1234,18 @@ class GenericLlmRequest
         mEstimatedReusableTokens = estimatedReusableTokens;
     }
 
+    //! Get the absolute context positions at which recurrent-state snapshots are expected.
+    [[nodiscard]] std::vector<SizeType32> const& getExpectedSnapshotPoints() const noexcept
+    {
+        return mExpectedSnapshotPoints;
+    }
+
+    //! Set the absolute context positions at which recurrent-state snapshots are expected.
+    void setExpectedSnapshotPoints(std::vector<SizeType32> expectedSnapshotPoints)
+    {
+        mExpectedSnapshotPoints = std::move(expectedSnapshotPoints);
+    }
+
     void setDraftTokens(std::shared_ptr<VecTokens> const& draftTokens)
     {
         mDraftTokens = draftTokens;
@@ -2097,6 +2109,9 @@ class GenericLlmRequest
     // the authoritative mPrepopulatedPromptLen and advances context position.
     mutable SizeType32 mEstimatedReusableTokens{0};
 
+    // Absolute context positions at which recurrent-state snapshots are expected.
+    std::vector<SizeType32> mExpectedSnapshotPoints;
+
     SizeType32 mMaxSentTokenLen;
 
     std::optional<TensorPtr> mEmbeddingBias{std::nullopt};

cpp/tensorrt_llm/batch_manager/microBatchScheduler.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
@@ -230,8 +230,9 @@ void MicroBatchScheduler::setCtxRequestsChunkSize<MicroBatchScheduler::ContextCh
 
 // Assigns chunk sizes to context requests under the kFORCE_CHUNK policy.
 //
-// Every request is assigned exactly min(contextRemainingLength, chunkUnitSize) tokens.
-// Requests whose chunk would push the running total past ctxTokensCapacity are zeroed.
+// Requests with expected snapshot points advance to their next absolute snapshot position.
+// Otherwise, every request consumes its full remaining context.
+// Capacity and context-length truncation are rounded down to a chunk-unit boundary.
 //
 // This policy is designed for linear attention state caching, so reusable KV-cache tokens are NOT
 // calculated because it's not supported yet.
@@ -248,16 +249,36 @@ void MicroBatchScheduler::setCtxRequestsChunkSize<MicroBatchScheduler::ContextCh
     SizeType32 totalTokens{0};
     for (auto& llmReq : contextsToBeChunked)
     {
-        SizeType32 const chunkSize = std::min(llmReq->getContextRemainingLength(), chunkUnitSize);
-        if (ctxTokensCapacity && totalTokens + chunkSize > ctxTokensCapacity.value())
+        SizeType32 chunkSize = llmReq->getContextRemainingLength();
+        auto const& expectedSnapshotPoints = llmReq->getExpectedSnapshotPoints();
+        if (!expectedSnapshotPoints.empty())
+        {
+            auto const currentPosition = llmReq->getContextCurrentPosition();
+            std::optional<SizeType32> nextSnapshotPoint;
+            for (auto const point : expectedSnapshotPoints)
+            {
+                if (point > currentPosition && (!nextSnapshotPoint || point < nextSnapshotPoint.value()))
+                {
+                    nextSnapshotPoint = point;
+                }
+            }
+            chunkSize = nextSnapshotPoint
+                ? std::max<SizeType32>(0, std::min(nextSnapshotPoint.value(), llmReq->getPromptLen()) - currentPosition)
+                : llmReq->getContextRemainingLength();
+        }
+
+        if (maxContextLength && chunkSize > maxContextLength.value())
         {
-            llmReq->setContextChunkSize(0);
+            chunkSize = maxContextLength.value() / chunkUnitSize * chunkUnitSize;
         }
-        else
+        if (ctxTokensCapacity && totalTokens + chunkSize > ctxTokensCapacity.value())
         {
-            llmReq->setContextChunkSize(chunkSize);
-            totalTokens += llmReq->getContextChunkSize();
+            auto const remainingCapacity = std::max<SizeType32>(0, ctxTokensCapacity.value() - totalTokens);
+            chunkSize = std::min(chunkSize, remainingCapacity) / chunkUnitSize * chunkUnitSize;
         }
+
+        llmReq->setContextChunkSize(chunkSize);
+        totalTokens += llmReq->getContextChunkSize();
     }
 }
 
@@ -267,8 +288,9 @@ void MicroBatchScheduler::setCtxRequestsChunkSize<MicroBatchScheduler::ContextCh
 //   kEQUAL_PROGRESS        — all requests advance together one chunkUnitSize at a time.
 //   kFIRST_COME_FIRST_SERVED — requests are served greedily in order until the budget
 //                              is exhausted.
-//   kFORCE_CHUNK           — every request gets exactly min(remaining, chunkUnitSize)
-//                              tokens; budget is charged at face value (no reuse discount).
+//   kFORCE_CHUNK           — requests advance to the next expected snapshot point, or consume
+//                              the remaining context when none are configured; budget is charged
+//                              at face value (no reuse discount).
 //
 // EQUAL_PROGRESS and FIRST_COME_FIRST_SERVED are compute-aware: tokens covered by the
 // reusable KV-cache prefix are not charged against ctxTokensCapacity.
@@ -436,7 +458,7 @@ std::tuple<RequestVector, RequestVector> MicroBatchScheduler::operator()(Request
         allContextRequestsFit = false;
     }
 
-    // For FORCE_CHUNK policy, always re-chunk regardless of whether all contexts fit.
+    // FORCE_CHUNK must always run boundary selection even when all contexts fit.
     if (mCtxChunkConfig && mCtxChunkConfig.value().chunkingPolicy == ContextChunkingPolicy::kFORCE_CHUNK)
     {
         allContextRequestsFit = false;

cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp

@@ -174,6 +174,8 @@ void initBindings(nb::module_& m)
             nb::arg("kv_tokens_per_block"))
         .def_prop_rw(
             "estimated_reusable_tokens", &GenLlmReq::getEstimatedReusableTokens, &GenLlmReq::setEstimatedReusableTokens)
+        .def_prop_rw(
+            "expect_snapshot_points", &GenLlmReq::getExpectedSnapshotPoints, &GenLlmReq::setExpectedSnapshotPoints)
         .def_prop_rw("guided_decoding_params", &GenLlmReq::getGuidedDecodingParams, &GenLlmReq::setGuidedDecodingParams)
         .def_prop_rw("context_phase_params", &GenLlmReq::getContextPhaseParams, &GenLlmReq::setContextPhaseParams)
         .def_prop_ro("is_context_only_request", &GenLlmReq::isContextOnlyRequest)

cpp/tests/unit_tests/batch_manager/microBatchSchedulerTest.cpp

@@ -2054,19 +2054,19 @@ class ForceChunkTest : public MicroBatchSchedulerTest
     }
 };
 
-TEST_F(ForceChunkTest, Basic)
+TEST_F(ForceChunkTest, NoSnapshotPointsUsesRemainingContext)
 {
-    // A single request with prompt_len > chunk_unit_size is chunked to unit_size.
+    // A request without snapshot points is not split at chunkUnitSize.
     auto reqs = initRequests({30});
     MicroBatchScheduler::setCtxRequestsChunkSize(reqs, Policy::kFORCE_CHUNK, /*ctxTokensCapacity=*/std::nullopt,
         /*chunkUnitSize=*/10, /*maxContextLength=*/std::nullopt);
 
-    EXPECT_EQ(reqs[0]->getContextChunkSize(), 10);
+    EXPECT_EQ(reqs[0]->getContextChunkSize(), 30);
 }
 
 TEST_F(ForceChunkTest, PromptSmallerThanUnit)
 {
-    // When prompt_len < chunk_unit_size, chunk_size = prompt_len (min).
+    // Without snapshot points, a short prompt is consumed in full.
     auto reqs = initRequests({8});
     MicroBatchScheduler::setCtxRequestsChunkSize(reqs, Policy::kFORCE_CHUNK, std::nullopt, 20, std::nullopt);
 
@@ -2075,7 +2075,7 @@ TEST_F(ForceChunkTest, PromptSmallerThanUnit)
 
 TEST_F(ForceChunkTest, ExactUnitSize)
 {
-    // When prompt_len == chunk_unit_size, chunk_size = prompt_len.
+    // Without snapshot points, an exact-unit prompt is consumed in full.
     auto reqs = initRequests({10});
     MicroBatchScheduler::setCtxRequestsChunkSize(reqs, Policy::kFORCE_CHUNK, std::nullopt, 10, std::nullopt);
 
@@ -2084,43 +2084,75 @@ TEST_F(ForceChunkTest, ExactUnitSize)
 
 TEST_F(ForceChunkTest, MultipleRequests)
 {
-    // Each request independently gets min(remaining, unit_size).
+    // Requests without snapshot points independently consume their remaining contexts.
     auto reqs = initRequests({25, 15, 5});
     MicroBatchScheduler::setCtxRequestsChunkSize(reqs, Policy::kFORCE_CHUNK, std::nullopt, 10, std::nullopt);
 
-    EXPECT_EQ(reqs[0]->getContextChunkSize(), 10);
-    EXPECT_EQ(reqs[1]->getContextChunkSize(), 10);
-    EXPECT_EQ(reqs[2]->getContextChunkSize(), 5); // min(5, 10) = 5
+    EXPECT_EQ(reqs[0]->getContextChunkSize(), 25);
+    EXPECT_EQ(reqs[1]->getContextChunkSize(), 15);
+    EXPECT_EQ(reqs[2]->getContextChunkSize(), 5);
 }
 
 TEST_F(ForceChunkTest, CapacityLimits)
 {
-    // When capacity is limited, later requests get chunk_size=0.
+    // Budget truncation is unit-aligned; later requests with less than one
+    // unit available are delayed.
     auto reqs = initRequests({30, 30});
     MicroBatchScheduler::setCtxRequestsChunkSize(
         reqs, Policy::kFORCE_CHUNK, /*ctxTokensCapacity=*/15, /*chunkUnitSize=*/10, std::nullopt);
 
-    // req0 gets 10, req1 would push total to 20 > 15 → 0
+    // req0 is budget-truncated to 10; only 5 remain, so req1 gets 0.
     EXPECT_EQ(reqs[0]->getContextChunkSize(), 10);
     EXPECT_EQ(reqs[1]->getContextChunkSize(), 0);
 }
 
 TEST_F(ForceChunkTest, CapacityExactFit)
 {
-    // When capacity exactly accommodates all chunks.
+    // Capacity exactly accommodates both requested snapshot chunks.
     auto reqs = initRequests({30, 30});
+    for (auto const& req : reqs)
+    {
+        req->setExpectedSnapshotPoints({10});
+    }
     MicroBatchScheduler::setCtxRequestsChunkSize(
         reqs, Policy::kFORCE_CHUNK, /*ctxTokensCapacity=*/20, /*chunkUnitSize=*/10, std::nullopt);
 
     EXPECT_EQ(reqs[0]->getContextChunkSize(), 10);
     EXPECT_EQ(reqs[1]->getContextChunkSize(), 10);
 }
 
+TEST_F(ForceChunkTest, ExpectedChunkingPoints)
+{
+    // Expected snapshot points are absolute context positions.
+    auto reqs = initRequests({30});
+    reqs[0]->setExpectedSnapshotPoints({12, 25});
+
+    chunkIteration(reqs, 10);
+    expectPositions(reqs, {12}, "iter 1");
+
+    chunkIteration(reqs, 10);
+    expectPositions(reqs, {25}, "iter 2");
+
+    chunkIteration(reqs, 10);
+    expectPositions(reqs, {30}, "iter 3");
+}
+
+TEST_F(ForceChunkTest, CapacityRoundsExpectedChunkDownToUnit)
+{
+    auto reqs = initRequests({50});
+    reqs[0]->setExpectedSnapshotPoints({30});
+
+    MicroBatchScheduler::setCtxRequestsChunkSize(
+        reqs, Policy::kFORCE_CHUNK, /*ctxTokensCapacity=*/25, /*chunkUnitSize=*/10, std::nullopt);
+
+    EXPECT_EQ(reqs[0]->getContextChunkSize(), 20);
+}
+
 TEST_F(ForceChunkTest, MultiIteration)
 {
-    // A request with prompt_len=25 and chunk_unit_size=10 processes in 3 iterations:
-    // chunk 1: 10, chunk 2: 10, chunk 3: 5.
+    // Snapshot points at 10 and 20 split a 25-token prompt into three iterations.
     auto reqs = initRequests({25});
+    reqs[0]->setExpectedSnapshotPoints({10, 20});
 
     // Iteration 1
     chunkIteration(reqs, 10);
@@ -2138,8 +2170,10 @@ TEST_F(ForceChunkTest, MultiIteration)
 TEST_F(ForceChunkTest, MultiRequestMultiIteration)
 {
     // Two requests with different lengths processed over multiple iterations.
-    // prompt_len={25, 12}, chunk_unit_size=10.
+    // Their expected snapshot points determine each boundary.
     auto reqs = initRequests({25, 12});
+    reqs[0]->setExpectedSnapshotPoints({10, 20});
+    reqs[1]->setExpectedSnapshotPoints({10});
 
     // Iteration 1: both get 10
     chunkIteration(reqs, 10);
@@ -2157,8 +2191,12 @@ TEST_F(ForceChunkTest, MultiRequestMultiIteration)
 TEST_F(ForceChunkTest, CapacityAcrossIterations)
 {
     // With limited capacity, some requests may be delayed to later iterations.
-    // prompt_len={25, 25}, chunk_unit_size=10, capacity=15.
+    // Both requests have snapshot points at 10 and 20; capacity is 15.
     auto reqs = initRequests({25, 25});
+    for (auto const& req : reqs)
+    {
+        req->setExpectedSnapshotPoints({10, 20});
+    }
 
     // Iteration 1: req0=10, req1=0 (10+10=20 > 15)
     chunkIteration(reqs, 10, /*ctxTokensCapacity=*/15);
@@ -2181,10 +2219,10 @@ TEST_F(ForceChunkTest, CapacityAcrossIterations)
     expectPositions(reqs, {25, 25}, "iter 5");
 }
 
-TEST_F(ForceChunkTest, FullSchedulerPath)
+TEST_F(ForceChunkTest, FullSchedulerWithoutSnapshotPoints)
 {
-    // Test via MicroBatchScheduler::operator() — FORCE_CHUNK always re-chunks
-    // even when all contexts fit within the token budget.
+    // Test via MicroBatchScheduler::operator(): without snapshot points, a
+    // context that fits is not split at chunkUnitSize.
     batch_scheduler::ContextChunkingConfig chunkConfig;
     chunkConfig.chunkingPolicy = Policy::kFORCE_CHUNK;
     chunkConfig.chunkUnitSize = 10;
@@ -2201,9 +2239,33 @@ TEST_F(ForceChunkTest, FullSchedulerPath)
     auto const [contextRequests, genRequests]
         = (*scheduler)(activeRequests, inflightReqIds, maxBatchSize, maxNumTokens);
 
-    // Despite budget=100 >> prompt=30, FORCE_CHUNK limits chunk to unit_size=10.
     ASSERT_EQ(contextRequests.size(), 1);
-    EXPECT_EQ(contextRequests[0]->getContextChunkSize(), 10);
+    EXPECT_EQ(contextRequests[0]->getContextChunkSize(), 30);
+    EXPECT_EQ(genRequests.size(), 0);
+}
+
+TEST_F(ForceChunkTest, FullSchedulerUsesExpectedChunkingPoints)
+{
+    batch_scheduler::ContextChunkingConfig chunkConfig;
+    chunkConfig.chunkingPolicy = Policy::kFORCE_CHUNK;
+    chunkConfig.chunkUnitSize = 10;
+
+    auto scheduler = std::make_shared<MicroBatchScheduler>(chunkConfig);
+
+    constexpr SizeType32 maxBatchSize = 4;
+    constexpr SizeType32 maxNumTokens = 100;
+
+    RequestVector activeRequests;
+    auto request = createRequest(/*promptLen=*/30, /*maxNewTokens=*/1, /*reqId=*/0);
+    request->setExpectedSnapshotPoints({12, 25});
+    activeRequests.push_back(request);
+
+    ReqIdsSet inflightReqIds;
+    auto const [contextRequests, genRequests]
+        = (*scheduler)(activeRequests, inflightReqIds, maxBatchSize, maxNumTokens);
+
+    ASSERT_EQ(contextRequests.size(), 1);
+    EXPECT_EQ(contextRequests[0]->getContextChunkSize(), 12);
     EXPECT_EQ(genRequests.size(), 0);
 }
 
@@ -2235,8 +2297,8 @@ TEST_F(ForceChunkTest, FullSchedulerMultipleRequests)
     {
         chunks[req->mRequestId] = req->getContextChunkSize();
     }
-    EXPECT_EQ(chunks[0], 10);
-    EXPECT_EQ(chunks[1], 10);
+    EXPECT_EQ(chunks[0], 25);
+    EXPECT_EQ(chunks[1], 15);
     EXPECT_EQ(chunks[2], 5);
 }
 
@@ -2268,6 +2330,7 @@ TEST_F(ForceChunkTest, FullSchedulerWithGeneration)
 
     EXPECT_EQ(genRequests.size(), 1);
     ASSERT_EQ(contextRequests.size(), 1);
-    // Budget remaining = 15 - 1 (gen) = 14; chunk = min(30, 10) = 10
+    // Budget remaining is 14, so the context is rounded down to one
+    // 10-token chunk-unit boundary.
     EXPECT_EQ(contextRequests[0]->getContextChunkSize(), 10);
 }

tensorrt_llm/_torch/disaggregation/native/rank_info.py

@@ -59,6 +59,7 @@ def from_kv_cache_manager(
         m = kv_cache_manager.mapping
         kvm = kv_cache_manager
         enable_attention_dp = m.enable_attention_dp
+        kv_heads_per_rank = next((h for h in kvm.num_kv_heads_per_layer if h > 0), 0)
         return cls(
             instance_name=instance_name,
             instance_rank=m.rank,
@@ -77,7 +78,7 @@ def from_kv_cache_manager(
             self_endpoint="",
             transfer_engine_info=bytes(),
             attention=AttentionInfo(
-                kv_heads_per_rank=kvm.num_kv_heads_per_layer[0],
+                kv_heads_per_rank=kv_heads_per_rank,
                 tokens_per_block=kvm.tokens_per_block,
                 dims_per_head=kvm.head_dim,
                 element_bytes=get_size_in_bytes(1, kvm.dtype),

tensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.py

@@ -15,7 +15,8 @@
 
 from .llm_request import LlmRequest
 from .mamba_cache_manager import (BaseMambaCacheManager,
-                                  CppMambaHybridCacheManager)
+                                  CppMambaHybridCacheManager,
+                                  MixedMambaHybridCacheManager)
 from .resource_manager import KVCacheManager
 
 CacheTransceiverCpp = tensorrt_llm.bindings.internal.batch_manager.CacheTransceiver
@@ -129,6 +130,12 @@ def create_kv_cache_transceiver(
     if cache_transceiver_config.transceiver_runtime == "auto":
         cache_transceiver_config.transceiver_runtime = None
 
+    if (cache_transceiver_config.transceiver_runtime != "PYTHON"
+            and isinstance(mamba_cache_manager, MixedMambaHybridCacheManager)):
+        raise ValueError(
+            "MixedMambaHybridCacheManager requires the Python transceiver "
+            "runtime in disaggregated serving.")
+
     _validate_disagg_inflight_cancel_config(cache_transceiver_config)
 
     if cache_transceiver_config.backend == "DEFAULT":