已合并
【refactor】: 删除dflow中废弃的KvCacheManager #3928
lining23666创建于 7月8日
【refactor】: 删除dflow中废弃的KvCacheManager #3928
已合并
共 7 个文件变更+0-492
| @@ -17,7 +17,6 @@ | |||
| 17 | 17 | ||
| 18 | 18 | ||
| 19 | 19 | ||
| 20 | - | ||
| 21 | 20 | ||
| 22 | 21 | ||
| 23 | 22 | ||
| @@ -1104,10 +1103,6 @@ FsmStatus LlmCommEntity::ReceiveKvCache(const PullKvReqInfo &req_info, const Syn | |||
| 1104 | } | 1103 | } |
| 1105 | UDF_LOG_INFO("enable_paged_attention_ = %d, cache_id = %ld", static_cast<int32_t>(enable_paged_attention_), | 1104 | UDF_LOG_INFO("enable_paged_attention_ = %d, cache_id = %ld", static_cast<int32_t>(enable_paged_attention_), |
| 1106 | req_info.cache_id); | 1105 | req_info.cache_id); |
| 1107 | - if ((!enable_paged_attention_) && (req_info.cache_id <= 0)) { | ||
| 1108 | - const auto &cache = KvCacheManager::GetInstance().GetDecoderKvCache(req_info.model_id); | ||
| 1109 | - (void)KvCacheManager::GetInstance().SaveDecoderKvCache({req_info.req_id, req_info.model_id}, cache); | ||
| 1110 | - } | ||
| 1111 | bool max_time_cost_flag = false; | 1106 | bool max_time_cost_flag = false; |
| 1112 | UpdateTickCost(tick_cost, stat_info_.recv_kv_total_times, stat_info_.recv_kv_min_tick_cost, | 1107 | UpdateTickCost(tick_cost, stat_info_.recv_kv_total_times, stat_info_.recv_kv_min_tick_cost, |
| 1113 | stat_info_.recv_kv_max_tick_cost, stat_info_.recv_kv_total_tick_cost, max_time_cost_flag); | 1108 | stat_info_.recv_kv_max_tick_cost, stat_info_.recv_kv_total_tick_cost, max_time_cost_flag); |
| @@ -1253,22 +1248,6 @@ std::vector<uintptr_t> LlmCommEntity::GetRecvAddrs(const PullKvReqInfo &req_info | |||
| 1253 | } | 1248 | } |
| 1254 | return recv_addrs; | 1249 | return recv_addrs; |
| 1255 | } | 1250 | } |
| 1256 | - if (enable_paged_attention_) { | ||
| 1257 | - auto &kv_blocks_tensors = KvCacheManager::GetInstance().QueryKvBlocksTensors(req_info.model_id); | ||
| 1258 | - for (uint32_t index = 0U; index < transfer_count; ++index) { | ||
| 1259 | - auto kv_index = index / sync_req_info->buffer_count_per_layer; | ||
| 1260 | - auto kv_blocks_addr = reinterpret_cast<uintptr_t>(kv_blocks_tensors[kv_index]->GetTensor()->GetData()); | ||
| 1261 | - recv_addrs.emplace_back(kv_blocks_addr); | ||
| 1262 | - } | ||
| 1263 | - } else { | ||
| 1264 | - // get kv cache: no need check kvCache nullptr | ||
| 1265 | - std::shared_ptr<FlowMsg> kv_cache = KvCacheManager::GetInstance().GetDecoderKvCache(req_info.model_id); | ||
| 1266 | - const auto kv_cache_addr = reinterpret_cast<uintptr_t>(kv_cache->GetTensor()->GetData()); | ||
| 1267 | - for (uint32_t index = 0U; index < transfer_count; ++index) { | ||
| 1268 | - auto kv_index = index / sync_req_info->buffer_count_per_layer; | ||
| 1269 | - recv_addrs.emplace_back(kv_cache_addr + kv_index * req_info.block_len); | ||
| 1270 | - } | ||
| 1271 | - } | ||
| 1272 | return recv_addrs; | 1251 | return recv_addrs; |
| 1273 | } | 1252 | } |
| 1274 | 1253 | ||
| @@ -1476,27 +1455,6 @@ FsmStatus LlmCommEntity::CheckSyncKvMetaInfo(const PullKvReqInfo &req_info, Sync | |||
| 1476 | } | 1455 | } |
| 1477 | return FsmStatus::kFsmSuccess; | 1456 | return FsmStatus::kFsmSuccess; |
| 1478 | } | 1457 | } |
| 1479 | - if (!enable_paged_attention_) { | ||
| 1480 | - std::shared_ptr<FlowMsg> kv_cache = KvCacheManager::GetInstance().GetDecoderKvCache(req_info.model_id); | ||
| 1481 | - const uint64_t kv_cache_size = kv_cache->GetTensor()->GetDataSize(); | ||
| 1482 | - if (resp_info.transfer_count * req_info.block_len > kv_cache_size) { | ||
| 1483 | - UDF_LOG_ERROR( | ||
| 1484 | - "Invalid param, req_id:%lu, transfer_count:%u, block_len:%lu, kv_cache_size:%lu, " | ||
| 1485 | - "model_id:%lu, entity:%s.", | ||
| 1486 | - req_info.req_id, resp_info.transfer_count, req_info.block_len, kv_cache_size, req_info.model_id, | ||
| 1487 | - desc_.c_str()); | ||
| 1488 | - return FsmStatus::kFsmParamInvalid; | ||
| 1489 | - } | ||
| 1490 | - return FsmStatus::kFsmSuccess; | ||
| 1491 | - } | ||
| 1492 | - size_t send_count_per_layer = sync_req_info->buffer_count_per_layer; | ||
| 1493 | - const auto &kv_block_tensors = KvCacheManager::GetInstance().QueryKvBlocksTensors(req_info.model_id); | ||
| 1494 | - if (resp_info.transfer_count != (kv_block_tensors.size() * send_count_per_layer)) { | ||
| 1495 | - UDF_LOG_ERROR("Invalid param, req_id:%lu, transfer_count:%u, send_count_per_layer:%u, kvSize:%zu, entity:%s.", | ||
| 1496 | - req_info.req_id, resp_info.transfer_count, send_count_per_layer, kv_block_tensors.size(), | ||
| 1497 | - desc_.c_str()); | ||
| 1498 | - return FsmStatus::kFsmParamInvalid; | ||
| 1499 | - } | ||
| 1500 | return FsmStatus::kFsmSuccess; | 1458 | return FsmStatus::kFsmSuccess; |
| 1501 | } | 1459 | } |
| 1502 | 1460 | ||
| @@ -12,7 +12,6 @@ | |||
| 12 | 12 | ||
| 13 | 13 | ||
| 14 | 14 | ||
| 15 | - | ||
| 16 | 15 | ||
| 17 | 16 | ||
| 18 | 17 | ||
| @@ -167,33 +166,6 @@ FsmStatus SendState::QueryPromptKvCache(const SyncKvReqInfo *req_info, std::vect | |||
| 167 | } | 166 | } |
| 168 | return ret; | 167 | return ret; |
| 169 | } | 168 | } |
| 170 | - kv_tensors = KvCacheManager::GetInstance().QueryPromptKvCache({req_info->req_id, req_info->model_id}); | ||
| 171 | - if (!kv_tensors.empty()) { | ||
| 172 | - if (CheckKvCacheManagerReq(kv_tensors, req_info) != FsmStatus::kFsmSuccess) { | ||
| 173 | - return FsmStatus::kFsmParamInvalid; | ||
| 174 | - } | ||
| 175 | - } | ||
| 176 | - return FsmStatus::kFsmSuccess; | ||
| 177 | -} | ||
| 178 | - | ||
| 179 | -FsmStatus SendState::CheckKvCacheManagerReq(const std::vector<KvTensor> &kv_tensors, const SyncKvReqInfo *req_info) { | ||
| 180 | - uint64_t real_kv_len = kv_tensors.front().data_size; | ||
| 181 | - uint64_t req_kv_len = 0UL; | ||
| 182 | - for (uint32_t i = 0U; i < req_info->buffer_count_per_layer; ++i) { | ||
| 183 | - if (req_info->transfer_infos[i].buffer_info.block_index != UINT64_MAX) { | ||
| 184 | - UDF_LOG_ERROR("Invalid buffer info, i:%u, block_index:%lu", i, | ||
| 185 | - req_info->transfer_infos[i].buffer_info.block_index); | ||
| 186 | - return FsmStatus::kFsmParamInvalid; | ||
| 187 | - } | ||
| 188 | - if ((req_info->buffer_count_per_layer > 1) && (i < req_info->buffer_count_per_layer - 1U)) { | ||
| 189 | - req_kv_len += req_info->transfer_infos[i].buffer_info.buffer_len; | ||
| 190 | - } | ||
| 191 | - } | ||
| 192 | - if (req_kv_len >= real_kv_len) { | ||
| 193 | - UDF_LOG_ERROR("Invalid param, req_kv_len:%lu, real_kv_len:%lu, req_id:%lu, prefix_id:%lu, model_id:%lu.", | ||
| 194 | - req_kv_len, real_kv_len, req_info->req_id, req_info->prefix_id, req_info->model_id); | ||
| 195 | - return FsmStatus::kFsmParamInvalid; | ||
| 196 | - } | ||
| 197 | return FsmStatus::kFsmSuccess; | 169 | return FsmStatus::kFsmSuccess; |
| 198 | } | 170 | } |
| 199 | 171 | ||
| @@ -338,8 +310,6 @@ void SendState::ReleaseKvCacheForPrompt(const LlmCommEntity &entity) { | |||
| 338 | CacheManager::GetInstance().RemoveCacheIndex(std::make_pair(entity.GetCurReqId(), entity.GetCurModelId()), false, | 310 | CacheManager::GetInstance().RemoveCacheIndex(std::make_pair(entity.GetCurReqId(), entity.GetCurModelId()), false, |
| 339 | tensor_num_and_indices.first, tensor_num_and_indices.second); | 311 | tensor_num_and_indices.first, tensor_num_and_indices.second); |
| 340 | } | 312 | } |
| 341 | - } else { | ||
| 342 | - (void)KvCacheManager::GetInstance().ReleaseKvCacheForPrompt({entity.GetCurReqId(), entity.GetCurModelId()}); | ||
| 343 | } | 313 | } |
| 344 | } | 314 | } |
| 345 | 315 | ||
| @@ -29,7 +29,6 @@ class SendState : public BaseState { | |||
| 29 | SendState &operator=(const SendState &&) = delete; | 29 | SendState &operator=(const SendState &&) = delete; |
| 30 | 30 | ||
| 31 | private: | 31 | private: |
| 32 | - static FsmStatus CheckKvCacheManagerReq(const std::vector<KvTensor> &kv_tensors, const SyncKvReqInfo *req_info); | ||
| 33 | static FsmStatus CheckNotBlocksReq(const SyncKvReqInfo *req_info); | 32 | static FsmStatus CheckNotBlocksReq(const SyncKvReqInfo *req_info); |
| 34 | static FsmStatus GenerateSyncKvMetaInfo(LlmCommEntity &entity); | 33 | static FsmStatus GenerateSyncKvMetaInfo(LlmCommEntity &entity); |
| 35 | static FsmStatus SendSyncKvMetaAsync(LlmCommEntity &entity); | 34 | static FsmStatus SendSyncKvMetaAsync(LlmCommEntity &entity); |
| @@ -1,304 +0,0 @@ | |||
| 1 | -/** | ||
| 2 | - * Copyright (c) 2025 Huawei Technologies Co., Ltd. | ||
| 3 | - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | - * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | - * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | - * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | - */ | ||
| 10 | - | ||
| 11 | - | ||
| 12 | - | ||
| 13 | - | ||
| 14 | - | ||
| 15 | - | ||
| 16 | - | ||
| 17 | -namespace FlowFunc { | ||
| 18 | -KvCacheManager &KvCacheManager::GetInstance() { | ||
| 19 | - static KvCacheManager manager; | ||
| 20 | - return manager; | ||
| 21 | -} | ||
| 22 | - | ||
| 23 | -KvCacheManager::~KvCacheManager() { | ||
| 24 | - ClearDecoderKvCache(); | ||
| 25 | - ClearPromptKvCache(); | ||
| 26 | - decoder_req_to_kv_cache_.clear(); | ||
| 27 | - kv_blocks_tensors_.clear(); | ||
| 28 | -} | ||
| 29 | - | ||
| 30 | -std::shared_ptr<FlowMsg> KvCacheManager::QueryDecoderKvCache(const ReqKey &req_key) { | ||
| 31 | - std::unique_lock<std::mutex> lock(mutex_); | ||
| 32 | - const auto iter = decoder_req_to_kv_cache_.find(req_key); | ||
| 33 | - if (iter == decoder_req_to_kv_cache_.end()) { | ||
| 34 | - uint64_t cached_req_id = | ||
| 35 | - decoder_req_to_kv_cache_.empty() ? UINT64_MAX : decoder_req_to_kv_cache_.begin()->first.req_id; | ||
| 36 | - UDF_LOG_ERROR("Cannot find kv cache for req_id:%lu, current cached req_id:%lu.", req_key.req_id, cached_req_id); | ||
| 37 | - return nullptr; | ||
| 38 | - } | ||
| 39 | - return iter->second; | ||
| 40 | -} | ||
| 41 | - | ||
| 42 | -FsmStatus KvCacheManager::SaveDecoderKvCache(const ReqKey &req_key, const std::shared_ptr<FlowMsg> &kv_cache) { | ||
| 43 | - std::unique_lock<std::mutex> lock(mutex_); | ||
| 44 | - const auto iter = decoder_req_to_kv_cache_.find(req_key); | ||
| 45 | - if (iter != decoder_req_to_kv_cache_.end()) { | ||
| 46 | - UDF_LOG_ERROR("Exist kv cache in decoder, req_id:%lu, model_id:%lu.", req_key.req_id, req_key.model_id); | ||
| 47 | - } | ||
| 48 | - decoder_req_to_kv_cache_[req_key] = kv_cache; | ||
| 49 | - UDF_LOG_INFO("Success to save kv cache for decoder, req_id:%lu, model_id:%lu.", req_key.req_id, req_key.model_id); | ||
| 50 | - return FsmStatus::kFsmSuccess; | ||
| 51 | -} | ||
| 52 | - | ||
| 53 | -FsmStatus KvCacheManager::ReleaseDecoderKvCache(const ReqKey &req_key) { | ||
| 54 | - std::unique_lock<std::mutex> lock(mutex_); | ||
| 55 | - const auto iter = decoder_req_to_kv_cache_.find(req_key); | ||
| 56 | - if (iter == decoder_req_to_kv_cache_.end()) { | ||
| 57 | - UDF_LOG_INFO("Not exist kv cache in decoder, req_id:%lu, model_id:%lu.", req_key.req_id, req_key.model_id); | ||
| 58 | - return FsmStatus::kFsmSuccess; | ||
| 59 | - } | ||
| 60 | - (void)decoder_req_to_kv_cache_.erase(iter); | ||
| 61 | - return FsmStatus::kFsmSuccess; | ||
| 62 | -} | ||
| 63 | - | ||
| 64 | -FsmStatus KvCacheManager::ReleaseDecoderKvCaches() { | ||
| 65 | - std::unique_lock<std::mutex> lock(mutex_); | ||
| 66 | - (void)decoder_req_to_kv_cache_.clear(); | ||
| 67 | - return FsmStatus::kFsmSuccess; | ||
| 68 | -} | ||
| 69 | - | ||
| 70 | -void KvCacheManager::SaveKvCacheForPrompt(const ReqKey &req_key, uint64_t cache_id, | ||
| 71 | - std::pair<int64_t, int64_t> offset_and_size) { | ||
| 72 | - std::lock_guard<std::mutex> lock(mutex_); | ||
| 73 | - UDF_LOG_INFO("KV Cache save for prompt, req_id = %lu, model_id = %lu, cache_id = %lu, offset = %ld, size = %ld", | ||
| 74 | - req_key.req_id, req_key.model_id, cache_id, offset_and_size.first, offset_and_size.second); | ||
| 75 | - auto &cache_entry = cache_id_to_cache_entry_[cache_id]; | ||
| 76 | - cache_entry.id_to_offset_and_size[req_key.req_id] = std::move(offset_and_size); | ||
| 77 | - req_id_to_cache_id_[req_key] = cache_id; | ||
| 78 | -} | ||
| 79 | - | ||
| 80 | -void KvCacheManager::SaveKvCacheForPromptPrefix(const PrefixReqKey &prefix_req_key, uint64_t cache_id, | ||
| 81 | - std::pair<int64_t, int64_t> offset_and_size) { | ||
| 82 | - UDF_LOG_INFO("KV Cache save for prefix, prefix_id = %lu, cache_id = %lu, offset = %ld, size = %ld", | ||
| 83 | - prefix_req_key.prefix_id, cache_id, offset_and_size.first, offset_and_size.second); | ||
| 84 | - std::unique_lock<std::mutex> lock(mutex_); | ||
| 85 | - const auto iter = prefix_id_to_cache_id_.find(prefix_req_key); | ||
| 86 | - if (iter != prefix_id_to_cache_id_.end()) { | ||
| 87 | - UDF_LOG_INFO("Exist kv cache in prompt, prefix_id:%lu.", prefix_req_key.prefix_id); | ||
| 88 | - } | ||
| 89 | - auto &cache_entry = cache_id_to_cache_entry_[cache_id]; | ||
| 90 | - cache_entry.id_to_offset_and_size[prefix_req_key.prefix_id] = std::move(offset_and_size); | ||
| 91 | - prefix_id_to_cache_id_[prefix_req_key] = cache_id; | ||
| 92 | -} | ||
| 93 | - | ||
| 94 | -FsmStatus KvCacheManager::ReleaseKvCacheForPromptByKey(const ReqKey &req_key) { | ||
| 95 | - UDF_LOG_INFO("Release KV Cache for prompt, req_id=%lu, model_id=%lu", req_key.req_id, req_key.model_id); | ||
| 96 | - const auto iter = req_id_to_cache_id_.find(req_key); | ||
| 97 | - if (iter == req_id_to_cache_id_.end()) { | ||
| 98 | - UDF_LOG_INFO("Not exist kv cache in prompt, req_id:%lu.", req_key.req_id); | ||
| 99 | - return FsmStatus::kFsmSuccess; | ||
| 100 | - } | ||
| 101 | - const auto cache_id = iter->second; | ||
| 102 | - if (cache_id_to_cache_entry_.count(cache_id) > 0UL) { | ||
| 103 | - auto &cache_entry = cache_id_to_cache_entry_[cache_id]; | ||
| 104 | - if (cache_entry.id_to_offset_and_size.count(req_key.req_id) > 0UL) { | ||
| 105 | - (void)cache_entry.id_to_offset_and_size.erase(req_key.req_id); | ||
| 106 | - } | ||
| 107 | - if (cache_entry.id_to_offset_and_size.empty()) { | ||
| 108 | - (void)cache_id_to_cache_entry_.erase(cache_id); | ||
| 109 | - UDF_LOG_INFO("KV Cache removed, cache id = %lu", cache_id); | ||
| 110 | - } | ||
| 111 | - } | ||
| 112 | - (void)req_id_to_cache_id_.erase(iter); | ||
| 113 | - return FsmStatus::kFsmSuccess; | ||
| 114 | -} | ||
| 115 | - | ||
| 116 | -FsmStatus KvCacheManager::ReleaseKvCacheForPrompt(const ReqKey &req_key) { | ||
| 117 | - std::unique_lock<std::mutex> lock(mutex_); | ||
| 118 | - return ReleaseKvCacheForPromptByKey(req_key); | ||
| 119 | -} | ||
| 120 | - | ||
| 121 | -FsmStatus KvCacheManager::ReleaseKvCacheForPrompt(uint64_t req_id) { | ||
| 122 | - std::unique_lock<std::mutex> lock(mutex_); | ||
| 123 | - UDF_LOG_INFO("Release KV Cache for prompt, req_id=%lu", req_id); | ||
| 124 | - std::vector<ReqKey> need_delete_req_keys; | ||
| 125 | - for (const auto &it : req_id_to_cache_id_) { | ||
| 126 | - if (it.first.req_id == req_id) { | ||
| 127 | - need_delete_req_keys.emplace_back(it.first); | ||
| 128 | - } | ||
| 129 | - } | ||
| 130 | - for (const auto &req_key : need_delete_req_keys) { | ||
| 131 | - (void)ReleaseKvCacheForPromptByKey(req_key); | ||
| 132 | - } | ||
| 133 | - return FsmStatus::kFsmSuccess; | ||
| 134 | -} | ||
| 135 | - | ||
| 136 | -FsmStatus KvCacheManager::ReleaseKvCacheForPrefix(const PrefixReqKey &prefix_req_key) { | ||
| 137 | - std::unique_lock<std::mutex> lock(mutex_); | ||
| 138 | - UDF_LOG_INFO("Release KV Cache for prefix, prefix_id=%lu, model_id=%lu", prefix_req_key.prefix_id, | ||
| 139 | - prefix_req_key.model_id); | ||
| 140 | - const auto iter = prefix_id_to_cache_id_.find(prefix_req_key); | ||
| 141 | - if (iter == prefix_id_to_cache_id_.end()) { | ||
| 142 | - UDF_LOG_INFO("Not exist kv cache in prompt, prefix_id:%lu.", prefix_req_key.prefix_id); | ||
| 143 | - return FsmStatus::kFsmSuccess; | ||
| 144 | - } | ||
| 145 | - const auto cache_id = iter->second; | ||
| 146 | - auto &cache_entry = cache_id_to_cache_entry_[cache_id]; | ||
| 147 | - (void)cache_entry.id_to_offset_and_size.erase(prefix_req_key.prefix_id); | ||
| 148 | - if (cache_entry.id_to_offset_and_size.empty()) { | ||
| 149 | - (void)cache_id_to_cache_entry_.erase(cache_id); | ||
| 150 | - UDF_LOG_INFO("KV Cache removed, cache id = %lu", cache_id); | ||
| 151 | - } | ||
| 152 | - (void)prefix_id_to_cache_id_.erase(iter); | ||
| 153 | - return FsmStatus::kFsmSuccess; | ||
| 154 | -} | ||
| 155 | - | ||
| 156 | -void KvCacheManager::SetDecoderKvCache(const std::vector<std::shared_ptr<FlowMsg>> &kv_caches) { | ||
| 157 | - decoder_kv_caches_ = kv_caches; | ||
| 158 | -} | ||
| 159 | - | ||
| 160 | -std::shared_ptr<FlowMsg> KvCacheManager::GetDecoderKvCache(uint64_t model_id) const { | ||
| 161 | - return decoder_kv_caches_[model_id]; | ||
| 162 | -} | ||
| 163 | - | ||
| 164 | -void KvCacheManager::ClearDecoderKvCache() { | ||
| 165 | - decoder_kv_caches_.clear(); | ||
| 166 | -} | ||
| 167 | - | ||
| 168 | -uint64_t KvCacheManager::GetPromptKvCacheCount() { | ||
| 169 | - std::unique_lock<std::mutex> lock(mutex_); | ||
| 170 | - uint64_t count = 0; | ||
| 171 | - std::set<uint64_t> prompt_cache_ids; | ||
| 172 | - for (const auto &req_id_and_cache_id : req_id_to_cache_id_) { | ||
| 173 | - prompt_cache_ids.insert(req_id_and_cache_id.second); | ||
| 174 | - } | ||
| 175 | - for (const auto cache_id : prompt_cache_ids) { | ||
| 176 | - const auto batch_num = cache_id_to_cache_entry_[cache_id].batch_num; | ||
| 177 | - UDF_LOG_INFO("cache id = %lu, batch_num = %zu", cache_id, batch_num); | ||
| 178 | - count += batch_num; | ||
| 179 | - } | ||
| 180 | - return count; | ||
| 181 | -} | ||
| 182 | - | ||
| 183 | -uint64_t KvCacheManager::GetDecoderKvCacheCount(const ReqKey &req_key) { | ||
| 184 | - std::unique_lock<std::mutex> lock(mutex_); | ||
| 185 | - return decoder_req_to_kv_cache_.count(req_key); | ||
| 186 | -} | ||
| 187 | - | ||
| 188 | -std::vector<KvTensor> KvCacheManager::QueryKvTensor(uint64_t cache_id, uint64_t tensor_id) { | ||
| 189 | - auto &cache_entry = cache_id_to_cache_entry_[cache_id]; | ||
| 190 | - const auto &offset_and_size = cache_entry.id_to_offset_and_size[tensor_id]; | ||
| 191 | - std::vector<KvTensor> kv_tensors; | ||
| 192 | - if (!cache_entry.kv_tensors.empty()) { | ||
| 193 | - kv_tensors.reserve(cache_entry.kv_tensors.size() * cache_entry.kv_tensors[0]->GetTensorList().size()); | ||
| 194 | - } | ||
| 195 | - for (const auto &kv_tensor_buffer : cache_entry.kv_tensors) { | ||
| 196 | - const auto tensor_list = kv_tensor_buffer->GetTensorList(); | ||
| 197 | - for (const auto &tensor : tensor_list) { | ||
| 198 | - KvTensor kv_tensor; | ||
| 199 | - kv_tensor.tensor_buffer = kv_tensor_buffer; | ||
| 200 | - kv_tensor.data_addr = static_cast<uint8_t *>(tensor->GetData()) + offset_and_size.first; | ||
| 201 | - kv_tensor.data_size = offset_and_size.second; | ||
| 202 | - kv_tensor.block_len = kv_tensor.data_size; | ||
| 203 | - kv_tensors.emplace_back(std::move(kv_tensor)); | ||
| 204 | - } | ||
| 205 | - } | ||
| 206 | - return kv_tensors; | ||
| 207 | -} | ||
| 208 | - | ||
| 209 | -std::vector<KvTensor> KvCacheManager::QueryPromptKvCache(const ReqKey &req_key) { | ||
| 210 | - static std::vector<KvTensor> empty_vector = {}; | ||
| 211 | - std::unique_lock<std::mutex> lock(mutex_); | ||
| 212 | - const auto iter = req_id_to_cache_id_.find(req_key); | ||
| 213 | - if (iter == req_id_to_cache_id_.end()) { | ||
| 214 | - UDF_LOG_INFO("Not exist kv cache in prompt, req_id:%lu.", req_key.req_id); | ||
| 215 | - return empty_vector; | ||
| 216 | - } | ||
| 217 | - const auto cache_id = iter->second; | ||
| 218 | - return QueryKvTensor(cache_id, req_key.req_id); | ||
| 219 | -} | ||
| 220 | - | ||
| 221 | -std::vector<KvTensor> KvCacheManager::QueryPromptPrefixKvCache(const PrefixReqKey &prefix_req_key) { | ||
| 222 | - static std::vector<KvTensor> empty_vector = {}; | ||
| 223 | - std::unique_lock<std::mutex> lock(mutex_); | ||
| 224 | - const auto iter = prefix_id_to_cache_id_.find(prefix_req_key); | ||
| 225 | - if (iter == prefix_id_to_cache_id_.end()) { | ||
| 226 | - UDF_LOG_INFO("Not exist kv cache in prompt, req_id:%lu.", prefix_req_key.prefix_id); | ||
| 227 | - return empty_vector; | ||
| 228 | - } | ||
| 229 | - const auto cache_id = iter->second; | ||
| 230 | - return QueryKvTensor(cache_id, prefix_req_key.prefix_id); | ||
| 231 | -} | ||
| 232 | - | ||
| 233 | -void KvCacheManager::InitBlockTensors(uint64_t model_id) { | ||
| 234 | - kv_blocks_tensors_.resize(model_id, std::vector<std::shared_ptr<FlowMsg>>()); | ||
| 235 | -} | ||
| 236 | - | ||
| 237 | -void KvCacheManager::SaveKvBlocksTensor(uint64_t model_id, std::shared_ptr<FlowMsg> &kv_blocks_tensor) { | ||
| 238 | - kv_blocks_tensors_[model_id].emplace_back(kv_blocks_tensor); | ||
| 239 | -} | ||
| 240 | - | ||
| 241 | -std::vector<std::shared_ptr<FlowMsg>> &KvCacheManager::QueryKvBlocksTensors(uint64_t model_id) { | ||
| 242 | - return kv_blocks_tensors_[model_id]; | ||
| 243 | -} | ||
| 244 | - | ||
| 245 | -void KvCacheManager::ClearKvBlocksTensors() { | ||
| 246 | - kv_blocks_tensors_.clear(); | ||
| 247 | -} | ||
| 248 | - | ||
| 249 | -uint64_t KvCacheManager::AddKvCache(size_t batch_num, std::vector<std::shared_ptr<FlowMsg>> kv_tensors) { | ||
| 250 | - std::unique_lock<std::mutex> lock(mutex_); | ||
| 251 | - const auto cache_id = cache_id_gen_++; | ||
| 252 | - CacheEntry cache_entry; | ||
| 253 | - cache_entry.batch_num = batch_num; | ||
| 254 | - cache_entry.kv_tensors = std::move(kv_tensors); | ||
| 255 | - cache_id_to_cache_entry_.emplace(cache_id, std::move(cache_entry)); | ||
| 256 | - UDF_LOG_INFO("KV Cache added, cache id = %lu, batch_num = %zu", cache_id, batch_num); | ||
| 257 | - return cache_id; | ||
| 258 | -} | ||
| 259 | - | ||
| 260 | -FsmStatus KvCacheManager::ShrinkKvCacheForPrefix(const PrefixReqKey &prefix_req_key, | ||
| 261 | - const std::shared_ptr<MetaRunContext> &run_context) { | ||
| 262 | - UDF_LOG_INFO("Shrink start, prefix_id:%lu.", prefix_req_key.prefix_id); | ||
| 263 | - std::vector<std::shared_ptr<FlowMsg>> shrinked_kv_tensors; | ||
| 264 | - int64_t tensor_size = -1; | ||
| 265 | - { | ||
| 266 | - std::unique_lock<std::mutex> lock(mutex_); | ||
| 267 | - const auto &cache_id = prefix_id_to_cache_id_[prefix_req_key]; | ||
| 268 | - auto &cache_entry = cache_id_to_cache_entry_[cache_id]; | ||
| 269 | - const auto &offset_and_size = cache_entry.id_to_offset_and_size[prefix_req_key.prefix_id]; | ||
| 270 | - tensor_size = offset_and_size.second; | ||
| 271 | - shrinked_kv_tensors.reserve(cache_entry.kv_tensors.size()); | ||
| 272 | - for (auto &kv_tensor : cache_entry.kv_tensors) { | ||
| 273 | - auto new_kv_tensor = run_context->AllocTensorMsg({tensor_size}, TensorDataType::DT_UINT8); | ||
| 274 | - if (new_kv_tensor == nullptr) { | ||
| 275 | - UDF_LOG_ERROR("Shink kv cache failed. prefix_id = %lu, Currently cache size:%lu", prefix_req_key.prefix_id, | ||
| 276 | - KvCacheManager::GetInstance().GetPromptKvCacheCount()); | ||
| 277 | - return FsmStatus::kFsmOutOfMemory; | ||
| 278 | - } | ||
| 279 | - auto rt_ret = halSdmaCopy(reinterpret_cast<uintptr_t>(new_kv_tensor->GetTensor()->GetData()), tensor_size, | ||
| 280 | - reinterpret_cast<uintptr_t>(static_cast<uint8_t *>(kv_tensor->GetTensor()->GetData()) + | ||
| 281 | - offset_and_size.first), | ||
| 282 | - tensor_size); | ||
| 283 | - if (rt_ret != DRV_ERROR_NONE) { | ||
| 284 | - UDF_LOG_ERROR("Fail to call halSdmaCopy, kv cache size:%ld, ret:%d.", tensor_size, rt_ret); | ||
| 285 | - return FsmStatus::kFsmFailed; | ||
| 286 | - } | ||
| 287 | - shrinked_kv_tensors.emplace_back(std::move(new_kv_tensor)); | ||
| 288 | - kv_tensor.reset(); | ||
| 289 | - } | ||
| 290 | - } | ||
| 291 | - UDF_LOG_INFO("copy tensors ended, prefix_id:%lu.", prefix_req_key.prefix_id); | ||
| 292 | - auto new_cache_id = AddKvCache(1U, shrinked_kv_tensors); | ||
| 293 | - (void)ReleaseKvCacheForPrefix(prefix_req_key); | ||
| 294 | - SaveKvCacheForPromptPrefix(prefix_req_key, new_cache_id, std::pair<int64_t, int64_t>(0L, tensor_size)); | ||
| 295 | - return FsmStatus::kFsmSuccess; | ||
| 296 | -} | ||
| 297 | - | ||
| 298 | -void KvCacheManager::ClearPromptKvCache() { | ||
| 299 | - std::unique_lock<std::mutex> lock(mutex_); | ||
| 300 | - cache_id_to_cache_entry_.clear(); | ||
| 301 | - req_id_to_cache_id_.clear(); | ||
| 302 | - prefix_id_to_cache_id_.clear(); | ||
| 303 | -} | ||
| 304 | -} // namespace FlowFunc | ||
| @@ -1,113 +0,0 @@ | |||
| 1 | -/** | ||
| 2 | - * Copyright (c) 2025 Huawei Technologies Co., Ltd. | ||
| 3 | - * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | - * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | - * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | - * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | - * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | - * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | - */ | ||
| 10 | - | ||
| 11 | - | ||
| 12 | - | ||
| 13 | - | ||
| 14 | - | ||
| 15 | - | ||
| 16 | - | ||
| 17 | - | ||
| 18 | - | ||
| 19 | - | ||
| 20 | - | ||
| 21 | -namespace FlowFunc { | ||
| 22 | -constexpr uint64_t kHashShiftLeft = 32UL; | ||
| 23 | - | ||
| 24 | -class KvCacheManager { | ||
| 25 | - public: | ||
| 26 | - struct ReqKey { | ||
| 27 | - uint64_t req_id; | ||
| 28 | - uint64_t model_id; | ||
| 29 | - bool operator==(const ReqKey &req_key) const { | ||
| 30 | - return ((req_id == req_key.req_id) && (model_id == req_key.model_id)); | ||
| 31 | - } | ||
| 32 | - }; | ||
| 33 | - struct ReqKeyHash { | ||
| 34 | - std::size_t operator()(const ReqKey &req_key) const { | ||
| 35 | - std::hash<uint64_t> hash_fn; | ||
| 36 | - return hash_fn(req_key.req_id) ^ hash_fn(req_key.model_id << kHashShiftLeft); | ||
| 37 | - } | ||
| 38 | - }; | ||
| 39 | - struct PrefixReqKey { | ||
| 40 | - uint64_t prefix_id; | ||
| 41 | - uint64_t model_id; | ||
| 42 | - bool operator==(const PrefixReqKey &req_key) const { | ||
| 43 | - return ((prefix_id == req_key.prefix_id) && (model_id == req_key.model_id)); | ||
| 44 | - } | ||
| 45 | - }; | ||
| 46 | - struct PrefixReqKeyHash { | ||
| 47 | - std::size_t operator()(const PrefixReqKey &req_key) const { | ||
| 48 | - std::hash<uint64_t> hash_fn; | ||
| 49 | - return hash_fn(req_key.prefix_id) ^ hash_fn(req_key.model_id << kHashShiftLeft); | ||
| 50 | - } | ||
| 51 | - }; | ||
| 52 | - | ||
| 53 | - static KvCacheManager &GetInstance(); | ||
| 54 | - virtual ~KvCacheManager(); | ||
| 55 | - std::shared_ptr<FlowMsg> QueryDecoderKvCache(const ReqKey &req_key); | ||
| 56 | - FsmStatus SaveDecoderKvCache(const ReqKey &req_key, const std::shared_ptr<FlowMsg> &kv_cache); | ||
| 57 | - FsmStatus ReleaseDecoderKvCache(const ReqKey &req_key); | ||
| 58 | - FsmStatus ReleaseDecoderKvCaches(); | ||
| 59 | - void SetDecoderKvCache(const std::vector<std::shared_ptr<FlowMsg>> &kv_caches); | ||
| 60 | - std::shared_ptr<FlowMsg> GetDecoderKvCache(uint64_t model_id) const; | ||
| 61 | - void ClearDecoderKvCache(); | ||
| 62 | - uint64_t GetPromptKvCacheCount(); | ||
| 63 | - uint64_t GetDecoderKvCacheCount(const ReqKey &req_key); | ||
| 64 | - std::vector<KvTensor> QueryPromptKvCache(const ReqKey &req_key); | ||
| 65 | - std::vector<KvTensor> QueryPromptPrefixKvCache(const PrefixReqKey &prefix_req_key); | ||
| 66 | - FsmStatus ShrinkKvCacheForPrefix(const PrefixReqKey &prefix_req_key, | ||
| 67 | - const std::shared_ptr<MetaRunContext> &run_context); | ||
| 68 | - FsmStatus ReleaseKvCacheForPrompt(const ReqKey &req_key); | ||
| 69 | - FsmStatus ReleaseKvCacheForPrompt(uint64_t req_id); | ||
| 70 | - FsmStatus ReleaseKvCacheForPrefix(const PrefixReqKey &prefix_req_key); | ||
| 71 | - uint64_t AddKvCache(size_t batch_num, std::vector<std::shared_ptr<FlowMsg>> kv_tensors); | ||
| 72 | - void SaveKvCacheForPrompt(const ReqKey &req_key, uint64_t cache_id, std::pair<int64_t, int64_t> offset_and_size); | ||
| 73 | - void SaveKvCacheForPromptPrefix(const PrefixReqKey &prefix_req_key, uint64_t cache_id, | ||
| 74 | - std::pair<int64_t, int64_t> offset_and_size); | ||
| 75 | - void InitBlockTensors(uint64_t model_id); | ||
| 76 | - void SaveKvBlocksTensor(uint64_t model_id, std::shared_ptr<FlowMsg> &kv_blocks_tensor); | ||
| 77 | - std::vector<std::shared_ptr<FlowMsg>> &QueryKvBlocksTensors(uint64_t model_id); | ||
| 78 | - void ClearKvBlocksTensors(); | ||
| 79 | - void ClearPromptKvCache(); | ||
| 80 | - KvCacheManager(const KvCacheManager &) = delete; | ||
| 81 | - KvCacheManager(const KvCacheManager &&) = delete; | ||
| 82 | - KvCacheManager &operator=(const KvCacheManager &) = delete; | ||
| 83 | - KvCacheManager &operator=(const KvCacheManager &&) = delete; | ||
| 84 | - | ||
| 85 | - private: | ||
| 86 | - KvCacheManager() = default; | ||
| 87 | - std::vector<KvTensor> QueryKvTensor(uint64_t cache_id, uint64_t tensor_id); | ||
| 88 | - FsmStatus ReleaseKvCacheForPromptByKey(const ReqKey &req_key); | ||
| 89 | - | ||
| 90 | - private: | ||
| 91 | - struct CacheEntry { | ||
| 92 | - size_t batch_num; | ||
| 93 | - std::vector<std::shared_ptr<FlowMsg>> kv_tensors; | ||
| 94 | - std::unordered_map<uint64_t, std::pair<int64_t, int64_t>> id_to_offset_and_size; | ||
| 95 | - }; | ||
| 96 | - | ||
| 97 | - // only one reserved kv cache for decoder | ||
| 98 | - std::vector<std::shared_ptr<FlowMsg>> decoder_kv_caches_; | ||
| 99 | - // mutex for decoder and prompt kv map | ||
| 100 | - std::mutex mutex_; | ||
| 101 | - // synced kv cache for decoder | ||
| 102 | - std::unordered_map<ReqKey, std::shared_ptr<FlowMsg>, ReqKeyHash> decoder_req_to_kv_cache_; | ||
| 103 | - // kv physical blocks tensor list, count = layerNum * 2 | ||
| 104 | - std::vector<std::vector<std::shared_ptr<FlowMsg>>> kv_blocks_tensors_; | ||
| 105 | - | ||
| 106 | - uint64_t cache_id_gen_{0}; | ||
| 107 | - std::unordered_map<uint64_t, CacheEntry> cache_id_to_cache_entry_; | ||
| 108 | - std::unordered_map<ReqKey, uint64_t, ReqKeyHash> req_id_to_cache_id_; | ||
| 109 | - std::unordered_map<PrefixReqKey, uint64_t, PrefixReqKeyHash> prefix_id_to_cache_id_; | ||
| 110 | -}; | ||
| 111 | -} // namespace FlowFunc | ||
| 112 | - | ||
| 113 | - | ||
| @@ -30,7 +30,6 @@ | |||
| 30 | 30 | ||
| 31 | 31 | ||
| 32 | 32 | ||
| 33 | - | ||
| 34 | 33 | ||
| 35 | 34 | ||
| 36 | namespace FlowFunc { | 35 | namespace FlowFunc { |
| @@ -31,7 +31,6 @@ | |||
| 31 | 31 | ||
| 32 | 32 | ||
| 33 | 33 | ||
| 34 | - | ||
| 35 | 34 | ||
| 36 | 35 | ||
| 37 | namespace FlowFunc { | 36 | namespace FlowFunc { |