bugfix: make DSV4 SWA capacity PD-aware.

tests/core/framework/batch/batch_test.cpp

@@ -3089,6 +3089,7 @@ TEST(BatchTest, OverlapMTPReplacementKeepsCompositeKvBlocks) {
       .sliding_window_size(window_size)
       .swa_blocks_per_seq(static_cast<uint32_t>(
           get_swa_blocks_per_seq(window_size, base_block_size)))
+      .swa_num_blocks(20)
       .max_tokens_per_batch(1280)
       .max_seqs_per_batch(max_seqs_per_batch)
       .manager_types({1, 0, 0})

tests/core/framework/block/composite_block_manager_test.cpp

@@ -47,11 +47,16 @@ BlockManager::Options MakeCompositeOptions(uint32_t base_num_blocks,
                                            uint32_t max_seqs_per_batch) {
   const uint32_t swa_blocks_per_seq =
       static_cast<uint32_t>(get_swa_blocks_per_seq(window_size, block_size));
+  const uint32_t burst_blocks =
+      (kMaxTokensPerBatch + block_size - 1) / block_size;
+  const uint32_t swa_num_blocks = swa_blocks_per_seq * max_seqs_per_batch +
+                                  burst_blocks + max_seqs_per_batch + 2;
   BlockManager::Options opts;
   opts.num_blocks(base_num_blocks)
       .block_size(block_size)
       .sliding_window_size(window_size)
       .swa_blocks_per_seq(swa_blocks_per_seq)
+      .swa_num_blocks(swa_num_blocks)
       .max_tokens_per_batch(kMaxTokensPerBatch)
       .max_seqs_per_batch(max_seqs_per_batch)
       .manager_types({kManagerTypeSlidingWindowBlockManager,
@@ -61,6 +66,20 @@ BlockManager::Options MakeCompositeOptions(uint32_t base_num_blocks,
   return opts;
 }
 
+void set_swa_capacity_for_token_budget(BlockManager::Options* options,
+                                       uint32_t max_tokens_per_batch) {
+  ASSERT_NE(options, nullptr);
+  const uint32_t block_size = options->block_size();
+  ASSERT_GT(block_size, 0u);
+  const uint32_t burst_blocks =
+      (max_tokens_per_batch + block_size - 1) / block_size;
+  const uint32_t max_seqs = std::max(options->max_seqs_per_batch(), 1u);
+  const uint32_t swa_num_blocks =
+      options->swa_blocks_per_seq() * max_seqs + burst_blocks + max_seqs + 2;
+  options->max_tokens_per_batch(max_tokens_per_batch)
+      .swa_num_blocks(swa_num_blocks);
+}
+
 constexpr uint32_t kBaseBlockSize = 128;
 constexpr uint32_t kCompressRatio4 = 4;
 constexpr uint32_t kCompressRatio128 = 128;
@@ -548,7 +567,7 @@ TEST(CompositeBlockManagerTest, Dsv4PrefixCacheHitOnRepeatedPrefix) {
       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);
+  set_swa_capacity_for_token_budget(&opts, 3 * kBlockSizeRatio128);
   ASSERT_TRUE(opts.enable_prefix_cache());
   CompositeBlockManager manager(build_composite_leaves(opts));
 
@@ -596,7 +615,7 @@ TEST(CompositeBlockManagerTest, Dsv4PrefixCacheMissCleanly) {
   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);
+  set_swa_capacity_for_token_budget(&opts, 3 * kBlockSizeRatio128);
   CompositeBlockManager manager(build_composite_leaves(opts));
 
   const size_t num_tokens = 2 * kBlockSizeRatio128;
@@ -629,7 +648,7 @@ TEST(CompositeBlockManagerTest, Dsv4PrefixCacheEvictsAtC128Capacity) {
   const uint32_t window_size = 4 * kBaseBlockSize;
   BlockManager::Options opts = MakeCompositeOptions(
       base_num_blocks, kBaseBlockSize, window_size, /*max_seqs_per_batch=*/1);
-  opts.max_tokens_per_batch(2 * kBlockSizeRatio128);
+  set_swa_capacity_for_token_budget(&opts, 2 * kBlockSizeRatio128);
   CompositeBlockManager manager(build_composite_leaves(opts));
 
   const size_t num_tokens = 2 * kBlockSizeRatio128;
@@ -663,7 +682,7 @@ TEST(CompositeBlockManagerTest, SlidingWindowSlidOutBlocksEnterPrefixCache) {
   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);
+  set_swa_capacity_for_token_budget(&opts, 3 * kBlockSizeRatio128);
   ASSERT_TRUE(opts.enable_prefix_cache());
   CompositeBlockManager manager(build_composite_leaves(opts));
 
@@ -691,6 +710,44 @@ TEST(CompositeBlockManagerTest, SlidingWindowSlidOutBlocksEnterPrefixCache) {
   manager.deallocate_for_sequence(&seq_hit);
 }
 
+TEST(CompositeBlockManagerTest,
+     SlidingWindowReclaimsUncachedBlockBeforeFullCacheUnit) {
+  const uint32_t window_size = 2 * kBaseBlockSize;
+  BlockManager::Options opts = MakeCompositeOptions(
+      /*base_num_blocks=*/4096,
+      kBaseBlockSize,
+      window_size,
+      /*max_seqs_per_batch=*/1);
+  opts.enable_prefix_cache(true).swa_num_blocks(
+      /*three live blocks plus padding=*/4);
+  CompositeBlockManager manager(build_composite_leaves(opts));
+
+  const size_t first_chunk_tokens = 3 * kBaseBlockSize;
+  const size_t second_chunk_tokens = 4 * kBaseBlockSize;
+  Sequence seq =
+      MakeTestSequence(0, std::vector<int32_t>(second_chunk_tokens, 7));
+
+  ASSERT_TRUE(manager.allocate_sequence(&seq, first_chunk_tokens));
+  BlockManager* swa_leaf = manager.leaf_entries().at(BlockType::SWA).leaf.get();
+  ASSERT_NE(swa_leaf, nullptr);
+  EXPECT_EQ(swa_leaf->num_free_blocks(), 0u);
+
+  seq.kv_state().incr_kv_cache_tokens_num(first_chunk_tokens);
+
+  // The first block is outside the two-block window. No C128 cache unit is
+  // complete yet, so the next growth releases it directly and reuses its
+  // physical id without publishing a partial DSV4 prefix.
+  ASSERT_TRUE(manager.allocate_sequence(&seq, second_chunk_tokens));
+  const std::vector<Block> swa_blocks = SwaBlocks(seq);
+  ASSERT_EQ(swa_blocks.size(), 4u);
+  EXPECT_FALSE(swa_blocks.front().is_valid());
+  EXPECT_TRUE(swa_blocks.back().is_valid());
+  EXPECT_EQ(swa_leaf->num_free_blocks(), 0u);
+  EXPECT_EQ(swa_leaf->num_blocks_in_prefix_cache(), 0u);
+
+  manager.deallocate_for_sequence(&seq);
+}
+
 // The post-grow hook advances KVCacheState::num_cached_blocks incrementally.
 // Newly allocated blocks are present by then, but the token cursor limits the
 // published range to blocks completed by the preceding forward.
@@ -700,7 +757,7 @@ TEST(CompositeBlockManagerTest, Dsv4PrefixCachePostGrowCursorAdvances) {
   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);
+  set_swa_capacity_for_token_budget(&opts, 4 * kBlockSizeRatio128);
   CompositeBlockManager manager(build_composite_leaves(opts));
 
   // Two-chunk prompt (2*C128). Chunk 1 is a single C128 block wide.
@@ -725,6 +782,55 @@ TEST(CompositeBlockManagerTest, Dsv4PrefixCachePostGrowCursorAdvances) {
             chunk / kBlockSizeRatio4);
   EXPECT_EQ(seq.kv_state().num_cached_blocks(BlockType::C128),
             chunk / kBlockSizeRatio128);
+  EXPECT_EQ(manager.leaf_entries()
+                .at(BlockType::SWA)
+                .leaf->num_blocks_in_prefix_cache(),
+            4u);
+  EXPECT_EQ(manager.leaf_entries()
+                .at(BlockType::C4)
+                .leaf->num_blocks_in_prefix_cache(),
+            32u);
+  EXPECT_EQ(manager.leaf_entries()
+                .at(BlockType::C128)
+                .leaf->num_blocks_in_prefix_cache(),
+            1u);
+
+  manager.deallocate_for_sequence(&seq);
+}
+
+TEST(CompositeBlockManagerTest, Dsv4PrefixCacheSkipsPartialCacheUnitTail) {
+  const uint32_t base_num_blocks = 4096;
+  const uint32_t window_size = kBaseBlockSize;
+  const uint32_t max_seqs_per_batch = 4;
+  BlockManager::Options opts = MakeCompositeOptions(
+      base_num_blocks, kBaseBlockSize, window_size, max_seqs_per_batch);
+  set_swa_capacity_for_token_budget(&opts, 2 * kBlockSizeRatio128);
+  CompositeBlockManager manager(build_composite_leaves(opts));
+
+  const size_t completed_tokens = kBlockSizeRatio128 + kBlockSizeRatio4;
+  Sequence seq =
+      MakeTestSequence(0, std::vector<int32_t>(completed_tokens, 17));
+  ASSERT_TRUE(manager.allocate_sequence(&seq, completed_tokens));
+  seq.kv_state().incr_kv_cache_tokens_num(completed_tokens);
+  manager.cache_for_sequence(&seq);
+
+  EXPECT_EQ(manager.leaf_entries()
+                .at(BlockType::SWA)
+                .leaf->num_blocks_in_prefix_cache(),
+            1u);
+  EXPECT_EQ(manager.leaf_entries()
+                .at(BlockType::C4)
+                .leaf->num_blocks_in_prefix_cache(),
+            32u);
+  EXPECT_EQ(manager.leaf_entries()
+                .at(BlockType::C128)
+                .leaf->num_blocks_in_prefix_cache(),
+            1u);
+  EXPECT_EQ(seq.kv_state().num_cached_blocks(BlockType::SWA),
+            kBlockSizeRatio128 / kBaseBlockSize);
+  EXPECT_EQ(seq.kv_state().num_cached_blocks(BlockType::C4),
+            kBlockSizeRatio128 / kBlockSizeRatio4);
+  EXPECT_EQ(seq.kv_state().num_cached_blocks(BlockType::C128), 1u);
 
   manager.deallocate_for_sequence(&seq);
 }
@@ -739,7 +845,7 @@ TEST(CompositeBlockManagerTest, Dsv4PrefixCacheExactRepeatPopsOneC128) {
   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);
+  set_swa_capacity_for_token_budget(&opts, 4 * kBlockSizeRatio128);
   CompositeBlockManager manager(build_composite_leaves(opts));
 
   const size_t num_tokens = 3 * kBlockSizeRatio128;
@@ -770,7 +876,7 @@ TEST(CompositeBlockManagerTest, DecodeRoleSkipsSwaPrefixCache) {
   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);
+  set_swa_capacity_for_token_budget(&opts, 3 * kBlockSizeRatio128);
   // First seed the cache under a PREFILL role so all three leaves publish
   // their blocks -- then swap in a DECODE-role composite that shares the
   // hash space via the same leaf construction path.
@@ -818,6 +924,33 @@ TEST(CompositeBlockManagerTest, DecodeRoleSkipsSwaPrefixCache) {
   decode_manager.deallocate_for_sequence(&seq_d);
 }
 
+TEST(CompositeBlockManagerTest, DecodeInitialSwaAllocationKeepsOnlyWindowTail) {
+  const uint32_t window_size = 2 * kBaseBlockSize;
+  BlockManager::Options opts = MakeCompositeOptions(
+      /*base_num_blocks=*/4096,
+      kBaseBlockSize,
+      window_size,
+      /*max_seqs_per_batch=*/1);
+  opts.instance_is_decode(true).enable_prefix_cache(false).swa_num_blocks(
+      /*two windows plus padding=*/5);
+  CompositeBlockManager manager(build_composite_leaves(opts));
+
+  const size_t logical_blocks = 10;
+  const size_t num_tokens = logical_blocks * kBaseBlockSize;
+  Sequence seq = MakeTestSequence(0, std::vector<int32_t>(num_tokens, 7));
+
+  ASSERT_TRUE(manager.allocate_sequence(&seq, num_tokens));
+  const std::vector<Block> swa = SwaBlocks(seq);
+  ASSERT_EQ(swa.size(), logical_blocks);
+  for (size_t i = 0; i < logical_blocks - 2; ++i) {
+    EXPECT_FALSE(swa[i].is_valid());
+  }
+  EXPECT_TRUE(swa[logical_blocks - 2].is_valid());
+  EXPECT_TRUE(swa[logical_blocks - 1].is_valid());
+
+  manager.deallocate_for_sequence(&seq);
+}
+
 // Qwen3.5 GDN D-side (instance_is_decode=true, LINEAR present): the LINEAR
 // leaf should stop advertising prefix cache, so build_composite_leaves
 // classifies FLAT_KV_LINEAR down to FLAT_KV and no restore-source mount

tests/core/framework/block/hierarchy_block_manager_pool_test.cpp

@@ -163,6 +163,7 @@ BlockManagerPool::Options make_typed_cache_options() {
       .enable_host_offload(true)
       .sliding_window_size(kWindow)
       .swa_blocks_per_seq(swa_blocks_per_seq)
+      .swa_num_blocks(266)
       .max_tokens_per_batch(32768)
       .max_seqs_per_batch(4)
       // SlidingWindow + BlockManagerImpl (C4) + BlockManagerImpl (C128).
@@ -385,9 +386,10 @@ TEST(HierarchyBlockManagerPoolTest,
   sequence.kv_state().set_kv_cache_tokens_num(kPromptTokens);
   pool.deallocate(&sequence);
 
-  // Only completed blocks are offloaded: 156 SWA blocks, 39 C4 blocks, and
-  // one C128 checkpoint. The partial SWA tail is not inserted or offloaded.
-  EXPECT_EQ(HierarchyPoolTestPeer::pending_offload_pair_count(pool), 196u);
+  // Decode allocates only the active SWA window. Here that window is the
+  // partial tail block, so no SWA block is offloaded. Only the complete C128
+  // cache unit (32 C4 blocks plus one C128 checkpoint) is offloaded.
+  EXPECT_EQ(HierarchyPoolTestPeer::pending_offload_pair_count(pool), 33u);
 }
 
 TEST(HierarchyBlockManagerPoolTest, AllocateSharedMountsMatchesWithoutH2d) {
@@ -1307,7 +1309,7 @@ TEST(HierarchyBlockManagerPoolTest,
 
   pool.prefetch_from_storage(request);
   ASSERT_NE(engine.result(), nullptr);
-  EXPECT_EQ(engine.transfer_infos().size(), 128u + 32u + 1u);
+  EXPECT_EQ(engine.transfer_infos().size(), 1u + 32u + 1u);
 
   const std::vector<uint8_t> hits(engine.transfer_infos().size(), 1);
   ASSERT_TRUE(engine.result()->set_batch_result(
@@ -1321,7 +1323,13 @@ TEST(HierarchyBlockManagerPoolTest,
   ASSERT_TRUE(pool.update_prefetch_result(request, /*timeout=*/0));
 
   pool.allocate_shared(sequence);
-  EXPECT_EQ(sequence->host_kv_state().num_blocks(BlockType::SWA), 128u);
+  const Slice<Block> swa_blocks =
+      sequence->host_kv_state().blocks(BlockType::SWA);
+  ASSERT_EQ(swa_blocks.size(), 128u);
+  for (size_t i = 0; i + 1 < swa_blocks.size(); ++i) {
+    EXPECT_FALSE(swa_blocks[i].is_valid());
+  }
+  EXPECT_TRUE(swa_blocks.back().is_valid());
   EXPECT_EQ(sequence->host_kv_state().num_blocks(BlockType::C4), 32u);
   EXPECT_EQ(sequence->host_kv_state().num_blocks(BlockType::C128), 1u);
   EXPECT_EQ(sequence->kv_cache_tokens_num(), 16384u);
@@ -1331,7 +1339,8 @@ TEST(HierarchyBlockManagerPoolTest,
 TEST(HierarchyBlockManagerPoolTest,
      TypedStoragePrefetchKeepsSwaHitsAfterMiddleMiss) {
   constexpr size_t kPromptTokens = 32769;
-  constexpr size_t kMissedSwaOrdinal = 100;
+  constexpr size_t kMissedSwaOrdinal = 0;
+  constexpr size_t kMissedSwaBlock = 127;
   BlockManagerPool::Options options = make_typed_cache_options();
   options.enable_kvcache_store(true);
   FakePrefetchEngine engine(/*worker_count=*/2);
@@ -1356,7 +1365,7 @@ TEST(HierarchyBlockManagerPoolTest,
     }
     ++swa_ordinal;
   }
-  ASSERT_EQ(swa_ordinal, 256u);
+  ASSERT_EQ(swa_ordinal, 2u);
   ASSERT_LT(missed_result_index, engine.transfer_infos().size());
 
   const size_t split = engine.transfer_infos().size() / 2;
@@ -1388,8 +1397,7 @@ TEST(HierarchyBlockManagerPoolTest,
   const Slice<Block> swa_blocks =
       sequence->host_kv_state().blocks(BlockType::SWA);
   ASSERT_EQ(swa_blocks.size(), 256u);
-  EXPECT_FALSE(swa_blocks[kMissedSwaOrdinal].is_valid());
-  EXPECT_TRUE(swa_blocks[kMissedSwaOrdinal + 1].is_valid());
+  EXPECT_FALSE(swa_blocks[kMissedSwaBlock].is_valid());
   EXPECT_TRUE(swa_blocks.back().is_valid());
   EXPECT_EQ(sequence->host_kv_state().num_blocks(BlockType::C4), 64u);
   EXPECT_EQ(sequence->host_kv_state().num_blocks(BlockType::C128), 2u);

tests/core/framework/config/disagg_pd_config_test.cpp

@@ -12,80 +12,3 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
-
-#include "core/framework/config/disagg_pd_config.h"
-
-#include <gtest/gtest.h>
-
-#include <algorithm>
-#include <string>
-
-#include "core/framework/config/kv_cache_config.h"
-#include "core/framework/config/scheduler_config.h"
-
-namespace xllm {
-namespace {
-
-struct PrefixRoleCase {
-  std::string role;
-  bool keep_prefix_cache;
-};
-
-void set_values_requiring_mlu_normalization(DisaggPDConfig& disagg_pd_config,
-                                            KVCacheConfig& kv_cache_config,
-                                            SchedulerConfig& scheduler_config) {
-  disagg_pd_config.kv_cache_transfer_mode("PULL").enable_pd_ooc(true);
-  kv_cache_config.kv_cache_dtype("fp8").enable_prefix_cache(true);
-  scheduler_config.enable_schedule_overlap(true);
-}
-
-void expect_normalized_values(const DisaggPDConfig& disagg_pd_config,
-                              const KVCacheConfig& kv_cache_config,
-                              const SchedulerConfig& scheduler_config) {
-  EXPECT_EQ(disagg_pd_config.kv_cache_transfer_mode(), "PULL");
-  EXPECT_FALSE(disagg_pd_config.enable_pd_ooc());
-  EXPECT_EQ(kv_cache_config.kv_cache_dtype(), "auto");
-  EXPECT_FALSE(scheduler_config.enable_schedule_overlap());
-}
-
-TEST(DisaggPDConfigTest, KeepsMluPrefixCacheForPrefillSideRoles) {
-  const PrefixRoleCase cases[] = {
-      {"PREFILL", true},
-      {"MIX", true},
-      {"DECODE", false},
-      {"DEFAULT", false},
-  };
-
-  for (const PrefixRoleCase& test_case : cases) {
-    DisaggPDConfig disagg_pd_config;
-    KVCacheConfig kv_cache_config;
-    SchedulerConfig scheduler_config;
-    disagg_pd_config.instance_role(test_case.role);
-    set_values_requiring_mlu_normalization(
-        disagg_pd_config, kv_cache_config, scheduler_config);
-
-    disagg_pd_config.normalize_mlu(kv_cache_config, scheduler_config);
-
-    SCOPED_TRACE(test_case.role);
-    expect_normalized_values(
-        disagg_pd_config, kv_cache_config, scheduler_config);
-    EXPECT_EQ(kv_cache_config.enable_prefix_cache(),
-              test_case.keep_prefix_cache);
-  }
-}
-
-TEST(DisaggPDConfigTest, OmitsRemovedHeterogeneousPullOptions) {
-  const std::vector<std::string>& option_names =
-      DisaggPDConfig::option_category().option_names;
-  EXPECT_EQ(
-      std::find(
-          option_names.begin(), option_names.end(), "enable_heterogeneous_pd"),
-      option_names.end());
-  EXPECT_EQ(std::find(option_names.begin(),
-                      option_names.end(),
-                      "enable_pd_parallel_shard_pull"),
-            option_names.end());
-}
-
-}  // namespace
-}  // namespace xllm