feat: add linear-state prefix cache batch and scheduler plumbing (5/n).

tests/core/framework/batch/batch_packed_input_test.cpp

@@ -16,17 +16,32 @@ limitations under the License.
 
 #include <gtest/gtest.h>
 
+#include <cstddef>
+#include <cstdint>
+#include <utility>
 #include <vector>
 
 #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 <typename T>
 bool tensor_equals_vector(const torch::Tensor& tensor,
                           const std::vector<T>& 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<Sequence*> sequences = {&seq};
+  std::vector<uint32_t> 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;

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<int32_t>& token_ids,
+    std::vector<Block>& 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<int32_t>(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<int32_t>& 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<int32_t> 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<Block> 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<Block> 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<Block> restore_src_blocks = manager.allocate(1);
+  restore_seq.set_linear_restore_src_block(std::move(restore_src_blocks[0]));
+
+  const std::vector<int32_t> 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<int32_t> 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<XXH3Key> 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<XXH3Key> 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<int32_t>({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<int32_t>({-1, -1}));
 }
 
 TEST(BatchTest, SharedMemoryRoundTripPreservesEmptyRankTensors) {

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> 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> request;