perf(k3): improve recurrent-state prefix reuse (#1173)

Signed-off-by: chenht2022 <chenht2022@gmail.com>

docs/design/cache-concepts.md

@@ -146,13 +146,16 @@ Consumers outside the cache layer treat the ids as opaque.
 
 Logical width does not imply dense physical residency. Full-history KV and
 retained sliding-window rows materialize every block their kernels read, but a
-full-history snapshot-state prefill needs only its input checkpoint, final
-output checkpoint. The next decode admission allocates its destination after
-prefill scheduling and rolls the expired input page forward, including under
-overlap scheduling. Its table therefore keeps absolute slot positions while
-representing skipped intermediate checkpoints as null holes (`0`). State
-consumers may gather only the declared input/output slots; compacting the row or
-publishing an unwritten intermediate checkpoint would break position identity.
+full-history snapshot-state prefill normally needs only its input and final
+output checkpoints. With prefix caching and an off-page final tail, the aligned
+body materializes its endpoint and atomically reserves the tail storage; the
+tail consumes that reservation instead of materializing another sparse input
+checkpoint. The next decode admission allocates its destination after prefill
+scheduling and rolls the expired input page forward, including under overlap
+scheduling. The table keeps absolute slot positions while representing other
+skipped intermediate checkpoints as null holes (`0`). State consumers may
+gather only the declared input/output slots; compacting the row or publishing an
+unwritten intermediate checkpoint would break position identity.
 
 Speculative KDA verification stores no per-position recurrent states: it
 captures each window's raw projections in a compact payload and commits by

tokenspeed-scheduler/csrc/fsm/forward_states.h

@@ -44,6 +44,8 @@ struct CacheProgress {
     std::uint64_t access_epoch{0};
     // Pending closed-prefix boundary; zero once published or when absent.
     std::int32_t promotion_boundary_tokens{0};
+    // Whether cache storage for the final state-checkpoint tail was reserved.
+    bool state_checkpoint_tail_reserved{false};
 };
 
 inline std::vector<std::int32_t> ComputeShiftedInputIds(const TokenContainer* token_container,

tokenspeed-scheduler/csrc/scheduler/operations/cache.cpp

@@ -44,6 +44,23 @@ std::int32_t AlignPrefillChunk(std::int32_t first_pos, std::int32_t unscheduled,
     return chunk_size - chunk_size % prefix_granularity;
 }
 
+std::optional<std::int32_t> FinalAlignedTailTokens(std::int32_t first_pos, std::int32_t unscheduled,
+                                                   std::int32_t token_budget, std::int32_t prefix_granularity,
+                                                   std::int32_t promotion_boundary_tokens) {
+    _assert(first_pos >= 0 && unscheduled >= 0 && token_budget >= 0, "prefill positions must be non-negative");
+    _assert(prefix_granularity > 0, "prefix_granularity must be > 0");
+    std::int32_t chunk_size = std::min(unscheduled, token_budget);
+    if (promotion_boundary_tokens > first_pos) {
+        chunk_size = std::min(chunk_size, promotion_boundary_tokens - first_pos);
+    }
+    if (chunk_size != unscheduled) {
+        return std::nullopt;
+    }
+
+    const std::int32_t tail_tokens = (first_pos + chunk_size) % prefix_granularity;
+    return tail_tokens != 0 && chunk_size - tail_tokens > 0 ? std::optional{tail_tokens} : std::nullopt;
+}
+
 std::vector<CacheGroupSpec> MakeSpecsFromConfig(const SchedulerConfig& config) {
     std::vector<CacheGroupSpec> specs;
     specs.reserve(config.cache_groups.size());

tokenspeed-scheduler/csrc/scheduler/operations/cache.h

@@ -21,6 +21,7 @@
 #pragma once
 
 #include <map>
+#include <optional>
 #include <span>
 #include <string>
 #include <vector>
@@ -41,6 +42,10 @@ std::vector<CacheGroupSpec> MakeSpecsFromConfig(const SchedulerConfig& config);
 std::int32_t AlignPrefillChunk(std::int32_t first_pos, std::int32_t unscheduled, std::int32_t token_budget,
                                std::int32_t prefix_granularity, std::int32_t promotion_boundary_tokens);
 
+std::optional<std::int32_t> FinalAlignedTailTokens(std::int32_t first_pos, std::int32_t unscheduled,
+                                                   std::int32_t token_budget, std::int32_t prefix_granularity,
+                                                   std::int32_t promotion_boundary_tokens);
+
 void FreeRequest(CacheCoordinator& coordinator, std::vector<BlockTable>& tables);
 
 // One row per config group_id. Each group allocator resolves the LCM placement

tokenspeed-scheduler/csrc/scheduler/operations/forward.cpp

@@ -78,16 +78,22 @@ void makeSnapshotStatePrefillSparse(std::span<GroupDemand> demands, std::span<co
     }
 }
 
-void deferSnapshotStateDecodeReservation(std::span<GroupDemand> demands,
-                                         std::span<const CacheGroupConfig> cache_groups) {
+void setSnapshotStatePrefillReserve(std::span<GroupDemand> demands, std::span<const CacheGroupConfig> cache_groups,
+                                    std::int32_t reserve_tokens) {
     _assert(demands.size() == cache_groups.size(), "demands/cache groups size mismatch");
     for (std::size_t i = 0; i < demands.size(); ++i) {
         if (cache_groups[i].IsSnapshotStateGroup()) {
-            demands[i].reserve_tokens = 0;
+            demands[i].reserve_tokens = reserve_tokens;
         }
     }
 }
 
+bool shouldSplitFinalStateCheckpoint(const SchedulerConfig& config, const CacheCoordinator& coordinator,
+                                     fsm::PrefillSource source) {
+    return config.role != Role::kD && !config.disable_prefix_cache && source == fsm::PrefillSource::kLocal &&
+           coordinator.HasMambaStateGroup();
+}
+
 void appendCompletedPrefixHashes(std::vector<std::string>& prefix_hashes,
                                  const std::vector<std::span<const std::int32_t>>& prefix_pages,
                                  std::int32_t filled_prefix_pages) {
@@ -287,11 +293,21 @@ std::optional<fsm::SchedulePrefillFirstChunkEvent> Scheduler::schedulePrefillFir
                  promotion_boundary_tokens > hit_tokens && promotion_boundary_tokens < request->PrefillSize()),
             "promotion boundary must be page-aligned and inside the unmatched prompt");
 
+    const fsm::PrefillSource source = config_.role == Role::kD && request->Is<fsm::Submitted>()
+                                          ? fsm::PrefillSource::kRemote
+                                          : fsm::PrefillSource::kLocal;
     const std::int32_t unscheduled = request->PrefillSize() - hit_tokens;
     std::int32_t tokens_this_round = std::min(remaining, unscheduled);
+    std::optional<std::int32_t> final_tail_tokens;
     if (coordinator_.HasMambaStateGroup() || promotion_boundary_tokens > 0) {
-        tokens_this_round = AlignPrefillChunk(hit_tokens, unscheduled, remaining, coordinator_.PrefixGranularity(),
-                                              promotion_boundary_tokens);
+        if (shouldSplitFinalStateCheckpoint(config_, coordinator_, source)) {
+            final_tail_tokens = FinalAlignedTailTokens(hit_tokens, unscheduled, remaining,
+                                                       coordinator_.PrefixGranularity(), promotion_boundary_tokens);
+        }
+        tokens_this_round = final_tail_tokens
+                                ? unscheduled - *final_tail_tokens
+                                : AlignPrefillChunk(hit_tokens, unscheduled, remaining,
+                                                    coordinator_.PrefixGranularity(), promotion_boundary_tokens);
         if (tokens_this_round == 0) {
             return std::nullopt;
         }
@@ -299,12 +315,12 @@ std::optional<fsm::SchedulePrefillFirstChunkEvent> Scheduler::schedulePrefillFir
 
     const bool completes_prefill = tokens_this_round == unscheduled;
     const std::int32_t decode_reserve = completes_prefill ? decode_input_tokens : 0;
-    const fsm::PrefillSource source = config_.role == Role::kD && request->Is<fsm::Submitted>()
-                                          ? fsm::PrefillSource::kRemote
-                                          : fsm::PrefillSource::kLocal;
+    const std::int32_t split_tail_tokens = final_tail_tokens.value_or(0);
+    const std::int32_t admission_reserve =
+        split_tail_tokens > 0 ? split_tail_tokens + decode_input_tokens : decode_reserve;
     std::vector<BlockTable> tables(static_cast<std::size_t>(coordinator_.NumGroups()));
     std::vector<GroupDemand> demands =
-        makeGroupDemands(tables, GroupDemand{.num_tokens = tokens_this_round, .reserve_tokens = decode_reserve});
+        makeGroupDemands(tables, GroupDemand{.num_tokens = tokens_this_round, .reserve_tokens = admission_reserve});
     if (source == fsm::PrefillSource::kLocal) {
         makeSnapshotStatePrefillSparse(demands, config_.cache_groups, coordinator_, hit_tokens + tokens_this_round);
     }
@@ -325,7 +341,7 @@ std::optional<fsm::SchedulePrefillFirstChunkEvent> Scheduler::schedulePrefillFir
             }
         }
     }
-    deferSnapshotStateDecodeReservation(demands, config_.cache_groups);
+    setSnapshotStatePrefillReserve(demands, config_.cache_groups, split_tail_tokens);
     std::vector<CacheKey> event_keys = registerKvEventPrefixPages(*request, match.candidate_prefix_hashes, 0);
     std::optional<CacheCoordinator::AdmissionResult> admission = admit(context, std::move(match.probe), demands);
     if (!admission) {
@@ -353,6 +369,7 @@ std::optional<fsm::SchedulePrefillFirstChunkEvent> Scheduler::schedulePrefillFir
             .prefix_hashes = std::move(match.prefix_hashes),
             .access_epoch = admission->access_epoch,
             .promotion_boundary_tokens = admission->promotion_boundary_tokens,
+            .state_checkpoint_tail_reserved = split_tail_tokens > 0,
         },
         std::move(admission->load_pairs),
     };
@@ -364,10 +381,23 @@ std::optional<fsm::SchedulePrefillEvent> Scheduler::schedulePrefill(
     const std::int32_t unscheduled = request->UnscheduledPrefillSize();
     const std::int32_t first_pos = request->PrefillSize() - unscheduled;
     fsm::CacheProgress cache_progress = request->CacheProgress();
+    const bool consumes_reserved_tail = cache_progress.state_checkpoint_tail_reserved;
     std::int32_t tokens_this_round = std::min(remaining, unscheduled);
+    std::optional<std::int32_t> final_tail_tokens;
     if (coordinator_.HasMambaStateGroup() || cache_progress.promotion_boundary_tokens > 0) {
-        tokens_this_round = AlignPrefillChunk(first_pos, unscheduled, remaining, coordinator_.PrefixGranularity(),
-                                              cache_progress.promotion_boundary_tokens);
+        if (shouldSplitFinalStateCheckpoint(config_, coordinator_, request->PrefillSource())) {
+            final_tail_tokens =
+                FinalAlignedTailTokens(first_pos, unscheduled, remaining, coordinator_.PrefixGranularity(),
+                                       cache_progress.promotion_boundary_tokens);
+        }
+        tokens_this_round = final_tail_tokens
+                                ? unscheduled - *final_tail_tokens
+                                : AlignPrefillChunk(first_pos, unscheduled, remaining, coordinator_.PrefixGranularity(),
+                                                    cache_progress.promotion_boundary_tokens);
+        if (final_tail_tokens) {
+            _assert(!consumes_reserved_tail, "cannot nest reserved state-checkpoint tails");
+            cache_progress.state_checkpoint_tail_reserved = true;
+        }
         if (tokens_this_round == 0) {
             return std::nullopt;
         }
@@ -375,6 +405,13 @@ std::optional<fsm::SchedulePrefillEvent> Scheduler::schedulePrefill(
 
     const bool completes_prefill = tokens_this_round == unscheduled;
     const std::int32_t decode_reserve = completes_prefill ? reserve_num_tokens_in_next_schedule_event : 0;
+    const std::int32_t checkpoint_tail_reserve = final_tail_tokens.value_or(0);
+    const std::int32_t admission_reserve =
+        final_tail_tokens ? checkpoint_tail_reserve + reserve_num_tokens_in_next_schedule_event : decode_reserve;
+    if (consumes_reserved_tail) {
+        _assert(tokens_this_round == unscheduled, "reserved state-checkpoint tail must complete in one round");
+        cache_progress.state_checkpoint_tail_reserved = false;
+    }
     const PrefillInfo previous = request->CurrentPrefillInfo();
     const std::int32_t num_computed_tokens = previous.already_scheduled_len + previous.extend_len;
     const CompletedPrefixPages completed =
@@ -388,12 +425,12 @@ std::optional<fsm::SchedulePrefillEvent> Scheduler::schedulePrefill(
                                      .new_prefix_hash_begin = completed.first_new_prefix_page,
                                      .completed_boundary_kind = completed.boundary_kind,
                                      .num_computed_tokens = num_computed_tokens,
-                                     .reserve_tokens = decode_reserve,
+                                     .reserve_tokens = admission_reserve,
                                  });
-    if (request->PrefillSource() == fsm::PrefillSource::kLocal) {
+    if (request->PrefillSource() == fsm::PrefillSource::kLocal && !consumes_reserved_tail) {
         makeSnapshotStatePrefillSparse(demands, config_.cache_groups, coordinator_, first_pos + tokens_this_round);
     }
-    deferSnapshotStateDecodeReservation(demands, config_.cache_groups);
+    setSnapshotStatePrefillReserve(demands, config_.cache_groups, checkpoint_tail_reserve);
     if (!admitWithKvEventTracking(context, *request, cache_progress, completed.first_new_prefix_page, demands)) {
         context.capacity_blocker = request->Id();
         return std::nullopt;
@@ -692,7 +729,7 @@ std::pair<std::vector<ForwardOperation>, std::vector<LoadBackOperation>> Schedul
                     // Decode-side recovery prefill runs in its own batch.
                     break;
                 }
-                if (request->Is<fsm::Prefilling>()) {
+                if (request->Is<fsm::Prefilling>() && !request->CacheProgress().state_checkpoint_tail_reserved) {
                     // Admission reserves only this chunk, not the request's
                     // remaining prompt. Keep one incomplete prefill as the
                     // head of line so another request cannot strand it by
@@ -733,7 +770,7 @@ std::pair<std::vector<ForwardOperation>, std::vector<LoadBackOperation>> Schedul
                     pd_transfer_pins_.insert(request->Id());
                 }
                 trackPendingForwardResult(request);
-                if (request->Is<fsm::Prefilling>()) {
+                if (request->Is<fsm::Prefilling>() && !request->CacheProgress().state_checkpoint_tail_reserved) {
                     break;
                 }
             }

tokenspeed-scheduler/csrc/scheduler/scheduler.cpp

@@ -115,6 +115,18 @@ std::int64_t Scheduler::singleRequestLcmBlocksRequired(std::int32_t token_limit)
     const std::int64_t max_prompt_tokens =
         std::max<std::int64_t>(static_cast<std::int64_t>(token_limit) - decode_width, 0);
     const std::int64_t chunk_tokens = config_.max_scheduled_tokens;
+    const std::int64_t prefix_granularity = config_.prefix_granularity;
+    // A final sub-page tail can follow the first aligned body, or a later body
+    // that also retains an input checkpoint. Bound both cases independently.
+    const bool splits_final_state_checkpoint = config_.role != Role::kD && !config_.disable_prefix_cache;
+    const auto max_split_tail_after = [&](std::int64_t minimum_body_end) {
+        return splits_final_state_checkpoint
+                   ? std::max<std::int64_t>(0, std::min({prefix_granularity - 1, chunk_tokens - prefix_granularity,
+                                                         max_prompt_tokens - minimum_body_end}))
+                   : 0;
+    };
+    const std::int64_t max_first_chunk_tail_tokens = max_split_tail_after(prefix_granularity);
+    const std::int64_t max_later_chunk_tail_tokens = max_split_tail_after(2 * prefix_granularity);
 
     std::vector<std::int64_t> group_pages(static_cast<std::size_t>(coordinator_.NumGroups()));
     for (std::int32_t i = 0; i < coordinator_.NumGroups(); ++i) {
@@ -125,7 +137,14 @@ std::int64_t Scheduler::singleRequestLcmBlocksRequired(std::int32_t token_limit)
                 if (token_limit == 0) return std::int64_t{0};
                 const std::int64_t input_lookback =
                     max_prompt_tokens > chunk_tokens ? coordinator_.GroupBoundaryLookbackPages(i) : 0;
-                return std::max<std::int64_t>(2, input_lookback + 1);
+                const std::int64_t first_split_checkpoint_peak =
+                    max_first_chunk_tail_tokens == 0 ? 0 : 1 + ceilDiv(max_first_chunk_tail_tokens, block_granularity);
+                const std::int64_t later_split_checkpoint_peak =
+                    max_later_chunk_tail_tokens == 0 ? 0
+                                                     : coordinator_.GroupBoundaryLookbackPages(i) + 1 +
+                                                           ceilDiv(max_later_chunk_tail_tokens, block_granularity);
+                return std::max(
+                    {std::int64_t{2}, input_lookback + 1, first_split_checkpoint_peak, later_split_checkpoint_peak});
             }
             // Across every prompt up to max_prompt_tokens, retain the largest
             // resident window seen by either the first chunk or a later chunk.

tokenspeed-scheduler/python/tests/test_kv_cache.py

@@ -199,24 +199,25 @@ def _make_k3_128k_config(num_device_pages: int) -> ts.SchedulerConfig:
 
 
 def test_k3_reports_group_aware_single_request_capacity() -> None:
-    # Each sparse State group needs two rolling checkpoints. The three groups
-    # therefore leave 278 of the 284 usable parents for Full KV.
-    # K_full=12 and P=128 expose 278 * 12 * 128 tokens.
+    # Each sparse State group needs an input checkpoint, an aligned-body
+    # checkpoint, and its reserved tail. The three groups therefore leave 275
+    # of the 284 usable parents for Full KV. K_full=12 and P=128 expose
+    # 275 * 12 * 128 tokens.
     scheduler = ts.Scheduler(_make_k3_128k_config(285))
-    assert scheduler.max_single_request_tokens() == 427_008
+    assert scheduler.max_single_request_tokens() == 422_400
 
 
 def test_k3_128k_requires_group_aware_shared_pool_geometry() -> None:
     prompt = _spec("128k", list(range(131_072)))
 
-    # Six State parents plus 86 Full parents admit 128K; one fewer Full parent
+    # Nine State parents plus 86 Full parents admit 128K; one fewer Full parent
     # is 512 tokens short because each Full parent carries 12 * 128 tokens.
-    undersized = ts.Scheduler(_make_k3_128k_config(92))
+    undersized = ts.Scheduler(_make_k3_128k_config(95))
     assert undersized.max_single_request_tokens() < 131_072
 
-    corrected = ts.Scheduler(_make_k3_128k_config(93))
+    corrected = ts.Scheduler(_make_k3_128k_config(96))
     before = corrected.available_kv_pages()
-    assert before == 92
+    assert before == 95
     corrected.submit_requests([prompt])
     completed_tokens = 0
     for chunk in range(32):
@@ -284,9 +285,7 @@ def _drive_k3_to_retract(scheduler) -> dict[str, dict[int, int]]:
 
 
 def test_k3_readmit_rebuilds_all_four_tables_and_restores_pages() -> None:
-    """Binding-marshalling smoke for readmit: the only python test that reads
-    ``op.prefill_lengths`` through the real nanobind property (the C++ suite
-    covers the retract/readmit scheduler scenarios themselves)."""
+    """Binding smoke for a split readmit and its reserved state tail."""
     scheduler = ts.Scheduler(_make_k3_config())
     before = scheduler.available_kv_pages()
     pre_retract_pages = _drive_k3_to_retract(scheduler)
@@ -294,19 +293,20 @@ def test_k3_readmit_rebuilds_all_four_tables_and_restores_pages() -> None:
     for request_id in ("b", "c", "d"):
         _finish(scheduler, request_id)
 
-    readmit = _find_forward_op(scheduler.next_execution_plan())
-    assert readmit is not None
-    assert tuple(readmit.request_ids) == ("a",)
-    assert tuple(readmit.prefill_lengths) == (11,)
-    assert readmit.extend_prefix_lens[0] + readmit.input_lengths[0] == 11
-    tables = dict(readmit.block_tables)
+    body = _find_forward_op(scheduler.next_execution_plan())
+    assert body is not None
+    assert tuple(body.request_ids) == ("a",)
+    assert tuple(body.prefill_lengths) == (11,)
+    assert tuple(body.extend_prefix_lens) == (8,)
+    assert tuple(body.input_lengths) == (2,)
+    tables = dict(body.block_tables)
     assert tuple(tables) == K3_GROUP_IDS
     prefix_granularity = _make_k3_config().prefix_granularity
-    assert readmit.extend_prefix_lens[0] % prefix_granularity == 0
-    prefix_slots = readmit.extend_prefix_lens[0] // prefix_granularity
+    assert body.extend_prefix_lens[0] % prefix_granularity == 0
+    prefix_slots = body.extend_prefix_lens[0] // prefix_granularity
     assert prefix_slots == 4
     expected_slots = (
-        readmit.prefill_lengths[0] + prefix_granularity - 1
+        body.prefill_lengths[0] + prefix_granularity - 1
     ) // prefix_granularity
     assert expected_slots == 6
 
@@ -331,22 +331,22 @@ def test_k3_readmit_rebuilds_all_four_tables_and_restores_pages() -> None:
         tail = row[prefix_slots:]
         assert len(tail) == 2
         group_tail = _positive_pages(tail)
-        if group_id == K3_GROUP_IDS[0]:
-            # Full history remains dense.
-            assert all(page > 0 for page in tail)
-            assert len(group_tail) == 2
-        else:
-            # Sparse State recovery materializes only the endpoint checkpoint;
-            # the first logical tail slot remains a null hole.
-            assert tail[0] == 0
-            assert tail[1] > 0
-            assert len(group_tail) == 1
+        # The aligned body materializes its state checkpoint and atomically
+        # reserves the final prompt-tail slot, so every group owns both pages.
+        assert all(page > 0 for page in tail)
+        assert len(group_tail) == 2
         fresh_tail_entries.extend(group_tail)
 
     assert len(set(all_positive_entries)) == len(all_positive_entries)
     assert len(set(fresh_tail_entries)) == len(fresh_tail_entries)
     assert set(fresh_tail_entries).isdisjoint(restored_pages)
 
+    tail = _find_forward_op(scheduler.next_execution_plan())
+    assert tail is not None
+    assert tuple(tail.request_ids) == ("a",)
+    assert tuple(tail.extend_prefix_lens) == (10,)
+    assert tuple(tail.input_lengths) == (1,)
+
     _advance_tokens(scheduler, "a", [3000])
     scheduler.next_execution_plan()
     _advance_tokens(scheduler, "a", [3001])

tokenspeed-scheduler/tests/cpp/test_cache_scenarios.cpp

@@ -192,6 +192,176 @@ TEST_F(MambaChunkAlignmentSuite, PartialPrefillEndsAtStatePageBoundary) {
     }
 }
 
+class MambaStateCheckpointSplitSuite : public MambaChunkAlignmentSuite {
+protected:
+    SchedulerConfig MakeConfig() override {
+        SchedulerConfig cfg = MambaChunkAlignmentSuite::MakeConfig();
+        cfg.max_scheduled_tokens = 64;
+        cfg.disable_prefix_cache = false;
+        return cfg;
+    }
+};
+
+TEST_F(MambaStateCheckpointSplitSuite, ReservesAndBatchesDependentTails) {
+    RequestSpec first = MakeRequestSpec("a", /*num_pages=*/3);
+    RequestSpec second = MakeRequestSpec("b", /*num_pages=*/3, /*start=*/100);
+    first.tokens.resize(10);
+    second.tokens.resize(10);
+    Submit({first, second});
+
+    ExecutionPlan bodies = PlanOnce();
+    const ForwardBatch* body_op = FindForwardBatch(bodies);
+    ASSERT_NE(body_op, nullptr);
+    EXPECT_EQ(body_op->request_ids, (std::vector<std::string>{"a", "b"}));
+    EXPECT_EQ(body_op->input_lengths, (std::vector<std::int32_t>{8, 8}));
+    for (const auto& row : body_op->block_tables.at("state")) {
+        EXPECT_EQ(row.size(), 3u);
+    }
+
+    ExecutionPlan tails = PlanOnce();
+    const ForwardBatch* tail_op = FindForwardBatch(tails);
+    ASSERT_NE(tail_op, nullptr);
+    EXPECT_EQ(tail_op->request_ids, (std::vector<std::string>{"a", "b"}));
+    EXPECT_EQ(tail_op->extend_prefix_lens, (std::vector<std::int32_t>{8, 8}));
+    EXPECT_EQ(tail_op->input_lengths, (std::vector<std::int32_t>{2, 2}));
+}
+
+TEST(MambaStateCheckpointCapacityTest, CountsInputBodyAndReservedTailCheckpoints) {
+    SchedulerConfig cfg{};
+    cfg.prefix_granularity = 4;
+    cfg.device_allocator.total_pages = 3;  // null + two usable state blocks
+    cfg.host_allocator.total_pages = 0;
+    cfg.max_scheduled_tokens = 8;
+    cfg.max_batch_size = 1;
+    cfg.disable_l2_cache = true;
+    cfg.cache_groups = {
+        MakeGroup("state", cfg.prefix_granularity, cfg.device_allocator.total_pages,
+                  CacheGroupConfig::Retention::FullHistory, CacheGroupFamily::State),
+    };
+
+    Scheduler scheduler{std::move(cfg)};
+
+    // Up to eight prompt tokens plus the decode reservation fit in two
+    // blocks. A longer prompt can reuse one input checkpoint while exposing an
+    // aligned body and reserving its sub-page tail, which needs three.
+    EXPECT_EQ(scheduler.MaxSingleRequestTokens(), 9);
+    RequestSpec too_long{
+        .request_id = "too-long",
+        .tokens = std::vector<std::int32_t>(14, 1),
+    };
+    EXPECT_THROW(scheduler.SubmitRequests({too_long}), std::invalid_argument);
+}
+
+TEST(MambaStateCheckpointCapacityTest, CountsFirstChunkBodyAndSubPageTail) {
+    SchedulerConfig cfg{};
+    cfg.prefix_granularity = 4;
+    cfg.device_allocator.total_pages = 4;  // null + three usable state blocks
+    cfg.host_allocator.total_pages = 0;
+    cfg.max_scheduled_tokens = 8;
+    cfg.max_batch_size = 1;
+    cfg.disable_l2_cache = true;
+    cfg.cache_groups = {
+        MakeGroup("state", /*block_granularity=*/1, cfg.device_allocator.total_pages,
+                  CacheGroupConfig::Retention::FullHistory, CacheGroupFamily::State),
+    };
+
+    Scheduler scheduler{std::move(cfg)};
+
+    // A seven-token prompt would split into a four-token body and a
+    // three-token tail. The body checkpoint plus the reserved tail need four
+    // state blocks, so three usable blocks cannot admit it and its decode.
+    EXPECT_EQ(scheduler.MaxSingleRequestTokens(), 7);
+    RequestSpec too_long{
+        .request_id = "too-long",
+        .tokens = std::vector<std::int32_t>(7, 1),
+    };
+    EXPECT_THROW(scheduler.SubmitRequests({too_long}), std::invalid_argument);
+}
+
+class MambaStateCheckpointNoPrefixCacheSuite : public MambaStateCheckpointSplitSuite {
+protected:
+    SchedulerConfig MakeConfig() override {
+        SchedulerConfig cfg = MambaStateCheckpointSplitSuite::MakeConfig();
+        cfg.disable_prefix_cache = true;
+        return cfg;
+    }
+};
+
+TEST_F(MambaStateCheckpointNoPrefixCacheSuite, KeepsSingleFinalChunk) {
+    RequestSpec spec = MakeRequestSpec("r1", /*num_pages=*/3);
+    spec.tokens.resize(10);
+    Submit(spec);
+
+    ExecutionPlan plan = PlanOnce();
+    const ForwardBatch* op = FindForwardBatch(plan);
+    ASSERT_NE(op, nullptr);
+    EXPECT_EQ(op->input_lengths, std::vector<std::int32_t>{10});
+}
+
+class MambaStateCheckpointPrefillRoleSuite : public MambaStateCheckpointSplitSuite {
+protected:
+    SchedulerConfig MakeConfig() override {
+        SchedulerConfig cfg = MambaStateCheckpointSplitSuite::MakeConfig();
+        cfg.role = Role::kP;
+        cfg.enable_pd_cache = true;
+        for (CacheGroupConfig& group : cfg.cache_groups) {
+            group.transfer_policy =
+                group.IsSnapshotStateGroup() ? CacheTransferPolicy::LatestSnapshot : CacheTransferPolicy::FullSuffix;
+        }
+        return cfg;
+    }
+
+    void SendBootstrapped(const std::string& request_id) {
+        ExecutionEvent event;
+        event.With(pd::BootstrappedEvent{request_id});
+        scheduler_->Advance(std::move(event));
+    }
+};
+
+TEST_F(MambaStateCheckpointPrefillRoleSuite, SplitsLocalPrefillOnPrefillWorker) {
+    RequestSpec spec = MakeRequestSpec("r1", /*num_pages=*/3);
+    spec.tokens.resize(10);
+    Submit(spec);
+    SendBootstrapped("r1");
+
+    ExecutionPlan body_plan = PlanOnce();
+    const ForwardBatch* body = FindForwardBatch(body_plan);
+    ASSERT_NE(body, nullptr);
+    EXPECT_TRUE(body->IsLocalPrefill());
+    EXPECT_EQ(body->extend_prefix_lens, std::vector<std::int32_t>{0});
+    EXPECT_EQ(body->input_lengths, std::vector<std::int32_t>{8});
+
+    ExecutionPlan tail_plan = PlanOnce();
+    const ForwardBatch* tail = FindForwardBatch(tail_plan);
+    ASSERT_NE(tail, nullptr);
+    EXPECT_TRUE(tail->IsLocalPrefill());
+    EXPECT_EQ(tail->extend_prefix_lens, std::vector<std::int32_t>{8});
+    EXPECT_EQ(tail->input_lengths, std::vector<std::int32_t>{2});
+}
+
+class MambaStateCheckpointDecodeRoleSuite : public MambaStateCheckpointPrefillRoleSuite {
+protected:
+    SchedulerConfig MakeConfig() override {
+        SchedulerConfig cfg = MambaStateCheckpointPrefillRoleSuite::MakeConfig();
+        cfg.role = Role::kD;
+        return cfg;
+    }
+};
+
+TEST_F(MambaStateCheckpointDecodeRoleSuite, KeepsRemoteAdmissionWhole) {
+    RequestSpec spec = MakeRequestSpec("r1", /*num_pages=*/3);
+    spec.tokens.resize(10);
+    Submit(spec);
+    SendBootstrapped("r1");
+
+    ExecutionPlan admission_plan = PlanOnce();
+    const ForwardBatch* admission = FindForwardBatch(admission_plan);
+    ASSERT_NE(admission, nullptr);
+    EXPECT_FALSE(admission->IsLocalPrefill());
+    EXPECT_EQ(admission->extend_prefix_lens, std::vector<std::int32_t>{0});
+    EXPECT_EQ(admission->input_lengths, std::vector<std::int32_t>{10});
+}
+
 class MambaSparsePrefillSuite : public MambaChunkAlignmentSuite {
 protected:
     SchedulerConfig MakeConfig() override {

tokenspeed-scheduler/tests/cpp/test_forward_cache_ops.cpp

@@ -98,6 +98,19 @@ TEST(AlignPrefillChunkTest, ReachedPromotionUsesOrdinaryPageAlignment) {
               8);
 }
 
+TEST(FinalAlignedTailTokensTest, FindsSubPageTailAfterAlignedBody) {
+    const std::optional<std::int32_t> tail =
+        FinalAlignedTailTokens(/*first_pos=*/16, /*unscheduled=*/11, /*token_budget=*/16,
+                               /*prefix_granularity=*/4, /*promotion_boundary_tokens=*/0);
+    EXPECT_EQ(tail, 3);
+}
+
+TEST(FinalAlignedTailTokensTest, LeavesAStandaloneSubPageWhole) {
+    EXPECT_EQ(FinalAlignedTailTokens(/*first_pos=*/24, /*unscheduled=*/3, /*token_budget=*/16,
+                                     /*prefix_granularity=*/4, /*promotion_boundary_tokens=*/0),
+              std::nullopt);
+}
+
 TEST(ForwardCacheOpsPrefill, FirstChunkAcquiresPagesForTokens) {
     BlockPool pool(/*num_lcm_blocks=*/32);
     CacheCoordinator coordinator = MakeTwoGroup(pool);

tokenspeed-scheduler/tests/cpp/test_scheduler_plan.cpp

@@ -368,7 +368,14 @@ protected:
 
 TEST_F(HybridPrefixPromotionTestSuite, ThirdRequestReusesPromotedStateBoundary) {
     Submit(MakeHybridRequest("seed", 100));
-    PlanOnce();
+    const ExecutionPlan seed_body_plan = PlanOnce();
+    const ForwardBatch* seed_body = FindForwardBatch(seed_body_plan);
+    ASSERT_NE(seed_body, nullptr);
+    EXPECT_EQ(seed_body->input_lengths, std::vector<std::int32_t>{10});
+    const ExecutionPlan seed_tail_plan = PlanOnce();
+    const ForwardBatch* seed_tail = FindForwardBatch(seed_tail_plan);
+    ASSERT_NE(seed_tail, nullptr);
+    EXPECT_EQ(seed_tail->input_lengths, std::vector<std::int32_t>{1});
     SendForwardDone("seed", {900});
     PlanOnce();
     SendFinish("seed");
@@ -380,7 +387,14 @@ TEST_F(HybridPrefixPromotionTestSuite, ThirdRequestReusesPromotedStateBoundary)
     ASSERT_NE(promotion, nullptr);
     ASSERT_EQ(promotion->request_ids, std::vector<std::string>{"promote"});
     EXPECT_EQ(promotion->input_lengths, std::vector<std::int32_t>{8});
-    PlanOnce();
+    const ExecutionPlan promotion_body_plan = PlanOnce();
+    const ForwardBatch* promotion_body = FindForwardBatch(promotion_body_plan);
+    ASSERT_NE(promotion_body, nullptr);
+    EXPECT_EQ(promotion_body->input_lengths, std::vector<std::int32_t>{2});
+    const ExecutionPlan promotion_tail_plan = PlanOnce();
+    const ForwardBatch* promotion_tail = FindForwardBatch(promotion_tail_plan);
+    ASSERT_NE(promotion_tail, nullptr);
+    EXPECT_EQ(promotion_tail->input_lengths, std::vector<std::int32_t>{1});
     SendForwardDone("promote", {901});
     PlanOnce();
     SendFinish("promote");