Skip to content

Commit ae8ba1a

Browse files
Kang-MengKang-Meng
committed
feat: allocate host KV cache as page-aligned pinned memory in KVCache impls.
1 parent 4a76203 commit ae8ba1a

15 files changed

Lines changed: 615 additions & 97 deletions

‎tests/core/framework/kv_cache/kv_cache_test.cpp‎

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,11 @@ limitations under the License.
1919

2020
#include <vector>
2121

22+
#include "framework/block/block.h"
23+
#include "framework/kv_cache/deepseek_v4_cache_policy.h"
24+
#include "framework/kv_cache/kv_cache_utils.h"
2225
#include "kv_cache_shape.h"
26+
#include "platform/device.h"
2327

2428
namespace xllm {
2529

@@ -42,6 +46,17 @@ std::vector<int64_t> dsv4_block_shape(int64_t block_count,
4246

4347
} // namespace
4448

49+
// Host prefix-cache allocation registers page-aligned host memory with the NPU
50+
// via aclrtHostRegister, which requires a live device context. Set one up once.
51+
class HostKVCacheTest : public ::testing::Test {
52+
protected:
53+
void SetUp() override {
54+
Device device(/*device_index=*/0);
55+
device.set_device();
56+
device.init_device_context();
57+
}
58+
};
59+
4560
TEST(KVCacheTest, DeepSeekV4FourDimCachesUseDeviceLayout) {
4661
constexpr int64_t kSwaCount = 10;
4762
constexpr int64_t kC4Count = 32;
@@ -109,4 +124,111 @@ TEST(KVCacheTest, DeepSeekV4FourDimCachesUseDeviceLayout) {
109124
(std::vector<int64_t>{kSwaCount, kBlockSize, kHeadDim}));
110125
}
111126

127+
TEST_F(HostKVCacheTest, HostKVCacheNormalLayoutAddsLayerDim) {
128+
constexpr int64_t kNumBlocks = 16;
129+
constexpr int64_t kBlockSize = 128;
130+
constexpr int64_t kHeadDim = 64;
131+
constexpr int64_t kNumHeads = 4;
132+
constexpr int64_t kLayerCount = 5;
133+
constexpr double kHostFactor = 2.0;
134+
135+
KVCacheCapacity capacity;
136+
capacity.n_blocks(kNumBlocks).block_size(kBlockSize);
137+
138+
ModelArgs model_args;
139+
model_args.model_type("qwen");
140+
model_args.n_kv_heads(kNumHeads);
141+
model_args.head_dim(kHeadDim);
142+
KVCacheShape shape(capacity, model_args, /*world_size=*/1);
143+
ASSERT_TRUE(shape.has_key_cache_shape());
144+
145+
KVCacheCreateOptions options;
146+
options.device(torch::Device(torch::kCPU))
147+
.dtype(torch::kFloat32)
148+
.num_layers(kLayerCount)
149+
.model_type("qwen")
150+
.host_blocks_factor(kHostFactor);
151+
152+
KVCache host_cache(shape, options, BlockType::KV, kLayerCount);
153+
154+
const BlockTypeTensorMap tensors =
155+
host_cache.get_block_type_tensors(BlockType::KV);
156+
ASSERT_TRUE(tensors.count(KVCacheTensorRole::KEY) > 0);
157+
158+
const std::vector<int64_t> base_key_shape = shape.key_cache_shape();
159+
const torch::Tensor& host_key = tensors.at(KVCacheTensorRole::KEY);
160+
EXPECT_TRUE(host_key.is_contiguous());
161+
EXPECT_EQ(host_key.device().type(), torch::kCPU);
162+
// host shape == [scaled_blocks, layer_count, ...per_block_dims]
163+
ASSERT_EQ(host_key.dim(), static_cast<int64_t>(base_key_shape.size()) + 1);
164+
EXPECT_EQ(host_key.size(0),
165+
scale_host_block_count(base_key_shape[0], kHostFactor));
166+
EXPECT_EQ(host_key.size(1), kLayerCount);
167+
for (size_t i = 1; i < base_key_shape.size(); ++i) {
168+
EXPECT_EQ(host_key.size(static_cast<int64_t>(i) + 1), base_key_shape[i]);
169+
}
170+
}
171+
172+
TEST_F(HostKVCacheTest, HostKVCacheDeepSeekV4PerBlockType) {
173+
constexpr int64_t kSwaCount = 10;
174+
constexpr int64_t kC4Count = 32;
175+
constexpr int64_t kC128Count = 4;
176+
constexpr int64_t kBlockSize = 128;
177+
constexpr int64_t kHeadDim = 16;
178+
constexpr int64_t kIndexHeadDim = 8;
179+
constexpr double kHostFactor = 3.0;
180+
181+
KVCacheCapacity capacity;
182+
capacity.block_size(kBlockSize)
183+
.swa_count(kSwaCount)
184+
.c4_count(kC4Count)
185+
.c128_count(kC128Count);
186+
187+
ModelArgs model_args;
188+
model_args.model_type("deepseek_v4");
189+
KVCacheShape shape(capacity, model_args, /*world_size=*/1);
190+
191+
KVCacheCreateOptions options;
192+
options.device(torch::Device(torch::kCPU))
193+
.dtype(torch::kFloat32)
194+
.num_layers(3)
195+
.model_type("deepseek_v4")
196+
.block_size(kBlockSize)
197+
.head_dim(kHeadDim)
198+
.index_head_dim(kIndexHeadDim)
199+
.window_size(/*window_size=*/512)
200+
.compress_ratios({1, 4, 128})
201+
.host_blocks_factor(kHostFactor);
202+
203+
// SWA host cache: 1 layer in this 3-layer config (compress_ratio == 1).
204+
KVCache swa_host(shape, options, BlockType::SWA, /*layer_count=*/1);
205+
const BlockTypeTensorMap swa_tensors =
206+
swa_host.get_block_type_tensors(BlockType::SWA);
207+
ASSERT_TRUE(swa_tensors.count(KVCacheTensorRole::SWA) > 0);
208+
const torch::Tensor& swa = swa_tensors.at(KVCacheTensorRole::SWA);
209+
EXPECT_TRUE(swa.is_contiguous());
210+
EXPECT_EQ(swa.size(0), scale_host_block_count(kSwaCount, kHostFactor));
211+
EXPECT_EQ(swa.size(1), 1);
212+
213+
// C4 host cache: key + index, index uses the DSV4 index dtype.
214+
KVCache c4_host(shape, options, BlockType::C4, /*layer_count=*/1);
215+
const BlockTypeTensorMap c4_tensors =
216+
c4_host.get_block_type_tensors(BlockType::C4);
217+
ASSERT_TRUE(c4_tensors.count(KVCacheTensorRole::KEY) > 0);
218+
ASSERT_TRUE(c4_tensors.count(KVCacheTensorRole::INDEX) > 0);
219+
EXPECT_EQ(c4_tensors.at(KVCacheTensorRole::KEY).size(0),
220+
scale_host_block_count(kC4Count, kHostFactor));
221+
EXPECT_EQ(c4_tensors.at(KVCacheTensorRole::INDEX).scalar_type(),
222+
get_dsv4_cache_policy(options.dtype()).index_dtype);
223+
224+
// C128 host cache: key only (no index).
225+
KVCache c128_host(shape, options, BlockType::C128, /*layer_count=*/1);
226+
const BlockTypeTensorMap c128_tensors =
227+
c128_host.get_block_type_tensors(BlockType::C128);
228+
ASSERT_TRUE(c128_tensors.count(KVCacheTensorRole::KEY) > 0);
229+
EXPECT_TRUE(c128_tensors.count(KVCacheTensorRole::INDEX) == 0);
230+
EXPECT_EQ(c128_tensors.at(KVCacheTensorRole::KEY).size(0),
231+
scale_host_block_count(kC128Count, kHostFactor));
232+
}
233+
112234
} // namespace xllm

‎xllm/core/framework/kv_cache/deepseek_v4_kv_cache_impl.cpp‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,79 @@ DeepSeekV4KVCacheImpl::DeepSeekV4KVCacheImpl(
9595
compress_index_kv_state_(tensors.compress_index_kv_state),
9696
compress_index_score_state_(tensors.compress_index_score_state) {}
9797

98+
DeepSeekV4KVCacheImpl::DeepSeekV4KVCacheImpl(
99+
const KVCacheShape& kv_cache_shape,
100+
const KVCacheCreateOptions& create_options,
101+
BlockType type,
102+
int64_t layer_count) {
103+
CHECK(kv_cache_shape.has_key_cache_shape())
104+
<< "DeepSeek V4 host kv cache shape must contain pool counts.";
105+
const std::vector<int64_t>& pool_counts = kv_cache_shape.key_cache_shape();
106+
CHECK_GE(pool_counts.size(), 3) << "DeepSeek V4 host kv cache shape must be "
107+
<< "[swa_count, c4_count, c128_count].";
108+
CHECK_GT(create_options.block_size(), 0)
109+
<< "DeepSeek V4 block_size must be positive.";
110+
CHECK_GT(create_options.head_dim(), 0)
111+
<< "DeepSeek V4 head_dim must be positive.";
112+
CHECK_GT(layer_count, 0)
113+
<< "DeepSeek V4 host prefix cache layer count must be positive.";
114+
115+
const double factor = create_options.host_blocks_factor();
116+
const int64_t host_swa_count = scale_host_block_count(pool_counts[0], factor);
117+
const int64_t host_c4_count = scale_host_block_count(pool_counts[1], factor);
118+
const int64_t host_c128_count =
119+
scale_host_block_count(pool_counts[2], factor);
120+
const int64_t block_size = create_options.block_size();
121+
const int64_t head_dim = create_options.head_dim();
122+
const int64_t index_head_dim =
123+
std::max<int64_t>(create_options.index_head_dim(), 1);
124+
const int64_t n_heads = 1;
125+
const int64_t index_n_heads = 1;
126+
const DeepSeekV4CachePolicy cache_policy =
127+
get_dsv4_cache_policy(create_options.dtype());
128+
129+
// Host tensor shape: device per-block dims with a layer dimension inserted at
130+
// index 1, i.e. [host_block_count, layer_count, ...per_block_dims].
131+
auto host_group_shape = [&](int64_t block_count, int64_t heads, int64_t dim) {
132+
std::vector<int64_t> shape =
133+
dsv4_block_shape(block_count, block_size, heads, dim);
134+
shape.insert(shape.begin() + 1, layer_count);
135+
return shape;
136+
};
137+
138+
switch (type) {
139+
case BlockType::SWA:
140+
host_page_aligned_regions_.reserve(1);
141+
create_host_tensor(host_group_shape(host_swa_count, n_heads, head_dim),
142+
create_options.dtype(),
143+
&swa_cache_,
144+
nullptr);
145+
break;
146+
case BlockType::C4:
147+
host_page_aligned_regions_.reserve(2);
148+
create_host_tensor(host_group_shape(host_c4_count, n_heads, head_dim),
149+
create_options.dtype(),
150+
&key_cache_,
151+
nullptr);
152+
create_host_tensor(
153+
host_group_shape(host_c4_count, index_n_heads, index_head_dim),
154+
cache_policy.index_dtype,
155+
&index_cache_,
156+
nullptr);
157+
break;
158+
case BlockType::C128:
159+
host_page_aligned_regions_.reserve(1);
160+
create_host_tensor(host_group_shape(host_c128_count, n_heads, head_dim),
161+
create_options.dtype(),
162+
&key_cache_,
163+
nullptr);
164+
break;
165+
default:
166+
LOG(FATAL) << "Unsupported DeepSeek V4 host block type: "
167+
<< static_cast<int32_t>(type);
168+
}
169+
}
170+
98171
torch::Tensor DeepSeekV4KVCacheImpl::get_k_cache() const { return key_cache_; }
99172

100173
torch::Tensor DeepSeekV4KVCacheImpl::get_index_cache() const {
@@ -125,6 +198,34 @@ torch::Tensor DeepSeekV4KVCacheImpl::get_compress_index_score_state() const {
125198
return compress_index_score_state_;
126199
}
127200

201+
BlockTypeTensorMap DeepSeekV4KVCacheImpl::get_block_type_tensors(
202+
BlockType type) const {
203+
BlockTypeTensorMap tensor_map;
204+
switch (type) {
205+
case BlockType::SWA:
206+
if (swa_cache_.defined() && swa_cache_.numel() > 0) {
207+
tensor_map.emplace(KVCacheTensorRole::SWA, swa_cache_);
208+
}
209+
break;
210+
case BlockType::C4:
211+
if (key_cache_.defined() && key_cache_.numel() > 0 &&
212+
index_cache_.defined() && index_cache_.numel() > 0) {
213+
tensor_map.emplace(KVCacheTensorRole::KEY, key_cache_);
214+
tensor_map.emplace(KVCacheTensorRole::INDEX, index_cache_);
215+
}
216+
break;
217+
case BlockType::C128:
218+
if (key_cache_.defined() && key_cache_.numel() > 0 &&
219+
(!index_cache_.defined() || index_cache_.numel() == 0)) {
220+
tensor_map.emplace(KVCacheTensorRole::KEY, key_cache_);
221+
}
222+
break;
223+
default:
224+
break;
225+
}
226+
return tensor_map;
227+
}
228+
128229
bool DeepSeekV4KVCacheImpl::empty() const { return !swa_cache_.defined(); }
129230

130231
std::vector<std::vector<int64_t>> DeepSeekV4KVCacheImpl::get_shapes() const {

‎xllm/core/framework/kv_cache/deepseek_v4_kv_cache_impl.h‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,10 @@ class KVCacheShape;
2828
class DeepSeekV4KVCacheImpl final : public KVCacheImpl {
2929
public:
3030
explicit DeepSeekV4KVCacheImpl(const DeepSeekV4KVCacheTensors& tensors);
31+
DeepSeekV4KVCacheImpl(const KVCacheShape& kv_cache_shape,
32+
const KVCacheCreateOptions& create_options,
33+
BlockType type,
34+
int64_t layer_count);
3135

3236
torch::Tensor get_k_cache() const override;
3337
torch::Tensor get_index_cache() const override;
@@ -38,6 +42,8 @@ class DeepSeekV4KVCacheImpl final : public KVCacheImpl {
3842
torch::Tensor get_compress_index_kv_state() const override;
3943
torch::Tensor get_compress_index_score_state() const override;
4044

45+
BlockTypeTensorMap get_block_type_tensors(BlockType type) const override;
46+
4147
bool empty() const override;
4248

4349
std::vector<std::vector<int64_t>> get_shapes() const override;

‎xllm/core/framework/kv_cache/indexed_kv_cache_impl.cpp‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,10 +46,60 @@ IndexedKVCacheImpl::IndexedKVCacheImpl(
4646
index_cache_shape_ = kv_cache_shape.index_cache_shape();
4747
}
4848

49+
IndexedKVCacheImpl::IndexedKVCacheImpl(
50+
const KVCacheShape& kv_cache_shape,
51+
const KVCacheCreateOptions& create_options,
52+
BlockType type,
53+
int64_t layer_count)
54+
: KVCacheImpl() {
55+
CHECK(type == BlockType::KV)
56+
<< "IndexedKVCacheImpl host cache only supports BlockType::KV.";
57+
host_page_aligned_regions_.reserve(3);
58+
if (kv_cache_shape.has_key_cache_shape()) {
59+
create_host_tensor(
60+
build_host_group_tensor_shape(kv_cache_shape.key_cache_shape(),
61+
create_options.host_blocks_factor(),
62+
layer_count),
63+
create_options.dtype(),
64+
&key_cache_,
65+
&key_cache_shape_);
66+
}
67+
if (kv_cache_shape.has_value_cache_shape()) {
68+
create_host_tensor(
69+
build_host_group_tensor_shape(kv_cache_shape.value_cache_shape(),
70+
create_options.host_blocks_factor(),
71+
layer_count),
72+
create_options.dtype(),
73+
&value_cache_,
74+
&value_cache_shape_);
75+
}
76+
if (kv_cache_shape.has_index_cache_shape()) {
77+
create_host_tensor(
78+
build_host_group_tensor_shape(kv_cache_shape.index_cache_shape(),
79+
create_options.host_blocks_factor(),
80+
layer_count),
81+
create_options.dtype(),
82+
&index_cache_,
83+
&index_cache_shape_);
84+
}
85+
}
86+
4987
torch::Tensor IndexedKVCacheImpl::get_index_cache() const {
5088
return index_cache_;
5189
}
5290

91+
BlockTypeTensorMap IndexedKVCacheImpl::get_block_type_tensors(
92+
BlockType type) const {
93+
BlockTypeTensorMap tensor_map = KVCacheImpl::get_block_type_tensors(type);
94+
if (type != BlockType::KV) {
95+
return tensor_map;
96+
}
97+
if (index_cache_.defined() && index_cache_.numel() > 0) {
98+
tensor_map.emplace(KVCacheTensorRole::INDEX, index_cache_);
99+
}
100+
return tensor_map;
101+
}
102+
53103
bool IndexedKVCacheImpl::empty() const {
54104
return !key_cache_.defined() || !value_cache_.defined() ||
55105
!index_cache_.defined();

‎xllm/core/framework/kv_cache/indexed_kv_cache_impl.h‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,9 +24,15 @@ class IndexedKVCacheImpl final : public KVCacheImpl {
2424
explicit IndexedKVCacheImpl(const IndexedKVCacheTensors& tensors);
2525
IndexedKVCacheImpl(const KVCacheShape& kv_cache_shape,
2626
const KVCacheCreateOptions& create_options);
27+
IndexedKVCacheImpl(const KVCacheShape& kv_cache_shape,
28+
const KVCacheCreateOptions& create_options,
29+
BlockType type,
30+
int64_t layer_count);
2731

2832
torch::Tensor get_index_cache() const override;
2933

34+
BlockTypeTensorMap get_block_type_tensors(BlockType type) const override;
35+
3036
bool empty() const override;
3137

3238
std::vector<std::vector<int64_t>> get_shapes() const override;

0 commit comments

Comments
 (0)