diff --git a/tests/core/framework/kv_cache/kv_cache_estimation_test.cpp b/tests/core/framework/kv_cache/kv_cache_estimation_test.cpp index 96dceed0af..399e51c373 100644 --- a/tests/core/framework/kv_cache/kv_cache_estimation_test.cpp +++ b/tests/core/framework/kv_cache/kv_cache_estimation_test.cpp @@ -93,6 +93,23 @@ TEST(KVCacheEstimationTest, UserIndexerCacheDtypeDirectlyControlsQuantization) { } #if defined(USE_MLU) +TEST(KVCacheEstimationTest, SharedDsaLayersDoNotConsumeIndexerCacheBudget) { + ModelArgs model_args = make_standard_args(); + model_args.model_type("glm_moe_dsa") + .index_n_heads(1) + .index_head_dim(16) + .index_topk(8) + .index_topk_pattern("FSFS"); + KVCacheEstimateOptions options = make_estimate_options(); + + const KVCacheCapacity capacity = + estimate_kv_cache_capacity(model_args, options); + + EXPECT_EQ(capacity.num_full_attention_layers(), 4); + EXPECT_EQ(capacity.num_indexer_layers(), 2); + EXPECT_EQ(capacity.n_blocks(), 113); +} + TEST(KVCacheEstimationTest, RdmaScalePaddingReducesCapacityToLargestAllocatableBlockCount) { ModelArgs model_args = make_standard_args(); diff --git a/tests/core/framework/kv_cache/kv_cache_test.cpp b/tests/core/framework/kv_cache/kv_cache_test.cpp index 2a31cb6817..d627712eb1 100644 --- a/tests/core/framework/kv_cache/kv_cache_test.cpp +++ b/tests/core/framework/kv_cache/kv_cache_test.cpp @@ -25,10 +25,15 @@ limitations under the License. #include "framework/block/block.h" #include "framework/kv_cache/deepseek_v4_cache_policy.h" #include "framework/kv_cache/deepseek_v4_kv_cache_impl.h" +#include "framework/kv_cache/kv_cache_tensor_allocator.h" #include "framework/kv_cache/kv_cache_utils.h" #include "kv_cache_estimation.h" #include "kv_cache_shape.h" #include "platform/device.h" +#if defined(USE_MLU) +#include "platform/mlu/mlu_tensor_alloc.h" +#endif +#include "platform/platform.h" #include "worker.pb.h" namespace xllm { @@ -66,6 +71,28 @@ class IndexerCacheDtypeConfigGuard final { std::string old_indexer_cache_dtype_; }; +#if defined(USE_MLU) +class RecordingKVCacheTensorAllocator final : public KVCacheTensorAllocator { + public: + struct Request { + KVCacheTensorRole role; + std::vector shape; + torch::ScalarType dtype; + torch::Device device; + }; + + torch::Tensor allocate(KVCacheTensorRole role, + const std::vector& shape, + torch::ScalarType dtype, + const torch::Device& device) override { + requests.emplace_back(Request{role, shape, dtype, device}); + return torch::zeros(shape, torch::dtype(dtype).device(device)); + } + + std::vector requests; +}; +#endif + } // namespace TEST(Dsv4StateCacheTest, SplitStateReturnsInputsAndSwapsBoth) { @@ -577,6 +604,156 @@ TEST(KVCacheTest, MluIndexerAutoUsesDefaultCacheShapeWithoutScale) { EXPECT_EQ(caches[0].get_index_cache().scalar_type(), torch::kBFloat16); EXPECT_FALSE(caches[0].get_indexer_cache_scale().has_value()); } + +TEST(KVCacheTest, MluSharedDsaLayersAllocateOnlyMlaCache) { + constexpr int64_t kBlockCount = 8; + constexpr int64_t kBlockSize = 16; + + KVCacheCapacity capacity; + capacity.n_blocks(kBlockCount).block_size(kBlockSize); + + ModelArgs model_args; + model_args.model_type("glm_moe_dsa") + .enable_mla(true) + .n_heads(8) + .n_kv_heads(2) + .head_dim(64) + .kv_lora_rank(64) + .qk_rope_head_dim(16) + .index_n_heads(1) + .index_head_dim(32); + const KVCacheShape shape(capacity, model_args, /*world_size=*/1); + + KVCacheCreateOptions options; + options.device(torch::Device(torch::kCPU)) + .dtype(torch::kBFloat16) + .num_layers(4) + .model_type("glm_moe_dsa") + .enable_lighting_indexer(true) + .indexer_cache_enabled_layers( + std::vector{true, false, false, true}); + + std::vector caches; + allocate_kv_caches(caches, shape, options); + + ASSERT_EQ(caches.size(), 4U); + EXPECT_TRUE(caches[0].get_index_cache().defined()); + EXPECT_FALSE(caches[1].get_index_cache().defined()); + EXPECT_FALSE(caches[2].get_index_cache().defined()); + EXPECT_TRUE(caches[3].get_index_cache().defined()); + for (const KVCache& cache : caches) { + EXPECT_TRUE(cache.get_k_cache().defined()); + EXPECT_FALSE(cache.get_v_cache().defined()); + EXPECT_FALSE(cache.empty()); + } + + caches[1].get_k_cache()[0].fill_(1); + torch::Tensor source_block = torch::tensor({0}, torch::kLong); + torch::Tensor destination_block = torch::tensor({1}, torch::kLong); + caches[1].swap_blocks(source_block, destination_block); + EXPECT_TRUE( + torch::equal(caches[1].get_k_cache()[0], caches[1].get_k_cache()[1])); +} + +TEST(KVCacheTest, MluMixedDsaLayersUseInjectedAllocatorForActualRoles) { + constexpr int64_t kBlockCount = 2; + constexpr int64_t kBlockSize = 4; + + KVCacheCapacity capacity; + capacity.n_blocks(kBlockCount) + .block_size(kBlockSize) + .enable_indexer_cache_quant(true); + + ModelArgs model_args; + model_args.model_type("glm_moe_dsa") + .enable_mla(true) + .n_heads(8) + .n_kv_heads(2) + .head_dim(8) + .kv_lora_rank(8) + .qk_rope_head_dim(4) + .index_n_heads(1) + .index_head_dim(4); + const KVCacheShape shape(capacity, model_args, /*world_size=*/1); + + auto allocator = std::make_shared(); + KVCacheCreateOptions options; + options.device(torch::Device(torch::kCPU)) + .dtype(torch::kBFloat16) + .num_layers(4) + .model_type("glm_moe_dsa") + .enable_lighting_indexer(true) + .enable_indexer_cache_quant(true) + .indexer_cache_enabled_layers({true, false, true, false}) + .tensor_allocator(allocator); + + std::vector caches; + allocate_kv_caches(caches, shape, options); + + const std::vector expected_roles = { + KVCacheTensorRole::KEY, + KVCacheTensorRole::INDEX, + KVCacheTensorRole::INDEX_SCALE, + KVCacheTensorRole::KEY, + KVCacheTensorRole::KEY, + KVCacheTensorRole::INDEX, + KVCacheTensorRole::INDEX_SCALE, + KVCacheTensorRole::KEY}; + ASSERT_EQ(allocator->requests.size(), expected_roles.size()); + for (size_t i = 0; i < expected_roles.size(); ++i) { + EXPECT_EQ(allocator->requests[i].role, expected_roles[i]); + EXPECT_EQ(allocator->requests[i].device, torch::Device(torch::kCPU)); + } + EXPECT_EQ(allocator->requests[1].dtype, torch::kChar); + EXPECT_EQ(allocator->requests[2].dtype, torch::kFloat32); + EXPECT_EQ(allocator->requests[2].shape, shape.index_cache_scale_shape()); + + ASSERT_EQ(caches.size(), 4U); + EXPECT_EQ(caches[0].get_cache_tensors().size(), 3U); + EXPECT_EQ(caches[1].get_cache_tensors().size(), 1U); + EXPECT_EQ(caches[2].get_cache_tensors().size(), 3U); + EXPECT_EQ(caches[3].get_cache_tensors().size(), 1U); + EXPECT_FALSE(caches[1].get_index_cache().defined()); + EXPECT_FALSE(caches[3].get_indexer_cache_scale().has_value()); +} + +TEST(KVCacheTensorAllocatorTest, + MluMooncakePadsOnlyIndexScaleAndDoesNotOwnTensorLifetime) { + if (Platform::device_count() < 1) { + GTEST_SKIP() << "MLU device is required for allocator storage checks."; + } + + Device device(/*device_index=*/0); + device.set_device(); + const torch::Device torch_device = device.unwrap(); + std::shared_ptr allocator = + mlu_mooncake_tensor_allocator(); + std::weak_ptr allocator_lifetime = allocator; + + torch::Tensor key = allocator->allocate( + KVCacheTensorRole::KEY, {2, 4}, torch::kBFloat16, torch_device); + torch::Tensor index_scale = + allocator->allocate(KVCacheTensorRole::INDEX_SCALE, + {2, 96, 1}, + torch::kFloat32, + torch_device); + + EXPECT_EQ(key.storage().nbytes(), key.nbytes()); + EXPECT_EQ(index_scale.sizes().vec(), (std::vector{2, 96, 1})); + EXPECT_EQ(index_scale.scalar_type(), torch::kFloat32); + EXPECT_EQ(index_scale.nbytes(), 2 * 96 * sizeof(float)); + EXPECT_GE(index_scale.storage().nbytes(), index_scale.nbytes()); + EXPECT_EQ(mlu::get_rdma_registerable_nbytes(index_scale), + index_scale.storage().nbytes()); + + allocator.reset(); + EXPECT_TRUE(allocator_lifetime.expired()); + key.fill_(1); + index_scale.fill_(2); + device.synchronize_default_stream(); + EXPECT_TRUE(key.defined()); + EXPECT_TRUE(index_scale.defined()); +} #endif TEST_F(HostKVCacheTest, HostKVCacheNormalLayoutAddsLayerDim) { diff --git a/tests/core/framework/kv_cache_transfer/mooncake_transfer_engine_test.cpp b/tests/core/framework/kv_cache_transfer/mooncake_transfer_engine_test.cpp index 4b5fe0e7de..64ef28dce0 100644 --- a/tests/core/framework/kv_cache_transfer/mooncake_transfer_engine_test.cpp +++ b/tests/core/framework/kv_cache_transfer/mooncake_transfer_engine_test.cpp @@ -20,6 +20,7 @@ limitations under the License. #include #include +#include #include #include #include @@ -88,6 +89,35 @@ void expect_same_merge( } #if defined(USE_MLU) +class RecordingMooncakeTransferEngine final : public MooncakeTransferEngine { + public: + RecordingMooncakeTransferEngine(uint16_t listen_port, + const torch::Device& device) + : MooncakeTransferEngine(listen_port, device) {} + + bool register_memory(std::vector addrs, + std::vector lens, + std::vector buf_bytes) override { + registered_addrs.emplace_back(std::move(addrs)); + registered_lens.emplace_back(std::move(lens)); + registered_block_bytes.emplace_back(std::move(buf_bytes)); + return true; + } + + bool pull_memory_blocks(const std::string& /*remote_addr*/, + const std::vector& /*src_blocks*/, + const std::vector& /*dst_blocks*/, + const std::vector& buf_ids) override { + pulled_buf_ids = buf_ids; + return true; + } + + std::vector> registered_addrs; + std::vector> registered_lens; + std::vector> registered_block_bytes; + std::vector pulled_buf_ids; +}; + KVCacheShape make_indexer_int8_transfer_shape() { proto::KVCacheShape proto_shape; for (int64_t dim : std::vector{2, 1, 1, 16}) { @@ -102,6 +132,37 @@ KVCacheShape make_indexer_int8_transfer_shape() { } return KVCacheShape::from_proto(proto_shape); } + +std::vector make_mixed_transfer_caches(const torch::Device& device) { + std::shared_ptr allocator = + mlu_mooncake_tensor_allocator(); + auto make_full_cache = [&allocator, &device]() { + torch::Tensor key = allocator->allocate( + KVCacheTensorRole::KEY, {2, 1, 1, 16}, torch::kBFloat16, device); + torch::Tensor index = allocator->allocate( + KVCacheTensorRole::INDEX, {2, 96, 1, 8}, torch::kChar, device); + torch::Tensor index_scale = allocator->allocate( + KVCacheTensorRole::INDEX_SCALE, {2, 96, 1}, torch::kFloat32, device); + return KVCache(IndexedKVCacheTensors{ + KVCacheTensors{key, torch::Tensor()}, index, index_scale}); + }; + auto make_shared_cache = [&allocator, &device]() { + torch::Tensor key = allocator->allocate( + KVCacheTensorRole::KEY, {2, 1, 1, 16}, torch::kChar, device); + torch::Tensor key_scale = allocator->allocate( + KVCacheTensorRole::KEY_SCALE, {2, 1, 1}, torch::kFloat32, device); + return KVCache(QuantizedKVCacheTensors{ + KVCacheTensors{key, torch::Tensor()}, key_scale, torch::Tensor()}); + }; + + std::vector caches; + caches.reserve(4); + caches.emplace_back(make_full_cache()); + caches.emplace_back(make_shared_cache()); + caches.emplace_back(make_full_cache()); + caches.emplace_back(make_shared_cache()); + return caches; +} #endif } // namespace @@ -250,6 +311,20 @@ TEST(MooncakeKVCacheTransferDefaultTest, SpecDraftBufIdsUseSpecOffset) { EXPECT_EQ(transfer.get_buf_ids({0}, true), (std::vector{80, 81})); } +TEST(MooncakeKVCacheTransferDefaultTest, + VariableLayerBufIdsUseRegistrationOffsets) { + MooncakeKVCacheTransferDefault transfer( + 0, 0, torch::Device(torch::kCPU), "test"); + transfer.main_layout_.num_layers = 4; + transfer.main_layout_.offset = 10; + transfer.main_layout_.layer_offsets = {0, 4, 6, 8, 12}; + transfer.main_layout_.total_buf_cnt = 12; + transfer.main_layout_.registered = true; + + EXPECT_EQ(transfer.get_buf_ids({1, 3}, false), + (std::vector{14, 15, 18, 19, 20, 21})); +} + TEST(MooncakeKVCacheTransferDefaultTest, AddBufUsesRdmaRegisterableLengthWithoutChangingBlockBytes) { if (Platform::device_count() < 1) { @@ -282,6 +357,88 @@ TEST(MooncakeKVCacheTransferDefaultTest, EXPECT_EQ(block_bytes[0], kScaleBlockBytes); } +TEST(MooncakeKVCacheTransferDefaultTest, + RegistersMixedLayersFromProtocolRolesInStableOrder) { + if (Platform::device_count() < 1) { + GTEST_SKIP() << "MLU device is required for Mooncake registration tests."; + } + Device device(/*device_id=*/0); + device.set_device(); + const torch::Device torch_device = device.unwrap(); + auto engine = std::make_unique( + /*listen_port=*/0, torch_device); + RecordingMooncakeTransferEngine* engine_observer = engine.get(); + MooncakeKVCacheTransferDefault transfer(/*device_id=*/0, + /*listen_port=*/0, + torch_device, + /*model_type=*/"glm_moe_dsa", + std::move(engine)); + std::vector caches = make_mixed_transfer_caches(torch_device); + const KVCacheShape shape = make_indexer_int8_transfer_shape(); + + transfer.register_kv_cache(caches, shape, torch::kBFloat16); + + ASSERT_EQ(engine_observer->registered_addrs.size(), 1U); + ASSERT_EQ(engine_observer->registered_addrs[0].size(), 8U); + const std::vector expected_addrs = { + caches[0].get_k_cache().data_ptr(), + caches[0].get_index_cache().data_ptr(), + caches[0].get_indexer_cache_scale()->data_ptr(), + caches[1].get_k_cache().data_ptr(), + caches[2].get_k_cache().data_ptr(), + caches[2].get_index_cache().data_ptr(), + caches[2].get_indexer_cache_scale()->data_ptr(), + caches[3].get_k_cache().data_ptr()}; + EXPECT_EQ(engine_observer->registered_addrs[0], expected_addrs); + EXPECT_EQ(engine_observer->registered_lens[0][2], + caches[0].get_indexer_cache_scale()->storage().nbytes()); + EXPECT_EQ(engine_observer->registered_block_bytes[0][2], + caches[0].get_indexer_cache_scale()->nbytes() / 2); + + EXPECT_TRUE(transfer.pull_kv_blocks(/*src_cluster_id=*/1, + /*src_addr=*/"remote", + /*src_blocks=*/{0}, + /*dst_blocks=*/{1}, + /*src_linear_state_ids=*/{}, + /*dst_linear_state_ids=*/{})); + EXPECT_EQ(engine_observer->pulled_buf_ids, + (std::vector{0, 1, 2, 3, 4, 5, 6, 7})); +} + +TEST(MooncakeKVCacheTransferDefaultTest, + SpecRegistrationStartsAfterActualMainBufferCount) { + if (Platform::device_count() < 1) { + GTEST_SKIP() << "MLU device is required for Mooncake registration tests."; + } + Device device(/*device_id=*/0); + device.set_device(); + const torch::Device torch_device = device.unwrap(); + auto engine = std::make_unique( + /*listen_port=*/0, torch_device); + RecordingMooncakeTransferEngine* engine_observer = engine.get(); + MooncakeKVCacheTransferDefault transfer(/*device_id=*/0, + /*listen_port=*/0, + torch_device, + /*model_type=*/"glm_moe_dsa", + std::move(engine)); + std::vector main_caches = make_mixed_transfer_caches(torch_device); + std::vector draft_source = make_mixed_transfer_caches(torch_device); + std::vector draft_caches; + draft_caches.reserve(2); + draft_caches.emplace_back(std::move(draft_source[1])); + draft_caches.emplace_back(std::move(draft_source[0])); + const KVCacheShape shape = make_indexer_int8_transfer_shape(); + + transfer.register_kv_cache(main_caches, shape, torch::kBFloat16); + transfer.register_kv_cache_spec(draft_caches, shape, torch::kBFloat16); + + ASSERT_EQ(engine_observer->registered_addrs.size(), 2U); + EXPECT_EQ(engine_observer->registered_addrs[0].size(), 8U); + EXPECT_EQ(engine_observer->registered_addrs[1].size(), 4U); + EXPECT_EQ(transfer.get_buf_ids({0, 1}, /*is_spec_draft=*/true), + (std::vector{8, 9, 10, 11})); +} + TEST(MooncakeKVCacheTransferDefaultTest, AddBufRejectsNonContiguousTensor) { GTEST_FLAG_SET(death_test_style, "threadsafe"); MooncakeKVCacheTransferDefault transfer( @@ -299,6 +456,35 @@ TEST(MooncakeKVCacheTransferDefaultTest, AddBufRejectsNonContiguousTensor) { "contiguous"); } +TEST(MooncakeKVCacheTransferDefaultTest, AddBufRejectsRdmaLengthBeyondStorage) { + if (Platform::device_count() < 1) { + GTEST_SKIP() << "MLU device is required for Mooncake registration tests."; + } + GTEST_FLAG_SET(death_test_style, "threadsafe"); + Device device(/*device_id=*/0); + device.set_device(); + const torch::Device torch_device = device.unwrap(); + MooncakeKVCacheTransferDefault transfer( + /*device_id=*/0, + /*listen_port=*/0, + torch_device, + /*model_type=*/"test"); + const torch::Tensor tensor = + mlu::alloc_zero_tensor({2, 96, 1}, torch::kFloat32, torch_device); + std::vector addrs; + std::vector lens; + std::vector block_bytes; + + EXPECT_DEATH( + transfer.add_buf(tensor, + addrs, + lens, + block_bytes, + MooncakeKVCacheTransferDefault::RegisterLengthPolicy:: + RDMA_REGISTERABLE_BYTES), + "exceeds tensor storage capacity"); +} + TEST(MooncakeKVCacheTransferDefaultTest, IndexScaleRegistersAndRoundTripsWithKvBlocks) { if (Platform::device_count() < 1) { @@ -319,8 +505,16 @@ TEST(MooncakeKVCacheTransferDefaultTest, transfer.initialize(/*device_id=*/0); const KVCacheShape shape = make_indexer_int8_transfer_shape(); + KVCacheCreateOptions options; + options.device(torch_device) + .dtype(torch::kBFloat16) + .num_layers(1) + .model_type("deepseek_v32") + .enable_lighting_indexer(true) + .enable_indexer_cache_quant(true) + .tensor_allocator(mlu_mooncake_tensor_allocator()); std::vector caches; - transfer.allocate_kv_cache(caches, /*num_layers=*/1, shape, torch::kBFloat16); + allocate_kv_caches(caches, shape, options); ASSERT_EQ(caches.size(), 1U); KVCache& cache = caches[0]; diff --git a/xllm/core/distributed_runtime/llm_engine.cpp b/xllm/core/distributed_runtime/llm_engine.cpp index 7aa1f42f0d..388ef59a6e 100644 --- a/xllm/core/distributed_runtime/llm_engine.cpp +++ b/xllm/core/distributed_runtime/llm_engine.cpp @@ -501,6 +501,7 @@ bool LLMEngine::allocate_kv_cache(const KVCacheCapacity& kv_cache_cap) { << ", blocks: " << kv_cache_cap.n_blocks() << ", slot_size: " << kv_cache_cap.slot_size() << ", index_slot_size: " << kv_cache_cap.index_slot_size() + << ", indexer_layers: " << kv_cache_cap.num_indexer_layers() << ", scale_slot_size: " << kv_cache_cap.scale_slot_size() << ", linear_slot_size: " << kv_cache_cap.linear_slot_size() << ", linear_blocks: " << kv_cache_cap.num_linear_state_blocks() diff --git a/xllm/core/distributed_runtime/speculative_engine.cpp b/xllm/core/distributed_runtime/speculative_engine.cpp index c17abdc3fd..a121a33808 100644 --- a/xllm/core/distributed_runtime/speculative_engine.cpp +++ b/xllm/core/distributed_runtime/speculative_engine.cpp @@ -218,11 +218,20 @@ int64_t SpeculativeEngine::calculate_kv_cache( // Draft model has no linear-attention layers in the current MTP/Eagle path. const int64_t draft_full_attention_layers = draft_kv_cache_cap.n_layers(); const int64_t target_full_attention_block_size_in_bytes = - block_size * target_full_attention_layers * - target_full_attention_slot_size; + block_size * + (target_full_attention_layers * (target_kv_cache_cap.slot_size() + + target_kv_cache_cap.scale_slot_size()) + + target_kv_cache_cap.num_indexer_layers() * + target_kv_cache_cap.index_slot_size()); const int64_t draft_full_attention_block_size_in_bytes = - block_size * draft_full_attention_layers * - draft_allocated_full_attention_slot_size; + draft_body_uses_tp1 + ? block_size * (draft_full_attention_layers * + (draft_kv_cache_cap.slot_size() + + draft_kv_cache_cap.scale_slot_size()) + + draft_kv_cache_cap.num_indexer_layers() * + draft_kv_cache_cap.index_slot_size()) + : block_size * draft_full_attention_layers * + draft_allocated_full_attention_slot_size; const int64_t full_attention_block_size_in_bytes = target_full_attention_block_size_in_bytes + draft_full_attention_block_size_in_bytes; diff --git a/xllm/core/distributed_runtime/vlm_engine.cpp b/xllm/core/distributed_runtime/vlm_engine.cpp index 3fc9bc1471..e871cef05d 100644 --- a/xllm/core/distributed_runtime/vlm_engine.cpp +++ b/xllm/core/distributed_runtime/vlm_engine.cpp @@ -295,6 +295,7 @@ bool VLMEngine::allocate_kv_cache(const KVCacheCapacity& kv_cache_cap) { << ", blocks: " << kv_cache_cap.n_blocks() << ", slot_size: " << kv_cache_cap.slot_size() << ", index_slot_size: " << kv_cache_cap.index_slot_size() + << ", indexer_layers: " << kv_cache_cap.num_indexer_layers() << ", scale_slot_size: " << kv_cache_cap.scale_slot_size() << ", linear_slot_size: " << kv_cache_cap.linear_slot_size() << ", linear_blocks: " << kv_cache_cap.num_linear_state_blocks() diff --git a/xllm/core/framework/kv_cache/CMakeLists.txt b/xllm/core/framework/kv_cache/CMakeLists.txt index c9514d9cc0..a5c96d9e62 100644 --- a/xllm/core/framework/kv_cache/CMakeLists.txt +++ b/xllm/core/framework/kv_cache/CMakeLists.txt @@ -22,6 +22,7 @@ cc_library( :deepseek_v4_cache_policy :block_utils :mlu_rdma_memory_plan + :platform glog::glog proto::xllm_proto torch @@ -37,6 +38,7 @@ cc_library( kv_cache.h kv_cache_impl.h kv_cache_shape.h + kv_cache_tensor_allocator.h kv_cache_tensor_role.h kv_cache_utils.h linear_attention_kv_cache_impl.h @@ -48,6 +50,7 @@ cc_library( kv_cache.cpp kv_cache_impl.cpp kv_cache_shape.cpp + kv_cache_tensor_allocator.cpp kv_cache_utils.cpp linear_attention_kv_cache_impl.cpp quantized_kv_cache_impl.cpp diff --git a/xllm/core/framework/kv_cache/indexed_kv_cache_impl.cpp b/xllm/core/framework/kv_cache/indexed_kv_cache_impl.cpp index 93fcec921f..2725c4d180 100644 --- a/xllm/core/framework/kv_cache/indexed_kv_cache_impl.cpp +++ b/xllm/core/framework/kv_cache/indexed_kv_cache_impl.cpp @@ -169,8 +169,7 @@ BlockTypeTensorMap IndexedKVCacheImpl::get_block_type_tensors( } bool IndexedKVCacheImpl::empty() const { - return !key_cache_.defined() || !value_cache_.defined() || - !index_cache_.defined(); + return KVCacheImpl::empty() || !index_cache_.defined(); } std::vector> IndexedKVCacheImpl::get_shapes() const { diff --git a/xllm/core/framework/kv_cache/kv_cache.cpp b/xllm/core/framework/kv_cache/kv_cache.cpp index f22b5c7f55..5676304340 100644 --- a/xllm/core/framework/kv_cache/kv_cache.cpp +++ b/xllm/core/framework/kv_cache/kv_cache.cpp @@ -73,10 +73,27 @@ std::unique_ptr create_kv_cache_impl( create_options); } - if (create_options.enable_lighting_indexer()) { + bool enable_indexer_cache = create_options.enable_lighting_indexer(); + const std::vector& indexer_cache_enabled_layers = + create_options.indexer_cache_enabled_layers(); + if (!indexer_cache_enabled_layers.empty()) { + CHECK_EQ(indexer_cache_enabled_layers.size(), + static_cast(create_options.num_layers())) + << "Indexer cache layer mask must match num_layers."; + enable_indexer_cache = + enable_indexer_cache && + indexer_cache_enabled_layers[static_cast(layer_id)]; + } + + if (enable_indexer_cache) { return std::make_unique(kv_cache_shape, create_options); } + if (create_options.enable_kv_cache_quant()) { + return std::make_unique(kv_cache_shape, + create_options); + } + return std::make_unique(kv_cache_shape, create_options); } diff --git a/xllm/core/framework/kv_cache/kv_cache_capacity.h b/xllm/core/framework/kv_cache/kv_cache_capacity.h index 6f1d05b5d1..0892529e41 100644 --- a/xllm/core/framework/kv_cache/kv_cache_capacity.h +++ b/xllm/core/framework/kv_cache/kv_cache_capacity.h @@ -30,6 +30,7 @@ class KVCacheCapacity final { // for index cache PROPERTY(int64_t, index_slot_size) = 0; + PROPERTY(int64_t, num_indexer_layers) = 0; PROPERTY(bool, enable_indexer_cache_quant) = false; // for kv cache quantization scale cache diff --git a/xllm/core/framework/kv_cache/kv_cache_estimation.cpp b/xllm/core/framework/kv_cache/kv_cache_estimation.cpp index fdb0d7458c..c49921d692 100644 --- a/xllm/core/framework/kv_cache/kv_cache_estimation.cpp +++ b/xllm/core/framework/kv_cache/kv_cache_estimation.cpp @@ -20,6 +20,8 @@ limitations under the License. #include #include +#include "core/layers/common/dsa_topk_share_plan.h" +#include "core/platform/platform.h" #include "framework/block/block_utils.h" #include "framework/kv_cache/deepseek_v4_cache_policy.h" #include "framework/model/model_args.h" @@ -120,24 +122,34 @@ size_t checked_product(size_t lhs, size_t rhs, const char* description) { return lhs * rhs; } -size_t standard_full_cache_allocation_bytes(const KVCacheCapacity& kv_cache_cap, - int64_t n_blocks, - bool enable_rdma_scale_padding) { - CHECK_GT(n_blocks, 0) << "n_blocks must be positive"; +int64_t standard_full_cache_block_size_in_bytes( + const KVCacheCapacity& kv_cache_cap) { const int64_t full_attention_layers = std::max(kv_cache_cap.num_full_attention_layers(), 1); + const int64_t indexer_layers = kv_cache_cap.num_indexer_layers(); + CHECK_GE(indexer_layers, 0) << "num_indexer_layers must be non-negative"; + CHECK_LE(indexer_layers, full_attention_layers) + << "num_indexer_layers cannot exceed full-attention layers"; const int64_t logical_block_bytes = kv_cache_cap.block_size() * - (kv_cache_cap.slot_size() + kv_cache_cap.index_slot_size() + - kv_cache_cap.scale_slot_size()); + (full_attention_layers * + (kv_cache_cap.slot_size() + kv_cache_cap.scale_slot_size()) + + indexer_layers * kv_cache_cap.index_slot_size()); CHECK_GT(logical_block_bytes, 0) << "logical block bytes must be positive"; + return logical_block_bytes; +} + +size_t standard_full_cache_allocation_bytes(const KVCacheCapacity& kv_cache_cap, + int64_t n_blocks, + bool enable_rdma_scale_padding) { + CHECK_GT(n_blocks, 0) << "n_blocks must be positive"; + const int64_t logical_block_bytes = + standard_full_cache_block_size_in_bytes(kv_cache_cap); - const size_t logical_bytes = checked_product( + const size_t logical_bytes = checked_product(static_cast(n_blocks), - static_cast(full_attention_layers), - "full cache block count"), - static_cast(logical_block_bytes), - "full cache logical bytes"); + static_cast(logical_block_bytes), + "full cache logical bytes"); if (!enable_rdma_scale_padding) { return logical_bytes; } @@ -155,7 +167,7 @@ size_t standard_full_cache_allocation_bytes(const KVCacheCapacity& kv_cache_cap, const size_t padding_per_layer = scale_plan.registered_bytes - scale_plan.logical_bytes; const size_t total_padding = - checked_product(static_cast(full_attention_layers), + checked_product(static_cast(kv_cache_cap.num_indexer_layers()), padding_per_layer, "indexer scale RDMA padding"); CHECK_LE(total_padding, std::numeric_limits::max() - logical_bytes) @@ -207,28 +219,23 @@ int64_t linear_slot_size(const ModelArgs& model_args, 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) { + int64_t full_cache_block_size_in_bytes) { 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); + CHECK_GT(full_cache_block_size_in_bytes, 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(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); + kPaddingLinearStateBlocks * full_cache_block_size_in_bytes) / + (linear_bytes_per_block + full_cache_block_size_in_bytes); max_linear_blocks = std::min(max_linear_blocks, balanced_max_linear_blocks); return std::max(max_linear_blocks, kPaddingLinearStateBlocks); @@ -237,8 +244,7 @@ int64_t max_linear_state_blocks(int64_t cache_size_in_bytes, 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 full_cache_block_size_in_bytes, int64_t max_seqs_per_batch, int64_t max_linear_state_cache_slots, bool enable_prefix_cache) { @@ -251,11 +257,7 @@ int64_t calculate_linear_state_blocks(int64_t cache_size_in_bytes, max_linear_state_blocks(cache_size_in_bytes, num_linear_attention_layers, linear_slot_size, - num_full_attention_layers, - full_attention_block_size); - if (linear_slot_size <= 0 || num_linear_attention_layers <= 0) { - return kPaddingLinearStateBlocks; - } + full_cache_block_size_in_bytes); if (max_linear_state_cache_slots > 0) { const int64_t requested_blocks = max_linear_state_cache_slots + kPaddingLinearStateBlocks; @@ -483,17 +485,26 @@ void init_standard_counts(const ModelArgs& model_args, } } + if (kv_cache_cap->index_slot_size() > 0) { + int64_t num_indexer_layers = kv_cache_cap->num_full_attention_layers(); + const std::vector indexer_layer_mask = + resolve_indexer_cache_enabled_layers(model_args, + kv_cache_cap->n_layers()); + if (!indexer_layer_mask.empty()) { + num_indexer_layers = static_cast(std::count( + indexer_layer_mask.begin(), indexer_layer_mask.end(), true)); + } + kv_cache_cap->num_indexer_layers(num_indexer_layers); + } + const int64_t block_size = kv_cache_cap->block_size(); - const int64_t block_size_in_bytes = - block_size * - (kv_cache_cap->slot_size() + kv_cache_cap->index_slot_size() + - kv_cache_cap->scale_slot_size()); + const int64_t full_cache_block_size_in_bytes = + standard_full_cache_block_size_in_bytes(*kv_cache_cap); 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, + full_cache_block_size_in_bytes, options.max_seqs_per_batch, options.max_linear_state_cache_slots, options.enable_prefix_cache)); @@ -519,11 +530,8 @@ void init_standard_counts(const ModelArgs& model_args, CHECK_GT(available_full_cache_size_in_bytes, 0) << "no memory left for full-attention kv cache after reserving linear " "state cache"; - const int64_t full_attention_layers = - std::max(kv_cache_cap->num_full_attention_layers(), 1); const int64_t logical_n_blocks = - available_full_cache_size_in_bytes / - (full_attention_layers * block_size_in_bytes); + available_full_cache_size_in_bytes / full_cache_block_size_in_bytes; const bool enable_rdma_scale_padding = use_rdma_indexer_scale_padding(options, *kv_cache_cap); if (!enable_rdma_scale_padding) { @@ -546,7 +554,8 @@ void init_standard_counts(const ModelArgs& model_args, << "no memory for one KV cache block with RDMA-registerable indexer " "scales: available_bytes=" << available_full_cache_size_in_bytes - << ", full_attention_layers=" << full_attention_layers + << ", full_attention_layers=" << kv_cache_cap->num_full_attention_layers() + << ", indexer_layers=" << kv_cache_cap->num_indexer_layers() << ", scale_registered_bytes_per_layer=" << minimum_scale_plan.registered_bytes << ", minimum_allocation_bytes=" << minimum_allocation_bytes; @@ -575,13 +584,23 @@ void init_standard_counts(const ModelArgs& model_args, << " to " << kv_cache_cap->n_blocks() << " blocks, scale_logical_bytes_per_layer=" << scale_plan.logical_bytes << ", scale_registered_bytes_per_layer=" - << scale_plan.registered_bytes - << ", full_attention_layers=" << full_attention_layers; + << scale_plan.registered_bytes << ", full_attention_layers=" + << kv_cache_cap->num_full_attention_layers() + << ", indexer_layers=" << kv_cache_cap->num_indexer_layers(); CHECK_GT(kv_cache_cap->n_blocks(), 0) << "no n_blocks for kv cache"; } } // namespace +std::vector resolve_indexer_cache_enabled_layers( + const ModelArgs& model_args, + int64_t num_cache_layers) { + if (!Platform::supports_dsa_indexer_cache_elision()) { + return {}; + } + return layer::get_dsa_indexer_layer_mask(model_args, num_cache_layers); +} + KVCacheCapacity estimate_kv_cache_capacity( const ModelArgs& model_args, const KVCacheEstimateOptions& options) { diff --git a/xllm/core/framework/kv_cache/kv_cache_estimation.h b/xllm/core/framework/kv_cache/kv_cache_estimation.h index acb2dd7f4f..3107edffa4 100644 --- a/xllm/core/framework/kv_cache/kv_cache_estimation.h +++ b/xllm/core/framework/kv_cache/kv_cache_estimation.h @@ -19,6 +19,7 @@ limitations under the License. #include #include +#include #include "framework/kv_cache/kv_cache_capacity.h" @@ -56,6 +57,10 @@ struct Dsv4KVCacheEstimateCost { int64_t manager_blocks_per_unit = 1; }; +std::vector resolve_indexer_cache_enabled_layers( + const ModelArgs& model_args, + int64_t num_cache_layers); + KVCacheCapacity estimate_kv_cache_capacity( const ModelArgs& model_args, const KVCacheEstimateOptions& options); diff --git a/xllm/core/framework/kv_cache/kv_cache_impl.cpp b/xllm/core/framework/kv_cache/kv_cache_impl.cpp index 0eae09e8fb..4b5eaec89e 100644 --- a/xllm/core/framework/kv_cache/kv_cache_impl.cpp +++ b/xllm/core/framework/kv_cache/kv_cache_impl.cpp @@ -177,7 +177,8 @@ BlockTypeTensorMap KVCacheImpl::get_block_type_tensors(BlockType type) const { } bool KVCacheImpl::empty() const { - return !key_cache_.defined() || !value_cache_.defined(); + return !key_cache_.defined() || + (!value_cache_shape_.empty() && !value_cache_.defined()); } std::vector> KVCacheImpl::get_shapes() const { @@ -190,13 +191,13 @@ std::vector> KVCacheImpl::get_shapes() const { void KVCacheImpl::swap_blocks(torch::Tensor& src_tensor, torch::Tensor& dst_tensor) { - // batch select keys and values - auto selected_keys = torch::index_select(key_cache_, 0, src_tensor); - auto selected_values = torch::index_select(value_cache_, 0, src_tensor); - - // batch copy keys and values to dst indices + torch::Tensor selected_keys = torch::index_select(key_cache_, 0, src_tensor); key_cache_.index_copy_(0, dst_tensor, selected_keys); - value_cache_.index_copy_(0, dst_tensor, selected_values); + if (value_cache_.defined() && value_cache_.numel() > 0) { + torch::Tensor selected_values = + torch::index_select(value_cache_, 0, src_tensor); + value_cache_.index_copy_(0, dst_tensor, selected_values); + } } } // namespace xllm diff --git a/xllm/core/framework/kv_cache/kv_cache_tensor_allocator.cpp b/xllm/core/framework/kv_cache/kv_cache_tensor_allocator.cpp new file mode 100644 index 0000000000..3b27088e97 --- /dev/null +++ b/xllm/core/framework/kv_cache/kv_cache_tensor_allocator.cpp @@ -0,0 +1,67 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +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 "framework/kv_cache/kv_cache_tensor_allocator.h" + +#if defined(USE_MLU) +#include "platform/mlu/mlu_tensor_alloc.h" +#endif + +namespace xllm { +namespace { + +class DefaultKVCacheTensorAllocator final : public KVCacheTensorAllocator { + public: + torch::Tensor allocate(KVCacheTensorRole /*role*/, + const std::vector& shape, + torch::ScalarType dtype, + const torch::Device& device) override { + return torch::zeros(shape, torch::dtype(dtype).device(device)); + } +}; + +#if defined(USE_MLU) +class MluMooncakeTensorAllocator final : public KVCacheTensorAllocator { + public: + torch::Tensor allocate(KVCacheTensorRole role, + const std::vector& shape, + torch::ScalarType dtype, + const torch::Device& device) override { + // INDEX_SCALE is currently the only Mooncake role whose registered length + // can exceed its logical tensor bytes. Extend this branch explicitly if a + // future role requires RDMA padding. + if (role == KVCacheTensorRole::INDEX_SCALE) { + return mlu::alloc_rdma_registerable_zero_tensor(shape, dtype, device); + } + return mlu::alloc_zero_tensor(shape, dtype, device); + } +}; +#endif + +} // namespace + +std::shared_ptr default_kv_tensor_allocator() { + static std::shared_ptr allocator = + std::make_shared(); + return allocator; +} + +#if defined(USE_MLU) +std::shared_ptr mlu_mooncake_tensor_allocator() { + return std::make_shared(); +} +#endif + +} // namespace xllm diff --git a/xllm/core/framework/kv_cache/kv_cache_tensor_allocator.h b/xllm/core/framework/kv_cache/kv_cache_tensor_allocator.h new file mode 100644 index 0000000000..37a3b9aec1 --- /dev/null +++ b/xllm/core/framework/kv_cache/kv_cache_tensor_allocator.h @@ -0,0 +1,45 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +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. +==============================================================================*/ + +#pragma once + +#include + +#include +#include + +#include "framework/kv_cache/kv_cache_tensor_role.h" + +namespace xllm { + +// Controls only the physical allocation of tensors selected by the KV cache +// factory. Implementations must preserve the requested role, shape, and dtype. +class KVCacheTensorAllocator { + public: + virtual ~KVCacheTensorAllocator() = default; + + virtual torch::Tensor allocate(KVCacheTensorRole role, + const std::vector& shape, + torch::ScalarType dtype, + const torch::Device& device) = 0; +}; + +std::shared_ptr default_kv_tensor_allocator(); + +#if defined(USE_MLU) +std::shared_ptr mlu_mooncake_tensor_allocator(); +#endif + +} // namespace xllm diff --git a/xllm/core/framework/kv_cache/kv_cache_utils.cpp b/xllm/core/framework/kv_cache/kv_cache_utils.cpp index 6c2f58e64f..61e42b7ea9 100644 --- a/xllm/core/framework/kv_cache/kv_cache_utils.cpp +++ b/xllm/core/framework/kv_cache/kv_cache_utils.cpp @@ -26,9 +26,6 @@ limitations under the License. #include "core/framework/config/kv_cache_config.h" #include "framework/kv_cache/kv_cache_shape.h" -#if defined(USE_MLU) -#include "platform/mlu/mlu_tensor_alloc.h" -#endif #if defined(USE_NPU) #include "acl/acl_rt.h" @@ -81,14 +78,17 @@ torch::Tensor alloc_scale(const std::vector& cache_shape, torch::dtype(torch::kFloat32).device(device)); } -#if defined(USE_MLU) -torch::Tensor alloc_mlu_scale_for_transfer( - const std::vector& scale_shape, - const torch::Device& device) { - return mlu::alloc_rdma_registerable_zero_tensor( - scale_shape, torch::kFloat32, device); +torch::Tensor alloc_cache_tensor(KVCacheTensorRole role, + const std::vector& shape, + torch::ScalarType dtype, + const KVCacheCreateOptions& create_options) { + std::shared_ptr allocator = + create_options.tensor_allocator(); + if (allocator == nullptr) { + allocator = default_kv_tensor_allocator(); + } + return allocator->allocate(role, shape, dtype, create_options.device()); } -#endif } // namespace @@ -142,25 +142,15 @@ KVCacheTensors create_kv_cache_tensors( const KVCacheCreateOptions& create_options) { KVCacheTensors tensors; #if defined(USE_MLU) - if (create_options.enable_raw_device_allocator()) { - tensors.key_cache = mlu::alloc_zero_tensor(kv_cache_shape.key_cache_shape(), - create_options.dtype(), - create_options.device()); - if (kv_cache_shape.has_value_cache_shape()) { - tensors.value_cache = - mlu::alloc_zero_tensor(kv_cache_shape.value_cache_shape(), - create_options.dtype(), - create_options.device()); - } - } else { - tensors.key_cache = torch::zeros( - kv_cache_shape.key_cache_shape(), - torch::dtype(create_options.dtype()).device(create_options.device())); - if (!kv_cache_shape.value_cache_shape().empty()) { - tensors.value_cache = torch::zeros( - kv_cache_shape.value_cache_shape(), - torch::dtype(create_options.dtype()).device(create_options.device())); - } + tensors.key_cache = alloc_cache_tensor(KVCacheTensorRole::KEY, + kv_cache_shape.key_cache_shape(), + create_options.dtype(), + create_options); + if (kv_cache_shape.has_value_cache_shape()) { + tensors.value_cache = alloc_cache_tensor(KVCacheTensorRole::VALUE, + kv_cache_shape.value_cache_shape(), + create_options.dtype(), + create_options); } #elif defined(USE_NPU) const aclFormat npu_format_type = @@ -187,15 +177,17 @@ KVCacheTensors create_kv_cache_tensors( npu_format_type); } #else - tensors.key_cache = torch::zeros( - kv_cache_shape.key_cache_shape(), - torch::dtype(create_options.dtype()).device(create_options.device())); + tensors.key_cache = alloc_cache_tensor(KVCacheTensorRole::KEY, + kv_cache_shape.key_cache_shape(), + create_options.dtype(), + create_options); // deepseek_v3 model has no value cache on mlu device if (!kv_cache_shape.value_cache_shape().empty()) { - tensors.value_cache = torch::zeros( - kv_cache_shape.value_cache_shape(), - torch::dtype(create_options.dtype()).device(create_options.device())); + tensors.value_cache = alloc_cache_tensor(KVCacheTensorRole::VALUE, + kv_cache_shape.value_cache_shape(), + create_options.dtype(), + create_options); } #endif return tensors; @@ -222,16 +214,10 @@ IndexedKVCacheTensors create_indexed_kv_cache_tensors( const torch::ScalarType index_dtype = create_options.enable_indexer_cache_quant() ? torch::kChar : create_options.dtype(); - if (create_options.enable_raw_device_allocator()) { - tensors.index_cache = - mlu::alloc_zero_tensor(kv_cache_shape.index_cache_shape(), - index_dtype, - create_options.device()); - } else { - tensors.index_cache = - torch::zeros(kv_cache_shape.index_cache_shape(), - torch::dtype(index_dtype).device(create_options.device())); - } + tensors.index_cache = alloc_cache_tensor(KVCacheTensorRole::INDEX, + kv_cache_shape.index_cache_shape(), + index_dtype, + create_options); #elif defined(USE_NPU) const aclFormat npu_format_type = get_npu_kv_cache_format(create_options.model_type()); @@ -258,14 +244,11 @@ IndexedKVCacheTensors create_indexed_kv_cache_tensors( #endif if (kv_cache_shape.has_index_cache_scale_shape()) { #if defined(USE_MLU) - if (create_options.enable_raw_device_allocator()) { - tensors.index_cache_scale = alloc_mlu_scale_for_transfer( - kv_cache_shape.index_cache_scale_shape(), create_options.device()); - } else { - tensors.index_cache_scale = torch::zeros( - kv_cache_shape.index_cache_scale_shape(), - torch::dtype(torch::kFloat32).device(create_options.device())); - } + tensors.index_cache_scale = + alloc_cache_tensor(KVCacheTensorRole::INDEX_SCALE, + kv_cache_shape.index_cache_scale_shape(), + torch::kFloat32, + create_options); #else tensors.index_cache_scale = torch::zeros( kv_cache_shape.index_cache_scale_shape(), @@ -273,16 +256,14 @@ IndexedKVCacheTensors create_indexed_kv_cache_tensors( #endif } else { #if defined(USE_MLU) - if (create_options.enable_raw_device_allocator()) { - std::vector scale_shape( - kv_cache_shape.index_cache_shape().begin(), - kv_cache_shape.index_cache_shape().end() - 1); - tensors.index_cache_scale = - alloc_mlu_scale_for_transfer(scale_shape, create_options.device()); - } else { - tensors.index_cache_scale = alloc_scale( - kv_cache_shape.index_cache_shape(), create_options.device()); - } + std::vector scale_shape( + kv_cache_shape.index_cache_shape().begin(), + kv_cache_shape.index_cache_shape().end() - 1); + tensors.index_cache_scale = + alloc_cache_tensor(KVCacheTensorRole::INDEX_SCALE, + scale_shape, + torch::kFloat32, + create_options); #else tensors.index_cache_scale = alloc_scale( kv_cache_shape.index_cache_shape(), create_options.device()); @@ -312,17 +293,20 @@ QuantizedKVCacheTensors create_quantized_kv_cache_tensors( key_cache_shape.end() - 1); // float32 scale tensor for quantized KV cache (int8) - tensors.key_cache_scale = torch::zeros( - key_scale_shape, - torch::dtype(torch::kFloat32).device(create_options.device())); + tensors.key_cache_scale = alloc_cache_tensor(KVCacheTensorRole::KEY_SCALE, + key_scale_shape, + torch::kFloat32, + create_options); if (!kv_cache_shape.value_cache_shape().empty()) { const std::vector& value_cache_shape = kv_cache_shape.value_cache_shape(); std::vector value_scale_shape(value_cache_shape.begin(), value_cache_shape.end() - 1); - tensors.value_cache_scale = torch::zeros( - value_scale_shape, - torch::dtype(torch::kFloat32).device(create_options.device())); + tensors.value_cache_scale = + alloc_cache_tensor(KVCacheTensorRole::VALUE_SCALE, + value_scale_shape, + torch::kFloat32, + create_options); } return tensors; @@ -362,12 +346,14 @@ LinearAttentionKVCacheTensors create_linear_attention_kv_cache_tensors( ACL_FORMAT_ND); } #else - tensors.conv_cache = torch::zeros( - kv_cache_shape.conv_cache_shape(), - torch::dtype(create_options.dtype()).device(create_options.device())); - tensors.ssm_cache = torch::zeros( - kv_cache_shape.ssm_cache_shape(), - torch::dtype(create_options.ssm_dtype()).device(create_options.device())); + tensors.conv_cache = alloc_cache_tensor(KVCacheTensorRole::CONV, + kv_cache_shape.conv_cache_shape(), + create_options.dtype(), + create_options); + tensors.ssm_cache = alloc_cache_tensor(KVCacheTensorRole::SSM, + kv_cache_shape.ssm_cache_shape(), + create_options.ssm_dtype(), + create_options); #endif return tensors; diff --git a/xllm/core/framework/kv_cache/kv_cache_utils.h b/xllm/core/framework/kv_cache/kv_cache_utils.h index 4719a2a800..8b3b1d594d 100644 --- a/xllm/core/framework/kv_cache/kv_cache_utils.h +++ b/xllm/core/framework/kv_cache/kv_cache_utils.h @@ -41,6 +41,7 @@ limitations under the License. #include "framework/block/block.h" #include "framework/kv_cache/kv_cache_capacity.h" +#include "framework/kv_cache/kv_cache_tensor_allocator.h" #include "framework/kv_cache/kv_cache_tensor_role.h" namespace xllm { @@ -66,8 +67,11 @@ struct KVCacheCreateOptions { PROPERTY(bool, enable_sleep_mode) = false; PROPERTY(bool, enable_linear_attention) = false; PROPERTY(bool, enable_lighting_indexer) = false; + // Empty keeps the legacy all-layer behavior. Otherwise each entry controls + // whether that layer owns indexer cache tensors. + PROPERTY(std::vector, indexer_cache_enabled_layers); PROPERTY(bool, enable_kv_cache_quant) = false; - PROPERTY(bool, enable_raw_device_allocator) = false; + PROPERTY(std::shared_ptr, tensor_allocator); #if defined(USE_NPU) PROPERTY(bool, enable_kv_cache_huge_page_allocator) = false; #endif diff --git a/xllm/core/framework/kv_cache/quantized_kv_cache_impl.cpp b/xllm/core/framework/kv_cache/quantized_kv_cache_impl.cpp index 1261f9a706..25391adae9 100644 --- a/xllm/core/framework/kv_cache/quantized_kv_cache_impl.cpp +++ b/xllm/core/framework/kv_cache/quantized_kv_cache_impl.cpp @@ -47,22 +47,22 @@ std::optional QuantizedKVCacheImpl::get_v_cache_scale() const { void QuantizedKVCacheImpl::swap_blocks(torch::Tensor& src_tensor, torch::Tensor& dst_tensor) { - // batch select keys and values - auto selected_keys = torch::index_select(key_cache_, 0, src_tensor); - auto selected_values = torch::index_select(value_cache_, 0, src_tensor); - - // batch copy keys and values to dst indices + torch::Tensor selected_keys = torch::index_select(key_cache_, 0, src_tensor); key_cache_.index_copy_(0, dst_tensor, selected_keys); - value_cache_.index_copy_(0, dst_tensor, selected_values); + if (value_cache_.defined() && value_cache_.numel() > 0) { + torch::Tensor selected_values = + torch::index_select(value_cache_, 0, src_tensor); + value_cache_.index_copy_(0, dst_tensor, selected_values); + } // batch copy scale tensors if (key_cache_scale_.defined() && key_cache_scale_.numel() > 0) { - auto selected_k_scales = + torch::Tensor selected_k_scales = torch::index_select(key_cache_scale_, 0, src_tensor); key_cache_scale_.index_copy_(0, dst_tensor, selected_k_scales); } if (value_cache_scale_.defined() && value_cache_scale_.numel() > 0) { - auto selected_v_scales = + torch::Tensor selected_v_scales = torch::index_select(value_cache_scale_, 0, src_tensor); value_cache_scale_.index_copy_(0, dst_tensor, selected_v_scales); } diff --git a/xllm/core/framework/kv_cache_transfer/hierarchy_kv_cache_transfer.cpp b/xllm/core/framework/kv_cache_transfer/hierarchy_kv_cache_transfer.cpp index a13fa2bb57..6dea4848b7 100644 --- a/xllm/core/framework/kv_cache_transfer/hierarchy_kv_cache_transfer.cpp +++ b/xllm/core/framework/kv_cache_transfer/hierarchy_kv_cache_transfer.cpp @@ -207,7 +207,7 @@ void HierarchyKVCacheTransfer::create_host_cache() { KVCacheCreateOptions host_opts = create_options_; host_opts.device(torch::Device(torch::kCPU)) .enable_xtensor(false) - .enable_raw_device_allocator(false) + .tensor_allocator(nullptr) .host_blocks_factor(options_.host_blocks_factor()); #if defined(USE_NPU) host_opts.enable_kv_cache_huge_page_allocator(false); diff --git a/xllm/core/framework/kv_cache_transfer/kv_cache_transfer.cpp b/xllm/core/framework/kv_cache_transfer/kv_cache_transfer.cpp index 67c0fd2b74..48d615671f 100644 --- a/xllm/core/framework/kv_cache_transfer/kv_cache_transfer.cpp +++ b/xllm/core/framework/kv_cache_transfer/kv_cache_transfer.cpp @@ -385,7 +385,8 @@ std::shared_ptr KVCacheTransferFactory::create( transfer->initialize(device_id); CHECK(allocate_kv_cache_func(kv_cache_shape, - /*use_huge_page_allocator=*/true)) + /*use_huge_page_allocator=*/true, + /*tensor_allocator=*/nullptr)) << "Allocate KV cache failed."; transfer->register_kv_cache(kv_caches, kv_cache_shape, dtype); #else @@ -414,8 +415,20 @@ std::shared_ptr KVCacheTransferFactory::create( #endif mooncake_transfer->initialize(device_id); +#if defined(USE_MLU) + CHECK(allocate_kv_cache_func(kv_cache_shape, + /*use_huge_page_allocator=*/false, + mlu_mooncake_tensor_allocator())) + << "Allocate KV cache failed."; +#else + // TODO(xllm-kv-allocator): NPU/DCU/XTensor remains on its existing + // physical allocation path in the MLU Mooncake migration. A follow-up + // must route this backend through + // KVCacheCreateOptions::tensor_allocator without moving cache structure + // decisions into Transfer. Do not add indexer layer-mask handling here. mooncake_transfer->allocate_kv_cache( kv_caches, num_layers, kv_cache_shape, dtype); +#endif mooncake_transfer->register_kv_cache(kv_caches, kv_cache_shape, dtype); transfer = mooncake_transfer; diff --git a/xllm/core/framework/kv_cache_transfer/kv_cache_transfer.h b/xllm/core/framework/kv_cache_transfer/kv_cache_transfer.h index c631c0d5b1..aa9f845b79 100644 --- a/xllm/core/framework/kv_cache_transfer/kv_cache_transfer.h +++ b/xllm/core/framework/kv_cache_transfer/kv_cache_transfer.h @@ -19,6 +19,7 @@ limitations under the License. #include "common/types.h" #include "framework/kv_cache/kv_cache.h" +#include "framework/kv_cache/kv_cache_tensor_allocator.h" #if defined(USE_NPU) #include "platform/npu/npu_layer_synchronizer.h" #endif @@ -175,8 +176,10 @@ class KVCacheTransfer { class KVCacheTransferFactory { public: - using AllocateKVCacheFunc = - std::function; + using AllocateKVCacheFunc = std::function tensor_allocator)>; static std::shared_ptr create( const std::string& transfer_type, diff --git a/xllm/core/framework/kv_cache_transfer/mooncake_kv_cache_transfer.cpp b/xllm/core/framework/kv_cache_transfer/mooncake_kv_cache_transfer.cpp index dc1d02a61b..28df4305bd 100644 --- a/xllm/core/framework/kv_cache_transfer/mooncake_kv_cache_transfer.cpp +++ b/xllm/core/framework/kv_cache_transfer/mooncake_kv_cache_transfer.cpp @@ -83,6 +83,26 @@ void merge_xtensor_offsets( } } +std::vector get_mooncake_tensors(const KVCache& cache) { + std::vector transfer_tensors; + for (const KVCacheTensor& cache_tensor : cache.get_cache_tensors()) { + switch (cache_tensor.role) { + case KVCacheTensorRole::KEY: + case KVCacheTensorRole::VALUE: + case KVCacheTensorRole::INDEX: + case KVCacheTensorRole::INDEX_SCALE: + transfer_tensors.emplace_back(cache_tensor); + break; + default: + // Mooncake roles form an explicit protocol whitelist. A new cache role + // must not become transferable without a corresponding protocol + // decision and registration-order test. + break; + } + } + return transfer_tensors; +} + void merge_kv_info( std::unordered_map& merged_kv_infos, @@ -194,6 +214,18 @@ MooncakeKVCacheTransferDefault::MooncakeKVCacheTransferDefault( std::make_unique(listen_port, device)), model_type_(model_type) {} +MooncakeKVCacheTransferDefault::MooncakeKVCacheTransferDefault( + const int32_t device_id, + const uint16_t listen_port, + const torch::Device& device, + const std::string& model_type, + std::unique_ptr engine) + : MooncakeKVCacheTransferBase(device_id, + listen_port, + device, + std::move(engine)), + model_type_(model_type) {} + void MooncakeKVCacheTransferDefault::allocate_kv_cache( std::vector& kv_caches, const int64_t num_layers, @@ -223,22 +255,10 @@ void MooncakeKVCacheTransferDefault::register_kv_cache( const std::vector& key_cache_shape = kv_cache_shape.key_cache_shape(); bool has_v_cache = true; - bool has_index_cache = false; - bool has_index_cache_scale = false; if (!kv_caches.empty()) { torch::Tensor value_cache = kv_caches[0].get_v_cache(); - torch::Tensor index_cache = kv_caches[0].get_index_cache(); - std::optional index_cache_scale = - kv_caches[0].get_indexer_cache_scale(); has_v_cache = value_cache.defined() && value_cache.numel() > 0; - has_index_cache = index_cache.defined() && index_cache.numel() > 0; - has_index_cache_scale = index_cache_scale.has_value() && - index_cache_scale->defined() && - index_cache_scale->numel() > 0; } - const int64_t buf_cnt_per_layer = 1 + static_cast(has_v_cache) + - static_cast(has_index_cache) + - static_cast(has_index_cache_scale); int64_t data_size = torch::scalarTypeToTypeMeta(dtype).itemsize(); int64_t count_per_block = 1; @@ -255,10 +275,22 @@ void MooncakeKVCacheTransferDefault::register_kv_cache( BufLayout layout; layout.num_layers = num_layers; - layout.buf_cnt = buf_cnt_per_layer; + layout.layer_offsets.reserve(static_cast(num_layers) + 1); + layout.layer_offsets.emplace_back(0); + for (const KVCache& cache : kv_caches) { + const std::vector transfer_tensors = + get_mooncake_tensors(cache); + const int64_t buffer_count = static_cast(transfer_tensors.size()); + if (layout.layer_offsets.size() == 1) { + layout.buf_cnt = buffer_count; + } else if (layout.buf_cnt != buffer_count) { + layout.buf_cnt = 0; + } + layout.total_buf_cnt += buffer_count; + layout.layer_offsets.emplace_back(layout.total_buf_cnt); + } if (is_spec_draft) { - layout.offset = - main_layout_.offset + main_layout_.num_layers * main_layout_.buf_cnt; + layout.offset = main_layout_.offset + main_layout_.total_buf_cnt; } layout.registered = true; @@ -273,49 +305,39 @@ void MooncakeKVCacheTransferDefault::register_kv_cache( register_kv_cache_impl(kv_caches); } +void MooncakeKVCacheTransferDefault::register_kv_cache_spec( + std::vector& kv_caches, + const KVCacheShape& kv_cache_shape, + torch::ScalarType dtype) { + CHECK(main_layout_.registered) + << "Main KV cache must be registered before spec draft KV cache."; + register_kv_cache(kv_caches, kv_cache_shape, dtype); +} + void MooncakeKVCacheTransferDefault::allocate_kv_cache_impl( std::vector& kv_caches, int64_t num_layers, const KVCacheShape& kv_cache_shape, torch::ScalarType dtype) { - const std::vector& key_cache_shape = - kv_cache_shape.key_cache_shape(); - const std::vector& value_cache_shape = - kv_cache_shape.value_cache_shape(); #if defined(USE_MLU) - for (int64_t i = 0; i < num_layers; ++i) { - torch::Tensor key_cache = - mlu::alloc_zero_tensor(key_cache_shape, dtype, device_); - torch::Tensor value_cache; - torch::Tensor index_cache; - std::optional index_cache_scale; - if (kv_cache_shape.has_value_cache_shape()) { - value_cache = mlu::alloc_zero_tensor(value_cache_shape, dtype, device_); - } - if (kv_cache_shape.has_index_cache_shape()) { - const torch::ScalarType index_dtype = - kv_cache_shape.has_index_cache_scale_shape() ? torch::kChar : dtype; - index_cache = mlu::alloc_zero_tensor( - kv_cache_shape.index_cache_shape(), index_dtype, device_); - } - if (kv_cache_shape.has_index_cache_scale_shape()) { - index_cache_scale = mlu::alloc_rdma_registerable_zero_tensor( - kv_cache_shape.index_cache_scale_shape(), torch::kFloat32, device_); - } - if (index_cache.defined()) { - kv_caches.emplace_back( - IndexedKVCacheTensors{KVCacheTensors{key_cache, value_cache}, - index_cache, - index_cache_scale}); - } else { - kv_caches.emplace_back(KVCacheTensors{key_cache, value_cache}); - } - } + (void)kv_caches; + (void)num_layers; + (void)kv_cache_shape; + (void)dtype; + LOG(FATAL) << "MLU Mooncake cache allocation must use the KV cache factory."; #elif defined(USE_DCU) + // TODO(xllm-kv-allocator): DCU remains on its existing physical allocation + // path in the MLU Mooncake migration. A follow-up must route it through + // KVCacheCreateOptions::tensor_allocator without moving cache structure + // decisions into Transfer. Do not add indexer layer-mask handling here. CHECK(kv_cache_shape.has_value_cache_shape()) << "DCU Mooncake KV transfer requires a value cache shape."; CHECK(!kv_cache_shape.has_index_cache_shape()) << "DCU Mooncake KV transfer does not support index cache yet."; + const std::vector& key_cache_shape = + kv_cache_shape.key_cache_shape(); + const std::vector& value_cache_shape = + kv_cache_shape.value_cache_shape(); for (int64_t i = 0; i < num_layers; ++i) { torch::Tensor key_cache = @@ -325,6 +347,14 @@ void MooncakeKVCacheTransferDefault::allocate_kv_cache_impl( kv_caches.emplace_back(KVCacheTensors{key_cache, value_cache}); } #else + // TODO(xllm-kv-allocator): NPU remains on its existing physical allocation + // path in the MLU Mooncake migration. A follow-up must route it through + // KVCacheCreateOptions::tensor_allocator without moving cache structure + // decisions into Transfer. Do not add indexer layer-mask handling here. + const std::vector& key_cache_shape = + kv_cache_shape.key_cache_shape(); + const std::vector& value_cache_shape = + kv_cache_shape.value_cache_shape(); // Original mode: allocate device memory using aclrtMalloc // calculate the size of kv cache for each layer auto data_size = torch::elementSize(dtype); @@ -467,15 +497,44 @@ std::vector MooncakeKVCacheTransferDefault::get_buf_ids( } std::vector buf_ids; - buf_ids.reserve(active_layer_ids.size() * - static_cast(layout.buf_cnt)); + const bool has_variable_layout = !layout.layer_offsets.empty(); + if (has_variable_layout) { + CHECK_EQ(layout.layer_offsets.size(), + static_cast(layout.num_layers) + 1) + << "KV cache buffer layout offsets are invalid."; + } else { + CHECK_GT(layout.buf_cnt, 0) << "KV cache uniform buffer layout is invalid."; + } for (int64_t layer_id : active_layer_ids) { CHECK_GE(layer_id, 0) << "layer_id must be non-negative"; CHECK_LT(layer_id, layout.num_layers) << "layer_id out of range"; + } - int64_t buf_id = layout.offset + layer_id * layout.buf_cnt; - for (int64_t buf_idx = 0; buf_idx < layout.buf_cnt; ++buf_idx) { - buf_ids.emplace_back(buf_id++); + size_t buffer_count = 0; + for (int64_t layer_id : active_layer_ids) { + const int64_t begin = + has_variable_layout + ? layout.layer_offsets[static_cast(layer_id)] + : layer_id * layout.buf_cnt; + const int64_t end = + has_variable_layout + ? layout.layer_offsets[static_cast(layer_id) + 1] + : begin + layout.buf_cnt; + buffer_count += static_cast(end - begin); + } + buf_ids.reserve(buffer_count); + + for (int64_t layer_id : active_layer_ids) { + const int64_t begin = + has_variable_layout + ? layout.layer_offsets[static_cast(layer_id)] + : layer_id * layout.buf_cnt; + const int64_t end = + has_variable_layout + ? layout.layer_offsets[static_cast(layer_id) + 1] + : begin + layout.buf_cnt; + for (int64_t relative_id = begin; relative_id < end; ++relative_id) { + buf_ids.emplace_back(layout.offset + relative_id); } } return buf_ids; @@ -490,22 +549,23 @@ void MooncakeKVCacheTransferDefault::register_kv_cache_impl( lens.reserve(kv_caches.size() * 4); buf_bytes.reserve(kv_caches.size() * 4); - for (int64_t i = 0; i < static_cast(kv_caches.size()); ++i) { - add_buf(kv_caches[i].get_k_cache(), addrs, lens, buf_bytes); - add_buf(kv_caches[i].get_v_cache(), addrs, lens, buf_bytes); - add_buf(kv_caches[i].get_index_cache(), addrs, lens, buf_bytes); - std::optional index_cache_scale = - kv_caches[i].get_indexer_cache_scale(); - if (index_cache_scale.has_value()) { + for (const KVCache& cache : kv_caches) { + const std::vector transfer_tensors = + get_mooncake_tensors(cache); + for (const KVCacheTensor& cache_tensor : transfer_tensors) { + if (cache_tensor.role == KVCacheTensorRole::INDEX_SCALE) { #if defined(USE_MLU) - add_buf(index_cache_scale.value(), - addrs, - lens, - buf_bytes, - RegisterLengthPolicy::RDMA_REGISTERABLE_BYTES); + add_buf(cache_tensor.tensor, + addrs, + lens, + buf_bytes, + RegisterLengthPolicy::RDMA_REGISTERABLE_BYTES); #else - add_buf(index_cache_scale.value(), addrs, lens, buf_bytes); + add_buf(cache_tensor.tensor, addrs, lens, buf_bytes); #endif + continue; + } + add_buf(cache_tensor.tensor, addrs, lens, buf_bytes); } } diff --git a/xllm/core/framework/kv_cache_transfer/mooncake_kv_cache_transfer.h b/xllm/core/framework/kv_cache_transfer/mooncake_kv_cache_transfer.h index 2951131e96..d96794bc7f 100644 --- a/xllm/core/framework/kv_cache_transfer/mooncake_kv_cache_transfer.h +++ b/xllm/core/framework/kv_cache_transfer/mooncake_kv_cache_transfer.h @@ -62,6 +62,12 @@ class MooncakeKVCacheTransferDefault final const uint16_t listen_port, const torch::Device& device, const std::string& model_type); + MooncakeKVCacheTransferDefault( + const int32_t device_id, + const uint16_t listen_port, + const torch::Device& device, + const std::string& model_type, + std::unique_ptr engine); void allocate_kv_cache(std::vector& kv_caches, const int64_t num_layers, @@ -77,6 +83,10 @@ class MooncakeKVCacheTransferDefault final const KVCacheShape& kv_cache_shape, const torch::ScalarType dtype) override; + void register_kv_cache_spec(std::vector& kv_caches, + const KVCacheShape& kv_cache_shape, + torch::ScalarType dtype) override; + bool pull_kv_blocks( const uint64_t src_cluster_id, const std::string& src_addr, @@ -103,10 +113,16 @@ class MooncakeKVCacheTransferDefault final struct BufLayout { // Number of layers owned by this layout. int64_t num_layers = 0; - // Number of registered buffers per layer, such as K/V/index cache. - int64_t buf_cnt = 0; // Starting buffer id of this layout in the Mooncake registration table. int64_t offset = 0; + // Registration-order offsets for each layer plus one terminal offset. + // Shared DSA layers omit indexer buffers, so counts may vary by layer. + std::vector layer_offsets; + // Legacy uniform buffers-per-layer view. It is zero when buffer counts + // vary. Keep this for callers that construct a uniform layout directly. + int64_t buf_cnt = 0; + // Total buffers registered by this layout. + int64_t total_buf_cnt = 0; // True after the corresponding KV cache memory has been registered. bool registered = false; }; diff --git a/xllm/core/framework/kv_cache_transfer/mooncake_transfer_engine.h b/xllm/core/framework/kv_cache_transfer/mooncake_transfer_engine.h index 93b7eb32ef..4f58178b8e 100644 --- a/xllm/core/framework/kv_cache_transfer/mooncake_transfer_engine.h +++ b/xllm/core/framework/kv_cache_transfer/mooncake_transfer_engine.h @@ -94,7 +94,7 @@ class MooncakeTransferEngineCore { stub_map_; }; -class MooncakeTransferEngine final { +class MooncakeTransferEngine { public: enum class MoveOpcode { READ = 0, WRITE = 1 }; @@ -104,9 +104,9 @@ class MooncakeTransferEngine final { std::string initialize(); - bool register_memory(std::vector addrs, - std::vector lens, - std::vector buf_bytes); + virtual bool register_memory(std::vector addrs, + std::vector lens, + std::vector buf_bytes); bool move_memory_blocks(const std::string& remote_addr, const std::vector& src_blocks, @@ -114,15 +114,15 @@ class MooncakeTransferEngine final { const std::vector& buf_ids, MoveOpcode move_opcode); - bool pull_memory_blocks(const std::string& remote_addr, - const std::vector& src_blocks, - const std::vector& dst_blocks, - const std::vector& buf_ids); + virtual bool pull_memory_blocks(const std::string& remote_addr, + const std::vector& src_blocks, + const std::vector& dst_blocks, + const std::vector& buf_ids); - bool push_memory_blocks(const std::string& remote_addr, - const std::vector& src_blocks, - const std::vector& dst_blocks, - const std::vector& buf_ids); + virtual bool push_memory_blocks(const std::string& remote_addr, + const std::vector& src_blocks, + const std::vector& dst_blocks, + const std::vector& buf_ids); // XTensor mode uses raw offsets in the GlobalXTensor region in buffer[0]. bool move_memory_by_global_offsets(const std::string& remote_addr, diff --git a/xllm/core/layers/common/dsa_topk_share_plan.h b/xllm/core/layers/common/dsa_topk_share_plan.h index c8a2ba84e7..2d11bd87cd 100644 --- a/xllm/core/layers/common/dsa_topk_share_plan.h +++ b/xllm/core/layers/common/dsa_topk_share_plan.h @@ -153,8 +153,6 @@ inline std::vector get_dsa_indexer_layer_mask(const ModelArgs& args, return layer_mask; } - CHECK_EQ(num_cache_layers, args.n_layers()) - << "Main-model DSA cache layers must match num_hidden_layers."; const DsaTopkSharePlan topk_share_plan(args); for (int32_t layer_id = 0; layer_id < num_cache_layers; ++layer_id) { layer_mask[static_cast(layer_id)] = diff --git a/xllm/core/platform/platform.h b/xllm/core/platform/platform.h index 67c2ed77c6..906614ce01 100644 --- a/xllm/core/platform/platform.h +++ b/xllm/core/platform/platform.h @@ -56,6 +56,13 @@ class Platform final { // after MLU moves CP input preparation into WorkerImpl. static constexpr bool uses_model_cp_partition() { return is_mlu(); } + // MLU can reuse DSA top-k results across layers without keeping an indexer + // cache for every layer. Other backends retain the legacy all-layer cache + // allocation until they implement the same cache-elision contract. + static constexpr bool supports_dsa_indexer_cache_elision() { + return is_mlu(); + } + static constexpr bool is_ilu() { #if defined(USE_ILU) return true; diff --git a/xllm/core/runtime/CMakeLists.txt b/xllm/core/runtime/CMakeLists.txt index 18ac236cc3..0922fe1e3a 100644 --- a/xllm/core/runtime/CMakeLists.txt +++ b/xllm/core/runtime/CMakeLists.txt @@ -88,6 +88,7 @@ cc_library( :model :parallel_state :kv_cache + :kv_cache_estimation :kv_cache_transfer :linear_state_restore :kv_transfer_completion diff --git a/xllm/core/runtime/worker_impl.cpp b/xllm/core/runtime/worker_impl.cpp index 203e62e6ea..5fa33ae766 100644 --- a/xllm/core/runtime/worker_impl.cpp +++ b/xllm/core/runtime/worker_impl.cpp @@ -51,6 +51,7 @@ limitations under the License. #include "core/framework/config/profile_config.h" #include "core/framework/config/scheduler_config.h" #include "core/framework/config/speculative_config.h" +#include "core/framework/kv_cache/kv_cache_estimation.h" #include "core/platform/platform.h" #include "core/platform/sleepable_allocator.h" #if defined(USE_NPU) @@ -298,9 +299,10 @@ WorkerImpl::WorkerImpl(const ParallelArgs& parallel_args, WorkerImpl::~WorkerImpl() = default; -bool WorkerImpl::allocate_kv_cache_storage(const KVCacheShape& kv_cache_shape, - bool use_huge_page_allocator, - bool enable_raw_device_allocator) { +bool WorkerImpl::allocate_kv_cache_storage( + const KVCacheShape& kv_cache_shape, + bool use_huge_page_allocator, + std::shared_ptr tensor_allocator) { CHECK(model_ != nullptr) << "Model is not initialized."; CHECK(kv_caches_.empty()) << "KV caches are already initialized."; @@ -331,6 +333,8 @@ bool WorkerImpl::allocate_kv_cache_storage(const KVCacheShape& kv_cache_shape, << "simultaneously."; const int64_t num_layers = get_num_layers(); + std::vector indexer_cache_enabled_layers = + resolve_indexer_cache_enabled_layers(args, num_layers); // Check if KV cache quantization is enabled // "auto" (default): cache dtype aligns with model dtype (no quantization) @@ -379,9 +383,10 @@ bool WorkerImpl::allocate_kv_cache_storage(const KVCacheShape& kv_cache_shape, .enable_sleep_mode(options_.enable_sleep_mode()) .enable_linear_attention(enable_linear_attention) .enable_lighting_indexer(enable_lighting_indexer) + .indexer_cache_enabled_layers(std::move(indexer_cache_enabled_layers)) .enable_kv_cache_quant(enable_kv_cache_quant) .enable_indexer_cache_quant(enable_indexer_cache_quant) - .enable_raw_device_allocator(enable_raw_device_allocator) + .tensor_allocator(std::move(tensor_allocator)) .block_size(options_.block_size()) .head_dim(args.head_dim()) .index_head_dim(std::max(args.index_head_dim(), 1)) @@ -419,9 +424,9 @@ bool WorkerImpl::allocate_kv_cache_with_transfer( CHECK(kv_caches_.empty()) << "KV caches are already initialized."; // create a KVCache for each layer - const int64_t num_layers = context_.get_model_args().n_layers(); - const bool enable_lighting_indexer = - context_.get_model_args().index_n_heads() > 0; + const ModelArgs& model_args = context_.get_model_args(); + const int64_t num_layers = model_args.n_layers(); + const bool enable_lighting_indexer = model_args.index_n_heads() > 0; kv_cache_transfer_ = KVCacheTransferFactory::create( ::xllm::DisaggPDConfig::get_instance().kv_cache_transfer_type(), options_.transfer_listen_port(), @@ -431,11 +436,14 @@ bool WorkerImpl::allocate_kv_cache_with_transfer( dtype_, kv_caches_, num_layers, - [this](const KVCacheShape& shape, bool use_huge_page_allocator) { - return this->allocate_kv_cache_storage(shape, use_huge_page_allocator); + [this](const KVCacheShape& shape, + bool use_huge_page_allocator, + std::shared_ptr tensor_allocator) { + return this->allocate_kv_cache_storage( + shape, use_huge_page_allocator, std::move(tensor_allocator)); }, enable_lighting_indexer, - context_.get_model_args().model_type(), + model_args.model_type(), options_.model_id()); status_ = Status::READY; @@ -451,22 +459,22 @@ bool WorkerImpl::allocate_kv_cache_with_transfer( kv_cache_transfer_ = kv_cache_transfer; + std::shared_ptr tensor_allocator; +#if defined(USE_MLU) + tensor_allocator = mlu_mooncake_tensor_allocator(); +#endif if (!allocate_kv_cache_storage(kv_cache_shape, /*use_huge_page_allocator=*/true, - /*enable_raw_device_allocator=*/true)) { + std::move(tensor_allocator))) { return false; } -#if defined(USE_NPU) if (is_spec_draft_) { kv_cache_transfer_->register_kv_cache_spec( kv_caches_, kv_cache_shape, dtype_); } else { kv_cache_transfer_->register_kv_cache(kv_caches_, kv_cache_shape, dtype_); } -#else - kv_cache_transfer_->register_kv_cache(kv_caches_, kv_cache_shape, dtype_); -#endif status_ = Status::READY; return true; diff --git a/xllm/core/runtime/worker_impl.h b/xllm/core/runtime/worker_impl.h index 62f8b97266..a1da04c390 100644 --- a/xllm/core/runtime/worker_impl.h +++ b/xllm/core/runtime/worker_impl.h @@ -224,9 +224,10 @@ class WorkerImpl { bool can_skip_npu_graph_decode_sync( const ModelInputParams& input_params) const; - bool allocate_kv_cache_storage(const KVCacheShape& kv_cache_shape, - bool use_huge_page_allocator = false, - bool enable_raw_device_allocator = false); + bool allocate_kv_cache_storage( + const KVCacheShape& kv_cache_shape, + bool use_huge_page_allocator = false, + std::shared_ptr tensor_allocator = nullptr); // Get the effective number of layers based on whether this is a spec draft // model diff --git a/xllm/models/llm/deepseek_v2.h b/xllm/models/llm/deepseek_v2.h index f371e72b94..0e180b19b5 100644 --- a/xllm/models/llm/deepseek_v2.h +++ b/xllm/models/llm/deepseek_v2.h @@ -17,6 +17,7 @@ limitations under the License. #include #include +#include #include #include @@ -32,8 +33,24 @@ limitations under the License. namespace xllm { class DeepseekV2ModelImpl : public torch::nn::Module { + protected: + using DecoderLayerFactory = + std::function; + + static DecoderLayerFactory default_decoder_layer_factory() { + return [](const ModelContext& context, int32_t layer_id) { + return layer::DeepseekV2DecoderLayer(context, layer_id); + }; + } + public: - DeepseekV2ModelImpl(const ModelContext& context) + explicit DeepseekV2ModelImpl(const ModelContext& context) + : DeepseekV2ModelImpl(context, default_decoder_layer_factory()) {} + + protected: + DeepseekV2ModelImpl(const ModelContext& context, + const DecoderLayerFactory& decoder_layer_factory) : model_args_(context.get_model_args()), device_(context.get_tensor_options().device()) { auto options = context.get_tensor_options(); @@ -55,7 +72,7 @@ class DeepseekV2ModelImpl : public torch::nn::Module { // create decoder layers for (int32_t i = 0; i < model_args_.n_layers(); ++i) { - auto block = layer::DeepseekV2DecoderLayer(context, i); + auto block = decoder_layer_factory(context, i); layers_.push_back(block); blocks_->push_back(block); } @@ -70,6 +87,7 @@ class DeepseekV2ModelImpl : public torch::nn::Module { } } + public: ModelOutput forward_native(torch::Tensor tokens, torch::Tensor positions, std::vector& kv_caches, @@ -104,12 +122,14 @@ class DeepseekV2ModelImpl : public torch::nn::Module { auto& layer = layers_[i]; prepare_decoder_layer_for_forward(i, layer, attn_metadata); - hidden_states = layer(hidden_states, - residual, - positions, - attn_metadata, - kv_caches[i], - modified_input_params); + hidden_states = forward_decoder_layer(i, + layer, + hidden_states, + residual, + positions, + attn_metadata, + kv_caches[i], + modified_input_params); if (!modified_input_params.record_layer(static_cast(i), hidden_states.device())) { return ModelOutput(); @@ -157,6 +177,26 @@ class DeepseekV2ModelImpl : public torch::nn::Module { layer::DeepseekV2DecoderLayer& /*layer*/, const layer::AttentionMetadata& /*attn_metadata*/) {} + // Single extension point for the per-layer invocation so subclasses can + // attach model-specific per-layer plumbing (e.g. GLM5.2 DSA cross-layer + // top-k sharing in glm52.h) without forking the whole forward loop. + virtual torch::Tensor forward_decoder_layer( + size_t layer_id, + layer::DeepseekV2DecoderLayer& layer, + torch::Tensor& hidden_states, + std::optional& residual, + torch::Tensor& positions, + layer::AttentionMetadata& attn_metadata, + KVCache& kv_cache, + const ModelInputParams& input_params) { + return layer(hidden_states, + residual, + positions, + attn_metadata, + kv_cache, + input_params); + } + layer::WordEmbedding& embed_mod() { return embed_tokens_; } std::vector& layers_ref() { return layers_; } diff --git a/xllm/models/llm/deepseek_v32.h b/xllm/models/llm/deepseek_v32.h index 0acaf94651..4c2fcd7701 100644 --- a/xllm/models/llm/deepseek_v32.h +++ b/xllm/models/llm/deepseek_v32.h @@ -46,7 +46,12 @@ inline std::optional validate_deepseek_v32_cp_config( class DeepseekV32ModelImpl : public DeepseekV2ModelImpl { public: explicit DeepseekV32ModelImpl(const ModelContext& context) - : DeepseekV2ModelImpl(context), + : DeepseekV32ModelImpl(context, default_decoder_layer_factory()) {} + + protected: + DeepseekV32ModelImpl(const ModelContext& context, + const DecoderLayerFactory& decoder_layer_factory) + : DeepseekV2ModelImpl(context, decoder_layer_factory), device_(context.get_tensor_options().device()), cp_group_(context.get_parallel_args().cp_group_), parallel_world_size_(context.get_parallel_args().world_size()), @@ -56,6 +61,7 @@ class DeepseekV32ModelImpl : public DeepseekV2ModelImpl { CHECK(!cp_config_error.has_value()) << cp_config_error.value(); } + public: ModelOutput forward(const torch::Tensor& tokens, const torch::Tensor& positions, std::vector& kv_caches, diff --git a/xllm/models/llm/glm5.h b/xllm/models/llm/glm5.h index e72489e3da..7eeece2aef 100644 --- a/xllm/models/llm/glm5.h +++ b/xllm/models/llm/glm5.h @@ -22,6 +22,11 @@ class Glm5ModelImpl : public DeepseekV32ModelImpl { public: explicit Glm5ModelImpl(const ModelContext& context) : DeepseekV32ModelImpl(context) {} + + protected: + Glm5ModelImpl(const ModelContext& context, + const DecoderLayerFactory& decoder_layer_factory) + : DeepseekV32ModelImpl(context, decoder_layer_factory) {} }; TORCH_MODULE(Glm5Model); @@ -39,8 +44,9 @@ class Glm5ForCausalLMImpl : public LlmForCausalLMImplBase { }; TORCH_MODULE(Glm5ForCausalLM); -// register the causal model -REGISTER_CAUSAL_MODEL(glm_moe_dsa, Glm5ForCausalLM); +// GLM5.0/5.1 and GLM5.2 share model_type "glm_moe_dsa". Register it once in +// glm52.h; Glm52ForCausalLM falls back to the legacy GLM5 behavior when the +// resolved DSA top-k share plan contains no shared layers. // register the model args // example config: @@ -95,6 +101,9 @@ REGISTER_MODEL_ARGS( LOAD_ARG_OR(index_topk_freq, "index_topk_freq", 1); LOAD_ARG_OR(index_topk_pattern, "index_topk_pattern", ""); LOAD_ARG_OR(index_skip_topk_offset, "index_skip_topk_offset", 0); + LOAD_ARG_OR(index_share_for_mtp_iteration, + "index_share_for_mtp_iteration", + false); // Computed parameters // the original head_dim in glm5 config seem useless diff --git a/xllm/models/llm/glm52.h b/xllm/models/llm/glm52.h new file mode 100644 index 0000000000..fbf22b7987 --- /dev/null +++ b/xllm/models/llm/glm52.h @@ -0,0 +1,105 @@ +/* Copyright 2025-2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +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. +==============================================================================*/ +#pragma once + +#include "core/layers/common/dsa_topk_share_plan.h" +#include "core/layers/mlu/dsa_topk_relay.h" +#include "models/llm/glm5.h" +#include "platform/platform.h" + +// GLM5.2 shares model_type "glm_moe_dsa" with GLM5.0/5.1 and is told apart by +// the resolved indexer top-k share plan. GLM5.0/5.1 configs carry no reuse, so +// every layer stays non-sharing and the relay below bypasses itself entirely -- +// Glm52 then behaves exactly like Glm5. glm_moe_dsa is therefore registered +// once here on Glm52ForCausalLM, while REGISTER_MODEL_ARGS(glm_moe_dsa, ...) +// stays in glm5.h. + +namespace xllm { + +class Glm52ModelImpl : public Glm5ModelImpl { + public: + explicit Glm52ModelImpl(const ModelContext& context) + : Glm5ModelImpl(context, create_decoder_layer_factory(context)), + dsa_topk_share_plan_(context.get_model_args()) { + const bool enable_prefill_cp = context.get_parallel_args().cp_size() > 1 && + Platform::uses_model_cp_partition(); + if (layer::cp_conflicts_with_dsa_topk_share(enable_prefill_cp, + dsa_topk_share_plan_)) { + LOG(FATAL) << "Prefill CP is not supported together with GLM5.2 " + "DSA cross-layer top-k sharing. Disable one of them: unset " + "cp_size, or run a model whose DSA top-k plan contains " + "no Shared layers."; + } + } + + protected: + // The relay is reset at layer 0 so its state remains forward-scoped. Decoder + // layers own role interpretation and the attention transfer protocol. + torch::Tensor forward_decoder_layer( + size_t layer_id, + layer::DeepseekV2DecoderLayer& layer, + torch::Tensor& hidden_states, + std::optional& residual, + torch::Tensor& positions, + layer::AttentionMetadata& attn_metadata, + KVCache& kv_cache, + const ModelInputParams& input_params) override { + if (layer_id == 0) { + topk_relay_.reset(); + } + return layer(hidden_states, + residual, + positions, + attn_metadata, + kv_cache, + input_params, + /*input_ids=*/std::nullopt, + &topk_relay_); + } + + private: + static DecoderLayerFactory create_decoder_layer_factory( + const ModelContext& context) { + const layer::DsaTopkSharePlan topk_share_plan(context.get_model_args()); + return + [topk_share_plan](const ModelContext& layer_context, int32_t layer_id) { + return layer::DeepseekV2DecoderLayer( + layer_context, layer_id, topk_share_plan); + }; + } + + layer::DsaTopkSharePlan dsa_topk_share_plan_; + layer::DsaTopkRelay topk_relay_; +}; +TORCH_MODULE(Glm52Model); + +class Glm52ForCausalLMImpl : public LlmForCausalLMImplBase { + public: + explicit Glm52ForCausalLMImpl(const ModelContext& context) + : LlmForCausalLMImplBase(context) {} + + void load_model( + std::unique_ptr loader, + std::string prefix = "model." /*llm model weight prefix*/) override { + LlmForCausalLMImplBase::load_model(std::move(loader), prefix); + model_->verify_loaded_weights(); + } +}; +TORCH_MODULE(Glm52ForCausalLM); + +// register the causal model +REGISTER_CAUSAL_MODEL(glm_moe_dsa, Glm52ForCausalLM); + +} // namespace xllm diff --git a/xllm/models/models.h b/xllm/models/models.h index 4e5e74cf33..a94510eb08 100644 --- a/xllm/models/models.h +++ b/xllm/models/models.h @@ -71,6 +71,7 @@ limitations under the License. #include "llm/deepseek_v3.h" // IWYU pragma: keep #include "llm/deepseek_v32.h" // IWYU pragma: keep #include "llm/glm5.h" // IWYU pragma: keep +#include "llm/glm52.h" // IWYU pragma: keep #include "llm/glm5_mtp.h" // IWYU pragma: keep #include "llm/joyai_llm_flash.h" // IWYU pragma: keep #include "llm/joyai_llm_flash_mtp.h" // IWYU pragma: keep