feat(scheduler): live-tail allocation for sliding groups

python/tokenspeed/runtime/configs/paged_cache_spec.py

@@ -35,6 +35,8 @@ class PagedCacheGroupSpec:
     family: Family = "history"
     # Per-group page tokens; None -> scheduler global block_size, else a multiple of it.
     block_size: int | None = None
+    # Sliding-only: real pages for the live tail + LCM-boundary resume pages, holes elsewhere (paged conv state only).
+    live_tail_alloc: bool = False
 
 
 _PAGED_CACHE_GROUP_DUMMY_PAGES = 1

python/tokenspeed/runtime/engine/scheduler_utils.py

@@ -171,6 +171,9 @@ def pool_to_paged_cache_groups(pool: Any) -> list:
         # Ctor default 0 = global base; a spec block_size sets the per-group granularity.
         if getattr(spec, "block_size", None):
             cfg.block_size = int(spec.block_size)
+        # Attribute-set (not ctor kwarg) so an older compiled ext fails loudly here.
+        if getattr(spec, "live_tail_alloc", False):
+            cfg.live_tail_alloc = True
         out.append(cfg)
     return out
 

python/tokenspeed/runtime/layers/attention/registry.py

@@ -472,6 +472,8 @@ def _publish_inkling_conv_groups(config, model_config, server_args) -> tuple[int
             sliding_window_tokens=conv_bt + tc.sconv_kernel_size,
             family="history",
             block_size=conv_bt,
+            # Conv reads the chunk from packed activations, so slid-out pages can be holes (degenerates to full at this block size).
+            live_tail_alloc=True,
         )
         for label in dict.fromkeys(tc.paged_cache_layer_types)
     )
@@ -488,6 +490,8 @@ def _publish_inkling_conv_groups(config, model_config, server_args) -> tuple[int
                 sliding_window_tokens=hbt + tc.sconv_kernel_size,
                 family="history",
                 block_size=hbt,
+                # Drops the ~17 GB/GPU 8k-chunk hiddenconv admission transient ~16x.
+                live_tail_alloc=True,
             )
             for label in dict.fromkeys(tc.paged_cache_layer_types)
         )

tokenspeed-scheduler/CMakeLists.txt

@@ -127,6 +127,7 @@ if(TOKENSPEED_SCHEDULER_BUILD_TESTS)
         tests/cpp/test_block_ref.cpp
         tests/cpp/test_full_attn_manager.cpp
         tests/cpp/test_swa_manager.cpp
+        tests/cpp/test_swa_live_alloc.cpp
         tests/cpp/test_match_host_pages.cpp
         tests/cpp/test_kv_cache_coordinator.cpp
         tests/cpp/test_forward_cache_ops.cpp

tokenspeed-scheduler/bindings/python_module.cpp

@@ -174,15 +174,16 @@ NB_MODULE(tokenspeed_scheduler_ext, m) {
             [](tokenspeed::PagedCacheGroupConfig* self, std::string group_id, std::int32_t rows_per_page,
                std::int32_t entry_stride_tokens, std::int32_t total_pages,
                tokenspeed::PagedCacheGroupConfig::Retention retention,
-               std::optional<std::int32_t> sliding_window_tokens, tokenspeed::PagedCacheGroupFamily family) {
+               std::optional<std::int32_t> sliding_window_tokens, tokenspeed::PagedCacheGroupFamily family,
+               bool live_tail_alloc) {
                 new (self) tokenspeed::PagedCacheGroupConfig{
                     std::move(group_id), rows_per_page, entry_stride_tokens,   total_pages,
-                    /*block_size=*/0,    retention,     sliding_window_tokens, family};
+                    /*block_size=*/0,    retention,     sliding_window_tokens, family,      live_tail_alloc};
             },
             nb::arg("group_id"), nb::arg("rows_per_page"), nb::arg("entry_stride_tokens"), nb::arg("total_pages"),
             nb::arg("retention") = tokenspeed::PagedCacheGroupConfig::Retention::FullHistory,
             nb::arg("sliding_window_tokens") = std::nullopt,
-            nb::arg("family") = tokenspeed::PagedCacheGroupFamily::History)
+            nb::arg("family") = tokenspeed::PagedCacheGroupFamily::History, nb::arg("live_tail_alloc") = false)
         .def_rw("group_id", &tokenspeed::PagedCacheGroupConfig::group_id)
         .def_rw("rows_per_page", &tokenspeed::PagedCacheGroupConfig::rows_per_page)
         .def_rw("entry_stride_tokens", &tokenspeed::PagedCacheGroupConfig::entry_stride_tokens)
@@ -191,6 +192,7 @@ NB_MODULE(tokenspeed_scheduler_ext, m) {
         .def_rw("retention", &tokenspeed::PagedCacheGroupConfig::retention)
         .def_rw("sliding_window_tokens", &tokenspeed::PagedCacheGroupConfig::sliding_window_tokens)
         .def_rw("family", &tokenspeed::PagedCacheGroupConfig::family)
+        .def_rw("live_tail_alloc", &tokenspeed::PagedCacheGroupConfig::live_tail_alloc)
         .def("raw_tokens_per_page", &tokenspeed::PagedCacheGroupConfig::RawTokensPerPage)
         .def("validate", &tokenspeed::PagedCacheGroupConfig::Validate);
 

tokenspeed-scheduler/csrc/cache/cache_types.h

@@ -39,6 +39,9 @@ struct KvCacheSpec {
     AttnKind kind;
     std::int32_t block_size;
     std::int32_t sliding_window;  // 0 for full attention
+    // Sliding-only: real pages for the live tail + LCM-boundary resume pages, holes elsewhere; MakeCoordinator sets the
+    // alignment.
+    bool live_tail_alloc{false};
 };
 
 // Per-request logical-page -> physical-page mapping.
@@ -74,6 +77,9 @@ class BlockTable {
 
     std::vector<BlockRef> blocks_{};
     std::int32_t tail_avail_{0};
+    // Reclaim scans never revisit slots below this. NOT a null guarantee: claimed/extension holes can strand
+    // real pages under an alignment==0 early break (same pages the pre-floor scan stranded); Free releases them.
+    std::int32_t reclaim_floor_{0};
 };
 
 // The single flattening authority: BlockId() per logical slot, null holes written as 0, no compaction.

tokenspeed-scheduler/csrc/cache/forward_cache_ops.cpp

@@ -84,6 +84,7 @@ std::vector<KvCacheSpec> MakeSpecsFromConfig(const SchedulerConfig& config) {
             .kind = is_swa ? AttnKind::kSlidingWindow : AttnKind::kFull,
             .block_size = block_size,
             .sliding_window = is_swa ? group.sliding_window_tokens.value_or(0) : 0,
+            .live_tail_alloc = is_swa && group.live_tail_alloc,
         });
     }
     return specs;

tokenspeed-scheduler/csrc/cache/kv_cache_coordinator.cpp

@@ -215,6 +215,25 @@ std::int32_t KvCacheCoordinator::BlocksNeededFor(std::int32_t num_tokens) const
     return total_needed;
 }
 
+std::int32_t KvCacheCoordinator::BlocksNeededForSequential(std::span<const BlockTable> tables,
+                                                           std::int32_t first_tokens, std::int32_t extra_tokens) const {
+    _assert(tables.size() == groups_.size(), "tables/groups size mismatch");
+    std::int32_t total_needed = 0;
+    for (std::size_t i = 0; i < groups_.size(); ++i) {
+        total_needed += groups_[i].Manager().BlocksNeededForSequential(tables[i], first_tokens, extra_tokens);
+    }
+    return total_needed;
+}
+
+std::int32_t KvCacheCoordinator::BlocksNeededForSequential(std::int32_t first_tokens, std::int32_t extra_tokens) const {
+    const BlockTable fresh;
+    std::int32_t total_needed = 0;
+    for (const CacheGroup& group : groups_) {
+        total_needed += group.Manager().BlocksNeededForSequential(fresh, first_tokens, extra_tokens);
+    }
+    return total_needed;
+}
+
 bool KvCacheCoordinator::Acquire(std::span<BlockTable> tables, std::int32_t num_tokens) {
     // Check-then-act: no group is ever left in a partial/unaligned state.
     if (BlocksNeededFor(tables, num_tokens) > pool_.NumFreeBlocks()) {
@@ -296,7 +315,10 @@ KvCacheCoordinator MakeCoordinator(std::span<const KvCacheSpec> specs, BlockPool
         } else if (spec.kind == AttnKind::kMambaState) {
             manager = std::make_unique<MambaStateManager>(spec.block_size);
         } else {
-            manager = std::make_unique<SwaManager>(spec.block_size, spec.sliding_window);
+            // Live-tail groups keep resume pages behind every LCM-aligned boundary (matches the engine lcm_align
+            // persist set).
+            manager =
+                std::make_unique<SwaManager>(spec.block_size, spec.sliding_window, spec.live_tail_alloc ? lcm : 0);
         }
         groups.emplace_back(spec, static_cast<std::uint32_t>(i), std::move(manager));
     }

tokenspeed-scheduler/csrc/cache/kv_cache_coordinator.h

@@ -83,6 +83,12 @@ class KvCacheCoordinator {
     // Fresh-table overload for a not-yet-allocated request (no tail credit).
     std::int32_t BlocksNeededFor(std::int32_t num_tokens) const;
 
+    // Peak demand of Acquire(first) then Acquire(extra) -- the chunk + decode-reserve admission
+    // charge (live-tail groups are non-monotonic, so the combined query can under-charge).
+    std::int32_t BlocksNeededForSequential(std::span<const BlockTable> tables, std::int32_t first_tokens,
+                                           std::int32_t extra_tokens) const;
+    std::int32_t BlocksNeededForSequential(std::int32_t first_tokens, std::int32_t extra_tokens) const;
+
     // end_tokens = the chunk's end position (-1 = unknown/legacy): aligned-final-page-only
     // groups register nothing without it, since only an aligned chunk end holds a real snapshot.
     // first_slot may reach back to the fold grid (decode re-covers so coarse groups can fold);

tokenspeed-scheduler/csrc/cache/kv_cache_manager.h

@@ -64,7 +64,9 @@ class KvCacheManager {
     }
 
     // All-or-nothing (tail-page room first, then fresh pages): on shortfall the table is unchanged, returns false.
-    bool Acquire(BlockPool& pool, BlockTable& table, std::int32_t num_tokens) {
+    // Virtual so live-tail overrides can hole dead slots; overrides keep the all-or-nothing contract + BlocksNeededFor
+    // mirror.
+    virtual bool Acquire(BlockPool& pool, BlockTable& table, std::int32_t num_tokens) {
         if (num_tokens <= 0) {
             return true;
         }
@@ -95,21 +97,29 @@ class KvCacheManager {
                 table.blocks_.push_back(BlockRef::Share(pool, pool.NullBlock()));
                 continue;
             }
-            const bool acquired = Acquire(pool, table, block_size_);
+            // Base Acquire on purpose: a loaded host page must land in a REAL slot, not a live-tail hole.
+            const bool acquired = KvCacheManager::Acquire(pool, table, block_size_);
             _assert(acquired, "pre-checked Acquire must succeed");
             load_pairs.emplace_back(host_block, table.blocks_.back().Get());
         }
     }
 
-    // Pure query mirroring Acquire's page math exactly.
-    std::int32_t BlocksNeededFor(const BlockTable& table, std::int32_t num_tokens) const {
+    // Pure query mirroring Acquire's page math exactly (virtual in lockstep with Acquire).
+    virtual std::int32_t BlocksNeededFor(const BlockTable& table, std::int32_t num_tokens) const {
         if (num_tokens <= table.tail_avail_) {
             return 0;
         }
         std::int32_t over = num_tokens - table.tail_avail_;
         return (over + block_size_ - 1) / block_size_;
     }
 
+    // Peak demand of Acquire(first) then Acquire(extra): live-tail holing makes the combined
+    // query non-monotonic, so gating on BlocksNeededFor(first + extra) can under-charge.
+    virtual std::int32_t BlocksNeededForSequential(const BlockTable& table, std::int32_t first_tokens,
+                                                   std::int32_t extra_tokens) const {
+        return BlocksNeededFor(table, first_tokens + extra_tokens);  // exact when pages never turn into holes
+    }
+
     // State snapshots are only boundary-correct where a forward call ended page-aligned:
     // such groups register just the final full page of an aligned range.
     virtual bool RegistersAlignedFinalPageOnly() const { return false; }
@@ -158,9 +168,25 @@ class KvCacheManager {
         pool.FreeBlocks(batch);
         table.blocks_.clear();
         table.tail_avail_ = 0;
+        table.reclaim_floor_ = 0;
     }
 
 protected:
+    // Table-mutation helpers for derived Acquire overrides (BlockTable befriends only this base).
+    static std::int32_t TableExtentTokens(const BlockTable& table, std::int32_t block_size) {
+        return table.NumBlocks() * block_size - table.tail_avail_;
+    }
+    static std::int32_t TableTailAvail(const BlockTable& table) { return table.tail_avail_; }
+    static void SetTableTailAvail(BlockTable& table, std::int32_t tail_avail) { table.tail_avail_ = tail_avail; }
+    static std::int32_t TableReclaimFloor(const BlockTable& table) { return table.reclaim_floor_; }
+    static void SetTableReclaimFloor(BlockTable& table, std::int32_t floor) { table.reclaim_floor_ = floor; }
+    static void AppendRealBlock(BlockPool& pool, BlockTable& table, CacheBlock* block) {
+        table.blocks_.push_back(BlockRef::Adopt(pool, block));
+    }
+    static void AppendHole(BlockPool& pool, BlockTable& table) {
+        table.blocks_.push_back(BlockRef::Share(pool, pool.NullBlock()));
+    }
+
     std::int32_t block_size_;
 };