Star1331
2076
代码
代码
Issues784
Pull Requests675
Actions
项目讨论
Wiki
分析
Star1331
2076
已关闭
[WIP]tqbmm_tiling #10797
jgx创建于 9月20日关闭于 24 天前
[WIP]tqbmm_tiling #10797
已关闭
jgx创建于 9月20日关闭于 24 天前
从jgx12/ops-nntqbmm_tiling合入到cann/ops-nnmaster
共 5 个文件变更+341-13
Mmatmul/transpose_quant_batch_mat_mul/op_host/op_tiling/arch35/transpose_quant_batch_mat_mul_asw_tiling.cpp+42-2
@@ -35,12 +35,12 @@ ge::graphStatus TransposeQuantBatchMatMulAswTiling::DoOpTiling()
35 runInfo_.baseK = BASIC_BLOCK_K_128_BYTE;35 runInfo_.baseK = BASIC_BLOCK_K_128_BYTE;
36 precisionMode_ = TQBMMPrecisionMode::PRECISION_MODE_MXFP4;36 precisionMode_ = TQBMMPrecisionMode::PRECISION_MODE_MXFP4;
37 apiLevel_ = TQBMMApiLevel::TENSOR_LEVEL;37 apiLevel_ = TQBMMApiLevel::TENSOR_LEVEL;
38- CalL1Tiling();38+ ApplyMxL1Tiling();
39 } else if (IsMxFp8(context_->GetInputDesc(X1_IDX), context_->GetInputDesc(X2_IDX),39 } else if (IsMxFp8(context_->GetInputDesc(X1_IDX), context_->GetInputDesc(X2_IDX),
40 context_->GetOptionalInputDesc(SCALE_X1_IDX), context_->GetOptionalInputDesc(SCALE_X2_IDX))) {40 context_->GetOptionalInputDesc(SCALE_X1_IDX), context_->GetOptionalInputDesc(SCALE_X2_IDX))) {
41 precisionMode_ = TQBMMPrecisionMode::PRECISION_MODE_MXFP8;41 precisionMode_ = TQBMMPrecisionMode::PRECISION_MODE_MXFP8;
42 apiLevel_ = TQBMMApiLevel::TENSOR_LEVEL;42 apiLevel_ = TQBMMApiLevel::TENSOR_LEVEL;
43- CalL1Tiling();43+ ApplyMxL1Tiling();
44 } else if (IsHIFP8(context_->GetInputDesc(X1_IDX), context_->GetInputDesc(X2_IDX))) {44 } else if (IsHIFP8(context_->GetInputDesc(X1_IDX), context_->GetInputDesc(X2_IDX))) {
45 precisionMode_ = TQBMMPrecisionMode::PRECISION_MODE_HIFP8;45 precisionMode_ = TQBMMPrecisionMode::PRECISION_MODE_HIFP8;
46 CalL1Tiling();46 CalL1Tiling();
@@ -92,6 +92,42 @@ void TransposeQuantBatchMatMulAswTiling::GetTransposeBatchMatMulInfo()
92 }92 }
93}93}
94 94 
95+void TransposeQuantBatchMatMulAswTiling::ApplyMxL1Tiling()
96+{
97+ TqbmmMxL1Params params;
98+ params.mSize = args_.mValue;
99+ params.kSize = args_.kValue;
100+ params.baseM = runInfo_.baseM;
101+ params.baseN = runInfo_.baseN;
102+ params.baseK = runInfo_.baseK;
103+ params.aDtype = args_.aType;
104+ params.bDtype = args_.bType;
105+ params.scaleDtype = context_->GetOptionalInputDesc(SCALE_X2_IDX)->GetDataType();
106+ params.perTokenScaleDtype = context_->GetOptionalInputDesc(SCALE_X1_IDX)->GetDataType();
107+ params.biasDtype = args_.biasType;
108+ params.cDtype = context_->GetOutputDesc(0)->GetDataType();
109+ params.hasBias = args_.hasBias;
110+ params.transB = args_.isBTrans;
111+ params.l1Size = compileInfo_.l1Size;
112+ params.l0CSize = compileInfo_.l0CSize;
113+ 
114+ const TqbmmMxL1Result res = CalcTqbmmMxL1Tiling(params);
115+ runInfo_.baseM = res.baseM;
116+ runInfo_.singleCoreM = res.baseM;
117+ runInfo_.stepKa = res.stepKa;
118+ runInfo_.stepKb = res.stepKb;
119+ runInfo_.depthA1 = res.depthA1;
120+ runInfo_.depthB1 = res.depthB1;
121+ runInfo_.dbL0C = res.dbL0C;
122+ runInfo_.l1BufferNum = res.nBufferNum;
123+ scaleKL1_ = static_cast<uint32_t>(res.scaleKL1);
124+ const uint64_t kL1 = static_cast<uint64_t>(res.stepKa) * runInfo_.baseK;
125+ const uint64_t scaleFactor = kL1 == 0UL ? SCALER_FACTOR_MIN : res.scaleKL1 / kL1;
126+ scaleFactorA_ = static_cast<uint32_t>(std::min(scaleFactor, static_cast<uint64_t>(SCALER_FACTOR_MAX)));
127+ scaleFactorB_ = scaleFactorA_;
128+ isMxL1Tiled_ = true;
129+}
130+ 
95void TransposeQuantBatchMatMulAswTiling::CalL1Tiling()131void TransposeQuantBatchMatMulAswTiling::CalL1Tiling()
96{132{
97 uint64_t leftL1Size = compileInfo_.l1Size;133 uint64_t leftL1Size = compileInfo_.l1Size;
@@ -265,6 +301,10 @@ ge::graphStatus TransposeQuantBatchMatMulAswTiling::GetTilingDataProcess(BatchMa
265 }301 }
266 }302 }
267 tilingData.batchSplitFactor = batchSplitFactor_;303 tilingData.batchSplitFactor = batchSplitFactor_;
304+ if (isMxL1Tiled_) {
305+ // tqbmm 无 streamk 模板，MX 借用空闲的 kTailCnt 承载全 K scale 预载窗口；0 表示随 kL1 窗口重载
306+ tilingData.matMulTilingData.kTailCnt = scaleKL1_;
307+ }
268 return ret;308 return ret;
269}309}
270 310 
Mmatmul/transpose_quant_batch_mat_mul/op_host/op_tiling/arch35/transpose_quant_batch_mat_mul_asw_tiling.h+4-0
@@ -15,6 +15,7 @@
15 15 
16#pragma once16#pragma once
17#include "matmul/mat_mul_v3/op_host/op_tiling/arch35/matmul_v3_base_tiling_advanced.h"17#include "matmul/mat_mul_v3/op_host/op_tiling/arch35/matmul_v3_base_tiling_advanced.h"
18+#include "transpose_quant_batch_mat_mul_mx_l1_tiling.h"
18#include "../../../op_kernel/arch35/transpose_quant_batch_mat_mul_tiling_key_public.h"19#include "../../../op_kernel/arch35/transpose_quant_batch_mat_mul_tiling_key_public.h"
19#include "../../../op_kernel/arch35/transpose_quant_batch_mat_mul_tiling_key.h"20#include "../../../op_kernel/arch35/transpose_quant_batch_mat_mul_tiling_key.h"
20 21 
@@ -41,6 +42,7 @@ protected:
41 42 
42 std::vector<size_t> GetWorkspaceSize() const override;43 std::vector<size_t> GetWorkspaceSize() const override;
43 void CalL1Tiling();44 void CalL1Tiling();
45+ void ApplyMxL1Tiling();
44 void CalStepKs();46 void CalStepKs();
45 uint64_t GetDepthA1B1(uint64_t leftSize, uint64_t perDepthSize, uint64_t depthInit);47 uint64_t GetDepthA1B1(uint64_t leftSize, uint64_t perDepthSize, uint64_t depthInit);
46 void CalScaleFactors(uint64_t baseASize, uint64_t baseBSize, uint64_t baseScaleASize, uint64_t baseScaleBSize);48 void CalScaleFactors(uint64_t baseASize, uint64_t baseBSize, uint64_t baseScaleASize, uint64_t baseScaleBSize);
@@ -55,7 +57,9 @@ private:
55 uint32_t batchSplitFactor_ = 1;57 uint32_t batchSplitFactor_ = 1;
56 uint32_t scaleFactorA_ = 1;58 uint32_t scaleFactorA_ = 1;
57 uint32_t scaleFactorB_ = 1;59 uint32_t scaleFactorB_ = 1;
60+ uint32_t scaleKL1_ = 0;
58 bool isMXFP4_ = false;61 bool isMXFP4_ = false;
62+ bool isMxL1Tiled_ = false;
59};63};
60} // namespace transpose_quant_batch_mat_mul_advanced64} // namespace transpose_quant_batch_mat_mul_advanced
61} // namespace optiling65} // namespace optiling
Amatmul/transpose_quant_batch_mat_mul/op_host/op_tiling/arch35/transpose_quant_batch_mat_mul_mx_l1_tiling.h+279-0
@@ -0,0 +1,279 @@
1+/**
2+ * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file transpose_quant_batch_mat_mul_mx_l1_tiling.h
13+ * \brief tqbmm MXFP8/MXFP4 的 L1 tiling，从 qbmm 基础模板
14+ * (matmul/quant_batch_matmul_v3/op_host/op_tiling/arch35/) 简化迁移。
15+ *
16+ * 迁移来源（函数级对应关系）：
17+ * CalcTqbmmMxL1Tiling <- AdaptiveSlidingWindowMXBasicAPITiling::CalL1Tiling(:151-186)
18+ * + CalculateNBufferNum(:220-268) 主干
19+ * CalcTqbmmMxDepth <- L1TilingDataCalculator::GetDepthA1B1(l1_tiling_data_calculator.cpp:541-564)
20+ * CalcTqbmmMxStepKs <- L1TilingDataCalculator::CalStepKs(l1_tiling_data_calculator.cpp:348-396)
21+ * CalcTqbmmMxScaleKL1 <- L1TilingDataCalculator::CalScaleFactors(l1_tiling_data_calculator.cpp:398-439)
22+ * CalcTqbmmMxDbL0C <- CalL1Tiling 的 dbL0c 三目(mx_basic_api_tiling.cpp:165-170)
23+ * + CheckBiasAndScale(adaptive_sliding_window_tiling.cpp:358-369)
24+ * CalcTqbmmMxUsedL1Size <- CalcUsedL1Size(mx_basic_api_tiling.cpp:368-389)
25+ * GetTqbmmMxHalfFallback <- GetHalfKFallbackScaleKL1(mx_basic_api_tiling.cpp:306-320)
26+ * GetTqbmmMxFullCoverKL1 <- GetFullCoverScaleKL1IfPossible(mx_basic_api_tiling.cpp:322-333)
27+ * CanTqbmmMxReduceStepK <- CanReduceStepKToTwo(mx_basic_api_tiling.cpp:286-304)
28+ *
29+ * 相对 qbmm 删除的逻辑（tqbmm 用不到）：
30+ * 1. A 全载（isAFullLoad，tqbmm 无 FULL_LOAD_MODE=1 内核实例，batchC 也不受限为 1）
31+ * 2. 3-buffer 分支（qbmm 用 IsMxMte2Bound 带宽模型门控；tqbmm MVP 不接，只做 4->2 兜底）
32+ * 3. isTensorapiCapable 判断（tqbmm MX 恒为 TENSOR_LEVEL/Blaze）
33+ * 4. transA 分支（tqbmm perm_x1 恒 [1,0,2]，isATrans 恒 false）
34+ * 5. L0C 乒乓模板的 scaleKL1 收缩（AdjustScaleFactorForL0CPingpong，qbmm KERNELTYPE=9/10 专用）
35+ * 6. isPertoken（MX 恒非 pertoken，isUbQuant 只看 cDtype==BF16）
36+ *
37+ * 出参约定：tiling data 沿用 BatchMatMulV3TilingData 不新增字段，MX 模式下用
38+ * matMulTilingData.kTailCnt（tqbmm 无 streamk 模板，该字段空闲）承载 scaleKL1，
39+ * kernel 侧 transpose_quant_batch_mat_mul_mx.h 读 kTailCnt，0 表示退化为 kL1。
40+ */
41+ 
42+#pragma once
43+ 
44+#include <algorithm>
45+#include <cstdint>
46+ 
47+#include "graph/types.h"
48+ 
49+namespace optiling {
50+namespace transpose_quant_batch_mat_mul_advanced {
51+ 
52+constexpr uint64_t TQBMM_MX_GROUP_SIZE = 32UL;
53+constexpr uint64_t TQBMM_MX_DIVISOR_SIZE = 64UL;
54+constexpr uint64_t TQBMM_MX_MULTI_BASE_SIZE = 2UL;
55+constexpr uint64_t TQBMM_MX_ALIGN_128B = 128UL;
56+constexpr uint64_t TQBMM_MX_ALIGN_256B = 256UL;
57+constexpr uint64_t TQBMM_MX_SCALE_K_ESTIMATE = 4096UL;
58+constexpr uint64_t TQBMM_MX_MAX_STEP_K = 4UL;
59+constexpr uint64_t TQBMM_MX_DOUBLE_BUFFER = 2UL;
60+constexpr uint64_t TQBMM_MX_FOUR_BUFFER = 4UL;
61+constexpr uint64_t TQBMM_MX_CUBE_BLOCK = 16UL;
62+ 
63+struct TqbmmMxL1Params {
64+ uint64_t mSize = 0UL;
65+ uint64_t kSize = 0UL;
66+ uint64_t baseM = 0UL;
67+ uint64_t baseN = 0UL;
68+ uint64_t baseK = 0UL;
69+ ge::DataType aDtype = ge::DT_FLOAT8_E4M3FN;
70+ ge::DataType bDtype = ge::DT_FLOAT8_E4M3FN;
71+ ge::DataType scaleDtype = ge::DT_FLOAT8_E8M0;
72+ ge::DataType perTokenScaleDtype = ge::DT_FLOAT8_E8M0;
73+ ge::DataType biasDtype = ge::DT_FLOAT;
74+ ge::DataType cDtype = ge::DT_BF16;
75+ bool hasBias = false;
76+ bool transB = false;
77+ uint64_t l1Size = 0UL;
78+ uint64_t l0CSize = 0UL;
79+};
80+ 
81+struct TqbmmMxL1Result {
82+ uint64_t baseM = 0UL;
83+ uint32_t stepKa = 1U;
84+ uint32_t stepKb = 1U;
85+ uint32_t depthA1 = 2U;
86+ uint32_t depthB1 = 2U;
87+ uint64_t scaleKL1 = 0UL;
88+ uint32_t nBufferNum = 2U;
89+ uint32_t dbL0C = 1U;
90+};
91+ 
92+inline uint64_t TqbmmMxCeilDiv(uint64_t a, uint64_t b) { return (a + b - 1UL) / b; }
93+ 
94+inline uint64_t TqbmmMxCeilAlign(uint64_t a, uint64_t b) { return TqbmmMxCeilDiv(a, b) * b; }
95+ 
96+inline uint64_t TqbmmMxDataSize(uint64_t elemCnt, ge::DataType dtype)
97+{
98+ if (dtype == ge::DT_INT4 || dtype == ge::DT_FLOAT4_E2M1 || dtype == ge::DT_FLOAT4_E1M2) {
99+ return (elemCnt + 1UL) >> 1UL;
100+ }
101+ return elemCnt * static_cast<uint64_t>(ge::GetSizeByDataType(dtype));
102+}
103+ 
104+inline uint64_t CalcTqbmmMxDepth(const TqbmmMxL1Params& p, uint64_t leftL1)
105+{
106+ const uint64_t baseABSize = TqbmmMxDataSize(p.baseM * p.baseK, p.aDtype) +
107+ TqbmmMxDataSize(p.baseN * p.baseK, p.bDtype);
108+ const uint64_t baseScaleSize = TqbmmMxDataSize(p.baseM, p.perTokenScaleDtype) +
109+ TqbmmMxDataSize(p.baseN, p.scaleDtype);
110+ uint64_t scaleKL1 = std::min(p.kSize, TQBMM_MX_SCALE_K_ESTIMATE);
111+ uint64_t scaleL1Size = baseScaleSize * TqbmmMxCeilDiv(scaleKL1, TQBMM_MX_DIVISOR_SIZE) * TQBMM_MX_MULTI_BASE_SIZE *
112+ TQBMM_MX_DOUBLE_BUFFER;
113+ uint64_t depth = 1UL;
114+ while (depth * baseABSize + scaleL1Size <= leftL1) {
115+ depth *= 2UL;
116+ const uint64_t kL1 = depth / TQBMM_MX_DOUBLE_BUFFER * p.baseK;
117+ if (kL1 > scaleKL1) {
118+ scaleKL1 = kL1;
119+ scaleL1Size = baseScaleSize * TqbmmMxCeilDiv(scaleKL1, TQBMM_MX_DIVISOR_SIZE) * TQBMM_MX_MULTI_BASE_SIZE *
120+ TQBMM_MX_DOUBLE_BUFFER;
121+ }
122+ }
123+ return depth == 1UL ? depth : depth / 2UL;
124+}
125+ 
126+inline void CalcTqbmmMxStepKs(const TqbmmMxL1Params& p, uint64_t depth, uint32_t& stepKa, uint32_t& stepKb,
127+ uint32_t& depthA1, uint32_t& depthB1)
128+{
129+ uint64_t ka = std::max(depth / TQBMM_MX_DOUBLE_BUFFER, 1UL);
130+ uint64_t kb = ka;
131+ if (ka * p.baseK > p.kSize) {
132+ ka = TqbmmMxCeilDiv(p.kSize, p.baseK);
133+ }
134+ if (kb * p.baseK > p.kSize) {
135+ kb = TqbmmMxCeilDiv(p.kSize, p.baseK);
136+ }
137+ if (ka > kb) {
138+ ka = ka / kb * kb;
139+ }
140+ if (kb > ka) {
141+ kb = kb / ka * ka;
142+ }
143+ ka = std::min(ka, TQBMM_MX_MAX_STEP_K);
144+ kb = std::min(kb, TQBMM_MX_MAX_STEP_K);
145+ stepKa = static_cast<uint32_t>(ka);
146+ stepKb = static_cast<uint32_t>(kb);
147+ depthA1 = static_cast<uint32_t>(ka * TQBMM_MX_DOUBLE_BUFFER);
148+ depthB1 = static_cast<uint32_t>(kb * TQBMM_MX_DOUBLE_BUFFER);
149+}
150+ 
151+inline uint64_t CalcTqbmmMxScaleKL1(const TqbmmMxL1Params& p, uint64_t leftL1, uint32_t stepKa, uint32_t stepKb,
152+ uint32_t depthA1, uint32_t depthB1)
153+{
154+ const uint64_t usedABSize = TqbmmMxDataSize(p.baseM * p.baseK, p.aDtype) * depthA1 +
155+ TqbmmMxDataSize(p.baseN * p.baseK, p.bDtype) * depthB1;
156+ if (leftL1 <= usedABSize) {
157+ return 0UL;
158+ }
159+ const uint64_t leftAfterAB = leftL1 - usedABSize;
160+ const uint64_t kL1 = std::min(stepKa, stepKb) * p.baseK;
161+ if (kL1 == 0UL) {
162+ return 0UL;
163+ }
164+ const uint64_t scaleGroupL1Size = (TqbmmMxDataSize(p.baseM, p.perTokenScaleDtype) +
165+ TqbmmMxDataSize(p.baseN, p.scaleDtype)) *
166+ TQBMM_MX_MULTI_BASE_SIZE * TQBMM_MX_DOUBLE_BUFFER;
167+ const uint64_t maxScaleKL1ByL1 = leftAfterAB / scaleGroupL1Size * TQBMM_MX_DIVISOR_SIZE;
168+ const uint64_t maxScaleKL1 = TqbmmMxCeilAlign(p.kSize, kL1);
169+ uint64_t scaleKL1 = std::min(maxScaleKL1ByL1, maxScaleKL1);
170+ return scaleKL1 / kL1 * kL1;
171+}
172+ 
173+inline uint32_t CalcTqbmmMxDbL0C(const TqbmmMxL1Params& p)
174+{
175+ constexpr uint32_t db = 2U;
176+ constexpr uint64_t maxBiasBaseN = 512UL;
177+ constexpr uint64_t maxScaleBaseN = 256UL;
178+ const bool isBiasInvalid = p.hasBias && p.biasDtype != ge::DT_BF16 && p.baseN > maxBiasBaseN;
179+ const bool isUbQuant = p.cDtype == ge::DT_BF16;
180+ const bool isScaleInvalid = !isUbQuant && p.baseN > maxScaleBaseN;
181+ const bool l0cEnough = p.baseM * p.baseN * 4UL * TQBMM_MX_DOUBLE_BUFFER <= p.l0CSize;
182+ return (l0cEnough && !isBiasInvalid && !isScaleInvalid) ? db : 1U;
183+}
184+ 
185+inline uint64_t CalcTqbmmMxUsedL1Size(const TqbmmMxL1Params& p, uint64_t kL1, uint64_t scaleKL1, uint32_t l1BufferNum)
186+{
187+ uint64_t usedL1Size = TqbmmMxDataSize(p.baseN * kL1, p.bDtype) * l1BufferNum;
188+ usedL1Size += TqbmmMxDataSize(p.baseN * TqbmmMxCeilDiv(scaleKL1, TQBMM_MX_GROUP_SIZE), p.scaleDtype) *
189+ TQBMM_MX_DOUBLE_BUFFER;
190+ if (p.hasBias) {
191+ usedL1Size += TqbmmMxDataSize(p.baseN, p.biasDtype) * TQBMM_MX_DOUBLE_BUFFER;
192+ }
193+ usedL1Size += TqbmmMxDataSize(p.baseM * kL1, p.aDtype) * l1BufferNum;
194+ usedL1Size += TqbmmMxDataSize(p.baseM * TqbmmMxCeilDiv(scaleKL1, TQBMM_MX_GROUP_SIZE), p.perTokenScaleDtype) *
195+ TQBMM_MX_DOUBLE_BUFFER;
196+ return usedL1Size;
197+}
198+ 
199+inline uint64_t GetTqbmmMxHalfFallback(const TqbmmMxL1Params& p, uint64_t scaleKL1Init, uint64_t kL1)
200+{
201+ uint64_t scaleKL1 = scaleKL1Init;
202+ if (scaleKL1 % TQBMM_MX_SCALE_K_ESTIMATE == 0UL) {
203+ return scaleKL1;
204+ }
205+ const uint64_t halfK = TqbmmMxCeilDiv(p.kSize, 2UL);
206+ if (scaleKL1 > halfK && scaleKL1 < p.kSize) {
207+ const uint64_t adjusted = TqbmmMxCeilAlign(halfK, kL1);
208+ scaleKL1 = adjusted < scaleKL1 ? adjusted : scaleKL1;
209+ }
210+ return scaleKL1;
211+}
212+ 
213+inline uint64_t GetTqbmmMxFullCoverKL1(const TqbmmMxL1Params& p, uint64_t kL1, uint64_t scaleKL1, uint32_t l1BufferNum)
214+{
215+ const uint64_t fullCover = TqbmmMxCeilAlign(p.kSize, kL1);
216+ if (fullCover <= scaleKL1) {
217+ return scaleKL1;
218+ }
219+ return CalcTqbmmMxUsedL1Size(p, kL1, fullCover, l1BufferNum) <= p.l1Size ? fullCover : scaleKL1;
220+}
221+ 
222+inline bool CanTqbmmMxReduceStepK(const TqbmmMxL1Params& p, uint64_t stepKTwoKL1)
223+{
224+ const bool isAInnerKAligned = TqbmmMxDataSize(p.kSize, p.aDtype) % TQBMM_MX_ALIGN_128B == 0UL;
225+ const bool isBInnerKAligned = !p.transB || TqbmmMxDataSize(p.kSize, p.bDtype) % TQBMM_MX_ALIGN_128B == 0UL;
226+ const bool isStepKTwoAAligned = TqbmmMxDataSize(stepKTwoKL1, p.aDtype) % TQBMM_MX_ALIGN_256B == 0UL;
227+ const bool isStepKTwoBAligned = !p.transB || TqbmmMxDataSize(stepKTwoKL1, p.bDtype) % TQBMM_MX_ALIGN_256B == 0UL;
228+ return isAInnerKAligned && isBInnerKAligned && isStepKTwoAAligned && isStepKTwoBAligned;
229+}
230+ 
231+inline TqbmmMxL1Result CalcTqbmmMxL1Tiling(const TqbmmMxL1Params& p)
232+{
233+ TqbmmMxL1Result res;
234+ res.baseM = std::min(p.baseM, TqbmmMxCeilAlign(p.mSize, TQBMM_MX_CUBE_BLOCK));
235+ 
236+ TqbmmMxL1Params aligned = p;
237+ aligned.baseM = res.baseM;
238+ res.dbL0C = CalcTqbmmMxDbL0C(aligned);
239+ 
240+ uint64_t leftL1 = aligned.l1Size;
241+ if (aligned.hasBias) {
242+ leftL1 -= TqbmmMxDataSize(aligned.baseN, aligned.biasDtype) * TQBMM_MX_DOUBLE_BUFFER;
243+ }
244+ const uint64_t depth = CalcTqbmmMxDepth(aligned, leftL1);
245+ CalcTqbmmMxStepKs(aligned, depth, res.stepKa, res.stepKb, res.depthA1, res.depthB1);
246+ const uint64_t scaleKL1Init = CalcTqbmmMxScaleKL1(aligned, leftL1, res.stepKa, res.stepKb, res.depthA1,
247+ res.depthB1);
248+ 
249+ const uint64_t stepK = std::min(res.stepKa, res.stepKb);
250+ const uint64_t currentKL1 = stepK * aligned.baseK;
251+ uint64_t currentScaleKL1 = GetTqbmmMxHalfFallback(aligned, scaleKL1Init, currentKL1);
252+ if (CalcTqbmmMxUsedL1Size(aligned, currentKL1, currentScaleKL1, TQBMM_MX_FOUR_BUFFER) <= aligned.l1Size) {
253+ res.scaleKL1 = GetTqbmmMxFullCoverKL1(aligned, currentKL1, currentScaleKL1, TQBMM_MX_FOUR_BUFFER);
254+ res.nBufferNum = TQBMM_MX_FOUR_BUFFER;
255+ return res;
256+ }
257+ 
258+ const uint64_t stepKTwoKL1 = 2UL * aligned.baseK;
259+ const bool canReduceStepK = currentKL1 * TQBMM_MX_DOUBLE_BUFFER < aligned.kSize && (stepK == 3UL || stepK == 4UL) &&
260+ CanTqbmmMxReduceStepK(aligned, stepKTwoKL1);
261+ if (canReduceStepK) {
262+ const uint64_t twoScaleKL1 = GetTqbmmMxHalfFallback(aligned, scaleKL1Init, stepKTwoKL1);
263+ if (CalcTqbmmMxUsedL1Size(aligned, stepKTwoKL1, twoScaleKL1, TQBMM_MX_FOUR_BUFFER) <= aligned.l1Size) {
264+ res.stepKa = static_cast<uint32_t>(stepKTwoKL1 / aligned.baseK);
265+ res.stepKb = res.stepKa;
266+ res.depthA1 = static_cast<uint32_t>(res.stepKa * TQBMM_MX_DOUBLE_BUFFER);
267+ res.depthB1 = res.depthA1;
268+ res.scaleKL1 = GetTqbmmMxFullCoverKL1(aligned, stepKTwoKL1, twoScaleKL1, TQBMM_MX_FOUR_BUFFER);
269+ res.nBufferNum = TQBMM_MX_FOUR_BUFFER;
270+ return res;
271+ }
272+ }
273+ 
274+ res.scaleKL1 = currentScaleKL1;
275+ res.nBufferNum = 2U;
276+ return res;
277+}
278+} // namespace transpose_quant_batch_mat_mul_advanced
279+} // namespace optiling
Mmatmul/transpose_quant_batch_mat_mul/op_kernel/arch35/transpose_quant_batch_mat_mul_mx.h+5-1
@@ -48,9 +48,13 @@ __aicore__ inline void TqbmmMxTensorApiKernel(GM_ADDR aGM, GM_ADDR bGM, GM_ADDR
48 using MatmulKernel = Blaze::Gemm::Kernel::GemmUniversal<ProblemShape, BlockMmad, BlockEpilogue, BlockScheduler>;48 using MatmulKernel = Blaze::Gemm::Kernel::GemmUniversal<ProblemShape, BlockMmad, BlockEpilogue, BlockScheduler>;
49 using Params = typename MatmulKernel::Params;49 using Params = typename MatmulKernel::Params;
50 const auto& tCubeTiling = tilingData.matMulTilingData.tCubeTiling;50 const auto& tCubeTiling = tilingData.matMulTilingData.tCubeTiling;
51+ // host 侧 MX 借用空闲的 kTailCnt 承载 scaleKL1（全 K scale 预载窗口）；0 表示随 kL1 窗口重载
52+ const uint64_t mxScaleKL1 = tilingData.matMulTilingData.kTailCnt != 0 ?
53+ static_cast<uint64_t>(tilingData.matMulTilingData.kTailCnt) :
54+ static_cast<uint64_t>(tilingData.kL1);
51 Params params = {{tCubeTiling.M, tCubeTiling.N, tCubeTiling.Ka, tilingData.cBatchDimAll},55 Params params = {{tCubeTiling.M, tCubeTiling.N, tCubeTiling.Ka, tilingData.cBatchDimAll},
52 {aGM, bGM, cGM, bias, perTokenScale, scale},56 {aGM, bGM, cGM, bias, perTokenScale, scale},
53- {tilingData.kL1, tilingData.kL1, tilingData.l1BufferNum},57+ {tilingData.kL1, mxScaleKL1, tilingData.l1BufferNum},
54 {static_cast<int64_t>(tCubeTiling.baseM), static_cast<int64_t>(tCubeTiling.baseN),58 {static_cast<int64_t>(tCubeTiling.baseM), static_cast<int64_t>(tCubeTiling.baseN),
55 tilingData.matMulTilingData.mTailCnt, tilingData.matMulTilingData.nTailCnt,59 tilingData.matMulTilingData.mTailCnt, tilingData.matMulTilingData.nTailCnt,
56 tilingData.matMulTilingData.mBaseTailSplitCnt, tilingData.matMulTilingData.nBaseTailSplitCnt,60 tilingData.matMulTilingData.mBaseTailSplitCnt, tilingData.matMulTilingData.nBaseTailSplitCnt,
Mmatmul/transpose_quant_batch_mat_mul/tests/ut/op_host/test_transpose_quant_batch_mat_mul_tiling.cpp+11-10
@@ -437,9 +437,9 @@ static TilingTestParam ascend950_cases_params[] = {
437 0,437 0,
438 32,438 32,
439 135185UL,439 135185UL,
440- "32 35 744 192 192 256 256 192 256 256 128 4 4 1 1 0 0 0 0 18432 8192 0 1 1 1 1 2 2 0 0 2 2 1 0 0 0 0 0 0 0 0 "440+ "32 35 744 192 192 48 256 192 48 256 128 4 4 1 1 0 0 0 0 18432 8192 0 1 1 1 1 2 2 0 0 2 2 2 0 0 0 0 0 0 0 0 "
441- "0 0 0 0 0 0 0 0 16843009 1 1 1 1 1 0 0 0 4 0 32 32 32 1 1 1 1 1 1 1 1 1 1 32 32 32 0 1 1 48 256 256 16843264 "441+ "0 0 0 0 0 0 0 0 16843009 1 1 256 1 1 0 0 0 4 0 32 32 32 1 1 1 1 1 1 1 1 1 1 32 32 32 0 1 1 48 256 256 16843776 "
442- "256 1 0 1 1 4 4 ",442+ "48 1 0 1 1 4 4 ",
443 1,443 1,
444 32,444 32,
445 {1, 0, 2},445 {1, 0, 2},
@@ -542,9 +542,9 @@ static TilingTestParam ascend950_cases_params[] = {
542 0,542 0,
543 32,543 32,
544 143377UL,544 143377UL,
545- "32 35 744 192 192 256 256 192 256 256 128 4 4 1 1 0 0 0 0 9216 8192 0 1 1 1 1 2 2 0 0 2 2 1 0 0 0 0 0 0 0 0 "545+ "32 35 744 192 192 48 256 192 48 256 128 4 4 1 1 0 0 0 0 9216 8192 0 1 1 1 1 2 2 0 0 2 2 2 0 0 0 0 0 0 0 0 "
546- "0 0 0 0 0 0 0 0 16843009 1 1 1 1 1 0 0 0 4 2 32 32 32 1 1 1 1 1 1 1 1 1 1 32 32 32 0 1 1 48 256 256 16843264 "546+ "0 0 0 0 0 0 0 0 16843009 1 1 256 1 1 0 0 0 4 2 32 32 32 1 1 1 1 1 1 1 1 1 1 32 32 32 0 1 1 48 256 256 16843776 "
547- "256 1 0 1 1 4 4 ",547+ "48 1 0 1 1 4 4 ",
548 1,548 1,
549 32,549 32,
550 {1, 0, 2},550 {1, 0, 2},
@@ -632,9 +632,10 @@ TEST(TransposeQuantBatchMatMulTilingRuntimeExtra, WeightNzMxFp8Success_0)
632 ASSERT_EQ(opImpl->tiling(tiling_context), ge::GRAPH_SUCCESS);632 ASSERT_EQ(opImpl->tiling(tiling_context), ge::GRAPH_SUCCESS);
633 uint64_t tiling_key = tiling_context->GetTilingKey();633 uint64_t tiling_key = tiling_context->GetTilingKey();
634 auto tiling_data_result = TilingData2Str(tiling_context->GetRawTilingData());634 auto tiling_data_result = TilingData2Str(tiling_context->GetRawTilingData());
635- string golden_tiling_data = "32 425 64 640 640 256 256 640 256 256 128 4 4 1 1 0 0 0 0 98304 8192 0 1 1 1 1 2 2 0 "635+ string
636- "0 2 2 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 16843266 1 1 1 1 1 0 0 0 4 1 2057 2057 2057 1 "636+ golden_tiling_data = "32 425 64 640 640 256 256 640 256 256 128 4 4 1 1 0 0 0 0 98304 8192 0 1 1 1 1 2 2 0 "
637- "1 1 1 1 1 1 1 1 1 2057 2057 2057 0 1 1 256 64 256 16843264 256 1 0 1 1 4 4 ";637+ "0 2 2 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 16843523 1 1 768 1 1 0 0 0 4 1 2057 2057 2057 1 "
638+ "1 1 1 1 1 1 1 1 1 2057 2057 2057 0 1 1 256 64 256 16843264 256 1 0 1 1 4 4 ";
638 ASSERT_EQ(tiling_key, 135169);639 ASSERT_EQ(tiling_key, 135169);
639 ASSERT_EQ(tiling_data_result, golden_tiling_data);640 ASSERT_EQ(tiling_data_result, golden_tiling_data);
640}641}
@@ -714,7 +715,7 @@ TEST(TransposeQuantBatchMatMulTilingRuntimeExtra, WeightNzMxFp8Success_1)
714 uint64_t tiling_key = tiling_context->GetTilingKey();715 uint64_t tiling_key = tiling_context->GetTilingKey();
715 auto tiling_data_result = TilingData2Str(tiling_context->GetRawTilingData());716 auto tiling_data_result = TilingData2Str(tiling_context->GetRawTilingData());
716 string golden_tiling_data = "32 284 192 1472 1472 256 256 1472 256 256 128 4 4 1 1 0 0 0 0 172032 12288 0 1 1 1 "717 string golden_tiling_data = "32 284 192 1472 1472 256 256 1472 256 256 128 4 4 1 1 0 0 0 0 172032 12288 0 1 1 1 "
717- "1 2 2 0 0 2 2 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 16844037 1 1 1 1 1 0 0 0 4 0 1172 "718+ "1 2 2 0 0 2 2 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 16844294 1 1 1536 1 1 0 0 0 4 0 1172 "
718 "1172 1172 1 1 1 1 1 1 1 1 1 1 1172 1172 1172 0 1 1 256 192 256 16843264 256 1 0 1 1 4 "719 "1172 1172 1 1 1 1 1 1 1 1 1 1 1172 1172 1172 0 1 1 256 192 256 16843264 256 1 0 1 1 4 "
719 "4 ";720 "4 ";
720 ASSERT_EQ(tiling_key, 135185);721 ASSERT_EQ(tiling_key, 135185);