异步不分层-dev分支-MindIE-LLM-AtomGit

Mdocs/zh/user_guide/feature/kv_cache_pool.md+1-1

	@@ -98,7 +98,7 @@
98	},	98	},
99	```	99	```
100		100
101	-3. 拉起池化后端对应的中心化服务Master Service，具体安装和拉起命令，请参考[KV Cache池化使用指导](../../../../mindie_llm/text_generator/mempool/README.md)。	101	+3. 拉起池化后端对应的中心化服务Master Service，具体安装和拉起命令，请参考[KV Cache池化使用指导](mempool.md)。
102	4. 启动服务。	102	4. 启动服务。
103		103
104	```bash	104	```bash

Rmindie_llm/text_generator/mempool/README.md→docs/zh/user_guide/feature/mempool.md+15-1

	@@ -11,13 +11,18 @@ KV Cache池化特性依赖于Prefix Cache特性。此外，通过在MindIE的con
11	```json	11	```json
12	"kvPoolConfig" : {"backend":"", "configPath":""}	12	"kvPoolConfig" : {"backend":"", "configPath":""}
13	```	13	```
		14	+或者
		15	+```
		16	+"kvPoolConfig" : {"backend":"", "configPath":"", "asyncWrite": true}
		17	+```
14		18
15	配置说明：	19	配置说明：
16		20
17	- `backend`：指定使用的池化后端。	21	- `backend`：指定使用的池化后端。
18	- `configPath`：池化后端所需要的配置文件路径。	22	- `configPath`：池化后端所需要的配置文件路径。
		23	+- `asyncWrite`：是否开启池化异步写，值为true或false。若不设该字段，则默认为false，即同步写。
19		24
20	-在开启Prefix Cache特性的前提下，当配置了上述两个字段后，则开启KV Cache池化特性。不同的池化后端需要自行安装。	25	+在开启Prefix Cache特性的前提下，当配置了上述字段后，则开启KV Cache池化特性。不同的池化后端需要自行安装。
21		26
22	## 已支持的池化后端	27	## 已支持的池化后端
23		28
	@@ -100,6 +105,15 @@ export ASCEND_BUFFER_POOL=4:8
100		105
101	</details>	106	</details>
102		107
		108	+## 限制与约束
		109	+
		110	+异步写特性当前仅支持Qwen稠密模型（非MOE）和DeepSeek模型（V3/V3.1/R1）。当异步写特性开启时，支持与以下特性叠加组合，叠加其他特性或模型可能导致不可预期的报错：
		111	+- Qwen稠密：异步调度、Prefix Cache、Function Call、思考解析、Yarn
		112	+- DeepSeek V3/V3.1/R1：异步推理、Prefix Cache、Context Parallel、Sequence Parallel
		113	+
		114	+已确定不支持与以下特性叠加：
		115	+- SplitFuse、Micro Batch、Multi-Lora
		116	+
103	## 声明	117	## 声明
104		118
105	- 本代码仓提到的不同池化后端仅作为示例，仅供您用于非商业目的。如您使用这些池化后端完成示例，请您特别注意应遵守对应池化后端的License，如您因使用池化后端而产生侵权纠纷，华为不承担任何责任。	119	- 本代码仓提到的不同池化后端仅作为示例，仅供您用于非商业目的。如您使用这些池化后端完成示例，请您特别注意应遵守对应池化后端的License，如您因使用池化后端而产生侵权纠纷，华为不承担任何责任。

Mexamples/atb_models/atb_framework/core/base/event_manager.cpp+79-0

	@@ -74,6 +74,13 @@ EventManager::~EventManager()
74	}	74	}
75	}	75	}
76	ATB_SPEED_LOG_DEBUG("EventManager destroyed.");	76	ATB_SPEED_LOG_DEBUG("EventManager destroyed.");
		77	+ for (auto& eventsTuple : eventsForExternal_) {
		78	+ auto& eventsForExtel = std::get<1>(eventsTuple.second);
		79	+ for (auto& eventE : eventsForExtel) {
		80	+ DestroyEvent(eventE);
		81	+ }
		82	+ }
		83	+ ATB_SPEED_LOG_DEBUG("EventManager destroyed.");
77	}	84	}
78		85
79	void EventManager::SetWaitOperationTimeout(uint32_t timeout)	86	void EventManager::SetWaitOperationTimeout(uint32_t timeout)
	@@ -219,4 +226,76 @@ atb::Status EventManager::WaitEvent(atb::Operation*& op, EventAction eventAction
219	return EventInternal(eventAction, EventType::WAIT, op, pipeKey);	226	return EventInternal(eventAction, EventType::WAIT, op, pipeKey);
220	}	227	}
221		228
		229	+EventManagerStatus EventManager::CheckPipeKey(const std::string &pipeKey)
		230	+{
		231	+ if (eventsForExternal_.find(pipeKey) == eventsForExternal_.end()) {
		232	+ if (eventQueues_.find(pipeKey) == eventQueues_.end()) {
		233	+ ATB_SPEED_LOG_DEBUG("Event for external error: no such pipekey" << pipeKey);
		234	+ return EM_INVALID_ACTION;
		235	+ }
		236	+ std::vector<aclrtEvent> queue;
		237	+ while (!eventQueues_[pipeKey].empty()) {
		238	+ queue.push_back(eventQueues_[pipeKey].front());
		239	+ eventQueues_[pipeKey].pop();
		240	+ }
		241	+ aclrtStream subStream;
		242	+ aclError ret = aclrtCreateStream(&subStream);
		243	+ if (ret != ACL_ERROR_NONE) {
		244	+ ATB_SPEED_LOG_ERROR("Failed to create aclrtStream: " << ret);
		245	+ return EM_INVALID_ACTION;
		246	+ }
		247	+ ret = aclrtSetStreamFailureMode(subStream, ACL_STOP_ON_FAILURE);
		248	+ if (ret != 0) {
		249	+ ATB_SPEED_LOG_ERROR("Failed to aclrtSetStreamFailureMode: " << ret);
		250	+ return EM_INVALID_ACTION;
		251	+ }
		252	+ eventsForExternal_[pipeKey] = std::make_tuple(0, queue, subStream);
		253	+ }
		254	+ return EM_SUCCESS;
		255	+}
		256	+
		257	+EventManagerStatus EventManager::RecordEvent(const std::string &pipeKey)
		258	+{
		259	+ auto rt = CheckPipeKey(pipeKey);
		260	+ if (rt != EM_SUCCESS) {
		261	+ return rt;
		262	+ }
		263	+ auto& currentEventIdx = std::get<0>(eventsForExternal_[pipeKey]);
		264	+ auto& eventsForExtel = std::get<1>(eventsForExternal_[pipeKey]);
		265	+ auto& stream = std::get<2>(eventsForExternal_[pipeKey]);
		266	+ aclError ret = aclrtRecordEvent(eventsForExtel[currentEventIdx], stream);
		267	+ currentEventIdx = (currentEventIdx + 1) % eventsForExtel.size();
		268	+ if (ret != ACL_SUCCESS) {
		269	+ ATB_SPEED_LOG_ERROR("aclrtRecordEvent fail, ret: " << ret);
		270	+ return EM_INVALID_ACTION;
		271	+ }
		272	+ return EM_SUCCESS;
		273	+}
		274	+
		275	+EventManagerStatus EventManager::WaitEvent(const std::string &pipeKey)
		276	+{
		277	+ auto rt = CheckPipeKey(pipeKey);
		278	+ if (rt != EM_SUCCESS) {
		279	+ return rt;
		280	+ }
		281	+ auto& currentEventIdx = std::get<0>(eventsForExternal_[pipeKey]);
		282	+ auto& eventsForExtel = std::get<1>(eventsForExternal_[pipeKey]);
		283	+ auto& stream = std::get<2>(eventsForExternal_[pipeKey]);
		284	+ aclError ret = aclrtStreamWaitEvent(stream, eventsForExtel[currentEventIdx]);
		285	+ if (ret != ACL_SUCCESS) {
		286	+ ATB_SPEED_LOG_ERROR("aclrtWaitEvent fail, ret: " << ret);
		287	+ return EM_INVALID_ACTION;
		288	+ }
		289	+ ret = aclrtResetEvent(eventsForExtel[currentEventIdx], stream);
		290	+ if (ret != ACL_SUCCESS) {
		291	+ ATB_SPEED_LOG_ERROR("aclrtResetEvent fail, ret: " << ret);
		292	+ return EM_INVALID_ACTION;
		293	+ }
		294	+ if (aclrtSynchronizeStream(stream) != 0) {
		295	+ ATB_SPEED_LOG_ERROR("EventManager::WaitEvent: aclrtSynchronizeStream fail");
		296	+ return EM_INVALID_ACTION;
		297	+ }
		298	+ currentEventIdx = (currentEventIdx + 1) % eventsForExtel.size();
		299	+ return EM_SUCCESS;
		300	+}
222	} // namespace atb_speed	301	} // namespace atb_speed

Mexamples/atb_models/atb_framework/core/include/atb_speed/base/event_manager.h+20-1

	@@ -13,7 +13,6 @@
13	#ifndef ATB_SPEED_EVENT_MANAGER_H	13	#ifndef ATB_SPEED_EVENT_MANAGER_H
14	#define ATB_SPEED_EVENT_MANAGER_H	14	#define ATB_SPEED_EVENT_MANAGER_H
15		15
16	-#include <acl/acl.h>
17	#include <atb/context.h>	16	#include <atb/context.h>
18	#include <atb/operation.h>	17	#include <atb/operation.h>
19	#include <atb/atb_infer.h>	18	#include <atb/atb_infer.h>
	@@ -96,6 +95,22 @@ public:
96	*/	95	*/
97	atb::Status WaitEvent(atb::Operation*& op, EventAction eventAction, const std::string &pipeKey = "default");	96	atb::Status WaitEvent(atb::Operation*& op, EventAction eventAction, const std::string &pipeKey = "default");
98		97
		98	+ /**
		99	+ * @brief 在多流场景下，用于图外部与图内部的流间同步，取pipe中的current event record一次，current event指向下一个event，循环event队列。
		100	+ 在同一张图里RecordEvent和WaitEvent只能使用同一个pipekey，否则会有时序问题。
		101	+ * @param pipeKey 事件管道的标识符（默认为 "default"）
		102	+ * @return 返回操作的状态码（atb::Status）
		103	+ */
		104	+ EventManagerStatus RecordEvent(const std::string &pipeKey = "default");
		105	+
		106	+ /**
		107	+ * @brief 在多流场景下，用于图外部与图内部的流间同步，取pipe中的current event wait一次然后sync一次，current event指向下一个event，循环event队列。
		108	+ 在同一张图里RecordEvent和WaitEvent只能使用同一个pipekey，否则会有时序问题。
		109	+ * @param pipeKey 事件管道的标识符（默认为 "default"）
		110	+ * @return 返回操作的状态码（atb::Status）
		111	+ */
		112	+ EventManagerStatus WaitEvent(const std::string &pipeKey = "default");
		113	+
99	private:	114	private:
100	// 构造和析构函数设为 private，确保单例模式	115	// 构造和析构函数设为 private，确保单例模式
101	// 设置 ACL 的操作等待超时时间，默认 180 秒	116	// 设置 ACL 的操作等待超时时间，默认 180 秒
	@@ -145,6 +160,8 @@ private:
145	EventType eventType,	160	EventType eventType,
146	atb::Operation*& op,	161	atb::Operation*& op,
147	const std::string &pipeKey);	162	const std::string &pipeKey);
		163	+
		164	+ EventManagerStatus CheckPipeKey(const std::string &pipeKey);
148		165
149	private:	166	private:
150	// 条件变量，用于在事件队列为空时阻塞等待，并在有新事件入队时通知等待线程	167	// 条件变量，用于在事件队列为空时阻塞等待，并在有新事件入队时通知等待线程
	@@ -152,6 +169,8 @@ private:
152	// 事件队列：仅用于对外提供 push/pop 接口	169	// 事件队列：仅用于对外提供 push/pop 接口
153	std::map<std::string, std::queue<aclrtEvent>> eventQueues_;	170	std::map<std::string, std::queue<aclrtEvent>> eventQueues_;
154	std::map<std::string, std::queue<std::pair<atb::Operation*, atb::common::EventParam>>> opsWithoutEvent_;	171	std::map<std::string, std::queue<std::pair<atb::Operation*, atb::common::EventParam>>> opsWithoutEvent_;
		172	+ // 用于图外部与图内部的流间同步
		173	+ std::map<std::string, std::tuple<int, std::vector<aclrtEvent>, aclrtStream>> eventsForExternal_;
155	// 保护 eventQueue_ 的互斥锁	174	// 保护 eventQueue_ 的互斥锁
156	std::mutex queueMutex_;	175	std::mutex queueMutex_;
157	// 记录当前队列中 event 的数量	176	// 记录当前队列中 event 的数量

Mexamples/atb_models/atb_framework/models/base/model/decoder_model.cpp+18-0

	@@ -18,6 +18,7 @@
18		18
19	#include "models/base/model/decoder_model.h"	19	#include "models/base/model/decoder_model.h"
20	#include "models/base/layer/decoder_layer.h"	20	#include "models/base/layer/decoder_layer.h"
		21	+#include "models/base/param/model_param.h"
21	#include "operations/aclnn/ops/split_with_size_operation.h"	22	#include "operations/aclnn/ops/split_with_size_operation.h"
22	#include "operations/fusion/infer_shape_functions.h"	23	#include "operations/fusion/infer_shape_functions.h"
23		24
	@@ -725,6 +726,7 @@ atb::Status DecoderModel::AddLayer()
725		726
726	uint32_t nodeCount = graph_.nodes.size();	727	uint32_t nodeCount = graph_.nodes.size();
727	this->AddSingleLayer(trueLayerId);	728	this->AddSingleLayer(trueLayerId);
		729	+ RecordEventBeforePrefixCacheSave();
728	for (uint32_t index = nodeCount; index < graph_.nodes.size(); index++) {	730	for (uint32_t index = nodeCount; index < graph_.nodes.size(); index++) {
729	if (GetSingleton<common::DapManager>().GetRole() == common::DapRole::SUCCESSOR) {	731	if (GetSingleton<common::DapManager>().GetRole() == common::DapRole::SUCCESSOR) {
730	CHECK_OPERATION_STATUS_RETURN(SetNodeStreamId(graph_.nodes.at(index), 1));	732	CHECK_OPERATION_STATUS_RETURN(SetNodeStreamId(graph_.nodes.at(index), 1));
	@@ -740,6 +742,22 @@ atb::Status DecoderModel::AddLayer()
740	return atb::NO_ERROR;	742	return atb::NO_ERROR;
741	}	743	}
742		744
		745	+atb::Status DecoderModel::RecordEventBeforePrefixCacheSave()
		746	+{
		747	+ if (param.memPoolType == atb_speed::base::MemPoolType::ASYNC_WRITE && param.isPrefill) {
		748	+ atb::Operation *op = nullptr;
		749	+ atb_speed::Model::Node recordSaveNode;
		750	+ CHECK_OPERATION_STATUS_RETURN(atb_speed::EventManager::GetInstance().RecordEvent(
		751	+ op, atb_speed::EventAction::PUSH, param.memPoolEventPipeKey));
		752	+ recordSaveNode.inTensors = {};
		753	+ recordSaveNode.outTensors = {};
		754	+ recordSaveNode.operation.reset(op);
		755	+ graph_.nodes.push_back(recordSaveNode);
		756	+ ATB_SPEED_LOG_DEBUG("[Events] [PUSH] [RECORD] will be pushed to the graph later");
		757	+ }
		758	+ return atb::NO_ERROR;
		759	+}
		760	+
743	atb::Status DecoderModel::AddSingleLayer(uint32_t layerId)	761	atb::Status DecoderModel::AddSingleLayer(uint32_t layerId)
744	{	762	{
745	atb::Operation *op = nullptr;	763	atb::Operation *op = nullptr;

Mexamples/atb_models/atb_framework/models/base/model/decoder_model.h+2-0

	@@ -261,6 +261,8 @@ private:
261	std::map<uint32_t, uint32_t> CopyMapWithSuffix(std::map<std::string, uint32_t>& tensorMap) const;	261	std::map<uint32_t, uint32_t> CopyMapWithSuffix(std::map<std::string, uint32_t>& tensorMap) const;
262	std::map<atb::Tensor , atb::Tensor > precederToSuccessorTensorMap = {};	262	std::map<atb::Tensor , atb::Tensor > precederToSuccessorTensorMap = {};
263	void ReplaceDapTensors(std::vector<atb::Tensor *>& tensors);	263	void ReplaceDapTensors(std::vector<atb::Tensor *>& tensors);
		264	+ atb::Status WaitEventAfterPrefixCacheLoad();
		265	+ atb::Status RecordEventBeforePrefixCacheSave();
264	};	266	};
265		267
266	} // namespace base	268	} // namespace base

Mexamples/atb_models/atb_framework/models/base/param/model_param.cpp+6-0

	@@ -86,6 +86,12 @@ void ModelParam::ParseParam(const nlohmann::json &paramJson)
86	if (paramJson.contains("enableFlashComm")) {	86	if (paramJson.contains("enableFlashComm")) {
87	this->enableFlashComm = paramJson["enableFlashComm"].get<bool>();	87	this->enableFlashComm = paramJson["enableFlashComm"].get<bool>();
88	}	88	}
		89	+ if (paramJson.contains("memPoolType")) {
		90	+ this->memPoolType = FetchJsonParam<MemPoolType>(paramJson, "memPoolType");
		91	+ }
		92	+ if (paramJson.contains("pipeKey")) {
		93	+ this->memPoolEventPipeKey = FetchJsonParam<std::string>(paramJson, "pipeKey");
		94	+ }
89	ParseNormParam(paramJson);	95	ParseNormParam(paramJson);
90	ParseAttentionParam(paramJson);	96	ParseAttentionParam(paramJson);
91	ParseMlpParam(paramJson);	97	ParseMlpParam(paramJson);

Mexamples/atb_models/atb_framework/models/base/param/model_param.h+9-0

	@@ -24,6 +24,12 @@ const int LWD_EDGE_LAST = 2;
24	namespace atb_speed {	24	namespace atb_speed {
25	namespace base {	25	namespace base {
26		26
		27	+enum MemPoolType {
		28	+ DISABLED = 0,
		29	+ SYNC_WRITE = 1,
		30	+ ASYNC_WRITE = 2
		31	+};
		32	+
27	/// Parse string to `nlohmann::json`	33	/// Parse string to `nlohmann::json`
28	/// \param param parameters in JSON string format passed from the Python side	34	/// \param param parameters in JSON string format passed from the Python side
29	/// \return parameters in `nlohmann::json` format	35	/// \return parameters in `nlohmann::json` format
	@@ -83,6 +89,9 @@ public:
83	bool reuseEmbedTable = false;	89	bool reuseEmbedTable = false;
84	bool outputEmbedTable = false;	90	bool outputEmbedTable = false;
85		91
		92	+ MemPoolType memPoolType = MemPoolType::DISABLED;
		93	+ // When `mempool_type` is ASYNC_WRITE, a event pipeKey is set for async layer-wise prefix cache transfer.
		94	+ std::string memPoolEventPipeKey = "default";
86	std::string backend = "hccl";	95	std::string backend = "hccl";
87	std::string tpDomain = "";	96	std::string tpDomain = "";
88	std::string rankTableFile = "";	97	std::string rankTableFile = "";

Aexamples/atb_models/atb_framework/pytorch/adapter/event/event.cpp+24-0

@@ -0,0 +1,24 @@
	1	+/*
	2	+ * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
	3	+ *
	4	+ * Licensed under the Apache License, Version 2.0 (the "License");
	5	+ * you may not use this file except in compliance with the License.
	6	+ * You may obtain a copy of the License at
	7	+ *
	8	+ * http://www.apache.org/licenses/LICENSE-2.0
	9	+ *
	10	+ * Unless required by applicable law or agreed to in writing, software
	11	+ * distributed under the License is distributed on an "AS IS" BASIS,
	12	+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	13	+ * See the License for the specific language governing permissions and
	14	+ * limitations under the License.
	15	+ */
	16	+
	17	+#include "event.h"
	18	+
	19	+namespace atb_speed {
	20	+
	21	+void Event::Record(const std::string& pipeKey) { EventManager::GetInstance().RecordEvent(pipeKey); }
	22	+void Event::Wait(const std::string& pipeKey) { EventManager::GetInstance().WaitEvent(pipeKey); }
	23	+
	24	+} // namespace atb_speed

Aexamples/atb_models/atb_framework/pytorch/adapter/event/event.h+32-0

@@ -0,0 +1,32 @@
	1	+/*
	2	+ * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
	3	+ *
	4	+ * Licensed under the Apache License, Version 2.0 (the "License");
	5	+ * you may not use this file except in compliance with the License.
	6	+ * You may obtain a copy of the License at
	7	+ *
	8	+ * http://www.apache.org/licenses/LICENSE-2.0
	9	+ *
	10	+ * Unless required by applicable law or agreed to in writing, software
	11	+ * distributed under the License is distributed on an "AS IS" BASIS,
	12	+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	13	+ * See the License for the specific language governing permissions and
	14	+ * limitations under the License.
	15	+ */
	16	+#ifndef ATB_SPEED_UTILS_EVENT_H
	17	+#define ATB_SPEED_UTILS_EVENT_H
	18	+
	19	+#include <torch/torch.h>
	20	+#include <torch_npu/csrc/framework/OpCommand.h>
	21	+#include <torch/custom_class.h>
	22	+#include <torch/script.h>
	23	+#include "atb_speed/base/event_manager.h"
	24	+
	25	+namespace atb_speed {
	26	+class Event : public torch::CustomClassHolder {
	27	+public:
	28	+ static void Record(const std::string& pipeKey);
	29	+ static void Wait(const std::string& pipeKey);
	30	+};
	31	+} // namespace atb_speed
	32	+#endif

Mexamples/atb_models/atb_framework/pytorch/adapter/model/model_torch.cpp+5-0

	@@ -27,6 +27,7 @@
27	#include "pytorch/adapter/workspace/workspace.h"	27	#include "pytorch/adapter/workspace/workspace.h"
28	#include "atb_speed/utils/model_factory.h"	28	#include "atb_speed/utils/model_factory.h"
29	#include "pytorch/adapter/context/context.h"	29	#include "pytorch/adapter/context/context.h"
		30	+#include "pytorch/adapter/event/event.h"
30		31
31	namespace atb_speed {	32	namespace atb_speed {
32	static std::mutex g_taskQueueLock;	33	static std::mutex g_taskQueueLock;
	@@ -343,5 +344,9 @@ TORCH_LIBRARY(ModelTorch, m)
343	.def_static("enable_cache_workspace", &Context::EnableCacheWorkspace)	344	.def_static("enable_cache_workspace", &Context::EnableCacheWorkspace)
344	.def_static("disable_cache_workspace", &Context::DisableCacheWorkspace)	345	.def_static("disable_cache_workspace", &Context::DisableCacheWorkspace)
345	.def_static("resume_hccl_comm", &Context::ResumeHcclComm);	346	.def_static("resume_hccl_comm", &Context::ResumeHcclComm);
		347	+
		348	+ m.class_<Event>("Event")
		349	+ .def_static("record", &Event::Record)
		350	+ .def_static("wait", &Event::Wait);
346	}	351	}
347	}	352	}

Mexamples/atb_models/atb_llm/models/base/flash_causal_lm.py+8-1

	@@ -19,6 +19,7 @@ from transformers.configuration_utils import PretrainedConfig
19		19
20	from atb_llm.utils.eplb_expert_data_collect import EplbExpertDataCollect	20	from atb_llm.utils.eplb_expert_data_collect import EplbExpertDataCollect
21	from atb_llm.utils.moe_utils import EPLBType	21	from atb_llm.utils.moe_utils import EPLBType
		22	+from mindie_llm.text_generator.plugins.plugin_manager import MemPoolType
22	from .model_utils import BaseModel	23	from .model_utils import BaseModel
23	from ...models import InferenceMode	24	from ...models import InferenceMode
24	from ...utils.env import ENV	25	from ...utils.env import ENV
	@@ -93,7 +94,7 @@ class FlashForCausalLM(BaseModel):
93	self.layerwise = LayerWiseAttr(edge_start_layer_count, edge_end_layer_count, split_type)	94	self.layerwise = LayerWiseAttr(edge_start_layer_count, edge_end_layer_count, split_type)
94		95
95	self.inference_mode = kwargs.get("inference_mode")	96	self.inference_mode = kwargs.get("inference_mode")
96	-	97	+ self.mempool_type: MemPoolType = kwargs.get('mempool_type', MemPoolType.DISABLED)
97	self.num_attention_heads = config.num_attention_heads	98	self.num_attention_heads = config.num_attention_heads
98	if hasattr(config, 'num_key_value_heads'):	99	if hasattr(config, 'num_key_value_heads'):
99	self.num_key_value_heads = config.num_key_value_heads	100	self.num_key_value_heads = config.num_key_value_heads
	@@ -620,3 +621,9 @@ class FlashForCausalLM(BaseModel):
620	logits = self.execute_dap_ascend_operator(	621	logits = self.execute_dap_ascend_operator(
621	all_inputs, json.dumps(acl_param_dict), is_prefill[0])	622	all_inputs, json.dumps(acl_param_dict), is_prefill[0])
622	return logits	623	return logits
		624	+
		625	+ def wait_model_event(self, event_pipe_key: str):
		626	+ torch.classes.ModelTorch.Event.wait(event_pipe_key)
		627	+
		628	+ def record_model_event(self, event_pipe_key: str):
		629	+ torch.classes.ModelTorch.Event.record(event_pipe_key)

Mexamples/atb_models/atb_llm/models/base/graph_manager/__init__.py+1-0

	@@ -16,6 +16,7 @@ from atb_llm.models.base.graph_manager.single_lora_graph_wrapper import SingleLo
16	from atb_llm.models.base.graph_manager.multi_lora_graph_wrapper import MultiLoraGraphWrapper	16	from atb_llm.models.base.graph_manager.multi_lora_graph_wrapper import MultiLoraGraphWrapper
17	from atb_llm.models.base.graph_manager.speculate_graph_wrapper import SpeculateGraphWrapper	17	from atb_llm.models.base.graph_manager.speculate_graph_wrapper import SpeculateGraphWrapper
18	from atb_llm.models.base.graph_manager.splitfuse_graph_wrapper import SplitFuseGraphWrapper	18	from atb_llm.models.base.graph_manager.splitfuse_graph_wrapper import SplitFuseGraphWrapper
		19	+from atb_llm.models.base.graph_manager.mem_pool_graph_wrapper import MemPoolGraphWrapper
19	from atb_llm.models.base.graph_manager.layerwise_decode_graph_wrapper import get_layerwise_decode_graph	20	from atb_llm.models.base.graph_manager.layerwise_decode_graph_wrapper import get_layerwise_decode_graph
20	from atb_llm.models.base.graph_manager.layerwise_prefill_graph_wrapper import get_layerwise_prefill_graph	21	from atb_llm.models.base.graph_manager.layerwise_prefill_graph_wrapper import get_layerwise_prefill_graph
21	from atb_llm.models.base.graph_manager.layerwise_combined_graph_wrapper import LayerwiseCombinedATBGraphWrapper	22	from atb_llm.models.base.graph_manager.layerwise_combined_graph_wrapper import LayerwiseCombinedATBGraphWrapper

Mexamples/atb_models/atb_llm/models/base/graph_manager/compatible_matrix.py+3-2

	@@ -21,11 +21,11 @@ class FeatureType(str, Enum):
21	FLASHCOMM = "flashcomm"	21	FLASHCOMM = "flashcomm"
22	SPECULATE = "speculate"	22	SPECULATE = "speculate"
23	SPLITFUSE = "splitfuse"	23	SPLITFUSE = "splitfuse"
24	-	24	+ MOONCAKE = "mempool"
25		25
26	COMPATIBLE_MATRIX: dict[FeatureType, list[FeatureType]] = {	26	COMPATIBLE_MATRIX: dict[FeatureType, list[FeatureType]] = {
27	FeatureType.PREFILL: [FeatureType.SINGLE_LORA, FeatureType.MULTI_LORA, FeatureType.FLASHCOMM, \	27	FeatureType.PREFILL: [FeatureType.SINGLE_LORA, FeatureType.MULTI_LORA, FeatureType.FLASHCOMM, \
28	- FeatureType.DAP, FeatureType.SPLITFUSE],	28	+ FeatureType.DAP, FeatureType.SPLITFUSE, FeatureType.MOONCAKE],
29	FeatureType.DECODE: [FeatureType.SINGLE_LORA, FeatureType.MULTI_LORA, FeatureType.SPECULATE],	29	FeatureType.DECODE: [FeatureType.SINGLE_LORA, FeatureType.MULTI_LORA, FeatureType.SPECULATE],
30	FeatureType.SINGLE_LORA: [FeatureType.PREFILL, FeatureType.DECODE, FeatureType.SPLITFUSE, FeatureType.FLASHCOMM],	30	FeatureType.SINGLE_LORA: [FeatureType.PREFILL, FeatureType.DECODE, FeatureType.SPLITFUSE, FeatureType.FLASHCOMM],
31	FeatureType.MULTI_LORA: [FeatureType.PREFILL, FeatureType.DECODE, FeatureType.SPLITFUSE, FeatureType.FLASHCOMM],	31	FeatureType.MULTI_LORA: [FeatureType.PREFILL, FeatureType.DECODE, FeatureType.SPLITFUSE, FeatureType.FLASHCOMM],
	@@ -37,4 +37,5 @@ COMPATIBLE_MATRIX: dict[FeatureType, list[FeatureType]] = {
37	FeatureType.LAYERWISE_PREFILL, FeatureType.SINGLE_LORA, FeatureType.MULTI_LORA],	37	FeatureType.LAYERWISE_PREFILL, FeatureType.SINGLE_LORA, FeatureType.MULTI_LORA],
38	FeatureType.LAYERWISE_PREFILL: [FeatureType.SPLITFUSE],	38	FeatureType.LAYERWISE_PREFILL: [FeatureType.SPLITFUSE],
39	FeatureType.LAYERWISE_DECODE: [],	39	FeatureType.LAYERWISE_DECODE: [],
		40	+ FeatureType.MOONCAKE: [FeatureType.PREFILL],
40	}	41	}

Aexamples/atb_models/atb_llm/models/base/graph_manager/mem_pool_graph_wrapper.py+27-0

@@ -0,0 +1,27 @@
	1	+# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
	2	+# MindIE is licensed under Mulan PSL v2.
	3	+# You can use this software according to the terms and conditions of the Mulan PSL v2.
	4	+# You may obtain a copy of Mulan PSL v2 at:
	5	+# http://license.coscl.org.cn/MulanPSL2
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
	7	+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
	8	+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
	9	+# See the Mulan PSL v2 for more details.
	10	+from atb_llm.models.base.graph_manager.graph_wrapper import ATBGraphWrapper
	11	+from atb_llm.models.base.graph_manager.compatible_matrix import FeatureType
	12	+from atb_llm.models.base.flash_causal_lm import FlashForCausalLM
	13	+from atb_llm.runner.model_runner import generate_mem_pool_event_key
	14	+from mindie_llm.text_generator.plugins.plugin_manager import MemPoolType
	15	+
	16	+
	17	+class MemPoolGraphWrapper(ATBGraphWrapper):
	18	+ pipe_key: str = generate_mem_pool_event_key(only_save_kv=True)
	19	+
	20	+ def __init__(self):
	21	+ super().__init__()
	22	+ self.feature_name = FeatureType.MOONCAKE
	23	+ self.feature_params = {"memPoolType": int(MemPoolType.ASYNC_WRITE), "pipeKey": self.pipe_key}
	24	+
	25	+ def activate(self, context: FlashForCausalLM, runtime_params, **kwargs) -> bool:
	26	+ mempool_type = kwargs.get("mempool_type", MemPoolType.DISABLED)
	27	+ return mempool_type == MemPoolType.ASYNC_WRITE

Mexamples/atb_models/atb_llm/models/base/graph_manager/splitfuse_graph_wrapper.py+6-4

Mexamples/atb_models/atb_llm/models/deepseekv2/flash_causal_deepseekv2.py+110-17

	@@ -42,6 +42,7 @@ from atb_llm.utils.weights import ProcessGroupType
42	from atb_llm.utils.log import print_log	42	from atb_llm.utils.log import print_log
43	from atb_llm.utils import file_utils	43	from atb_llm.utils import file_utils
44	from atb_llm.utils.eplb_expert_data_collect import EplbExpertDataCollect	44	from atb_llm.utils.eplb_expert_data_collect import EplbExpertDataCollect
		45	+from mindie_llm.text_generator.plugins.plugin_manager import MemPoolType
45	from ...models import InferenceMode	46	from ...models import InferenceMode
46		47
47	try:	48	try:
	@@ -79,6 +80,12 @@ KVCACHE_QUANT_LAYERS = "kvcacheQuantLayers"
79	MOE_PACK_QUANT_TYPE = "moePackQuantType"	80	MOE_PACK_QUANT_TYPE = "moePackQuantType"
80	ALLTOALL_LONG_SEQLEN_THRESHOLD = 65536	81	ALLTOALL_LONG_SEQLEN_THRESHOLD = 65536
81	MAX_ALLTOALL_BUFF_SCALE = 3	82	MAX_ALLTOALL_BUFF_SCALE = 3
		83	+WARMUP_IS_END = "warmup_is_end"
		84	+
		85	+
		86	+# 直接引用model_runner中的同名实现会导致循环调用和拉起卡死, 故本文件中直接引用这里
		87	+def generate_mem_pool_event_key(only_save_kv: bool) -> str:
		88	+ return "only_save_kv" if only_save_kv else "both_save_kv"
82		89
83		90
84	class MaskType(int, Enum):	91	class MaskType(int, Enum):
	@@ -600,16 +607,26 @@ class FlashDeepseekv2ForCausalLM(FlashForCausalLM):
600	def init_ascend_operations(self, config: DeepseekV2Config):	607	def init_ascend_operations(self, config: DeepseekV2Config):
601	if not self.layerwise_disaggregated:	608	if not self.layerwise_disaggregated:
602	self.acl_encoder_operation = torch.classes.ModelTorch.ModelTorch(CPP_DEEPSEEKV2_MODEL_CLASS_NAME)	609	self.acl_encoder_operation = torch.classes.ModelTorch.ModelTorch(CPP_DEEPSEEKV2_MODEL_CLASS_NAME)
		610	+ if self.mempool_type == MemPoolType.ASYNC_WRITE:
		611	+ self.acl_encoder_operation_prefixcache_mempool = \
		612	+ torch.classes.ModelTorch.ModelTorch(CPP_DEEPSEEKV2_MODEL_CLASS_NAME)
603	logger.info(f"when init_ascend_operations, self.prefix_cache_enable is : {self.prefix_cache_enable}")	613	logger.info(f"when init_ascend_operations, self.prefix_cache_enable is : {self.prefix_cache_enable}")
604	if self.prefix_cache_enable:	614	if self.prefix_cache_enable:
605	self.acl_encoder_operation_prefixcache = \	615	self.acl_encoder_operation_prefixcache = \
606	torch.classes.ModelTorch.ModelTorch(CPP_DEEPSEEKV2_MODEL_CLASS_NAME)	616	torch.classes.ModelTorch.ModelTorch(CPP_DEEPSEEKV2_MODEL_CLASS_NAME)
		617	+ self.acl_encoder_operation_prefixcache_async_write = \
		618	+ torch.classes.ModelTorch.ModelTorch(CPP_DEEPSEEKV2_MODEL_CLASS_NAME)
607	if self.num_speculative_tokens:	619	if self.num_speculative_tokens:
608	self.acl_encoder_operation_mtp = torch.classes.ModelTorch.ModelTorch(	620	self.acl_encoder_operation_mtp = torch.classes.ModelTorch.ModelTorch(
609	- CPP_DEEPSEEKV2_MTP_MODEL_CLASS_NAME)	621	+ CPP_DEEPSEEKV2_MTP_MODEL_CLASS_NAME)
		622	+ if self.mempool_type == MemPoolType.ASYNC_WRITE:
		623	+ self.acl_encoder_operation_prefixcache_mempool_mtp = \
		624	+ torch.classes.ModelTorch.ModelTorch(CPP_DEEPSEEKV2_MTP_MODEL_CLASS_NAME)
610	if self.prefix_cache_enable:	625	if self.prefix_cache_enable:
611	self.acl_encoder_operation_prefixcache_mtp = \	626	self.acl_encoder_operation_prefixcache_mtp = \
612	torch.classes.ModelTorch.ModelTorch(CPP_DEEPSEEKV2_MTP_MODEL_CLASS_NAME)	627	torch.classes.ModelTorch.ModelTorch(CPP_DEEPSEEKV2_MTP_MODEL_CLASS_NAME)
		628	+ self.acl_encoder_operation_prefixcache_mtp_async_write = \
		629	+ torch.classes.ModelTorch.ModelTorch(CPP_DEEPSEEKV2_MTP_MODEL_CLASS_NAME)
613	self.acl_decoder_operation = torch.classes.ModelTorch.ModelTorch(CPP_DEEPSEEKV2_MODEL_CLASS_NAME)	630	self.acl_decoder_operation = torch.classes.ModelTorch.ModelTorch(CPP_DEEPSEEKV2_MODEL_CLASS_NAME)
614	if self.num_speculative_tokens:	631	if self.num_speculative_tokens:
615	self.acl_decoder_operation_mtp = torch.classes.ModelTorch.ModelTorch(	632	self.acl_decoder_operation_mtp = torch.classes.ModelTorch.ModelTorch(
	@@ -1095,9 +1112,18 @@ class FlashDeepseekv2ForCausalLM(FlashForCausalLM):
1095	"enableLcocTp": self.enable_lcoc_tp,	1112	"enableLcocTp": self.enable_lcoc_tp,
1096	"enableFusedMLA": self.enable_fused_mla,	1113	"enableFusedMLA": self.enable_fused_mla,
1097	}	1114	}
		1115	+ encoder_param_mempool = {**encoder_param,
		1116	+ "memPoolType": int(self.mempool_type),
		1117	+ "pipeKey": generate_mem_pool_event_key(only_save_kv=True),
		1118	+ }
1098	encoder_param_prefixcache = {**encoder_param,	1119	encoder_param_prefixcache = {**encoder_param,
1099	- "enablePrefixCache": self.prefix_cache_enable,	1120	+ "enablePrefixCache": self.prefix_cache_enable,
1100	- }	1121	+ }
		1122	+ encoder_param_prefixcache_async_write = {**encoder_param,
		1123	+ "enablePrefixCache": self.prefix_cache_enable,
		1124	+ "memPoolType": int(self.mempool_type),
		1125	+ "pipeKey": generate_mem_pool_event_key(only_save_kv=False),
		1126	+ }
1101	decoder_param = {**coder_param, "isPrefill": False, "supportLcoc": False,	1127	decoder_param = {**coder_param, "isPrefill": False, "supportLcoc": False,
1102	"expertParallelDegree": self.ep_level \	1128	"expertParallelDegree": self.ep_level \
1103	if self.ep_level != ExpertParallelDegree.MIX_EP \	1129	if self.ep_level != ExpertParallelDegree.MIX_EP \
	@@ -1120,9 +1146,16 @@ class FlashDeepseekv2ForCausalLM(FlashForCausalLM):
1120	self.acl_encoder_operation.set_format_to_nz(weight_id)	1146	self.acl_encoder_operation.set_format_to_nz(weight_id)
1121	if self.enable_atlas_gmm_fused:	1147	if self.enable_atlas_gmm_fused:
1122	self.acl_encoder_operation.set_format_to_nz(weight_id + 6)	1148	self.acl_encoder_operation.set_format_to_nz(weight_id + 6)
		1149	+ if self.mempool_type == MemPoolType.ASYNC_WRITE:
		1150	+ self.acl_encoder_operation_prefixcache_mempool.set_param(json.dumps({**encoder_param_mempool}))
		1151	+ self.acl_encoder_operation_prefixcache_mempool.set_weight(self.ascend_weight)
1123	if self.prefix_cache_enable:	1152	if self.prefix_cache_enable:
1124	self.acl_encoder_operation_prefixcache.set_param(json.dumps({**encoder_param_prefixcache}))	1153	self.acl_encoder_operation_prefixcache.set_param(json.dumps({**encoder_param_prefixcache}))
1125	self.acl_encoder_operation_prefixcache.set_weight(self.ascend_weight)	1154	self.acl_encoder_operation_prefixcache.set_weight(self.ascend_weight)
		1155	+ self.acl_encoder_operation_prefixcache_async_write.set_param(
		1156	+ json.dumps({**encoder_param_prefixcache_async_write})
		1157	+ )
		1158	+ self.acl_encoder_operation_prefixcache_async_write.set_weight(self.ascend_weight)
1126		1159
1127	if self.acl_decoder_operation is not None:	1160	if self.acl_decoder_operation is not None:
1128	self.acl_decoder_operation.set_param(json.dumps({**decoder_param}))	1161	self.acl_decoder_operation.set_param(json.dumps({**decoder_param}))
	@@ -1257,11 +1290,28 @@ class FlashDeepseekv2ForCausalLM(FlashForCausalLM):
1257	self.acl_decoder_operation_mtp.set_format_to_nz(weight_id)	1290	self.acl_decoder_operation_mtp.set_format_to_nz(weight_id)
1258	if self.enable_atlas_gmm_fused:	1291	if self.enable_atlas_gmm_fused:
1259	self.acl_decoder_operation_mtp.set_format_to_nz(weight_id + 6)	1292	self.acl_decoder_operation_mtp.set_format_to_nz(weight_id + 6)
1260	-	1293	+ if self.mempool_type == MemPoolType.ASYNC_WRITE:
		1294	+ encoder_param_mempool_mtp = {
		1295	+ **encoder_param_mtp,
		1296	+ "memPoolType": int(self.mempool_type),
		1297	+ "pipeKey": generate_mem_pool_event_key(only_save_kv=True)
		1298	+ }
		1299	+ self.acl_encoder_operation_prefixcache_mempool_mtp.set_param(json.dumps({**encoder_param_mempool_mtp}))
		1300	+ self.acl_encoder_operation_prefixcache_mempool_mtp.set_weight(self.ascend_weight_mtp)
1261	if self.prefix_cache_enable:	1301	if self.prefix_cache_enable:
1262	- encoder_param_mtp = {**encoder_param_mtp, "enablePrefixCache": self.prefix_cache_enable}	1302	+ encoder_param_mtp = {
		1303	+ **encoder_param_mtp,
		1304	+ "enablePrefixCache": self.prefix_cache_enable,
		1305	+ }
1263	self.acl_encoder_operation_prefixcache_mtp.set_param(json.dumps({**encoder_param_mtp}))	1306	self.acl_encoder_operation_prefixcache_mtp.set_param(json.dumps({**encoder_param_mtp}))
1264	self.acl_encoder_operation_prefixcache_mtp.set_weight(self.ascend_weight_mtp)	1307	self.acl_encoder_operation_prefixcache_mtp.set_weight(self.ascend_weight_mtp)
		1308	+ encoder_param_mtp = {
		1309	+ **encoder_param_mtp,
		1310	+ "memPoolType": int(self.mempool_type),
		1311	+ "pipeKey": generate_mem_pool_event_key(only_save_kv=False)
		1312	+ }
		1313	+ self.acl_encoder_operation_prefixcache_mtp_async_write.set_param(json.dumps({**encoder_param_mtp}))
		1314	+ self.acl_encoder_operation_prefixcache_mtp_async_write.set_weight(self.ascend_weight_mtp)
1265		1315
1266	if self.enable_dap and self.acl_dap_operation is not None:	1316	if self.enable_dap and self.acl_dap_operation is not None:
1267	self.acl_dap_operation.set_param(	1317	self.acl_dap_operation.set_param(
	@@ -1663,7 +1713,7 @@ class FlashDeepseekv2ForCausalLM(FlashForCausalLM):
1663	perf_time_start = 0	1713	perf_time_start = 0
1664	if ENV.benchmark_enable:	1714	if ENV.benchmark_enable:
1665	import time	1715	import time
1666	- torch.npu.synchronize()	1716	+ torch.npu.current_stream().synchronize()
1667	perf_time_start = time.time()	1717	perf_time_start = time.time()
1668	self.expert_array = self.placeholder	1718	self.expert_array = self.placeholder
1669		1719
	@@ -1955,10 +2005,19 @@ class FlashDeepseekv2ForCausalLM(FlashForCausalLM):
1955	"""Execute the Ascend acl operator."""	2005	"""Execute the Ascend acl operator."""
1956	split_part = kwargs.get("split_part", None)	2006	split_part = kwargs.get("split_part", None)
1957	layer_index = kwargs.get("layer_index", None)	2007	layer_index = kwargs.get("layer_index", None)
		2008	+ warmup_is_end = kwargs.get(WARMUP_IS_END, True)
		2009	+ allow_async_write = self.mempool_type == MemPoolType.ASYNC_WRITE and warmup_is_end
1958	if not self.num_speculative_tokens:	2010	if not self.num_speculative_tokens:
1959	if is_prefill and self.prefix_cache_enable and self.has_prefixcache:	2011	if is_prefill and self.prefix_cache_enable and self.has_prefixcache:
1960	if not self.layerwise_disaggregated:	2012	if not self.layerwise_disaggregated:
1961	- acl_model_out = self.acl_encoder_operation_prefixcache.execute(acl_inputs, acl_param)	2013	+ if allow_async_write:
		2014	+ # skip event for warmup
		2015	+ self.acl_encoder_operation_prefixcache_async_write.skip_event(not warmup_is_end)
		2016	+ acl_model_out = self.acl_encoder_operation_prefixcache_async_write.execute(
		2017	+ acl_inputs, acl_param
		2018	+ )
		2019	+ else:
		2020	+ acl_model_out = self.acl_encoder_operation_prefixcache.execute(acl_inputs, acl_param)
1962	else:	2021	else:
1963	if split_part == LwdLayerStatus.EDGE_START_LAYER:	2022	if split_part == LwdLayerStatus.EDGE_START_LAYER:
1964	acl_model_out = self.acl_head_encoder_operation_prefixcache.execute(acl_inputs, acl_param)	2023	acl_model_out = self.acl_head_encoder_operation_prefixcache.execute(acl_inputs, acl_param)
	@@ -1967,7 +2026,10 @@ class FlashDeepseekv2ForCausalLM(FlashForCausalLM):
1967	else:	2026	else:
1968	acl_model_out = self.acl_tail_encoder_operation_prefixcache.execute(acl_inputs, acl_param)	2027	acl_model_out = self.acl_tail_encoder_operation_prefixcache.execute(acl_inputs, acl_param)
1969	self.has_prefixcache = False	2028	self.has_prefixcache = False
1970	-	2029	+ elif is_prefill and allow_async_write and not self.layerwise_disaggregated:
		2030	+ # skip event for warmup
		2031	+ self.acl_encoder_operation_prefixcache_mempool.skip_event(not warmup_is_end)
		2032	+ acl_model_out = self.acl_encoder_operation_prefixcache_mempool.execute(acl_inputs, acl_param)
1971	elif is_prefill:	2033	elif is_prefill:
1972	if self.layerwise_disaggregated:	2034	if self.layerwise_disaggregated:
1973	if split_part == LwdLayerStatus.EDGE_START_LAYER:	2035	if split_part == LwdLayerStatus.EDGE_START_LAYER:
	@@ -2000,13 +2062,27 @@ class FlashDeepseekv2ForCausalLM(FlashForCausalLM):
2000	if is_mtp:	2062	if is_mtp:
2001	encoder_operation = self.acl_encoder_operation_mtp	2063	encoder_operation = self.acl_encoder_operation_mtp
2002	decoder_operation = self.acl_decoder_operation_mtp	2064	decoder_operation = self.acl_decoder_operation_mtp
		2065	+ if allow_async_write:
		2066	+ self.acl_encoder_operation_prefixcache_mempool_mtp.skip_event(not warmup_is_end)
		2067	+ encoder_operation = self.acl_encoder_operation_prefixcache_mempool_mtp
2003	if self.has_prefixcache:	2068	if self.has_prefixcache:
2004	- encoder_operation = self.acl_encoder_operation_prefixcache_mtp	2069	+ if allow_async_write:
		2070	+ self.acl_encoder_operation_prefixcache_mtp_async_write.skip_event(not warmup_is_end)
		2071	+ encoder_operation = self.acl_encoder_operation_prefixcache_mtp_async_write
		2072	+ else:
		2073	+ encoder_operation = self.acl_encoder_operation_prefixcache_mtp
2005	else:	2074	else:
2006	encoder_operation = self.acl_encoder_operation	2075	encoder_operation = self.acl_encoder_operation
2007	decoder_operation = self.acl_decoder_operation	2076	decoder_operation = self.acl_decoder_operation
		2077	+ if allow_async_write:
		2078	+ self.acl_encoder_operation_prefixcache_mempool.skip_event(not warmup_is_end)
		2079	+ encoder_operation = self.acl_encoder_operation_prefixcache_mempool
2008	if self.has_prefixcache:	2080	if self.has_prefixcache:
2009	- encoder_operation = self.acl_encoder_operation_prefixcache	2081	+ if allow_async_write:
		2082	+ self.acl_encoder_operation_prefixcache_async_write.skip_event(not warmup_is_end)
		2083	+ encoder_operation = self.acl_encoder_operation_prefixcache_async_write
		2084	+ else:
		2085	+ encoder_operation = self.acl_encoder_operation_prefixcache
2010		2086
2011	if is_prefill:	2087	if is_prefill:
2012	acl_model_out = encoder_operation.execute(acl_inputs, acl_param)	2088	acl_model_out = encoder_operation.execute(acl_inputs, acl_param)
	@@ -2171,7 +2247,7 @@ class FlashDeepseekv2ForCausalLM(FlashForCausalLM):
2171	llm_hidden_states = (hidden_states[0][lm_head_indices[0]], hidden_states[1][lm_head_indices[1]])	2247	llm_hidden_states = (hidden_states[0][lm_head_indices[0]], hidden_states[1][lm_head_indices[1]])
2172		2248
2173	# logits 是一个tuple，里面是两个logits，logits是一个两维NPU tensor [ntokens, vocabSize]	2249	# logits 是一个tuple，里面是两个logits，logits是一个两维NPU tensor [ntokens, vocabSize]
2174	- torch.npu.synchronize()	2250	+ torch.npu.current_stream().synchronize()
2175	logits_mtp, hidden_states_mtp = logits, hidden_states	2251	logits_mtp, hidden_states_mtp = logits, hidden_states
2176	for mtp_idx in range(self.num_speculative_tokens):	2252	for mtp_idx in range(self.num_speculative_tokens):
2177	self.acl_dap_operation_mtp.set_kv_cache(self.mtp_k_caches[mtp_idx: mtp_idx + 1],	2253	self.acl_dap_operation_mtp.set_kv_cache(self.mtp_k_caches[mtp_idx: mtp_idx + 1],
	@@ -2474,7 +2550,8 @@ class FlashDeepseekv2ForCausalLM(FlashForCausalLM):
2474	span_end(prof, True)	2550	span_end(prof, True)
2475		2551
2476	prof = span_start("operatorExecute", True)	2552	prof = span_start("operatorExecute", True)
2477	- logits = self.execute_ascend_operator(acl_inputs, acl_param, is_prefill)	2553	+ logits = self.execute_ascend_operator(acl_inputs, acl_param, is_prefill,
		2554	+ warmup_is_end=kwargs.get(WARMUP_IS_END, True))
2478	if not is_prefill and self.mapping.has_attn_cp():	2555	if not is_prefill and self.mapping.has_attn_cp():
2479	# During the CP decode stage, each CP domain receives the same token as input,	2556	# During the CP decode stage, each CP domain receives the same token as input,
2480	# resulting in the same output token. As a result, duplicates exist in the aggregated next_token.	2557	# resulting in the same output token. As a result, duplicates exist in the aggregated next_token.
	@@ -2522,16 +2599,23 @@ class FlashDeepseekv2ForCausalLM(FlashForCausalLM):
2522	input_ids, position_ids, is_prefill, kv_cache,	2599	input_ids, position_ids, is_prefill, kv_cache,
2523	block_tables, slots, input_lengths, max_seq_len,	2600	block_tables, slots, input_lengths, max_seq_len,
2524	lm_head_indices, **kwargs)	2601	lm_head_indices, **kwargs)
2525	- logits, hidden_states = self.execute_ascend_operator(acl_inputs, acl_param, is_prefill)	2602	+ logits, hidden_states = self.execute_ascend_operator(acl_inputs, acl_param, is_prefill,
		2603	+ warmup_is_end=kwargs.get(WARMUP_IS_END, True))
2526	llm_logits, llm_hidden_states = logits, hidden_states[lm_head_indices]	2604	llm_logits, llm_hidden_states = logits, hidden_states[lm_head_indices]
2527	q_lens = kwargs.get(Q_LENS, None)	2605	q_lens = kwargs.get(Q_LENS, None)
2528	logits_mtp, hidden_states_mtp = logits, hidden_states	2606	logits_mtp, hidden_states_mtp = logits, hidden_states
2529	acl_inputs_mtp, acl_param_mtp = acl_inputs, acl_param	2607	acl_inputs_mtp, acl_param_mtp = acl_inputs, acl_param
2530	- self.acl_encoder_operation_mtp.set_kv_cache(self.mtp_k_caches[0: 1],	2608	+ if self.acl_encoder_operation_mtp is not None:
2531	- self.mtp_v_caches[0: 1])	2609	+ self.acl_encoder_operation_mtp.set_kv_cache(self.mtp_k_caches[0: 1],
		2610	+ self.mtp_v_caches[0: 1])
		2611	+ if self.mempool_type == MemPoolType.ASYNC_WRITE:
		2612	+ self.acl_encoder_operation_prefixcache_mempool_mtp.set_kv_cache(self.mtp_k_caches[0: 1],
		2613	+ self.mtp_v_caches[0: 1])
2532	if self.prefix_cache_enable:	2614	if self.prefix_cache_enable:
2533	self.acl_encoder_operation_prefixcache_mtp.set_kv_cache(self.mtp_k_caches[0: 1],	2615	self.acl_encoder_operation_prefixcache_mtp.set_kv_cache(self.mtp_k_caches[0: 1],
2534	self.mtp_v_caches[0: 1])	2616	self.mtp_v_caches[0: 1])
		2617	+ self.acl_encoder_operation_prefixcache_mtp_async_write.set_kv_cache(self.mtp_k_caches[0: 1],
		2618	+ self.mtp_v_caches[0: 1])
2535		2619
2536	acl_inputs_mtp = self.delete_local_tp_mtp_inputs(acl_inputs_mtp)	2620	acl_inputs_mtp = self.delete_local_tp_mtp_inputs(acl_inputs_mtp)
2537	if q_lens is not None:	2621	if q_lens is not None:
	@@ -2543,11 +2627,12 @@ class FlashDeepseekv2ForCausalLM(FlashForCausalLM):
2543	acl_inputs_mtp, acl_param_mtp = \	2627	acl_inputs_mtp, acl_param_mtp = \
2544	self.update_mtp_inputs(logits_mtp, hidden_states_mtp, input_lengths,	2628	self.update_mtp_inputs(logits_mtp, hidden_states_mtp, input_lengths,
2545	acl_inputs_mtp, acl_param_mtp, 0, kwargs)	2629	acl_inputs_mtp, acl_param_mtp, 0, kwargs)
2546	- torch.npu.synchronize()	2630	+ torch.npu.current_stream().synchronize()
2547		2631
2548	acl_inputs_mtp = self.add_local_tp_mtp_inputs(acl_inputs_mtp)	2632	acl_inputs_mtp = self.add_local_tp_mtp_inputs(acl_inputs_mtp)
2549	logits_mtp, hidden_states_mtp = \	2633	logits_mtp, hidden_states_mtp = \
2550	- self.execute_ascend_operator(acl_inputs_mtp, acl_param_mtp, is_prefill, is_mtp=True)	2634	+ self.execute_ascend_operator(acl_inputs_mtp, acl_param_mtp, is_prefill, is_mtp=True,
		2635	+ warmup_is_end=kwargs.get(WARMUP_IS_END, True))
2551		2636
2552	if self.distributed_enable:	2637	if self.distributed_enable:
2553	llm_logits = self.select_logits(llm_logits, **kwargs)	2638	llm_logits = self.select_logits(llm_logits, **kwargs)
	@@ -2903,20 +2988,28 @@ class FlashDeepseekv2ForCausalLM(FlashForCausalLM):
2903	if self.acl_dap_operation is not None:	2988	if self.acl_dap_operation is not None:
2904	self.acl_dap_operation.set_kv_cache(k_caches[:-1],	2989	self.acl_dap_operation.set_kv_cache(k_caches[:-1],
2905	v_caches[:-1])	2990	v_caches[:-1])
		2991	+ if self.mempool_type == MemPoolType.ASYNC_WRITE:
		2992	+ self.acl_encoder_operation_prefixcache_mempool.set_kv_cache(k_caches[:-1],
		2993	+ v_caches[:-1])
2906	if self.prefix_cache_enable:	2994	if self.prefix_cache_enable:
2907	self.acl_encoder_operation_prefixcache.set_kv_cache(k_caches[:-1],	2995	self.acl_encoder_operation_prefixcache.set_kv_cache(k_caches[:-1],
2908	v_caches[:-1])	2996	v_caches[:-1])
		2997	+ self.acl_encoder_operation_prefixcache_async_write.set_kv_cache(k_caches[:-1],
		2998	+ v_caches[:-1])
2909	self.mtp_k_caches = k_caches[-1:]	2999	self.mtp_k_caches = k_caches[-1:]
2910	self.mtp_v_caches = v_caches[-1:]	3000	self.mtp_v_caches = v_caches[-1:]
2911	else:	3001	else:
2912	if self.acl_encoder_operation is not None:	3002	if self.acl_encoder_operation is not None:
2913	self.acl_encoder_operation.set_kv_cache(k_caches, v_caches)	3003	self.acl_encoder_operation.set_kv_cache(k_caches, v_caches)
		3004	+ if self.mempool_type == MemPoolType.ASYNC_WRITE:
		3005	+ self.acl_encoder_operation_prefixcache_mempool.set_kv_cache(k_caches, v_caches)
2914	if self.acl_decoder_operation is not None:	3006	if self.acl_decoder_operation is not None:
2915	self.acl_decoder_operation.set_kv_cache(k_caches, v_caches)	3007	self.acl_decoder_operation.set_kv_cache(k_caches, v_caches)
2916	if self.acl_dap_operation is not None:	3008	if self.acl_dap_operation is not None:
2917	self.acl_dap_operation.set_kv_cache(k_caches, v_caches)	3009	self.acl_dap_operation.set_kv_cache(k_caches, v_caches)
2918	if self.prefix_cache_enable:	3010	if self.prefix_cache_enable:
2919	self.acl_encoder_operation_prefixcache.set_kv_cache(k_caches, v_caches)	3011	self.acl_encoder_operation_prefixcache.set_kv_cache(k_caches, v_caches)
		3012	+ self.acl_encoder_operation_prefixcache_async_write.set_kv_cache(k_caches, v_caches)
2920	else:	3013	else:
2921	if self.layerwise.split_type == DistributedType.EDGE:	3014	if self.layerwise.split_type == DistributedType.EDGE:
2922	self.acl_head_encoder_operation.set_kv_cache(k_caches_sp1, v_caches_sp1)	3015	self.acl_head_encoder_operation.set_kv_cache(k_caches_sp1, v_caches_sp1)

Mexamples/atb_models/atb_llm/models/qwen2/flash_causal_qwen2.py+31-8

	@@ -18,12 +18,13 @@ import torch_npu
18	from atb_llm.utils.data.layer_adapter import ParallelLMHead	18	from atb_llm.utils.data.layer_adapter import ParallelLMHead
19	from atb_llm.utils.data.quant_method_adapter import LinearMethodSupportAtbGraph	19	from atb_llm.utils.data.quant_method_adapter import LinearMethodSupportAtbGraph
20	from mindie_llm.runtime.layers.custom_layer import CustomLayer	20	from mindie_llm.runtime.layers.custom_layer import CustomLayer
		21	+from mindie_llm.text_generator.plugins.plugin_manager import MemPoolType
21	from .modeling_qwen2 import FlashQwenModel	22	from .modeling_qwen2 import FlashQwenModel
22	from .modeling_qwen2_refactor import Qwen2Model	23	from .modeling_qwen2_refactor import Qwen2Model
23	from ..base.flash_causal_lm import FlashForCausalLM, DistributedType, LwdLayerStatus	24	from ..base.flash_causal_lm import FlashForCausalLM, DistributedType, LwdLayerStatus
24	from ..base.graph_manager import ATBGraphManager, DapGraphWrapper, SpeculateGraphWrapper, \	25	from ..base.graph_manager import ATBGraphManager, DapGraphWrapper, SpeculateGraphWrapper, \
25	SplitFuseGraphWrapper, SingleLoraGraphWrapper, MultiLoraGraphWrapper, FlashCommGraphWrapper, \	26	SplitFuseGraphWrapper, SingleLoraGraphWrapper, MultiLoraGraphWrapper, FlashCommGraphWrapper, \
26	- get_layerwise_decode_graph, get_layerwise_prefill_graph	27	+ MemPoolGraphWrapper, get_layerwise_decode_graph, get_layerwise_prefill_graph
27	from ..base.graph_manager.layerwise_combined_graph_wrapper import LayerwiseCombinedATBGraphWrapper	28	from ..base.graph_manager.layerwise_combined_graph_wrapper import LayerwiseCombinedATBGraphWrapper
28	from ..base.inputs_modifier.flash_comm_modifier import FlashCommModifier	29	from ..base.inputs_modifier.flash_comm_modifier import FlashCommModifier
29	from ..base.inputs_modifier.long_seq_modifier import LongSeqModifier	30	from ..base.inputs_modifier.long_seq_modifier import LongSeqModifier
	@@ -182,7 +183,7 @@ class FlashQwen2ForCausalLM(FlashForCausalLM):
182	# 若开启,则冒烟测试卡50ms数据需重新调整(layer多一个输出,内存占用变大)	183	# 若开启,则冒烟测试卡50ms数据需重新调整(layer多一个输出,内存占用变大)
183	self.enable_intra_layer_add_norm = False	184	self.enable_intra_layer_add_norm = False
184	self.enable_inter_layer_add_norm = False	185	self.enable_inter_layer_add_norm = False
185	- self.enable_swiglu_quant = not self.soc_info.need_nz	186	+ self.enable_swiglu_quant = not (self.soc_info.need_nz or self.mempool_type == MemPoolType.ASYNC_WRITE)
186	# Multi engines management	187	# Multi engines management
187	if self.layerwise_disaggregated:	188	if self.layerwise_disaggregated:
188	prefill_graph = get_layerwise_prefill_graph(self.config, self.layerwise)	189	prefill_graph = get_layerwise_prefill_graph(self.config, self.layerwise)
	@@ -227,7 +228,9 @@ class FlashQwen2ForCausalLM(FlashForCausalLM):
227	if self.config.use_qk_norm:	228	if self.config.use_qk_norm:
228	weight_wrapper.register_model_norm(layer.attn.q_norm) # q_norm	229	weight_wrapper.register_model_norm(layer.attn.q_norm) # q_norm
229	weight_wrapper.register_model_norm(layer.attn.k_norm) # k_norm	230	weight_wrapper.register_model_norm(layer.attn.k_norm) # k_norm
230	- if self.enable_intra_layer_add_norm or self.enable_inter_layer_add_norm:	231	+ # not support mempool asyncWrite + add_norm
		232	+ if self.mempool_type != MemPoolType.ASYNC_WRITE and \
		233	+ (self.enable_intra_layer_add_norm or self.enable_inter_layer_add_norm):
231	weight_wrapper.register_layer_addrmsnormquant(layer, attn_wrapper, mlp_wrapper, self.quantize)	234	weight_wrapper.register_layer_addrmsnormquant(layer, attn_wrapper, mlp_wrapper, self.quantize)
232	if self.soc_info.need_nz and self.adapter_manager is None:	235	if self.soc_info.need_nz and self.adapter_manager is None:
233	del layer.attn	236	del layer.attn
	@@ -373,7 +376,9 @@ class FlashQwen2ForCausalLM(FlashForCausalLM):
373	if self.config.use_qk_norm:	376	if self.config.use_qk_norm:
374	weight_wrapper.register_model_norm(layer.attn.q_norm) # q_norm	377	weight_wrapper.register_model_norm(layer.attn.q_norm) # q_norm
375	weight_wrapper.register_model_norm(layer.attn.k_norm) # k_norm	378	weight_wrapper.register_model_norm(layer.attn.k_norm) # k_norm
376	- if self.enable_intra_layer_add_norm or self.enable_inter_layer_add_norm:	379	+ # not support mempool asyncWrite + add_norm
		380	+ if self.mempool_type != MemPoolType.ASYNC_WRITE and \
		381	+ (self.enable_intra_layer_add_norm or self.enable_inter_layer_add_norm):
377	weight_wrapper.register_layer_addrmsnormquant(layer, attn_wrapper, mlp_wrapper, self.quantize)	382	weight_wrapper.register_layer_addrmsnormquant(layer, attn_wrapper, mlp_wrapper, self.quantize)
378	if self.soc_info.need_nz and self.adapter_manager is None:	383	if self.soc_info.need_nz and self.adapter_manager is None:
379	del layer.attn	384	del layer.attn
	@@ -485,8 +490,10 @@ class FlashQwen2ForCausalLM(FlashForCausalLM):
485	LINEAR_HAS_BIAS: linear_has_bias * self.config.num_hidden_layers	490	LINEAR_HAS_BIAS: linear_has_bias * self.config.num_hidden_layers
486	if not self.layerwise_disaggregated else None,	491	if not self.layerwise_disaggregated else None,
487	"matmulBackend": OpBackend.ACLNN if self.aclnn_matmul_backend else OpBackend.ATB,	492	"matmulBackend": OpBackend.ACLNN if self.aclnn_matmul_backend else OpBackend.ATB,
488	- "enableIntraLayerAddNorm": self.enable_intra_layer_add_norm,	493	+ "enableIntraLayerAddNorm": self.enable_intra_layer_add_norm and \
489	- "enableInterLayerAddNorm": self.enable_inter_layer_add_norm,	494	+ self.mempool_type != MemPoolType.ASYNC_WRITE,
		495	+ "enableInterLayerAddNorm": self.enable_inter_layer_add_norm and \
		496	+ self.mempool_type != MemPoolType.ASYNC_WRITE,
490	"enableGreedySearchOpt": self.enable_greedy_search_opt,	497	"enableGreedySearchOpt": self.enable_greedy_search_opt,
491	"enableOmniAttention": self.omni_attention_enable,	498	"enableOmniAttention": self.omni_attention_enable,
492	"enableQScale": (self.config.transformers_version == "4.43.1" or	499	"enableQScale": (self.config.transformers_version == "4.43.1" or
	@@ -520,6 +527,16 @@ class FlashQwen2ForCausalLM(FlashForCausalLM):
520	}	527	}
521		528
522	if not self.layerwise_disaggregated:	529	if not self.layerwise_disaggregated:
		530	+ #Mooncake池化与lora、dap、flashcomm不适配
		531	+ if self.adapter_manager is not None and self.mempool_type == MemPoolType.ASYNC_WRITE:
		532	+ raise ValueError("Feature composition not supported: If lora is activated, "
		533	+ "mempool_type must be DISABLED or SYNC_WRITE.")
		534	+ if self.enable_dap and self.mempool_type == MemPoolType.ASYNC_WRITE:
		535	+ raise ValueError("Feature composition not supported: If dap is activated, "
		536	+ "mempool_type must be DISABLED or SYNC_WRITE.")
		537	+ if self.mempool_type == MemPoolType.ASYNC_WRITE:
		538	+ self.flash_comm_modifier.enable_flash_comm = False
		539	+
523	if self.adapter_manager is not None:	540	if self.adapter_manager is not None:
524	self.graph_manager.register_graph(MultiLoraGraphWrapper())	541	self.graph_manager.register_graph(MultiLoraGraphWrapper())
525	self.graph_manager.register_graph(SingleLoraGraphWrapper())	542	self.graph_manager.register_graph(SingleLoraGraphWrapper())
	@@ -536,6 +553,10 @@ class FlashQwen2ForCausalLM(FlashForCausalLM):
536	if self.flash_comm_modifier.enable_flash_comm:	553	if self.flash_comm_modifier.enable_flash_comm:
537	self.graph_manager.register_graph(FlashCommGraphWrapper())	554	self.graph_manager.register_graph(FlashCommGraphWrapper())
538		555
		556	+ #Mooncake池化
		557	+ if self.mempool_type == MemPoolType.ASYNC_WRITE:
		558	+ self.graph_manager.register_graph(MemPoolGraphWrapper())
		559	+
539	specified_params = {"decode": decoder_param}	560	specified_params = {"decode": decoder_param}
540	specified_weight = {"decode": self.decode_weight}	561	specified_weight = {"decode": self.decode_weight}
541	self.graph_manager.set_param(CPP_QWEN_MODEL_CLASS_NAME, encoder_param, specified_params)	562	self.graph_manager.set_param(CPP_QWEN_MODEL_CLASS_NAME, encoder_param, specified_params)
	@@ -734,9 +755,10 @@ class FlashQwen2ForCausalLM(FlashForCausalLM):
734	acl_param,	755	acl_param,
735	is_prefill, **kwargs):	756	is_prefill, **kwargs):
736	exe_stage = kwargs.get("layerwise_disaggregated_exe_stage", None)	757	exe_stage = kwargs.get("layerwise_disaggregated_exe_stage", None)
		758	+ runtime_mempool_type = self.mempool_type if self.warmup_is_end else MemPoolType.DISABLED
737	acl_model_out = self.graph_manager.select_and_execute(	759	acl_model_out = self.graph_manager.select_and_execute(
738	- self, acl_inputs, acl_param, is_prefill=is_prefill, layerwise_disaggregated_exe_stage=exe_stage	760	+ self, acl_inputs, acl_param, is_prefill=is_prefill, layerwise_disaggregated_exe_stage=exe_stage,
739	- )	761	+ mempool_type=runtime_mempool_type)
740	try:	762	try:
741	acl_model_out = self.layerwise_modifier.process_out(acl_model_out, is_prefill=is_prefill, **kwargs)	763	acl_model_out = self.layerwise_modifier.process_out(acl_model_out, is_prefill=is_prefill, **kwargs)
742	acl_hidden_state = acl_model_out[0]	764	acl_hidden_state = acl_model_out[0]
	@@ -855,6 +877,7 @@ class FlashQwen2ForCausalLM(FlashForCausalLM):
855	Returns:	877	Returns:
856	torch.Tensor: Output logits.	878	torch.Tensor: Output logits.
857	"""	879	"""
		880	+ self.warmup_is_end = kwargs.get("warmup_is_end", True)
858	if not self.weight_initialized:	881	if not self.weight_initialized:
859	self.get_adapter_ids(**kwargs)	882	self.get_adapter_ids(**kwargs)
860	from mindie_llm.runtime.utils.torch_utils import set_default_torch_dtype	883	from mindie_llm.runtime.utils.torch_utils import set_default_torch_dtype

Mexamples/atb_models/atb_llm/runner/model_runner.py+18-5

	@@ -25,6 +25,7 @@ import atb_llm.nn.distributed as dist
25	from atb_llm.nn.network_manager import get_default_net	25	from atb_llm.nn.network_manager import get_default_net
26	from atb_llm.nn.tensor import Tensor	26	from atb_llm.nn.tensor import Tensor
27	from atb_llm.nn.functional import gather, split	27	from atb_llm.nn.functional import gather, split
		28	+from mindie_llm.text_generator.plugins.plugin_manager import MemPoolType
28		29
29	from ..models import get_model	30	from ..models import get_model
30	from ..models.base.config import LoraModelConfig	31	from ..models.base.config import LoraModelConfig
	@@ -80,13 +81,18 @@ class TruncationSide(int, Enum):
80	RIGHT = -1	81	RIGHT = -1
81		82
82		83
		84	+# 专用于mempool异步分层传输特性的event pipe_key
		85	+def generate_mem_pool_event_key(only_save_kv: bool) -> str:
		86	+ return "only_save_kv" if only_save_kv else "both_save_kv"
		87	+
		88	+
83	def exception_handler(cls):	89	def exception_handler(cls):
84	"""	90	"""
85	Class decorator for ModelRunner that applies various handlers to methods.	91	Class decorator for ModelRunner that applies various handlers to methods.
86	Currently applies:	92	Currently applies:
87	1. _torch_oom_handler: Catches and logs PyTorch OOM errors.	93	1. _torch_oom_handler: Catches and logs PyTorch OOM errors.
88	"""	94	"""
89	-	95	+
90	def _torch_oom_handler(func):	96	def _torch_oom_handler(func):
91	"""Handler specifically for PyTorch OOM errors."""	97	"""Handler specifically for PyTorch OOM errors."""
92	@wraps(func)	98	@wraps(func)
	@@ -107,11 +113,11 @@ def exception_handler(cls):
107	raise RuntimeError(f"{error_msg}. Error_code: {error_code}") from e	113	raise RuntimeError(f"{error_msg}. Error_code: {error_code}") from e
108	raise	114	raise
109	return wrapper	115	return wrapper
110	-	116	+
111	def _apply_handlers(func):	117	def _apply_handlers(func):
112	"""Apply the chain of handlers to a function."""	118	"""Apply the chain of handlers to a function."""
113	return _torch_oom_handler(func)	119	return _torch_oom_handler(func)
114	-	120	+
115	def _is_target_method(name):	121	def _is_target_method(name):
116	"""Filter methods that need handling."""	122	"""Filter methods that need handling."""
117	if name == "generate_position_ids":	123	if name == "generate_position_ids":
	@@ -225,6 +231,7 @@ class ModelRunner:
225	except Exception:	231	except Exception:
226	print_log(rank, logger.info, "deserialized unknow error")	232	print_log(rank, logger.info, "deserialized unknow error")
227		233
		234	+ self.mempool_type = kwargs.get('mempool_type', MemPoolType.DISABLED)
228	load_atb_speed()	235	load_atb_speed()
229		236
230	if ENV.bind_cpu:	237	if ENV.bind_cpu:
	@@ -370,6 +377,10 @@ class ModelRunner:
370	def resume_hccl_comm(cls):	377	def resume_hccl_comm(cls):
371	torch.classes.ModelTorch.Context.resume_hccl_comm()	378	torch.classes.ModelTorch.Context.resume_hccl_comm()
372		379
		380	+ @classmethod
		381	+ def wait_event(cls, pipe_key: str):
		382	+ torch.classes.ModelTorch.Event.wait(pipe_key)
		383	+
373	def load_weights(self, **kwargs):	384	def load_weights(self, **kwargs):
374	"""Load weights from file."""	385	"""Load weights from file."""
375	enable_v3 = False	386	enable_v3 = False
	@@ -418,9 +429,11 @@ class ModelRunner:
418	with self.device:	429	with self.device:
419	self.init_model(self.config, weights,	430	self.init_model(self.config, weights,
420	quant_config=mindie_llm_config_v2.quant_config,	431	quant_config=mindie_llm_config_v2.quant_config,
421	- prealloc_weight_mem_on_npu=True)	432	+ prealloc_weight_mem_on_npu=True,
		433	+ mempool_type=self.mempool_type)
422	else:	434	else:
423	- self.init_model(self.config if not enable_v3 else mindie_llm_config, weights)	435	+ self.init_model(self.config if not enable_v3 else mindie_llm_config, weights,
		436	+ mempool_type=self.mempool_type)
424		437
425	except TypeError as e:	438	except TypeError as e:
426	logger.warning(	439	logger.warning(

Mexamples/atb_models/tests/pythontest/atb_llm/models/base/graph_manager/test_graph_wrapper.py+10-1

	@@ -22,6 +22,8 @@ from atb_llm.models.base.graph_manager.single_lora_graph_wrapper import SingleLo
22	from atb_llm.models.base.graph_manager.multi_lora_graph_wrapper import MultiLoraGraphWrapper	22	from atb_llm.models.base.graph_manager.multi_lora_graph_wrapper import MultiLoraGraphWrapper
23	from atb_llm.models.base.graph_manager.speculate_graph_wrapper import SpeculateGraphWrapper	23	from atb_llm.models.base.graph_manager.speculate_graph_wrapper import SpeculateGraphWrapper
24	from atb_llm.models.base.graph_manager.splitfuse_graph_wrapper import SplitFuseGraphWrapper	24	from atb_llm.models.base.graph_manager.splitfuse_graph_wrapper import SplitFuseGraphWrapper
		25	+from atb_llm.models.base.graph_manager.mem_pool_graph_wrapper import MemPoolGraphWrapper
		26	+from mindie_llm.text_generator.plugins.plugin_manager import MemPoolType
25	from tests.pythontest.atb_llm.models.base.graph_manager.test_graph_manager import MockATBGraphWrapper	27	from tests.pythontest.atb_llm.models.base.graph_manager.test_graph_manager import MockATBGraphWrapper
26		28
27		29
	@@ -121,4 +123,11 @@ class TestATBGraphWrapper(unittest.TestCase):
121	self.assertFalse(graph_wrapper.activate(mock_context, json.dumps({"seqLen": 19}), is_prefill=True))	123	self.assertFalse(graph_wrapper.activate(mock_context, json.dumps({"seqLen": 19}), is_prefill=True))
122		124
123	mock_context.inference_mode = None	125	mock_context.inference_mode = None
124	- self.assertFalse(graph_wrapper.activate(mock_context, json.dumps({"qLen": 17}), is_prefill=True))	126	+ self.assertFalse(graph_wrapper.activate(mock_context, json.dumps({"qLen": 17}), is_prefill=True))
		127	+
		128	+ def test_mem_pool(self):
		129	+ graph_wrapper = MemPoolGraphWrapper()
		130	+ mock_context = MagicMock()
		131	+ self.assertTrue(graph_wrapper.activate(mock_context, {}, mempool_type=MemPoolType.ASYNC_WRITE))
		132	+ self.assertFalse(graph_wrapper.activate(mock_context, {}, mempool_type=MemPoolType.DISABLED))
		133	+ self.assertFalse(graph_wrapper.activate(mock_context, {}, mempool_type=MemPoolType.SYNC_WRITE))

Aexamples/atb_models/tests/pythontest/atb_llm/models/base/graph_manager/test_mempool_combined_graph_wrapper.py+174-0

@@ -0,0 +1,174 @@
	1	+# Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
	2	+# MindIE is licensed under Mulan PSL v2.
	3	+# You can use this software according to the terms and conditions of the Mulan PSL v2.
	4	+# You may obtain a copy of Mulan PSL v2 at:
	5	+# http://license.coscl.org.cn/MulanPSL2
	6	+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
	7	+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
	8	+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
	9	+# See the Mulan PSL v2 for more details.
	10	+
	11	+import unittest
	12	+from unittest.mock import MagicMock, patch
	13	+
	14	+import json
	15	+import torch
	16	+
	17	+from atb_llm.models.qwen2.config_qwen2 import Qwen2Config
	18	+from atb_llm.models.qwen2.flash_causal_qwen2 import FlashQwen2ForCausalLM
	19	+from atb_llm.utils.mapping import Mapping
	20	+from atb_llm.utils.quantize.pack_type import TransposeType
	21	+
	22	+from mindie_llm.runtime.utils.distributed import set_parallel_info_manager
	23	+from mindie_llm.text_generator.plugins.plugin_manager import MemPoolType
	24	+from tests.pythontest.atb_llm.models.base.mock_class import MockTorchClasses
	25	+
	26	+
	27	+LOAD_ATB_SPEED = "atb_llm.models.base.flash_causal_lm.load_atb_speed"
	28	+FLASH_QWEN2 = "atb_llm.models.qwen2.flash_causal_qwen2"
	29	+
	30	+
	31	+class TestMempoolCombinedATBGraphWrapper(unittest.TestCase):
	32	+ def setUp(self) -> None:
	33	+ self.torch_classes = MockTorchClasses()
	34	+ torch.classes = self.torch_classes
	35	+
	36	+ self.config = Qwen2Config(
	37	+ model_type="qwen2",
	38	+ hidden_size=1024,
	39	+ max_position_embeddings=1024,
	40	+ num_attention_heads=16,
	41	+ num_key_value_heads=4,
	42	+ num_hidden_layers=28,
	43	+ rms_norm_eps=1e-6,
	44	+ torch_dtype=torch.float16,
	45	+ vocab_size=125696,
	46	+ tie_word_embeddings=False,
	47	+ )
	48	+
	49	+ self.weights = MagicMock()
	50	+ self.weights.device = torch.device("npu")
	51	+ self.weights.dtype = torch.float16
	52	+ self.weights.mapping = Mapping(world_size=2, rank=0)
	53	+ self.weights.mapping.attn_tp.rank = 1
	54	+ self.weights.process_group = MagicMock()
	55	+ self.weights.process_group.rank.return_value = 0
	56	+ self.weights.process_group.size.return_value = 2
	57	+ self.weights.quant_desc = None
	58	+
	59	+ # Create mock parallel info
	60	+ self.mock_parallel_info = MagicMock()
	61	+ self.mock_parallel_info.rank = 0
	62	+ self.mock_parallel_info.group_size = 2
	63	+ self.mock_parallel_info.process_group = None
	64	+
	65	+ # Create mock parallel info manager
	66	+ self.mock_parallel_info_manager = MagicMock()
	67	+ self.mock_parallel_info_manager.world_size = 2
	68	+ self.mock_parallel_info_manager.word_embed_tp = self.mock_parallel_info
	69	+ self.mock_parallel_info_manager.attn_tp = self.mock_parallel_info
	70	+ self.mock_parallel_info_manager.lm_head_tp = self.mock_parallel_info
	71	+
	72	+ # Set the global parallel info manager
	73	+ set_parallel_info_manager(self.mock_parallel_info_manager)
	74	+
	75	+ def tearDown(self):
	76	+ """Clean up after tests."""
	77	+ set_parallel_info_manager(None)
	78	+
	79	+ @patch(f"{LOAD_ATB_SPEED}")
	80	+ @patch(f"{FLASH_QWEN2}.FlashQwenModel", return_value=MagicMock())
	81	+ @patch(f"{FLASH_QWEN2}.Qwen2Model", return_value=MagicMock())
	82	+ @patch(f"{FLASH_QWEN2}.load_column_multi")
	83	+ @patch(f"{FLASH_QWEN2}.TensorHead")
	84	+ def test_init(
	85	+ self,
	86	+ mock_tensor_head,
	87	+ mock_load_column_multi,
	88	+ mock_new_qwen_model,
	89	+ mock_qwen_model,
	90	+ _mock_init_so
	91	+ ) -> None:
	92	+ FlashQwen2ForCausalLM(self.config, self.weights, prealloc_weight_mem_on_npu=True)
	93	+ mock_new_qwen_model.assert_called_once_with(
	94	+ self.config, "model", quant_config=None
	95	+ )
	96	+
	97	+ self.config.quantize = "w8a8sc"
	98	+ FlashQwen2ForCausalLM(self.config, self.weights)
	99	+ mock_tensor_head.load_weight.assert_called_once_with(
	100	+ self.config,
	101	+ prefix="lm_head",
	102	+ weights=self.weights,
	103	+ is_norm=False,
	104	+ )
	105	+
	106	+ self.config.quantize = ""
	107	+ self.config.tie_word_embeddings = True
	108	+ FlashQwen2ForCausalLM(self.config, self.weights)
	109	+ mock_load_column_multi.assert_called_with(
	110	+ self.config,
	111	+ prefixes=["model.embed_tokens"],
	112	+ weights=self.weights,
	113	+ head_size=1,
	114	+ lm_head=True
	115	+ )
	116	+
	117	+ @patch(f"{LOAD_ATB_SPEED}")
	118	+ @patch(f"{FLASH_QWEN2}.FlashQwenModel", return_value=MagicMock())
	119	+ @patch(f"{FLASH_QWEN2}.load_column_multi")
	120	+ @patch(f"{FLASH_QWEN2}.WeightWrapper")
	121	+ def test_init_ascend_weight(
	122	+ self, mock_weight_wrapper, _mock_load_column_multi, _mock_qwen_model, _mock_init_so):
	123	+ mock_weight_wrapper_ins = mock_weight_wrapper.return_value
	124	+ mock_weight_wrapper_ins.register_embedding = MagicMock()
	125	+ mock_weight_wrapper_ins.register_layer = MagicMock()
	126	+ mock_weight_wrapper_ins.register_model_norm = MagicMock()
	127	+ mock_weight_wrapper_ins.register_model_lmhead = MagicMock()
	128	+ mock_weight_wrapper_ins.weights = []
	129	+ mock_weight_wrapper_ins.linear_type = {}
	130	+ mock_weight_wrapper_ins.pack_quant_type = {}
	131	+ mock_weight_wrapper_ins.linear_transpose_types = {}
	132	+ mock_weight_wrapper_ins.linear_descs = []
	133	+
	134	+ # prefill_splitfuse
	135	+ ins = FlashQwen2ForCausalLM(self.config, self.weights, prealloc_weight_mem_on_npu=True)
	136	+ ins.lm_head.get_weight_transpose_type()[0] = TransposeType.NOT_TRANSPOSE
	137	+ ins.prefix_cache_enable = True
	138	+ ins.mempool_type = MemPoolType.ASYNC_WRITE
	139	+ ins.init_ascend_weight()
	140	+ mock_context = MagicMock()
	141	+ mock_context.inference_mode = MagicMock()
	142	+ mock_context.inference_mode.enable_prefill_pa = True
	143	+ selected_graph = None
	144	+ for graph in ins.graph_manager._graph_list:
	145	+ if graph.activate(mock_context, json.dumps({"qLen": 9}), is_prefill=True, mempool_type=MemPoolType.ASYNC_WRITE):
	146	+ selected_graph = graph
	147	+ break
	148	+ self.assertEqual(selected_graph.feature_name, "prefill_splitfuse")
	149	+
	150	+ # prefill_mempool
	151	+ ins = FlashQwen2ForCausalLM(self.config, self.weights, prealloc_weight_mem_on_npu=True)
	152	+ ins.lm_head.get_weight_transpose_type()[0] = TransposeType.NOT_TRANSPOSE
	153	+ ins.mempool_type = MemPoolType.ASYNC_WRITE
	154	+ ins.init_ascend_weight()
	155	+ mock_context = MagicMock()
	156	+ selected_graph = None
	157	+ for graph in ins.graph_manager._graph_list:
	158	+ if graph.activate(mock_context, {}, is_prefill=True, mempool_type=MemPoolType.ASYNC_WRITE):
	159	+ selected_graph = graph
	160	+ break
	161	+ self.assertEqual(selected_graph.feature_name, "prefill_mempool")
	162	+
	163	+ # base_prefill
	164	+ ins = FlashQwen2ForCausalLM(self.config, self.weights, prealloc_weight_mem_on_npu=True)
	165	+ ins.lm_head.get_weight_transpose_type()[0] = TransposeType.NOT_TRANSPOSE
	166	+ ins.mempool_type = MemPoolType.SYNC_WRITE
	167	+ ins.init_ascend_weight()
	168	+ mock_context = MagicMock()
	169	+ selected_graph = None
	170	+ for graph in ins.graph_manager._graph_list:
	171	+ if graph.activate(mock_context, {}, is_prefill=True, mempool_type=MemPoolType.SYNC_WRITE):
	172	+ selected_graph = graph
	173	+ break
	174	+ self.assertEqual(str(selected_graph.feature_name), "FeatureType.PREFILL")

Mexamples/atb_models/tests/run_all_tests.sh+1-1

	@@ -157,7 +157,7 @@ function fn_build_coverage()
157	cd $GCOV_DIR	157	cd $GCOV_DIR
158	$LCOV_PATH -c -d $CACHE_DIR -o $GCOV_INFO_DIR/cover.info --rc lcov_branch_coverage=1 >> $GCOV_DIR/log.txt --rc lcov_excl_br_line='ATB_SPEED_LOG_*'	158	$LCOV_PATH -c -d $CACHE_DIR -o $GCOV_INFO_DIR/cover.info --rc lcov_branch_coverage=1 >> $GCOV_DIR/log.txt --rc lcov_excl_br_line='ATB_SPEED_LOG_*'
159	$LCOV_PATH -a $GCOV_INFO_DIR/init.info -a $GCOV_INFO_DIR/cover.info -o $GCOV_INFO_DIR/total.info --rc lcov_branch_coverage=1 >> $GCOV_DIR/log.txt	159	$LCOV_PATH -a $GCOV_INFO_DIR/init.info -a $GCOV_INFO_DIR/cover.info -o $GCOV_INFO_DIR/total.info --rc lcov_branch_coverage=1 >> $GCOV_DIR/log.txt
160	- $LCOV_PATH --remove $GCOV_INFO_DIR/total.info '/third_party/' 'torch/' 'c10/' 'ATen/' '/c++/7' 'tests/' '/include/utils/' '/utils/log/' 'tools/' '/usr/' 'ascend-transformer-boost/' 'qkv_linear_split' 'dequant_bias_operation' 'w8a8_operation' 'linear' 'model.cpp' '/moe/layer/decoder_layer' '/moe/model/decoder_model' '/sparse_moe' '/obfuscation_setup_operation.cpp' '/obfuscation_calculate_operation.cpp' '/deepseekv2/' '/glm/' '/mapping' '/fusion/moe/ep/' 'add_rms_norm_' 'check_util.cpp' '/base/layer/decoder_layer' 'base/param/model_param' 'base/param/param' '/concat_operation.cpp' 'mlp' '/qwen//moe_decoder' '/core/base/context_factory.cpp' '/fusion_attention.cpp' '/rotary_pos_emb_operation.cpp' '/attn_v3_operation.cpp' '/self_attention.cpp' '/operations/aclrt/ops/aclrt_cmo_async.cpp' '/operations/aclnn/core/acl_nn_operation.cpp' '/len_operation.cpp' '/repeat_operation.cpp' '/inplacemasked_filltensor_operation.cpp' '/decoder_model_edge.cpp' 'edge' '/operations/aclnn/core/acl_nn_global_cache.cpp' '/operations/aclnn/ops/' '/operations/aclnn/utils/' '/fusion/utils.cpp' -o $GCOV_INFO_DIR/final.info --rc lcov_branch_coverage=1 >> $GCOV_DIR/log.txt	160	+ $LCOV_PATH --remove $GCOV_INFO_DIR/total.info '/third_party/' 'torch/' 'c10/' 'ATen/' '/c++/7' 'tests/' '/include/utils/' '/utils/log/' 'tools/' '/usr/' 'ascend-transformer-boost/' 'qkv_linear_split' 'dequant_bias_operation' 'w8a8_operation' 'linear' 'model.cpp' '/moe/layer/decoder_layer' '/moe/model/decoder_model' '/sparse_moe' '/obfuscation_setup_operation.cpp' '/obfuscation_calculate_operation.cpp' '/deepseekv2/' '/glm/' '/mapping' '/fusion/moe/ep/' 'add_rms_norm_' 'check_util.cpp' '/base/layer/decoder_layer' 'base/param/model_param' 'base/param/param' '/concat_operation.cpp' 'mlp' '/qwen//moe_decoder' '/core/base/context_factory.cpp' '/fusion_attention.cpp' '/rotary_pos_emb_operation.cpp' '/attn_v3_operation.cpp' '/self_attention.cpp' '/operations/aclrt/ops/aclrt_cmo_async.cpp' '/operations/aclnn/core/acl_nn_operation.cpp' '/len_operation.cpp' '/repeat_operation.cpp' '/inplacemasked_filltensor_operation.cpp' '/decoder_model_edge.cpp' 'edge' '/operations/aclnn/core/acl_nn_global_cache.cpp' '/operations/aclnn/ops/' '/operations/aclnn/utils/' '/fusion/utils.cpp' 'event_manager.cpp' -o $GCOV_INFO_DIR/final.info --rc lcov_branch_coverage=1 >> $GCOV_DIR/log.txt
161	$GENHTML_PATH --rc lcov_branch_coverage=1 -o cover_result $GCOV_INFO_DIR/final.info -o cover_result >> $GCOV_DIR/log.txt	161	$GENHTML_PATH --rc lcov_branch_coverage=1 -o cover_result $GCOV_INFO_DIR/final.info -o cover_result >> $GCOV_DIR/log.txt
162	tail -n 4 $GCOV_DIR/log.txt	162	tail -n 4 $GCOV_DIR/log.txt
163	cd $OUTPUT_DIR	163	cd $OUTPUT_DIR

Mexamples/atb_models/tests/unittest/core/event_manager_test.cpp+39-0

	@@ -334,4 +334,43 @@ TEST(EventManager, RecordEventInvalidActionTest)
334	EXPECT_EQ(op, nullptr);	334	EXPECT_EQ(op, nullptr);
335	}	335	}
336		336
		337	+// 测试CheckPipeKey函数
		338	+TEST(EventManager, CheckPipeKeyTest)
		339	+{
		340	+ EventManager& manager = EventManager::GetInstance();
		341	+ const std::string pipeKey = "test_pipe_1";
		342	+ aclrtStream subStream;
		343	+ std::vector<aclrtEvent> queue;
		344	+ manager.eventsForExternal_[pipeKey] = std::make_tuple(0, queue, subStream);
		345	+ EXPECT_EQ(manager.CheckPipeKey(pipeKey), EM_SUCCESS);
		346	+ aclrtEvent event = nullptr;
		347	+ manager.CreateAndPushEvent(event, pipeKey);
		348	+ EXPECT_EQ(manager.CheckPipeKey(pipeKey), EM_SUCCESS);
		349	+ EXPECT_TRUE(manager.eventsForExternal_.count(pipeKey) > 0);
		350	+
		351	+ // 验证数据结构是否正确初始化
		352	+ auto& tuple = manager.eventsForExternal_[pipeKey];
		353	+ const std::string pipeKey2 = "test_pipe_2";
		354	+ EXPECT_EQ(manager.CheckPipeKey(pipeKey2), EM_INVALID_ACTION);
		355	+}
		356	+
		357	+// 测试RecordEvent函数
		358	+TEST(EventManager, RecordEventTest) {
		359	+ EventManager& manager = EventManager::GetInstance();
		360	+ const std::string pipeKey = "record_normal";
		361	+ aclrtEvent event = nullptr;
		362	+ manager.CreateAndPushEvent(event, pipeKey);
		363	+ EXPECT_EQ(manager.RecordEvent(pipeKey), EM_SUCCESS);
		364	+}
		365	+
		366	+// 测试RecordEvent函数
		367	+TEST(EventManager, WaitEventTest) {
		368	+ EventManager& manager = EventManager::GetInstance();
		369	+ const std::string pipeKey = "wait_normal";
		370	+ aclrtEvent event = nullptr;
		371	+ manager.CreateAndPushEvent(event, pipeKey);
		372	+ EXPECT_EQ(manager.RecordEvent(pipeKey), EM_SUCCESS);
		373	+ EXPECT_EQ(manager.WaitEvent(pipeKey), EM_SUCCESS);
		374	+}
		375	+
337	} // namespace atb_speed	376	} // namespace atb_speed

Mmindie_llm/modeling/model_wrapper/atb/atb_model_wrapper.py+8-3

	@@ -14,13 +14,14 @@ import torch
14	import numpy as np	14	import numpy as np
15		15
16	from atb_llm.models import InferenceMode	16	from atb_llm.models import InferenceMode
17	-from atb_llm.runner.model_runner import ModelRunner	17	+from atb_llm.runner.model_runner import ModelRunner, generate_mem_pool_event_key
18	from atb_llm.utils.env import ENV	18	from atb_llm.utils.env import ENV
19	from atb_llm.utils.eplb_expert_data_collect import EplbExpertDataCollect	19	from atb_llm.utils.eplb_expert_data_collect import EplbExpertDataCollect
20	from atb_llm.utils.moe_utils import EPLBType, save_eplb_data	20	from atb_llm.utils.moe_utils import EPLBType, save_eplb_data
21	from ..model_info import ModelInfo	21	from ..model_info import ModelInfo
22	from ..wrapper import ModelWrapper	22	from ..wrapper import ModelWrapper
23	from ....utils.log.logging import logger	23	from ....utils.log.logging import logger
		24	+from ....text_generator.plugins.plugin_manager import MemPoolType
24		25
25	ASCEND_310B = 240	26	ASCEND_310B = 240
26		27
	@@ -83,7 +84,8 @@ class ATBModelWrapper(ModelWrapper):
83	tls_crl_path=kwargs.get("interNodeTlsCrlPath", ""),	84	tls_crl_path=kwargs.get("interNodeTlsCrlPath", ""),
84	tls_crl_files=kwargs.get("interNodeTlsCrlFiles", ""),	85	tls_crl_files=kwargs.get("interNodeTlsCrlFiles", ""),
85	batch_p_num=2 if kwargs.get('lwdNextPHeadPrior', False) else 1,	86	batch_p_num=2 if kwargs.get('lwdNextPHeadPrior', False) else 1,
86	- lwd_comm_args=kwargs.get('lwd_comm_args', None)	87	+ lwd_comm_args=kwargs.get('lwd_comm_args', None),
		88	+ mempool_type=kwargs.get('mempool_type', MemPoolType.DISABLED)
87	)	89	)
88	self.config = self.model_runner.config	90	self.config = self.model_runner.config
89	self.config_dict = self.model_runner.config_dict	91	self.config_dict = self.model_runner.config_dict
	@@ -438,4 +440,7 @@ class ATBModelWrapper(ModelWrapper):
438	return context	440	return context
439		441
440	def resume_hccl_comm(self):	442	def resume_hccl_comm(self):
441	- self.model_runner.resume_hccl_comm()	443	+ self.model_runner.resume_hccl_comm()
		444	+
		445	+ def generate_mem_pool_event_key(self, only_save_kv: bool) -> str:
		446	+ return generate_mem_pool_event_key(only_save_kv)

Mmindie_llm/text_generator/adapter/generator_backend.py+13-5

Mmindie_llm/text_generator/adapter/generator_torch.py+5-0

	@@ -222,6 +222,11 @@ class GeneratorTorch(GeneratorBackend):
222	return logits	222	return logits
223		223
224	def update_cache_policy(self, kvcache_settings, sepd_worker=None):	224	def update_cache_policy(self, kvcache_settings, sepd_worker=None):
		225	+ if hasattr(self, 'cache_pool') and self.cache_pool is not None:
		226	+ del self.cache_pool
		227	+ torch.npu.empty_cache()
		228	+ gc.collect()
		229	+
225	self.cache_pool = KVCachePool(kvcache_settings, self.device, enable_kv_pool=self.enable_kv_pool)	230	self.cache_pool = KVCachePool(kvcache_settings, self.device, enable_kv_pool=self.enable_kv_pool)
226	self.cache_pool.allocate_cpu_kvcache()	231	self.cache_pool.allocate_cpu_kvcache()
227	self.cache_pool.allocate_npu_kvcache()	232	self.cache_pool.allocate_npu_kvcache()

Mmindie_llm/text_generator/adapter/torch_utils/kvcache_pool.py+5-0

	@@ -87,6 +87,11 @@ class KVCachePool:
87		87
88	def allocate_npu_kvcache(self):	88	def allocate_npu_kvcache(self):
89	self.npu_cache.clear()	89	self.npu_cache.clear()
		90	+ self.npu_blocks_addrs.clear()
		91	+ self.k_blocks_addrs.clear()
		92	+ self.v_blocks_addrs.clear()
		93	+ self.k_blocks_quant_addrs.clear()
		94	+ self.index_blocks_addrs.clear()
90	if self.kvcache_settings.num_npu_blocks < 0:	95	if self.kvcache_settings.num_npu_blocks < 0:
91	message = (	96	message = (
92	"Num_npu_blocks must be non-negative.\n"	97	"Num_npu_blocks must be non-negative.\n"

Mmindie_llm/text_generator/generator.py+13-2

	@@ -256,6 +256,8 @@ class Generator(PDInterface):
256	self.world_size = parse_config(model_config, 'world_size', required=True, parse_type=ParseType.TO_INT)	256	self.world_size = parse_config(model_config, 'world_size', required=True, parse_type=ParseType.TO_INT)
257	self.local_rank = parse_config(model_config, 'local_rank', required=True, parse_type=ParseType.TO_INT)	257	self.local_rank = parse_config(model_config, 'local_rank', required=True, parse_type=ParseType.TO_INT)
258	self.npu_device_id = parse_config(model_config, 'npu_device_id', required=True, parse_type=ParseType.TO_INT)	258	self.npu_device_id = parse_config(model_config, 'npu_device_id', required=True, parse_type=ParseType.TO_INT)
		259	+ kv_pool_async_write = parse_config(model_config, 'kv_pool_async_write', required=False,
		260	+ parse_type=ParseType.TO_BOOL, default_value=False)
259		261
260	async_inference_key = 'async_inference'	262	async_inference_key = 'async_inference'
261	model_config[async_inference_key] = ENV.async_inference	263	model_config[async_inference_key] = ENV.async_inference
	@@ -265,6 +267,9 @@ class Generator(PDInterface):
265	validator.check_async_inference_and_plugin_type(True, plugin_config.get("plugin_type"))	267	validator.check_async_inference_and_plugin_type(True, plugin_config.get("plugin_type"))
266	model_config['splitfuse_enabled'] = self.is_mix_model	268	model_config['splitfuse_enabled'] = self.is_mix_model
267		269
		270	+ if kv_pool_async_write and "splitfuse" in model_config.get("plugin_params", ""):
		271	+ raise ValueError("Async mempool does not support plugin_type: splitfuse!")
		272	+
268	self.layerwise_disaggregated = parse_config(model_config, 'layerwiseDisaggregated', required=False,	273	self.layerwise_disaggregated = parse_config(model_config, 'layerwiseDisaggregated', required=False,
269	parse_type=ParseType.TO_BOOL, default_value=False)	274	parse_type=ParseType.TO_BOOL, default_value=False)
270	self.layerwise_disaggregated_role_type = parse_config(model_config, 'layerwiseDisaggregatedRoleType',	275	self.layerwise_disaggregated_role_type = parse_config(model_config, 'layerwiseDisaggregatedRoleType',
	@@ -389,7 +394,8 @@ class Generator(PDInterface):
389		394
390	# NOTE: Warmup async inference will lead to lower mtp acceptance rate with unknown reason,	395	# NOTE: Warmup async inference will lead to lower mtp acceptance rate with unknown reason,
391	# so we disable async inference here.	396	# so we disable async inference here.
392	- with self._temporarily_disable(async_inference=self.async_inference):	397	+ with self._temporarily_disable(async_inference=self.async_inference,
		398	+ mem_pool=self.generator_backend.kv_pool_backend):
393	block_mem_size_gb = gb(calc_block_mem(self.model_info, self.block_size, self.num_speculative_tokens))	399	block_mem_size_gb = gb(calc_block_mem(self.model_info, self.block_size, self.num_speculative_tokens))
394	print_log(self.rank, logger.info,	400	print_log(self.rank, logger.info,
395	f'One block during warmup needs npu memory(GiB): {block_mem_size_gb}')	401	f'One block during warmup needs npu memory(GiB): {block_mem_size_gb}')
	@@ -822,20 +828,25 @@ class Generator(PDInterface):
822	return ret_dict	828	return ret_dict
823		829
824	@contextmanager	830	@contextmanager
825	- def _temporarily_disable(self, dap: bool = False, async_inference: bool = False):	831	+ def _temporarily_disable(self, dap: bool = False, async_inference: bool = False, mem_pool: str = ""):
826	origin_enable_dap = self.generator_backend.enable_dap	832	origin_enable_dap = self.generator_backend.enable_dap
827	origin_async_inference = self.async_inference	833	origin_async_inference = self.async_inference
		834	+ origin_mem_pool = self.generator_backend.kv_pool_backend
828	try:	835	try:
829	if dap:	836	if dap:
830	self.generator_backend.enable_dap = False	837	self.generator_backend.enable_dap = False
831	if async_inference and not self.backend_type == "torch":	838	if async_inference and not self.backend_type == "torch":
832	self.async_inference = False	839	self.async_inference = False
		840	+ if len(mem_pool) != 0:
		841	+ self.generator_backend.kv_pool_backend = ""
833	yield	842	yield
834	finally:	843	finally:
835	if dap:	844	if dap:
836	self.generator_backend.enable_dap = origin_enable_dap	845	self.generator_backend.enable_dap = origin_enable_dap
837	if async_inference:	846	if async_inference:
838	self.async_inference = origin_async_inference	847	self.async_inference = origin_async_inference
		848	+ if len(mem_pool) != 0:
		849	+ self.generator_backend.kv_pool_backend = origin_mem_pool
839		850
840	def _init_plugin_manager(	851	def _init_plugin_manager(
841	self,	852	self,

Mmindie_llm/text_generator/plugins/plugin_manager.py+61-14

	@@ -13,8 +13,10 @@ from __future__ import annotations
13	import importlib	13	import importlib
14	import queue	14	import queue
15	import threading	15	import threading
		16	+import time
16	import copy	17	import copy
17	from dataclasses import fields	18	from dataclasses import fields
		19	+from enum import IntEnum
18	from typing import Iterable, Optional, Any, TYPE_CHECKING	20	from typing import Iterable, Optional, Any, TYPE_CHECKING
19		21
20	import numpy as np	22	import numpy as np
	@@ -52,6 +54,12 @@ LAUNCH_DONE_TIMEOUT = 20 * 60 # unit: second
52	MEM_DETECT_INTERVAL = 1000 # unit: second	54	MEM_DETECT_INTERVAL = 1000 # unit: second
53		55
54		56
		57	+class MemPoolType(IntEnum):
		58	+ DISABLED = 0
		59	+ SYNC_WRITE = 1
		60	+ ASYNC_WRITE = 2
		61	+
		62	+
55	class PluginManager:	63	class PluginManager:
56	def __init__(	64	def __init__(
57	self,	65	self,
	@@ -96,6 +104,8 @@ class PluginManager:
96	self.is_inference_pause = False	104	self.is_inference_pause = False
97	self.mem_det_trigger_counter = 0	105	self.mem_det_trigger_counter = 0
98	self.error_code_collected_in_async = None	106	self.error_code_collected_in_async = None
		107	+ self.mempool_type = MemPoolType.DISABLED
		108	+ self.warmup_is_end = True
99	# 结构化输出管理器 (延迟初始化)	109	# 结构化输出管理器 (延迟初始化)
100	self._structured_output_manager: Optional[Any] = None	110	self._structured_output_manager: Optional[Any] = None
101	self._structured_output_enabled = kwargs.get('enable_structured_output', True)	111	self._structured_output_enabled = kwargs.get('enable_structured_output', True)
	@@ -138,7 +148,7 @@ class PluginManager:
138	host_array = field_value.cpu().numpy()	148	host_array = field_value.cpu().numpy()
139	setattr(new_instance, field.name, host_array)	149	setattr(new_instance, field.name, host_array)
140	return new_instance	150	return new_instance
141	-	151	+
142	def clear_cache(	152	def clear_cache(
143	self,	153	self,
144	sequence_ids: Iterable[int],	154	sequence_ids: Iterable[int],
	@@ -169,10 +179,21 @@ class PluginManager:
169	**self.kwargs,	179	**self.kwargs,
170	)	180	)
171	setattr(self, plugin, plugin_tmp)	181	setattr(self, plugin, plugin_tmp)
172	-	182	+ if "prefix_cache" in self.plugin_list:
		183	+ self.mempool_type = self.prefix_cache.mempool_type
		184	+
173	# 初始化结构化输出管理器	185	# 初始化结构化输出管理器
174	self._init_structured_output_manager()	186	self._init_structured_output_manager()
175		187
		188	+ def wait_put_finish(self, input_metadata):
		189	+ if "prefix_cache" in self.plugin_list and input_metadata.is_prefill:
		190	+ logger.info("Waiting save to finished")
		191	+ start_t, timeout_t = time.time(), self.prefix_cache.save_timeout
		192	+ if self.prefix_cache.save_event.wait(timeout=timeout_t):
		193	+ logger.info(f"Save finished in {(time.time() - start_t)*1000:.1f} ms")
		194	+ else:
		195	+ logger.error(f"[TIMEOUT] Save unfinished after {timeout_t} seconds. Exit")
		196	+
176	def mem_det_trigger_counter_acc(self):	197	def mem_det_trigger_counter_acc(self):
177	if self.mem_det_trigger_counter < MEM_DETECT_INTERVAL:	198	if self.mem_det_trigger_counter < MEM_DETECT_INTERVAL:
178	self.mem_det_trigger_counter = self.mem_det_trigger_counter + 1	199	self.mem_det_trigger_counter = self.mem_det_trigger_counter + 1
	@@ -191,6 +212,9 @@ class PluginManager:
191	model_inputs, input_metadata, sampling_metadata, cache_ids)	212	model_inputs, input_metadata, sampling_metadata, cache_ids)
192	self.plugin_data_param.q_len = qlen if qlen is not None else self.plugin_data_param.q_len	213	self.plugin_data_param.q_len = qlen if qlen is not None else self.plugin_data_param.q_len
193	self.plugin_data_param.mask = mask if mask is not None else self.plugin_data_param.mask	214	self.plugin_data_param.mask = mask if mask is not None else self.plugin_data_param.mask
		215	+ if not warmup and "prefix_cache" in self.plugin_list and \
		216	+ self.prefix_cache.mempool_type == MemPoolType.ASYNC_WRITE:
		217	+ self.prefix_cache.async_put_prefix_kvcache_to_mempool(input_metadata, cache_ids)
194	span_end(prof)	218	span_end(prof)
195	self.watcher.watch_npu_mem(self.rank, f'After preprocess',	219	self.watcher.watch_npu_mem(self.rank, f'After preprocess',
196	trigger_count=self.mem_det_trigger_counter)	220	trigger_count=self.mem_det_trigger_counter)
	@@ -203,15 +227,20 @@ class PluginManager:
203		227
204	if ENV.framework_backend == BackendType.ATB:	228	if ENV.framework_backend == BackendType.ATB:
205	self.model_wrapper.model_runner.clear_internal_tensors()	229	self.model_wrapper.model_runner.clear_internal_tensors()
		230	+ forward_extra_kwargs = {}
		231	+ if warmup:
		232	+ forward_extra_kwargs["warmup_is_end"] = False
206	if (self.plugin_list and "mtp" not in self.plugin_list) or self.is_mix_model:	233	if (self.plugin_list and "mtp" not in self.plugin_list) or self.is_mix_model:
207	result = self.generator_backend.forward(model_inputs, q_lens=self.plugin_data_param.q_len,	234	result = self.generator_backend.forward(model_inputs, q_lens=self.plugin_data_param.q_len,
208	- attn_mask=self.plugin_data_param.mask) # q_len spec_mask	235	+ attn_mask=self.plugin_data_param.mask,
		236	+ **forward_extra_kwargs) # q_len spec_mask
209	# old graph forward	237	# old graph forward
210	else:	238	else:
211	result = self.generator_backend.forward(model_inputs, q_lens=self.plugin_data_param.q_len,	239	result = self.generator_backend.forward(model_inputs, q_lens=self.plugin_data_param.q_len,
212	spec_mask=self.plugin_data_param.mask,	240	spec_mask=self.plugin_data_param.mask,
213	sub_model_inputs=self.plugin_data_param.mtp_model_inputs,	241	sub_model_inputs=self.plugin_data_param.mtp_model_inputs,
214	- hidden_states=self.plugin_data_param.hidden_states)	242	+ hidden_states=self.plugin_data_param.hidden_states,
		243	+ **forward_extra_kwargs)
215	else:	244	else:
216	result = self.generator_backend.forward(model_inputs, q_lens=self.plugin_data_param.q_len,	245	result = self.generator_backend.forward(model_inputs, q_lens=self.plugin_data_param.q_len,
217	spec_mask=self.plugin_data_param.mask) # q_len spec_mask	246	spec_mask=self.plugin_data_param.mask) # q_len spec_mask
	@@ -231,7 +260,10 @@ class PluginManager:
231	self.watcher.watch_npu_mem(self.rank, f'After sample', trigger_count=self.mem_det_trigger_counter)	260	self.watcher.watch_npu_mem(self.rank, f'After sample', trigger_count=self.mem_det_trigger_counter)
232	logger.info("sample end", extra={"handler_ids": HandlerType.TOKEN})	261	logger.info("sample end", extra={"handler_ids": HandlerType.TOKEN})
233	prof = span_start("postprocess")	262	prof = span_start("postprocess")
234	- self.put_prefix_kvcache_to_mempool(input_metadata, cache_ids)	263	+ if self.mempool_type == MemPoolType.SYNC_WRITE:
		264	+ self.put_prefix_kvcache_to_mempool(input_metadata, cache_ids)
		265	+ elif not warmup and self.mempool_type == MemPoolType.ASYNC_WRITE:
		266	+ self.wait_put_finish(input_metadata)
235	generation_output = self.postprocess(	267	generation_output = self.postprocess(
236	cache_ids, input_metadata, result, sampling_metadata, sampling_output)	268	cache_ids, input_metadata, result, sampling_metadata, sampling_output)
237	generation_output.trace_ids = trace_ids	269	generation_output.trace_ids = trace_ids
	@@ -287,6 +319,13 @@ class PluginManager:
287	sub_model_inputs=self.plugin_data_param.mtp_model_inputs,	319	sub_model_inputs=self.plugin_data_param.mtp_model_inputs,
288	hidden_states=self.plugin_data_param.hidden_states	320	hidden_states=self.plugin_data_param.hidden_states
289	)	321	)
		322	+
		323	+ self.warmup_is_end = True
		324	+ if warmup:
		325	+ if model_kwargs is None:
		326	+ model_kwargs = {}
		327	+ model_kwargs["warmup_is_end"] = False
		328	+ self.warmup_is_end = False
290		329
291	if self.generator_backend.dp > 1:	330	if self.generator_backend.dp > 1:
292	cur_dp_rank_id_per_token_mask = model_input.dp_rank_ids == self.generator_backend.mapping.attn_dp.rank	331	cur_dp_rank_id_per_token_mask = model_input.dp_rank_ids == self.generator_backend.mapping.attn_dp.rank
	@@ -330,6 +369,10 @@ class PluginManager:
330	self.generator_backend.synchronize()	369	self.generator_backend.synchronize()
331	span_end(prof)	370	span_end(prof)
332		371
		372	+ if not warmup and "prefix_cache" in self.plugin_list and \
		373	+ self.prefix_cache.mempool_type == MemPoolType.ASYNC_WRITE:
		374	+ self.prefix_cache.async_put_prefix_kvcache_to_mempool(input_metadata, cache_ids)
		375	+
333	prof = span_start("put_into_input_queue")	376	prof = span_start("put_into_input_queue")
334	self.input_queue.put(model_input_wrapper)	377	self.input_queue.put(model_input_wrapper)
335	span_end(prof)	378	span_end(prof)
	@@ -653,18 +696,22 @@ class PluginManager:
653	if not self.is_inference_pause:	696	if not self.is_inference_pause:
654	model_input_wrapper.postprocess_done.wait()	697	model_input_wrapper.postprocess_done.wait()
655		698
656	- prof = span_start("verify")	699	+ prof = span_start("verify")
657	self.plugin_verify_manager(	700	self.plugin_verify_manager(
658	sampling_output, model_input_wrapper.cache_ids, model_output.original_result)	701	sampling_output, model_input_wrapper.cache_ids, model_output.original_result)
659	span_end(prof)	702	span_end(prof)
660		703
661	- prof = span_start("put_prefix_kvcache_to_mempool")	704	+ if self.mempool_type == MemPoolType.SYNC_WRITE:
662	- if model_input_wrapper.cache_ids is not None and not model_input_wrapper.input_metadata.is_dummy_batch:	705	+ prof = span_start("put_prefix_kvcache_to_mempool")
663	- self.put_prefix_kvcache_to_mempool(	706	+ if (
664	- model_input_wrapper.input_metadata,	707	+ model_input_wrapper.cache_ids is not None
665	- model_input_wrapper.cache_ids	708	+ and not model_input_wrapper.input_metadata.is_dummy_batch
666	- )	709	+ ):
667	- span_end(prof)	710	+ self.put_prefix_kvcache_to_mempool(
		711	+ model_input_wrapper.input_metadata, model_input_wrapper.cache_ids)
		712	+ span_end(prof)
		713	+ elif self.warmup_is_end and self.mempool_type == MemPoolType.ASYNC_WRITE:
		714	+ self.wait_put_finish(model_input_wrapper.input_metadata)
668		715
669	launch_done = threading.Event()	716	launch_done = threading.Event()
670	model_output_wrapper = ModelOutputWrapper(	717	model_output_wrapper = ModelOutputWrapper(
	@@ -804,7 +851,7 @@ class PluginManager:
804	config=config,	851	config=config,
805	)	852	)
806	self.infer_context.set_structured_output_manager(self._structured_output_manager)	853	self.infer_context.set_structured_output_manager(self._structured_output_manager)
807	-	854	+
808	except ImportError as e:	855	except ImportError as e:
809	logger.warning(f"Failed to import structured output module: {e}")	856	logger.warning(f"Failed to import structured output module: {e}")
810	self._structured_output_enabled = False	857	self._structured_output_enabled = False

Mmindie_llm/text_generator/plugins/prefix_cache/prefix_cache_plugin.py+65-10

	@@ -7,10 +7,13 @@
7	# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,	7	# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
8	# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.	8	# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
9	# See the Mulan PSL v2 for more details.	9	# See the Mulan PSL v2 for more details.
		10	+import threading
		11	+import queue
10	import numpy as np	12	import numpy as np
11		13
12	from .prefix_cache_preprocess import PrefixCachePreprocess	14	from .prefix_cache_preprocess import PrefixCachePreprocess
13	from ..plugin import Plugin	15	from ..plugin import Plugin
		16	+from ..plugin_manager import MemPoolType
14	from ....modeling.backend_type import BackendType	17	from ....modeling.backend_type import BackendType
15	from ....utils.env import ENV	18	from ....utils.env import ENV
16	from ....utils.log.logging import logger, print_log	19	from ....utils.log.logging import logger, print_log
	@@ -63,6 +66,7 @@ class PrefixCachePlugin(Plugin):
63	self.tp_size = 1	66	self.tp_size = 1
64	self.tp_rank = 0	67	self.tp_rank = 0
65	self.rank = generator_backend.rank	68	self.rank = generator_backend.rank
		69	+ self.device_id = self.generator_backend.npu_device_id
66	if hasattr(self.model_wrapper, "mapping"):	70	if hasattr(self.model_wrapper, "mapping"):
67	if self.model_wrapper.mapping.attn_tp.group_size > 1:	71	if self.model_wrapper.mapping.attn_tp.group_size > 1:
68	self.tp_size = self.model_wrapper.mapping.attn_tp.group_size	72	self.tp_size = self.model_wrapper.mapping.attn_tp.group_size
	@@ -84,20 +88,32 @@ class PrefixCachePlugin(Plugin):
84	self.is_300i = self.model_wrapper.model_runner.soc_info.is_300i()	88	self.is_300i = self.model_wrapper.model_runner.soc_info.is_300i()
85		89
86	## for kvcache pool	90	## for kvcache pool
		91	+ self.mempool_type = MemPoolType.DISABLED
		92	+ self.num_put_layers = self.kvcache_settings.num_layers
87	if len(self.generator_backend.kv_pool_backend) != 0 and len(self.generator_backend.kv_pool_config_path) != 0:	93	if len(self.generator_backend.kv_pool_backend) != 0 and len(self.generator_backend.kv_pool_config_path) != 0:
		94	+ self.mempool_type = \
		95	+ MemPoolType.ASYNC_WRITE if self.generator_backend.kv_pool_async_write else MemPoolType.SYNC_WRITE
88	from mindie_llm.text_generator.mempool import MemPool	96	from mindie_llm.text_generator.mempool import MemPool
89	self.m_store = MemPool.create_pool(	97	self.m_store = MemPool.create_pool(
90	backend=self.generator_backend.kv_pool_backend,	98	backend=self.generator_backend.kv_pool_backend,
91	config_path=self.generator_backend.kv_pool_config_path,	99	config_path=self.generator_backend.kv_pool_config_path,
92	role=MEM_POOL_ROLE_KEY,	100	role=MEM_POOL_ROLE_KEY,
93	- device_id=self.model_wrapper.device.index,	101	+ device_id=self.device_id,
94	kv_caches=self.generator_backend.cache_pool.npu_cache	102	kv_caches=self.generator_backend.cache_pool.npu_cache
95	)	103	)
96	- self.enable_mem_pool = self.m_store is not None	104	+ if self.m_store is None:
97	- if self.enable_mem_pool:	105	+ self.mempool_type = MemPoolType.DISABLED
98	- logger.info("Init mem pool successfully!!!")	106	+ if self.mempool_type != MemPoolType.DISABLED:
99	- else:	107	+ logger.info("Init mem pool successfully.")
100	- self.enable_mem_pool = False	108	+ if self.mempool_type == MemPoolType.ASYNC_WRITE:
		109	+ self.put_input_queue = queue.Queue()
		110	+ self.put_task_queue = queue.Queue()
		111	+ self.put_prefix_kvcache_thread = threading.Thread(target=self._put_prefix_kvcache_thread, daemon=True)
		112	+ self.put_prefix_kvcache_thread.start()
		113	+ self.save_event = threading.Event()
		114	+ self.save_event.set() # True
		115	+ self.save_timeout = 10 # 落盘等待时限,10s
		116	+ logger.info("Create prefix cache async save threads successfully.")
101		117
102	self.is_300i = False	118	self.is_300i = False
103	if self.generator_backend.backend_type == BackendType.ATB:	119	if self.generator_backend.backend_type == BackendType.ATB:
	@@ -228,7 +244,7 @@ class PrefixCachePlugin(Plugin):
228	return prefix_keys	244	return prefix_keys
229		245
230	def get_prefix_kvcache_from_mempool(self, input_metadata):	246	def get_prefix_kvcache_from_mempool(self, input_metadata):
231	- if not self.enable_mem_pool:	247	+ if self.mempool_type == MemPoolType.DISABLED:
232	return	248	return
233	computed_blocks = input_metadata.computed_blocks	249	computed_blocks = input_metadata.computed_blocks
234	remote_computed_blocks = input_metadata.remote_computed_blocks	250	remote_computed_blocks = input_metadata.remote_computed_blocks
	@@ -274,7 +290,7 @@ class PrefixCachePlugin(Plugin):
274	else:	290	else:
275	req_blocks_id = input_metadata.batch_block_tables[i, req_computed_blocks_id]	291	req_blocks_id = input_metadata.batch_block_tables[i, req_computed_blocks_id]
276	one_block_kvcache_tensors = []	292	one_block_kvcache_tensors = []
277	- for layer_id in range(self.kvcache_settings.num_layers):	293	+ for layer_id in range(self.num_put_layers):
278	k_cache = self.generator_backend.cache_pool.npu_cache[layer_id][0][req_blocks_id]	294	k_cache = self.generator_backend.cache_pool.npu_cache[layer_id][0][req_blocks_id]
279	v_cache = self.generator_backend.cache_pool.npu_cache[layer_id][1][req_blocks_id]	295	v_cache = self.generator_backend.cache_pool.npu_cache[layer_id][1][req_blocks_id]
280		296
	@@ -289,8 +305,28 @@ class PrefixCachePlugin(Plugin):
289	# 调用mempool get api接口，刷新有前缀复用block的kvcache	305	# 调用mempool get api接口，刷新有前缀复用block的kvcache
290	self.m_store.get(prefix_keys, kvcache_tensors)	306	self.m_store.get(prefix_keys, kvcache_tensors)
291		307
		308	+ def async_put_prefix_kvcache_to_mempool(self, input_metadata, cache_ids):
		309	+ if self.mempool_type == MemPoolType.DISABLED or not input_metadata.is_prefill:
		310	+ return
		311	+ self.put_input_queue.put((input_metadata, cache_ids))
		312	+
		313	+ def put_prefix_kvcache_put_task_queue(self, input_metadata, cache_ids):
		314	+ only_save_kv = False
		315	+ remote_computed_blocks = input_metadata.remote_computed_blocks
		316	+ if remote_computed_blocks is None:
		317	+ only_save_kv = True
		318	+ elif self.scp_size == 1:
		319	+ attn_dp_rank = self.generator_backend.mapping.attn_dp.rank
		320	+ cur_dp_remote_blocks_hits = 0
		321	+ for batch_dp_rank_id, num_computed_blocks in zip(input_metadata.batch_dp_rank_ids, remote_computed_blocks):
		322	+ if attn_dp_rank == batch_dp_rank_id:
		323	+ cur_dp_remote_blocks_hits += num_computed_blocks
		324	+ only_save_kv = cur_dp_remote_blocks_hits == 0
		325	+ for layer_id in range(self.num_put_layers):
		326	+ self.put_task_queue.put((layer_id == 0, layer_id == (self.num_put_layers - 1), only_save_kv))
		327	+
292	def put_prefix_kvcache_to_mempool(self, input_metadata, cache_ids):	328	def put_prefix_kvcache_to_mempool(self, input_metadata, cache_ids):
293	- if not self.enable_mem_pool or not input_metadata.is_prefill or \	329	+ if self.mempool_type == MemPoolType.DISABLED or not input_metadata.is_prefill or \
294	sum(input_metadata.batch_dp_rank_ids == self.generator_backend.mapping.attn_dp.rank) <= 0:	330	sum(input_metadata.batch_dp_rank_ids == self.generator_backend.mapping.attn_dp.rank) <= 0:
295	return	331	return
296	batch_input_ids = self.infer_context.get_all_input_ids(cache_ids)	332	batch_input_ids = self.infer_context.get_all_input_ids(cache_ids)
	@@ -339,7 +375,7 @@ class PrefixCachePlugin(Plugin):
339	else:	375	else:
340	req_blocks_id = input_metadata.batch_block_tables[i, req_uncomputed_blocks_id]	376	req_blocks_id = input_metadata.batch_block_tables[i, req_uncomputed_blocks_id]
341	one_block_kvcache_tensors = []	377	one_block_kvcache_tensors = []
342	- for layer_id in range(self.kvcache_settings.num_layers):	378	+ for layer_id in range(self.num_put_layers):
343	k_cache = self.generator_backend.cache_pool.npu_cache[layer_id][0][req_blocks_id]	379	k_cache = self.generator_backend.cache_pool.npu_cache[layer_id][0][req_blocks_id]
344	v_cache = self.generator_backend.cache_pool.npu_cache[layer_id][1][req_blocks_id]	380	v_cache = self.generator_backend.cache_pool.npu_cache[layer_id][1][req_blocks_id]
345	one_block_kvcache_tensors.append([k_cache, v_cache])	381	one_block_kvcache_tensors.append([k_cache, v_cache])
	@@ -351,3 +387,22 @@ class PrefixCachePlugin(Plugin):
351	if len(prefix_keys) > 0:	387	if len(prefix_keys) > 0:
352	# 调用mempool put api接口，将新计算的kvcache传到mempool	388	# 调用mempool put api接口，将新计算的kvcache传到mempool
353	self.m_store.put(prefix_keys, kvcache_tensors)	389	self.m_store.put(prefix_keys, kvcache_tensors)
		390	+
		391	+ def _put_prefix_kvcache_thread(self):
		392	+ import torch
		393	+ torch.npu.set_device(f"npu:{self.device_id}")
		394	+ stream = torch.npu.Stream()
		395	+ torch.npu.set_stream(stream)
		396	+ logger.info("Create _put_prefix_kvcache_thread")
		397	+ while True:
		398	+ input_metadata, cache_ids = self.put_input_queue.get()
		399	+ self.put_prefix_kvcache_put_task_queue(input_metadata, cache_ids)
		400	+ while not self.put_task_queue.empty():
		401	+ is_first, is_last, only_save_kv = self.put_task_queue.get()
		402	+ if is_first:
		403	+ self.save_event.clear() # False
		404	+ pipe_key = self.model_wrapper.generate_mem_pool_event_key(only_save_kv)
		405	+ self.model_wrapper.model_runner.wait_event(pipe_key)
		406	+ if is_last:
		407	+ self.put_prefix_kvcache_to_mempool(input_metadata, cache_ids)
		408	+ self.save_event.set() # True

Msrc/config_manager/backend_config.cpp+5-0

	@@ -43,6 +43,7 @@ void BackendConfigManager::InitKvPoolConfigFromJson(Json &backendConfigData)
43	{	43	{
44	std::string backend{};	44	std::string backend{};
45	std::string configPath{};	45	std::string configPath{};
		46	+ bool asyncWrite = false;
46	if (backendConfigData.contains("kvPoolConfig")) {	47	if (backendConfigData.contains("kvPoolConfig")) {
47	Json& kvPoolConfig = backendConfigData["kvPoolConfig"];	48	Json& kvPoolConfig = backendConfigData["kvPoolConfig"];
48	if (kvPoolConfig.contains("backend")) {	49	if (kvPoolConfig.contains("backend")) {
	@@ -51,9 +52,13 @@ void BackendConfigManager::InitKvPoolConfigFromJson(Json &backendConfigData)
51	if (kvPoolConfig.contains("configPath")) {	52	if (kvPoolConfig.contains("configPath")) {
52	configPath = kvPoolConfig["configPath"];	53	configPath = kvPoolConfig["configPath"];
53	}	54	}
		55	+ if (kvPoolConfig.contains("asyncWrite")) {
		56	+ asyncWrite = kvPoolConfig["asyncWrite"];
		57	+ }
54	}	58	}
55	backendConfig_.kvPoolConfig.backend = backend;	59	backendConfig_.kvPoolConfig.backend = backend;
56	backendConfig_.kvPoolConfig.configPath = configPath;	60	backendConfig_.kvPoolConfig.configPath = configPath;
		61	+ backendConfig_.kvPoolConfig.asyncWrite = asyncWrite;
57	}	62	}
58		63
59	bool BackendConfigManager::InitTlsConfigFromJson(Json &backendConfigData)	64	bool BackendConfigManager::InitTlsConfigFromJson(Json &backendConfigData)

Msrc/include/config/config_info.h+1-0

	@@ -37,6 +37,7 @@ enum class WorkFlowType : uint32_t {
37	struct KvPoolConfig {	37	struct KvPoolConfig {
38	std::string backend;	38	std::string backend;
39	std::string configPath;	39	std::string configPath;
		40	+ bool asyncWrite;
40	};	41	};
41		42
42	struct ServerConfig {	43	struct ServerConfig {

Msrc/llm_manager_v2/impl/llm_manager_impl.cpp+3-4

	@@ -944,7 +944,8 @@ static void LLMSetModelConfig(std::map<std::string, std::string> &modelConfig, c
944	modelConfig["max_beam_width"] = std::to_string(engineConfig.maxBeamWidth);	944	modelConfig["max_beam_width"] = std::to_string(engineConfig.maxBeamWidth);
945	modelConfig["kv_pool_backend"] = engineConfig.kvPoolConfig.backend;	945	modelConfig["kv_pool_backend"] = engineConfig.kvPoolConfig.backend;
946	modelConfig["kv_pool_config_path"] = engineConfig.kvPoolConfig.configPath;	946	modelConfig["kv_pool_config_path"] = engineConfig.kvPoolConfig.configPath;
947	-	947	+ modelConfig["kv_pool_async_write"] = engineConfig.kvPoolConfig.asyncWrite ? "true" : "false";
		948	+
948	std::string npuIds;	949	std::string npuIds;
949	if (!modelParam.npuDeviceIds.empty()) {	950	if (!modelParam.npuDeviceIds.empty()) {
950	for (auto &item : modelParam.npuDeviceIds) {	951	for (auto &item : modelParam.npuDeviceIds) {
	@@ -977,9 +978,7 @@ static void LLMSetModelConfig(std::map<std::string, std::string> &modelConfig, c
977	modelConfig["threadNum"] = (modelConfig["asyncBatchscheduler"] == "true") ? "2" : "1";	978	modelConfig["threadNum"] = (modelConfig["asyncBatchscheduler"] == "true") ? "2" : "1";
978		979
979	auto &configManager = mindie_llm::ConfigManager::GetInstance();	980	auto &configManager = mindie_llm::ConfigManager::GetInstance();
980	- if (configManager.IslayerwiseDisaggregated()) {	981	+ if (configManager.IslayerwiseDisaggregated()) LLMSetLayerwiseDisaggregatedModelConfig(modelConfig, engineConfig);
981	- LLMSetLayerwiseDisaggregatedModelConfig(modelConfig, engineConfig);
982	- }
983	}	982	}
984		983
985	static void InitPolicyConfig(SchedulerConfig &schedulerConfig, const EngineConfig &engineConfig)	984	static void InitPolicyConfig(SchedulerConfig &schedulerConfig, const EngineConfig &engineConfig)

Mtests/pythontest/npu/text_generator/test_plugins/test_prefix_cache_plugin.py+152-5

	@@ -39,10 +39,14 @@ class TestPrefixCahcePlugin(unittest.TestCase):
39	@classmethod	39	@classmethod
40	def setUpClass(cls):	40	def setUpClass(cls):
41	sys.modules['_libatb_torch'] = MagicMock()	41	sys.modules['_libatb_torch'] = MagicMock()
		42	+ mock_cpu_handler = SimpleNamespace()
		43	+ mock_cpu_handler._PostProcessingManager = MagicMock()
		44	+ sys.modules['_cpu_logits_handler'] = mock_cpu_handler
42		45
43	@classmethod	46	@classmethod
44	def tearDownClass(cls):	47	def tearDownClass(cls):
45	del sys.modules['_libatb_torch']	48	del sys.modules['_libatb_torch']
		49	+ del sys.modules['_cpu_logits_handler']
46		50
47	def setUp(self):	51	def setUp(self):
48	self.model_config = {	52	self.model_config = {
	@@ -70,6 +74,15 @@ class TestPrefixCahcePlugin(unittest.TestCase):
70	plugin_dict = {'plugin_params': PLUGIN_PARAMS}	74	plugin_dict = {'plugin_params': PLUGIN_PARAMS}
71	self.model_config.update(plugin_dict)	75	self.model_config.update(plugin_dict)
72		76
		77	+ if hasattr(self, 'model_config') and self.model_config:
		78	+ # 创建一个模拟的 generator_backend
		79	+ self.generator_backend = MagicMock()
		80	+ self.generator_backend.mapping = MagicMock()
		81	+ self.generator_backend.mapping.attn_dp = MagicMock()
		82	+ self.generator_backend.mapping.attn_dp.rank = np.array([1])
		83	+ self.generator_backend.rank = int(self.model_config['rank'])
		84	+ self.generator_backend.local_rank = int(self.model_config['local_rank'])
		85	+
73	fake_parallel_info = FakeParallelInfo(	86	fake_parallel_info = FakeParallelInfo(
74	dp=int(self.model_config['dp']),	87	dp=int(self.model_config['dp']),
75	tp=int(self.model_config['tp']),	88	tp=int(self.model_config['tp']),
	@@ -77,6 +90,7 @@ class TestPrefixCahcePlugin(unittest.TestCase):
77	cp=int(self.model_config['cp'])	90	cp=int(self.model_config['cp'])
78	)	91	)
79	self.fake_model_runner = FakeModelRunner(parallel_info=fake_parallel_info)	92	self.fake_model_runner = FakeModelRunner(parallel_info=fake_parallel_info)
		93	+ return super().setUp()
80		94
81	@patch('torch.npu.synchronize', return_value=None)	95	@patch('torch.npu.synchronize', return_value=None)
82	@patch('atb_llm.runner.model_runner.ModelRunner')	96	@patch('atb_llm.runner.model_runner.ModelRunner')
	@@ -111,7 +125,7 @@ class TestPrefixCahcePlugin(unittest.TestCase):
111		125
112	self.model_config['kv_pool_backend'] = "mooncake"	126	self.model_config['kv_pool_backend'] = "mooncake"
113	self.model_config['kv_pool_config_path'] = "a.json"	127	self.model_config['kv_pool_config_path'] = "a.json"
114	-	128	+
115	generator = Generator(self.model_config)	129	generator = Generator(self.model_config)
116		130
117	prefix_cache_plugin = generator.plugin_manager	131	prefix_cache_plugin = generator.plugin_manager
	@@ -122,8 +136,8 @@ class TestPrefixCahcePlugin(unittest.TestCase):
122	mock_kv_block_num = [1] * 100	136	mock_kv_block_num = [1] * 100
123	prefix_cache_plugin.generator_backend.cache_pool.npu_cache = [(mock_kv_block_num, mock_kv_block_num)] * 100	137	prefix_cache_plugin.generator_backend.cache_pool.npu_cache = [(mock_kv_block_num, mock_kv_block_num)] * 100
124	gen_len = 2	138	gen_len = 2
125	- req = Request.request_from_token(input1,	139	+ req = Request.request_from_token(input1,
126	- sampling_params=greedy_param,	140	+ sampling_params=greedy_param,
127	generation_params=GenerationParams(max_new_tokens=gen_len))	141	generation_params=GenerationParams(max_new_tokens=gen_len))
128	meta_data = InputMetadata.from_requests([req], block_tables, True)	142	meta_data = InputMetadata.from_requests([req], block_tables, True)
129	meta_data.block_tables = block_tables	143	meta_data.block_tables = block_tables
	@@ -132,11 +146,11 @@ class TestPrefixCahcePlugin(unittest.TestCase):
132	meta_data.sp_tokens = np.array([128, 72]).reshape(1, -1)	146	meta_data.sp_tokens = np.array([128, 72]).reshape(1, -1)
133	# 无复用，使用fa算子做prefill	147	# 无复用，使用fa算子做prefill
134	generation_output = prefix_cache_plugin.generate_token(meta_data)	148	generation_output = prefix_cache_plugin.generate_token(meta_data)
135	-	149	+
136	# 有复用，使用qlen > 1 的 pa算子做prefill	150	# 有复用，使用qlen > 1 的 pa算子做prefill
137	remote_computed_blocks = np.ones(CP, dtype=np.int64).reshape(1, -1)	151	remote_computed_blocks = np.ones(CP, dtype=np.int64).reshape(1, -1)
138	remote_computed_blocks[0, 1] = 0	152	remote_computed_blocks[0, 1] = 0
139	-	153	+
140	# 都在本地命中	154	# 都在本地命中
141	meta_data.computed_blocks = np.zeros(CP, dtype=np.int64).reshape(1, -1)	155	meta_data.computed_blocks = np.zeros(CP, dtype=np.int64).reshape(1, -1)
142	meta_data.computed_blocks[0, 1] = 1	156	meta_data.computed_blocks[0, 1] = 1
	@@ -168,6 +182,139 @@ class TestPrefixCahcePlugin(unittest.TestCase):
168	break	182	break
169	self.assertEqual(check_greedy, 1)	183	self.assertEqual(check_greedy, 1)
170		184
		185	+ def test_async_put_prefix_kvcache_to_mempool_function(self):
		186	+ def side_effect_initialize_distributed(rank, npu_id, world_size):
		187	+ return torch.device("cpu")
		188	+
		189	+ def side_effect_forward(model_inputs, **kwargs):
		190	+ if model_inputs.is_prefill:
		191	+ token_num = model_inputs.prefill_head_indices.shape[0]
		192	+ else:
		193	+ token_num = model_inputs.input_ids.shape[0]
		194	+ logits = torch.zeros(token_num, 10)
		195	+ return logits
		196	+
		197	+ with patch.object(GeneratorTorch, 'forward') as mock_forward, \
		198	+ patch('atb_llm.utils.dist.initialize_distributed') as mock_initialize_distributed, \
		199	+ patch('atb_llm.runner.model_runner.ModelRunner', return_value=self.fake_model_runner) as _, \
		200	+ patch('mindie_llm.text_generator.utils.kvcache_settings.NPUSocInfo.support_nz', return_value=True) as _, \
		201	+ patch('mindie_llm.text_generator.utils.kvcache_settings.KVCacheSettings') as mock_kvcache_settings_class, \
		202	+ patch('torch.npu.synchronize', return_value=None) as _, \
		203	+ patch('torch.npu.set_device', return_value=None) as _, \
		204	+ patch('torch.npu.set_stream', return_value=None) as _, \
		205	+ patch('torch.npu.Stream', return_value=MagicMock()) as mock_stream, \
		206	+ patch('mindie_llm.text_generator.adapter.generator_torch.GeneratorTorch._get_obfuscation_func', \
		207	+ return_value=None) as _:
		208	+
		209	+ mock_initialize_distributed.side_effect = side_effect_initialize_distributed
		210	+ mock_forward.side_effect = side_effect_forward
		211	+ mock_kvcache_settings = MagicMock(dtype=None)
		212	+ mock_kvcache_settings_class.return_value = mock_kvcache_settings
		213	+ self.model_config['kv_pool_backend'] = "mooncake"
		214	+ self.model_config['kv_pool_config_path'] = "a.json"
		215	+ self.model_config['kv_pool_async_write'] = "true"
		216	+ generator = Generator(self.model_config)
		217	+
		218	+ prefix_cache_plugin = generator.plugin_manager
		219	+
		220	+ # 创建测试输入元数据
		221	+ mock_metadata = MagicMock()
		222	+ mock_metadata.is_prefill = True
		223	+ mock_metadata.computed_blocks = np.array([[1, 0]])
		224	+ mock_metadata.remote_computed_blocks = np.array([[1, 0]])
		225	+
		226	+ # 创建测试 cache_ids
		227	+ cache_ids = [0]
		228	+
		229	+ # 测试 async_put_prefix_kvcache_to_mempool
		230	+ prefix_cache_plugin.prefix_cache.async_put_prefix_kvcache_to_mempool(mock_metadata, cache_ids)
		231	+
		232	+ # 测试 put_prefix_kvcache_put_task_queue
		233	+ mock_metadata.batch_dp_rank_ids = [0]
		234	+ mock_metadata.remote_computed_blocks = np.array([1])
		235	+ prefix_cache_plugin.prefix_cache.put_prefix_kvcache_put_task_queue(mock_metadata, cache_ids)
		236	+
		237	+ # 参数构建
		238	+ mock_metadata.batch_dp_rank_ids = np.array([0, 1])
		239	+ mock_metadata.remote_computed_blocks = np.array([[1, 1]])
		240	+ mock_metadata.batch_size = 1
		241	+ mock_metadata.batch_seq_len = [128]
		242	+ mock_metadata.max_block_size = 128
		243	+ mock_metadata.batch_block_tables = np.array([[[0, 1]]])
		244	+ prefix_cache_plugin.infer_context = MagicMock()
		245	+ prefix_cache_plugin.infer_context.get_all_input_ids.return_value = np.array([[1, 2, 3, 4, 5] * 256])
		246	+ prefix_cache_plugin.infer_context.get_seq_lens.return_value = [128]
		247	+ mock_kv_block_num = [1] * 100
		248	+ prefix_cache_plugin.generator_backend.cache_pool.npu_cache = [(mock_kv_block_num, mock_kv_block_num)] * 100
		249	+
		250	+ # 测试 put_prefix_kvcache_to_mempool
		251	+ prefix_cache_plugin.put_prefix_kvcache_to_mempool(mock_metadata, cache_ids)
		252	+
		253	+ def test_get_prefix_kvcache_from_mempool(self):
		254	+ from mindie_llm.text_generator.plugins.plugin_manager import MemPoolType
		255	+ from mindie_llm.text_generator.plugins.prefix_cache.prefix_cache_plugin import PrefixCachePlugin
		256	+
		257	+ plugin = PrefixCachePlugin.__new__(PrefixCachePlugin)
		258	+ plugin.mempool_type = MemPoolType.SYNC_WRITE
		259	+ plugin.scp_size = 2
		260	+ plugin.scp_rank = 0
		261	+ plugin.num_put_layers = 61
		262	+ plugin.generator_backend = self.generator_backend
		263	+ plugin.model_name = "llama"
		264	+
		265	+ metadata = SimpleNamespace(
		266	+ is_prefill=True,
		267	+ computed_blocks=None,
		268	+ remote_computed_blocks=np.array([[0]]),
		269	+ batch_size=1,
		270	+ batch_dp_rank_ids=[0],
		271	+ batch_seq_len=[10],
		272	+ input_ids=np.array([1, 2, 3, 4, 5]),
		273	+ max_block_size=128,
		274	+ batch_block_tables=np.array([[0, 1]]).reshape(1, 1, -1),
		275	+ block_tables=np.array([[0, 1]]).reshape(1, 1, -1)
		276	+ )
		277	+
		278	+ plugin.get_prefix_kvcache_from_mempool(metadata)
		279	+
		280	+ def test_hash_and_prefix_key(self):
		281	+ from mindie_llm.text_generator.plugins.prefix_cache.prefix_cache_plugin import (
		282	+ cpp_style_hash, hash_combine, PrefixCachePlugin
		283	+ )
		284	+
		285	+ self.assertEqual(cpp_style_hash(123), 123)
		286	+ self.assertNotEqual(cpp_style_hash("abc"), 0)
		287	+
		288	+ seed = hash_combine(0, 10)
		289	+ self.assertNotEqual(seed, 0)
		290	+
		291	+ plugin = PrefixCachePlugin.__new__(PrefixCachePlugin)
		292	+ plugin.scp_size = 1
		293	+ plugin.tp_rank = 0
		294	+ plugin.tp_size = 1
		295	+ plugin.model_name = "llama"
		296	+
		297	+ hash_val = plugin.hash_block(0, [1, 2, 3, 4])
		298	+ self.assertNotEqual(hash_val, 0)
		299	+
		300	+ key = plugin.get_prefix_keys(hash_val)
		301	+ self.assertIn("llama", key)
		302	+
		303	+ def test_enable_prefixcache_flags(self):
		304	+ from mindie_llm.text_generator.plugins.prefix_cache.prefix_cache_plugin import PrefixCachePlugin
		305	+
		306	+ meta = SimpleNamespace(
		307	+ is_prefill=True,
		308	+ computed_blocks=np.array([1]),
		309	+ remote_computed_blocks=np.array([1])
		310	+ )
		311	+
		312	+ self.assertTrue(PrefixCachePlugin.enable_local_prefixcache(meta))
		313	+ self.assertTrue(PrefixCachePlugin.enable_remmote_prefixcache(meta))
		314	+
		315	+ meta.computed_blocks = None
		316	+ self.assertFalse(PrefixCachePlugin.enable_local_prefixcache(meta))
		317	+
171		318
172	if __name__ == "__main__":	319	if __name__ == "__main__":
173	unittest.main()	320	unittest.main()

	@@ -10,20 +10,22 @@
10	from atb_llm.models.base.graph_manager.graph_wrapper import ATBGraphWrapper	10	from atb_llm.models.base.graph_manager.graph_wrapper import ATBGraphWrapper
11	from atb_llm.models.base.graph_manager.compatible_matrix import FeatureType	11	from atb_llm.models.base.graph_manager.compatible_matrix import FeatureType
12	from atb_llm.models.base.flash_causal_lm import FlashForCausalLM	12	from atb_llm.models.base.flash_causal_lm import FlashForCausalLM
		13	+from atb_llm.runner.model_runner import generate_mem_pool_event_key
13		14
14		15
15	class SplitFuseGraphWrapper(ATBGraphWrapper):	16	class SplitFuseGraphWrapper(ATBGraphWrapper):
16	"ATBGraphWrapper class for prefixcache and splitfuse"	17	"ATBGraphWrapper class for prefixcache and splitfuse"
		18	+ pipe_key: str = generate_mem_pool_event_key(only_save_kv=False)
		19	+
17	def __init__(self):	20	def __init__(self):
18	super().__init__()	21	super().__init__()
19	-
20	self.feature_name = FeatureType.SPLITFUSE	22	self.feature_name = FeatureType.SPLITFUSE
21	- self.feature_params = {"enableSplitFuse": True, "isPrefill": True}	23	+ self.feature_params = {"enableSplitFuse": True, "isPrefill": True, "pipekey": self.pipe_key}
22	-	24	+
23	def activate(self, context: FlashForCausalLM, runtime_params, **kwargs) -> bool:	25	def activate(self, context: FlashForCausalLM, runtime_params, **kwargs) -> bool:
24	pa_enable = False if context.inference_mode is None else context.inference_mode.enable_prefill_pa	26	pa_enable = False if context.inference_mode is None else context.inference_mode.enable_prefill_pa
25	q_lens = "\"qLen\"" in runtime_params	27	q_lens = "\"qLen\"" in runtime_params
26	is_prefill = kwargs.get("is_prefill", False)	28	is_prefill = kwargs.get("is_prefill", False)
27	if q_lens and is_prefill and pa_enable:	29	if q_lens and is_prefill and pa_enable:
28	return True	30	return True
29	- return False	31	+ return False

	@@ -28,6 +28,7 @@ from ...utils.log.logging import logger
28	from ...utils.tensor import op	28	from ...utils.tensor import op
29	from ...utils.validation import parse_config, ParseType, MODEL_CONFIG_KEY_TYPE	29	from ...utils.validation import parse_config, ParseType, MODEL_CONFIG_KEY_TYPE
30	from .recovery_utils import check_and_recover_uce_in_cache	30	from .recovery_utils import check_and_recover_uce_in_cache
		31	+from ...text_generator.plugins.plugin_manager import MemPoolType
31		32
32	MAX_WORLD_SIZE = 1048576	33	MAX_WORLD_SIZE = 1048576
33	MAX_KEY_LENGTH = 256	34	MAX_KEY_LENGTH = 256
	@@ -49,7 +50,7 @@ class GeneratorBackend:
49	backend_type = parse_config(model_config, 'backend_type', required=True)	50	backend_type = parse_config(model_config, 'backend_type', required=True)
50	num_threads = parse_config(model_config, 'num_threads', parse_type=ParseType.TO_INT, default_value=8)	51	num_threads = parse_config(model_config, 'num_threads', parse_type=ParseType.TO_INT, default_value=8)
51	self.npu_device_id = parse_config(model_config, 'npu_device_id', required=True, parse_type=ParseType.TO_INT)	52	self.npu_device_id = parse_config(model_config, 'npu_device_id', required=True, parse_type=ParseType.TO_INT)
52	- local_rank = parse_config(model_config, 'local_rank', required=True, parse_type=ParseType.TO_INT)	53	+ self.local_rank = parse_config(model_config, 'local_rank', required=True, parse_type=ParseType.TO_INT)
53	self.rank = parse_config(model_config, 'rank', required=True, parse_type=ParseType.TO_INT)	54	self.rank = parse_config(model_config, 'rank', required=True, parse_type=ParseType.TO_INT)
54	self.world_size = parse_config(model_config, 'world_size', required=True, parse_type=ParseType.TO_INT)	55	self.world_size = parse_config(model_config, 'world_size', required=True, parse_type=ParseType.TO_INT)
55	self.trust_remote_code = parse_config(model_config, 'trust_remote_code', required=True,	56	self.trust_remote_code = parse_config(model_config, 'trust_remote_code', required=True,
	@@ -68,6 +69,8 @@ class GeneratorBackend:
68	parse_type=ParseType.TO_STR, default_value='')	69	parse_type=ParseType.TO_STR, default_value='')
69	self.kv_pool_config_path = parse_config(model_config, 'kv_pool_config_path', required=False,	70	self.kv_pool_config_path = parse_config(model_config, 'kv_pool_config_path', required=False,
70	parse_type=ParseType.TO_STR, default_value='')	71	parse_type=ParseType.TO_STR, default_value='')
		72	+ self.kv_pool_async_write = parse_config(model_config, 'kv_pool_async_write', required=False,
		73	+ parse_type=ParseType.TO_BOOL, default_value=False)
71		74
72	if self.world_size < 1 or self.world_size > MAX_WORLD_SIZE:	75	if self.world_size < 1 or self.world_size > MAX_WORLD_SIZE:
73	raise ValueError("World size should be in the range of 1 to 1048576.")	76	raise ValueError("World size should be in the range of 1 to 1048576.")
	@@ -86,7 +89,7 @@ class GeneratorBackend:
86	model_config, 'num_lccl_comm_shards', parse_type=ParseType.TO_INT, default_value=1)	89	model_config, 'num_lccl_comm_shards', parse_type=ParseType.TO_INT, default_value=1)
87	lccl_comm_shard_id = parse_config(	90	lccl_comm_shard_id = parse_config(
88	model_config, 'lccl_comm_shard_id', parse_type=ParseType.TO_INT, default_value=0)	91	model_config, 'lccl_comm_shard_id', parse_type=ParseType.TO_INT, default_value=0)
89	- if local_rank < 0 or local_rank >= self.world_size:	92	+ if self.local_rank < 0 or self.local_rank >= self.world_size:
90	raise ValueError("Local rank should be in the range of 0 to world_size - 1.")	93	raise ValueError("Local rank should be in the range of 0 to world_size - 1.")
91	max_loras = parse_config(model_config, 'max_loras', required=False, parse_type=ParseType.TO_INT, default_value=0)	94	max_loras = parse_config(model_config, 'max_loras', required=False, parse_type=ParseType.TO_INT, default_value=0)
92	max_lora_rank = parse_config(model_config, 'max_lora_rank', required=False, parse_type=ParseType.TO_INT, default_value=0)	95	max_lora_rank = parse_config(model_config, 'max_lora_rank', required=False, parse_type=ParseType.TO_INT, default_value=0)
	@@ -106,7 +109,7 @@ class GeneratorBackend:
106	model_config["rank"] = self.rank	109	model_config["rank"] = self.rank
107	model_config["world_size"] = self.world_size	110	model_config["world_size"] = self.world_size
108	model_config['npu_device_id'] = self.npu_device_id	111	model_config['npu_device_id'] = self.npu_device_id
109	- model_config['local_rank'] = local_rank	112	+ model_config['local_rank'] = self.local_rank
110	model_config['dp'] = dp	113	model_config['dp'] = dp
111	model_config['tp'] = tp	114	model_config['tp'] = tp
112	model_config['attn_inner_sp'] = attn_inner_sp	115	model_config['attn_inner_sp'] = attn_inner_sp
	@@ -123,6 +126,11 @@ class GeneratorBackend:
123	model_config['max_lora_rank'] = max_lora_rank	126	model_config['max_lora_rank'] = max_lora_rank
124	model_config['sampler_config'] = sampler_config	127	model_config['sampler_config'] = sampler_config
125	model_config['lwdNextPHeadPrior'] = lwd_next_p_head_prior	128	model_config['lwdNextPHeadPrior'] = lwd_next_p_head_prior
		129	+ if bool(self.kv_pool_config_path) and bool(self.kv_pool_backend):
		130	+ model_config['mempool_type'] = \
		131	+ MemPoolType.ASYNC_WRITE if self.kv_pool_async_write else MemPoolType.SYNC_WRITE
		132	+ else:
		133	+ model_config['mempool_type'] = MemPoolType.DISABLED
126		134
127	self.backend_type = backend_type	135	self.backend_type = backend_type
128	self.model_wrapper = get_model_wrapper(model_config, backend_type)	136	self.model_wrapper = get_model_wrapper(model_config, backend_type)
	@@ -176,7 +184,7 @@ class GeneratorBackend:
176	'''	184	'''
177	self.force_stop_exception_occurred.set()	185	self.force_stop_exception_occurred.set()
178	logger.info(f"FORCE STOP exception detected and notified for device {self.npu_device_id}")	186	logger.info(f"FORCE STOP exception detected and notified for device {self.npu_device_id}")
179	-	187	+
180	def execute_recover_command(self, command: str) -> dict:	188	def execute_recover_command(self, command: str) -> dict:
181	'''	189	'''
182	Execute recover related command.	190	Execute recover related command.
	@@ -282,7 +290,7 @@ class GeneratorBackend:
282	command_result = 0	290	command_result = 0
283	error_msg = ""	291	error_msg = ""
284	return command_result, error_msg	292	return command_result, error_msg
285	-	293	+
286	def _execute_cmd_reinit_npu(self):	294	def _execute_cmd_reinit_npu(self):
287	'''Reinitialize NPU. Subclasses must override with backend-specific logic.'''	295	'''Reinitialize NPU. Subclasses must override with backend-specific logic.'''
288	raise NotImplementedError("Subclasses must implement _execute_cmd_reinit_npu")	296	raise NotImplementedError("Subclasses must implement _execute_cmd_reinit_npu")