From f9bfa79a86347c248772c1cf9c33fdb01bdbb266 Mon Sep 17 00:00:00 2001 From: yingxudeng Date: Wed, 15 Jul 2026 15:09:39 +0800 Subject: [PATCH] feat: add linear-state prefix cache batch and scheduler plumbing (5/n). --- .../batch/batch_packed_input_test.cpp | 93 ++++++++ tests/core/framework/batch/batch_test.cpp | 172 +++++++++++++- .../request_prefix_cache_tokens_test.cpp | 3 + .../core/runtime/acl_graph_executor_test.cpp | 41 +++- tests/core/scheduler/CMakeLists.txt | 1 + .../chunked_prefill_scheduler_test.cpp | 219 ++++++++++++++++-- .../scheduler/continuous_scheduler_test.cpp | 5 +- ...sagg_pd_chunked_prefill_scheduler_test.cpp | 7 +- .../scheduler/disagg_pd_scheduler_test.cpp | 7 +- .../scheduler/fixed_steps_scheduler_test.cpp | 1 + .../profile/profile_graph_warmup_test.cpp | 2 +- xllm/core/framework/batch/CMakeLists.txt | 1 + .../framework/batch/batch_input_builder.cpp | 114 ++++++++- .../framework/batch/batch_input_builder.h | 8 + .../rec_multi_round_batch_input_builder.cpp | 6 +- .../scheduler/chunked_prefill_scheduler.cpp | 38 ++- .../scheduler/chunked_prefill_scheduler.h | 9 + xllm/core/scheduler/continuous_scheduler.cpp | 3 + xllm/core/scheduler/continuous_scheduler.h | 5 + xllm/core/scheduler/mix_scheduler.cpp | 15 +- .../core/scheduler/prefill_only_scheduler.cpp | 28 ++- 21 files changed, 738 insertions(+), 40 deletions(-) diff --git a/tests/core/framework/batch/batch_packed_input_test.cpp b/tests/core/framework/batch/batch_packed_input_test.cpp index fdcba9716a..3c6c686baf 100644 --- a/tests/core/framework/batch/batch_packed_input_test.cpp +++ b/tests/core/framework/batch/batch_packed_input_test.cpp @@ -16,17 +16,32 @@ limitations under the License. #include +#include +#include +#include #include #include "core/common/global_flags.h" #include "core/framework/batch/batch_input_builder.h" #include "core/framework/block/block_manager_impl.h" +#include "core/framework/model/model_input_params.h" #include "core/framework/request/stopping_checker.h" #include "core/runtime/forward_params.h" #include "core/runtime/params_utils.h" namespace xllm { +namespace { + +void expect_linear_state_cache_op_eq(const LinearStateCacheOp& actual, + const LinearStateCacheOp& expected) { + EXPECT_EQ(actual.linear_state_id, expected.linear_state_id); + EXPECT_EQ(actual.restore_requested, expected.restore_requested); + EXPECT_EQ(actual.restore_src_slot_id, expected.restore_src_slot_id); +} + +} // namespace + template bool tensor_equals_vector(const torch::Tensor& tensor, const std::vector& values) { @@ -42,6 +57,84 @@ bool tensor_equals_vector(const torch::Tensor& tensor, return true; } +TEST(BatchPackedInputTest, PackedProtoLazyUnpackPreservesLinearStateCacheOps) { + RequestSamplingParam sampling_param; + sampling_param.logprobs = true; + + StoppingChecker stopping_checker; + stopping_checker.set_max_generated_tokens(4); + + SequenceParams seq_params; + seq_params.seq_capacity = 32; + seq_params.stopping_checker = &stopping_checker; + seq_params.sampling_param = &sampling_param; + seq_params.skip_special_tokens = true; + seq_params.echo = false; + seq_params.logprobs = true; + seq_params.enable_schedule_overlap = true; + + torch::Tensor input_embedding; + MMData mm_data; + BlockManager::Options options; + options.num_blocks(2).block_size(4); + BlockManagerImpl manager(options); + + IncrementalDecoder decoder("", 1, false, false); + Sequence seq(/*index=*/0, + /*token_ids=*/{1, 2, 3, 4}, + input_embedding, + mm_data, + std::move(decoder), + seq_params); + + seq.add_blocks(BlockType::KV, manager.allocate(1)); + + std::vector sequences = {&seq}; + std::vector budgets = {4}; + BatchInputBuilder builder(sequences, + budgets, + {}, + {}, + nullptr, + /*batch_id=*/1, + nullptr, + BatchForwardType::DECODE); + + ForwardInput input = + builder.build_forward_input(/*num_decoding_tokens=*/1, + /*min_decoding_batch_size=*/0); + LinearStateCacheOp restore_op; + restore_op.linear_state_id = 7; + restore_op.restore_requested = true; + restore_op.restore_src_slot_id = 3; + + LinearStateCacheOp no_restore_op; + no_restore_op.linear_state_id = 8; + + input.input_params.linear_state_cache_ops = {restore_op, no_restore_op}; + + proto::PackedForwardInput packed_input; + ASSERT_TRUE(forward_input_to_packed_proto(input, &packed_input)); + + ForwardInput lazy_input; + packed_proto_to_forward_input( + packed_input, lazy_input, torch::Device(torch::kCPU), nullptr); + EXPECT_TRUE(lazy_input.input_params.linear_state_cache_ops.empty()); + EXPECT_TRUE(lazy_input.input_host_buffer_has_layout); + + ForwardInput unpacked_input; + ASSERT_TRUE(detail::unpack_from_input_host_buffer(lazy_input, + torch::Device(torch::kCPU), + torch::kFloat32, + unpacked_input, + false)); + ASSERT_EQ(unpacked_input.input_params.linear_state_cache_ops.size(), 2u); + expect_linear_state_cache_op_eq( + unpacked_input.input_params.linear_state_cache_ops[0], restore_op); + expect_linear_state_cache_op_eq( + unpacked_input.input_params.linear_state_cache_ops[1], no_restore_op); +} + TEST(BatchPackedInputTest, PackedProtoLazyUnpackRestoresSampleIdxes) { RequestSamplingParam sampling_param; sampling_param.logprobs = true; diff --git a/tests/core/framework/batch/batch_test.cpp b/tests/core/framework/batch/batch_test.cpp index 8e0beaf8bb..9f8dd4fced 100644 --- a/tests/core/framework/batch/batch_test.cpp +++ b/tests/core/framework/batch/batch_test.cpp @@ -33,13 +33,13 @@ limitations under the License. #include "framework/config/beam_search_config.h" #include "framework/kv_cache/kv_cache.h" #include "framework/model/model_args.h" +#include "framework/prefix_cache/block_hasher.h" #include "framework/request/stopping_checker.h" #include "framework/sampling/sampling_params.h" #include "platform/device.h" #include "platform/platform.h" #include "runtime/cp_input_partition.h" #include "runtime/forward_shared_memory_manager.h" -#include "runtime/params_utils.h" #include "util/tensor_helper.h" namespace xllm { @@ -121,6 +121,33 @@ void expect_blocks(const TransferKVInfo& info, EXPECT_EQ(info.dp_rank, 3); } +LinearStatePrefixHash compute_linear_state_prefix_hash_for_test( + const std::vector& token_ids, + std::vector& blocks, + size_t boundary_tokens) { + LinearStatePrefixHash hash{}; + if (blocks.empty()) { + return hash; + } + uint32_t block_size = blocks[0].size(); + if (block_size == 0 || boundary_tokens % block_size != 0) { + return hash; + } + size_t boundary_blocks = boundary_tokens / block_size; + if (boundary_blocks == 0 || boundary_tokens > token_ids.size()) { + return hash; + } + const uint8_t* previous_hash = nullptr; + for (size_t block_idx = 0; block_idx < boundary_blocks; ++block_idx) { + xxh3_128bits_hash(previous_hash, + Slice(token_ids).slice( + block_idx * block_size, (block_idx + 1) * block_size), + hash.data()); + previous_hash = hash.data(); + } + return hash; +} + Sequence make_basic_sequence(const std::vector& prompt_token_ids) { static RequestSamplingParam sampling_param; static StoppingChecker stopping_checker; @@ -348,6 +375,7 @@ TEST(BatchTest, ProcessRawOutputStoresMtpBootstrapEmbedding) { TEST(BatchTest, DecodeForwardInputConsumesMtpBootstrap) { BlockManagerPool::Options options; options.num_blocks(8).block_size(4).enable_disagg_pd(true); + options.max_seqs_per_batch(1024); BlockManagerPool manager(options, /*dp_size=*/1); Sequence sequence = make_basic_sequence({1, 2, 3}); @@ -374,6 +402,7 @@ TEST(BatchTest, DecodeForwardInputConsumesMtpBootstrap) { TEST(BatchTest, DecodeForwardInputMapsSparseMtpBootstrapRows) { BlockManagerPool::Options options; options.num_blocks(8).block_size(4).enable_disagg_pd(true); + options.max_seqs_per_batch(1024); BlockManagerPool manager(options, /*dp_size=*/1); Sequence first = make_basic_sequence({1, 2, 3}); @@ -1292,6 +1321,139 @@ TEST(BatchTest, DecodeSingleBlockIdsStaySplitInTransportButShareSlotValue) { expected_slot_id); } +TEST(BatchTest, LinearStateCheckpointSavesOnlyAtPrefillStepBoundary) { + // Linear-state checkpoints are hashed per chunk-end boundary. This test's + // save/restore boundaries (16, 20) are multiples of the KV block_size (4), + // so set the chunk stride to block_size to keep them chunk-aligned and the + // per-chunk hash chain identical to the per-block helper below. + const int32_t prev_chunk_stride = + SchedulerConfig::get_instance().max_tokens_per_chunk_for_prefill(); + SchedulerConfig::get_instance().max_tokens_per_chunk_for_prefill(4); + + torch::Device device(Platform::type_torch(), 0); + const uint32_t n_blocks = 24; + const uint32_t block_size = 4; + BlockManager::Options options; + options.num_blocks(n_blocks).block_size(block_size); + BlockManagerImpl manager(options); + + RequestSamplingParam sampling_param; + StoppingChecker stopping_checker; + stopping_checker.set_max_generated_tokens(4); + SequenceParams seq_params; + seq_params.seq_capacity = 32; + seq_params.stopping_checker = &stopping_checker; + seq_params.sampling_param = &sampling_param; + seq_params.skip_special_tokens = true; + seq_params.echo = false; + seq_params.logprobs = false; + seq_params.enable_schedule_overlap = false; + + torch::Tensor input_embedding; + MMData mm_data; + const std::vector aligned_tokens = { + 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20}; + IncrementalDecoder aligned_decoder("", 20, false, false); + Sequence aligned_seq(/*index=*/0, + aligned_tokens, + input_embedding, + mm_data, + std::move(aligned_decoder), + seq_params); + std::vector aligned_blocks = manager.allocate(5); + aligned_seq.add_blocks(BlockType::KV, aligned_blocks); + + IncrementalDecoder restore_decoder("", 20, false, false); + Sequence restore_seq(/*index=*/1, + aligned_tokens, + input_embedding, + mm_data, + std::move(restore_decoder), + seq_params); + std::vector restore_blocks = manager.allocate(5); + restore_seq.add_blocks(BlockType::KV, restore_blocks); + restore_seq.kv_state().incr_kv_cache_tokens_num(/*size=*/16); + // A sequence that restores at a chunk boundary carries a mounted restore + // source: production mounts it in allocate_shared_for_sequence (class A) when + // the reused KV prefix maps onto a committed linear-state checkpoint. This + // hand-built sequence bypasses the block manager, so mount an explicit source + // slot to reproduce that invariant -- the builder keys the restore emission + // off its presence. + std::vector restore_src_blocks = manager.allocate(1); + restore_seq.set_linear_restore_src_block(std::move(restore_src_blocks[0])); + + const std::vector off_boundary_tokens = {21, 22, 23, 24, 25, 26, 27, + 28, 29, 30, 31, 32, 33, 34, + 35, 36, 37, 38, 39, 40}; + IncrementalDecoder off_boundary_decoder("", 20, false, false); + Sequence off_boundary_seq(/*index=*/2, + off_boundary_tokens, + input_embedding, + mm_data, + std::move(off_boundary_decoder), + seq_params); + off_boundary_seq.add_blocks(BlockType::KV, manager.allocate(5)); + + const std::vector decode_tokens = { + 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56}; + IncrementalDecoder decode_decoder("", 16, false, false); + Sequence decode_seq(/*index=*/3, + decode_tokens, + input_embedding, + mm_data, + std::move(decode_decoder), + seq_params); + decode_seq.add_blocks(BlockType::KV, manager.allocate(5)); + decode_seq.kv_state().incr_kv_cache_tokens_num(/*size=*/16); + decode_seq.append_token(57); + + Batch batch; + batch.add(&aligned_seq, /*allowed_max_token=*/16); + batch.add(&restore_seq, /*allowed_max_token=*/4); + batch.add(&off_boundary_seq, /*allowed_max_token=*/15); + batch.add(&decode_seq, /*allowed_max_token=*/1); + + ModelArgs args; + args.layer_types({"linear_attention"}); + ForwardInput forward_input = batch.prepare_forward_input( + /*num_decoding_tokens=*/1, /*min_decoding_bach_size=*/0, args); + + const auto& cache_ops = forward_input.input_params.linear_state_cache_ops; + ASSERT_EQ(cache_ops.size(), 4u); + // The save decision now lands on the sequence (pending save), not the cache + // op; the LINEAR leaf executes it at the next step. The restore no longer + // rides a hash on the cache op: the builder resolves the checkpoint to a + // source slot and sets restore_requested, mirroring KV's resolved swap + // descriptor (the worker's copy-in keys off that bit + src slot). + const LinearStatePrefixHash aligned_save_expected = + compute_linear_state_prefix_hash_for_test( + aligned_tokens, aligned_blocks, /*boundary_tokens=*/16); + EXPECT_FALSE(is_zero_prefix_hash(aligned_save_expected)); + std::optional aligned_save = aligned_seq.take_pending_linear_save(); + ASSERT_TRUE(aligned_save.has_value()); + EXPECT_EQ(*aligned_save, XXH3Key(aligned_save_expected.data())); + + // restore_seq restores at boundary 16 (the same hash aligned_seq saved) and + // saves at boundary 20. + EXPECT_TRUE(cache_ops[1].restore_requested); + EXPECT_GE(cache_ops[1].restore_src_slot_id, 0); + const LinearStatePrefixHash restore_save_expected = + compute_linear_state_prefix_hash_for_test( + aligned_tokens, restore_blocks, /*boundary_tokens=*/20); + std::optional restore_save = restore_seq.take_pending_linear_save(); + ASSERT_TRUE(restore_save.has_value()); + EXPECT_EQ(*restore_save, XXH3Key(restore_save_expected.data())); + + // off_boundary_seq and decode_seq neither restore nor save. + EXPECT_FALSE(cache_ops[2].restore_requested); + EXPECT_FALSE(off_boundary_seq.has_pending_linear_save()); + EXPECT_FALSE(cache_ops[3].restore_requested); + EXPECT_FALSE(decode_seq.has_pending_linear_save()); + + SchedulerConfig::get_instance().max_tokens_per_chunk_for_prefill( + prev_chunk_stride); +} + TEST(BatchTest, SharedMemoryRoundTripPreservesLinearStateIds) { ForwardInput forward_input; auto int_options = torch::TensorOptions() @@ -1346,6 +1508,14 @@ TEST(BatchTest, SharedMemoryRoundTripPreservesLinearStateIds) { reader_manager.input_read(from_shm, torch::Device(torch::kCPU)); EXPECT_EQ(from_shm.input_params.embedding.linear_state_ids, std::vector({4, 6})); + + forward_input.input_params.embedding.linear_state_ids.clear(); + ASSERT_TRUE(writer_manager.input_write(forward_input)); + + ForwardInput legacy_from_shm; + reader_manager.input_read(legacy_from_shm, torch::Device(torch::kCPU)); + EXPECT_EQ(legacy_from_shm.input_params.embedding.linear_state_ids, + std::vector({-1, -1})); } TEST(BatchTest, SharedMemoryRoundTripPreservesEmptyRankTensors) { diff --git a/tests/core/framework/request/request_prefix_cache_tokens_test.cpp b/tests/core/framework/request/request_prefix_cache_tokens_test.cpp index 6e66226c56..36c6586ecb 100644 --- a/tests/core/framework/request/request_prefix_cache_tokens_test.cpp +++ b/tests/core/framework/request/request_prefix_cache_tokens_test.cpp @@ -114,6 +114,7 @@ TEST(RequestCachedTokensTest, RecordsCachedTokensFromPrefixCache) { BlockManagerPool::Options options; options.num_blocks(16).host_num_blocks(0).block_size(4).enable_prefix_cache( true); + options.max_seqs_per_batch(1024); BlockManagerPool pool(options, /*dp_size=*/1); { @@ -145,6 +146,7 @@ TEST(RequestCachedTokensTest, CachedTokensSurvivesKVBlockRelease) { BlockManagerPool::Options options; options.num_blocks(16).host_num_blocks(0).block_size(4).enable_prefix_cache( true); + options.max_seqs_per_batch(1024); BlockManagerPool pool(options, /*dp_size=*/1); std::shared_ptr request; @@ -197,6 +199,7 @@ TEST(RequestCachedTokensTest, RecordIsIdempotentWithMaxSemantics) { BlockManagerPool::Options options; options.num_blocks(16).host_num_blocks(0).block_size(4).enable_prefix_cache( true); + options.max_seqs_per_batch(1024); BlockManagerPool pool(options, /*dp_size=*/1); std::shared_ptr request; diff --git a/tests/core/runtime/acl_graph_executor_test.cpp b/tests/core/runtime/acl_graph_executor_test.cpp index 0bc704e73f..7db963de46 100644 --- a/tests/core/runtime/acl_graph_executor_test.cpp +++ b/tests/core/runtime/acl_graph_executor_test.cpp @@ -21,6 +21,7 @@ limitations under the License. #include #include +#include #include #include "core/framework/batch/batch.h" @@ -29,6 +30,7 @@ limitations under the License. #include "core/framework/config/execution_config.h" #include "core/framework/config/speculative_config.h" #include "core/framework/kv_cache/kv_cache.h" +#include "core/framework/kv_cache/kv_cache_utils.h" #include "core/framework/kv_cache/linear_state_restore.h" #include "core/framework/model/model_args.h" #include "core/framework/model/model_output.h" @@ -735,11 +737,20 @@ TEST_F(AclGraphExecutorTest, BatchInputCarriesLinearStateIds) { ASSERT_FALSE(batch->empty()); ASSERT_FALSE(sequences_.empty()); + // embedding_ids come from the single_block slot, while linear_state_ids come + // from the dedicated linear_state_slot slot; the two are decoupled and carry + // independent ids through transport. auto& seq = sequences_.back(); - auto linear_state_block = block_manager_->allocate(1); - ASSERT_EQ(linear_state_block.size(), 1); - const int32_t expected_linear_state_id = linear_state_block[0].id(); - seq.add_blocks(BlockType::SINGLE, linear_state_block); + auto embedding_block = block_manager_->allocate(1); + ASSERT_EQ(embedding_block.size(), 1); + const int32_t expected_embedding_id = embedding_block[0].id(); + seq.add_blocks(BlockType::SINGLE, embedding_block); + + auto linear_state_slot = block_manager_->allocate(1); + ASSERT_EQ(linear_state_slot.size(), 1); + const int32_t expected_linear_state_id = linear_state_slot[0].id(); + seq.add_blocks(BlockType::LINEAR, linear_state_slot); + ASSERT_NE(expected_embedding_id, expected_linear_state_id); auto forward_input = batch->prepare_forward_input( options_.num_decoding_tokens(), 0, model_args_); @@ -749,7 +760,27 @@ TEST_F(AclGraphExecutorTest, BatchInputCarriesLinearStateIds) { expected_linear_state_id); ASSERT_EQ(forward_input.input_params.embedding.embedding_ids.size(), 1); EXPECT_EQ(forward_input.input_params.embedding.embedding_ids[0], - expected_linear_state_id); + expected_embedding_id); + + forward_input.input_params.parallel.has_initial_state = {0}; + forward_input = forward_input.to(*device_, torch::kFloat32); + npu::GraphPersistentParam persistent_param(model_args_, *device_, options_); + std::optional params_for_capture = persistent_param.update( + forward_input.token_ids, + first_full_attention_cache(kv_caches_).get_k_cache(), + first_full_attention_cache(kv_caches_).get_v_cache(), + forward_input.positions, + forward_input.input_params, + /*padded_num_tokens=*/2, + /*return_capture_params=*/true); + ASSERT_TRUE(params_for_capture.has_value()); + EXPECT_EQ(params_for_capture->meta.num_sequences, 2); + EXPECT_EQ( + params_for_capture->embedding.linear_state_ids, + std::vector({expected_linear_state_id, kPaddingLinearStateId})); + ASSERT_EQ(params_for_capture->parallel.has_initial_state.size(), 2); + EXPECT_EQ(params_for_capture->parallel.has_initial_state[0], 0); + EXPECT_EQ(params_for_capture->parallel.has_initial_state[1], 0); } TEST(LinearStateRestoreTest, RestoreCopiesCheckpointSlot) { diff --git a/tests/core/scheduler/CMakeLists.txt b/tests/core/scheduler/CMakeLists.txt index 92529b019e..4506ccae05 100644 --- a/tests/core/scheduler/CMakeLists.txt +++ b/tests/core/scheduler/CMakeLists.txt @@ -16,6 +16,7 @@ cc_test( :config :scheduler :xllm_server + :block GTest::gtest_main $<$:nnopbase> ) diff --git a/tests/core/scheduler/chunked_prefill_scheduler_test.cpp b/tests/core/scheduler/chunked_prefill_scheduler_test.cpp index 62e3c257f1..445aebece3 100644 --- a/tests/core/scheduler/chunked_prefill_scheduler_test.cpp +++ b/tests/core/scheduler/chunked_prefill_scheduler_test.cpp @@ -19,15 +19,39 @@ limitations under the License. #include #include +#include #include +#include +#include +#include +#include "common/global_flags.h" #include "core/framework/config/scheduler_config.h" #include "distributed_runtime/engine.h" +#include "framework/block/block_manager_pool_test_peer.h" +#include "framework/block/linear_state_block_manager.h" #include "util/utils.h" namespace xllm { namespace { + +template +class ScopedValue final { + public: + ScopedValue(T* target, T value) : target_(target), old_(*target) { + *target_ = value; + } + ~ScopedValue() { *target_ = old_; } + + ScopedValue(const ScopedValue&) = delete; + ScopedValue& operator=(const ScopedValue&) = delete; + + private: + T* target_; + T old_; +}; + class FakeTokenizer : public Tokenizer { public: bool encode(const std::string_view& text, @@ -51,11 +75,27 @@ class FakeTokenizer : public Tokenizer { class FakeEngine : public Engine { public: - FakeEngine(int32_t num_blocks, int32_t block_size) { + FakeEngine(int32_t num_blocks, + int32_t block_size, + bool enable_linear_attention = false, + bool enable_prefix_cache = false, + int32_t linear_chunk_stride = -1) { BlockManagerPool::Options opt; opt.num_blocks_ = num_blocks; opt.block_size_ = block_size; - opt.enable_prefix_cache_ = false; // we dont consider prefix cache here + opt.max_seqs_per_batch_ = 1024; + opt.enable_prefix_cache_ = enable_prefix_cache; + opt.enable_linear_state_ = enable_linear_attention; + if (enable_linear_attention) { + model_args_.layer_types({"linear_attention"}); + // The unified linear-state slot pool needs a positive physical capacity; + // size it generously so tests never hit slot pressure. + opt.linear_state_num_slots_ = num_blocks + 2; + // Chunk stride the LINEAR checkpoint index probes with. Mirrors the + // engine, which captures it from the scheduler config at construction so + // the override never reads a global singleton. + opt.linear_chunk_stride_ = linear_chunk_stride; + } fake_tokenizer_ = std::make_unique(); fake_block_manager_ = std::make_unique(opt, 1); } @@ -65,7 +105,7 @@ class FakeEngine : public Engine { BlockManagerPool* block_manager_pool() const { return fake_block_manager_.get(); } - const ModelArgs& model_args() const { NOT_IMPLEMENTED(); } + const ModelArgs& model_args() const { return model_args_; } const TokenizerArgs& tokenizer_args() const { NOT_IMPLEMENTED(); } std::vector get_active_activation_memory() const { NOT_IMPLEMENTED(); @@ -75,20 +115,17 @@ class FakeEngine : public Engine { private: std::unique_ptr fake_tokenizer_; std::unique_ptr fake_block_manager_; + ModelArgs model_args_; }; -template -class ScopedConfigValue final { +class TestableChunkedPrefillScheduler : public ChunkedPrefillScheduler { public: - ScopedConfigValue(T& value, T new_value) : value_(value), old_(value) { - value_ = new_value; - } + using ChunkedPrefillScheduler::allocate_shared_blocks_for; + using ChunkedPrefillScheduler::get_prefill_handle_tokens; - ~ScopedConfigValue() { value_ = old_; } - - private: - T& value_; - T old_; + TestableChunkedPrefillScheduler(Engine* engine, + const ContinuousScheduler::Options& options) + : ChunkedPrefillScheduler(engine, options) {} }; ContinuousScheduler::Options create_scheduler_options( @@ -486,9 +523,9 @@ TEST(ChunkedPrefillSchedulerTest, OnPrefillPreemptOffDecode) { // set chunked max_tokens budgets 10000 per step ContinuousScheduler::Options opt = create_scheduler_options( 10000, 256, 0, max_tokens_per_chunk_for_prefill, 1); - ScopedConfigValue memory_threshold( - SchedulerConfig::get_instance() - .prefill_scheduling_memory_usage_threshold(), + ScopedValue memory_threshold( + &SchedulerConfig::get_instance() + .prefill_scheduling_memory_usage_threshold(), 2.0); { @@ -714,4 +751,154 @@ TEST(ChunkedPrefillSchedulerTest, LatencySchedule) { EXPECT_EQ(running_sequences_budgets[3], max_tokens_per_chunk_for_prefill); } +TEST(ChunkedPrefillSchedulerTest, LinearStateSafePrefillTokens) { + struct TestCase { + const char* name; + int32_t block_num; + int32_t block_size; + int32_t max_tokens_per_batch; + int32_t max_tokens_per_chunk_for_prefill; + int32_t prompt_len; + size_t allocated_tokens; + size_t cached_tokens; + size_t token_budget; + size_t expected_handle_tokens; + }; + + const std::vector test_cases = { + {"DefersWhenBudgetCannotReachBlockBoundary", + /*block_num=*/32, + /*block_size=*/128, + /*max_tokens_per_batch=*/64, + /*max_tokens_per_chunk_for_prefill=*/64, + /*prompt_len=*/512, + /*allocated_tokens=*/128, + /*cached_tokens=*/128, + /*token_budget=*/64, + /*expected_handle_tokens=*/0}, + {"StopsAtLastBlockBoundaryBeforeTail", + /*block_num=*/32, + /*block_size=*/128, + /*max_tokens_per_batch=*/2048, + /*max_tokens_per_chunk_for_prefill=*/2048, + /*prompt_len=*/1200, + /*allocated_tokens=*/0, + /*cached_tokens=*/0, + /*token_budget=*/2048, + /*expected_handle_tokens=*/1152}, + {"KeepsChunkBudgetAtBlockBoundary", + /*block_num=*/32, + /*block_size=*/4, + /*max_tokens_per_batch=*/64, + /*max_tokens_per_chunk_for_prefill=*/16, + /*prompt_len=*/64, + /*allocated_tokens=*/0, + /*cached_tokens=*/0, + /*token_budget=*/16, + /*expected_handle_tokens=*/16}, + {"RoundsChunkBudgetDownToBlockBoundary", + /*block_num=*/32, + /*block_size=*/4, + /*max_tokens_per_batch=*/64, + /*max_tokens_per_chunk_for_prefill=*/17, + /*prompt_len=*/64, + /*allocated_tokens=*/0, + /*cached_tokens=*/0, + /*token_budget=*/17, + /*expected_handle_tokens=*/16}, + }; + + ScopedValue prefix_cache_guard(&FLAGS_enable_prefix_cache, true); + for (const TestCase& test_case : test_cases) { + SCOPED_TRACE(test_case.name); + ContinuousScheduler::Options opt = + create_scheduler_options(test_case.max_tokens_per_batch, + /*max_seqs_per_batch=*/256, + /*num_speculative_tokens=*/0, + test_case.max_tokens_per_chunk_for_prefill, + /*dp_size=*/1); + auto engine = std::make_unique(test_case.block_num, + test_case.block_size, + /*enable_linear_attention=*/true, + /*enable_prefix_cache=*/true); + auto scheduler = + std::make_unique(engine.get(), opt); + + auto requests = generate_request( + {test_case.prompt_len}, {10}, std::nullopt, std::nullopt, 30000); + Sequence* sequence = requests[0]->sequences()[0].get(); + if (test_case.allocated_tokens > 0) { + ASSERT_TRUE(engine->block_manager_pool()->allocate( + sequence, test_case.allocated_tokens)); + } + sequence->kv_state().set_kv_cache_tokens_num(test_case.cached_tokens); + + EXPECT_EQ(scheduler->get_prefill_handle_tokens( + sequence, test_case.token_budget, test_case.cached_tokens), + test_case.expected_handle_tokens); + } +} + +TEST(ChunkedPrefillSchedulerTest, + LinearStateChunkedPrefillMatchesAtEveryChunkBoundary) { + ScopedValue prefix_cache_guard(&FLAGS_enable_prefix_cache, true); + ScopedValue match_frequency_guard(&FLAGS_chunked_match_frequency, 2); + + int32_t block_num = 32; + int32_t block_size = 128; + int32_t max_tokens_per_chunk_for_prefill = 128; + ContinuousScheduler::Options opt = + create_scheduler_options(max_tokens_per_chunk_for_prefill, + 256, + 0, + max_tokens_per_chunk_for_prefill, + 1); + // The LINEAR checkpoint index probes its hash domain with the chunk stride + // the engine captures at construction. Thread the same stride into the pool + // here so match() runs with the real domain instead of bailing on the -1 + // default. + auto engine = std::make_unique(block_num, + block_size, + /*enable_linear_attention=*/true, + /*enable_prefix_cache=*/true, + max_tokens_per_chunk_for_prefill); + auto scheduler = + std::make_unique(engine.get(), opt); + EXPECT_TRUE(scheduler != nullptr); + + auto active_requests = + generate_request({512}, {10}, std::nullopt, std::nullopt, 30000); + auto* active_seq = active_requests[0]->sequences()[0].get(); + ASSERT_TRUE(engine->block_manager_pool()->allocate(active_seq, 128)); + active_seq->kv_state().set_kv_cache_tokens_num(128); + + auto cached_requests = + generate_request({256}, {10}, std::nullopt, std::nullopt, 30000); + auto* cached_seq = cached_requests[0]->sequences()[0].get(); + ASSERT_TRUE(engine->block_manager_pool()->allocate(cached_seq, + cached_seq->num_tokens())); + cached_seq->kv_state().set_kv_cache_tokens_num(cached_seq->num_tokens()); + engine->block_manager_pool()->cache(cached_seq); + const XXH3Key checkpoint_hash(cached_seq->kv_state() + .blocks(BlockType::KV)[1] + .get_immutable_hash_value()); + LinearStateBlockManager* linear_leaf = BlockManagerPoolTestPeer::linear_leaf( + *engine->block_manager_pool(), /*dp_rank=*/0); + ASSERT_NE(linear_leaf, nullptr); + Block checkpoint_block = linear_leaf->allocate(); + ASSERT_TRUE(checkpoint_block.is_valid()); + const int32_t checkpoint_slot = checkpoint_block.id(); + ASSERT_GE(checkpoint_slot, 1); + checkpoint_block.set_hash_value(checkpoint_hash.data); + std::vector checkpoint; + checkpoint.emplace_back(std::move(checkpoint_block)); + linear_leaf->cache(checkpoint); + engine->block_manager_pool()->deallocate_without_cache(cached_seq); + + scheduler->allocate_shared_blocks_for(active_seq); + EXPECT_EQ(active_seq->kv_state().shared_blocks_num(BlockType::KV), 2u); + EXPECT_EQ(active_seq->kv_state().num_blocks(BlockType::KV), 2u); + EXPECT_EQ(active_seq->kv_state().kv_cache_tokens_num(), 256u); +} + } // namespace xllm diff --git a/tests/core/scheduler/continuous_scheduler_test.cpp b/tests/core/scheduler/continuous_scheduler_test.cpp index 6535de92ab..8599171295 100644 --- a/tests/core/scheduler/continuous_scheduler_test.cpp +++ b/tests/core/scheduler/continuous_scheduler_test.cpp @@ -13,6 +13,7 @@ #include "core/framework/config/scheduler_config.h" #include "core/platform/platform.h" #include "distributed_runtime/engine.h" +#include "framework/model/model_args.h" #include "prefill_only_scheduler.h" #include "scheduler_factory.h" #include "util/utils.h" @@ -50,6 +51,7 @@ class FakeEngine : public Engine { opt.num_blocks_ = num_blocks; opt.block_size_ = block_size; opt.enable_prefix_cache_ = enable_prefix_cache; + opt.max_seqs_per_batch_ = 1024; fake_tokenizer_ = std::make_unique(); fake_block_manager_ = std::make_unique(opt, 1); } @@ -59,7 +61,7 @@ class FakeEngine : public Engine { BlockManagerPool* block_manager_pool() const { return fake_block_manager_.get(); } - const ModelArgs& model_args() const { NOT_IMPLEMENTED(); } + const ModelArgs& model_args() const { return model_args_; } const TokenizerArgs& tokenizer_args() const { NOT_IMPLEMENTED(); } std::vector get_active_activation_memory() const { return {0}; } bool init() override { return true; } @@ -67,6 +69,7 @@ class FakeEngine : public Engine { private: std::unique_ptr fake_tokenizer_; std::unique_ptr fake_block_manager_; + ModelArgs model_args_; }; class TestContinuousScheduler final : public ContinuousScheduler { diff --git a/tests/core/scheduler/disagg_pd_chunked_prefill_scheduler_test.cpp b/tests/core/scheduler/disagg_pd_chunked_prefill_scheduler_test.cpp index bf925bf5ca..db4ce36e4b 100644 --- a/tests/core/scheduler/disagg_pd_chunked_prefill_scheduler_test.cpp +++ b/tests/core/scheduler/disagg_pd_chunked_prefill_scheduler_test.cpp @@ -27,6 +27,7 @@ limitations under the License. #include "core/framework/config/kv_cache_config.h" #include "distributed_runtime/engine.h" #include "framework/block/block_manager_pool.h" +#include "framework/model/model_args.h" #include "framework/request/request.h" #include "framework/request/request_state.h" #include "framework/request/sequence.h" @@ -42,7 +43,8 @@ BlockManagerPool::Options make_block_options(int32_t num_blocks, options.num_blocks(num_blocks) .block_size(block_size) .enable_prefix_cache(true) - .enable_disagg_pd(true); + .enable_disagg_pd(true) + .max_seqs_per_batch(1024); return options; } @@ -117,7 +119,7 @@ class FakeEngine final : public Engine { return block_manager_.get(); } - const ModelArgs& model_args() const override { NOT_IMPLEMENTED(); } + const ModelArgs& model_args() const override { return model_args_; } const TokenizerArgs& tokenizer_args() const override { NOT_IMPLEMENTED(); } @@ -130,6 +132,7 @@ class FakeEngine final : public Engine { private: std::unique_ptr tokenizer_; std::unique_ptr block_manager_; + ModelArgs model_args_; }; template diff --git a/tests/core/scheduler/disagg_pd_scheduler_test.cpp b/tests/core/scheduler/disagg_pd_scheduler_test.cpp index 539d100ce4..87cb33b365 100644 --- a/tests/core/scheduler/disagg_pd_scheduler_test.cpp +++ b/tests/core/scheduler/disagg_pd_scheduler_test.cpp @@ -27,6 +27,7 @@ limitations under the License. #include "common/metrics.h" #include "distributed_runtime/engine.h" #include "framework/block/block_manager_pool.h" +#include "framework/model/model_args.h" #include "framework/request/request.h" #include "framework/request/request_state.h" #include "framework/tokenizer/tokenizer.h" @@ -68,7 +69,8 @@ class FakeEngine final : public Engine { options.num_blocks(num_blocks) .block_size(block_size) .enable_prefix_cache(true) - .enable_disagg_pd(true); + .enable_disagg_pd(true) + .max_seqs_per_batch(1024); tokenizer_ = std::make_unique(); block_manager_ = std::make_unique(options, /*dp_size=*/1); } @@ -87,7 +89,7 @@ class FakeEngine final : public Engine { return block_manager_.get(); } - const ModelArgs& model_args() const override { NOT_IMPLEMENTED(); } + const ModelArgs& model_args() const override { return model_args_; } const TokenizerArgs& tokenizer_args() const override { NOT_IMPLEMENTED(); } @@ -100,6 +102,7 @@ class FakeEngine final : public Engine { private: std::unique_ptr tokenizer_; std::unique_ptr block_manager_; + ModelArgs model_args_; }; class TestDisaggPDScheduler final : public DisaggPDScheduler { diff --git a/tests/core/scheduler/fixed_steps_scheduler_test.cpp b/tests/core/scheduler/fixed_steps_scheduler_test.cpp index a38b8fc4a1..31262f1ef4 100644 --- a/tests/core/scheduler/fixed_steps_scheduler_test.cpp +++ b/tests/core/scheduler/fixed_steps_scheduler_test.cpp @@ -68,6 +68,7 @@ class FakeEngine : public Engine { opt.num_blocks_ = num_blocks; opt.block_size_ = block_size; opt.enable_prefix_cache_ = false; + opt.max_seqs_per_batch_ = 1024; fake_tokenizer_ = std::make_unique(); fake_block_manager_ = std::make_unique(opt, 1); } diff --git a/tests/core/scheduler/profile/profile_graph_warmup_test.cpp b/tests/core/scheduler/profile/profile_graph_warmup_test.cpp index e118fb23df..2cf4247340 100644 --- a/tests/core/scheduler/profile/profile_graph_warmup_test.cpp +++ b/tests/core/scheduler/profile/profile_graph_warmup_test.cpp @@ -160,7 +160,7 @@ TEST(GraphWarmupTest, PresetDpRankControlsBlockAllocation) { options.num_blocks(/*num_blocks=*/8) .block_size(/*block_size=*/2) .enable_prefix_cache(/*enable_prefix_cache=*/false) - .max_concurrent_requests(/*max_concurrent_requests=*/4); + .max_seqs_per_batch(/*max_seqs_per_batch=*/4); BlockManagerPool pool(options, /*dp_size=*/2); Sequence sequence = make_sequence(/*index=*/0, /*tokens=*/{1, 2, 3}); sequence.set_dp_rank(/*dp_rank=*/1); diff --git a/xllm/core/framework/batch/CMakeLists.txt b/xllm/core/framework/batch/CMakeLists.txt index 33246a27db..dbb022b5b0 100644 --- a/xllm/core/framework/batch/CMakeLists.txt +++ b/xllm/core/framework/batch/CMakeLists.txt @@ -29,4 +29,5 @@ cc_library( :common :mposition glog::glog + xxHash ) diff --git a/xllm/core/framework/batch/batch_input_builder.cpp b/xllm/core/framework/batch/batch_input_builder.cpp index c4165343ce..495673bf75 100644 --- a/xllm/core/framework/batch/batch_input_builder.cpp +++ b/xllm/core/framework/batch/batch_input_builder.cpp @@ -170,6 +170,23 @@ torch::Tensor build_pinned_int_tensor(const std::vector& values) { .pinned_memory(true)); } +// Whether the current prefill step end should hold a linear-state checkpoint. +// Checkpoints are saved at prefill step ends that land on a chunk-end boundary +// (stride = max_tokens_per_chunk_for_prefill). The linear-state cache is a +// sparse per-chunk overlay: KV may cache every block boundary while +// linear-state saves only at chunk ends. +bool should_save_linear_checkpoint(Sequence* sequence, + uint32_t boundary_tokens, + uint32_t chunk_stride) { + if (sequence == nullptr || !sequence->is_prefill_stage()) { + return false; + } + if (boundary_tokens == 0 || chunk_stride == 0) { + return false; + } + return boundary_tokens % chunk_stride == 0; +} + } // namespace BatchInputBuilder::BatchInputBuilder( @@ -432,6 +449,9 @@ void BatchInputBuilder::process_sequences_multithreaded() { state_.linear_state_ids.insert(state_.linear_state_ids.end(), state.linear_state_ids.begin(), state.linear_state_ids.end()); + state_.linear_state_cache_ops.insert(state_.linear_state_cache_ops.end(), + state.linear_state_cache_ops.begin(), + state.linear_state_cache_ops.end()); state_.request_ids.insert(state_.request_ids.end(), state.request_ids.begin(), state.request_ids.end()); @@ -620,9 +640,7 @@ void BatchInputBuilder::extract_tokens_and_positions(Sequence* sequence, } } - // `linear_state_ids` is sequence-scoped metadata and must stay aligned with - // logical batch rows even for non-terminal chunked-prefill slices. - state.linear_state_ids.emplace_back(sequence->get_single_block_id()); + append_linear_state_row(sequence, n_kv_cache_tokens, seq_len, state); // Add extra token id int32_t extra_token_id = -1; @@ -685,6 +703,94 @@ void BatchInputBuilder::extract_tokens_and_positions(Sequence* sequence, } } +void BatchInputBuilder::append_linear_state_row(Sequence* sequence, + uint32_t n_kv_cache_tokens, + uint32_t seq_len, + BuilderState& state) { + // linear_state_ids must stay aligned with logical batch rows even when the + // model has no linear-attention layers, because downstream consumers index by + // batch row. Prefer the dedicated linear-state slot. Linear-attention decode + // paths that only carry a scheduler-side single block still share that live + // slot value across embedding and linear-state transport fields. + const bool has_linear_attention = + args_ && has_linear_attention_layers(*args_); + int32_t linear_state_id = sequence->get_linear_state_slot_id(); + if (linear_state_id < 0 && has_linear_attention) { + linear_state_id = sequence->get_single_block_id(); + } + state.linear_state_ids.emplace_back(linear_state_id); + if (!has_linear_attention) { + return; + } + + LinearStateCacheOp linear_state_cache_op; + linear_state_cache_op.linear_state_id = state.linear_state_ids.back(); + // Linear-state checkpoints live on chunk-end boundaries, so the prefix hash + // is chained per chunk (stride = max_tokens_per_chunk_for_prefill), not per + // KV block. The engine enforces this stride is a positive multiple of + // block_size when linear prefix cache is on (llm_engine.cpp); guard against + // an unset (<= 0) stride so a misconfigured run simply skips cache ops. + const int32_t chunk_stride = ::xllm::SchedulerConfig::get_instance() + .max_tokens_per_chunk_for_prefill(); + // Cold-start restore: emit a restore hash only when a restore source + // checkpoint is mounted on this sequence -- class A at admission + // (allocate_shared_for_sequence) or class B at the previous step's + // save-rotation (allocate_for_sequence) -- AND the reused prefix lands + // on a chunk-end boundary, where the recurrent state lives in a checkpoint. + // A mounted source is present exactly on a slot that is cold and needs + // copy-in; continued forwards keep their live slot warm with no source + // mounted, so they emit no restore and are not reset to cold by the worker. + // The source slot id is taken from that mounted block below. + const bool needs_restore_hash = sequence->has_linear_restore_src_block() && + n_kv_cache_tokens > 0 && chunk_stride > 0 && + n_kv_cache_tokens % chunk_stride == 0; + // Exit-boundary save: persist the live state only when this prefill step + // lands on a chunk-end boundary, so the linear-state cache stays a sparse + // per-chunk overlay on top of the per-block KV cache. + const bool needs_save_hash = + should_save_linear_checkpoint(sequence, seq_len, chunk_stride); + // Refresh the sequence's cached chunk hashes to cover this step's deepest + // boundary, then read them back. The cache is chained and incremental, so + // this only hashes chunks not seen on a previous step; the match probe and + // this builder now share the one hash source instead of each recomputing. + Slice linear_state_hashes; + if (needs_restore_hash || needs_save_hash) { + sequence->update_linear_state_hashes(static_cast(chunk_stride)); + linear_state_hashes = sequence->linear_state_hashes(); + } + // Restore source (block-carried): allocate_shared_for_sequence mounts the + // deepest-hit checkpoint at admission (class A); allocate_for_sequence + // mounts the slot it just checkpointed at the previous step's save-rotation + // (class B). Take it unconditionally so the pin never outlives the step + // that consumed the match; its id is used below only when a restore hash is + // actually emitted, otherwise the handle drops here and releases the pin. + std::optional mounted_restore_src = + sequence->take_linear_restore_src_block(); + if (needs_restore_hash) { + const size_t restore_chunk_idx = + static_cast(n_kv_cache_tokens) / chunk_stride - 1; + if (restore_chunk_idx < linear_state_hashes.size()) { + linear_state_cache_op.restore_requested = true; + if (mounted_restore_src.has_value()) { + linear_state_cache_op.restore_src_slot_id = mounted_restore_src->id(); + } + } + } + if (needs_save_hash) { + const size_t save_chunk_idx = + static_cast(seq_len) / chunk_stride - 1; + if (save_chunk_idx < linear_state_hashes.size()) { + // Record the boundary hash on the sequence. The LINEAR leaf executes + // the save at the next step's allocate_for_sequence, after this step's + // forward writes the boundary state into the live slot. Writing only + // the sequence's own pending-save field keeps this safe inside the + // parallel build loop. + sequence->set_pending_linear_save(linear_state_hashes[save_chunk_idx]); + } + } + state.linear_state_cache_ops.emplace_back(std::move(linear_state_cache_op)); +} + void BatchInputBuilder::handle_sampling_parameters(Sequence* sequence, BuilderState* state_ptr) { BuilderState& state = state_ptr ? *state_ptr : state_; @@ -951,6 +1057,8 @@ ForwardInput BatchInputBuilder::state_to_forward_input() { input_params.embedding.embedding_ids = std::move(state_.embedding_ids); input_params.embedding.linear_state_ids = std::move(state_.linear_state_ids); + input_params.linear_state_cache_ops = + std::move(state_.linear_state_cache_ops); if (!input_params.embedding.linear_state_ids.empty()) { input_params.embedding.linear_state_indices = torch::tensor(input_params.embedding.linear_state_ids, torch::kInt); diff --git a/xllm/core/framework/batch/batch_input_builder.h b/xllm/core/framework/batch/batch_input_builder.h index 5180e708b5..30f9739aca 100644 --- a/xllm/core/framework/batch/batch_input_builder.h +++ b/xllm/core/framework/batch/batch_input_builder.h @@ -123,6 +123,7 @@ class BatchInputBuilder { // Additional data std::vector embedding_ids; std::vector linear_state_ids; + std::vector linear_state_cache_ops; std::vector request_ids; std::vector extra_token_ids; std::vector mtp_shifted_token_ids; @@ -150,6 +151,13 @@ class BatchInputBuilder { uint32_t seq_len, uint32_t padded_seq_len, BuilderState* state_ptr = nullptr); + // Append this batch row's linear-state transport fields: the live slot id + // (always, so rows stay aligned) plus a LinearStateCacheOp carrying the + // resolved restore/save plan for linear-attention models. + void append_linear_state_row(Sequence* sequence, + uint32_t n_kv_cache_tokens, + uint32_t seq_len, + BuilderState& state); torch::Tensor get_mrope_positions(Sequence* sequence, uint32_t start, uint32_t end); diff --git a/xllm/core/framework/batch/rec_multi_round_batch_input_builder.cpp b/xllm/core/framework/batch/rec_multi_round_batch_input_builder.cpp index a71b727132..9ca192ea63 100644 --- a/xllm/core/framework/batch/rec_multi_round_batch_input_builder.cpp +++ b/xllm/core/framework/batch/rec_multi_round_batch_input_builder.cpp @@ -261,8 +261,10 @@ void RecMultiRoundBatchInputBuilder::extract_tokens_and_positions( } // `linear_state_ids` is consumed per sequence, so preserve one entry per - // logical batch row even when this slice is not the last prefill chunk. - base_state.linear_state_ids.push_back(sequence->get_single_block_id()); + // logical batch row even when this slice is not the last prefill chunk. The + // slot is drawn from the dedicated LinearStateBlockManager and is -1 when + // linear state is disabled. + base_state.linear_state_ids.push_back(sequence->get_linear_state_slot_id()); // Add extra token id if (n_tokens == seq_len) { diff --git a/xllm/core/scheduler/chunked_prefill_scheduler.cpp b/xllm/core/scheduler/chunked_prefill_scheduler.cpp index 85f2035eb0..8e8d30acea 100644 --- a/xllm/core/scheduler/chunked_prefill_scheduler.cpp +++ b/xllm/core/scheduler/chunked_prefill_scheduler.cpp @@ -25,7 +25,9 @@ limitations under the License. #include "core/framework/config/parallel_config.h" #include "core/framework/config/scheduler_config.h" #include "core/platform/platform.h" +#include "distributed_runtime/engine.h" #include "framework/batch/batch_factory.h" +#include "framework/model/model_args.h" #include "util/timer.h" #include "util/utils.h" @@ -832,6 +834,30 @@ std::vector ChunkedPrefillScheduler::prepare_batch() { return batches; } +size_t ChunkedPrefillScheduler::get_prefill_handle_tokens( + Sequence* sequence, + size_t token_budget, + size_t kv_cache_tokens_num) const { + size_t max_handle_num_tokens = + std::min(kv_cache_tokens_num + token_budget, sequence->num_tokens()); + if (!has_linear_attention_layers_ || !enable_prefix_cache_ || + !sequence->is_prefill_stage()) { + return max_handle_num_tokens; + } + + const size_t block_size = kv_cache_manager_->block_size(); + const size_t aligned_handle_num_tokens = + block_size == 0 ? max_handle_num_tokens + : (max_handle_num_tokens / block_size) * block_size; + if (aligned_handle_num_tokens <= kv_cache_tokens_num) { + if (max_handle_num_tokens == sequence->num_tokens()) { + return max_handle_num_tokens; + } + return 0; + } + return aligned_handle_num_tokens; +} + bool ChunkedPrefillScheduler::allocate_blocks_for( Sequence* sequence, size_t token_budget, @@ -846,7 +872,10 @@ bool ChunkedPrefillScheduler::allocate_blocks_for( // the total number tokens for the sequence can be handled till now. // there may some tokens can not be handled once when enable chunked prefill. size_t max_handle_num_tokens = - std::min(kv_cache_tokens_num + token_budget, sequence->num_tokens()); + get_prefill_handle_tokens(sequence, token_budget, kv_cache_tokens_num); + if (max_handle_num_tokens == 0) { + return false; + } // speculative decoding specific logic, // prefill stage don't need speculative decoding. @@ -874,6 +903,13 @@ void ChunkedPrefillScheduler::allocate_shared_blocks_for(Sequence* sequence) { return; } if (sequence->is_chunked_prefill_stage()) { + if (has_linear_attention_layers_ && enable_prefix_cache_) { + // Linear-state prefix cache can only resume at saved state checkpoints. + // Re-match at every chunk boundary and let BlockManagerPool truncate the + // hit to checkpoint-backed blocks. + kv_cache_manager_->allocate_shared(sequence); + return; + } const size_t max_tokens_per_chunk_for_prefill = std::max(options_.max_tokens_per_chunk_for_prefill(), 64); size_t total_chunked_size = diff --git a/xllm/core/scheduler/chunked_prefill_scheduler.h b/xllm/core/scheduler/chunked_prefill_scheduler.h index e2a3489c6b..83d931f456 100644 --- a/xllm/core/scheduler/chunked_prefill_scheduler.h +++ b/xllm/core/scheduler/chunked_prefill_scheduler.h @@ -37,6 +37,15 @@ class ChunkedPrefillScheduler : public ContinuousScheduler { virtual ~ChunkedPrefillScheduler(); protected: + // Returns 0 when no safe boundary can make forward progress. + // `kv_cache_tokens_num` is the count of already-processed tokens as seen by + // the caller's scheduling accounting; callers pass their own value so this + // stays correct under host<->device kv swap, where it may exceed the + // sequence's live device kv_cache_tokens_num. + size_t get_prefill_handle_tokens(Sequence* sequence, + size_t token_budget, + size_t kv_cache_tokens_num) const; + // build a batch of requests from the priority queue virtual std::vector prepare_batch() override; // 1. for prefill sequence: the allocated_tokens will be within diff --git a/xllm/core/scheduler/continuous_scheduler.cpp b/xllm/core/scheduler/continuous_scheduler.cpp index a6717426e8..2dbca78784 100644 --- a/xllm/core/scheduler/continuous_scheduler.cpp +++ b/xllm/core/scheduler/continuous_scheduler.cpp @@ -38,6 +38,7 @@ limitations under the License. #include "core/platform/platform.h" #include "distributed_runtime/engine.h" #include "framework/batch/batch_factory.h" +#include "framework/model/model_args.h" #include "framework/request/priority_comparator.h" #include "framework/request/request.h" #include "framework/request/sequence.h" @@ -123,6 +124,8 @@ ContinuousScheduler::ContinuousScheduler(Engine* engine, const Options& options) enable_prefix_cache_ = ::xllm::KVCacheConfig::get_instance().enable_prefix_cache(); + has_linear_attention_layers_ = + ::xllm::has_linear_attention_layers(engine_->model_args()); enable_in_batch_prefix_cache_ = ::xllm::KVCacheConfig::get_instance().enable_in_batch_prefix_cache(); diff --git a/xllm/core/scheduler/continuous_scheduler.h b/xllm/core/scheduler/continuous_scheduler.h index dfd83db819..1efba0ee56 100644 --- a/xllm/core/scheduler/continuous_scheduler.h +++ b/xllm/core/scheduler/continuous_scheduler.h @@ -294,6 +294,11 @@ class ContinuousScheduler : public Scheduler { bool enable_prefix_cache_ = false; + // Cached once at construction (mirrors enable_prefix_cache_): the model's + // layer set is immutable, so the hot scheduling path must not re-scan + // layer_types() with std::any_of on every allocate. + bool has_linear_attention_layers_ = false; + bool enable_in_batch_prefix_cache_ = false; // the number of requests that are waiting to be scheduled diff --git a/xllm/core/scheduler/mix_scheduler.cpp b/xllm/core/scheduler/mix_scheduler.cpp index 7cc48ef087..b412c9fe9a 100644 --- a/xllm/core/scheduler/mix_scheduler.cpp +++ b/xllm/core/scheduler/mix_scheduler.cpp @@ -679,8 +679,21 @@ bool MixScheduler::allocate_blocks_for(Sequence* sequence, // the total number tokens for the sequence can be handled till now. // there may some tokens can not be handled once when enable chunked prefill. + // + // Route through the shared linear-state clamp so this overload keeps the same + // guard as ChunkedPrefillScheduler::allocate_blocks_for: for GDN + prefix + // cache it rounds the prefill step down to a chunk-end boundary (or defers + // via a 0 return) so each step lands where a linear-state checkpoint is + // saved. The caller-supplied kv_cache_tokens_num (which may be inflated by + // pending host->device kv swaps) is threaded through so the returned budget + // and the downstream current_step_handle_tokens accounting stay consistent on + // this path. For every non-GDN / non-prefix-cache case the helper returns the + // plain std::min, preserving the previous behavior. size_t max_handle_num_tokens = - std::min(kv_cache_tokens_num + token_budget, sequence->num_tokens()); + get_prefill_handle_tokens(sequence, token_budget, kv_cache_tokens_num); + if (max_handle_num_tokens == 0) { + return false; + } // speculative decoding specific logic, // prefill stage don't need speculative decoding. diff --git a/xllm/core/scheduler/prefill_only_scheduler.cpp b/xllm/core/scheduler/prefill_only_scheduler.cpp index 77ab8dd759..537bf57cb6 100644 --- a/xllm/core/scheduler/prefill_only_scheduler.cpp +++ b/xllm/core/scheduler/prefill_only_scheduler.cpp @@ -27,7 +27,19 @@ limitations under the License. namespace xllm { namespace { -size_t get_prefill_allocate_tokens(Sequence* sequence, size_t step_tokens) { +// Linear-state models (Qwen3.5 GDN) must size the KV allocation by the full +// prompt: a prefix-cache hit clamps the committed KV position down to the +// linear-recoverable length, so a later chunk would advance +// kv_cache_tokens_num() past a capacity grown only for this step's budget +// (tripping batch_input_builder's current_max_tokens_capacity >= seq_len +// invariant). Pure-KV models keep the original per-step sizing, so their +// scheduling / memory behavior is unchanged. +size_t get_prefill_allocate_tokens(Sequence* sequence, + size_t step_tokens, + bool is_linear_state_prefix_cache) { + if (is_linear_state_prefix_cache) { + return sequence->num_tokens(); + } return sequence->kv_state().kv_cache_tokens_num() + step_tokens; } @@ -146,8 +158,11 @@ void PrefillOnlyScheduler::handle_prefill_requests( } // preempt offline decode - const size_t max_handle_num_tokens = - get_prefill_allocate_tokens(prefill_sequence.get(), num_tokens); + const size_t max_handle_num_tokens = get_prefill_allocate_tokens( + prefill_sequence.get(), + num_tokens, + /*is_linear_state_prefix_cache=*/ + has_linear_attention_layers_ && enable_prefix_cache_); if (!kv_cache_manager_->allocate(prefill_sequence.get(), max_handle_num_tokens)) { can_schedule = false; @@ -340,8 +355,11 @@ void PrefillOnlyScheduler::handle_last_step_prefill_requests( } // preempt offline decode - const size_t max_handle_num_tokens = - get_prefill_allocate_tokens(prefill_sequence.get(), num_tokens); + const size_t max_handle_num_tokens = get_prefill_allocate_tokens( + prefill_sequence.get(), + num_tokens, + /*is_linear_state_prefix_cache=*/ + has_linear_attention_layers_ && enable_prefix_cache_); if (!kv_cache_manager_->allocate(prefill_sequence.get(), max_handle_num_tokens)) { can_schedule = false;