engram反向性能优化-ops-transformer-AtomGit

Mmc2/common/torch_extension/csrc/elastic_buffer.cpp+103-19

	@@ -119,7 +119,7 @@ struct EngramCommContext {
119	uint32_t rankId = 0;	119	uint32_t rankId = 0;
120	uint32_t rankSize = 0;	120	uint32_t rankSize = 0;
121	uint64_t virtualAddrList[HCCL_MAX_RANK_SIZE] = {};	121	uint64_t virtualAddrList[HCCL_MAX_RANK_SIZE] = {};
122	- uint64_t hcommHandle[HCCL_MAX_RANK_SIZE] = {};	122	+ uint64_t hcommHandle[HCCL_MAX_RANK_SIZE * 2] = {};
123	uint32_t channelsPerRank = 1;	123	uint32_t channelsPerRank = 1;
124	};	124	};
125		125
	@@ -392,12 +392,22 @@ private:
392	static void QueryHcclBufferResource(const HcclComm &commHandle, EngramContextResources &resources)	392	static void QueryHcclBufferResource(const HcclComm &commHandle, EngramContextResources &resources)
393	{	393	{
394	uint32_t rankId = resources.context.rankId;	394	uint32_t rankId = resources.context.rankId;
395	- constexpr uint32_t channelsPerPeer = 2U;	395	+ uint32_t rankSize = resources.context.rankSize;
		396	+ constexpr uint32_t aivNum = 72U;
		397	+ uint32_t numSendCores = aivNum / 2U;
		398	+ if (numSendCores == 0U) {
		399	+ numSendCores = 1U;
		400	+ }
		401	+ uint32_t groupSize = numSendCores / rankSize;
		402	+ if (groupSize == 0U) {
		403	+ groupSize = 1U;
		404	+ }
		405	+ uint32_t channelsPerPeer = groupSize * 2U;
396	resources.context.channelsPerRank = channelsPerPeer;	406	resources.context.channelsPerRank = channelsPerPeer;
397	- ChannelHandle handlesByRank[HCCL_MAX_RANK_SIZE * 2] = {};	407	+
398	- GetHcclCommChannel(commHandle, resources.context.rankSize, rankId, channelsPerPeer, resources.memHandle,	408	+ std::vector<ChannelHandle> handlesByRank(rankSize * channelsPerPeer);
399	- handlesByRank);	409	+ GetHcclCommChannel(commHandle, rankSize, rankId, channelsPerPeer, resources.memHandle, handlesByRank.data());
400	- for (uint32_t peer = 0; peer < resources.context.rankSize; ++peer) {	410	+ for (uint32_t peer = 0; peer < rankSize; ++peer) {
401	if (peer == rankId)	411	if (peer == rankId)
402	continue;	412	continue;
403	for (uint32_t ch = 0; ch < channelsPerPeer; ++ch) {	413	for (uint32_t ch = 0; ch < channelsPerPeer; ++ch) {
	@@ -759,11 +769,26 @@ public:
759	void EngramBarrier(bool useCommStream = true, bool withCpuSync = false);	769	void EngramBarrier(bool useCommStream = true, bool withCpuSync = false);
760	void Destroy();	770	void Destroy();
761		771
762	- int64_t GetHostBufPtr() const { return reinterpret_cast<int64_t>(engramHostBufPtr_); }	772	+ int64_t GetHostBufPtr() const
763	- const at::Tensor &GetContextTensor() const { return engramContextTensor_; }	773	+ {
764	- const at::Tensor &GetLocalStorageAddrTensor() const { return localStorageAddrTensor_; }	774	+ return reinterpret_cast<int64_t>(engramHostBufPtr_);
765	- int64_t GetCommBufferSize() const { return commBufferSize_; }	775	+ }
766	- uint32_t GetRankSize() const { return engramCommContext_.rankSize; }	776	+ const at::Tensor &GetContextTensor() const
		777	+ {
		778	+ return engramContextTensor_;
		779	+ }
		780	+ const at::Tensor &GetLocalStorageAddrTensor() const
		781	+ {
		782	+ return localStorageAddrTensor_;
		783	+ }
		784	+ int64_t GetCommBufferSize() const
		785	+ {
		786	+ return commBufferSize_;
		787	+ }
		788	+ uint32_t GetRankSize() const
		789	+ {
		790	+ return engramCommContext_.rankSize;
		791	+ }
767		792
768	static at::Tensor EngramFetch(const at::Tensor &context, const at::Tensor &indices, int64_t hiddenSize,	793	static at::Tensor EngramFetch(const at::Tensor &context, const at::Tensor &indices, int64_t hiddenSize,
769	int64_t numEntries, int64_t dtypeEnum);	794	int64_t numEntries, int64_t dtypeEnum);
	@@ -778,7 +803,20 @@ public:
778	const at::Tensor &sendCounts, const at::Tensor &recvCounts,	803	const at::Tensor &sendCounts, const at::Tensor &recvCounts,
779	const at::Tensor &recvLocalEntry, const at::Tensor &numRecv);	804	const at::Tensor &recvLocalEntry, const at::Tensor &numRecv);
780		805
781	- bool IsWithGrad() const { return withGrad_; }	806	+ // Stateless static method for training backward (graph-mode compatible).
		807	+ // Outputs: gradUnique [maxR, H], uniqueLocalEntry [maxR], numUnique [1] (NOT narrowed).
		808	+ // Caller is responsible for narrowing by numUnique.item() outside the graph.
		809	+ using EngramFetchGradOutput = std::tuple<at::Tensor, at::Tensor, at::Tensor>;
		810	+ static EngramFetchGradOutput EngramFetchGrad(const at::Tensor &context, const at::Tensor &gradFetched,
		811	+ const at::Tensor &perm, const at::Tensor &sendCounts,
		812	+ const at::Tensor &recvCounts, const at::Tensor &recvLocalEntry,
		813	+ const at::Tensor &numRecv, int64_t numEntries, int64_t commBufferSize,
		814	+ int64_t numMaxTokensPerRank, int64_t rankSize);
		815	+
		816	+ bool IsWithGrad() const
		817	+ {
		818	+ return withGrad_;
		819	+ }
782		820
783	static int64_t GetEngramStorageSizeHint(int64_t numEntries, int64_t hiddenSize,	821	static int64_t GetEngramStorageSizeHint(int64_t numEntries, int64_t hiddenSize,
784	at::ScalarType dtype = at::kBFloat16);	822	at::ScalarType dtype = at::kBFloat16);
	@@ -1062,6 +1100,29 @@ std::tuple<at::Tensor, at::Tensor> ElasticBuffer::EngramFetchGrad(const at::Tens
1062	return {gradUnique, uniqueLocalEntry};	1100	return {gradUnique, uniqueLocalEntry};
1063	}	1101	}
1064		1102
		1103	+// EngramFetchGrad - stateless static method for training backward (graph-mode compatible).
		1104	+ElasticBuffer::EngramFetchGradOutput ElasticBuffer::EngramFetchGrad(
		1105	+ const at::Tensor &context, const at::Tensor &gradFetched, const at::Tensor &perm, const at::Tensor &sendCounts,
		1106	+ const at::Tensor &recvCounts, const at::Tensor &recvLocalEntry, const at::Tensor &numRecv, int64_t numEntries,
		1107	+ int64_t commBufferSize, int64_t numMaxTokensPerRank, int64_t rankSize)
		1108	+{
		1109	+ int64_t hidden = gradFetched.size(1);
		1110	+ TORCH_CHECK(rankSize > 0 && numMaxTokensPerRank <= INT64_MAX / rankSize,
		1111	+ "numMaxTokensPerRank * rankSize overflow, got numMaxTokensPerRank=", numMaxTokensPerRank,
		1112	+ ", rankSize=", rankSize);
		1113	+ int64_t maxR = numMaxTokensPerRank * rankSize;
		1114	+
		1115	+ auto gradUnique =
		1116	+ at::empty({maxR, hidden}, at::TensorOptions().dtype(gradFetched.dtype()).device(gradFetched.device()));
		1117	+ auto uniqueLocalEntry = at::empty({maxR}, at::TensorOptions().dtype(at::kInt).device(gradFetched.device()));
		1118	+ auto numUnique = at::empty({1}, at::TensorOptions().dtype(at::kInt).device(gradFetched.device()));
		1119	+
		1120	+ ACLNN_CMD(aclnnEngramFetchGrad, context, gradFetched, perm, sendCounts, recvCounts, recvLocalEntry, numRecv,
		1121	+ gradUnique, uniqueLocalEntry, numUnique, numEntries, commBufferSize);
		1122	+
		1123	+ return std::make_tuple(gradUnique, uniqueLocalEntry, numUnique);
		1124	+}
		1125	+
1065	// EngramBarrier - cross-rank synchronization	1126	// EngramBarrier - cross-rank synchronization
1066	void ElasticBuffer::EngramBarrier(bool useCommStream, bool withCpuSync)	1127	void ElasticBuffer::EngramBarrier(bool useCommStream, bool withCpuSync)
1067	{	1128	{
	@@ -1179,7 +1240,10 @@ public:
1179	}	1240	}
1180	}	1241	}
1181		1242
1182	- void Reset() { reinterpret_cast<volatile int64_t >(hostPtr_) = -1; }	1243	+ void Reset()
		1244	+ {
		1245	+ reinterpret_cast<volatile int64_t >(hostPtr_) = -1;
		1246	+ }
1183		1247
1184	int64_t SpinWait()	1248	int64_t SpinWait()
1185	{	1249	{
	@@ -1191,9 +1255,15 @@ public:
1191	}	1255	}
1192	}	1256	}
1193		1257
1194	- uintptr_t DevicePtr() const { return reinterpret_cast<uintptr_t>(devPtr_); }	1258	+ uintptr_t DevicePtr() const
		1259	+ {
		1260	+ return reinterpret_cast<uintptr_t>(devPtr_);
		1261	+ }
1195		1262
1196	- uintptr_t HostPtr() const { return reinterpret_cast<uintptr_t>(hostPtr_); }	1263	+ uintptr_t HostPtr() const
		1264	+ {
		1265	+ return reinterpret_cast<uintptr_t>(hostPtr_);
		1266	+ }
1197		1267
1198	private:	1268	private:
1199	void *hostPtr_ = nullptr;	1269	void *hostPtr_ = nullptr;
	@@ -1371,10 +1441,24 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
1371	pybind11::arg("rank_size"))	1441	pybind11::arg("rank_size"))
1372	.def_static("engram_fetch_wait", &Mc2Api::ElasticBuffer::EngramFetchWait, pybind11::arg("context"),	1442	.def_static("engram_fetch_wait", &Mc2Api::ElasticBuffer::EngramFetchWait, pybind11::arg("context"),
1373	pybind11::arg("fetched"))	1443	pybind11::arg("fetched"))
1374	- .def("engram_fetch_grad", &Mc2Api::ElasticBuffer::EngramFetchGrad, pybind11::arg("gradFetched").noconvert(),	1444	+ .def("engram_fetch_grad",
1375	- pybind11::arg("perm").noconvert(), pybind11::arg("sendCounts").noconvert(),	1445	+ static_cast<std::tuple<at::Tensor, at::Tensor> (Mc2Api::ElasticBuffer::*)(
1376	- pybind11::arg("recvCounts").noconvert(), pybind11::arg("recvLocalEntry").noconvert(),	1446	+ const at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &,
1377	- pybind11::arg("numRecv").noconvert())	1447	+ const at::Tensor &)>(&Mc2Api::ElasticBuffer::EngramFetchGrad),
		1448	+ pybind11::arg("gradFetched").noconvert(), pybind11::arg("perm").noconvert(),
		1449	+ pybind11::arg("sendCounts").noconvert(), pybind11::arg("recvCounts").noconvert(),
		1450	+ pybind11::arg("recvLocalEntry").noconvert(), pybind11::arg("numRecv").noconvert())
		1451	+ .def_static(
		1452	+ "engram_fetch_grad_op",
		1453	+ static_cast<Mc2Api::ElasticBuffer::EngramFetchGradOutput (*)(
		1454	+ const at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &, const at::Tensor &,
		1455	+ const at::Tensor &, const at::Tensor &, int64_t, int64_t, int64_t, int64_t)>(
		1456	+ &Mc2Api::ElasticBuffer::EngramFetchGrad),
		1457	+ pybind11::arg("context"), pybind11::arg("gradFetched").noconvert(), pybind11::arg("perm").noconvert(),
		1458	+ pybind11::arg("sendCounts").noconvert(), pybind11::arg("recvCounts").noconvert(),
		1459	+ pybind11::arg("recvLocalEntry").noconvert(), pybind11::arg("numRecv").noconvert(),
		1460	+ pybind11::arg("numEntries"), pybind11::arg("commBufferSize"), pybind11::arg("numMaxTokensPerRank"),
		1461	+ pybind11::arg("rankSize"))
1378	.def("engram_barrier", &Mc2Api::ElasticBuffer::EngramBarrier, pybind11::arg("useCommStream") = true,	1462	.def("engram_barrier", &Mc2Api::ElasticBuffer::EngramBarrier, pybind11::arg("useCommStream") = true,
1379	pybind11::arg("withCpuSync") = false)	1463	pybind11::arg("withCpuSync") = false)
1380	.def("destroy", &Mc2Api::ElasticBuffer::Destroy)	1464	.def("destroy", &Mc2Api::ElasticBuffer::Destroy)

Mmc2/common/torch_extension/elastic_buffer.py+86-7

	@@ -71,6 +71,10 @@ class ElasticBufferOpBuilder(OpBuilder):
71	"int num_max_tokens_per_rank, int comm_buffer_size, int rank_size) "	71	"int num_max_tokens_per_rank, int comm_buffer_size, int rank_size) "
72	"-> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)",	72	"-> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)",
73	"engram_fetch_wait(Tensor context, Tensor fetched) -> Tensor",	73	"engram_fetch_wait(Tensor context, Tensor fetched) -> Tensor",
		74	+ "engram_fetch_grad(Tensor context, Tensor grad_fetched, Tensor perm, Tensor send_counts, "
		75	+ "Tensor recv_counts, Tensor recv_local_entry, Tensor num_recv, int num_entries, "
		76	+ "int comm_buffer_size, int num_max_tokens_per_rank, int rank_size) "
		77	+ "-> (Tensor, Tensor, Tensor)",
74	]	78	]
75		79
76	def register_meta(self):	80	def register_meta(self):
	@@ -124,6 +128,34 @@ class ElasticBufferOpBuilder(OpBuilder):
124	def engram_fetch_wait_meta(context, fetched):	128	def engram_fetch_wait_meta(context, fetched):
125	return torch.empty_like(fetched, device="meta")	129	return torch.empty_like(fetched, device="meta")
126		130
		131	+ @impl(get_as_library(), "engram_fetch_grad", "Meta")
		132	+ def engram_fetch_grad_meta(
		133	+ context,
		134	+ grad_fetched,
		135	+ perm,
		136	+ send_counts,
		137	+ recv_counts,
		138	+ recv_local_entry,
		139	+ num_recv,
		140	+ num_entries,
		141	+ comm_buffer_size,
		142	+ num_max_tokens_per_rank,
		143	+ rank_size,
		144	+ ):
		145	+ # Outputs are NOT narrowed; full maxR rows. Caller narrows by num_unique.
		146	+ num_max_tokens_per_rank = (
		147	+ 1 if num_max_tokens_per_rank <= 0 else num_max_tokens_per_rank
		148	+ )
		149	+ max_r = num_max_tokens_per_rank * rank_size
		150	+ grad_unique = torch.empty(
		151	+ (max_r, grad_fetched.size(1)),
		152	+ dtype=grad_fetched.dtype,
		153	+ device="meta",
		154	+ )
		155	+ unique_local_entry = torch.empty((max_r,), dtype=torch.int32, device="meta")
		156	+ num_unique = torch.empty((1,), dtype=torch.int32, device="meta")
		157	+ return (grad_unique, unique_local_entry, num_unique)
		158	+
127	def extra_ldflags(self):	159	def extra_ldflags(self):
128	"""Extra link flags for HCCL and ACL libraries."""	160	"""Extra link flags for HCCL and ACL libraries."""
129	flags = super().extra_ldflags()	161	flags = super().extra_ldflags()
	@@ -167,6 +199,36 @@ def engram_fetch_wait(context, fetched):
167	return op_module.ElasticBuffer.engram_fetch_wait(context, fetched)	199	return op_module.ElasticBuffer.engram_fetch_wait(context, fetched)
168		200
169		201
		202	+@impl(get_as_library(), "engram_fetch_grad", "PrivateUse1")
		203	+def engram_fetch_grad(
		204	+ context,
		205	+ grad_fetched,
		206	+ perm,
		207	+ send_counts,
		208	+ recv_counts,
		209	+ recv_local_entry,
		210	+ num_recv,
		211	+ num_entries,
		212	+ comm_buffer_size,
		213	+ num_max_tokens_per_rank,
		214	+ rank_size,
		215	+):
		216	+ op_module = _elastic_buffer_op_builder.load()
		217	+ return op_module.ElasticBuffer.engram_fetch_grad_op(
		218	+ context,
		219	+ grad_fetched,
		220	+ perm,
		221	+ send_counts,
		222	+ recv_counts,
		223	+ recv_local_entry,
		224	+ num_recv,
		225	+ num_entries,
		226	+ comm_buffer_size,
		227	+ num_max_tokens_per_rank,
		228	+ rank_size,
		229	+ )
		230	+
		231	+
170	def _inline_align(value: int, base: int) -> int:	232	def _inline_align(value: int, base: int) -> int:
171	return (value + base - 1) // base * base	233	return (value + base - 1) // base * base
172		234
	@@ -745,14 +807,31 @@ class ElasticBuffer:
745	f"fetch_ctx.num_recv length ({fetch_ctx.num_recv.size(0)}) must be 1"	807	f"fetch_ctx.num_recv length ({fetch_ctx.num_recv.size(0)}) must be 1"
746	),	808	),
747	)	809	)
748	- return self._runtime.engram_fetch_grad(	810	+ grad_unique_full, unique_local_entry_full, num_unique = (
749	- grad_fetched,	811	+ torch.ops.cann_ops_transformer.engram_fetch_grad(
750	- fetch_ctx.perm,	812	+ self._engram_context_tensor,
751	- fetch_ctx.send_counts,	813	+ grad_fetched,
752	- fetch_ctx.recv_counts,	814	+ fetch_ctx.perm,
753	- fetch_ctx.recv_local_entry,	815	+ fetch_ctx.send_counts,
754	- fetch_ctx.num_recv,	816	+ fetch_ctx.recv_counts,
		817	+ fetch_ctx.recv_local_entry,
		818	+ fetch_ctx.num_recv,
		819	+ self._engram_num_entries,
		820	+ self._comm_buffer_size,
		821	+ self._num_max_tokens_per_rank,
		822	+ self._rank_size,
		823	+ )
755	)	824	)
		825	+ # Narrow to actualK rows outside the graph (.item() syncs; not graph-safe).
		826	+ actual_k = int(num_unique.item())
		827	+ max_r = self._num_max_tokens_per_rank * self._rank_size
		828	+ if actual_k < max_r:
		829	+ grad_unique = grad_unique_full.narrow(0, 0, actual_k)
		830	+ unique_local_entry = unique_local_entry_full.narrow(0, 0, actual_k)
		831	+ else:
		832	+ grad_unique = grad_unique_full
		833	+ unique_local_entry = unique_local_entry_full
		834	+ return grad_unique, unique_local_entry
756		835
757	def dispatch(	836	def dispatch(
758	self,	837	self,

Mmc2/engram_fetch/op_kernel/arch35/engram_fetch_train_arch35.h+9-6

	@@ -197,8 +197,8 @@ __aicore__ inline void EngramFetchTrainArch35::DrainChecked(uint64_t handle)
197	__aicore__ inline void EngramFetchTrainArch35::TimeoutCheck(uint64_t startTime)	197	__aicore__ inline void EngramFetchTrainArch35::TimeoutCheck(uint64_t startTime)
198	{	198	{
199	uint64_t nowUs = static_cast<uint64_t>(AscendC::GetSystemCycle()) / ENGRAM_CYCLES_PER_US;	199	uint64_t nowUs = static_cast<uint64_t>(AscendC::GetSystemCycle()) / ENGRAM_CYCLES_PER_US;
200	- ascendc_assert((nowUs - startTime) < ENGRAM_TIMEOUT_US,	200	+ ascendc_assert((nowUs - startTime) < ENGRAM_TIMEOUT_US, "timeout, rankId=%u, aivId=%u, elapsed=%llu us", rankId_,
201	- "timeout, rankId=%u, aivId=%u, elapsed=%llu us", rankId_, aivId_, nowUs - startTime);	201	+ aivId_, nowUs - startTime);
202	}	202	}
203		203
204	__aicore__ inline GM_ADDR EngramFetchTrainArch35::GetRemoteWinAddr(uint32_t dstRank, uint64_t offset)	204	__aicore__ inline GM_ADDR EngramFetchTrainArch35::GetRemoteWinAddr(uint32_t dstRank, uint64_t offset)
	@@ -208,7 +208,10 @@ __aicore__ inline GM_ADDR EngramFetchTrainArch35::GetRemoteWinAddr(uint32_t dstR
208		208
209	__aicore__ inline uint64_t EngramFetchTrainArch35::GetCommHandle(uint32_t dstRank)	209	__aicore__ inline uint64_t EngramFetchTrainArch35::GetCommHandle(uint32_t dstRank)
210	{	210	{
211	- uint32_t channelIdx = isSender_ ? SENDER_CHANNEL_IDX : RECEIVER_CHANNEL_IDX;	211	+ uint32_t channelIdx = aivId_ / numRanks_ + 1U;
		212	+ if (channelIdx >= channelsPerRank_) {
		213	+ channelIdx = 1U;
		214	+ }
212	return ctxPtr_->hcommHandle[dstRank * channelsPerRank_ + channelIdx];	215	return ctxPtr_->hcommHandle[dstRank * channelsPerRank_ + channelIdx];
213	}	216	}
214		217
	@@ -1049,9 +1052,9 @@ __aicore__ inline void EngramFetchTrainArch35::LocalReadTable()
1049	if (globalIdx < 0) {	1052	if (globalIdx < 0) {
1050	continue;	1053	continue;
1051	}	1054	}
1052	- uint32_t localEntryIdx = static_cast<uint32_t>(	1055	+ uint32_t localEntryIdx =
1053	- static_cast<int64_t>(globalIdx) -	1056	+ static_cast<uint32_t>(static_cast<int64_t>(globalIdx) -
1054	- static_cast<int64_t>(rankId_) * static_cast<int64_t>(numEntriesPerRank_));	1057	+ static_cast<int64_t>(rankId_) * static_cast<int64_t>(numEntriesPerRank_));
1055	if (localEntryIdx >= static_cast<uint32_t>(numEntriesPerRank_)) {	1058	if (localEntryIdx >= static_cast<uint32_t>(numEntriesPerRank_)) {
1056	continue;	1059	continue;
1057	}	1060	}

Mmc2/engram_fetch/op_kernel/engram_fetch_utils.h+1-1

	@@ -45,7 +45,7 @@ struct EngramCommContext {
45	uint32_t rankId;	45	uint32_t rankId;
46	uint32_t rankSize;	46	uint32_t rankSize;
47	uint64_t commBuffer[HCCL_MAX_RANK_SIZE];	47	uint64_t commBuffer[HCCL_MAX_RANK_SIZE];
48	- uint64_t hcommHandle[HCCL_MAX_RANK_SIZE];	48	+ uint64_t hcommHandle[HCCL_MAX_RANK_SIZE * 2];
49	uint32_t channelsPerRank;	49	uint32_t channelsPerRank;
50	};	50	};
51		51

Mmc2/engram_fetch_grad/op_kernel/arch35/engram_fetch_grad_arch35.h+453-158

Mmc2/engram_fetch_grad/op_kernel/arch35/engram_fetch_grad_unique.h+70-69

	@@ -24,49 +24,49 @@ __aicore__ inline void SyncFunc(AscendC::TPipe &pipe)
24	AscendC::WaitFlag<event>(eventID);	24	AscendC::WaitFlag<event>(eventID);
25	}	25	}
26		26
		27	+constexpr uint32_t MAX_BLOCK_BYTES = 65535U;
		28	+
27	class EngramFetchGradUnique {	29	class EngramFetchGradUnique {
28	public:	30	public:
29	__aicore__ inline EngramFetchGradUnique() = default;	31	__aicore__ inline EngramFetchGradUnique() = default;
30		32
31	- __aicore__ inline void Init(	33	+ __aicore__ inline void Init(uint32_t aivId, uint32_t totalBlocks, uint32_t rankId, uint32_t numRanks,
32	- uint32_t aivId, uint32_t totalBlocks, uint32_t rankId, uint32_t numRanks,	34	+ int32_t numEntriesPerRank, int64_t hiddenDim, int64_t hiddenBytes, int32_t inputDtype,
33	- int32_t numEntriesPerRank, int64_t hiddenDim, int64_t hiddenBytes,	35	+ int32_t outputDtype, AscendC::TPipe *pipe, AscendC::TBuf<> &pingBuf,
34	- int32_t inputDtype, int32_t outputDtype, AscendC::TPipe *pipe,	36	+ AscendC::TBuf<> &pongBuf, AscendC::TBuf<> &indicesBuf, AscendC::TBuf<> &tempBuf,
35	- AscendC::TBuf<> &pingBuf, AscendC::TBuf<> &pongBuf, AscendC::TBuf<> &indicesBuf,	37	+ AscendC::TBuf<> &statusBuf, AscendC::TBuf<> &gradSumBuf);
36	- AscendC::TBuf<> &tempBuf, AscendC::TBuf<> &statusBuf, AscendC::TBuf<> &gradSumBuf);
37		38
38	- __aicore__ inline void Run(uint32_t numRecv,	39	+ __aicore__ inline void Run(uint32_t numRecv, GM_ADDR recvLocalEntryOutGM, GM_ADDR uniqueLocalEntryOutGM,
39	- GM_ADDR recvLocalEntryOutGM, GM_ADDR uniqueLocalEntryOutGM, GM_ADDR numUniqueOutGM,	40	+ GM_ADDR numUniqueOutGM, GM_ADDR gradUniqueOutGM, GM_ADDR gradUniqueFp32GM,
40	- GM_ADDR gradUniqueOutGM, GM_ADDR gradUniqueFp32GM, GM_ADDR recvGradGM,	41	+ GM_ADDR recvGradGM, GM_ADDR coreStartGM, GM_ADDR segCountGM, GM_ADDR sortCompanionGM);
41	- GM_ADDR coreStartGM, GM_ADDR segCountGM, GM_ADDR sortCompanionGM);
42		42
43	__aicore__ inline void WriteNumUniqueZero(GM_ADDR numUniqueOutGM);	43	__aicore__ inline void WriteNumUniqueZero(GM_ADDR numUniqueOutGM);
44		44
45	private:	45	private:
46	__aicore__ inline void ZeroGradUnique(uint32_t numRecv, GM_ADDR gradUniqueOutGM, GM_ADDR gradUniqueFp32GM);	46	__aicore__ inline void ZeroGradUnique(uint32_t numRecv, GM_ADDR gradUniqueOutGM, GM_ADDR gradUniqueFp32GM);
47	- __aicore__ inline void CountUniquesParallel(uint32_t numRecv, GM_ADDR recvLocalEntryOutGM,	47	+ __aicore__ inline void CountUniquesParallel(uint32_t numRecv, GM_ADDR recvLocalEntryOutGM, GM_ADDR coreStartGM,
48	- GM_ADDR coreStartGM, GM_ADDR segCountGM);	48	+ GM_ADDR segCountGM);
49	__aicore__ inline void FixCoreStartForSkippedCores(GM_ADDR coreStartGM, GM_ADDR segCountGM);	49	__aicore__ inline void FixCoreStartForSkippedCores(GM_ADDR coreStartGM, GM_ADDR segCountGM);
50	__aicore__ inline void ScatterAddAtomicParallel(uint32_t numRecv, GM_ADDR recvLocalEntryOutGM,	50	__aicore__ inline void ScatterAddAtomicParallel(uint32_t numRecv, GM_ADDR recvLocalEntryOutGM,
51	GM_ADDR uniqueLocalEntryOutGM, GM_ADDR gradUniqueOutGM,	51	GM_ADDR uniqueLocalEntryOutGM, GM_ADDR gradUniqueOutGM,
52	- GM_ADDR gradUniqueFp32GM, GM_ADDR recvGradGM,	52	+ GM_ADDR gradUniqueFp32GM, GM_ADDR recvGradGM, GM_ADDR coreStartGM,
53	- GM_ADDR coreStartGM, GM_ADDR segCountGM, GM_ADDR sortCompanionGM);	53	+ GM_ADDR segCountGM, GM_ADDR sortCompanionGM);
54	- __aicore__ inline void CastGradUniqueToOutput(GM_ADDR segCountGM, GM_ADDR numUniqueOutGM,	54	+ __aicore__ inline void CastGradUniqueToOutput(GM_ADDR segCountGM, GM_ADDR numUniqueOutGM, GM_ADDR gradUniqueOutGM,
55	- GM_ADDR gradUniqueOutGM, GM_ADDR gradUniqueFp32GM);	55	+ GM_ADDR gradUniqueFp32GM);
56	__aicore__ inline void LoadCoreRange(uint32_t numRecv, GM_ADDR coreStartGM, uint32_t &start, uint32_t &end);	56	__aicore__ inline void LoadCoreRange(uint32_t numRecv, GM_ADDR coreStartGM, uint32_t &start, uint32_t &end);
57	__aicore__ inline int32_t ComputePreCoreOffset(GM_ADDR segCountGM);	57	__aicore__ inline int32_t ComputePreCoreOffset(GM_ADDR segCountGM);
58	- __aicore__ inline uint32_t ProcessScatterBatch(	58	+ __aicore__ inline uint32_t ProcessScatterBatch(uint32_t cur, uint32_t end, int32_t &runningOffset,
59	- uint32_t cur, uint32_t end, int32_t &runningOffset, int32_t &runningUniqueOffset,	59	+ int32_t &runningUniqueOffset, int32_t &prevEntry,
60	- int32_t &prevEntry, bool &isFirstElement,	60	+ bool &isFirstElement, GM_ADDR recvLocalEntryOutGM,
61	- GM_ADDR recvLocalEntryOutGM, GM_ADDR uniqueLocalEntryOutGM, GM_ADDR gradUniqueOutGM,	61	+ GM_ADDR uniqueLocalEntryOutGM, GM_ADDR gradUniqueOutGM,
62	- GM_ADDR gradUniqueFp32GM, GM_ADDR recvGradGM, GM_ADDR sortCompanionGM);	62	+ GM_ADDR gradUniqueFp32GM, GM_ADDR recvGradGM,
63	- __aicore__ inline void ScatterGradSubBatch(uint32_t subStart, uint32_t subLen,	63	+ GM_ADDR sortCompanionGM);
64	- GM_ADDR gradUniqueOutGM, GM_ADDR gradUniqueFp32GM, GM_ADDR recvGradGM);	64	+ __aicore__ inline void ScatterGradSubBatch(uint32_t subStart, uint32_t subLen, GM_ADDR gradUniqueOutGM,
		65	+ GM_ADDR gradUniqueFp32GM, GM_ADDR recvGradGM);
65	__aicore__ inline void AtomicAddGrad(AscendC::GlobalTensor<float> &atomicDstGM, AscendC::LocalTensor<float> &srcT,	66	__aicore__ inline void AtomicAddGrad(AscendC::GlobalTensor<float> &atomicDstGM, AscendC::LocalTensor<float> &srcT,
66	uint32_t subStart, uint32_t subLen);	67	uint32_t subStart, uint32_t subLen);
67	__aicore__ inline uint32_t ProcessBatchUnique(uint32_t cur, uint32_t batchLen, int32_t runningOffset,	68	__aicore__ inline uint32_t ProcessBatchUnique(uint32_t cur, uint32_t batchLen, int32_t runningOffset,
68	- int32_t &prevEntry, bool &isFirstElement,	69	+ int32_t &prevEntry, bool &isFirstElement, int32_t &inclusiveSum,
69	- int32_t &inclusiveSum,
70	GM_ADDR recvLocalEntryOutGM, GM_ADDR sortCompanionGM);	70	GM_ADDR recvLocalEntryOutGM, GM_ADDR sortCompanionGM);
71	__aicore__ inline void WriteBatchUnique(uint32_t tileUniqueCnt, int32_t &runningUniqueOffset,	71	__aicore__ inline void WriteBatchUnique(uint32_t tileUniqueCnt, int32_t &runningUniqueOffset,
72	GM_ADDR uniqueLocalEntryOutGM);	72	GM_ADDR uniqueLocalEntryOutGM);
	@@ -89,12 +89,13 @@ private:
89	AscendC::TBuf<> *gradSumBuf_{nullptr};	89	AscendC::TBuf<> *gradSumBuf_{nullptr};
90	};	90	};
91		91
92	-__aicore__ inline void EngramFetchGradUnique::Init(	92	+__aicore__ inline void EngramFetchGradUnique::Init(uint32_t aivId, uint32_t totalBlocks, uint32_t rankId,
93	- uint32_t aivId, uint32_t totalBlocks, uint32_t rankId, uint32_t numRanks,	93	+ uint32_t numRanks, int32_t numEntriesPerRank, int64_t hiddenDim,
94	- int32_t numEntriesPerRank, int64_t hiddenDim, int64_t hiddenBytes,	94	+ int64_t hiddenBytes, int32_t inputDtype, int32_t outputDtype,
95	- int32_t inputDtype, int32_t outputDtype, AscendC::TPipe *pipe,	95	+ AscendC::TPipe *pipe, AscendC::TBuf<> &pingBuf,
96	- AscendC::TBuf<> &pingBuf, AscendC::TBuf<> &pongBuf, AscendC::TBuf<> &indicesBuf,	96	+ AscendC::TBuf<> &pongBuf, AscendC::TBuf<> &indicesBuf,
97	- AscendC::TBuf<> &tempBuf, AscendC::TBuf<> &statusBuf, AscendC::TBuf<> &gradSumBuf)	97	+ AscendC::TBuf<> &tempBuf, AscendC::TBuf<> &statusBuf,
		98	+ AscendC::TBuf<> &gradSumBuf)
98	{	99	{
99	aivId_ = aivId;	100	aivId_ = aivId;
100	totalBlocks_ = totalBlocks;	101	totalBlocks_ = totalBlocks;
	@@ -114,11 +115,10 @@ __aicore__ inline void EngramFetchGradUnique::Init(
114	gradSumBuf_ = &gradSumBuf;	115	gradSumBuf_ = &gradSumBuf;
115	}	116	}
116		117
117	-__aicore__ inline void EngramFetchGradUnique::Run(	118	+__aicore__ inline void EngramFetchGradUnique::Run(uint32_t numRecv, GM_ADDR recvLocalEntryOutGM,
118	- uint32_t numRecv,	119	+ GM_ADDR uniqueLocalEntryOutGM, GM_ADDR numUniqueOutGM,
119	- GM_ADDR recvLocalEntryOutGM, GM_ADDR uniqueLocalEntryOutGM, GM_ADDR numUniqueOutGM,	120	+ GM_ADDR gradUniqueOutGM, GM_ADDR gradUniqueFp32GM, GM_ADDR recvGradGM,
120	- GM_ADDR gradUniqueOutGM, GM_ADDR gradUniqueFp32GM, GM_ADDR recvGradGM,	121	+ GM_ADDR coreStartGM, GM_ADDR segCountGM, GM_ADDR sortCompanionGM)
121	- GM_ADDR coreStartGM, GM_ADDR segCountGM, GM_ADDR sortCompanionGM)
122	{	122	{
123	ZeroGradUnique(numRecv, gradUniqueOutGM, gradUniqueFp32GM);	123	ZeroGradUnique(numRecv, gradUniqueOutGM, gradUniqueFp32GM);
124	CountUniquesParallel(numRecv, recvLocalEntryOutGM, coreStartGM, segCountGM);	124	CountUniquesParallel(numRecv, recvLocalEntryOutGM, coreStartGM, segCountGM);
	@@ -127,9 +127,8 @@ __aicore__ inline void EngramFetchGradUnique::Run(
127	FixCoreStartForSkippedCores(coreStartGM, segCountGM);	127	FixCoreStartForSkippedCores(coreStartGM, segCountGM);
128	AscendC::SyncAll<true>();	128	AscendC::SyncAll<true>();
129		129
130	- ScatterAddAtomicParallel(numRecv, recvLocalEntryOutGM, uniqueLocalEntryOutGM,	130	+ ScatterAddAtomicParallel(numRecv, recvLocalEntryOutGM, uniqueLocalEntryOutGM, gradUniqueOutGM, gradUniqueFp32GM,
131	- gradUniqueOutGM, gradUniqueFp32GM, recvGradGM,	131	+ recvGradGM, coreStartGM, segCountGM, sortCompanionGM);
132	- coreStartGM, segCountGM, sortCompanionGM);
133	AscendC::SyncAll<true>();	132	AscendC::SyncAll<true>();
134		133
135	CastGradUniqueToOutput(segCountGM, numUniqueOutGM, gradUniqueOutGM, gradUniqueFp32GM);	134	CastGradUniqueToOutput(segCountGM, numUniqueOutGM, gradUniqueOutGM, gradUniqueFp32GM);
	@@ -168,6 +167,10 @@ __aicore__ inline void EngramFetchGradUnique::ZeroGradUnique(uint32_t numRecv, G
168	if (start < end) {	167	if (start < end) {
169	AscendC::LocalTensor<float> zeroBuf = pongBuf_->Get<float>();	168	AscendC::LocalTensor<float> zeroBuf = pongBuf_->Get<float>();
170	uint32_t tileFloats = Mc2Kernel::TILE_BYTES / sizeof(float);	169	uint32_t tileFloats = Mc2Kernel::TILE_BYTES / sizeof(float);
		170	+ constexpr uint32_t maxBlockFloats = MAX_BLOCK_BYTES / sizeof(float);
		171	+ if (tileFloats > maxBlockFloats) {
		172	+ tileFloats = maxBlockFloats;
		173	+ }
171	AscendC::Duplicate<float>(zeroBuf, 0.0f, tileFloats);	174	AscendC::Duplicate<float>(zeroBuf, 0.0f, tileFloats);
172	SyncFunc<AscendC::HardEvent::V_MTE3>(*pipe_);	175	SyncFunc<AscendC::HardEvent::V_MTE3>(*pipe_);
173		176
	@@ -186,8 +189,8 @@ __aicore__ inline void EngramFetchGradUnique::ZeroGradUnique(uint32_t numRecv, G
186	}	189	}
187	}	190	}
188		191
189	-__aicore__ inline void EngramFetchGradUnique::CountUniquesParallel(	192	+__aicore__ inline void EngramFetchGradUnique::CountUniquesParallel(uint32_t numRecv, GM_ADDR recvLocalEntryOutGM,
190	- uint32_t numRecv, GM_ADDR recvLocalEntryOutGM, GM_ADDR coreStartGM, GM_ADDR segCountGM)	193	+ GM_ADDR coreStartGM, GM_ADDR segCountGM)
191	{	194	{
192	uint32_t chunk = (numRecv + totalBlocks_ - 1U) / totalBlocks_;	195	uint32_t chunk = (numRecv + totalBlocks_ - 1U) / totalBlocks_;
193	uint32_t rawStart = aivId_ * chunk;	196	uint32_t rawStart = aivId_ * chunk;
	@@ -307,8 +310,8 @@ __aicore__ inline void EngramFetchGradUnique::FixCoreStartForSkippedCores(GM_ADD
307	}	310	}
308	}	311	}
309		312
310	-__aicore__ inline void EngramFetchGradUnique::LoadCoreRange(uint32_t numRecv, GM_ADDR coreStartGM,	313	+__aicore__ inline void EngramFetchGradUnique::LoadCoreRange(uint32_t numRecv, GM_ADDR coreStartGM, uint32_t &start,
311	- uint32_t &start, uint32_t &end)	314	+ uint32_t &end)
312	{	315	{
313	AscendC::GlobalTensor<int32_t> coreStartGM_;	316	AscendC::GlobalTensor<int32_t> coreStartGM_;
314	coreStartGM_.SetGlobalBuffer((__gm__ int32_t *)coreStartGM);	317	coreStartGM_.SetGlobalBuffer((__gm__ int32_t *)coreStartGM);
	@@ -351,12 +354,12 @@ __aicore__ inline int32_t EngramFetchGradUnique::ComputePreCoreOffset(GM_ADDR se
351	return preCoreOffset;	354	return preCoreOffset;
352	}	355	}
353		356
354	-__aicore__ inline void EngramFetchGradUnique::AtomicAddGrad(	357	+__aicore__ inline void EngramFetchGradUnique::AtomicAddGrad(AscendC::GlobalTensor<float> &atomicDstGM,
355	- AscendC::GlobalTensor<float> &atomicDstGM, AscendC::LocalTensor<float> &srcT, uint32_t subStart, uint32_t subLen)	358	+ AscendC::LocalTensor<float> &srcT, uint32_t subStart,
		359	+ uint32_t subLen)
356	{	360	{
357	AscendC::LocalTensor<int32_t> pingInt32 = pingBuf_->Get<int32_t>();	361	AscendC::LocalTensor<int32_t> pingInt32 = pingBuf_->Get<int32_t>();
358	AscendC::LocalTensor<int32_t> compactIdxBuf = pingInt32[3 * Mc2Kernel::ENTRY_BATCH_CAP];	362	AscendC::LocalTensor<int32_t> compactIdxBuf = pingInt32[3 * Mc2Kernel::ENTRY_BATCH_CAP];
359	- constexpr uint32_t MAX_BLOCK_BYTES = 65535U;
360	uint32_t totalBytes = static_cast<uint32_t>(hiddenDim_) * sizeof(float);	363	uint32_t totalBytes = static_cast<uint32_t>(hiddenDim_) * sizeof(float);
361	uint32_t numBlocks = (totalBytes + MAX_BLOCK_BYTES - 1U) / MAX_BLOCK_BYTES;	364	uint32_t numBlocks = (totalBytes + MAX_BLOCK_BYTES - 1U) / MAX_BLOCK_BYTES;
362	for (uint32_t j = 0; j < subLen; j++) {	365	for (uint32_t j = 0; j < subLen; j++) {
	@@ -366,8 +369,7 @@ __aicore__ inline void EngramFetchGradUnique::AtomicAddGrad(
366	uint32_t byteOffset = b * MAX_BLOCK_BYTES;	369	uint32_t byteOffset = b * MAX_BLOCK_BYTES;
367	uint32_t elemOffset = byteOffset / sizeof(float);	370	uint32_t elemOffset = byteOffset / sizeof(float);
368	uint32_t remaining = totalBytes - byteOffset;	371	uint32_t remaining = totalBytes - byteOffset;
369	- uint16_t blkBytes = static_cast<uint16_t>(	372	+ uint16_t blkBytes = static_cast<uint16_t>(remaining > MAX_BLOCK_BYTES ? MAX_BLOCK_BYTES : remaining);
370	- remaining > MAX_BLOCK_BYTES ? MAX_BLOCK_BYTES : remaining);
371	AscendC::DataCopyParams atomicParams{1U, blkBytes, 0U, 0U};	373	AscendC::DataCopyParams atomicParams{1U, blkBytes, 0U, 0U};
372	AscendC::DataCopyPad(atomicDstGM[static_cast<uint64_t>(compactIdx) * hiddenDim_ + elemOffset],	374	AscendC::DataCopyPad(atomicDstGM[static_cast<uint64_t>(compactIdx) * hiddenDim_ + elemOffset],
373	srcT[static_cast<uint32_t>(j) * static_cast<uint32_t>(hiddenDim_) + elemOffset],	375	srcT[static_cast<uint32_t>(j) * static_cast<uint32_t>(hiddenDim_) + elemOffset],
	@@ -377,8 +379,9 @@ __aicore__ inline void EngramFetchGradUnique::AtomicAddGrad(
377	}	379	}
378	}	380	}
379		381
380	-__aicore__ inline void EngramFetchGradUnique::ScatterGradSubBatch(	382	+__aicore__ inline void EngramFetchGradUnique::ScatterGradSubBatch(uint32_t subStart, uint32_t subLen,
381	- uint32_t subStart, uint32_t subLen, GM_ADDR gradUniqueOutGM, GM_ADDR gradUniqueFp32GM, GM_ADDR recvGradGM)	383	+ GM_ADDR gradUniqueOutGM, GM_ADDR gradUniqueFp32GM,
		384	+ GM_ADDR recvGradGM)
382	{	385	{
383	AscendC::LocalTensor<int32_t> pingInt32 = pingBuf_->Get<int32_t>();	386	AscendC::LocalTensor<int32_t> pingInt32 = pingBuf_->Get<int32_t>();
384	AscendC::LocalTensor<int32_t> recvIdxBuf = pingInt32[2 * Mc2Kernel::ENTRY_BATCH_CAP];	387	AscendC::LocalTensor<int32_t> recvIdxBuf = pingInt32[2 * Mc2Kernel::ENTRY_BATCH_CAP];
	@@ -395,8 +398,8 @@ __aicore__ inline void EngramFetchGradUnique::ScatterGradSubBatch(
395	AscendC::DataCopyExtParams gradParams{1U, static_cast<uint32_t>(hiddenBytes_), 0U, 0U, 0U};	398	AscendC::DataCopyExtParams gradParams{1U, static_cast<uint32_t>(hiddenBytes_), 0U, 0U, 0U};
396	AscendC::GlobalTensor<uint8_t> gradSrcGM;	399	AscendC::GlobalTensor<uint8_t> gradSrcGM;
397	gradSrcGM.SetGlobalBuffer((__gm__ uint8_t *)gradAddr);	400	gradSrcGM.SetGlobalBuffer((__gm__ uint8_t *)gradAddr);
398	- AscendC::DataCopyPad(gradRaw[static_cast<uint32_t>(j) * static_cast<uint32_t>(hiddenBytes_)],	401	+ AscendC::DataCopyPad(gradRaw[static_cast<uint32_t>(j) * static_cast<uint32_t>(hiddenBytes_)], gradSrcGM,
399	- gradSrcGM, gradParams, gradPad);	402	+ gradParams, gradPad);
400	}	403	}
401		404
402	if (inputDtype_ == Mc2Kernel::ENGRAM_DT_FLOAT) {	405	if (inputDtype_ == Mc2Kernel::ENGRAM_DT_FLOAT) {
	@@ -421,10 +424,11 @@ __aicore__ inline void EngramFetchGradUnique::ScatterGradSubBatch(
421	}	424	}
422	}	425	}
423		426
424	-__aicore__ inline uint32_t EngramFetchGradUnique::ProcessBatchUnique(	427	+__aicore__ inline uint32_t EngramFetchGradUnique::ProcessBatchUnique(uint32_t cur, uint32_t batchLen,
425	- uint32_t cur, uint32_t batchLen, int32_t runningOffset,	428	+ int32_t runningOffset, int32_t &prevEntry,
426	- int32_t &prevEntry, bool &isFirstElement, int32_t &inclusiveSum,	429	+ bool &isFirstElement, int32_t &inclusiveSum,
427	- GM_ADDR recvLocalEntryOutGM, GM_ADDR sortCompanionGM)	430	+ GM_ADDR recvLocalEntryOutGM,
		431	+ GM_ADDR sortCompanionGM)
428	{	432	{
429	AscendC::LocalTensor<int32_t> entryUb = indicesBuf_->Get<int32_t>();	433	AscendC::LocalTensor<int32_t> entryUb = indicesBuf_->Get<int32_t>();
430	AscendC::LocalTensor<int32_t> pingInt32 = pingBuf_->Get<int32_t>();	434	AscendC::LocalTensor<int32_t> pingInt32 = pingBuf_->Get<int32_t>();
	@@ -476,9 +480,8 @@ __aicore__ inline void EngramFetchGradUnique::WriteBatchUnique(uint32_t tileUniq
476	}	480	}
477		481
478	__aicore__ inline uint32_t EngramFetchGradUnique::ProcessScatterBatch(	482	__aicore__ inline uint32_t EngramFetchGradUnique::ProcessScatterBatch(
479	- uint32_t cur, uint32_t end, int32_t &runningOffset, int32_t &runningUniqueOffset,	483	+ uint32_t cur, uint32_t end, int32_t &runningOffset, int32_t &runningUniqueOffset, int32_t &prevEntry,
480	- int32_t &prevEntry, bool &isFirstElement,	484	+ bool &isFirstElement, GM_ADDR recvLocalEntryOutGM, GM_ADDR uniqueLocalEntryOutGM, GM_ADDR gradUniqueOutGM,
481	- GM_ADDR recvLocalEntryOutGM, GM_ADDR uniqueLocalEntryOutGM, GM_ADDR gradUniqueOutGM,
482	GM_ADDR gradUniqueFp32GM, GM_ADDR recvGradGM, GM_ADDR sortCompanionGM)	485	GM_ADDR gradUniqueFp32GM, GM_ADDR recvGradGM, GM_ADDR sortCompanionGM)
483	{	486	{
484	uint32_t batchLen = end - cur;	487	uint32_t batchLen = end - cur;
	@@ -487,9 +490,8 @@ __aicore__ inline uint32_t EngramFetchGradUnique::ProcessScatterBatch(
487	}	490	}
488		491
489	int32_t inclusiveSum = 0;	492	int32_t inclusiveSum = 0;
490	- uint32_t tileUniqueCnt = ProcessBatchUnique(	493	+ uint32_t tileUniqueCnt = ProcessBatchUnique(cur, batchLen, runningOffset, prevEntry, isFirstElement, inclusiveSum,
491	- cur, batchLen, runningOffset, prevEntry, isFirstElement, inclusiveSum,	494	+ recvLocalEntryOutGM, sortCompanionGM);
492	- recvLocalEntryOutGM, sortCompanionGM);
493		495
494	uint32_t maxGradPerBatch = Mc2Kernel::TILE_BYTES / static_cast<uint32_t>(hiddenBytes_);	496	uint32_t maxGradPerBatch = Mc2Kernel::TILE_BYTES / static_cast<uint32_t>(hiddenBytes_);
495	if (maxGradPerBatch < 1U) {	497	if (maxGradPerBatch < 1U) {
	@@ -514,9 +516,8 @@ __aicore__ inline uint32_t EngramFetchGradUnique::ProcessScatterBatch(
514	}	516	}
515		517
516	__aicore__ inline void EngramFetchGradUnique::ScatterAddAtomicParallel(	518	__aicore__ inline void EngramFetchGradUnique::ScatterAddAtomicParallel(
517	- uint32_t numRecv, GM_ADDR recvLocalEntryOutGM, GM_ADDR uniqueLocalEntryOutGM,	519	+ uint32_t numRecv, GM_ADDR recvLocalEntryOutGM, GM_ADDR uniqueLocalEntryOutGM, GM_ADDR gradUniqueOutGM,
518	- GM_ADDR gradUniqueOutGM, GM_ADDR gradUniqueFp32GM, GM_ADDR recvGradGM,	520	+ GM_ADDR gradUniqueFp32GM, GM_ADDR recvGradGM, GM_ADDR coreStartGM, GM_ADDR segCountGM, GM_ADDR sortCompanionGM)
519	- GM_ADDR coreStartGM, GM_ADDR segCountGM, GM_ADDR sortCompanionGM)
520	{	521	{
521	uint32_t start;	522	uint32_t start;
522	uint32_t end;	523	uint32_t end;
	@@ -534,13 +535,13 @@ __aicore__ inline void EngramFetchGradUnique::ScatterAddAtomicParallel(
534	uint32_t cur = start;	535	uint32_t cur = start;
535	while (cur < end) {	536	while (cur < end) {
536	cur = ProcessScatterBatch(cur, end, runningOffset, runningUniqueOffset, prevEntry, isFirstElement,	537	cur = ProcessScatterBatch(cur, end, runningOffset, runningUniqueOffset, prevEntry, isFirstElement,
537	- recvLocalEntryOutGM, uniqueLocalEntryOutGM, gradUniqueOutGM,	538	+ recvLocalEntryOutGM, uniqueLocalEntryOutGM, gradUniqueOutGM, gradUniqueFp32GM,
538	- gradUniqueFp32GM, recvGradGM, sortCompanionGM);	539	+ recvGradGM, sortCompanionGM);
539	}	540	}
540	}	541	}
541		542
542	-__aicore__ inline void EngramFetchGradUnique::CastGradUniqueToOutput(	543	+__aicore__ inline void EngramFetchGradUnique::CastGradUniqueToOutput(GM_ADDR segCountGM, GM_ADDR numUniqueOutGM,
543	- GM_ADDR segCountGM, GM_ADDR numUniqueOutGM, GM_ADDR gradUniqueOutGM, GM_ADDR gradUniqueFp32GM)	544	+ GM_ADDR gradUniqueOutGM, GM_ADDR gradUniqueFp32GM)
544	{	545	{
545	AscendC::GlobalTensor<int32_t> segCountGM_;	546	AscendC::GlobalTensor<int32_t> segCountGM_;
546	segCountGM_.SetGlobalBuffer((__gm__ int32_t *)segCountGM);	547	segCountGM_.SetGlobalBuffer((__gm__ int32_t *)segCountGM);

Mmc2/engram_fetch_grad/op_kernel/engram_fetch_grad_utils.h+3-3

	@@ -20,13 +20,13 @@
20	namespace Mc2Kernel {	20	namespace Mc2Kernel {
21	constexpr uint32_t MAX_QP_SIZE = 1024U;	21	constexpr uint32_t MAX_QP_SIZE = 1024U;
22	constexpr uint32_t UB_ALIGN = 32U;	22	constexpr uint32_t UB_ALIGN = 32U;
23	-constexpr uint32_t TILE_BYTES = 32U * 1024U;	23	+constexpr uint32_t TILE_BYTES = 64U * 1024U;
24	constexpr uint32_t HCOMM_INIT_SIZE = 512U;	24	constexpr uint32_t HCOMM_INIT_SIZE = 512U;
25	constexpr int32_t BITS_PER_BYTE = 8;	25	constexpr int32_t BITS_PER_BYTE = 8;
26	constexpr uint32_t ALIGNED_LEN_256 = 256U;	26	constexpr uint32_t ALIGNED_LEN_256 = 256U;
27		27
28	constexpr uint32_t STATE_OFFSET = 32U;	28	constexpr uint32_t STATE_OFFSET = 32U;
29	-constexpr uint32_t NUM_SLOTS = 4U;	29	+constexpr uint32_t NUM_SLOTS = 8U;
30	constexpr uint32_t SENDER_CHANNEL_IDX = 0U;	30	constexpr uint32_t SENDER_CHANNEL_IDX = 0U;
31	constexpr uint32_t RECEIVER_CHANNEL_IDX = 1U;	31	constexpr uint32_t RECEIVER_CHANNEL_IDX = 1U;
32		32
	@@ -42,7 +42,7 @@ struct EngramCommContext {
42	uint32_t rankId;	42	uint32_t rankId;
43	uint32_t rankSize;	43	uint32_t rankSize;
44	uint64_t commBuffer[MAX_QP_SIZE];	44	uint64_t commBuffer[MAX_QP_SIZE];
45	- uint64_t hcommHandle[MAX_QP_SIZE];	45	+ uint64_t hcommHandle[MAX_QP_SIZE * 2];
46	uint32_t channelsPerRank;	46	uint32_t channelsPerRank;
47	};	47	};
48		48