Skip to content

feat: add linear-state restore and capacity estimation (4/n). #1940

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
yingxudeng wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: add linear-state restore and capacity estimation (4/n). #1940

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
115 changes: 112 additions & 3 deletions tests/core/framework/kv_cache/kv_cache_estimation_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,9 @@ limitations under the License.

#include <gtest/gtest.h>

#include <cstdint>
#include <vector>

#include "framework/model/model_args.h"

namespace xllm {
Expand All @@ -37,7 +40,6 @@ KVCacheEstimateOptions make_estimate_options() {
options.world_size = 1;
options.n_local_kv_heads = 2;
options.max_seqs_per_batch = 8;
options.max_concurrent_requests = 8;
return options;
}

Expand All @@ -58,17 +60,32 @@ TEST(KVCacheEstimationTest, EstimatesStandardAttentionBlocks) {
EXPECT_EQ(capacity.n_blocks(), 128);
}

TEST(KVCacheEstimationTest, ReservesLinearAttentionState) {
namespace {

ModelArgs make_linear_attention_args(int64_t head_dim = 16) {
ModelArgs model_args = make_standard_args();
model_args.full_attention_interval(2)
model_args.head_dim(head_dim)
.full_attention_interval(2)
.linear_num_key_heads(2)
.linear_num_value_heads(2)
.linear_key_head_dim(4)
.linear_value_head_dim(8)
.linear_conv_kernel_dim(3);
return model_args;
}

KVCacheEstimateOptions make_linear_attention_options() {
KVCacheEstimateOptions options = make_estimate_options();
options.n_local_linear_k_heads = 2;
options.n_local_linear_v_heads = 2;
return options;
}

} // namespace

TEST(KVCacheEstimationTest, ReservesLinearAttentionState) {
ModelArgs model_args = make_linear_attention_args();
KVCacheEstimateOptions options = make_linear_attention_options();

KVCacheCapacity capacity = estimate_kv_cache_capacity(model_args, options);

Expand All @@ -80,6 +97,98 @@ TEST(KVCacheEstimationTest, ReservesLinearAttentionState) {
EXPECT_EQ(capacity.n_blocks(), 254);
}

TEST(KVCacheEstimationTest, LinearStateCapacityVariants) {
struct TestCase {
const char* name;
int64_t head_dim;
int64_t cache_size_in_bytes;
int64_t block_size;
int64_t n_local_kv_heads;
int64_t max_seqs_per_batch;
bool enable_prefix_cache;
int64_t max_linear_state_cache_slots;
int64_t expected_num_linear_state_blocks;
int64_t min_num_linear_state_blocks;
};

const std::vector<TestCase> test_cases = {
{"PrefixCacheGrowsLinearStateCheckpointPool",
/*head_dim=*/16,
/*cache_size_in_bytes=*/64LL << 30,
/*block_size=*/16,
/*n_local_kv_heads=*/2,
/*max_seqs_per_batch=*/200,
/*enable_prefix_cache=*/true,
/*max_linear_state_cache_slots=*/0,
/*expected_num_linear_state_blocks=*/-1,
/*min_num_linear_state_blocks=*/202},
{"PrefixCacheUsesLinearStateMemoryRatio",
/*head_dim=*/1,
/*cache_size_in_bytes=*/1024 * 1024,
/*block_size=*/1,
/*n_local_kv_heads=*/1,
/*max_seqs_per_batch=*/8,
/*enable_prefix_cache=*/true,
/*max_linear_state_cache_slots=*/0,
/*expected_num_linear_state_blocks=*/970,
/*min_num_linear_state_blocks=*/-1},
{"NoPrefixCacheCapsLinearStateBlocksByBudget",
/*head_dim=*/16,
/*cache_size_in_bytes=*/1024 * 1024,
/*block_size=*/16,
/*n_local_kv_heads=*/2,
/*max_seqs_per_batch=*/100000,
/*enable_prefix_cache=*/false,
/*max_linear_state_cache_slots=*/0,
/*expected_num_linear_state_blocks=*/229,
/*min_num_linear_state_blocks=*/-1},
{"UnlimitedConcurrencyFallsBackToPaddingSlots",
/*head_dim=*/16,
/*cache_size_in_bytes=*/1024 * 1024,
/*block_size=*/16,
/*n_local_kv_heads=*/2,
/*max_seqs_per_batch=*/0,
/*enable_prefix_cache=*/false,
/*max_linear_state_cache_slots=*/0,
/*expected_num_linear_state_blocks=*/2,
/*min_num_linear_state_blocks=*/-1},
{"ExplicitLinearStateSlotsOverrideAutoSizing",
/*head_dim=*/16,
/*cache_size_in_bytes=*/64LL << 30,
/*block_size=*/16,
/*n_local_kv_heads=*/2,
/*max_seqs_per_batch=*/200,
/*enable_prefix_cache=*/true,
/*max_linear_state_cache_slots=*/32,
/*expected_num_linear_state_blocks=*/34,
/*min_num_linear_state_blocks=*/-1},
};

for (const TestCase& test_case : test_cases) {
SCOPED_TRACE(test_case.name);
ModelArgs model_args = make_linear_attention_args(test_case.head_dim);
KVCacheEstimateOptions options = make_linear_attention_options();
options.cache_size_in_bytes = test_case.cache_size_in_bytes;
options.block_size = test_case.block_size;
options.n_local_kv_heads = test_case.n_local_kv_heads;
options.max_seqs_per_batch = test_case.max_seqs_per_batch;
options.enable_prefix_cache = test_case.enable_prefix_cache;
options.max_linear_state_cache_slots =
test_case.max_linear_state_cache_slots;

KVCacheCapacity capacity = estimate_kv_cache_capacity(model_args, options);

if (test_case.expected_num_linear_state_blocks >= 0) {
EXPECT_EQ(capacity.num_linear_state_blocks(),
test_case.expected_num_linear_state_blocks);
}
if (test_case.min_num_linear_state_blocks >= 0) {
EXPECT_GT(capacity.num_linear_state_blocks(),
test_case.min_num_linear_state_blocks);
}
}
}

TEST(KVCacheEstimationTest, Qwen35MtpExpandsConvStateLen) {
ModelArgs model_args = make_standard_args();
model_args.model_type("qwen3_5")
Expand Down
46 changes: 46 additions & 0 deletions tests/core/runtime/acl_graph_executor_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ limitations under the License.
#include "core/framework/block/block_manager_impl.h"
#include "core/framework/config/execution_config.h"
#include "core/framework/kv_cache/kv_cache.h"
#include "core/framework/kv_cache/linear_state_restore.h"
#include "core/framework/model/model_args.h"
#include "core/framework/model/model_output.h"
#include "core/framework/model_loader.h"
Expand Down Expand Up @@ -749,6 +750,51 @@ TEST_F(AclGraphExecutorTest, BatchInputCarriesLinearStateIds) {
expected_linear_state_id);
}

TEST(LinearStateRestoreTest, RestoreCopiesCheckpointSlot) {
std::vector<KVCache> kv_caches;
torch::Tensor conv_cache = torch::zeros({3, 2, 2}, torch::kFloat32);
torch::Tensor ssm_cache = torch::zeros({3, 2, 2}, torch::kFloat32);
conv_cache.select(0, 1).fill_(7.0);
ssm_cache.select(0, 1).fill_(11.0);
kv_caches.emplace_back(LinearAttentionKVCacheTensors{conv_cache, ssm_cache});

LinearStateCacheOp restore;
restore.linear_state_id = 2;
restore.restore_requested = true;
restore.restore_src_slot_id = 1;
std::vector<int64_t> has_initial_state(1, 0);
restore_linear_state_slots(kv_caches, {restore}, has_initial_state);

EXPECT_EQ(has_initial_state[0], 1);
EXPECT_TRUE(
torch::allclose(conv_cache.select(0, 2), conv_cache.select(0, 1)));
EXPECT_TRUE(torch::allclose(ssm_cache.select(0, 2), ssm_cache.select(0, 1)));
}

TEST(LinearStateRestoreTest, RestoreRequestedButUnresolvedColdStarts) {
std::vector<KVCache> kv_caches;
kv_caches.emplace_back(LinearAttentionKVCacheTensors{
torch::zeros({2, 4, 3}, torch::dtype(torch::kFloat32)),
torch::zeros({2, 4, 4}, torch::dtype(torch::kFloat32))});

LinearStateCacheOp no_restore;
no_restore.linear_state_id = 1;
LinearStateCacheOp missing_restore;
missing_restore.linear_state_id = 1;
missing_restore.restore_requested = true;
missing_restore.restore_src_slot_id = -1;

const std::vector<LinearStateCacheOp> cache_ops = {no_restore,
missing_restore};
// Pre-fill both entries with the kv-cache default of 1 so the test can
// distinguish "SKIPPED leaves the default untouched" from "COLD_START
// overrides it to 0".
std::vector<int64_t> has_initial_state = {1, 1};
restore_linear_state_slots(kv_caches, cache_ops, has_initial_state);
EXPECT_EQ(has_initial_state[0], 1);
EXPECT_EQ(has_initial_state[1], 0);
}

TEST_F(AclGraphExecutorTest, GraphDoubleBufferFlagControlsSlotCount) {
ExecutionConfig& execution_config = ExecutionConfig::get_instance();
const bool original_enable_graph_double_buffer =
Expand Down
24 changes: 23 additions & 1 deletion xllm/core/framework/kv_cache/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ cc_library(
quantized_kv_cache_impl.h
SRCS
deepseek_v4_kv_cache_impl.cpp
embedding_cache.cpp
embedding_cache.cpp
indexed_kv_cache_impl.cpp
kv_cache.cpp
kv_cache_impl.cpp
Expand All @@ -60,3 +60,25 @@ cc_library(
$<$<BOOL:${USE_NPU}>:ascendcl>
$<$<BOOL:${USE_NPU}>:torch_npu>
)

# linear_state_restore depends on both :kv_cache and :model, so it lives in a
# separate target to avoid pulling :model into the base :kv_cache library
# (which would create a cycle, since :model already depends on :kv_cache).
cc_library(
NAME
linear_state_restore
HDRS
linear_state_restore.h
SRCS
linear_state_restore.cpp
DEPS
:common
:kv_cache
:model
:platform
:util
glog::glog
torch
$<$<BOOL:${USE_NPU}>:ascendcl>
$<$<BOOL:${USE_NPU}>:torch_npu>
)
99 changes: 96 additions & 3 deletions xllm/core/framework/kv_cache/kv_cache_estimation.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,91 @@ int64_t linear_slot_size(const ModelArgs& model_args,
linear_ssm_slot_size * (num_speculative_tokens + 1);
}

int64_t max_linear_state_blocks(int64_t cache_size_in_bytes,
int64_t num_linear_attention_layers,
int64_t linear_slot_size,
int64_t num_full_attention_layers,
int64_t full_attention_block_size) {
if (linear_slot_size <= 0 || num_linear_attention_layers <= 0) {
return kPaddingLinearStateBlocks;
}

CHECK_GT(cache_size_in_bytes, 0);
CHECK_GT(full_attention_block_size, 0);
const int64_t linear_bytes_per_block =
num_linear_attention_layers * linear_slot_size;
const int64_t full_cache_bytes_per_block =
std::max<int64_t>(num_full_attention_layers, 1) *
full_attention_block_size;
CHECK_GT(linear_bytes_per_block, 0);
CHECK_GT(full_cache_bytes_per_block, 0);

int64_t max_linear_blocks =
(cache_size_in_bytes - 1) / linear_bytes_per_block;
const int64_t balanced_max_linear_blocks =
(cache_size_in_bytes +
kPaddingLinearStateBlocks * full_cache_bytes_per_block) /
(linear_bytes_per_block + full_cache_bytes_per_block);
max_linear_blocks = std::min(max_linear_blocks, balanced_max_linear_blocks);

return std::max<int64_t>(max_linear_blocks, kPaddingLinearStateBlocks);
}

int64_t calculate_linear_state_blocks(int64_t cache_size_in_bytes,
int64_t num_linear_attention_layers,
int64_t linear_slot_size,
int64_t num_full_attention_layers,
int64_t full_attention_block_size,
int64_t max_seqs_per_batch,
int64_t max_linear_state_cache_slots,
bool enable_prefix_cache) {
CHECK_GE(max_linear_state_cache_slots, 0)
<< "max_linear_state_cache_slots must be greater than or equal to 0.";
const int64_t max_blocks =
max_linear_state_blocks(cache_size_in_bytes,
num_linear_attention_layers,
linear_slot_size,
num_full_attention_layers,
full_attention_block_size);
if (max_linear_state_cache_slots > 0) {
const int64_t requested_blocks =
max_linear_state_cache_slots + kPaddingLinearStateBlocks;
CHECK_LE(requested_blocks, max_blocks)
<< "max_linear_state_cache_slots requires " << requested_blocks
<< " linear-state blocks, but only " << max_blocks
<< " fit in the configured KV cache budget.";
return requested_blocks;
}

if (!enable_prefix_cache) {
const int64_t live_slot_blocks =
max_seqs_per_batch + kPaddingLinearStateBlocks;
return std::max<int64_t>(std::min<int64_t>(live_slot_blocks, max_blocks),
kPaddingLinearStateBlocks);
}

// Auto-size: allocate ~47% of cache bytes to linear-state slots (ratio 0.9
// means linear fraction = 0.9 / 1.9).
constexpr double kLinearStateFullKvMemoryRatio = 0.9;
const int64_t linear_bytes_per_block =
num_linear_attention_layers * linear_slot_size;
int64_t auto_blocks = kPaddingLinearStateBlocks;
if (linear_slot_size > 0 && num_linear_attention_layers > 0 &&
linear_bytes_per_block > 0) {
const double linear_memory_fraction =
kLinearStateFullKvMemoryRatio / (1.0 + kLinearStateFullKvMemoryRatio);
const double linear_memory_bytes =
static_cast<double>(cache_size_in_bytes) * linear_memory_fraction;
auto_blocks = std::max<int64_t>(
static_cast<int64_t>(linear_memory_bytes / linear_bytes_per_block),
kPaddingLinearStateBlocks);
}
// Both bounds already sit at or above kPaddingLinearStateBlocks (auto_blocks
// is floored above; max_blocks is floored in max_linear_state_blocks), so the
// min of the two cannot drop below it.
return std::min<int64_t>(auto_blocks, max_blocks);
}

Dsv4KVCacheEstimateCost estimate_dsv4_kv_cache_cost(
const ModelArgs& model_args,
const KVCacheEstimateOptions& options) {
Expand Down Expand Up @@ -310,7 +395,6 @@ void init_dsv4_counts(const ModelArgs& model_args,
void init_standard_counts(const ModelArgs& model_args,
const KVCacheEstimateOptions& options,
KVCacheCapacity* kv_cache_cap) {
kv_cache_cap->num_linear_state_blocks(options.max_concurrent_requests + 2);
for (int64_t layer_id = 0; layer_id < kv_cache_cap->n_layers(); ++layer_id) {
if (is_full_attention_layer(model_args, layer_id)) {
++kv_cache_cap->num_full_attention_layers();
Expand All @@ -324,6 +408,15 @@ void init_standard_counts(const ModelArgs& model_args,
block_size *
(kv_cache_cap->slot_size() + kv_cache_cap->index_slot_size() +
kv_cache_cap->scale_slot_size());
kv_cache_cap->num_linear_state_blocks(
calculate_linear_state_blocks(kv_cache_cap->cache_size_in_bytes(),
kv_cache_cap->num_linear_attention_layers(),
kv_cache_cap->linear_slot_size(),
kv_cache_cap->num_full_attention_layers(),
block_size_in_bytes,
options.max_seqs_per_batch,
options.max_linear_state_cache_slots,
options.enable_prefix_cache));
kv_cache_cap->linear_cache_size_in_bytes(
kv_cache_cap->num_linear_attention_layers() *
kv_cache_cap->num_linear_state_blocks() *
Expand All @@ -336,8 +429,8 @@ void init_standard_counts(const ModelArgs& model_args,
kv_cache_cap->linear_cache_size_in_bytes())
<< "failed to reserve linear state cache for linear-attention "
"layers: "
<< "max_concurrent_requests (" << options.max_concurrent_requests
<< ") is too large. Please reduce max_concurrent_requests to less than "
<< "max_seqs_per_batch (" << options.max_seqs_per_batch
<< ") is too large. Please reduce max_seqs_per_batch to less than "
<< kv_cache_cap->cache_size_in_bytes() /
(kv_cache_cap->num_linear_attention_layers() *
kv_cache_cap->linear_slot_size()) -
Expand Down
4 changes: 4 additions & 0 deletions xllm/core/framework/kv_cache/kv_cache_estimation.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,10 @@ struct KVCacheEstimateOptions {
int64_t max_seqs_per_batch = 0;
int64_t num_speculative_tokens = 0;
int64_t max_tokens_per_batch = 0;
int64_t max_linear_state_cache_slots = 0;
// Legacy alias kept so the (not-yet-migrated) engine/worker callers still
// compile against the old field name. The engine/worker wiring PR removes
// this and switches those callers to max_linear_state_cache_slots.
int64_t max_concurrent_requests = 0;
bool is_draft_engine = false;
bool enable_prefix_cache = false;
Expand Down
1 change: 1 addition & 0 deletions xllm/core/framework/kv_cache/kv_cache_utils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ limitations under the License.
#include <vector>

#include "common/macros.h"
#include "core/common/constants.h"

@yingxudeng yingxudeng Jul 14, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

后面pr会使用

#include "util/tensor_helper.h"

#if defined(USE_NPU)
Expand Down
Loading
Loading