From 2948da62ab4b9245fe70ed653e842f5a17dc9064 Mon Sep 17 00:00:00 2001 From: kangmeng3 Date: Wed, 22 Jul 2026 15:03:53 +0800 Subject: [PATCH] feat: support DSV4 device prefix cache (SWA + C4 + C128). --- .../block/composite_block_manager_test.cpp | 204 +++++- xllm/core/distributed_runtime/master.cpp | 11 - xllm/core/framework/block/block_manager.h | 24 +- .../framework/block/block_manager_impl.cpp | 47 +- .../core/framework/block/block_manager_impl.h | 11 +- .../block/composite_block_manager.cpp | 630 +++++++++++++----- .../framework/block/composite_block_manager.h | 102 +-- .../block/concurrent_block_manager_impl.cpp | 5 +- .../block/concurrent_block_manager_impl.h | 3 +- .../block/linear_state_block_manager.cpp | 66 +- .../block/linear_state_block_manager.h | 36 +- .../framework/block/single_block_manager.cpp | 3 +- .../framework/block/single_block_manager.h | 3 +- .../block/sliding_window_block_manager.cpp | 75 ++- .../block/sliding_window_block_manager.h | 42 +- .../linear_state_prefix_cache.cpp | 229 +++++-- .../prefix_cache/linear_state_prefix_cache.h | 72 +- .../framework/prefix_cache/prefix_cache.cpp | 14 +- .../framework/prefix_cache/prefix_cache.h | 82 +-- .../prefix_cache/prefix_cache_factory.cpp | 14 +- xllm/core/framework/request/sequence.cpp | 33 +- xllm/core/framework/request/sequence.h | 42 +- .../framework/request/sequence_kv_state.cpp | 47 +- .../framework/request/sequence_kv_state.h | 24 + .../xtensor/xtensor_block_manager_impl.cpp | 3 +- .../xtensor/xtensor_block_manager_impl.h | 3 +- .../disagg_pd_chunked_prefill_scheduler.cpp | 9 +- xllm/core/scheduler/scheduler_policy.cpp | 9 +- .../scheduler/zero_eviction_scheduler.cpp | 8 +- 29 files changed, 1248 insertions(+), 603 deletions(-) diff --git a/tests/core/framework/block/composite_block_manager_test.cpp b/tests/core/framework/block/composite_block_manager_test.cpp index 7a72dd3618..0f9a9c924f 100644 --- a/tests/core/framework/block/composite_block_manager_test.cpp +++ b/tests/core/framework/block/composite_block_manager_test.cpp @@ -84,7 +84,10 @@ Sequence MakeTestSequence(size_t index, StoppingChecker stopping_checker; stopping_checker.set_max_generated_tokens(256); SequenceParams seq_params; - seq_params.seq_capacity = 5000; + // Large enough to hold DSV4-scale prompts (>= a full C128 block = 16384 + // tokens). Individual tests can still use short prompts; sequence capacity + // just has to bound `num_tokens + max_generated_tokens`. + seq_params.seq_capacity = 65536; seq_params.stopping_checker = &stopping_checker; seq_params.sampling_param = &sampling_param; seq_params.skip_special_tokens = true; @@ -528,4 +531,203 @@ TEST(CompositeBlockManagerTest, CapacityStatsUseFinestAdmissionLeaf) { manager.deallocate_for_sequence(&seq); } +// DSV4 prefix cache: a fresh sequence with the same prompt as a previously +// released sequence should mount shared blocks from all three prefix-cache +// leaves (SWA / C4 / C128). The composite takes the cross-leaf min hit length +// clamped to a C128 block, so the prompt has to span at least one C128 block +// (128 * base = 16384 tokens) for the hit to be non-trivial. Uses a 2*C128 +// prompt so we get a meaningful hit even after the exact-repeat pop. +TEST(CompositeBlockManagerTest, Dsv4PrefixCacheHitOnRepeatedPrefix) { + const uint32_t base_num_blocks = 4096; + // Sliding window covers a couple C128 blocks worth of tokens so the SWA + // tail-continuity check has room to succeed even after the exact-repeat pop. + const uint32_t window_size = 4 * kBaseBlockSize; + const uint32_t max_seqs_per_batch = 4; + BlockManager::Options opts = MakeCompositeOptions( + base_num_blocks, kBaseBlockSize, window_size, max_seqs_per_batch); + // max_tokens_per_batch has to accommodate the prompt so allocate_sequence + // does not exceed the SWA burst budget. + opts.max_tokens_per_batch(3 * kBlockSizeRatio128); + ASSERT_TRUE(opts.enable_prefix_cache()); + CompositeBlockManager manager(build_composite_leaves(opts)); + + // First sequence: a 2*C128 prompt (32768 tokens) so both C4 and C128 have + // multiple full blocks worth of cacheable prefix. Mark all tokens as + // forwarded so the pre-grow hook (fired during the NEXT allocate) has data + // to insert -- and we drive one extra allocate to trigger it explicitly. + const size_t num_tokens = 2 * kBlockSizeRatio128; + const std::vector prompt(num_tokens, 7); + Sequence seq_first = MakeTestSequence(0, prompt); + ASSERT_TRUE(manager.allocate_sequence(&seq_first, num_tokens)); + seq_first.kv_state().incr_kv_cache_tokens_num(num_tokens); + // deallocate_for_sequence runs the pre-grow hook once more (via + // cache_for_sequence's fan-out), flushing the residual full blocks. + manager.deallocate_for_sequence(&seq_first); + seq_first.reset(); + + // The prefix cache under SWA / C4 / C128 should now hold the prompt's full + // blocks. On a second sequence with the same prompt, admission mounts from + // all three leaves and pins kv_cache_tokens_num to the safe min hit. + Sequence seq_hit = MakeTestSequence(1, prompt); + manager.allocate_shared_for_sequence(&seq_hit); + EXPECT_GT(seq_hit.kv_state().shared_blocks_num(BlockType::SWA), 0u); + EXPECT_GT(seq_hit.kv_state().shared_blocks_num(BlockType::C4), 0u); + EXPECT_GT(seq_hit.kv_state().shared_blocks_num(BlockType::C128), 0u); + const size_t kv_tokens = seq_hit.kv_state().kv_cache_tokens_num(); + EXPECT_GT(kv_tokens, 0u); + EXPECT_LT(kv_tokens, num_tokens); // exact-repeat pop kept at least one c128 + EXPECT_EQ(kv_tokens % kBlockSizeRatio128, 0u) + << "safe_hit_tokens must be a multiple of the C128 block stride"; + // SWA vector length matches safe_hit / base and the tail carries valid + // blocks. + const std::vector hit_swa = SwaBlocks(seq_hit); + EXPECT_EQ(hit_swa.size(), kv_tokens / kBaseBlockSize); + EXPECT_TRUE(hit_swa.back().is_valid()); + + manager.deallocate_for_sequence(&seq_hit); +} + +// DSV4 prefix cache: a fresh sequence with a completely different prompt should +// miss cleanly (no shared mount, kv_cache_tokens_num unchanged). +TEST(CompositeBlockManagerTest, Dsv4PrefixCacheMissCleanly) { + const uint32_t base_num_blocks = 4096; + const uint32_t window_size = 4 * kBaseBlockSize; + const uint32_t max_seqs_per_batch = 4; + BlockManager::Options opts = MakeCompositeOptions( + base_num_blocks, kBaseBlockSize, window_size, max_seqs_per_batch); + opts.max_tokens_per_batch(3 * kBlockSizeRatio128); + CompositeBlockManager manager(build_composite_leaves(opts)); + + const size_t num_tokens = 2 * kBlockSizeRatio128; + const std::vector prompt_a(num_tokens, 7); + Sequence seq_a = MakeTestSequence(0, prompt_a); + ASSERT_TRUE(manager.allocate_sequence(&seq_a, num_tokens)); + seq_a.kv_state().incr_kv_cache_tokens_num(num_tokens); + manager.deallocate_for_sequence(&seq_a); + seq_a.reset(); + + // A sequence with a totally different prompt should not share any block. + std::vector prompt_b(num_tokens, 0); + for (size_t i = 0; i < prompt_b.size(); ++i) { + prompt_b[i] = static_cast(i + 100); + } + Sequence seq_b = MakeTestSequence(1, prompt_b); + manager.allocate_shared_for_sequence(&seq_b); + EXPECT_EQ(seq_b.kv_state().shared_blocks_num(BlockType::SWA), 0u); + EXPECT_EQ(seq_b.kv_state().shared_blocks_num(BlockType::C4), 0u); + EXPECT_EQ(seq_b.kv_state().shared_blocks_num(BlockType::C128), 0u); + EXPECT_EQ(seq_b.kv_state().kv_cache_tokens_num(), 0u); + + manager.deallocate_for_sequence(&seq_b); +} + +// SWA slid-out blocks should enter the prefix cache via the pre-grow hook (v2b: +// hook fires at every allocate_sequence, insertion is incremental and cached +// blocks survive slid-out release through the ref<=2u path). A follow-up +// sequence with the same prompt should hit those cached blocks. +TEST(CompositeBlockManagerTest, SlidingWindowSlidOutBlocksEnterPrefixCache) { + const uint32_t base_num_blocks = 4096; + // Window covers a small number of base blocks so slide-out happens quickly. + const uint32_t sliding_window_blocks_per_sequence = 3; + const uint32_t window_size = + sliding_window_blocks_per_sequence * kBaseBlockSize; + const uint32_t max_seqs_per_batch = 4; + BlockManager::Options opts = MakeCompositeOptions( + base_num_blocks, kBaseBlockSize, window_size, max_seqs_per_batch); + opts.max_tokens_per_batch(3 * kBlockSizeRatio128); + ASSERT_TRUE(opts.enable_prefix_cache()); + CompositeBlockManager manager(build_composite_leaves(opts)); + + // Prompt spans two full C128 blocks so the min-across-leaves gate can + // survive AND the exact-repeat pop (one c128 stride) still leaves shared + // blocks behind. A single-c128 prompt would end up entirely popped. + const size_t num_tokens = 2 * kBlockSizeRatio128; + const std::vector prompt(num_tokens, 7); + Sequence seq_first = MakeTestSequence(0, prompt); + ASSERT_TRUE(manager.allocate_sequence(&seq_first, num_tokens)); + seq_first.kv_state().incr_kv_cache_tokens_num(num_tokens); + // The SWA leaf should have released everything but the last window-sized + // slab back to the pool (as invalid placeholders), and those blocks landed + // in the prefix cache via the pre-grow hook before release. + manager.deallocate_for_sequence(&seq_first); + seq_first.reset(); + + // A second sequence with the same prompt should hit the cached SWA tail. + Sequence seq_hit = MakeTestSequence(1, prompt); + manager.allocate_shared_for_sequence(&seq_hit); + EXPECT_GT(seq_hit.kv_state().shared_blocks_num(BlockType::SWA), 0u); + EXPECT_GT(seq_hit.kv_state().shared_blocks_num(BlockType::C4), 0u); + EXPECT_GT(seq_hit.kv_state().shared_blocks_num(BlockType::C128), 0u); + + manager.deallocate_for_sequence(&seq_hit); +} + +// v2b: pre-grow hook advances KVCacheState::num_cached_blocks incrementally +// across chunked-prefill steps. First chunk: cursor at 0. Second chunk (after +// forwarding chunk 1's tokens): cursor at chunk-1's full-block count. +TEST(CompositeBlockManagerTest, Dsv4PrefixCachePreHookCursorAdvances) { + const uint32_t base_num_blocks = 4096; + const uint32_t window_size = 4 * kBaseBlockSize; + const uint32_t max_seqs_per_batch = 4; + BlockManager::Options opts = MakeCompositeOptions( + base_num_blocks, kBaseBlockSize, window_size, max_seqs_per_batch); + opts.max_tokens_per_batch(4 * kBlockSizeRatio128); + CompositeBlockManager manager(build_composite_leaves(opts)); + + // Two-chunk prompt (2*C128). Chunk 1 is a single C128 block wide. + const size_t chunk = kBlockSizeRatio128; + const size_t total = 2 * chunk; + Sequence seq = MakeTestSequence(0, std::vector(total, 7)); + + // Chunk 1: allocate + forward. + ASSERT_TRUE(manager.allocate_sequence(&seq, chunk)); + EXPECT_EQ(seq.kv_state().num_cached_blocks(BlockType::SWA), 0u) + << "pre-grow hook has nothing to cache on the very first allocate: kv=0"; + EXPECT_EQ(seq.kv_state().num_cached_blocks(BlockType::C4), 0u); + EXPECT_EQ(seq.kv_state().num_cached_blocks(BlockType::C128), 0u); + seq.kv_state().incr_kv_cache_tokens_num(chunk); + + // Chunk 2 allocation triggers the pre-grow hook, which now sees kv=chunk + // worth of forwarded tokens and inserts them. + ASSERT_TRUE(manager.allocate_sequence(&seq, total)); + EXPECT_EQ(seq.kv_state().num_cached_blocks(BlockType::SWA), + chunk / kBaseBlockSize); + EXPECT_EQ(seq.kv_state().num_cached_blocks(BlockType::C4), + chunk / kBlockSizeRatio4); + EXPECT_EQ(seq.kv_state().num_cached_blocks(BlockType::C128), + chunk / kBlockSizeRatio128); + + manager.deallocate_for_sequence(&seq); +} + +// v2b: exact-repeat safe-hit hits the pop path -- when a fresh sequence has a +// prompt whose entire length is cached, safe_hit_tokens is popped by one +// c128 block worth of tokens so the forward has data to process. The test +// verifies kv_cache_tokens_num after mount is strictly less than the prompt. +TEST(CompositeBlockManagerTest, Dsv4PrefixCacheExactRepeatPopsOneC128) { + const uint32_t base_num_blocks = 4096; + const uint32_t window_size = 4 * kBaseBlockSize; + const uint32_t max_seqs_per_batch = 4; + BlockManager::Options opts = MakeCompositeOptions( + base_num_blocks, kBaseBlockSize, window_size, max_seqs_per_batch); + opts.max_tokens_per_batch(4 * kBlockSizeRatio128); + CompositeBlockManager manager(build_composite_leaves(opts)); + + const size_t num_tokens = 3 * kBlockSizeRatio128; + const std::vector prompt(num_tokens, 42); + Sequence seq_first = MakeTestSequence(0, prompt); + ASSERT_TRUE(manager.allocate_sequence(&seq_first, num_tokens)); + seq_first.kv_state().incr_kv_cache_tokens_num(num_tokens); + manager.deallocate_for_sequence(&seq_first); + seq_first.reset(); + + Sequence seq_hit = MakeTestSequence(1, prompt); + manager.allocate_shared_for_sequence(&seq_hit); + // Full-length hit → pop one c128 block. Expect kv = num_tokens - c128. + EXPECT_EQ(seq_hit.kv_state().kv_cache_tokens_num(), + num_tokens - kBlockSizeRatio128); + + manager.deallocate_for_sequence(&seq_hit); +} + } // namespace xllm diff --git a/xllm/core/distributed_runtime/master.cpp b/xllm/core/distributed_runtime/master.cpp index b1743ddde2..f43be7f139 100644 --- a/xllm/core/distributed_runtime/master.cpp +++ b/xllm/core/distributed_runtime/master.cpp @@ -219,17 +219,6 @@ Master::Master(const Options& options, EngineType type) const std::optional cp_error = validate_model_cp(options_, type, cp_model_type, global_world_size); CHECK(!cp_error.has_value()) << cp_error.value(); - if (options_.enable_prefix_cache() && options_.backend() == "llm") { - const std::string model_type = util::get_model_type(model_path); - if (util::is_deepseek_v4_model_type(model_type)) { - LOG(WARNING) << model_type - << " does not support prefix cache with " - "CompositeBlockManager yet, fallback to " - "enable_prefix_cache=false"; - options_.enable_prefix_cache(false); - KVCacheConfig::get_instance().enable_prefix_cache(false); - } - } options_.enable_mla(util::should_enable_mla(model_path, options_.backend())); print_startup_banner(model_path, options_.backend(), options_.node_rank()); LOG(INFO) << "Master init options: " << options_.to_string(); diff --git a/xllm/core/framework/block/block_manager.h b/xllm/core/framework/block/block_manager.h index e73deaaa8a..01b8dee55e 100644 --- a/xllm/core/framework/block/block_manager.h +++ b/xllm/core/framework/block/block_manager.h @@ -102,17 +102,27 @@ class BlockManager { virtual std::vector allocate(size_t num_blocks) = 0; - // `matched_tokens`, when non-null, receives the matched prefix length in - // TOKENS (forwarded straight to PrefixCache::match: the flat-KV leaf writes - // blocks.size() * block_size; the LINEAR leaf writes the recoverable - // checkpoint prefix in its own chunk-strided hash domain). Callers that only - // want the blocks leave it null. + // Returns the shared-prefix vector for the caller's Sequence. The vector's + // shape is leaf-defined: + // - KV / C4 / C128 return a solid prefix `[valid, valid, ..., valid]`; + // length in blocks × block_size() is the matched-token count. + // - SWA returns a gap-tolerant `[opt_valid, ..., valid_last]` where the + // length equals last-hit-index + 1 in base blocks (attention only reads + // the last swa_blocks_per_seq base blocks; the composite enforces the + // tail-continuity check). + // - LINEAR (constraint leaf) returns `[inv, inv, ..., deepest_valid]` at + // chunk-stride granularity; the deepest slot is the class-A checkpoint + // restore source, and length in chunks × block_size() is the recoverable + // prefix length. The composite pulls that block out and never mounts + // LINEAR blocks into the sequence's live LINEAR vector. + // + // Matched-token count is `returned.size() * block_size()`; callers derive + // it directly, so no separate out-param is needed. virtual std::vector allocate_shared( const Slice& token_ids, const Slice& existed_shared_blocks = {}, const MMData& mm_data = MMData(), - const Slice& block_hashes = {}, - size_t* matched_tokens = nullptr) = 0; + const Slice& block_hashes = {}) = 0; virtual void cache(const Slice& token_ids, std::vector& blocks, diff --git a/xllm/core/framework/block/block_manager_impl.cpp b/xllm/core/framework/block/block_manager_impl.cpp index 090bc84934..962298ff0f 100644 --- a/xllm/core/framework/block/block_manager_impl.cpp +++ b/xllm/core/framework/block/block_manager_impl.cpp @@ -22,46 +22,38 @@ limitations under the License. namespace xllm { namespace { -bool mark_used(std::vector* usage_ids, int32_t block_id) { +bool clear_used(std::vector* usage_ids, int32_t block_id) { CHECK(usage_ids != nullptr); CHECK_GE(block_id, 0); CHECK_LT(static_cast(block_id), usage_ids->size()); - if ((*usage_ids)[block_id] != 0) { + if ((*usage_ids)[block_id] == 0) { return false; } - (*usage_ids)[block_id] = 1; + (*usage_ids)[block_id] = 0; return true; } -bool clear_used(std::vector* usage_ids, int32_t block_id) { +} // namespace + +bool BlockManagerImpl::mark_used(std::vector* usage_ids, + int32_t block_id) { CHECK(usage_ids != nullptr); CHECK_GE(block_id, 0); CHECK_LT(static_cast(block_id), usage_ids->size()); - if ((*usage_ids)[block_id] == 0) { + if ((*usage_ids)[block_id] != 0) { return false; } - (*usage_ids)[block_id] = 0; + (*usage_ids)[block_id] = 1; return true; } -} // namespace - BlockManagerImpl::BlockManagerImpl(const Options& options) : BlockManager(options) { CHECK_GT(options.num_blocks(), 0) << "No blocks to allocate"; CHECK_GT(options.block_size(), 0) << "Block size must be positive"; if (options_.enable_prefix_cache()) { PrefixCache::Options prefix_cache_options; - // A prefix cache indexes on a single block-boundary stride. For KV that is - // the KV block size; for the linear-state checkpoint index it is the - // prefill chunk stride. Feed the right one into the cache's single - // block_size slot so the LINEAR cache reuses the base by-block machinery - // verbatim. - const int32_t prefix_cache_block_size = - options.block_type() == BlockType::LINEAR - ? options.linear_chunk_stride() - : options.block_size(); - prefix_cache_options.block_size(prefix_cache_block_size) + prefix_cache_options.block_size(options.block_size()) .hasher_type(options.hasher_type()) .block_type(options.block_type()); prefix_cache_ = create_prefix_cache(prefix_cache_options); @@ -169,20 +161,16 @@ std::vector BlockManagerImpl::allocate_shared( const Slice& token_ids, const Slice& existed_shared_blocks, const MMData& mm_data, - const Slice& block_hashes, - size_t* matched_tokens) { - // only allocate shared blocks for prefill sequences + const Slice& block_hashes) { + // only allocate shared blocks for prefix caching prefill sequences if (options_.enable_prefix_cache()) { AUTO_COUNTER(prefix_cache_latency_seconds_match); - size_t prefix_length = 0; - std::vector shared_blocks = - prefix_cache_->match(token_ids, - existed_shared_blocks, - mm_data, - block_hashes, - &prefix_length); + std::vector shared_blocks = prefix_cache_->match( + token_ids, existed_shared_blocks, mm_data, block_hashes); + const size_t prefix_length = + shared_blocks.size() * static_cast(block_size_); COUNTER_ADD(prefix_cache_match_length_total, prefix_length); VLOG(1) << "Prefix cache matched " << shared_blocks.size() << " blocks, prefix_length=" << prefix_length; @@ -193,9 +181,6 @@ std::vector BlockManagerImpl::allocate_shared( num_used_blocks_.fetch_add(1, std::memory_order_relaxed); } } - if (matched_tokens != nullptr) { - *matched_tokens = prefix_length; - } return shared_blocks; } return {}; diff --git a/xllm/core/framework/block/block_manager_impl.h b/xllm/core/framework/block/block_manager_impl.h index ad39e1cc3c..3615d25d11 100644 --- a/xllm/core/framework/block/block_manager_impl.h +++ b/xllm/core/framework/block/block_manager_impl.h @@ -47,8 +47,7 @@ class BlockManagerImpl : public BlockManager { const Slice& token_ids, const Slice& existed_shared_blocks = {}, const MMData& mm_data = MMData(), - const Slice& block_hashes = {}, - size_t* matched_tokens = nullptr) override; + const Slice& block_hashes = {}) override; // cache blocks when enable prefix cache void cache(const Slice& token_ids, @@ -95,6 +94,14 @@ class BlockManagerImpl : public BlockManager { // total blocks num size_t num_total_blocks() const override { return free_blocks_.size() - 1; } + protected: + // Flip a block's entry in `usage_ids` from 0 to 1. Returns true if the flip + // happened; false if the entry was already 1 (i.e. block was already marked + // used). Shared with subclasses (e.g. SlidingWindowBlockManager) that need + // to reproduce the base allocate_shared refcount bookkeeping over their own + // custom probe path. Static-friendly signature keeps callers free of `this`. + static bool mark_used(std::vector* usage_ids, int32_t block_id); + private: // check if has enough slots, if not, try to evict some blocks // from the prefix cache diff --git a/xllm/core/framework/block/composite_block_manager.cpp b/xllm/core/framework/block/composite_block_manager.cpp index 384f9b717c..c78da4ebb7 100644 --- a/xllm/core/framework/block/composite_block_manager.cpp +++ b/xllm/core/framework/block/composite_block_manager.cpp @@ -17,6 +17,7 @@ limitations under the License. #include #include +#include #include #include "block_manager_impl.h" @@ -39,9 +40,8 @@ uint32_t ceil_div(uint32_t numerator, uint32_t denominator) { return (numerator + denominator - 1) / denominator; } -// Wrap a leaf in a concurrency adapter when sequence-level entry points may run -// off the scheduler thread (disagg PD / kvcache store prefill threadpools, or -// the host-offload D2H completion callback). +// Wrap the leaf in a concurrency adapter when sequence-level calls may run +// off the scheduler thread (disagg PD / kvcache store / host-offload). std::unique_ptr maybe_concurrent( std::unique_ptr leaf, const BlockManager::Options& options) { @@ -52,8 +52,8 @@ std::unique_ptr maybe_concurrent( return leaf; } -// Build the KV leaf: an xtensor VMM manager when enable_xtensor, otherwise the -// flat free-list BlockManagerImpl. xtensor does not support prefix cache. +// Xtensor VMM manager or flat free-list BlockManagerImpl. Xtensor has no +// prefix cache. std::unique_ptr make_kv_leaf(const BlockManager::Options& kv_opts, int32_t dp_rank) { if (!kv_opts.enable_xtensor()) { @@ -65,7 +65,7 @@ std::unique_ptr make_kv_leaf(const BlockManager::Options& kv_opts, << "slot_size must be set when enable_xtensor is true"; const size_t page_size = ::xllm::KVCacheConfig::get_instance().phy_page_granularity_size(); - // K and V are the same size in the current implementation, so divide by 2. + // K and V share block size; divide by 2. const size_t block_mem_size = static_cast(kv_opts.block_size()) * kv_opts.slot_size() / 2; return std::make_unique(kv_opts, @@ -83,17 +83,12 @@ std::map build_composite_leaves( int32_t dp_rank) { std::map leaves; - // Per-sequence resource leaf, additive on top of the KV family built below: - // a GDN model holds both KV (full-attention layers) and LINEAR (linear - // layers), so LINEAR is never an alternative KV shape. leaves is keyed by - // BlockType, so insertion order is irrelevant; emplacing here keeps the - // early-returning KV-family branches untouched. participates_in_admission - // stays false -- a true here would let capacity_leaf() pick this - // block_size==1 leaf and misreport pool capacity as the linear-slot pool. - // LINEAR is scheduler-thread only, so unlike the SINGLE leaf it is - // deliberately NOT wrapped in ConcurrentBlockManagerImpl. - // TODO(refactor): fold the SINGLE leaf (still emplaced by BlockManagerPool) - // in here too, once num_single_blocks rides on BlockManager::Options. + // LINEAR resource leaf (Qwen3.5-Next GDN). Additive on top of the KV + // family: a GDN model holds both KV and LINEAR. Not an admission leaf + // (block_size==1 would misreport pool capacity). Scheduler-thread only, so + // no ConcurrentBlockManagerImpl wrap. supports_prefix_cache=true so + // probe_prefix_leaves picks it up; the FLAT_KV_LINEAR trimmer pulls out + // the deepest checkpoint as a restore source. if (options.enable_linear_state()) { CHECK_GT(options.linear_state_num_slots(), 0) << "linear_state_num_slots must be set when linear state is enabled"; @@ -102,15 +97,14 @@ std::map build_composite_leaves( CompositeBlockManager::LeafEntry{ std::make_unique( static_cast(options.linear_state_num_slots()), - options.block_size(), options.linear_chunk_stride()), /*participates_in_admission=*/false, - /*supports_prefix_cache=*/false}); + /*supports_prefix_cache=*/options.enable_prefix_cache()}); } if (options.manager_types().empty()) { - // Normal / Qwen / xtensor model: a single KV leaf. It is the admission - // source; prefix cache is on only for the flat (non-xtensor) KV leaf. + // Normal / Qwen / xtensor: a single KV leaf. Prefix cache is on only for + // the flat (non-xtensor) KV leaf. BlockManager::Options kv_opts = options; kv_opts.block_type(BlockType::KV); leaves.emplace( @@ -123,8 +117,7 @@ std::map build_composite_leaves( return leaves; } - // DSV4: SWA + compressed (C4 / C128) leaves, derived from manager_types / - // compress_ratios (kept identical to the previous in-composite construction). + // DSV4: SWA + compressed (C4 / C128) leaves. const size_t n = options.manager_types().size(); CHECK_EQ(n, options.compress_ratios().size()) << "manager_types and compress_ratios must have the same size"; @@ -145,14 +138,19 @@ std::map build_composite_leaves( << " for composite BlockManagerImpl sub-manager"; const BlockType key = compress_ratio == 4 ? BlockType::C4 : BlockType::C128; - opts.block_type(key); - // Compressed groups are admission sources but do not serve prefix cache. - leaves.emplace(key, - CompositeBlockManager::LeafEntry{ - maybe_concurrent( - std::make_unique(opts), options), - /*participates_in_admission=*/true, - /*supports_prefix_cache=*/false}); + opts.block_type(key) + .enable_prefix_cache(options.enable_prefix_cache()) + .hasher_type(options.hasher_type()); + // C4/C128: full-history attention (no window), so hits must form a + // solid left-to-right prefix. The default BlockManagerImpl probe is + // exactly right. Cross-leaf min happens in the composite. + leaves.emplace( + key, + CompositeBlockManager::LeafEntry{ + maybe_concurrent(std::make_unique(opts), + options), + /*participates_in_admission=*/true, + /*supports_prefix_cache=*/options.enable_prefix_cache()}); } else if (type == kManagerTypeSlidingWindowBlockManager) { const uint32_t swa_blocks_per_seq = options.swa_blocks_per_seq(); CHECK_GT(swa_blocks_per_seq, 0u) << "swa_blocks_per_seq must be positive"; @@ -163,20 +161,24 @@ std::map build_composite_leaves( const uint32_t burst_blocks = ceil_div(std::max(options.max_tokens_per_batch(), 1u), static_cast(options.block_size())); + // Slack fits the peak "old blocks not yet released + new tail". const uint32_t swa_total_blocks = swa_blocks_per_seq * max_seqs + burst_blocks + max_seqs + 2; opts.num_blocks(swa_total_blocks) .swa_blocks_per_seq(swa_blocks_per_seq) .sliding_window_size(sliding_window_size) - .block_type(BlockType::SWA); - // SWA neither participates in admission nor serves prefix cache. + .block_type(BlockType::SWA) + .enable_prefix_cache(options.enable_prefix_cache()) + .hasher_type(options.hasher_type()); + // SWA is not an admission leaf (pool sized by ring, not token budget) + // but serves gap-tolerant prefix cache. leaves.emplace( BlockType::SWA, CompositeBlockManager::LeafEntry{ maybe_concurrent( std::make_unique(opts), options), /*participates_in_admission=*/false, - /*supports_prefix_cache=*/false}); + /*supports_prefix_cache=*/options.enable_prefix_cache()}); } else { LOG(FATAL) << "Unknown manager_type " << type; } @@ -186,15 +188,80 @@ std::map build_composite_leaves( CompositeBlockManager::CompositeBlockManager( std::map leaves) - : BlockManager(BlockManager::Options()), leaves_(std::move(leaves)) { + : BlockManager(BlockManager::Options()), + leaves_(std::move(leaves)), + combination_(classify_leaf_combination(leaves_)) { CHECK(!leaves_.empty()) << "CompositeBlockManager requires at least one leaf"; } +CompositeBlockManager::LeafCombination +CompositeBlockManager::classify_leaf_combination( + const std::map& leaves) { + auto has_prefix = [&](BlockType t) { + const auto it = leaves.find(t); + return it != leaves.end() && it->second.supports_prefix_cache; + }; + const bool has_kv = has_prefix(BlockType::KV); + const bool has_swa = has_prefix(BlockType::SWA); + const bool has_c4 = has_prefix(BlockType::C4); + const bool has_c128 = has_prefix(BlockType::C128); + const bool has_linear = has_prefix(BlockType::LINEAR); + + if (has_swa || has_c4 || has_c128) { + CHECK(has_swa && has_c4 && has_c128) + << "SWA_COMPRESSED requires all of {SWA, C4, C128}; got SWA=" << has_swa + << " C4=" << has_c4 << " C128=" << has_c128; + CHECK(!has_kv) << "SWA_COMPRESSED must not carry a KV leaf"; + CHECK(!has_linear) << "SWA_COMPRESSED must not carry a LINEAR leaf"; + return LeafCombination::SWA_COMPRESSED; + } + if (has_kv) { + return has_linear ? LeafCombination::FLAT_KV_LINEAR + : LeafCombination::FLAT_KV; + } + return LeafCombination::UNSUPPORTED; +} + BlockManager* CompositeBlockManager::leaf_of(BlockType type) const { auto it = leaves_.find(type); return it == leaves_.end() ? nullptr : it->second.leaf.get(); } +void CompositeBlockManager::cache_full_blocks_for_sequence(Sequence* seq) { + if (seq == nullptr) { + return; + } + KVCacheState& kv = seq->kv_state(); + for (auto& [type, entry] : leaves_) { + // KV owns its final flush via cache_for_sequence at deallocate time; + // SINGLE / LINEAR hold no token cache. + if (type == BlockType::KV || type == BlockType::SINGLE || + type == BlockType::LINEAR) { + continue; + } + if (!entry.supports_prefix_cache) { + continue; + } + BlockManager& leaf = *entry.leaf; + const size_t block_size = leaf.block_size(); + if (block_size == 0) { + continue; + } + std::vector* blocks = kv.mutable_blocks(type); + if (blocks == nullptr || blocks->empty()) { + continue; + } + const size_t num_full = kv.kv_cache_tokens_num() / block_size; + const size_t cached = kv.num_cached_blocks(type); + const size_t end = std::min(num_full, blocks->size()); + if (end <= cached) { + continue; + } + leaf.cache(seq->tokens(), *blocks, cached, seq->mm_data()); + kv.set_num_cached_blocks(type, end); + } +} + bool CompositeBlockManager::allocate_sequence(Sequence* seq, size_t num_tokens) { if (seq == nullptr) { @@ -202,11 +269,15 @@ bool CompositeBlockManager::allocate_sequence(Sequence* seq, } KVCacheState& kv_state = seq->kv_state(); - // Each leaf grows its own block_type() blocks for the sequence and returns - // the blocks it newly allocated; it does not insert them. The composite - // stages them keyed by BlockType and commits only after every leaf succeeds. - // Rollback = deallocate every staged run (the sequence is never grown before - // commit, so existing blocks stay intact). + // Pre-grow cache hook: give each leaf a chance to insert its already- + // forwarded full blocks into the prefix cache BEFORE this round's grow. + // Cursor lives in KVCacheState::num_cached_blocks_ so repeated calls only + // pay for the delta. + cache_full_blocks_for_sequence(seq); + + // Fan out growth. Each leaf returns its newly allocated blocks (or nullopt + // on failure). Stage keyed by BlockType; commit only after every leaf + // succeeds so a failure rolls back cleanly. std::vector>> staged; for (auto& [type, entry] : leaves_) { @@ -223,17 +294,12 @@ bool CompositeBlockManager::allocate_sequence(Sequence* seq, } } - // Post-condition (grow-or-fail): every cache-bearing leaf must now hold - // enough blocks to cover num_tokens. A leaf that returns an empty vector - // means "no growth needed", so we must verify that was true. SINGLE / LINEAR - // leaves are per-sequence resource slots (one block per sequence), not - // token-cache, and are exempt. Without this check, a leaf that mistakenly - // returned an empty vector under pool pressure (instead of nullopt) would - // let the pool report admission success while the device is under-provisioned - // for num_tokens -- batch_input_builder.cpp:589 CHECK then fires downstream - // with `current_max_tokens_capacity() >= seq_len`. Rolling back here lets the - // scheduler defer the sequence to the next tick, at which point pool state - // may have recovered. + // Grow-or-fail: every cache-bearing leaf must cover num_tokens now. + // Without this a leaf that mistakenly returned an empty vector under + // pressure would let the pool report admission success while the device + // is under-provisioned -- batch_input_builder's CHECK then fires + // downstream. Rolling back lets the scheduler defer to the next tick. + // SINGLE / LINEAR are per-sequence resource slots, not token cache. for (const auto& [type, entry] : leaves_) { if (type == BlockType::SINGLE || type == BlockType::LINEAR) { continue; @@ -259,13 +325,11 @@ bool CompositeBlockManager::allocate_sequence(Sequence* seq, } } - // Commit: append staged blocks to the sequence under each leaf's block type. + // Commit staged blocks, then release slid-out (SWA only; other leaves + // no-op). Post-commit only, so failures never touch existing blocks. for (auto& [type, blocks] : staged) { kv_state.add_blocks(type, blocks); } - // Post-commit: release blocks that have slid out of the window (SWA only; - // other leaves are a no-op). Runs only after a fully successful round, so a - // failed allocation never releases the sequence's existing blocks. for (auto& [type, entry] : leaves_) { entry.leaf->release_out_of_window(seq); } @@ -276,134 +340,365 @@ void CompositeBlockManager::deallocate_for_sequence(Sequence* seq) { if (seq == nullptr) { return; } - // Publish prefix cache first (prefix leaves only), then release every leaf's - // blocks by key. seq->reset() stays with the pool caller. + // Publish prefix cache first, then release blocks. seq->reset() belongs to + // the pool caller. cache_for_sequence(seq); for (auto& [type, entry] : leaves_) { entry.leaf->deallocate(seq->kv_state().blocks(type)); } } +namespace { + +// One leaf's allocate_shared() result carried through trim / mount. Vector +// shape is leaf-defined (see BlockManager::allocate_shared doc); reach in +// tokens is `blocks.size() * block_size`. +struct ProbeResult { + BlockType type; + BlockManager* leaf = nullptr; + std::vector blocks; + size_t block_size = 0; +}; + +// Trim outcome. `to_mount` feeds the shape's mount step; `to_drop` goes to +// leaf->deallocate() (cache release, not physical free). +// `linear_restore_src`, when set, is the LINEAR checkpoint that the caller +// stashes via Sequence::set_linear_restore_src_block. +struct TrimOutcome { + std::vector to_mount; + std::vector to_drop; + size_t safe_hit_tokens = 0; + std::optional linear_restore_src; +}; + +// Probe every leaf whose supports_prefix_cache is true. LINEAR probes with +// its chunk-strided hash chain (linear_state_hashes) and an empty existed +// slice; other leaves use the by-block chain (block_hashes) and the +// sequence's already-shared prefix for try_replace semantics. +std::vector probe_prefix_leaves( + Sequence* seq, + const std::map& leaves) { + std::vector probes; + probes.reserve(leaves.size()); + for (const auto& [type, entry] : leaves) { + if (!entry.supports_prefix_cache) { + continue; + } + BlockManager& leaf = *entry.leaf; + Slice hashes; + Slice existed; + if (type == BlockType::LINEAR) { + seq->update_linear_state_hashes(static_cast(leaf.block_size())); + hashes = seq->linear_state_hashes(); + } else { + seq->update_block_hashes(static_cast(leaf.block_size()), + leaf.options().hasher_type()); + hashes = seq->block_hashes(); + KVCacheState& kv_state = seq->kv_state(); + existed = + kv_state.blocks(type).slice(0, kv_state.shared_blocks_num(type)); + } + std::vector blocks = + leaf.allocate_shared(seq->tokens(), existed, seq->mm_data(), hashes); + probes.push_back({/*type=*/type, + /*leaf=*/&leaf, + /*blocks=*/std::move(blocks), + /*block_size=*/leaf.block_size()}); + } + return probes; +} + +std::optional take_probe(std::vector& probes, + BlockType type) { + for (auto it = probes.begin(); it != probes.end(); ++it) { + if (it->type == type) { + ProbeResult r = std::move(*it); + probes.erase(it); + return r; + } + } + return std::nullopt; +} + +// FLAT_KV: no trim; Sequence::add_shared_blocks owns replace + exact-repeat. +TrimOutcome trim_flat_kv(std::vector probes) { + CHECK_EQ(probes.size(), 1u) << "FLAT_KV expects a single KV probe"; + TrimOutcome out; + out.to_mount = std::move(probes); + return out; +} + +// FLAT_KV_LINEAR: extract LINEAR's deepest checkpoint as a restore source +// (its vector is `[inv, ..., deepest_valid]` at chunk stride, so reach in +// tokens is `blocks.size() * block_size`), then clamp KV shared blocks to +// that recoverable budget. +TrimOutcome trim_flat_kv_linear(std::vector probes) { + auto linear = take_probe(probes, BlockType::LINEAR); + CHECK_EQ(probes.size(), 1u) + << "FLAT_KV_LINEAR expects one KV probe (LINEAR removed above)"; + TrimOutcome out; + + size_t linear_recoverable_tokens = 0; + if (linear.has_value()) { + linear_recoverable_tokens = linear->blocks.size() * linear->block_size; + for (auto it = linear->blocks.rbegin(); it != linear->blocks.rend(); ++it) { + if (it->is_valid()) { + out.linear_restore_src = std::move(*it); + break; + } + } + } + + ProbeResult& kv = probes.front(); + const size_t kv_block_size = kv.block_size; + const size_t recoverable_blocks = + kv_block_size == 0 ? 0 : linear_recoverable_tokens / kv_block_size; + const size_t safe_count = std::min(kv.blocks.size(), recoverable_blocks); + if (safe_count < kv.blocks.size()) { + ProbeResult drop = {kv.type, + kv.leaf, + /*blocks=*/{}, + /*block_size=*/kv.block_size}; + drop.blocks.reserve(kv.blocks.size() - safe_count); + for (size_t i = safe_count; i < kv.blocks.size(); ++i) { + drop.blocks.emplace_back(std::move(kv.blocks[i])); + } + kv.blocks.resize(safe_count); + out.to_drop.emplace_back(std::move(drop)); + } + out.safe_hit_tokens = safe_count * kv_block_size; + out.to_mount = std::move(probes); + return out; +} + +// SWA_COMPRESSED: cross-leaf min -> C128-stride clamp -> SWA tail-continuity +// (fallback in C128 steps) -> exact-repeat pop -> per-leaf trim. +TrimOutcome trim_swa_compressed(std::vector probes, + size_t prompt_tokens) { + TrimOutcome out; + size_t c128_block_size = 0; + for (const auto& p : probes) { + if (p.type == BlockType::C128) { + c128_block_size = p.block_size; + } + } + CHECK_GT(c128_block_size, 0u) + << "SWA_COMPRESSED trim requires a C128 leaf with non-zero block_size"; + + size_t safe_hit_tokens = std::numeric_limits::max(); + for (const auto& p : probes) { + safe_hit_tokens = std::min(safe_hit_tokens, p.blocks.size() * p.block_size); + } + safe_hit_tokens = (safe_hit_tokens / c128_block_size) * c128_block_size; + + // SWA attention reads the last `swa_blocks_per_seq` base blocks; a middle + // invalid placeholder in that tail means garbage KV. Fall back in C128 + // steps until the tail is fully valid. + auto swa_it = + std::find_if(probes.begin(), probes.end(), [](const ProbeResult& p) { + return p.type == BlockType::SWA; + }); + if (swa_it != probes.end() && safe_hit_tokens > 0) { + const size_t swa_block_size = swa_it->block_size; + const size_t tail_required = + static_cast(swa_it->leaf->options().swa_blocks_per_seq()); + const std::vector& swa_vector = swa_it->blocks; + if (swa_block_size > 0 && tail_required > 0) { + while (safe_hit_tokens >= c128_block_size) { + const size_t trimmed_len = safe_hit_tokens / swa_block_size; + bool tail_ok = trimmed_len >= tail_required; + for (size_t i = trimmed_len - tail_required; tail_ok && i < trimmed_len; + ++i) { + if (i >= swa_vector.size() || !swa_vector[i].is_valid()) { + tail_ok = false; + } + } + if (tail_ok) { + break; + } + safe_hit_tokens -= c128_block_size; + } + if (safe_hit_tokens < c128_block_size) { + safe_hit_tokens = 0; + } + } + } + + // Exact-repeat pop: forward needs at least one C128 block to compute. + if (safe_hit_tokens == prompt_tokens && safe_hit_tokens >= c128_block_size) { + safe_hit_tokens -= c128_block_size; + } + + if (safe_hit_tokens == 0) { + out.to_drop = std::move(probes); + return out; + } + + out.to_mount.reserve(probes.size()); + for (auto& p : probes) { + const size_t target_len = safe_hit_tokens / p.block_size; + if (p.blocks.size() > target_len) { + ProbeResult drop = {p.type, + p.leaf, + /*blocks=*/{}, + /*block_size=*/p.block_size}; + drop.blocks.reserve(p.blocks.size() - target_len); + for (size_t i = target_len; i < p.blocks.size(); ++i) { + drop.blocks.emplace_back(std::move(p.blocks[i])); + } + p.blocks.resize(target_len); + out.to_drop.emplace_back(std::move(drop)); + } + out.to_mount.emplace_back(std::move(p)); + } + out.safe_hit_tokens = safe_hit_tokens; + return out; +} + +// Hand probes back to their leaf's cache and drop our aliases immediately. +// Clearing the vector after deallocate is deliberate: it destroys our Block +// aliases so refcount drops right away. Otherwise a concurrent +// leaf.deallocate() on the same block could observe ref_count > 2u +// (cache + our lingering alias) and skip the used-count decrement. +void release_probes(std::vector& probes) { + for (auto& p : probes) { + if (p.blocks.empty()) { + continue; + } + p.leaf->deallocate(p.blocks); + p.blocks.clear(); + } +} + +} // namespace + void CompositeBlockManager::allocate_shared_for_sequence(Sequence* seq) { - // Today only the flat-KV model supports prefix cache. Other shapes (DSV4 / - // xtensor) leave the KV entry's supports_prefix_cache=false (or have no KV - // leaf at all), so this is a no-op for them. - if (seq == nullptr) { + // Shared flow: probe -> trim -> mount. Trim strategy is per-shape (see + // trim_flat_kv / trim_flat_kv_linear / trim_swa_compressed). + if (seq == nullptr || combination_ == LeafCombination::UNSUPPORTED) { return; } - auto it = leaves_.find(BlockType::KV); - if (it == leaves_.end() || !it->second.supports_prefix_cache) { + std::vector probes = probe_prefix_leaves(seq, leaves_); + if (probes.empty()) { return; } - BlockManager& kv_leaf = *it->second.leaf; - seq->update_block_hashes(static_cast(kv_leaf.block_size()), - kv_leaf.options().hasher_type()); - KVCacheState& kv_state = seq->kv_state(); - const auto existed = kv_state.blocks(BlockType::KV) - .slice(0, kv_state.shared_blocks_num(BlockType::KV)); - std::vector shared = kv_leaf.allocate_shared( - seq->tokens(), existed, seq->mm_data(), seq->block_hashes()); - // Linear-state clamp: when a linear-state leaf is registered, the KV shared - // prefix must be trimmed to what a committed linear-state checkpoint can - // recover. The leaf runs right after KV through the SAME allocate_shared - // verb (plain virtual dispatch, no downcast): it reports its recoverable - // length (its own hash domain, in TOKENS) via matched_tokens and, on a hit, - // hands back the single deepest checkpoint block. Unlike KV's allocate_shared - // this call is a probe, not an admission: the returned block is a restore - // SOURCE and never enters the sequence or admission accounting. The leaf is - // KV-blind; the composite owns the cross-leaf clamp: convert the linear reach - // to KV blocks and take min(kv_match, .). The linear reach is a whole - // multiple of the prefill chunk stride, itself a multiple of the KV block - // size, so the division is exact and the result stays chunk-aligned - // (batch_input_builder relies on that alignment to emit the restore hash). - // The mounted checkpoint is stashed as the class-A restore source - // (started_empty first forwards reach the probe; continued chunks do not and - // resolve their own checkpoint). - auto linear_it = leaves_.find(BlockType::LINEAR); - if (linear_it != leaves_.end() && !shared.empty()) { - BlockManager& linear_leaf = *linear_it->second.leaf; - // Refresh the sequence's linear-state hash cache to cover the whole prefix - // before probing, mirroring update_block_hashes() for KV above. The stride - // is the leaf's prefill-chunk stride (own hash domain); the match probe - // consumes these cached hashes instead of recomputing the chain per call. - seq->update_linear_state_hashes( - static_cast(linear_leaf.options().linear_chunk_stride())); - size_t recoverable_tokens = 0; - std::vector mounted = - linear_leaf.allocate_shared(seq->tokens(), - /*existed_shared_blocks=*/{}, - MMData(), - seq->linear_state_hashes(), - &recoverable_tokens); - const size_t kv_block_size = static_cast(kv_leaf.block_size()); - const size_t recoverable_blocks = recoverable_tokens / kv_block_size; - const size_t safe_count = std::min(shared.size(), recoverable_blocks); - if (safe_count < shared.size()) { - kv_leaf.deallocate(Slice(shared).slice(safe_count)); - shared.resize(safe_count); + + TrimOutcome trimmed; + switch (combination_) { + case LeafCombination::FLAT_KV: + trimmed = trim_flat_kv(std::move(probes)); + break; + case LeafCombination::FLAT_KV_LINEAR: + trimmed = trim_flat_kv_linear(std::move(probes)); + break; + case LeafCombination::SWA_COMPRESSED: + trimmed = trim_swa_compressed(std::move(probes), seq->tokens().size()); + break; + case LeafCombination::UNSUPPORTED: + return; + } + + release_probes(trimmed.to_drop); + if (trimmed.linear_restore_src.has_value()) { + seq->set_linear_restore_src_block(std::move(*trimmed.linear_restore_src)); + } + + // Mount: FLAT_KV{,_LINEAR} defer to Sequence::add_shared_blocks (owns + // replace + exact-repeat); SWA_COMPRESSED mounts each leaf and advances + // kv_cache_tokens_num_ once with the shared safe_hit. + switch (combination_) { + case LeafCombination::FLAT_KV: + case LeafCombination::FLAT_KV_LINEAR: { + if (!trimmed.to_mount.empty()) { + ProbeResult& kv = trimmed.to_mount.front(); + seq->add_shared_blocks(kv.type, std::move(kv.blocks)); + } + break; } - // The mounted checkpoint is a restore SOURCE only; it must never enter the - // sequence's LINEAR block vector, whose slot[0] is the live, GDN - // in-place-updated slot. - if (!mounted.empty()) { - seq->set_linear_restore_src_block(std::move(mounted.front())); + case LeafCombination::SWA_COMPRESSED: { + if (trimmed.safe_hit_tokens == 0) { + break; + } + for (auto& p : trimmed.to_mount) { + seq->kv_state().mount_composite_shared(p.type, std::move(p.blocks)); + } + seq->kv_state().set_kv_cache_tokens_num(trimmed.safe_hit_tokens); + break; } + case LeafCombination::UNSUPPORTED: + break; } - seq->add_shared_blocks(BlockType::KV, std::move(shared)); } void CompositeBlockManager::cache_for_sequence(Sequence* seq) { - if (seq == nullptr) { + // Final flush at deallocate. KV shapes flush the tail via leaf->cache(); + // SWA_COMPRESSED re-runs the pre-grow hook (cursor-guarded, idempotent) so + // the window-tail block that never triggered another allocate_sequence + // still makes it into the cache. + if (seq == nullptr || combination_ == LeafCombination::UNSUPPORTED) { return; } - auto it = leaves_.find(BlockType::KV); - if (it == leaves_.end() || !it->second.supports_prefix_cache) { - return; + switch (combination_) { + case LeafCombination::FLAT_KV: + case LeafCombination::FLAT_KV_LINEAR: { + BlockManager& kv_leaf = *leaf_of(BlockType::KV); + seq->update_block_hashes(static_cast(kv_leaf.block_size()), + kv_leaf.options().hasher_type()); + KVCacheState& kv_state = seq->kv_state(); + std::vector* blocks = kv_state.mutable_blocks(BlockType::KV); + kv_leaf.cache(seq->cached_tokens(), + *blocks, + kv_state.shared_blocks_num(BlockType::KV), + seq->mm_data(), + seq->block_hashes()); + break; + } + case LeafCombination::SWA_COMPRESSED: { + cache_full_blocks_for_sequence(seq); + break; + } + case LeafCombination::UNSUPPORTED: + break; } - BlockManager& kv_leaf = *it->second.leaf; - seq->update_block_hashes(static_cast(kv_leaf.block_size()), - kv_leaf.options().hasher_type()); - KVCacheState& kv_state = seq->kv_state(); - std::vector* blocks = kv_state.mutable_blocks(BlockType::KV); - kv_leaf.cache(seq->cached_tokens(), - *blocks, - kv_state.shared_blocks_num(BlockType::KV), - seq->mm_data(), - seq->block_hashes()); } void CompositeBlockManager::cache_for_sequence(Sequence* seq, size_t num_tokens) { - if (seq == nullptr) { + // Chunked-prefill mid-step: only KV needs a token-clamped flush now; + // SWA / C4 / C128 get their flush from the pre-grow hook at the top of + // the next allocate_sequence. + if (seq == nullptr || combination_ == LeafCombination::UNSUPPORTED) { return; } - auto it = leaves_.find(BlockType::KV); - if (it == leaves_.end() || !it->second.supports_prefix_cache) { - return; - } - BlockManager& kv_leaf = *it->second.leaf; - KVCacheState& kv_state = seq->kv_state(); - const size_t block_size = kv_leaf.block_size(); - // Clamp to blocks actually allocated and to the sequence's own tokens; the - // last partial block is dropped by the prefix-cache insert. - const size_t available_tokens_num = - std::min({num_tokens, - kv_state.num_blocks(BlockType::KV) * block_size, - seq->tokens().size()}); - const size_t existed_shared_blocks_num = - kv_state.shared_blocks_num(BlockType::KV); - if (available_tokens_num <= existed_shared_blocks_num * block_size) { - return; + switch (combination_) { + case LeafCombination::FLAT_KV: + case LeafCombination::FLAT_KV_LINEAR: { + BlockManager& kv_leaf = *leaf_of(BlockType::KV); + KVCacheState& kv_state = seq->kv_state(); + const size_t block_size = kv_leaf.block_size(); + const size_t available_tokens_num = + std::min({num_tokens, + kv_state.num_blocks(BlockType::KV) * block_size, + seq->tokens().size()}); + const size_t existed_shared_blocks_num = + kv_state.shared_blocks_num(BlockType::KV); + if (available_tokens_num > existed_shared_blocks_num * block_size) { + seq->update_block_hashes(static_cast(block_size), + kv_leaf.options().hasher_type()); + std::vector* blocks = kv_state.mutable_blocks(BlockType::KV); + CHECK_GE(blocks->size(), existed_shared_blocks_num); + kv_leaf.cache(seq->tokens().slice(0, available_tokens_num), + *blocks, + existed_shared_blocks_num, + seq->mm_data(), + seq->block_hashes()); + } + break; + } + case LeafCombination::SWA_COMPRESSED: + case LeafCombination::UNSUPPORTED: + break; } - seq->update_block_hashes(static_cast(block_size), - kv_leaf.options().hasher_type()); - std::vector* blocks = kv_state.mutable_blocks(BlockType::KV); - CHECK_GE(blocks->size(), existed_shared_blocks_num); - kv_leaf.cache(seq->tokens().slice(0, available_tokens_num), - *blocks, - existed_shared_blocks_num, - seq->mm_data(), - seq->block_hashes()); } std::vector CompositeBlockManager::allocate_blocks(BlockType type, @@ -415,7 +710,7 @@ std::vector CompositeBlockManager::allocate_blocks(BlockType type, } void CompositeBlockManager::deallocate(const Slice& blocks) { - // Route each run of blocks to its owning leaf by Block::manager(). + // Route each contiguous run to its owning leaf by Block::manager(). if (blocks.empty()) { return; } @@ -456,9 +751,8 @@ std::vector CompositeBlockManager::allocate(size_t /*num_blocks*/) { std::optional> CompositeBlockManager::allocate_for_sequence( Sequence* /*seq*/, size_t /*num_tokens*/) { - // The composition is driven via allocate_sequence(), which fans out to each - // leaf's allocate_for_sequence. The leaf-level entry point is meaningless on - // the composite itself. + // The pool drives the composite via allocate_sequence(); the leaf-level + // entry point is meaningless on the composite itself. NOT_IMPLEMENTED(); return std::nullopt; } @@ -467,8 +761,7 @@ std::vector CompositeBlockManager::allocate_shared( const Slice& /*tokens_ids*/, const Slice& /*existed_shared_blocks*/, const MMData& /*mm_data*/, - const Slice& /*block_hashes*/, - size_t* /*matched_tokens*/) { + const Slice& /*block_hashes*/) { NOT_IMPLEMENTED(); return {}; } @@ -501,14 +794,13 @@ size_t CompositeBlockManager::num_blocks_in_prefix_cache() const { const CompositeBlockManager::LeafEntry* CompositeBlockManager::capacity_leaf() const { + // Smallest block_size = finest granularity = closest to the base block the + // scheduler assumes. KV for normal models; C4 for DSV4. const LeafEntry* chosen = nullptr; for (const auto& [type, entry] : leaves_) { if (!entry.participates_in_admission) { continue; } - // Smallest block_size = finest granularity = closest to the base block the - // scheduler assumes. For normal models this is the lone KV leaf; for DSV4 - // (C4 bs=4*base, C128 bs=128*base) this is C4. if (chosen == nullptr || entry.leaf->block_size() < chosen->leaf->block_size()) { chosen = &entry; @@ -518,12 +810,8 @@ const CompositeBlockManager::LeafEntry* CompositeBlockManager::capacity_leaf() } size_t CompositeBlockManager::num_free_blocks() const { - // Scheduler-facing capacity is reported as a single admission leaf's raw - // block count (see capacity_leaf): mixing leaves of different block_size - // would make num_free * block_size() meaningless. This reproduces the - // pre-refactor single-group (`sub_managers_[1]`) semantics and is identical - // to the old BlockManagerImpl for normal models (KV is the only admission - // leaf). + // Reports one admission leaf's raw block count. Mixing leaves of different + // block_size would make num_free * block_size() meaningless. const LeafEntry* leaf = capacity_leaf(); return leaf == nullptr ? 0 : leaf->leaf->num_free_blocks(); } diff --git a/xllm/core/framework/block/composite_block_manager.h b/xllm/core/framework/block/composite_block_manager.h index 7d21c06a5a..c0ef7279f6 100644 --- a/xllm/core/framework/block/composite_block_manager.h +++ b/xllm/core/framework/block/composite_block_manager.h @@ -15,6 +15,7 @@ limitations under the License. #pragma once +#include #include #include #include @@ -24,50 +25,59 @@ limitations under the License. namespace xllm { -// Generic composition of multiple BlockManager leaves keyed by BlockType. The -// map key decides which KVCacheState slot a leaf's blocks land in; the leaf -// itself is type-free. The composite is the only block-side class that touches -// Sequence: it extracts parameters, drives the leaves' pure planners and -// type-free primitives, and writes results back by key. It holds no -// model-specific logic. +// Composition of BlockManager leaves keyed by BlockType. The map key decides +// which KVCacheState slot a leaf's blocks land in; the leaf itself is +// type-free. The composite is the only block-side class that touches +// Sequence. class CompositeBlockManager : public BlockManager { public: - // Per-leaf entry. The admission / prefix roles live here (on the composite), - // not on the leaf: the same BlockManagerImpl is a prefix-capable admission - // leaf under the KV key but a non-prefix admission leaf under C4/C128. + // The admission / prefix roles live on the composite, not on the leaf: the + // same BlockManagerImpl is prefix-capable under KV but non-prefix under + // C4/C128. struct LeafEntry { std::unique_ptr leaf; bool participates_in_admission = true; bool supports_prefix_cache = false; }; + // Which prefix-cache-supported leaf shape this composite carries. Classified + // once at construction and cached in `combination_`; the sequence-level + // orchestrators dispatch on it to pick the trim / mount strategy. + // FLAT_KV Plain KV (normal / Qwen). No trim. + // FLAT_KV_LINEAR KV + LINEAR restore (Qwen3.5-Next GDN). Composite clamps + // KV to LINEAR's recoverable budget and pulls out its + // deepest checkpoint as a restore source. + // SWA_COMPRESSED SWA + C4 + C128 (DSV4). Cross-leaf min, C128-stride + // clamp, SWA tail-continuity, exact-repeat pop. + // UNSUPPORTED Prefix cache off (xtensor / --enable_prefix_cache=false). + enum class LeafCombination : int8_t { + FLAT_KV, + FLAT_KV_LINEAR, + SWA_COMPRESSED, + UNSUPPORTED, + }; + explicit CompositeBlockManager(std::map leaves); ~CompositeBlockManager() override = default; bool is_composite() const override { return true; } - // —— Sequence-level orchestration (the only Sequence-aware surface) —— - // Drives every leaf's allocate_for_sequence(seq, num_tokens), stages the - // blocks each leaf returns, and commits them into the sequence under the - // leaf's block_type() once all succeed (rolling back on any failure). - // Distinct from the leaf-level BlockManager::allocate_for_sequence, which - // returns the blocks for a single leaf without inserting them. + // Sequence-level orchestration (the only Sequence-aware surface). Drives + // each leaf's own primitives and writes results back into KVCacheState + // under the leaf's block_type(). bool allocate_sequence(Sequence* seq, size_t num_tokens); void deallocate_for_sequence(Sequence* seq); void allocate_shared_for_sequence(Sequence* seq); void cache_for_sequence(Sequence* seq); void cache_for_sequence(Sequence* seq, size_t num_tokens); - // Typed block-level allocation routed to the leaf under `type`. Used by the - // pool for beam copy-on-write (which needs exactly one KV block). + // Typed block-level allocation routed to the leaf under `type`. Used for + // beam copy-on-write. std::vector allocate_blocks(BlockType type, size_t num_blocks); // Type-ambiguous block-level primitives are not meaningful on a composition. void deallocate(const Slice& blocks) override; std::vector allocate(size_t num_blocks) override; - // Leaf-level growth is not meaningful on the composition: the pool drives the - // composite via allocate_sequence() (which fans out to each leaf's - // allocate_for_sequence). Satisfies the pure-virtual base; never called. std::optional> allocate_for_sequence( Sequence* seq, size_t num_tokens) override; @@ -75,8 +85,7 @@ class CompositeBlockManager : public BlockManager { const Slice& tokens_ids, const Slice& existed_shared_blocks = {}, const MMData& mm_data = MMData(), - const Slice& block_hashes = {}, - size_t* matched_tokens = nullptr) override; + const Slice& block_hashes = {}) override; void cache(const Slice& token_ids, std::vector& blocks, size_t existed_shared_blocks_num = 0, @@ -88,48 +97,49 @@ class CompositeBlockManager : public BlockManager { void reset_prefix_cache() override; // Stats reported from the single capacity leaf (see capacity_leaf()). - size_t num_blocks_in_prefix_cache() const override; // sum over all leaves - size_t num_free_blocks() const override; // from capacity leaf - size_t num_used_blocks() const override; // from capacity leaf + size_t num_blocks_in_prefix_cache() const override; + size_t num_free_blocks() const override; + size_t num_used_blocks() const override; double kv_cache_utilization() const override; void free(int32_t block_id) override; Block allocate() override; - size_t num_total_blocks() const override; // from capacity leaf + size_t num_total_blocks() const override; void reserve_xtensor_padding_blocks() override; size_t num_sub_managers() const { return leaves_.size(); } private: - // Test-only reach into the LINEAR leaf, mirroring the friend pattern the leaf - // itself uses: the test peer seeds checkpoints the way the scheduler does - // while resolving cache ops. No production caller reaches a leaf by type -- - // the composite drives every leaf through its own type-free orchestration -- - // so leaf_of stays private. friend class BlockManagerPoolTestPeer; - // Leaf serving `type`, or nullptr if none. Internal helper for the - // composite's own orchestration; not a Sequence-aware or externally driven - // surface. BlockManager* leaf_of(BlockType type) const; - // The single admission leaf whose raw block count defines the pool's - // scheduler-facing capacity unit. Schedulers treat num_free/used/total_blocks - // as counts of base (block_size()) blocks, so we must report one leaf's raw - // count rather than mixing leaves of different block sizes. Picks the - // admission leaf with the smallest block_size (the finest-grained, closest to - // base): KV for normal models, C4 for DSV4. Reproduces the pre-refactor - // single-`sub_managers_[1]` capacity semantics. nullptr if none. + // The admission leaf whose raw block count defines the pool's + // scheduler-facing capacity unit. Picks the finest-grained admission leaf + // (smallest block_size): KV for normal models, C4 for DSV4. const LeafEntry* capacity_leaf() const; + static LeafCombination classify_leaf_combination( + const std::map& leaves); + + // Pre-grow / final-flush cache hook. Inserts every leaf's newly-forwarded + // full blocks into its own prefix cache incrementally (cursor lives in + // KVCacheState::num_cached_blocks). Only fires for leaves that carry token + // cache and are not the KV leaf -- KV has its own final-flush semantics + // via cache_for_sequence at deallocate time; SINGLE / LINEAR hold no token + // cache. + void cache_full_blocks_for_sequence(Sequence* seq); + std::map leaves_; + LeafCombination combination_; }; -// Build the leaf map for one DP rank from the pool options (per-model: -// normal/Qwen -> {KV, SINGLE}; DSV4 -> {SWA, C4, C128, SINGLE}; xtensor -> -// {KV(XTensorBlockManagerImpl), SINGLE}). Leaves are wrapped in -// ConcurrentBlockManagerImpl when disagg-PD / kvcache store is enabled. The -// SINGLE entry is appended by the caller (pool). dp_rank is needed by the +// Build the leaf map for one DP rank. Per model: +// normal / Qwen -> {KV, SINGLE} +// DSV4 -> {SWA, C4, C128, SINGLE} +// xtensor -> {KV(XTensorBlockManagerImpl), SINGLE} +// Leaves are wrapped in ConcurrentBlockManagerImpl for disagg-PD / kvcache +// store. SINGLE is appended by the pool caller. dp_rank is used by the // xtensor KV leaf (per-rank VMM page pool). std::map build_composite_leaves( const BlockManager::Options& options, diff --git a/xllm/core/framework/block/concurrent_block_manager_impl.cpp b/xllm/core/framework/block/concurrent_block_manager_impl.cpp index f76b9dcece..34a5b5cd52 100644 --- a/xllm/core/framework/block/concurrent_block_manager_impl.cpp +++ b/xllm/core/framework/block/concurrent_block_manager_impl.cpp @@ -58,11 +58,10 @@ std::vector ConcurrentBlockManagerImpl::allocate_shared( const Slice& token_ids, const Slice& existed_shared_blocks, const MMData& mm_data, - const Slice& block_hashes, - size_t* matched_tokens) { + const Slice& block_hashes) { std::lock_guard lock(mutex_); std::vector blocks = inner_->allocate_shared( - token_ids, existed_shared_blocks, mm_data, block_hashes, matched_tokens); + token_ids, existed_shared_blocks, mm_data, block_hashes); for (Block& block : blocks) { block.set_manager(this); } diff --git a/xllm/core/framework/block/concurrent_block_manager_impl.h b/xllm/core/framework/block/concurrent_block_manager_impl.h index e5c4e448b6..c6fb6cf14c 100644 --- a/xllm/core/framework/block/concurrent_block_manager_impl.h +++ b/xllm/core/framework/block/concurrent_block_manager_impl.h @@ -45,8 +45,7 @@ class ConcurrentBlockManagerImpl : public BlockManager { const Slice& token_ids, const Slice& existed_shared_blocks = {}, const MMData& mm_data = MMData(), - const Slice& block_hashes = {}, - size_t* matched_tokens = nullptr) override; + const Slice& block_hashes = {}) override; void cache(const Slice& token_ids, std::vector& blocks, diff --git a/xllm/core/framework/block/linear_state_block_manager.cpp b/xllm/core/framework/block/linear_state_block_manager.cpp index 544fc58f04..08794d5008 100644 --- a/xllm/core/framework/block/linear_state_block_manager.cpp +++ b/xllm/core/framework/block/linear_state_block_manager.cpp @@ -28,32 +28,25 @@ namespace xllm { namespace { BlockManager::Options make_linear_state_options(uint32_t num_slots, - int32_t kv_block_size, int32_t chunk_stride) { BlockManager::Options options; options.num_blocks(num_slots); - // The slot pool is one id per slot; block_size carries the real KV block - // size so the checkpoint index can convert its KV-block budget into whole - // prefill chunks without the leaf holding a private copy. - options.block_size(kv_block_size); + options.block_size(chunk_stride); options.enable_prefix_cache(true); options.enable_disagg_pd(false); options.block_type(BlockType::LINEAR); - options.linear_chunk_stride(chunk_stride); return options; } } // namespace LinearStateBlockManager::LinearStateBlockManager(uint32_t num_slots, - int32_t kv_block_size, int32_t chunk_stride) - : BlockManagerImpl( - make_linear_state_options(num_slots, kv_block_size, chunk_stride)) { + : BlockManagerImpl(make_linear_state_options(num_slots, chunk_stride)) { CHECK_GT(num_slots, 1u) << "linear-state leaf needs at least one usable slot (plus padding)"; - CHECK_GT(kv_block_size, 0) - << "linear-state leaf needs the KV block size for safe-prefix clamping"; + CHECK_GT(chunk_stride, 0) + << "linear-state leaf needs a positive chunk stride"; } std::optional> @@ -132,33 +125,38 @@ std::vector LinearStateBlockManager::allocate_shared( const Slice& token_ids, const Slice& /*existed_shared_blocks*/, const MMData& /*mm_data*/, - const Slice& block_hashes, - size_t* matched_tokens) { - if (matched_tokens != nullptr) { - *matched_tokens = 0; - } + const Slice& block_hashes) { // Probe + mount through the base match virtual: prefix_cache_ is the - // LinearStatePrefixCache, whose override runs the chunk-strided probe over - // its own hash domain (chunk stride captured at construction from the - // scheduler config). |block_hashes| carries the sequence's precomputed - // chained per-chunk linear-state hashes - // (Sequence::update_linear_state_hashes), forwarded here so the probe - // consumes them instead of recomputing; the override writes the recoverable - // token length back through |matched_tokens| and returns the single deepest - // checkpoint (already pinned + LRU-promoted via find()). Linear thus travels - // the same allocate_shared verb as KV, and beneath it the same match verb -- - // no downcast, plain virtual dispatch. The index is KV-blind; the composite - // owns the cross-leaf min + chunk alignment. + // LinearStatePrefixCache, whose gap-tolerant walk over the chunk-strided + // hash domain (chunk stride captured at construction from the scheduler + // config; block_size_ of this cache) returns a vector of shape + // `[hit_or_inv, ..., last_hit]` per chunk. |block_hashes| carries the + // sequence's precomputed chained per-chunk linear-state hashes + // (Sequence::update_linear_state_hashes) so the probe consumes them + // instead of recomputing; the composite picks the deepest valid slot out + // as the class-A restore source. Linear thus travels the same + // allocate_shared verb as KV, and beneath it the same match verb -- no + // downcast, plain virtual dispatch. The index is KV-blind; the composite + // owns the cross-leaf min + chunk alignment + deepest-pick. // - // One-block insight: the deepest checkpoint's cumulatively-compressed state - // subsumes every earlier hit, so match() returns at most one block. That - // block, when present, is the class-A restore source; the caller stashes it - // WITHOUT adding it to the sequence's live LINEAR slot vector. - return prefix_cache_->match(token_ids, + // Reserve one fresh token by slicing off the last input token before the + // probe. The forward must compute at least one token, so a checkpoint that + // sits exactly at token_ids.size() would leave zero tokens to run; slicing + // to size-1 makes match() naturally stop at the deepest chunk STRICTLY + // below the final token. This is LINEAR-specific -- it cannot be delegated + // to KV's own size-1 (the add_shared_blocks last-block pop): that pop lands + // on an arbitrary KV block boundary, whereas linear checkpoints exist only + // on sparse chunk boundaries. If reuse were clamped by KV's boundary the + // restore would look for a checkpoint that need not be there; on a miss the + // sequence cold-starts on a non-empty KV prefix and its recurrent state is + // corrupt. + if (token_ids.size() == 0) { + return {}; + } + return prefix_cache_->match(token_ids.slice(0, token_ids.size() - 1), /*existed_shared_blocks=*/{}, MMData(), - block_hashes, - matched_tokens); + block_hashes); } void LinearStateBlockManager::cache(const Slice& /*token_ids*/, diff --git a/xllm/core/framework/block/linear_state_block_manager.h b/xllm/core/framework/block/linear_state_block_manager.h index efda4834c7..e71288cc67 100644 --- a/xllm/core/framework/block/linear_state_block_manager.h +++ b/xllm/core/framework/block/linear_state_block_manager.h @@ -42,12 +42,10 @@ class Sequence; class LinearStateBlockManager final : public BlockManagerImpl { public: // |chunk_stride| is the linear-state checkpoint stride in tokens (one prefill - // chunk), forwarded to the checkpoint index as its hash-domain step. It - // defaults to -1 (probe disabled) for the block-level unit tests that drive - // the slot pool directly and never exercise the token-based match() override. - LinearStateBlockManager(uint32_t num_slots, - int32_t kv_block_size, - int32_t chunk_stride = -1); + // chunk), used as block_size for both the slot pool and the prefix-cache hash + // domain. Defaults to -1 (probe disabled) for block-level unit tests that + // drive the slot pool directly and never exercise the token-based match(). + LinearStateBlockManager(uint32_t num_slots, int32_t chunk_stride = -1); ~LinearStateBlockManager() override = default; // ---- BlockManagerImpl overrides ---- @@ -79,24 +77,26 @@ class LinearStateBlockManager final : public BlockManagerImpl { // admission. |token_ids| is the sequence's full token span; |block_hashes| // carries the sequence's precomputed chained per-chunk linear-state hashes // (own hash domain, chunk-strided -- not the KV by-block chain), forwarded to - // the probe; |existed_shared_blocks|/|mm_data| are unused. The recoverable - // length in TOKENS (a multiple of the prefill chunk stride) is written - // through |matched_tokens|. The leaf is KV-blind: the composite owns the - // cross-leaf min + chunk alignment + KV-block clamp. + // the probe; |existed_shared_blocks|/|mm_data| are unused. Reach in tokens + // is derived from `returned.size() * block_size()` (block_size() is the + // chunk stride), symmetric with KV/SWA. The leaf is KV-blind: the composite + // owns the cross-leaf min + chunk alignment + KV-block clamp. // // The recurrent state is cumulatively compressed, so at most ONE checkpoint - // is mounted -- the single deepest hit, whose state subsumes every earlier - // one. On a hit the return vector holds that one checkpoint block, fetched - // via find() (refcount+1, LRU-promoted) so it stays pinned until restored; - // the caller stashes it as the class-A restore source WITHOUT adding it to - // the sequence's LINEAR block vector (that slot is the live, in-place-updated - // one). Returns an empty vector when nothing hits. + // survives -- the single deepest hit, whose state subsumes every earlier + // one. The returned vector has the shape `[inv, inv, ..., deepest_valid]` + // at chunk-stride granularity (deepest slot fetched via find(), so refcount + // is +1 and it is LRU-promoted, pinned until restored). The caller + // (composite) picks the deepest slot out as the class-A restore source and + // never mounts LINEAR blocks into the sequence's live LINEAR vector (that + // slot is the live, in-place-updated one). Reach in tokens is + // `returned.size() * block_size()` where block_size() is the chunk stride; + // an empty vector means nothing hit. std::vector allocate_shared( const Slice& token_ids, const Slice& existed_shared_blocks = {}, const MMData& mm_data = MMData(), - const Slice& block_hashes = {}, - size_t* matched_tokens = nullptr) override; + const Slice& block_hashes = {}) override; void cache(const Slice& token_ids, std::vector& blocks, size_t existed_shared_blocks_num = 0, diff --git a/xllm/core/framework/block/single_block_manager.cpp b/xllm/core/framework/block/single_block_manager.cpp index c742c55887..6f770c050f 100644 --- a/xllm/core/framework/block/single_block_manager.cpp +++ b/xllm/core/framework/block/single_block_manager.cpp @@ -86,8 +86,7 @@ std::vector SingleBlockManager::allocate_shared( const Slice& /*token_ids*/, const Slice& /*existed_shared_blocks*/, const MMData& /*mm_data*/, - const Slice& /*block_hashes*/, - size_t* /*matched_tokens*/) { + const Slice& /*block_hashes*/) { NOT_IMPLEMENTED(); return {}; } diff --git a/xllm/core/framework/block/single_block_manager.h b/xllm/core/framework/block/single_block_manager.h index 400cc37c44..ada344aa23 100644 --- a/xllm/core/framework/block/single_block_manager.h +++ b/xllm/core/framework/block/single_block_manager.h @@ -57,8 +57,7 @@ class SingleBlockManager final : public BlockManagerImpl { const Slice& token_ids, const Slice& existed_shared_blocks = {}, const MMData& mm_data = MMData(), - const Slice& block_hashes = {}, - size_t* matched_tokens = nullptr) override; + const Slice& block_hashes = {}) override; void cache(const Slice& token_ids, std::vector& blocks, size_t existed_shared_blocks_num = 0, diff --git a/xllm/core/framework/block/sliding_window_block_manager.cpp b/xllm/core/framework/block/sliding_window_block_manager.cpp index 7e298be407..6ba77b91bc 100644 --- a/xllm/core/framework/block/sliding_window_block_manager.cpp +++ b/xllm/core/framework/block/sliding_window_block_manager.cpp @@ -17,22 +17,21 @@ limitations under the License. #include -namespace xllm { - -namespace { - -// SWA never serves prefix cache; force it off so the base does not build one. -BlockManager::Options without_prefix_cache(BlockManager::Options options) { - options.enable_prefix_cache(false); - return options; -} +#include "framework/prefix_cache/prefix_cache.h" -} // namespace +namespace xllm { SlidingWindowBlockManager::SlidingWindowBlockManager(const Options& options) - : BlockManagerImpl(without_prefix_cache(options)) { + : BlockManagerImpl(options) { CHECK_GT(options_.swa_blocks_per_seq(), 0u) << "swa_blocks_per_seq must be positive"; + if (options_.enable_prefix_cache()) { + // SWA prefix cache uses Sequence::block_hashes_ (TEXT chain). VLM MM + // hasher is not compatible; fail loud instead of silently corrupting hits. + CHECK(options_.hasher_type() == BlockHasherType::TEXT) + << "SWA prefix cache does not yet support VLM (MM hasher). " + "Disable prefix cache for VLM DSV4 or wait for VLM support."; + } } void SlidingWindowBlockManager::release_out_of_window(Sequence* seq) { @@ -45,10 +44,6 @@ void SlidingWindowBlockManager::release_out_of_window(Sequence* seq) { if (block_size == 0 || swa_blocks.empty()) { return; } - // Leading logical blocks that have fully slid out of the window can be freed. - // Runs AFTER the composite has committed this round's growth, and only on a - // successful allocation, so a failed allocate_sequence never releases the - // sequence's existing SWA blocks. const size_t cached_tokens = kv_state.kv_cache_tokens_num(); const size_t num_spec_tokens = static_cast(options_.num_speculative_tokens()); @@ -64,9 +59,9 @@ void SlidingWindowBlockManager::release_out_of_window(Sequence* seq) { if (release_blocks == 0) { return; } - // Move slid-out blocks out (leaving invalid placeholders so logical-position - // indexing stays stable) and deallocate them through the base free list. The - // is_valid() guard makes re-entry a safe no-op. + // Move slid-out blocks out (leaving invalid placeholders so positional + // indexing stays stable). Cache alias still pins the physical block, so + // deallocate walks the ref<=2u branch and only clears usage bookkeeping. std::vector blocks_to_release; blocks_to_release.reserve(release_blocks); for (size_t j = 0; j < release_blocks; ++j) { @@ -80,25 +75,37 @@ void SlidingWindowBlockManager::release_out_of_window(Sequence* seq) { } std::vector SlidingWindowBlockManager::allocate_shared( - const Slice& /*token_ids*/, + const Slice& token_ids, const Slice& /*existed_shared_blocks*/, - const MMData& /*mm_data*/, - const Slice& /*block_hashes*/, - size_t* /*matched_tokens*/) { - NOT_IMPLEMENTED(); - return {}; -} + const MMData& mm_data, + const Slice& block_hashes) { + if (!options_.enable_prefix_cache() || options_.block_size() == 0 || + prefix_cache_ == nullptr) { + return {}; + } + AUTO_COUNTER(prefix_cache_latency_seconds_match); + std::vector result = prefix_cache_->match(token_ids, + /*existed_shared_blocks=*/{}, + mm_data, + block_hashes); + if (result.empty()) { + return {}; + } -void SlidingWindowBlockManager::cache(const Slice& /*token_ids*/, - std::vector& /*blocks*/, - size_t /*existed_shared_blocks_num*/, - const MMData& /*mm_data*/, - const Slice& /*block_hashes*/) { - NOT_IMPLEMENTED(); -} + // Bookkeeping: mark_used only for valid positions. mark_used is idempotent + // per block id, so blocks shared across sequences are only counted once. + size_t added = 0; + for (const auto& b : result) { + if (b.is_valid() && mark_used(&usage_accounted_ids_, b.id())) { + ++added; + } + } + num_used_blocks_.fetch_add(added, std::memory_order_relaxed); -void SlidingWindowBlockManager::cache(const std::vector& /*blocks*/) { - NOT_IMPLEMENTED(); + const size_t reach_tokens = + result.size() * static_cast(options_.block_size()); + COUNTER_ADD(prefix_cache_match_length_total, reach_tokens); + return result; } } // namespace xllm diff --git a/xllm/core/framework/block/sliding_window_block_manager.h b/xllm/core/framework/block/sliding_window_block_manager.h index 1a7a3f6f3e..adc05e10f7 100644 --- a/xllm/core/framework/block/sliding_window_block_manager.h +++ b/xllm/core/framework/block/sliding_window_block_manager.h @@ -19,45 +19,31 @@ limitations under the License. namespace xllm { -// Sliding-window leaf of CompositeBlockManager. -// -// The low-level physical block pool (free list, allocate / deallocate / free, -// num_* accounting, id-0 padding) is identical to BlockManagerImpl, so this -// derives from it and reuses that machinery unchanged. Growth is the base flat -// append (allocate_for_sequence inherited). The sliding-window specifics are: -// - release_out_of_window() drops leading blocks that have slid out of the -// window; the composite calls it AFTER committing a successful round, so a -// failed allocate_sequence never releases the sequence's existing blocks; -// - the SWA block list grows by logical position and never shrinks in size; -// released leading positions become invalid placeholders; -// - the pool is sized with slack (see build_composite_leaves) so the peak -// "old blocks not yet released + new tail" fits without borrowing capacity; -// - no prefix cache. +// Sliding-window leaf of CompositeBlockManager. Reuses BlockManagerImpl's +// physical pool and flat-append growth. SWA-specific behavior: +// - release_out_of_window() drops leading slid-out blocks; released +// positions stay as invalid placeholders so DSA modulo indexing +// (`(pos/block_size) % semantic_cols`) remains stable. +// - Prefix cache is gap-tolerant (LinearStatePrefixCache backend); the +// tail-continuity check lives on the composite. TEXT hasher only. class SlidingWindowBlockManager : public BlockManagerImpl { public: explicit SlidingWindowBlockManager(const Options& options); ~SlidingWindowBlockManager() override = default; - // Release the leading SWA blocks of the sequence that have fully slid out of - // the window: move them out (leaving invalid placeholders so logical-position - // indexing stays stable) and return their ids to the pool. Called by the - // composite after a successful allocate commit. Growth reuses the inherited - // BlockManagerImpl::allocate_for_sequence (flat append). + // Deallocate leading blocks that have slid out of the window; leaves + // invalid placeholders in their slots. Called by the composite after a + // successful allocate commit. void release_out_of_window(Sequence* seq) override; - // SWA never serves prefix cache; these must not be reached. + // Gap-tolerant SWA probe. Delegates to LinearStatePrefixCache::match; see + // that class for the shape returned. Composite owns the tail-continuity + // check and the min-across-leaves clamp. std::vector allocate_shared( const Slice& token_ids, const Slice& existed_shared_blocks = {}, const MMData& mm_data = MMData(), - const Slice& block_hashes = {}, - size_t* matched_tokens = nullptr) override; - void cache(const Slice& token_ids, - std::vector& blocks, - size_t existed_shared_blocks_num = 0, - const MMData& mm_data = MMData(), - const Slice& block_hashes = {}) override; - void cache(const std::vector& blocks) override; + const Slice& block_hashes = {}) override; uint32_t swa_blocks_per_seq() const { return options_.swa_blocks_per_seq(); } }; diff --git a/xllm/core/framework/prefix_cache/linear_state_prefix_cache.cpp b/xllm/core/framework/prefix_cache/linear_state_prefix_cache.cpp index 011294566f..73ddc0a631 100644 --- a/xllm/core/framework/prefix_cache/linear_state_prefix_cache.cpp +++ b/xllm/core/framework/prefix_cache/linear_state_prefix_cache.cpp @@ -15,74 +15,189 @@ limitations under the License. #include "linear_state_prefix_cache.h" +#include +#include + +#include + +#include "common/metrics.h" + namespace xllm { std::vector LinearStatePrefixCache::match( const Slice& token_ids, - const Slice& /*existed_shared_blocks*/, - const MMData& /*mm_data*/, - const Slice& block_hashes, - size_t* matched_tokens) { - if (matched_tokens != nullptr) { - *matched_tokens = 0; - } - // The chunk stride (one prefill chunk in tokens) is the linear hash domain's - // step. It is this cache's single block-boundary size (the inherited - // block_size_), fed in at construction from the scheduler config; the engine - // enforces a positive multiple of the KV block size when linear prefix cache - // is on. A misconfigured run leaves the stride at its -1 default, which - // arrives here as a saturated unsigned value that drives probe_chunk_count to - // 0 below (nothing probed); guard the zero case explicitly so the probe - // arithmetic never divides by zero. - if (block_size_ == 0) { + const Slice& existed_shared_blocks, + const MMData& mm_data, + const Slice& block_hashes) { + const size_t n_tokens = round_down(token_ids.size(), block_size_); + if (n_tokens == 0) { return {}; } - const size_t chunk_stride = block_size_; - // The chained per-chunk hashes are precomputed once on the sequence - // (update_linear_state_hashes) and passed in via |block_hashes|; they cover - // every whole chunk of the prefix, including one that may end exactly at the - // final token. Probe only chunks STRICTLY below the final token: the forward - // must compute at least one fresh token, so a checkpoint sitting exactly at - // token_ids.size() is unusable (restoring it would leave zero tokens to run). - // This is pure token-domain arithmetic -- still KV-blind -- but it cannot be - // delegated to KV's own size-1 (the add_shared_blocks pop): that pop lands on - // an arbitrary KV *block* boundary, which need not be a linear *checkpoint* - // boundary. When the two diverge the restore misses and the sequence silently - // cold-starts on a non-empty KV prefix -> corrupt recurrent state. Only this - // index knows where the checkpoints are, so snapping to the deepest one below - // the final token must happen here; the caller (composite) then takes the min - // with the KV match. - const size_t probe_chunk_count = - token_ids.size() == 0 ? 0 : (token_ids.size() - 1) / chunk_stride; - const size_t probe_limit = std::min(probe_chunk_count, block_hashes.size()); - size_t deepest_hit_chunk = probe_limit; // sentinel: nothing hit - for (size_t chunk_idx = 0; chunk_idx < probe_limit; ++chunk_idx) { - // Probe with contains() (no LRU touch): a probed-but-unused prefix must not - // pollute recency. Recurrent state is cumulatively compressed, so a later - // hit's checkpoint subsumes earlier ones -- keep only the deepest. Gaps are - // allowed: do NOT break on a miss, a later hit extends past it. - if (contains(block_hashes[chunk_idx])) { - deepest_hit_chunk = chunk_idx; + const size_t n_blocks = n_tokens / block_size_; + total_blocks_.fetch_add(n_blocks); + + std::vector blocks; + blocks.reserve(n_blocks); + blocks.insert( + blocks.end(), existed_shared_blocks.begin(), existed_shared_blocks.end()); + + DNodeList node_list; + const size_t start_block = existed_shared_blocks.size(); + + // Hit: emplace + LRU-touch. Miss: emplace invalid placeholder. + auto probe = [&](const XXH3Key& key) -> bool { + auto iter = cached_blocks_.find(key); + if (iter == cached_blocks_.end()) { + blocks.emplace_back(); + return false; + } + blocks.emplace_back(iter->second->block); + lru_lst_.remove_node(iter->second); + node_list.push_front(iter->second); + return true; + }; + + size_t last_hit_idx = 0; + bool has_any_hit = false; + if (block_hashes.size() >= n_blocks) { + for (size_t b = start_block; b < n_blocks; ++b) { + if (probe(block_hashes[b])) { + last_hit_idx = b; + has_any_hit = true; + } + } + } else { + // Fallback: compute chain on the fly, keeping it advancing across misses. + XXH3Key token_hash_key = + existed_shared_blocks.empty() + ? XXH3Key{} + : XXH3Key{existed_shared_blocks.back().get_immutable_hash_value()}; + auto hasher = + BlockHasher::create(hasher_type_, mm_data, start_block * block_size_); + for (size_t b = start_block; b < n_blocks; ++b) { + const size_t i = b * block_size_; + const uint8_t* pre_hash_value = (b == 0) ? nullptr : token_hash_key.data; + hasher->compute( + token_ids, i, i + block_size_, pre_hash_value, token_hash_key); + if (probe(token_hash_key)) { + last_hit_idx = b; + has_any_hit = true; + } } } - if (deepest_hit_chunk >= probe_limit) { - return {}; + + // Trim trailing placeholders; end at deepest hit. + const size_t keep = + has_any_hit ? (last_hit_idx + 1) : existed_shared_blocks.size(); + blocks.resize(keep); + + while (!node_list.is_empty()) { + Node* node = node_list.pop_front(); + lru_lst_.push_back(node); } - // One-block insight: mount only the single deepest checkpoint. Its - // cumulatively-compressed recurrent state covers the whole recoverable - // prefix. find() pins it (refcount+1) and promotes its LRU recency now that - // it is about to be restored, so the earlier probe stays side-effect free. - Block deepest = find(block_hashes[deepest_hit_chunk]); - if (!deepest.is_valid()) { - return {}; + + size_t valid_hits = 0; + for (const auto& b : blocks) { + if (b.is_valid()) { + ++valid_hits; + } } - if (matched_tokens != nullptr) { - *matched_tokens = (deepest_hit_chunk + 1) * chunk_stride; + matched_blocks_.fetch_add(valid_hits); + + const int64_t int_rate_percent = + static_cast(static_cast(valid_hits) * 100.0 / n_blocks); + HISTOGRAM_OBSERVE(prefix_cache_block_matched_rate, int_rate_percent); + HISTOGRAM_OBSERVE(prefix_cache_block_matched_num, valid_hits); + + return blocks; +} + +size_t LinearStatePrefixCache::insert(const Slice& token_ids, + std::vector& blocks, + size_t existed_shared_blocks_num, + const MMData& mm_data, + const Slice& block_hashes) { + const int64_t now = absl::ToUnixMicros(absl::Now()); + // allign tokens to block boundary + const size_t n_blocks = + std::min(token_ids.size() / block_size_, blocks.size()); + + if (n_blocks == 0) { + return 0; + } + CHECK_GE(n_blocks, existed_shared_blocks_num); + + DNodeList node_list; + + // Fill `token_hash_key` with the chained hash of block `block_idx`, reusing + // the precomputed hash when it covers all blocks, otherwise computing it. + // The block vector may contain invalid placeholders (SWA slid-out slots), so + // the compute path walks the chain from token 0 rather than seeding from + // blocks[existed_shared_blocks_num - 1] whose stamped hash is undefined when + // it is a placeholder. The pre-existed walk hashes tokens in order to catch + // the parent-hash cursor up. + const bool use_precomputed = block_hashes.size() >= n_blocks; + XXH3Key token_hash_key{}; + std::unique_ptr hasher; + if (!use_precomputed) { + hasher = BlockHasher::create(hasher_type_, mm_data, /*start=*/0); + const uint8_t* prev_hash = nullptr; + for (size_t b = 0; b < existed_shared_blocks_num; ++b) { + const size_t i = b * block_size_; + hasher->compute(token_ids, i, i + block_size_, prev_hash, token_hash_key); + prev_hash = token_hash_key.data; + } } - std::vector mounted; - mounted.reserve(1); - mounted.emplace_back(std::move(deepest)); - return mounted; + auto fill_block_hash = [&](size_t block_idx) { + if (use_precomputed) { + token_hash_key = block_hashes[block_idx]; + return; + } + const size_t i = block_idx * block_size_; + const uint8_t* pre_hash_value = + (block_idx == 0) ? nullptr : token_hash_key.data; + hasher->compute( + token_ids, i, i + block_size_, pre_hash_value, token_hash_key); + }; + + for (size_t block_idx = existed_shared_blocks_num; block_idx < n_blocks; + ++block_idx) { + fill_block_hash(block_idx); + // Skip invalid placeholders (e.g. SWA slid-out slots): the chain-hash + // cursor has already advanced above, but there is no physical block to + // stamp or emplace here. + if (!blocks[block_idx].is_valid()) { + continue; + } + blocks[block_idx].set_hash_value(token_hash_key.data); + + auto iter = cached_blocks_.find(token_hash_key); + if (iter != cached_blocks_.end()) { + iter->second->last_access_time = now; + + lru_lst_.remove_node(iter->second); + node_list.push_front(iter->second); + } else { + Node* new_node = new Node(); + + new_node->block = blocks[block_idx]; + new_node->last_access_time = now; + + node_list.push_front(new_node); + + cached_blocks_.emplace(std::make_pair(token_hash_key, new_node)); + + num_blocks_++; + } + } + + const size_t n_tokens = n_blocks * block_size_; + while (!node_list.is_empty()) { + Node* node = node_list.pop_front(); + lru_lst_.push_back(node); + } + + return n_tokens; } } // namespace xllm diff --git a/xllm/core/framework/prefix_cache/linear_state_prefix_cache.h b/xllm/core/framework/prefix_cache/linear_state_prefix_cache.h index 988ce97651..33738954e8 100644 --- a/xllm/core/framework/prefix_cache/linear_state_prefix_cache.h +++ b/xllm/core/framework/prefix_cache/linear_state_prefix_cache.h @@ -26,62 +26,34 @@ limitations under the License. namespace xllm { -// Checkpoint index for Qwen3.5 GDN linear-state slots. Same LRU / eviction / -// insert machinery as the base cache; the only addition is the admission-time -// prefix probe. Restore still goes through the base find(XXH3Key) point-lookup, -// so this subclass owns just the "how long a prefix can I recover" question. -// The chained per-chunk hashes it probes with are precomputed once on the -// sequence (Sequence::update_linear_state_hashes) and passed into match() via -// |block_hashes|, symmetric with how KV passes seq->block_hashes(). A prefix -// cache indexes on a single block-boundary stride: for KV that is the KV block -// size, for this index it is one prefill chunk in tokens. That stride IS the -// base block_size_ -- the factory feeds the chunk stride into the single -// block_size slot -- so the size-1 boundary rule below reads it directly and -// stays self-contained and unit-testable. +// Gap-tolerant variant of PrefixCache. Used by SWA (per base block; middle +// misses tolerated because attention only reads the tail) and LINEAR (per +// chunk; deeper hit subsumes earlier ones). Same LRU / insert / evict +// machinery as the base; only match() differs. class LinearStatePrefixCache final : public PrefixCache { public: - LinearStatePrefixCache(uint32_t chunk_stride, BlockHasherType hasher_type) - : PrefixCache(chunk_stride, hasher_type) {} + LinearStatePrefixCache(uint32_t block_size, BlockHasherType hasher_type) + : PrefixCache(block_size, hasher_type) {} - // Linear-state override of the KV match. |token_ids| is hashed into chained - // per-chunk keys (chunk = one prefill chunk in tokens, captured at - // construction from the scheduler config) for every whole chunk STRICTLY - // BELOW the final token, and the index is probed for the furthest recoverable - // prefix. Each hit extends the reach to (chunk_index + 1) * chunk_stride - // tokens. Gaps are allowed -- a later hit extends past an intervening miss. - // |matched_tokens| is written with that reach (a multiple of chunk_stride; 0 - // when nothing hits). - // - // The final token is reserved (probe limit = size - 1) so the forward always - // has at least one fresh token to compute. This size-1 is NOT redundant with - // KV's own size-1 (the add_shared_blocks last-block pop): that pop lands on - // an arbitrary KV *block* boundary, whereas linear checkpoints exist only on - // sparse *chunk* boundaries. If reuse were clamped by KV's boundary the - // restore would look for a checkpoint that need not be there; on a miss the - // sequence cold-starts on a non-empty KV prefix and its recurrent state is - // corrupt (per-token KV attention tolerates arbitrary-boundary cold starts, - // cumulatively-compressed recurrent state does not). Only this index knows - // where the checkpoints are, so the snap-to-deepest-checkpoint-below-final - // must happen here. The override stays KV-blind otherwise: it counts only in - // its own token / chunk domain and never sees the KV match length -- bounding - // reuse to what KV also matched is the composite's job. - // - // The probe uses contains() and does NOT touch LRU recency. Because recurrent - // state is cumulatively compressed, only the single deepest hit is mounted: - // its checkpoint subsumes every earlier one. On a hit the returned vector - // holds that one block, fetched via find() (refcount+1, LRU-promoted) so it - // stays pinned until restored; empty when nothing hits. - // - // |existed_shared_blocks|/|mm_data| are unused. |block_hashes| carries the - // sequence's precomputed chained per-chunk linear-state hashes (own hash - // domain, chunk-strided -- not the KV by-block chain); the probe consumes - // them directly instead of recomputing. |matched_tokens|, when non-null, - // receives the recovered token length. + // Walks the chain per block position: hit -> emplace + LRU-promote; miss + // -> emplace invalid placeholder. Trims trailing misses. Reach in tokens + // is `returned.size() * block_size_`. std::vector match(const Slice& token_ids, const Slice& existed_shared_blocks = {}, const MMData& mm_data = MMData(), - const Slice& block_hashes = {}, - size_t* matched_tokens = nullptr) override; + const Slice& block_hashes = {}) override; + + // Gap-tolerant insert for SWA slid-out placeholders. Unlike the base, the + // block vector may contain invalid placeholders (released window blocks), so + // the compute path walks the chain from token 0 rather than seeding from + // blocks[existed_shared_blocks_num - 1] (which may itself be a placeholder + // with an undefined stamped hash). Placeholders advance the chain cursor but + // are not stamped or emplaced. + size_t insert(const Slice& token_ids, + std::vector& blocks, + size_t existed_shared_blocks_num = 0, + const MMData& mm_data = MMData(), + const Slice& block_hashes = {}) override; }; } // namespace xllm diff --git a/xllm/core/framework/prefix_cache/prefix_cache.cpp b/xllm/core/framework/prefix_cache/prefix_cache.cpp index 1124f53690..983bacf44d 100644 --- a/xllm/core/framework/prefix_cache/prefix_cache.cpp +++ b/xllm/core/framework/prefix_cache/prefix_cache.cpp @@ -29,14 +29,10 @@ namespace xllm { std::vector PrefixCache::match(const Slice& token_ids, const Slice& existed_shared_blocks, const MMData& mm_data, - const Slice& block_hashes, - size_t* matched_tokens) { + const Slice& block_hashes) { // allign tokens to block boundary const size_t n_tokens = round_down(token_ids.size(), block_size_); if (n_tokens == 0) { - if (matched_tokens != nullptr) { - *matched_tokens = 0; - } return std::vector(); } @@ -99,10 +95,6 @@ std::vector PrefixCache::match(const Slice& token_ids, matched_blocks_.fetch_add(blocks.size()); - if (matched_tokens != nullptr) { - *matched_tokens = blocks.size() * block_size_; - } - int64_t int_rate_percent = static_cast( static_cast(blocks.size()) * 100.0 / n_blocks); HISTOGRAM_OBSERVE(prefix_cache_block_matched_rate, int_rate_percent); @@ -131,6 +123,10 @@ size_t PrefixCache::insert(const Slice& token_ids, // Fill `token_hash_key` with the chained hash of block `block_idx`, reusing // the precomputed hash when it covers all blocks, otherwise computing it. + // The KV / C4 / C128 prefix is solid (no placeholders), so the compute path + // seeds the chain in O(1) from the previous block's stamped hash. SWA's + // slid-out placeholders are handled by LinearStatePrefixCache::insert, which + // overrides this. const bool use_precomputed = block_hashes.size() >= n_blocks; XXH3Key token_hash_key = existed_shared_blocks_num == 0 ? XXH3Key{} diff --git a/xllm/core/framework/prefix_cache/prefix_cache.h b/xllm/core/framework/prefix_cache/prefix_cache.h index 55d854cdd5..e95bef219e 100644 --- a/xllm/core/framework/prefix_cache/prefix_cache.h +++ b/xllm/core/framework/prefix_cache/prefix_cache.h @@ -41,17 +41,18 @@ inline size_t round_down(size_t n, size_t multiple) { return (n / multiple) * multiple; } +// LRU-evicted prefix cache keyed by chained per-block hash. Solid-prefix +// match by default; LinearStatePrefixCache overrides match with a +// gap-tolerant walk for SWA / LINEAR leaves. class PrefixCache { public: struct Options { PROPERTY(int32_t, block_size) = 128; PROPERTY(BlockHasherType, hasher_type) = BlockHasherType::TEXT; - // Selects the concrete cache in create_prefix_cache: LINEAR builds the - // linear-state checkpoint index, everything else the base cache. A prefix - // cache indexes on a single block-boundary size (block_size above): for KV - // that is the KV block size, for the LINEAR index it is the checkpoint - // stride (one prefill chunk in tokens), routed into the same slot by - // BlockManagerImpl's constructor. + // Picks the concrete cache in create_prefix_cache: LINEAR / SWA build + // LinearStatePrefixCache (gap-tolerant), everything else this base class. + // A prefix cache indexes on a single stride (block_size above): KV block + // size for KV, prefill chunk stride for LINEAR, SWA base block for SWA. PROPERTY(BlockType, block_type) = BlockType::KV; }; @@ -69,55 +70,44 @@ class PrefixCache { sleep(2); }; - // When `block_hashes` covers all matchable blocks, the chained hash is reused - // directly and no hash is recomputed; otherwise it is computed on the fly - // from `token_ids`/`mm_data` (backward-compatible fallback). - // - // `matched_tokens`, when non-null, receives the matched prefix length in - // TOKENS (this KV implementation writes blocks.size() * block_size_; the - // LinearStatePrefixCache override writes the recoverable checkpoint prefix in - // its own chunk-strided domain). Callers that do not need the length leave it - // null. + // Solid-prefix probe: walks the chain per block position, stops on the + // first miss. Reach in tokens is `returned.size() * block_size_`. + // When `block_hashes` covers all matchable blocks it is consumed as-is; + // otherwise the chain is computed on the fly from `token_ids` / `mm_data`. virtual std::vector match( const Slice& token_ids, const Slice& existed_shared_blocks = {}, const MMData& mm_data = MMData(), - const Slice& block_hashes = {}, - size_t* matched_tokens = nullptr); - - // insert the token ids and blocks into the prefix tree - // and set hash key to the corresponding block - // return the length of new inserted tokens. - // `block_hashes` carries the precomputed chained hash and is reused when it - // covers all inserted blocks; otherwise the hash is computed on the fly. + const Slice& block_hashes = {}); + + // Insert blocks[existed_shared_blocks_num, n_blocks) into the cache and + // stamp their chain hash. Assumes a solid prefix (no placeholders): the + // compute path seeds the chain in O(1) from blocks[existed_shared_blocks_num + // - 1]'s stamped hash. When `block_hashes` covers all inserted blocks it is + // consumed as-is; otherwise the chain is computed on the fly. Returns the + // token span of the walked range. LinearStatePrefixCache overrides this to + // tolerate SWA slid-out placeholders. virtual size_t insert(const Slice& token_ids, std::vector& blocks, size_t existed_shared_blocks_num = 0, const MMData& mm_data = MMData(), const Slice& block_hashes = {}); - // insert the blocks with hash key into the prefix tree + // Insert already-hashed blocks (hash stamped by the caller). virtual size_t insert(Slice& blocks); virtual size_t insert(const std::vector& blocks); - // Point-lookup a single block by its chained hash key (no token-sequence - // walk). On hit, refresh the entry's LRU recency and return its block; on - // miss return an invalid Block. Used by callers whose hash domain is not the - // KV by-block-boundary chain (e.g. linear-state checkpoints with a different - // stride). + // Point-lookup by chained hash key. `find` hits refresh LRU recency and + // return the block; `contains` is LRU-neutral. Both return an invalid + // Block / false on miss. Block find(const XXH3Key& hash); - - // Whether a block is cached under `hash`, without touching LRU recency. bool contains(const XXH3Key& hash) const; - // evict blocks hold by the prefix cache - // return the actual number of evicted blocks + // Evict up to `n_blocks` LRU-oldest entries. Returns the number evicted. virtual size_t evict(size_t n_blocks); - // get the number of blocks in the prefix cache virtual size_t num_blocks() const { CHECK(num_blocks_ == cached_blocks_.size()) << "check block num failed"; - return num_blocks_; } @@ -136,10 +126,7 @@ class PrefixCache { protected: struct Node { Block block; - // the last access time of the node, used to evict blocks int64_t last_access_time = 0; - - // the previous and next nodes, used to maintain the LRU list Node* prev = nullptr; Node* next = nullptr; }; @@ -161,44 +148,34 @@ class PrefixCache { bool is_empty() { return lst_front.next == &lst_back; } - // remove the node from the LRU list, and return next node Node* remove_node(Node* node) { Node* next_node = node->next; - node->prev->next = next_node; next_node->prev = node->prev; - return next_node; } bool is_last(Node* node) { return node == &lst_back; } - // add a new node to the front of the LRU list void push_front(Node* node) { node->next = lst_front.next; lst_front.next->prev = node; - node->prev = &lst_front; lst_front.next = node; } Node* get_first() { return lst_front.next; } - // pop out node to the back of the LRU list Node* pop_front() { if (lst_front.next == &lst_back) { return nullptr; } - Node* node = lst_front.next; - lst_front.next = node->next; node->next->prev = &lst_front; - return node; } - // add a new node to the back of the LRU list void push_back(Node* node) { node->prev = lst_back.prev; node->next = &lst_back; @@ -206,28 +183,19 @@ class PrefixCache { lst_back.prev = node; } - // move the node to the back of the LRU list void move_back(Node* node) { remove_node(node); push_back(node); } - // Node lst_front; Node lst_front; Node lst_back; }; DNodeList lru_lst_; - - // the block size of the memory blocks uint32_t block_size_; - - // hasher type used to construct BlockHasher in match/insert. BlockHasherType hasher_type_; - - // the total number of blocks in the prefix cache size_t num_blocks_ = 0; - std::atomic_bool exited_{false}; std::unordered_map diff --git a/xllm/core/framework/prefix_cache/prefix_cache_factory.cpp b/xllm/core/framework/prefix_cache/prefix_cache_factory.cpp index 39be095b1e..804748f35e 100644 --- a/xllm/core/framework/prefix_cache/prefix_cache_factory.cpp +++ b/xllm/core/framework/prefix_cache/prefix_cache_factory.cpp @@ -10,9 +10,17 @@ namespace xllm { std::unique_ptr create_prefix_cache(PrefixCache::Options options) { int32_t block_size = options.block_size(); BlockHasherType hasher_type = options.hasher_type(); - if (options.block_type() == BlockType::LINEAR) { - // block_size here is the linear checkpoint stride (one prefill chunk), - // routed into the cache's single block-boundary slot by BlockManagerImpl. + // Two block types need the gap-tolerant probe: + // LINEAR -- block_size here is the linear checkpoint stride (one prefill + // chunk), routed into the cache's single block-boundary slot + // by BlockManagerImpl. + // SWA -- block_size here is the SWA base block. Middle misses are + // harmless (attention only reads the tail; the composite + // enforces the tail-continuity check). + // Both share the same subclass; the type name is retained to keep the + // SWA-adoption diff narrow. + if (options.block_type() == BlockType::LINEAR || + options.block_type() == BlockType::SWA) { return std::make_unique(block_size, hasher_type); } return std::make_unique(block_size, hasher_type); diff --git a/xllm/core/framework/request/sequence.cpp b/xllm/core/framework/request/sequence.cpp index 786429b71f..edb334e170 100644 --- a/xllm/core/framework/request/sequence.cpp +++ b/xllm/core/framework/request/sequence.cpp @@ -292,7 +292,7 @@ Sequence::Sequence(const Sequence& other) num_tokens_(other.num_tokens_), token_to_count_map_(other.token_to_count_map_), num_prompt_tokens_(other.num_prompt_tokens_), - block_hashes_(other.block_hashes_), + block_hashes_by_stride_(other.block_hashes_by_stride_), hash_block_size_(other.hash_block_size_), linear_state_hashes_(other.linear_state_hashes_), linear_hash_stride_(other.linear_hash_stride_), @@ -754,27 +754,44 @@ void Sequence::add_shared_host_blocks(BlockType type, host_kv_state_.add_shared_blocks(type, std::move(blocks), num_tokens_); } +Slice Sequence::block_hashes() const { + const auto it = block_hashes_by_stride_.find(hash_block_size_); + if (it == block_hashes_by_stride_.end()) { + return {}; + } + return it->second; +} + void Sequence::update_block_hashes(uint32_t block_size, BlockHasherType hasher_type) { if (block_size == 0) { return; } + // DSV4 admission probes SWA / C4 / C128 back-to-back with different strides + // (base / 4*base / 128*base). Each stride keeps its own chain in + // `block_hashes_by_stride_`, so switching strides extends that stride's chain + // incrementally instead of discarding and rebuilding the whole prompt chain + // every probe. Select this stride as the one block_hashes() returns. hash_block_size_ = block_size; extend_prefix_hashes(hasher_type, mm_data_, this->tokens(), block_size, /*boundary_blocks=*/num_tokens_ / block_size, - block_hashes_); + block_hashes_by_stride_[block_size]); } void Sequence::invalidate_block_hashes_from(size_t token_index) { - if (block_hashes_.empty() || hash_block_size_ == 0) { - return; - } - const size_t first_stale_block = token_index / hash_block_size_; - if (first_stale_block < block_hashes_.size()) { - block_hashes_.resize(first_stale_block); + // Truncate every stride's chain at the first block the rewrite touched; each + // stride recomputes lazily on its next update_block_hashes(). + for (auto& [block_size, hashes] : block_hashes_by_stride_) { + if (hashes.empty() || block_size == 0) { + continue; + } + const size_t first_stale_block = token_index / block_size; + if (first_stale_block < hashes.size()) { + hashes.resize(first_stale_block); + } } } diff --git a/xllm/core/framework/request/sequence.h b/xllm/core/framework/request/sequence.h index d1248cce15..16ebe99f4e 100644 --- a/xllm/core/framework/request/sequence.h +++ b/xllm/core/framework/request/sequence.h @@ -21,6 +21,7 @@ limitations under the License. #include #include +#include #include #include #include @@ -245,19 +246,24 @@ class Sequence final { // Precomputed chained block hashes used by the prefix cache. Covers all full // blocks of the current tokens; reused by match()/insert() so the hash is - // computed once per sequence instead of recomputed on every call. - Slice block_hashes() const { return block_hashes_; } - - // Extend `block_hashes_` to cover any newly completed full blocks. Cheap - // (no-op) when no new full block is available, so it is safe to call before - // every match()/cache(). + // computed once per sequence instead of recomputed on every call. Returns the + // chain for the stride set by the most recent update_block_hashes() call. + Slice block_hashes() const; + + // Extend the per-stride chain for `block_size` to cover any newly completed + // full blocks, and select it as the one block_hashes() returns. Cheap (no-op) + // when no new full block is available, so it is safe to call before every + // match()/cache(). DSV4 admission probes SWA / C4 / C128 back-to-back with + // different strides (base / 4*base / 128*base); each stride keeps its own + // chain so switching strides no longer discards and rebuilds the whole chain. void update_block_hashes(uint32_t block_size, BlockHasherType hasher_type); // Precomputed chained per-chunk hashes for the linear-state checkpoint index. - // Separate hash domain from `block_hashes_`: the stride is one prefill chunk - // (a multiple of the KV block size), not a KV block, so a linear checkpoint - // is a sparse overlay on the per-block KV cache. Consumed by the batch - // builder (save/restore boundaries) and the LINEAR leaf's match probe. + // Separate hash domain from the KV block-hash chains: the stride is one + // prefill chunk (a multiple of the KV block size), not a KV block, so a + // linear checkpoint is a sparse overlay on the per-block KV cache. Consumed + // by the batch builder (save/restore boundaries) and the LINEAR leaf's match + // probe. Slice linear_state_hashes() const { return linear_state_hashes_; } // Extend `linear_state_hashes_` to cover any newly completed full chunks at @@ -561,11 +567,17 @@ class Sequence final { // the length of the prompt tokens size_t num_prompt_tokens_ = 0; - // Precomputed chained block hashes covering all full blocks of `tokens_`. - // Extended incrementally; consumed by the prefix cache. - std::vector block_hashes_; - - // Block size used to compute `block_hashes_` (0 until first computed). + // Precomputed chained block hashes covering all full blocks of `tokens_`, + // keyed by block-size stride. DSV4 admission probes multiple strides (base / + // 4*base / 128*base) per tick; each keeps its own chain so a stride switch + // extends incrementally instead of discarding and rebuilding. Extended + // incrementally; consumed by the prefix cache. std::map node storage is + // pointer-stable, so a Slice handed out by block_hashes() survives inserts of + // other strides. + std::map> block_hashes_by_stride_; + + // Stride selected by the most recent update_block_hashes() call; keys + // block_hashes() into `block_hashes_by_stride_` (0 until first computed). uint32_t hash_block_size_ = 0; // Precomputed chained per-chunk hashes for the linear-state checkpoint index diff --git a/xllm/core/framework/request/sequence_kv_state.cpp b/xllm/core/framework/request/sequence_kv_state.cpp index d60f843aa9..6e50851436 100644 --- a/xllm/core/framework/request/sequence_kv_state.cpp +++ b/xllm/core/framework/request/sequence_kv_state.cpp @@ -136,6 +136,7 @@ void KVCacheState::incr_shared_blocks_num(BlockType type, size_t num) { void KVCacheState::erase_blocks(BlockType type) { composite_blocks_.erase(type); num_owned_shared_blocks_.erase(type); + num_cached_blocks_.erase(type); if (type == BlockType::LINEAR) { pending_linear_save_hash_.reset(); linear_restore_src_block_.reset(); @@ -158,16 +159,23 @@ size_t KVCacheState::shared_blocks_num(BlockType type) const { } size_t KVCacheState::shared_tokens_num() const { - // Shared token count is a sequence-level value: shared_blocks * block_size is - // the same across block types. Read it off whichever type carries shared - // blocks (KV is the canonical source). + // Shared token count is a sequence-level value: shared_blocks * block_size + // is the same across cache-bearing block types by construction (composite + // trim keeps every leaf aligned to a common safe_hit). Read it off any type + // that carries shared blocks. Use the first VALID block in the vector for + // block_size -- vec[0] may be an invalid placeholder in the gap-tolerant + // SWA case (last_hit at position K, vec_len = K+1, positions [0..K-1] may + // include invalid placeholders), and reading .size() off an invalid Block + // returns 0. for (const auto& [type, shared] : num_owned_shared_blocks_) { if (shared == 0) { continue; } const Slice bs = blocks(type); - if (!bs.empty()) { - return shared * bs[0].size(); + for (const auto& b : bs) { + if (b.is_valid()) { + return shared * b.size(); + } } } return 0; @@ -217,6 +225,34 @@ void KVCacheState::add_shared_blocks(BlockType type, kv_cache_tokens_num_ = num_shared_tokens; } +void KVCacheState::mount_composite_shared(BlockType type, + std::vector&& shared_blocks) { + if (shared_blocks.empty()) { + return; + } + std::vector& bs = composite_blocks_[type]; + CHECK(bs.empty()) << "mount_composite_shared: existing blocks under type " + << static_cast(type); + const size_t count = shared_blocks.size(); + bs = std::move(shared_blocks); + num_owned_shared_blocks_[type] = static_cast(count); + // Mounted blocks are already resident in the prefix cache (they came from + // find()); the pre-grow hook consults num_cached_blocks to skip them. Chain + // hashes are recomputed from tokens inside PrefixCache::insert (compute + // path), so we don't have to worry about trailing invalid placeholders + // here. + num_cached_blocks_[type] = count; +} + +size_t KVCacheState::num_cached_blocks(BlockType type) const { + const auto it = num_cached_blocks_.find(type); + return it == num_cached_blocks_.end() ? 0 : it->second; +} + +void KVCacheState::set_num_cached_blocks(BlockType type, size_t n) { + num_cached_blocks_[type] = n; +} + void KVCacheState::set_slice_window_size(uint32_t size) { CHECK(size > 0); CHECK(!blocks(BlockType::SWA).empty()); @@ -335,6 +371,7 @@ void KVCacheState::advance_group_transfer_block_idx(BlockType type, void KVCacheState::reset() { kv_cache_tokens_num_ = 0; num_owned_shared_blocks_.clear(); + num_cached_blocks_.clear(); pushed_local_block_count_ = 0; composite_blocks_.clear(); src_blocks_.clear(); diff --git a/xllm/core/framework/request/sequence_kv_state.h b/xllm/core/framework/request/sequence_kv_state.h index e5ce0a1403..3ae62d4b07 100644 --- a/xllm/core/framework/request/sequence_kv_state.h +++ b/xllm/core/framework/request/sequence_kv_state.h @@ -63,6 +63,15 @@ class KVCacheState { void add_shared_blocks(BlockType type, std::vector&& blocks, size_t current_total_num_tokens); + // Composite mount for DSV4 admission: install the (possibly gap-containing) + // shared block vector for `type` at logical positions [0, blocks.size()), + // set shared_blocks_num[type] and num_cached_blocks[type] to blocks.size() + // (the mounted blocks are already in the prefix cache -- no need to re-insert + // on the next pre-grow hook). Does NOT touch kv_cache_tokens_num_; the + // composite advances that once after all leaves have mounted, so all leaves + // observe a consistent shared-token count. + void mount_composite_shared(BlockType type, + std::vector&& shared_blocks); void incr_shared_blocks_num(BlockType type, size_t num); // Drop all blocks held under `type` (releases their Block refs and removes // the map entry). @@ -74,6 +83,18 @@ class KVCacheState { // same across block types, so it takes no BlockType. size_t shared_tokens_num() const; + // Pre-grow cache cursor: how many blocks under `type` have already been + // inserted into the prefix cache. The composite's pre-grow hook consults + // this to skip blocks already in the cache and only stamp+insert the delta + // that has been forwarded since the last hook run. Grows monotonically: + // - Admission mount: set to shared_blocks.size() (mounted blocks are + // already cache-resident, so no re-insert on the next pre-grow). + // - Pre-grow hook: after inserting a run [cursor, end), advance cursor to + // `end`. + // - reset(): cleared alongside the rest of the sequence's cache state. + size_t num_cached_blocks(BlockType type) const; + void set_num_cached_blocks(BlockType type, size_t n); + void set_slice_window_size(uint32_t size); void update_slice_window_pos(); @@ -199,6 +220,9 @@ class KVCacheState { // shared blocks number per block type. std::map num_owned_shared_blocks_; + // Pre-grow cache insert cursor per block type. See num_cached_blocks() above. + std::map num_cached_blocks_; + // Sliding-window cursor for legacy callers. CompositeBlockManager keeps DSA // SWA block vectors in absolute logical block order and leaves expired // logical positions invalid. diff --git a/xllm/core/framework/xtensor/xtensor_block_manager_impl.cpp b/xllm/core/framework/xtensor/xtensor_block_manager_impl.cpp index 793937e570..6c741efca6 100644 --- a/xllm/core/framework/xtensor/xtensor_block_manager_impl.cpp +++ b/xllm/core/framework/xtensor/xtensor_block_manager_impl.cpp @@ -269,8 +269,7 @@ std::vector XTensorBlockManagerImpl::allocate_shared( const Slice& /*token_ids*/, const Slice& /*existed_shared_blocks*/, const MMData& /*mm_data*/, - const Slice& /*block_hashes*/, - size_t* /*matched_tokens*/) { + const Slice& /*block_hashes*/) { // Prefix cache not supported VLOG(1) << "allocate_shared called but prefix cache is not supported"; return {}; diff --git a/xllm/core/framework/xtensor/xtensor_block_manager_impl.h b/xllm/core/framework/xtensor/xtensor_block_manager_impl.h index aaa693e145..df4ae23040 100644 --- a/xllm/core/framework/xtensor/xtensor_block_manager_impl.h +++ b/xllm/core/framework/xtensor/xtensor_block_manager_impl.h @@ -76,8 +76,7 @@ class XTensorBlockManagerImpl : public BlockManager { const Slice& token_ids, const Slice& existed_shared_blocks = {}, const MMData& mm_data = MMData(), - const Slice& block_hashes = {}, - size_t* matched_tokens = nullptr) override; + const Slice& block_hashes = {}) override; // Cache blocks (prefix cache not supported) void cache(const Slice& token_ids, diff --git a/xllm/core/scheduler/disagg_pd_chunked_prefill_scheduler.cpp b/xllm/core/scheduler/disagg_pd_chunked_prefill_scheduler.cpp index da75791fab..78be5122f4 100644 --- a/xllm/core/scheduler/disagg_pd_chunked_prefill_scheduler.cpp +++ b/xllm/core/scheduler/disagg_pd_chunked_prefill_scheduler.cpp @@ -102,10 +102,17 @@ void DisaggPDChunkedPrefillScheduler::match_prefix_blocks(Sequence* sequence) { return; } - if (sequence->kv_state().num_blocks(BlockType::KV) == 0) { + if (!sequence->kv_state().has_any_blocks()) { kv_cache_manager_->allocate_shared(sequence); return; } + // DSV4 (SWA_COMPRESSED) holds SWA/C4/C128 but never a KV leaf, so a + // num_blocks(KV)==0 guard alone would treat an already-mounted DSV4 sequence + // as fresh and re-run allocate_shared -> mount_composite_shared CHECK. Skip + // re-match for the KV-less composite; only flat-KV shapes re-match below. + if (sequence->kv_state().num_blocks(BlockType::KV) == 0) { + return; + } if (!sequence->is_chunked_prefill_stage()) { return; } diff --git a/xllm/core/scheduler/scheduler_policy.cpp b/xllm/core/scheduler/scheduler_policy.cpp index 9cf7bdb12b..4ccf30349a 100644 --- a/xllm/core/scheduler/scheduler_policy.cpp +++ b/xllm/core/scheduler/scheduler_policy.cpp @@ -416,10 +416,17 @@ bool SchedulerPolicy::allocate_for_prefill(Sequence* seq, void SchedulerPolicy::allocate_shared_blocks_for(Sequence* seq, SchedulerState& state) { - if (seq->kv_state().num_blocks(BlockType::KV) == 0) { + if (!seq->kv_state().has_any_blocks()) { state.kv_cache_manager->allocate_shared(seq); return; } + // DSV4 (SWA_COMPRESSED) holds SWA/C4/C128 but never a KV leaf, so a + // num_blocks(KV)==0 guard alone would treat an already-mounted DSV4 sequence + // as fresh and re-run allocate_shared -> mount_composite_shared CHECK. Skip + // re-match for the KV-less composite; only flat-KV shapes re-match below. + if (seq->kv_state().num_blocks(BlockType::KV) == 0) { + return; + } if (seq->is_chunked_prefill_stage()) { if (state.has_linear_attention_layers && state.enable_prefix_cache) { // Linear-state prefix cache can only resume at saved state checkpoints. diff --git a/xllm/core/scheduler/zero_eviction_scheduler.cpp b/xllm/core/scheduler/zero_eviction_scheduler.cpp index 6f7a208762..ee8f45a109 100644 --- a/xllm/core/scheduler/zero_eviction_scheduler.cpp +++ b/xllm/core/scheduler/zero_eviction_scheduler.cpp @@ -121,7 +121,13 @@ void BlockCapacityGuard::compute_reserved_block_num() { void BlockCapacityGuard::prefix_cache_for_candidate_sequences() { if (::xllm::KVCacheConfig::get_instance().enable_prefix_cache()) { for (auto* sequence : candidate_sequences_) { - kv_cache_manager_->allocate_shared(sequence); + // A candidate may be re-evaluated across capacity-guard passes after it + // already mounted its shared prefix. Re-running allocate_shared on a + // sequence that holds cache-bearing blocks re-mounts and trips + // mount_composite_shared's CHECK (DSV4), so match only fresh sequences. + if (!sequence->kv_state().has_any_blocks()) { + kv_cache_manager_->allocate_shared(sequence); + } } } }