matmul_dequant_pipeline-GVirt-AtomGit

已关闭

matmul_dequant_pipeline #436

liuweiqi创建于 8月14日关闭于 5 天前从lwq030909/GVirt： pr_branch 合入到 openeuler/GVirt： master

matmul_dequant_pipeline #436

已关闭

liuweiqi创建于 8月14日关闭于 5 天前

从lwq030909/GVirt pr_branch合入到openeuler/GVirt master

共 12 个文件变更+1003-0

Mxlite/CMakeLists.txt+1-0

	@@ -75,6 +75,7 @@ file(GLOB KERNEL_FILES_BF16
75	${CMAKE_CURRENT_SOURCE_DIR}/csrc/kernels/flash_attention_bfloat16_t.cpp	75	${CMAKE_CURRENT_SOURCE_DIR}/csrc/kernels/flash_attention_bfloat16_t.cpp
76	${CMAKE_CURRENT_SOURCE_DIR}/csrc/kernels/quant_bfloat16_t.cpp	76	${CMAKE_CURRENT_SOURCE_DIR}/csrc/kernels/quant_bfloat16_t.cpp
77	${CMAKE_CURRENT_SOURCE_DIR}/csrc/kernels/quant_dyn_bfloat16_t.cpp	77	${CMAKE_CURRENT_SOURCE_DIR}/csrc/kernels/quant_dyn_bfloat16_t.cpp
		78	+ ${CMAKE_CURRENT_SOURCE_DIR}/csrc/kernels/fusion_operator_matmul_dequant_pipeline_bfloat16_t.cpp
78	${CMAKE_CURRENT_SOURCE_DIR}/csrc/kernels/norm_bfloat16_t.cpp	79	${CMAKE_CURRENT_SOURCE_DIR}/csrc/kernels/norm_bfloat16_t.cpp
79	${CMAKE_CURRENT_SOURCE_DIR}/csrc/kernels/qk_rms_norm_bfloat16_t.cpp	80	${CMAKE_CURRENT_SOURCE_DIR}/csrc/kernels/qk_rms_norm_bfloat16_t.cpp
80	${CMAKE_CURRENT_SOURCE_DIR}/csrc/kernels/mla_prepare_bfloat16_t.cpp	81	${CMAKE_CURRENT_SOURCE_DIR}/csrc/kernels/mla_prepare_bfloat16_t.cpp

Mxlite/csrc/_C.cpp+24-0

	@@ -1808,6 +1808,26 @@ void MatmulDeQuant(XRuntime &rt, at::Tensor &x, at::Tensor &y, at::Tensor &bias,
1808	rt.Synchronize();	1808	rt.Synchronize();
1809	}	1809	}
1810		1810
		1811	+void FusionOperatorMatmulDequantPipeline(XRuntime &rt, at::Tensor &x, at::Tensor &weight,
		1812	+ at::Tensor &out, at::Tensor &bias, at::Tensor &deqScale,
		1813	+ at::Tensor &outScale, at::Tensor &num, bool weightNZ,
		1814	+ bool transpose)
		1815	+{
		1816	+ XTensor _x, _weight, _out, _bias, _deqScale, _outScale, _num;
		1817	+
		1818	+ InitXTensor(_x, x);
		1819	+ InitXTensor(_weight, weight);
		1820	+ InitXTensor(_out, out);
		1821	+ InitOptionalXTensor(_bias, bias);
		1822	+ InitXTensor(_deqScale, deqScale);
		1823	+ InitOptionalXTensor(_outScale, outScale);
		1824	+ InitOptionalXTensor(_num, num);
		1825	+
		1826	+ XliteOpFusionOperatorMatmulDequantPipeline(rt, _x, _weight, _out, _bias, _deqScale, weightNZ,
		1827	+ transpose, _outScale, _num);
		1828	+ rt.Synchronize();
		1829	+}
		1830	+
1811	void DeQuant(XRuntime &rt, at::Tensor &in, at::Tensor &scale, at::Tensor &out, bool hasScale)	1831	void DeQuant(XRuntime &rt, at::Tensor &in, at::Tensor &scale, at::Tensor &out, bool hasScale)
1812	{	1832	{
1813	XTensor _in, _scale, _out;	1833	XTensor _in, _scale, _out;
	@@ -2503,6 +2523,10 @@ PYBIND11_MODULE(_C, m)
2503	m.def("matmul_dequant", &MatmulDeQuant, "matmul_dequant", py::arg("rt"), py::arg("x"),	2523	m.def("matmul_dequant", &MatmulDeQuant, "matmul_dequant", py::arg("rt"), py::arg("x"),
2504	py::arg("y"), py::arg("bias"), py::arg("deq_scale"), py::arg("z"),	2524	py::arg("y"), py::arg("bias"), py::arg("deq_scale"), py::arg("z"),
2505	py::arg("weight_nz") = false, py::arg("transpose") = false);	2525	py::arg("weight_nz") = false, py::arg("transpose") = false);
		2526	+ m.def("fusion_operator_matmul_dequant_pipeline", &FusionOperatorMatmulDequantPipeline,
		2527	+ "fusion_operator_matmul_dequant_pipeline", py::arg("rt"), py::arg("x"), py::arg("weight"),
		2528	+ py::arg("out"), py::arg("bias"), py::arg("deq_scale"), py::arg("out_scale"),
		2529	+ py::arg("num"), py::arg("weight_nz") = false, py::arg("transpose") = false);
2506	m.def("dequant", &DeQuant, py::arg("rt"), py::arg("in_"), py::arg("scale"), py::arg("out"),	2530	m.def("dequant", &DeQuant, py::arg("rt"), py::arg("in_"), py::arg("scale"), py::arg("out"),
2507	py::arg("has_scale"));	2531	py::arg("has_scale"));
2508	m.def("mla_v2", &MLAV2, py::arg("rt"), py::arg("q_with_qr"), py::arg("qr"), py::arg("k_cache"),	2532	m.def("mla_v2", &MLAV2, py::arg("rt"), py::arg("q_with_qr"), py::arg("qr"), py::arg("k_cache"),

Axlite/csrc/kernels/fusion_operator_matmul_dequant_pipeline.h+622-0

@@ -0,0 +1,622 @@
	1	+/*
	2	+ * Copyright (C) 2026. Huawei Technologies Co., Ltd. All rights reserved.
	3	+ */
	4	+#pragma once
	5	+
	6	+#include "kernel_macro.h"
	7	+#include "kernel_operator.h"
	8	+
	9	+using namespace AscendC;
	10	+
	11	+#define FUSION_MATMUL_DQ_A2V_FLAG 8
	12	+#define FUSION_MATMUL_DQ_V2A_FLAG 9
	13	+#define FUSION_MATMUL_DQ_BATCH_TILE_NUM 2
	14	+#define FUSION_MATMUL_DQ_CREDIT_WINDOW 14
	15	+#define FUSION_MATMUL_DQ_DEQUANT_N_TILE 7168
	wangxiaoran_89368月14日切成 m0, n0 的话，也不需要进行越界检查了，n0 最大 256
	16	+
	17	+template <typename Dtype, typename MatDtype, typename OutDtype>
	18	+class FusionMatmulDequantPipeline
	19	+{
	20	+public:
	21	+ __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z, GM_ADDR bias, GM_ADDR deqScale,
	22	+ GM_ADDR outScale, GM_ADDR num, uint64_t m, uint64_t n, uint64_t k,
	23	+ uint64_t nz, uint64_t transpose, uint64_t m0, uint64_t n0,
	24	+ uint64_t k0, int xSrcDValue = -1, int zDstDValue = -1)
	25	+ {
	26	+ KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2);
	27	+
	28	+ const bool hasBiasArg = (bias != nullptr);
	29	+ const bool hasDeqScaleArg = (deqScale != nullptr);
	30	+ if (m0 == (uint64_t)-1) {
	31	+ m0 = ROUND_UP(m, 32);
	32	+ if (m0 > 128) {
	33	+ m0 = 128;
	34	+ }
	35	+ n0 = (hasBiasArg \|\| hasDeqScaleArg) ? 128 : 256;
	36	+ k0 = 512 / sizeof(Dtype);
	37	+ }
	38	+
	39	+ this->m = m;
	40	+ this->n = n;
	41	+ this->m0 = m0;
	42	+ this->n0 = n0;
	43	+ uint64_t nLoop = DIV_ROUND_UP(n, n0);
	44	+ uint64_t mLoop = DIV_ROUND_UP(m, m0);
	45	+ this->nLoop = nLoop;
	46	+ this->coreLoop = nLoop * mLoop;
	47	+
	48	+#if defined(__DAV_C220_CUBE__)
	49	+ this->aGmBuf.SetGlobalBuffer((__gm__ Dtype *)x);
	50	+ this->bGmBuf.SetGlobalBuffer((__gm__ Dtype *)y);
	51	+ this->biasGmBuf.SetGlobalBuffer((__gm__ MatDtype *)bias);
	52	+ this->deqScaleGmBuf.SetGlobalBuffer((__gm__ uint64_t *)deqScale);
	53	+ this->cGmBuf.SetGlobalBuffer((__gm__ OutDtype *)z);
	54	+
	55	+ this->k = k;
	56	+ this->nz = nz;
	57	+ this->transpose = transpose;
	58	+ this->k0 = k0;
	59	+ this->srcDValue = xSrcDValue == -1 ? k : xSrcDValue;
	wangxiaoran_89368月14日初始化过程混乱，未很好区分 AIC/AIV 的初始化差异
	60	+ this->dstDValue = zDstDValue == -1 ? n : zDstDValue;
	61	+ this->mBlockSize = 16;
	62	+ this->nBlockSize = 16;
	63	+ if (std::is_same<Dtype, int8_t>::value && transpose) {
	64	+ this->nBlockSize = 32;
	65	+ }
	66	+ this->kBlockSize = 32 / sizeof(Dtype);
	67	+ this->kDtileSize = k0 << 1;
	68	+ this->kQtileSize = k0 >> 2;
	69	+ this->hasBias = hasBiasArg;
	70	+ this->hasDeqScale = hasDeqScaleArg;
	71	+
	72	+ int l1ATileBytes = m0 * this->kDtileSize * sizeof(Dtype);
	73	+ int l1BTileBytes = n0 * k0 * sizeof(Dtype);
	74	+ int l1BiasTileBytes = n0 * sizeof(MatDtype);
	75	+ int l1DeqScaleTileBytes = n0 * sizeof(uint64_t);
	76	+ int l0ATileBytes = m0 * this->kQtileSize * sizeof(Dtype);
	77	+ int l0BTileBytes = n0 * this->kQtileSize * sizeof(Dtype);
	78	+ int l0BiasTileBytes = n0 * sizeof(MatDtype);
	79	+ int l0CTileBytes = m0 * n0 * sizeof(Dtype);
	80	+ uint64_t off = 0;
	81	+ for (int i = 0; i < PINGPONG_BUF_NUM; i++) {
	82	+ l1aBuf[i].address_.logicPos = static_cast<uint8_t>(TPosition::A1);
	83	+ l1aBuf[i].address_.bufferAddr = reinterpret_cast<uint64_t>(off);
	84	+ off += l1ATileBytes;
	85	+ }
	86	+ for (int i = 0; i < PINGPONG_BUF_NUM; i++) {
	87	+ l1bBuf[i].address_.logicPos = static_cast<uint8_t>(TPosition::B1);
	88	+ l1bBuf[i].address_.bufferAddr = reinterpret_cast<uint64_t>(off);
	89	+ off += l1BTileBytes;
	90	+ }
	91	+ if (this->hasBias) {
	92	+ l1BiasBuf.address_.logicPos = static_cast<uint8_t>(TPosition::C1);
	93	+ l1BiasBuf.address_.bufferAddr = reinterpret_cast<uint64_t>(off);
	94	+ off += l1BiasTileBytes;
	95	+ }
	96	+ if (this->hasDeqScale) {
	97	+ l1DeqScaleBuf.address_.logicPos = static_cast<uint8_t>(TPosition::C1);
	98	+ l1DeqScaleBuf.address_.bufferAddr = reinterpret_cast<uint64_t>(off);
	99	+ off += l1DeqScaleTileBytes;
	100	+ }
	101	+ off = 0;
	102	+ for (int i = 0; i < PINGPONG_BUF_NUM; i++) {
	103	+ l0aBuf[i].address_.logicPos = static_cast<uint8_t>(TPosition::A2);
	104	+ l0aBuf[i].address_.bufferAddr = reinterpret_cast<uint64_t>(off);
	105	+ off += l0ATileBytes;
	106	+ }
	107	+ off = 0;
	108	+ for (int i = 0; i < PINGPONG_BUF_NUM; i++) {
	109	+ l0bBuf[i].address_.logicPos = static_cast<uint8_t>(TPosition::B2);
	110	+ l0bBuf[i].address_.bufferAddr = reinterpret_cast<uint64_t>(off);
	111	+ off += l0BTileBytes;
	112	+ }
	113	+ off = 0;
	114	+ if (this->hasBias) {
	115	+ l0BiasBuf.address_.logicPos = static_cast<uint8_t>(TPosition::C2);
	116	+ l0BiasBuf.address_.bufferAddr = reinterpret_cast<uint64_t>(off);
	117	+ off = 0;
	118	+ }
	119	+ if (this->hasDeqScale) {
	120	+ fixpipeBuf.address_.logicPos = static_cast<uint8_t>(TPosition::C2PIPE2GM);
	121	+ fixpipeBuf.address_.bufferAddr = reinterpret_cast<uint64_t>(off);
	122	+ off = 0;
	123	+ }
	124	+ l0cBuf.address_.logicPos = static_cast<uint8_t>(TPosition::CO1);
	125	+ l0cBuf.address_.bufferAddr = reinterpret_cast<uint64_t>(off);
	126	+ (void)outScale;
	127	+ (void)num;
	128	+#endif
	129	+
	130	+#if defined(__DAV_C220_VEC__)
	131	+ this->dequantInGm = reinterpret_cast<__gm__ float16_t *>(z);
	132	+ this->dequantScaleGm = reinterpret_cast<__gm__ float32_t *>(outScale);
	133	+ this->dequantOutGm = reinterpret_cast<__gm__ bfloat16_t *>(z);
	134	+ this->dequantRowStride = static_cast<uint32_t>(n);
	135	+ this->hasOutScale = (outScale != nullptr);
	136	+ this->effectiveM = static_cast<uint32_t>(m);
	137	+ if (num != nullptr) {
	138	+ uint32_t numTokens = ((__gm__ uint32_t )num);
	139	+ this->effectiveM = numTokens < m ? numTokens : static_cast<uint32_t>(m);
	140	+ }
	141	+ this->dequantSubBlockIdx = get_subblockid();
	142	+
	143	+ uint64_t maxNActual = n < n0 ? n : n0;
	144	+ this->dequantNTile = maxNActual < FUSION_MATMUL_DQ_DEQUANT_N_TILE
	145	+ ? static_cast<uint32_t>(maxNActual)
	146	+ : FUSION_MATMUL_DQ_DEQUANT_N_TILE;
	147	+ uint32_t nPad = ROUND_UP(this->dequantNTile, (256 / sizeof(float16_t)));
	148	+ uint64_t maxMActual = this->effectiveM < m0 ? this->effectiveM : m0;
	149	+ uint32_t maxLocalRows = DIV_ROUND_UP(static_cast<uint32_t>(maxMActual), 2);
	150	+ uint32_t scaleVecLen = this->hasOutScale ? maxLocalRows : 1;
	151	+ uint32_t scaleVecPad = ROUND_UP(scaleVecLen, (256 / sizeof(float32_t)));
	152	+
	153	+ uint64_t off = 0;
	154	+ for (int i = 0; i < PINGPONG_BUF_NUM; i++) {
	155	+ this->dequantInUb[i] = reinterpret_cast<__ubuf__ float16_t *>((uintptr_t)off);
	156	+ off += nPad * sizeof(float16_t);
	157	+ }
	158	+ for (int i = 0; i < PINGPONG_BUF_NUM; i++) {
	159	+ this->dequantTmpUb[i] = reinterpret_cast<__ubuf__ float32_t *>((uintptr_t)off);
	160	+ off += nPad * sizeof(float32_t);
	161	+ }
	162	+ for (int i = 0; i < PINGPONG_BUF_NUM; i++) {
	163	+ this->dequantMulUb[i] = reinterpret_cast<__ubuf__ float32_t *>((uintptr_t)off);
	164	+ off += nPad * sizeof(float32_t);
	165	+ }
	166	+ for (int i = 0; i < PINGPONG_BUF_NUM; i++) {
	167	+ this->dequantOutUb[i] = reinterpret_cast<__ubuf__ bfloat16_t *>((uintptr_t)off);
	168	+ off += nPad * sizeof(bfloat16_t);
	169	+ }
	170	+ this->dequantScaleUb = reinterpret_cast<__ubuf__ float32_t *>((uintptr_t)off);
	171	+ off += scaleVecPad * sizeof(float32_t);
	172	+ assert(off <= UB_SIZE);
	173	+#endif
	174	+ }
	175	+
	176	+#if defined(__DAV_C220_VEC__)
	177	+ __aicore__ inline void DequantTile(__gm__ float16_t tileInGm, __gm__ float32_t tileScaleGm,
	178	+ __gm__ bfloat16_t *tileOutGm, uint32_t localRows,
	179	+ uint32_t nActual)
	180	+ {
	181	+ if (localRows == 0 \|\| nActual == 0) {
	182	+ return;
	183	+ }
	184	+
	185	+ uint32_t nTile = nActual < this->dequantNTile ? nActual : this->dequantNTile;
	186	+ uint32_t nLoop = DIV_ROUND_UP(nActual, nTile);
	187	+
	188	+ set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0);
	189	+ set_flag(PIPE_V, PIPE_MTE2, EVENT_ID1);
	190	+ set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0);
	191	+ set_flag(PIPE_MTE3, PIPE_V, EVENT_ID1);
	192	+
	193	+ if (this->hasOutScale) {
	194	+ copy_gm_to_ubuf_align_b16(this->dequantScaleUb, tileScaleGm, 0, 1,
	195	+ localRows * sizeof(float32_t), 0, 0, 0, 0);
	196	+ set_flag(PIPE_MTE2, PIPE_S, EVENT_ID0);
	197	+ wait_flag(PIPE_MTE2, PIPE_S, EVENT_ID0);
	198	+ }
	199	+
	200	+ int eventId = 0;
	201	+ for (uint32_t localRow = 0; localRow < localRows; localRow++) {
	202	+ __gm__ float16_t rowInGm = tileInGm + localRow this->dequantRowStride;
	203	+ __gm__ bfloat16_t rowOutGm = tileOutGm + localRow this->dequantRowStride;
	204	+ for (uint32_t loop = 0; loop < nLoop; loop++) {
	205	+ uint32_t localNOffset = loop * nTile;
	206	+ uint32_t nSize =
	207	+ (localNOffset + nTile > nActual) ? (nActual - localNOffset) : nTile;
	208	+ uint32_t nSizePad = ROUND_UP(nSize, (256 / sizeof(float16_t)));
	209	+ uint32_t nRepeats = DIV_ROUND_UP(nSizePad, VECTOR_MAX_NUM_OF_FP32);
	210	+
	211	+ wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0 + eventId);
	212	+ copy_gm_to_ubuf_align_b16(this->dequantInUb[eventId], rowInGm + localNOffset, 0, 1,
	213	+ nSize * sizeof(float16_t), 0, 0, 0, 0);
	214	+ set_flag(PIPE_MTE2, PIPE_V, EVENT_ID0 + eventId);
	215	+
	216	+ wait_flag(PIPE_MTE2, PIPE_V, EVENT_ID0 + eventId);
	217	+ wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0 + eventId);
	218	+ vconv_f162f32(this->dequantTmpUb[eventId], this->dequantInUb[eventId], nRepeats, 1,
	219	+ 1, 8, 4);
	220	+ pipe_barrier(PIPE_V);
	221	+ set_flag(PIPE_V, PIPE_MTE2, EVENT_ID0 + eventId);
	222	+
	223	+ __ubuf__ float32_t *srcForBf16 = this->dequantTmpUb[eventId];
	224	+ if (this->hasOutScale) {
	225	+ vmuls(this->dequantMulUb[eventId], this->dequantTmpUb[eventId],
	226	+ float(this->dequantScaleUb[localRow]), nRepeats, 1, 1, 8, 8);
	227	+ pipe_barrier(PIPE_V);
	228	+ srcForBf16 = this->dequantMulUb[eventId];
	229	+ }
	230	+
	231	+ vconv_f322bf16r(this->dequantOutUb[eventId], srcForBf16, nRepeats, 1, 1, 4, 8);
	232	+ set_flag(PIPE_V, PIPE_MTE3, EVENT_ID0 + eventId);
	233	+
	234	+ wait_flag(PIPE_V, PIPE_MTE3, EVENT_ID0 + eventId);
	235	+ copy_ubuf_to_gm_align_b16(rowOutGm + localNOffset, this->dequantOutUb[eventId], 0,
	236	+ 1, nSize * sizeof(bfloat16_t), 0, 0, 0, 0);
	237	+ set_flag(PIPE_MTE3, PIPE_V, EVENT_ID0 + eventId);
	238	+
	239	+ eventId = 1 - eventId;
	240	+ }
	241	+ }
	242	+
	243	+ wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID0);
	244	+ wait_flag(PIPE_V, PIPE_MTE2, EVENT_ID1);
	245	+ wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID0);
	246	+ wait_flag(PIPE_MTE3, PIPE_V, EVENT_ID1);
	247	+ pipe_barrier(PIPE_ALL);
	248	+ }
	249	+#endif
	250	+
	251	+#if defined(__DAV_C220_CUBE__)
	252	+ __aicore__ inline void RunAic()
	253	+ {
	254	+ set_padding(0);
	255	+ set_atomic_none();
	256	+ set_nd_para((uint64_t)1);
	257	+ SetHF32Mode(false);
	258	+ SetMaskNorm();
	259	+ SetAtomicNone();
	260	+
	261	+ uint32_t batchTileCount = 0;
	262	+ uint32_t publishedBatchCount = 0;
	263	+
	264	+ int kQtileBlockNum = kQtileSize / kBlockSize;
	265	+ int kLoop = DIV_ROUND_UP(k, kQtileSize);
	266	+ int nStride = ROUND_UP(n, nBlockSize);
	267	+ int kStride = ROUND_UP(k, kBlockSize);
	268	+
	269	+ int pingpongL1A = 0;
	270	+ int pingpongL1B = 0;
	271	+
	272	+ SetFlag<HardEvent::M_MTE1>(EVENT_ID0);
	273	+ SetFlag<HardEvent::M_MTE1>(EVENT_ID1);
	274	+ SetFlag<HardEvent::M_MTE1>(EVENT_ID4);
	275	+ SetFlag<HardEvent::MTE1_MTE2>(EVENT_ID0);
	276	+ SetFlag<HardEvent::MTE1_MTE2>(EVENT_ID1);
	277	+ SetFlag<HardEvent::MTE1_MTE2>(EVENT_ID2);
	278	+ SetFlag<HardEvent::MTE1_MTE2>(EVENT_ID3);
	279	+ SetFlag<HardEvent::MTE1_MTE2>(EVENT_ID4);
	280	+ SetFlag<HardEvent::FIX_MTE2>(EVENT_ID5);
	281	+ SetFlag<HardEvent::FIX_M>(EVENT_ID0);
	282	+
	283	+ for (int32_t loopIdx = block_idx; loopIdx < coreLoop; loopIdx += block_num) {
	284	+ int64_t midx = loopIdx / nLoop;
	285	+ int64_t nidx = loopIdx % nLoop;
	286	+ int nOffset = nidx * n0;
	287	+ int mOffset = midx * m0;
	288	+
	289	+ int mActual = m0;
	290	+ if (mOffset + m0 > m) {
	291	+ mActual = m - mOffset;
	292	+ }
	293	+ int mActualBlockPad = ROUND_UP(mActual, mBlockSize);
	294	+ int mActualBlockNum = DIV_ROUND_UP(mActual, mBlockSize);
	295	+
	296	+ int nActual = n0;
	297	+ if (nOffset + n0 > n) {
	298	+ nActual = n - nOffset;
	299	+ }
	300	+ int nActualBlockPad = ROUND_UP(nActual, nBlockSize);
	301	+ int nActualBlockNum = DIV_ROUND_UP(nActual, nBlockSize);
	302	+
	303	+ GlobalTensor<OutDtype> outGm = cGmBuf[mOffset * dstDValue + nOffset];
	304	+
	305	+ if (hasBias) {
	306	+ WaitFlag<HardEvent::MTE1_MTE2>(EVENT_ID4);
	307	+ DataCopy(l1BiasBuf, biasGmBuf[nOffset], nActualBlockPad);
	308	+ SetFlag<HardEvent::MTE2_MTE1>(EVENT_ID4);
	309	+ WaitFlag<HardEvent::MTE2_MTE1>(EVENT_ID4);
	310	+ WaitFlag<HardEvent::M_MTE1>(EVENT_ID4);
	311	+ DataCopy(
	312	+ l0BiasBuf, l1BiasBuf,
	313	+ {1,
	314	+ (uint16_t)(DIV_ROUND_UP((nActualBlockPad * sizeof(MatDtype)), C2_DATABLOCK)),
	315	+ 0, 0});
	316	+ SetFlag<HardEvent::MTE1_MTE2>(EVENT_ID4);
	317	+ SetFlag<HardEvent::MTE1_M>(EVENT_ID4);
	318	+ }
	319	+ if (hasDeqScale) {
	320	+ WaitFlag<HardEvent::FIX_MTE2>(EVENT_ID5);
	321	+ DataCopy(l1DeqScaleBuf, deqScaleGmBuf[nOffset], nActualBlockPad);
	322	+ SetFlag<HardEvent::MTE2_FIX>(EVENT_ID5);
	323	+
	324	+ WaitFlag<HardEvent::MTE2_FIX>(EVENT_ID5);
	325	+ DataCopy(fixpipeBuf, l1DeqScaleBuf,
	326	+ {1,
	327	+ (uint16_t)(DIV_ROUND_UP((nActualBlockPad * sizeof(uint64_t)),
	328	+ FIXPIPE_DATABLOCK)),
	329	+ 0, 0});
	330	+ PipeBarrier<PIPE_FIX>();
	331	+ }
	332	+
	333	+ WaitFlag<HardEvent::FIX_M>(EVENT_ID0);
	334	+ int kOffset = 0;
	335	+ int kIdx = 0;
	336	+ for (; kIdx < kLoop; kIdx++) {
	337	+ int kIdx8 = kIdx % 8;
	338	+ int kIdx4 = kIdx % 4;
	339	+ int kIdx2 = kIdx % 2;
	340	+
	341	+ if (kIdx8 == 0) {
	342	+ int kRemSize = kDtileSize;
	343	+ if (kOffset + kRemSize > k) {
	344	+ kRemSize = k - kOffset;
	345	+ }
	346	+ WaitFlag<HardEvent::MTE1_MTE2>(EVENT_ID0 + pingpongL1A);
	347	+ CopyGmToL1Nd2Nz(l1aBuf[pingpongL1A], aGmBuf[mOffset * srcDValue + kOffset],
	348	+ mActual, kRemSize, srcDValue, mActualBlockPad);
	349	+ SetFlag<HardEvent::MTE2_MTE1>(EVENT_ID0 + pingpongL1A);
	350	+ }
	351	+
	352	+ int k0ActualBlockNum;
	353	+ if (kIdx4 == 0) {
	354	+ int kRemSize = k0;
	355	+ if (kOffset + kRemSize > k) {
	356	+ kRemSize = k - kOffset;
	357	+ }
	358	+ k0ActualBlockNum = DIV_ROUND_UP(kRemSize, kBlockSize);
	359	+ WaitFlag<HardEvent::MTE1_MTE2>(EVENT_ID2 + pingpongL1B);
	360	+ if (transpose == 0 && nz == 0) {
	361	+ CopyGmToL1Nd2Nz(l1bBuf[pingpongL1B], bGmBuf[nOffset * k + kOffset], nActual,
	362	+ kRemSize, k, nActualBlockPad);
	363	+ } else if (transpose == 0 && nz == 1) {
	364	+ CopyGmToL1(l1bBuf[pingpongL1B],
	365	+ bGmBuf[kOffset * nStride + nOffset * kBlockSize], nActual,
	366	+ k0ActualBlockNum, nStride);
	367	+ } else if (transpose == 1 && nz == 0) {
	368	+ CopyGmToL1Nd2Nz(l1bBuf[pingpongL1B], bGmBuf[kOffset * n + nOffset],
	369	+ kRemSize, nActual, n, ROUND_UP(kRemSize, kBlockSize));
	370	+ } else if (transpose == 1 && nz == 1) {
	371	+ CopyGmToL1(l1bBuf[pingpongL1B],
	372	+ bGmBuf[nOffset * kStride + kOffset * nBlockSize], kRemSize,
	373	+ DIV_ROUND_UP(nActual, nBlockSize), kStride);
	374	+ }
	375	+ SetFlag<HardEvent::MTE2_MTE1>(EVENT_ID2 + pingpongL1B);
	376	+ }
	377	+
	378	+ int kActual = kQtileSize;
	379	+ if (kOffset + kActual > k) {
	380	+ kActual = k - kOffset;
	381	+ }
	382	+ int kActualBlockPad = ROUND_UP(kActual, kBlockSize);
	383	+ int kActualBlockNum = DIV_ROUND_UP(kActual, kBlockSize);
	384	+
	385	+ WaitFlag<HardEvent::M_MTE1>(EVENT_ID0 + kIdx2);
	386	+
	387	+ if (kIdx8 == 0) {
	388	+ WaitFlag<HardEvent::MTE2_MTE1>(EVENT_ID0 + pingpongL1A);
	389	+ }
	390	+ CopyToL0ACol(l0aBuf[kIdx2], l1aBuf[pingpongL1A], mActualBlockNum,
	391	+ kIdx8 * kQtileBlockNum, kActualBlockNum);
	392	+ if (kIdx8 == 7 \|\| kIdx == (kLoop - 1)) {
	393	+ SetFlag<HardEvent::MTE1_MTE2>(EVENT_ID0 + pingpongL1A);
	394	+ pingpongL1A ^= 1;
	395	+ }
	396	+
	397	+ if (kIdx4 == 0) {
	398	+ WaitFlag<HardEvent::MTE2_MTE1>(EVENT_ID2 + pingpongL1B);
	399	+ }
	400	+ if (transpose) {
	401	+ CopyToL0BTCol(l0bBuf[kIdx2], l1bBuf[pingpongL1B], nActualBlockNum,
	402	+ kIdx4 * kQtileBlockNum, kActualBlockNum, k0ActualBlockNum);
	403	+ } else {
	404	+ CopyToL0BCol(l0bBuf[kIdx2], l1bBuf[pingpongL1B], nActualBlockNum,
	405	+ kIdx4 * kQtileBlockNum, kActualBlockNum);
	406	+ }
	407	+ if (kIdx4 == 3 \|\| kIdx == (kLoop - 1)) {
	408	+ SetFlag<HardEvent::MTE1_MTE2>(EVENT_ID2 + pingpongL1B);
	409	+ pingpongL1B ^= 1;
	410	+ }
	411	+
	412	+ SetFlag<HardEvent::MTE1_M>(EVENT_ID0 + kIdx2);
	413	+ WaitFlag<HardEvent::MTE1_M>(EVENT_ID0 + kIdx2);
	414	+
	415	+ PipeBarrier<PIPE_M>();
	416	+ if (hasBias && kIdx == 0) {
	417	+ WaitFlag<HardEvent::MTE1_M>(EVENT_ID4);
	418	+ CalMmadWithBias(l0cBuf, l0aBuf[kIdx2], l0bBuf[kIdx2], l0BiasBuf,
	419	+ mActualBlockPad, nActualBlockPad, kActualBlockPad);
	420	+ SetFlag<HardEvent::M_MTE1>(EVENT_ID4);
	421	+ } else {
	422	+ CalMmad(l0cBuf, l0aBuf[kIdx2], l0bBuf[kIdx2], mActualBlockPad, nActualBlockPad,
	423	+ kActualBlockPad, kIdx == 0);
	424	+ }
	425	+ SetFlag<HardEvent::M_MTE1>(EVENT_ID0 + kIdx2);
	426	+ kOffset += kActual;
	427	+ }
	428	+ SetFlag<HardEvent::M_FIX>(EVENT_ID0);
	429	+ WaitFlag<HardEvent::M_FIX>(EVENT_ID0);
	430	+ CopyToGmWithDequant(outGm, l0cBuf, mActual, nActual, mActualBlockPad, dstDValue,
	431	+ hasDeqScale, fixpipeBuf);
	432	+ PipeBarrier<PIPE_FIX>();
	433	+ if (hasDeqScale) {
	434	+ SetFlag<HardEvent::FIX_MTE2>(EVENT_ID5);
	435	+ }
	436	+ SetFlag<HardEvent::FIX_M>(EVENT_ID0);
	437	+
	438	+ batchTileCount += 1;
	439	+ bool isBatchEnd = batchTileCount >= FUSION_MATMUL_DQ_BATCH_TILE_NUM;
	440	+ bool isLastTile = loopIdx + block_num >= coreLoop;
	441	+ if (isBatchEnd \|\| isLastTile) {
	442	+ WaitFlag<HardEvent::FIX_M>(EVENT_ID0);
	443	+ AscendC::CrossCoreSetFlag<0x2, PIPE_FIX>(FUSION_MATMUL_DQ_A2V_FLAG);
	444	+ SetFlag<HardEvent::FIX_M>(EVENT_ID0);
	445	+ batchTileCount = 0;
	446	+ publishedBatchCount += 1;
	447	+ if (publishedBatchCount % FUSION_MATMUL_DQ_CREDIT_WINDOW == 0) {
	448	+ AscendC::CrossCoreWaitFlag(FUSION_MATMUL_DQ_V2A_FLAG);
	449	+ }
	450	+ }
	451	+ }
	452	+
	453	+ WaitFlag<HardEvent::FIX_M>(EVENT_ID0);
	454	+ WaitFlag<HardEvent::FIX_MTE2>(EVENT_ID5);
	455	+ WaitFlag<HardEvent::MTE1_MTE2>(EVENT_ID4);
	456	+ WaitFlag<HardEvent::MTE1_MTE2>(EVENT_ID3);
	457	+ WaitFlag<HardEvent::MTE1_MTE2>(EVENT_ID2);
	458	+ WaitFlag<HardEvent::MTE1_MTE2>(EVENT_ID1);
	459	+ WaitFlag<HardEvent::MTE1_MTE2>(EVENT_ID0);
	460	+ WaitFlag<HardEvent::M_MTE1>(EVENT_ID4);
	461	+ WaitFlag<HardEvent::M_MTE1>(EVENT_ID1);
	462	+ WaitFlag<HardEvent::M_MTE1>(EVENT_ID0);
	463	+ }
	464	+#endif
	465	+
	466	+#if defined(__DAV_C220_VEC__)
	467	+ __aicore__ inline void RunAiv()
	468	+ {
	469	+ set_atomic_none();
	470	+ set_mask_norm();
	471	+ set_vector_mask((uint64_t)-1, (uint64_t)-1);
	472	+
	473	+ uint32_t batchTileCount = 0;
	474	+ uint32_t consumedBatchCount = 0;
	475	+ for (int32_t loopIdx = block_idx; loopIdx < coreLoop; loopIdx += block_num) {
	476	+ if (batchTileCount == 0) {
	477	+ AscendC::CrossCoreWaitFlag(FUSION_MATMUL_DQ_A2V_FLAG);
	478	+ }
	479	+
	480	+ int64_t midx = loopIdx / nLoop;
	481	+ int64_t nidx = loopIdx % nLoop;
	482	+
	483	+ uint32_t mOffset = static_cast<uint32_t>(midx * m0);
	484	+ uint32_t nOffset = static_cast<uint32_t>(nidx * n0);
	485	+ uint32_t mActual = static_cast<uint32_t>((mOffset + m0 > m) ? (m - mOffset) : m0);
	486	+ uint32_t nActual = static_cast<uint32_t>((nOffset + n0 > n) ? (n - nOffset) : n0);
	487	+
	488	+ uint32_t activeMActual = mActual;
	489	+ if (mOffset >= this->effectiveM) {
	490	+ activeMActual = 0;
	491	+ } else if (mOffset + activeMActual > this->effectiveM) {
	492	+ activeMActual = this->effectiveM - mOffset;
	493	+ }
	494	+
	495	+ uint32_t localRows = 0;
	496	+ __gm__ float16_t *tileInGm = nullptr;
	497	+ __gm__ float32_t *tileScaleGm = nullptr;
	498	+ __gm__ bfloat16_t *tileOutGm = nullptr;
	499	+ if (activeMActual > 0 && nActual > 0) {
	500	+ uint32_t rowsPerSubBlock = DIV_ROUND_UP(activeMActual, 2);
	501	+ uint32_t localRowStart = this->dequantSubBlockIdx * rowsPerSubBlock;
	502	+ if (localRowStart < activeMActual) {
	503	+ localRows = rowsPerSubBlock;
	504	+ if (localRowStart + localRows > activeMActual) {
	505	+ localRows = activeMActual - localRowStart;
	506	+ }
	507	+
	508	+ uint32_t firstRow = mOffset + localRowStart;
	509	+ tileInGm = this->dequantInGm + firstRow * this->dequantRowStride + nOffset;
	510	+ tileOutGm = this->dequantOutGm + firstRow * this->dequantRowStride + nOffset;
	511	+ if (this->hasOutScale) {
	512	+ tileScaleGm = this->dequantScaleGm + firstRow;
	513	+ }
	514	+ }
	515	+ }
	516	+
	517	+ if (localRows > 0) {
	518	+ DequantTile(tileInGm, tileScaleGm, tileOutGm, localRows, nActual);
	519	+ }
	520	+
	521	+ batchTileCount += 1;
	522	+ bool isBatchEnd = batchTileCount >= FUSION_MATMUL_DQ_BATCH_TILE_NUM;
	523	+ bool isLastTile = loopIdx + block_num >= coreLoop;
	524	+ if (isBatchEnd \|\| isLastTile) {
	525	+ batchTileCount = 0;
	526	+ consumedBatchCount += 1;
	527	+ if (consumedBatchCount % FUSION_MATMUL_DQ_CREDIT_WINDOW == 0) {
	528	+ AscendC::CrossCoreSetFlag<0x2, PIPE_MTE2>(FUSION_MATMUL_DQ_V2A_FLAG);
	529	+ }
	530	+ }
	531	+ }
	532	+ }
	533	+#endif
	534	+
	535	+ __aicore__ inline void Run()
	536	+ {
	537	+#if defined(__DAV_C220_CUBE__)
	538	+ RunAic();
	539	+#elif defined(__DAV_C220_VEC__)
	540	+ RunAiv();
	541	+#endif
	542	+ }
	543	+
	544	+private:
	545	+ uint64_t m;
	546	+ uint64_t n;
	547	+ uint64_t m0;
	548	+ uint64_t n0;
	549	+ int coreLoop;
	550	+ int nLoop;
	551	+
	552	+#if defined(__DAV_C220_CUBE__)
	553	+ GlobalTensor<Dtype> aGmBuf;
	wangxiaoran_89368月14日 AIC/AIV 的变量隔离开，tiling 信息保留
	554	+ GlobalTensor<Dtype> bGmBuf;
	555	+ GlobalTensor<OutDtype> cGmBuf;
	556	+ GlobalTensor<MatDtype> biasGmBuf;
	557	+ GlobalTensor<uint64_t> deqScaleGmBuf;
	558	+ LocalTensor<Dtype> l1aBuf[PINGPONG_BUF_NUM];
	559	+ LocalTensor<Dtype> l1bBuf[PINGPONG_BUF_NUM];
	560	+ LocalTensor<MatDtype> l1BiasBuf;
	561	+ LocalTensor<uint64_t> l1DeqScaleBuf;
	562	+ LocalTensor<Dtype> l0aBuf[PINGPONG_BUF_NUM];
	563	+ LocalTensor<Dtype> l0bBuf[PINGPONG_BUF_NUM];
	564	+ LocalTensor<MatDtype> l0BiasBuf;
	565	+ LocalTensor<uint64_t> fixpipeBuf;
	566	+ LocalTensor<MatDtype> l0cBuf;
	567	+
	568	+ uint64_t k;
	569	+ uint64_t nz;
	570	+ uint64_t transpose;
	571	+ uint64_t k0;
	572	+ uint64_t mBlockSize;
	573	+ uint64_t nBlockSize;
	574	+ uint64_t kBlockSize;
	575	+ int srcDValue;
	576	+ int dstDValue;
	577	+ int kDtileSize;
	578	+ int kQtileSize;
	579	+ bool hasBias = false;
	580	+ bool hasDeqScale = false;
	581	+#endif
	582	+
	583	+#if defined(__DAV_C220_VEC__)
	584	+ __gm__ float16_t *dequantInGm;
	585	+ __gm__ float32_t *dequantScaleGm;
	586	+ __gm__ bfloat16_t *dequantOutGm;
	587	+ __ubuf__ float16_t *dequantInUb[PINGPONG_BUF_NUM];
	588	+ __ubuf__ float32_t *dequantTmpUb[PINGPONG_BUF_NUM];
	589	+ __ubuf__ float32_t *dequantMulUb[PINGPONG_BUF_NUM];
	590	+ __ubuf__ bfloat16_t *dequantOutUb[PINGPONG_BUF_NUM];
	591	+ __ubuf__ float32_t *dequantScaleUb;
	592	+ uint32_t dequantNTile;
	593	+ uint32_t dequantRowStride;
	594	+ uint32_t dequantSubBlockIdx;
	595	+ uint32_t effectiveM;
	596	+ bool hasOutScale = false;
	597	+#endif
	598	+};
	599	+
	600	+template <typename MatInDtype, typename MatDtype, typename MatOutDtype>
	601	+__aicore__ inline void FusionOperatorMatmulDequantPipelineRun(
	602	+ GM_ADDR x, GM_ADDR weight, GM_ADDR out, GM_ADDR bias, GM_ADDR weightScale, GM_ADDR outScale,
	603	+ GM_ADDR num, uint64_t m, uint64_t n, uint64_t k, uint64_t weightNz, uint64_t transposeWeight,
	604	+ uint64_t m0, uint64_t n0, uint64_t k0)
	605	+{
	606	+ KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2);
	607	+ FusionMatmulDequantPipeline<MatInDtype, MatDtype, MatOutDtype> op;
	608	+ op.Init(x, weight, out, bias, weightScale, outScale, num, m, n, k, weightNz, transposeWeight,
	609	+ m0, n0, k0);
	610	+ op.Run();
	611	+}
	612	+
	613	+#define FUSION_OPERATOR_MATMUL_DEQUANT_PIPELINE_FUNC_DEFINE(dtype) \
	614	+ extern "C" __global__ __aicore__ void fusion_operator_matmul_dequant_pipeline_##dtype( \
	615	+ GM_ADDR x, GM_ADDR weight, GM_ADDR out, GM_ADDR bias, GM_ADDR weightScale, \
	616	+ GM_ADDR outScale, GM_ADDR num, uint64_t m, uint64_t n, uint64_t k, uint64_t weightNz, \
	617	+ uint64_t transposeWeight, uint64_t m0, uint64_t n0, uint64_t k0) \
	618	+ { \
	619	+ FusionOperatorMatmulDequantPipelineRun<int8_t, int32_t, half>( \
	620	+ x, weight, out, bias, weightScale, outScale, num, m, n, k, weightNz, transposeWeight, \
	621	+ m0, n0, k0); \
	622	+ }

Axlite/csrc/kernels/fusion_operator_matmul_dequant_pipeline_bfloat16_t.cpp+6-0

@@ -0,0 +1,6 @@
	1	+/*
	2	+ * Copyright (C) 2026. Huawei Technologies Co., Ltd. All rights reserved.
	3	+ */
	4	+#include "fusion_operator_matmul_dequant_pipeline.h"
	5	+
	6	+FUSION_OPERATOR_MATMUL_DEQUANT_PIPELINE_FUNC_DEFINE(bfloat16_t)

Mxlite/csrc/model.cpp+21-0

	@@ -303,6 +303,27 @@ void XModel::ForwardLinear(XRuntime &rt, uint32_t layer, XTensor &x,
303	isNz = _c.weightNZ \|\| _c.quantAttnWeightNz;	303	isNz = _c.weightNZ \|\| _c.quantAttnWeightNz;
304	isTransposed = _c.quantAttnWeightTrans;	304	isTransposed = _c.quantAttnWeightTrans;
305		305
		306	+ if ((quantType == DYNAMIC_QUANT \|\| quantType == STATIC_QUANT) && rt.enableFusedDenseW8A8 &&
		307	+ x.dtype == BF16 && w.weight.dtype == INT8 && out.dtype == BF16 &&
		308	+ w.deqScale.ptr != nullptr && x.shape.size() == 2 && out.shape.size() == 2 &&
		309	+ (quantType == DYNAMIC_QUANT \|\|
		310	+ (w.inputScale.ptr != nullptr && w.inputOffset.ptr != nullptr))) {
		311	+ XTensor &xQuanted = rt.GetTensor(x.shape, INT8, DBG_LOC);
		312	+ if (quantType == DYNAMIC_QUANT) {
		313	+ XTensor &scale = rt.GetTensor({x.shape[0]}, FP32, DBG_LOC);
		314	+ XliteOpQuantDyn(rt, x, scale, xQuanted);
		315	+ XliteOpFusionOperatorMatmulDequantPipeline(rt, xQuanted, w.weight, out, w.quantBias,
		316	+ w.deqScale, isNz, isTransposed, scale);
		317	+ rt.PutTensor(scale);
		318	+ } else {
		319	+ XliteOpQuant(rt, x, w.inputScale, w.inputOffset, xQuanted);
		320	+ XliteOpFusionOperatorMatmulDequantPipeline(rt, xQuanted, w.weight, out, w.quantBias,
		321	+ w.deqScale, isNz, isTransposed);
		322	+ }
		323	+ rt.PutTensor(xQuanted);
		324	+ return;
		325	+ }
		326	+
306	XTensor &xQuanted = rt.GetTensor(x.shape, INT8, DBG_LOC);	327	XTensor &xQuanted = rt.GetTensor(x.shape, INT8, DBG_LOC);
307	if (quantType == DYNAMIC_QUANT) {	328	if (quantType == DYNAMIC_QUANT) {
308	// quant(x) -> xQuanted, perChannelScale	329	// quant(x) -> xQuanted, perChannelScale

Mxlite/csrc/op.cpp+108-0

	@@ -665,6 +665,9 @@ void XliteOpMatmul(XRuntime &rt, XTensor &in, XTensor &weight, XTensor &out, boo
665	mLoop = DIV_ROUND_UP(m, m0);	665	mLoop = DIV_ROUND_UP(m, m0);
666	nLoop = DIV_ROUND_UP(n, n0);	666	nLoop = DIV_ROUND_UP(n, n0);
667	totalLoops = mLoop * nLoop;	667	totalLoops = mLoop * nLoop;
		668	+ // Keep the launch count aligned with active matmul tiles. The reused
		669	+ // Matmul::Run() path has final pipe waits that are only valid for blocks
		670	+ // that actually enter its tile loop.
668	uint32_t aicNum = totalLoops > rt.aicNum ? rt.aicNum : totalLoops;	671	uint32_t aicNum = totalLoops > rt.aicNum ? rt.aicNum : totalLoops;
669	if (aicNum == 0) {	672	if (aicNum == 0) {
670	aicNum = 1;	673	aicNum = 1;
	@@ -1387,6 +1390,111 @@ void XliteOpMatmulDeQuant(XRuntime &rt, XTensor &in, XTensor &weight, XTensor &o
1387	}	1390	}
1388	}	1391	}
1389		1392
		1393	+void XliteOpFusionOperatorMatmulDequantPipeline(XRuntime &rt, XTensor &in, XTensor &weight,
		1394	+ XTensor &out, const XTensor &quantBias,
		1395	+ const XTensor &weightScale, bool weightNZ,
		1396	+ bool transpose, const XTensor &outScale,
		1397	+ const XTensor &num, uint64_t m0, uint64_t n0,
		1398	+ uint64_t k0)
		1399	+{
		1400	+ if (IsDummyRuntime(rt) \|\| in.numel == 0) {
		1401	+ return;
		1402	+ }
		1403	+ const std::string dbgPrefix = DBG_PREFIX;
		1404	+ auto errStr = [&]() {
		1405	+ return dbgPrefix + XT_STR(in) + XT_STR(weight) + XT_STR(out) + XT_STR(weightScale) +
		1406	+ XT_STR(outScale) + XT_STR(num);
		1407	+ };
		1408	+ if (in.ptr == nullptr \|\| weight.ptr == nullptr \|\| out.ptr == nullptr \|\|
		1409	+ weightScale.ptr == nullptr) {
		1410	+ throw std::runtime_error(errStr() + " null pointer!");
		1411	+ }
		1412	+ if (in.shape.size() != 2 \|\| weight.shape.size() != 2 \|\| out.shape.size() != 2) {
		1413	+ throw std::runtime_error(errStr() + " shape unsupported!");
		1414	+ }
		1415	+
		1416	+ uint64_t m = in.shape[0];
		1417	+ uint64_t k = in.shape[1];
		1418	+ uint64_t n = transpose ? weight.shape[1] : weight.shape[0];
		1419	+ uint64_t weightK = transpose ? weight.shape[0] : weight.shape[1];
		1420	+ if (weightK != k \|\| out.shape[0] != m \|\| out.shape[1] != n) {
		1421	+ throw std::runtime_error(errStr() + " shape mismatch!");
		1422	+ }
		1423	+ if (in.dtype != INT8 \|\| weight.dtype != INT8 \|\| out.dtype != BF16) {
		1424	+ throw std::runtime_error(errStr() + " dtype unsupported!");
		1425	+ }
		1426	+ if (quantBias.ptr != nullptr && (quantBias.dtype != INT32 \|\| quantBias.numel < n)) {
		1427	+ throw std::runtime_error(errStr() + " bias unsupported!");
		1428	+ }
		1429	+ if (weightScale.dtype != FP32 \|\| weightScale.numel < 2 * n) {
		1430	+ throw std::runtime_error(errStr() + " weight scale unsupported!");
		1431	+ }
		1432	+ if (outScale.ptr != nullptr && (outScale.dtype != FP32 \|\| outScale.numel < m)) {
		1433	+ throw std::runtime_error(errStr() + " output scale unsupported!");
		1434	+ }
		1435	+ if (num.ptr != nullptr && (num.dtype != INT32 \|\| num.numel < 1)) {
		1436	+ throw std::runtime_error(errStr() + " num unsupported!");
		1437	+ }
		1438	+
		1439	+ bool needExtraSpace = (quantBias.ptr != nullptr \|\| weightScale.ptr != nullptr);
		1440	+
		1441	+ // Keep the AIC matmul tile selection identical to XliteOpMatmul for the
		1442	+ // int8 + fixpipe path. The fused kernel only moves the dequant work into
		1443	+ // the same MIX kernel; it should not change matmul's tiling policy.
		1444	+ if (m0 == MATMUL_M0_N0_K0_DEFAULT_VALUE \|\| n0 == MATMUL_M0_N0_K0_DEFAULT_VALUE \|\|
		1445	+ k0 == MATMUL_M0_N0_K0_DEFAULT_VALUE) {
		1446	+ m0 = ROUND_UP(m, 32);
		1447	+ if (m0 > 128) {
		1448	+ m0 = 128;
		1449	+ }
		1450	+ n0 = needExtraSpace ? 128 : 256;
		1451	+ k0 = 4096 / XDtypeBit(weight.dtype);
		1452	+
		1453	+ uint64_t mLoop = DIV_ROUND_UP(m, m0);
		1454	+ uint64_t nLoop = DIV_ROUND_UP(n, n0);
		1455	+ uint64_t totalLoops = mLoop * nLoop;
		1456	+ uint64_t lastLoops = totalLoops % rt.aicNum;
		1457	+
		1458	+ if (totalLoops < static_cast<uint64_t>(3) * rt.aicNum &&
		1459	+ (lastLoops != 0 && lastLoops < rt.aicNum / 2)) {
		1460	+ if (n <= static_cast<uint64_t>(32) * rt.aicNum) {
		1461	+ m0 = m0 > 64 ? 64 : m0;
		1462	+ n0 = 64;
		1463	+ } else if (n <= static_cast<uint64_t>(64) * rt.aicNum) {
		1464	+ n0 = 64;
		1465	+ } else if (n <= static_cast<uint64_t>(128) * rt.aicNum) {
		1466	+ n0 = 128;
		1467	+ } else if (n <= static_cast<uint64_t>(256) * rt.aicNum) {
		1468	+ n0 = needExtraSpace ? 128 : 256;
		1469	+ } else {
		1470	+ m0 = m0 > 64 ? 64 : m0;
		1471	+ n0 = needExtraSpace ? 256 : 384;
		1472	+ k0 /= 2;
		1473	+ }
		1474	+ }
		1475	+ }
		1476	+
		1477	+ uint64_t totalLoops = DIV_ROUND_UP(m, m0) * DIV_ROUND_UP(n, n0);
		1478	+ if (totalLoops == 0) {
		1479	+ return;
		1480	+ }
		1481	+
		1482	+ if (out.numel < m * n) {
		1483	+ throw std::runtime_error(errStr() + " output is too small!");
		1484	+ }
		1485	+
		1486	+ uint32_t aicNum = totalLoops > rt.aicNum ? rt.aicNum : totalLoops;
		1487	+ if (aicNum == 0) {
		1488	+ aicNum = 1;
		1489	+ }
		1490	+
		1491	+ out.View(FP16);
		1492	+ aclrtlaunch_fusion_operator_matmul_dequant_pipeline_bfloat16_t(
		1493	+ aicNum, rt.stream, in.ptr, weight.ptr, out.ptr, quantBias.ptr, weightScale.ptr,
		1494	+ outScale.ptr, num.ptr, m, n, k, weightNZ, transpose, m0, n0, k0);
		1495	+ out.View(BF16);
		1496	+}
		1497	+
1390	void XliteOpGroupMatmulDeQuant(XRuntime &rt, XTensor &in, XTensor &weights, XTensor &deqScales,	1498	void XliteOpGroupMatmulDeQuant(XRuntime &rt, XTensor &in, XTensor &weights, XTensor &deqScales,
1391	XTensor &counts, uint32_t start, uint32_t end, XDtype weightDtype,	1499	XTensor &counts, uint32_t start, uint32_t end, XDtype weightDtype,
1392	long outDim, long inDim, XTensor &output, XTensor &outScale,	1500	long outDim, long inDim, XTensor &output, XTensor &outScale,

Mxlite/csrc/op.h+6-0

	@@ -147,6 +147,12 @@ void XliteOpMatmulDeQuant(XRuntime &rt, XTensor &in, XTensor &weight, XTensor &o
147	const XTensor &weightScale = XTensor(), bool weightNZ = false,	147	const XTensor &weightScale = XTensor(), bool weightNZ = false,
148	bool transpose = false, const XTensor &outScale = XTensor(),	148	bool transpose = false, const XTensor &outScale = XTensor(),
149	const XTensor &num = XTensor());	149	const XTensor &num = XTensor());
		150	+void XliteOpFusionOperatorMatmulDequantPipeline(
		151	+ XRuntime &rt, XTensor &in, XTensor &weight, XTensor &out, const XTensor &quantBias = XTensor(),
		152	+ const XTensor &weightScale = XTensor(), bool weightNZ = false, bool transpose = false,
		153	+ const XTensor &outScale = XTensor(), const XTensor &num = XTensor(),
		154	+ uint64_t m0 = MATMUL_M0_N0_K0_DEFAULT_VALUE, uint64_t n0 = MATMUL_M0_N0_K0_DEFAULT_VALUE,
		155	+ uint64_t k0 = MATMUL_M0_N0_K0_DEFAULT_VALUE);
150	void XliteOpGroupMatmulDeQuant(XRuntime &rt, XTensor &in, XTensor &weights, XTensor &deqScales,	156	void XliteOpGroupMatmulDeQuant(XRuntime &rt, XTensor &in, XTensor &weights, XTensor &deqScales,
151	XTensor &counts, uint32_t start, uint32_t end, XDtype weightDtype,	157	XTensor &counts, uint32_t start, uint32_t end, XDtype weightDtype,
152	long outDim, long inDim, XTensor &output, XTensor &outScale,	158	long outDim, long inDim, XTensor &output, XTensor &outScale,

Mxlite/csrc/runtime.cpp+11-0

	@@ -92,6 +92,13 @@ void XRuntime::Init(size_t sizeMB)
92	}	92	}
93	}	93	}
94		94
		95	+ if (isEnvironmentVariableTrue(std::getenv("XLITE_FUSED_DENSE_W8A8"))) {
		96	+ enableFusedDenseW8A8 = true;
		97	+ if (_rankId == 0) {
		98	+ std::cout << "Xlite fused dense W8A8 enabled!" << std::endl;
		99	+ }
		100	+ }
		101	+
95	const char *ratioPerEPEnv = std::getenv("XLITE_ACTIVE_TOKENS_RATIO_PER_EP");	102	const char *ratioPerEPEnv = std::getenv("XLITE_ACTIVE_TOKENS_RATIO_PER_EP");
96	if (ratioPerEPEnv) {	103	if (ratioPerEPEnv) {
97	char *endPtr = nullptr;	104	char *endPtr = nullptr;
	@@ -839,6 +846,10 @@ void XDummyRuntime::InitDummyRuntime(size_t sizeMB)
839	originAicNum = aicNum;	846	originAicNum = aicNum;
840	originAivNum = aivNum;	847	originAivNum = aivNum;
841		848
		849	+ if (isEnvironmentVariableTrue(std::getenv("XLITE_FUSED_DENSE_W8A8"))) {
		850	+ enableFusedDenseW8A8 = true;
		851	+ }
		852	+
842	const char *envCommOptimizeLen = std::getenv("XLITE_COMM_OPTIMIZE_LEN");	853	const char *envCommOptimizeLen = std::getenv("XLITE_COMM_OPTIMIZE_LEN");
843	if (envCommOptimizeLen) {	854	if (envCommOptimizeLen) {
844	char *endPtr = nullptr;	855	char *endPtr = nullptr;

Mxlite/csrc/runtime.h+1-0

	@@ -166,6 +166,7 @@ public:
166	uint32_t defaultMatmulSwizzle = 0x600;	166	uint32_t defaultMatmulSwizzle = 0x600;
167	bool disableSwizzleTable = false;	167	bool disableSwizzleTable = false;
168	bool enableMoEAllToAll = false;	168	bool enableMoEAllToAll = false;
		169	+ bool enableFusedDenseW8A8 = false;
169	double activeTokensRatioPerEp = 1.0f;	170	double activeTokensRatioPerEp = 1.0f;
170		171
171	XcclComm *_tpXcclComm = nullptr;	172	XcclComm *_tpXcclComm = nullptr;

Mxlite/doc/ENVIRONMENT_VARIABLES.md+1-0

	@@ -13,6 +13,7 @@
13	\| `XLITE_DISABLE_XCCL` \| 布尔 \| `false` \| 是否禁用 XCCL（XLite 自定义通信算子）。设置为 `true` 时禁用，回退到 HCCL。注意：外部环境变量 `HCCL_DETERMINISTIC` 被设为 `true` 时同样会禁用 XCCL（因 xccl 非确定性，与确定性模式冲突，见 `runtime.cpp` 的 `InitXcclComm`/`InitDummyXcclComm`）。 \|	13	\| `XLITE_DISABLE_XCCL` \| 布尔 \| `false` \| 是否禁用 XCCL（XLite 自定义通信算子）。设置为 `true` 时禁用，回退到 HCCL。注意：外部环境变量 `HCCL_DETERMINISTIC` 被设为 `true` 时同样会禁用 XCCL（因 xccl 非确定性，与确定性模式冲突，见 `runtime.cpp` 的 `InitXcclComm`/`InitDummyXcclComm`）。 \|
14	\| `XLITE_ENABLE_GRAPH_COMM` \| 布尔 \| `true` \| 是否启用通信算子入图（HCCL AllGather/ReduceScatter 的 ACL Graph 捕获+重放，主流上 capture 与 replay。默认启用；设为 `false`/`0`/`no`/`off`（不区分大小写）时关闭，回退到 eager HCCL 单算子路径。注意：即便启用，仅当对应 (type, dtype) 无可用 xccl comm 时才实际走入图路径（xccl self-kernel 不可被 ACL Graph 捕获）；dummy runtime 评估 `get_tensor_pool_size` 时恒走 eager（不入图）。 \|	14	\| `XLITE_ENABLE_GRAPH_COMM` \| 布尔 \| `true` \| 是否启用通信算子入图（HCCL AllGather/ReduceScatter 的 ACL Graph 捕获+重放，主流上 capture 与 replay。默认启用；设为 `false`/`0`/`no`/`off`（不区分大小写）时关闭，回退到 eager HCCL 单算子路径。注意：即便启用，仅当对应 (type, dtype) 无可用 xccl comm 时才实际走入图路径（xccl self-kernel 不可被 ACL Graph 捕获）；dummy runtime 评估 `get_tensor_pool_size` 时恒走 eager（不入图）。 \|
15	\| `XLITE_MOE_ALLTOALL` \| 布尔 \| `false` \| 是否启用 MoE AlltoAll 通信模式。启用后，MoE 的 dispatch/combine 阶段使用 AlltoAllV 集合通信替代默认的 AllGather + ReduceScatter 方式，跨 EP（Expert Parallel）rank 分发/收集 token，并消除 MoE 后的 TP AllReduce。 \|	15	\| `XLITE_MOE_ALLTOALL` \| 布尔 \| `false` \| 是否启用 MoE AlltoAll 通信模式。启用后，MoE 的 dispatch/combine 阶段使用 AlltoAllV 集合通信替代默认的 AllGather + ReduceScatter 方式，跨 EP（Expert Parallel）rank 分发/收集 token，并消除 MoE 后的 TP AllReduce。 \|
		16	+\| `XLITE_FUSED_DENSE_W8A8` \| 布尔 \| `false` \| 是否启用 dense 动态 W8A8 的实验 fused 路径，将 activation dynamic quant、INT8 matmul 和 per-token dequant 合并到一个 MIX AIC/AIV kernel 中。 \|
16	\| `XLITE_ACTIVE_TOKENS_RATIO_PER_EP` \| double \| `1.0f` \| 影响开启EP后单卡按专家排序后激活Tensor的内存占用。该环境变量有效配置范围为[1/ep_size ... 1.0f]。默认1.0f含义为单卡极端情况需要处理所有激活的Token，因此专家排序后激活Tensor的shape为[token个数 * num_experts_per_tok, hidden_size]; 配置为1/ep_size含义为单卡只需要处理所有激活的Token的1/ep_size，因此专家排序后激活Tensor的shape为[token个数 * num_experts_per_tok /ep_size, hidden_size]。该值越小代表EP负载越均衡，排序后激活Tensor预占内存越小，配置过小可能导致算子报错，建议按照实际负载均衡度评估合理值。注意：该配置仅当输入token个数（所有DP的batched token总和）>= 1024个时才生效，小于1024的情况使用默认1.0f进行计算。 \|	17	\| `XLITE_ACTIVE_TOKENS_RATIO_PER_EP` \| double \| `1.0f` \| 影响开启EP后单卡按专家排序后激活Tensor的内存占用。该环境变量有效配置范围为[1/ep_size ... 1.0f]。默认1.0f含义为单卡极端情况需要处理所有激活的Token，因此专家排序后激活Tensor的shape为[token个数 * num_experts_per_tok, hidden_size]; 配置为1/ep_size含义为单卡只需要处理所有激活的Token的1/ep_size，因此专家排序后激活Tensor的shape为[token个数 * num_experts_per_tok /ep_size, hidden_size]。该值越小代表EP负载越均衡，排序后激活Tensor预占内存越小，配置过小可能导致算子报错，建议按照实际负载均衡度评估合理值。注意：该配置仅当输入token个数（所有DP的batched token总和）>= 1024个时才生效，小于1024的情况使用默认1.0f进行计算。 \|
17		18
18	## 布尔值解析规则	19	## 布尔值解析规则

Axlite/tests/kernels/fusion_operator_matmul_dequant_pipeline.py+183-0

@@ -0,0 +1,183 @@
	1	+#!/usr/bin/python3
	2	+# coding=utf-8
	3	+#
	4	+# Copyright (C) 2026. Huawei Technologies Co., Ltd. All rights reserved.
	5	+#
	6	+# Correctness test for fusion_operator_matmul_dequant_pipeline.
	7	+#
	8	+# The reference path is the exact serial model path we want to replace:
	9	+# 1. matmul_dequant(rt, ...) -> launches XliteOpMatmul and writes FP16.
	10	+# 2. dequant(rt, ...) -> converts that FP16 output to BF16.
	11	+# The fused path receives the same x/weight/bias/deq_scale tensors and should
	12	+# produce the same final BF16 tensor.
	13	+
	14	+import argparse
	15	+from dataclasses import dataclass
	16	+
	17	+import torch
	18	+import torch_npu
	19	+
	20	+from xlite._C import (
	21	+ Runtime,
	22	+ dequant,
	23	+ fusion_operator_matmul_dequant_pipeline,
	24	+ matmul_dequant,
	25	+)
	26	+
	27	+
	28	+ACL_FORMAT_FRACTAL_NZ = 29
	29	+
	30	+
	31	+@dataclass(frozen=True)
	32	+class TestCase:
	33	+ m: int
	34	+ n: int
	35	+ k: int
	36	+
	37	+
	38	+TEST_CASES = (
	39	+ TestCase(1, 256, 1024),
	40	+ TestCase(17, 768, 1024),
	41	+ TestCase(103, 1024, 2048),
	42	+)
	43	+
	44	+
	45	+def make_fixpipe_deq_scale(n: int, device: str) -> torch.Tensor:
	46	+ # Matmul fixpipe consumes the FP32 scale as a uint64 stream. Existing
	47	+ # matmul_dequant tests represent each channel as float32 pair [scale, 0].
	48	+ values = torch.rand(n, dtype=torch.float32, device=device) * 0.015 + 0.001
	49	+ values = values.to(torch.bfloat16).to(torch.float32)
	50	+ packed = torch.zeros(n * 2, dtype=torch.float32, device=device)
	51	+ packed[0::2] = values
	52	+ return packed
	53	+
	54	+
	55	+def make_case_tensors(
	56	+ case: TestCase,
	57	+ transpose: bool,
	58	+ weight_nz: bool,
	59	+ device: str,
	60	+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
	61	+ x = torch.randint(-8, 8, (case.m, case.k), dtype=torch.int8, device=device)
	62	+ weight_ref = torch.randint(-8, 8, (case.n, case.k), dtype=torch.int8, device=device)
	63	+ weight = weight_ref.t().contiguous() if transpose else weight_ref.contiguous()
	64	+ if weight_nz:
	65	+ weight = torch_npu.npu_format_cast(weight, ACL_FORMAT_FRACTAL_NZ)
	66	+
	67	+ bias = torch.randint(-64, 64, (case.n,), dtype=torch.int32, device=device)
	68	+ deq_scale = make_fixpipe_deq_scale(case.n, device)
	69	+ return x, weight, bias, deq_scale
	70	+
	71	+
	72	+def assert_close(
	73	+ name: str,
	74	+ expected: torch.Tensor,
	75	+ actual: torch.Tensor,
	76	+ atol: float,
	77	+ rtol: float,
	78	+) -> None:
	79	+ expected_cpu = expected.cpu()
	80	+ actual_cpu = actual.cpu()
	81	+ try:
	82	+ torch.testing.assert_close(actual_cpu, expected_cpu, atol=atol, rtol=rtol)
	83	+ except AssertionError:
	84	+ diff = (actual_cpu.float() - expected_cpu.float()).abs()
	85	+ print(f"{name} failed: max_abs_diff={float(diff.max().item())}")
	86	+ print(f"expected={expected_cpu}")
	87	+ print(f"actual={actual_cpu}")
	88	+ raise
	89	+
	90	+
	91	+def run_one(
	92	+ rt: Runtime,
	93	+ case: TestCase,
	94	+ transpose: bool,
	95	+ weight_nz: bool,
	96	+ has_out_scale: bool,
	97	+ device: str,
	98	+ atol: float,
	99	+ rtol: float,
	100	+) -> None:
	101	+ x, weight, bias, deq_scale = make_case_tensors(case, transpose, weight_nz, device)
	102	+
	103	+ serial_tmp = torch.empty(case.m, case.n, dtype=torch.float16, device=device)
	104	+ serial_out = torch.empty(case.m, case.n, dtype=torch.bfloat16, device=device)
	105	+ fused_out = torch.empty(case.m, case.n, dtype=torch.bfloat16, device=device)
	106	+
	107	+ if has_out_scale:
	108	+ out_scale = torch.rand(case.m, dtype=torch.float32, device=device) * 0.25 + 0.75
	109	+ else:
	110	+ out_scale = torch.empty(0, dtype=torch.float32, device=device)
	111	+ empty_num = torch.empty(0, dtype=torch.int32, device=device)
	112	+
	113	+ name = (
	114	+ f"fusion_operator_matmul_dequant_pipeline "
	115	+ f"m={case.m} n={case.n} k={case.k} "
	116	+ f"weight_nz={weight_nz} transpose={transpose} "
	117	+ f"has_out_scale={has_out_scale}"
	118	+ )
	119	+ print(f"{name} running", flush=True)
	120	+
	121	+ torch.npu.synchronize()
	122	+ print(f"{name} serial_matmul start", flush=True)
	123	+ matmul_dequant(rt, x, weight, bias, deq_scale, serial_tmp, weight_nz, transpose)
	124	+ print(f"{name} serial_dequant start", flush=True)
	125	+ dequant(rt, serial_tmp, out_scale, serial_out, has_out_scale)
	126	+ torch.npu.synchronize()
	127	+ print(f"{name} serial done", flush=True)
	128	+
	129	+ print(f"{name} fused start", flush=True)
	130	+ fusion_operator_matmul_dequant_pipeline(
	131	+ rt,
	132	+ x,
	133	+ weight,
	134	+ fused_out,
	135	+ bias,
	136	+ deq_scale,
	137	+ out_scale,
	138	+ empty_num,
	139	+ weight_nz,
	140	+ transpose,
	141	+ )
	142	+ torch.npu.synchronize()
	143	+ print(f"{name} fused done", flush=True)
	144	+
	145	+ assert_close(name, serial_out, fused_out, atol, rtol)
	146	+ print(f"{name} passed")
	147	+
	148	+
	149	+def main() -> None:
	150	+ parser = argparse.ArgumentParser()
	151	+ parser.add_argument("--device", type=int, default=0)
	152	+ parser.add_argument("--pool-mb", type=int, default=500)
	153	+ parser.add_argument("--seed", type=int, default=0)
	154	+ parser.add_argument("--quick", action="store_true")
	155	+ parser.add_argument("--atol", type=float, default=1e-2)
	156	+ parser.add_argument("--rtol", type=float, default=1e-2)
	157	+ args = parser.parse_args()
	158	+
	159	+ device = f"npu:{args.device}"
	160	+ torch.npu.set_device(args.device)
	161	+ torch.npu.set_option({"ALLOW_INTERNAL_FORMAT": True})
	162	+ torch.manual_seed(args.seed)
	163	+
	164	+ rt = Runtime(args.device, args.pool_mb)
	165	+ cases = TEST_CASES[:1] if args.quick else TEST_CASES
	166	+ for weight_nz in (False, True):
	167	+ for transpose in (False, True):
	168	+ for has_out_scale in (False, True):
	169	+ for case in cases:
	170	+ run_one(
	171	+ rt,
	172	+ case,
	173	+ transpose,
	174	+ weight_nz,
	175	+ has_out_scale,
	176	+ device,
	177	+ args.atol,
	178	+ args.rtol,
	179	+ )
	180	+
	181	+
	182	+if __name__ == "__main__":
	183	+ main()

Mxlite/xlite/_C.pyi+19-0

	@@ -1958,6 +1958,25 @@ def msd_merge_dequant(
1958	"""	1958	"""
1959	...	1959	...
1960		1960
		1961	+def fusion_operator_matmul_dequant_pipeline(
		1962	+ rt: Runtime,
		1963	+ x: torch.Tensor,
		1964	+ weight: torch.Tensor,
		1965	+ out: torch.Tensor,
		1966	+ bias: torch.Tensor,
		1967	+ deq_scale: torch.Tensor,
		1968	+ out_scale: torch.Tensor,
		1969	+ num: torch.Tensor,
		1970	+ weight_nz: bool = False,
		1971	+ transpose: bool = False,
		1972	+) -> None:
		1973	+ """Run the INT8 matmul + BF16 dequant tile-pipeline fusion kernel.
		1974	+
		1975	+ ``x`` is expected to be an already-quantized INT8 activation tensor.
		1976	+ The native kernel writes BF16 output in place to ``out``.
		1977	+ """
		1978	+ ...
		1979	+
1961	def dequant(rt: Runtime, in_: torch.Tensor, scale: torch.Tensor, out: torch.Tensor, has_scale: bool) -> None:	1980	def dequant(rt: Runtime, in_: torch.Tensor, scale: torch.Tensor, out: torch.Tensor, has_scale: bool) -> None:
1962	"""Dequantize tensor values into output precision.	1981	"""Dequantize tensor values into output precision.
1963		1982