[INFO] Initializing environment for https://gitcode.com/pre-commit/pre-commit-hooks. [WARNING] repo `https://gitcode.com/pre-commit/pre-commit-hooks` uses deprecated stage names (commit, push) which will be removed in a future version. Hint: often `pre-commit autoupdate --repo https://gitcode.com/pre-commit/pre-commit-hooks` will fix this. if it does not -- consider reporting an issue to that repo. [INFO] Initializing environment for https://gitcode.com/pre-commit-clang/mirrors-clang-format. [INFO] Initializing environment for https://gitcode.com/gh_mirrors/ru/ruff-pre-commit. [INFO] Initializing environment for https://gitcode.com/gh_mirrors/co/codespell. [INFO] Installing environment for https://gitcode.com/pre-commit/pre-commit-hooks. [INFO] Once installed this environment will be reused. [INFO] This may take a few minutes... [INFO] Installing environment for https://gitcode.com/pre-commit-clang/mirrors-clang-format. [INFO] Once installed this environment will be reused. [INFO] This may take a few minutes... [INFO] Installing environment for https://gitcode.com/gh_mirrors/ru/ruff-pre-commit. [INFO] Once installed this environment will be reused. [INFO] This may take a few minutes... [INFO] Installing environment for https://gitcode.com/gh_mirrors/co/codespell. [INFO] Once installed this environment will be reused. [INFO] This may take a few minutes... trim trailing whitespace.................................................Failed - hook id: trailing-whitespace - exit code: 1 - files were modified by this hook Fixing include/blaze/epilogue/block/block_epilogue_gelu_mx_quant.h fix end of files.........................................................Passed check yaml...........................................(no files to check)Skipped check for added large files..............................................Passed check for merge conflicts................................................Passed detect private key.......................................................Passed check json...........................................(no files to check)Skipped clang-format.............................................................Failed - hook id: clang-format - files were modified by this hook Formatting [1/4] examples/quant_matmul_activation_quant/quant_matmul_activation_quant/quant_matmul_activation_quant.cpp Formatting [2/4] include/blaze/gemm/block/block_scheduler_qbmm.h Formatting [3/4] tests/ut/op_kernel/quant_matmul_activation_quant/test_quant_matmul_activation_quant.cpp Formatting [4/4] include/blaze/epilogue/block/block_epilogue_gelu_mx_quant.h Formatting [1/2] include/blaze/gemm/kernel/kernel_qbmm_mx_activation_quant.h Formatting [2/2] tests/ut/op_kernel/quant_matmul_activation_quant/quant_matmul_activation_quant.h ruff check...............................................................Failed - hook id: ruff-check - exit code: 1 ::error title=Ruff (F841),file=/opt/cloud/agent_1766032928179_phenv/workspace/j_Ji5eolyu/examples/quant_matmul_activation_quant/scripts/gen_data.py,line=108,col=5,endLine=108,endColumn=12::examples/quant_matmul_activation_quant/scripts/gen_data.py:108:5: F841 Local variable `trans_a` is assigned to but never used ::error title=Ruff (F841),file=/opt/cloud/agent_1766032928179_phenv/workspace/j_Ji5eolyu/examples/quant_matmul_activation_quant/scripts/gen_data.py,line=109,col=5,endLine=109,endColumn=12::examples/quant_matmul_activation_quant/scripts/gen_data.py:109:5: F841 Local variable `trans_b` is assigned to but never used ruff format..............................................................Failed - hook id: ruff-format - files were modified by this hook 3 files reformatted codespell................................................................Failed - hook id: codespell - exit code: 65 examples/quant_matmul_activation_quant/quant_matmul_activation_quant/README.md:39: nd ==> and, 2nd examples/quant_matmul_activation_quant/quant_matmul_activation_quant/README.md:40: nd ==> and, 2nd All changes made by hooks: diff --git a/examples/quant_matmul_activation_quant/quant_matmul_activation_quant/parse_csv.py b/examples/quant_matmul_activation_quant/quant_matmul_activation_quant/parse_csv.py index b8474b3..9e8cc4b 100644 --- a/examples/quant_matmul_activation_quant/quant_matmul_activation_quant/parse_csv.py +++ b/examples/quant_matmul_activation_quant/quant_matmul_activation_quant/parse_csv.py @@ -41,10 +41,20 @@ def main(): for row in csv.DictReader(stream): case_dir = os.path.join(data_root, row["casename"]) generate = [ - sys.executable, os.path.join(scripts_dir, "gen_data.py"), - "--m", row["m"], "--k", row["k"], "--n", row["n"], - "--bias", row["bias"], "--layout", row["layout"], - "--output-dir", case_dir, + sys.executable, + os.path.join(scripts_dir, "gen_data.py"), + "--m", + row["m"], + "--k", + row["k"], + "--n", + row["n"], + "--bias", + row["bias"], + "--layout", + row["layout"], + "--output-dir", + case_dir, ] status = "PASS" stage = "verify" @@ -52,10 +62,23 @@ def main(): if run(generate) != 0: status, stage, message = "FAIL", "gen_data", "data generation failed" else: - launch = [args.executable, row["m"], row["k"], row["n"], row["bias"], row["layout"]] + launch = [ + args.executable, + row["m"], + row["k"], + row["n"], + row["bias"], + row["layout"], + ] tiling_keys = ( - "base_m", "base_n", "base_k", "tile_k_l1", "scale_k_l1", - "l1_buffers", "block_num") + "base_m", + "base_n", + "base_k", + "tile_k_l1", + "scale_k_l1", + "l1_buffers", + "block_num", + ) for key in tiling_keys: launch.append(row[key]) launch.append(case_dir) @@ -63,22 +86,33 @@ def main(): status, stage, message = "FAIL", "kernel", "kernel execution failed" else: verify = [ - sys.executable, os.path.join(scripts_dir, "verify_result.py"), + sys.executable, + os.path.join(scripts_dir, "verify_result.py"), os.path.join(case_dir, "golden_y.bin"), os.path.join(case_dir, "npu_y.bin"), os.path.join(case_dir, "golden_y_scale.bin"), os.path.join(case_dir, "npu_y_scale.bin"), ] if run(verify) != 0: - status, stage, message = "FAIL", "verify", "golden comparison failed" + status, stage, message = ( + "FAIL", + "verify", + "golden comparison failed", + ) print(f"[{status}] {row['casename']}") - results.append({ - "casename": row["casename"], "status": status, - "stage": stage, "message": message}) + results.append( + { + "casename": row["casename"], + "status": status, + "stage": stage, + "message": message, + } + ) with open(args.result_path, "w", newline="") as stream: writer = csv.DictWriter( - stream, fieldnames=("casename", "status", "stage", "message")) + stream, fieldnames=("casename", "status", "stage", "message") + ) writer.writeheader() writer.writerows(results) if any(item["status"] != "PASS" for item in results): diff --git a/examples/quant_matmul_activation_quant/quant_matmul_activation_quant/quant_matmul_activation_quant.cpp b/examples/quant_matmul_activation_quant/quant_matmul_activation_quant/quant_matmul_activation_quant.cpp index 827855c..3e1cac4 100644 --- a/examples/quant_matmul_activation_quant/quant_matmul_activation_quant/quant_matmul_activation_quant.cpp +++ b/examples/quant_matmul_activation_quant/quant_matmul_activation_quant/quant_matmul_activation_quant.cpp @@ -74,9 +74,9 @@ bool ParseArgs(int argc, const char** argv, ExampleConfig& config) { if (argc != 14) { std::fprintf(stderr, - "Usage: %s " - " \n", - argv[0]); + "Usage: %s " + " \n", + argv[0]); return false; } @@ -156,10 +156,7 @@ public: DeviceBuffer(const DeviceBuffer&) = delete; DeviceBuffer& operator=(const DeviceBuffer&) = delete; - uint8_t* Get() const - { - return data_; - } + uint8_t* Get() const { return data_; } void CopyFromFile(const std::string& path) const { @@ -185,11 +182,12 @@ private: }; template -__global__ __aicore__ void QuantMatmulActivationQuantKernel( - GM_ADDR x1Gm, GM_ADDR x2Gm, GM_ADDR biasGm, GM_ADDR x1ScaleGm, GM_ADDR x2ScaleGm, - GM_ADDR yGm, GM_ADDR yScaleGm, - int64_t m, int64_t k, int64_t n, uint64_t baseM, uint64_t baseN, uint64_t baseK, - uint64_t tileShapeKL1, uint64_t tileShapeScaleKL1, uint64_t l1BufferNum, uint64_t biasElements) +__global__ __aicore__ void QuantMatmulActivationQuantKernel(GM_ADDR x1Gm, GM_ADDR x2Gm, GM_ADDR biasGm, + GM_ADDR x1ScaleGm, GM_ADDR x2ScaleGm, GM_ADDR yGm, + GM_ADDR yScaleGm, int64_t m, int64_t k, int64_t n, + uint64_t baseM, uint64_t baseN, uint64_t baseK, + uint64_t tileShapeKL1, uint64_t tileShapeScaleKL1, + uint64_t l1BufferNum, uint64_t biasElements) { KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); AscendC::InitSocState(); @@ -204,34 +202,43 @@ __global__ __aicore__ void QuantMatmulActivationQuantKernel( using LayoutC = AscendC::Te::NDExtLayoutPtn; using ProblemShape = AscendC::Te::Shape; - using DispatchPolicy = Blaze::Gemm::MatmulWithScaleMx<0, false, - Blaze::Gemm::KernelMmadWithScaleMxActivationQuant, Blaze::Gemm::L0C2UB_MODE_DUAL_DST_SPLIT_M>; - using BlockScheduler = Blaze::Gemm::Block::BlockSchedulerQuantBatchMatmulV3< - ProblemShape, 0, LayoutA, LayoutB, AType>; - using BlockMmad = Blaze::Gemm::Block::BlockMmad< - DispatchPolicy, AType, LayoutA, BType, LayoutB, MatmulOutType, LayoutC, BiasType, LayoutC>; + using DispatchPolicy = Blaze::Gemm::MatmulWithScaleMx<0, false, Blaze::Gemm::KernelMmadWithScaleMxActivationQuant, + Blaze::Gemm::L0C2UB_MODE_DUAL_DST_SPLIT_M>; + using BlockScheduler = Blaze::Gemm::Block::BlockSchedulerQuantBatchMatmulV3; + using BlockMmad = Blaze::Gemm::Block::BlockMmad; using BlockEpilogue = Blaze::Epilogue::Block::BlockEpilogueGeluMxQuant; using Kernel = Blaze::Gemm::Kernel::GemmUniversal; - typename Kernel::QBMMTiling qbmmTiling{ - 1, 1, 1, 1, - 1, 1, 1, 1, - 1, 1, 1, 1, - 0, - static_cast(baseM), static_cast(baseN), static_cast(baseK), - biasElements != 0U ? 1U : 0U, 1U}; - - typename Kernel::Params params{ - {m, n, k, 1}, - {x1Gm, x2Gm, yGm, biasGm, x1ScaleGm, x2ScaleGm}, - {yGm, yScaleGm, static_cast(baseM), static_cast(baseN), - Blaze::Epilogue::Block::GeluAlg::TANH, - Blaze::Epilogue::Block::QuantAlg::OCP, - Blaze::Epilogue::Block::ROUND_MODE_FP4::RINT}, - {static_cast(tileShapeKL1), static_cast(tileShapeScaleKL1), - static_cast(l1BufferNum)}, - {static_cast(baseM), static_cast(baseN), 1, 1, 1, 1, 0, 0}, - qbmmTiling}; + typename Kernel::QBMMTiling qbmmTiling{1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 0, + static_cast(baseM), + static_cast(baseN), + static_cast(baseK), + biasElements != 0U ? 1U : 0U, + 1U}; + + typename Kernel::Params params{{m, n, k, 1}, + {x1Gm, x2Gm, yGm, biasGm, x1ScaleGm, x2ScaleGm}, + {yGm, yScaleGm, static_cast(baseM), static_cast(baseN), + Blaze::Epilogue::Block::GeluAlg::TANH, Blaze::Epilogue::Block::QuantAlg::OCP, + Blaze::Epilogue::Block::ROUND_MODE_FP4::RINT}, + {static_cast(tileShapeKL1), static_cast(tileShapeScaleKL1), + static_cast(l1BufferNum)}, + {static_cast(baseM), static_cast(baseN), 1, 1, 1, 1, 0, 0}, + qbmmTiling}; Kernel kernel; kernel(params); } @@ -240,17 +247,16 @@ template void RunCase(const ExampleConfig& config, aclrtStream stream) { const uint64_t scaleK = AlignUp(static_cast(config.k), 64U) / 32U; - const uint64_t bElements = WeightNz - ? AlignUp(static_cast(config.n), 32U) * AlignUp(static_cast(config.k), 16U) - : static_cast(config.k * config.n); + const uint64_t bElements = WeightNz ? AlignUp(static_cast(config.n), 32U) * + AlignUp(static_cast(config.k), 16U) : + static_cast(config.k * config.n); const size_t aSize = static_cast(config.m * config.k); const size_t bSize = static_cast(bElements); const size_t biasSize = static_cast(config.n) * sizeof(float); const size_t x1ScaleSize = static_cast(config.m) * scaleK; const size_t x2ScaleSize = static_cast(config.n) * scaleK; const size_t ySize = static_cast(config.m * config.n); - const size_t yScaleSize = static_cast(config.m) * - AlignUp(static_cast(config.n), 32U) / 32U; + const size_t yScaleSize = static_cast(config.m) * AlignUp(static_cast(config.n), 32U) / 32U; DeviceBuffer a(aSize); DeviceBuffer b(bSize); @@ -265,12 +271,10 @@ void RunCase(const ExampleConfig& config, aclrtStream stream) x1Scale.CopyFromFile(config.dataDir + "/scale_a.bin"); x2Scale.CopyFromFile(config.dataDir + "/scale_b.bin"); - QuantMatmulActivationQuantKernel - <<>>( - a.Get(), b.Get(), bias.Get(), x1Scale.Get(), x2Scale.Get(), - y.Get(), yScale.Get(), config.m, config.k, config.n, - config.baseM, config.baseN, config.baseK, config.tileShapeKL1, - config.tileShapeScaleKL1, config.l1BufferNum, config.biasElements); + QuantMatmulActivationQuantKernel<<>>( + a.Get(), b.Get(), bias.Get(), x1Scale.Get(), x2Scale.Get(), y.Get(), yScale.Get(), config.m, config.k, config.n, + config.baseM, config.baseN, config.baseK, config.tileShapeKL1, config.tileShapeScaleKL1, config.l1BufferNum, + config.biasElements); ACL_CHECK(aclrtSynchronizeStream(stream)); y.CopyToFile(config.dataDir + "/npu_y.bin"); yScale.CopyToFile(config.dataDir + "/npu_y_scale.bin"); diff --git a/examples/quant_matmul_activation_quant/scripts/gen_data.py b/examples/quant_matmul_activation_quant/scripts/gen_data.py index 4295fc3..ca2c176 100644 --- a/examples/quant_matmul_activation_quant/scripts/gen_data.py +++ b/examples/quant_matmul_activation_quant/scripts/gen_data.py @@ -53,9 +53,13 @@ def build_scale_broadcast(scale, target_shape, chunk_axis): scale_repeat = np.repeat(scale.astype(np.float32), MX_GROUP_SIZE, axis=-1) if chunk_axis == 1: - scale_broadcast = scale_repeat.reshape(scale.shape[0], -1)[..., : target_shape[1]] + scale_broadcast = scale_repeat.reshape(scale.shape[0], -1)[ + ..., : target_shape[1] + ] elif chunk_axis == 0: - scale_broadcast = np.transpose(scale_repeat, (0, 2, 1)).reshape(-1, scale.shape[1])[: target_shape[0], ...] + scale_broadcast = np.transpose(scale_repeat, (0, 2, 1)).reshape( + -1, scale.shape[1] + )[: target_shape[0], ...] else: raise ValueError(f"Invalid chunk_axis={chunk_axis}, expected 0 or 1.") @@ -68,7 +72,7 @@ def dequant_mxfp8(fp8_input, scale, chunk_axis): def gelu_tanh(x): - return 0.5 * x * (1.0 + np.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * x ** 3))) + return 0.5 * x * (1.0 + np.tanh(np.sqrt(2.0 / np.pi) * (x + 0.044715 * x**3))) def ocp_mx_quantize(data, emax): @@ -139,14 +143,17 @@ def generate(args): a_ori.view(np.uint8).tofile(os.path.join(args.output_dir, "input_a.bin")) if args.layout == "nz": to_weight_nz_layout(b_ori.view(np.uint8)).tofile( - os.path.join(args.output_dir, "input_b.bin")) + os.path.join(args.output_dir, "input_b.bin") + ) else: b_ori.view(np.uint8).tofile(os.path.join(args.output_dir, "input_b.bin")) a_scale.tofile(os.path.join(args.output_dir, "scale_a.bin")) b_scale.tofile(os.path.join(args.output_dir, "scale_b.bin")) bias.tofile(os.path.join(args.output_dir, "bias.bin")) golden_y.view(np.uint8).tofile(os.path.join(args.output_dir, "golden_y.bin")) - golden_y_scale.view(np.uint8).tofile(os.path.join(args.output_dir, "golden_y_scale.bin")) + golden_y_scale.view(np.uint8).tofile( + os.path.join(args.output_dir, "golden_y_scale.bin") + ) def main(): diff --git a/examples/quant_matmul_activation_quant/scripts/verify_result.py b/examples/quant_matmul_activation_quant/scripts/verify_result.py index 0ad6552..d47aa34 100644 --- a/examples/quant_matmul_activation_quant/scripts/verify_result.py +++ b/examples/quant_matmul_activation_quant/scripts/verify_result.py @@ -29,15 +29,18 @@ def compare(name, golden_path, actual_path, dtype, rtol, atol): if golden.shape != actual.shape: raise ValueError(f"{name}: size mismatch: {actual.size} != {golden.size}") close = np.isclose( - actual.astype(np.float32), golden.astype(np.float32), rtol=rtol, atol=atol) + actual.astype(np.float32), golden.astype(np.float32), rtol=rtol, atol=atol + ) if not np.all(close): index = int(np.flatnonzero(~close)[0]) error = abs(float(actual[index]) - float(golden[index])) raise ValueError( f"{name} mismatch at {index}: expected {golden[index]}, " - f"got {actual[index]}, abs_error={error}") - max_error = float(np.max(np.abs( - actual.astype(np.float32) - golden.astype(np.float32)))) + f"got {actual[index]}, abs_error={error}" + ) + max_error = float( + np.max(np.abs(actual.astype(np.float32) - golden.astype(np.float32))) + ) print(f"[PASS] {name}: {actual.size} elements, max_abs_error={max_error}") @@ -53,7 +56,14 @@ def main(): args = parser.parse_args() compare("y", args.golden_y, args.actual_y, FP8_E4M3FN, args.rtol, args.atol) - compare("y_scale", args.golden_y_scale, args.actual_y_scale, FP8_E8M0, args.rtol, args.scale_atol) + compare( + "y_scale", + args.golden_y_scale, + args.actual_y_scale, + FP8_E8M0, + args.rtol, + args.scale_atol, + ) if __name__ == "__main__": diff --git a/include/blaze/epilogue/block/block_epilogue_gelu_mx_quant.h b/include/blaze/epilogue/block/block_epilogue_gelu_mx_quant.h index ae29609..fa94797 100644 --- a/include/blaze/epilogue/block/block_epilogue_gelu_mx_quant.h +++ b/include/blaze/epilogue/block/block_epilogue_gelu_mx_quant.h @@ -54,18 +54,18 @@ constexpr float ONE_OVER_SQRT_TWO = 0.707106781f; constexpr uint32_t MAX_SINGLE_MN = 128 * 256; constexpr uint32_t MAX_SINGLE_SCALE_NUM = MAX_SINGLE_MN / AscendC::ONE_BLK_SIZE; -constexpr uint16_t MAX_EXP_FOR_BF16 = 0x7f80; // 0b0111 1111 1000 0000 +constexpr uint16_t MAX_EXP_FOR_BF16 = 0x7f80; // 0b0111 1111 1000 0000 constexpr uint16_t MAX_EXP_FOR_FP8 = 0x00ff; -constexpr uint16_t BF16_EXP_BIAS = 0x7f00; // 0b0111 1111 0000 0000 +constexpr uint16_t BF16_EXP_BIAS = 0x7f00; // 0b0111 1111 0000 0000 constexpr int16_t SHR_NUM_FOR_BF16 = 7; constexpr int16_t SHR_NUM_FOR_FP32 = 23; constexpr uint16_t NAN_CUSTOMIZATION = 0x7f81; constexpr uint16_t SPECIAL_EXP_THRESHOLD = 0x0040; // elem_emax右移7位(BF16E8M7) -constexpr uint16_t FP8_E4M3_MAX_EXP = 0x0400; // 0b 0000 0100 0000 0000 右移7位为8 -constexpr uint16_t FP8_E5M2_MAX_EXP = 0x0780; // 0b 0000 0111 1000 0000 右移7位为15 -constexpr uint16_t FP4_E2M1_MAX_EXP = 0x0100; // 0b 0000 0001 0000 0000 右移7位为2 -constexpr uint16_t FP4_E1M2_MAX_EXP = 0x0000; // 右移7位为0 +constexpr uint16_t FP8_E4M3_MAX_EXP = 0x0400; // 0b 0000 0100 0000 0000 右移7位为8 +constexpr uint16_t FP8_E5M2_MAX_EXP = 0x0780; // 0b 0000 0111 1000 0000 右移7位为15 +constexpr uint16_t FP4_E2M1_MAX_EXP = 0x0100; // 0b 0000 0001 0000 0000 右移7位为2 +constexpr uint16_t FP4_E1M2_MAX_EXP = 0x0000; // 右移7位为0 constexpr uint16_t ABS_MASK_FOR_16BIT = 0x7fff; constexpr uint32_t FP8_E5M2_MAX = 0x37924925; // 1/57344的float32表示 57334是E5M2所能表示的最大值 @@ -83,9 +83,7 @@ constexpr int8_t FLOAT_OVERFLOW_MODE_CTRL = 60; template class BlockEpilogueGeluMxQuant { public: - __aicore__ inline BlockEpilogueGeluMxQuant() - { - } + __aicore__ inline BlockEpilogueGeluMxQuant() {} struct Params { GM_ADDR yGmAddr{nullptr}; @@ -108,43 +106,42 @@ public: using ProblemShape = AscendC::Te::Shape; public: - __aicore__ inline void Init(Params const ¶ms); + __aicore__ inline void Init(Params const& params); __aicore__ inline auto GetTensor(); - __aicore__ inline void operator()(const BlockShape &blockShape, const BlockCoord &blockCoord); - __aicore__ inline void UpdateGlobalAddr(const BlockCoord &baseOffset); - __aicore__ inline void UpdateNextProblem(const ProblemShape &problemShape); + __aicore__ inline void operator()(const BlockShape& blockShape, const BlockCoord& blockCoord); + __aicore__ inline void UpdateGlobalAddr(const BlockCoord& baseOffset); + __aicore__ inline void UpdateNextProblem(const ProblemShape& problemShape); private: __aicore__ inline void VFDoGeluForMX(uint16_t mSize); __aicore__ inline void TransMxScaleLayout(uint16_t mSize); - __aicore__ inline void VFDoGeluAndQuantForMX(__ubuf__ int8_t *outputDst, __ubuf__ uint16_t *scaleDst, - uint16_t mSize, uint16_t nSize); - __aicore__ inline void GeluTanh(__ubuf__ bfloat16_t *geluResAddr, uint16_t mSize, uint16_t nSize, - uint32_t nAligned); - __aicore__ inline void GeluErf(__ubuf__ bfloat16_t *geluResAddr, uint16_t mSize, uint16_t nSize, + __aicore__ inline void VFDoGeluAndQuantForMX(__ubuf__ int8_t* outputDst, __ubuf__ uint16_t* scaleDst, + uint16_t mSize, uint16_t nSize); + __aicore__ inline void GeluTanh(__ubuf__ bfloat16_t* geluResAddr, uint16_t mSize, uint16_t nSize, uint32_t nAligned); - __aicore__ inline void ComputeScaleOCP(__ubuf__ uint16_t *maxExpAddr, __ubuf__ uint16_t *mxScaleLocalAddr, - __ubuf__ uint16_t *halfScaleLocalAddr, uint32_t totalScaleInUB, - uint16_t loopNumScale); - __aicore__ inline void ComputeScalecuBLAS(__ubuf__ uint16_t *maxExpAddr, __ubuf__ uint16_t *mxScaleLocalAddr, - __ubuf__ uint16_t *halfScaleLocalAddr, uint32_t totalScaleInUB, - uint16_t loopNumScale); - __aicore__ inline void ComputeMaxExpOCP(__ubuf__ bfloat16_t *srcAddr, __ubuf__ uint16_t *maxExpAddr, - uint16_t loopNum); - __aicore__ inline void ComputeMaxExpcuBLAS(__ubuf__ bfloat16_t *srcAddr, __ubuf__ uint16_t *maxExpAddr, - uint16_t loopNum); - __aicore__ inline void ComputeDataForQuantTargetFp8(__ubuf__ bfloat16_t *srcAddr, - __ubuf__ uint16_t *halfScaleLocalAddr, - __ubuf__ int8_t *outLocalAddr, uint32_t totalCountInUB, + __aicore__ inline void GeluErf(__ubuf__ bfloat16_t* geluResAddr, uint16_t mSize, uint16_t nSize, uint32_t nAligned); + __aicore__ inline void ComputeScaleOCP(__ubuf__ uint16_t* maxExpAddr, __ubuf__ uint16_t* mxScaleLocalAddr, + __ubuf__ uint16_t* halfScaleLocalAddr, uint32_t totalScaleInUB, + uint16_t loopNumScale); + __aicore__ inline void ComputeScalecuBLAS(__ubuf__ uint16_t* maxExpAddr, __ubuf__ uint16_t* mxScaleLocalAddr, + __ubuf__ uint16_t* halfScaleLocalAddr, uint32_t totalScaleInUB, + uint16_t loopNumScale); + __aicore__ inline void ComputeMaxExpOCP(__ubuf__ bfloat16_t* srcAddr, __ubuf__ uint16_t* maxExpAddr, + uint16_t loopNum); + __aicore__ inline void ComputeMaxExpcuBLAS(__ubuf__ bfloat16_t* srcAddr, __ubuf__ uint16_t* maxExpAddr, + uint16_t loopNum); + __aicore__ inline void ComputeDataForQuantTargetFp8(__ubuf__ bfloat16_t* srcAddr, + __ubuf__ uint16_t* halfScaleLocalAddr, + __ubuf__ int8_t* outLocalAddr, uint32_t totalCountInUB, uint16_t loopNum); template - __aicore__ inline void ComputeDataForQuantTargetFp4(__ubuf__ bfloat16_t *srcAddr, - __ubuf__ uint16_t *halfScaleLocalAddr, - __ubuf__ int8_t *outLocalAddr, uint32_t totalCountInUB, + __aicore__ inline void ComputeDataForQuantTargetFp4(__ubuf__ bfloat16_t* srcAddr, + __ubuf__ uint16_t* halfScaleLocalAddr, + __ubuf__ int8_t* outLocalAddr, uint32_t totalCountInUB, uint16_t loopNum); - __aicore__ inline void CopyOutputFromUb2Gm(uint64_t blockCount, uint64_t offset, AscendC::LocalTensor &src); - __aicore__ inline void CopyScaleFromUb2Gm(uint64_t blockCount, uint64_t offset, AscendC::LocalTensor &src); + __aicore__ inline void CopyOutputFromUb2Gm(uint64_t blockCount, uint64_t offset, AscendC::LocalTensor& src); + __aicore__ inline void CopyScaleFromUb2Gm(uint64_t blockCount, uint64_t offset, AscendC::LocalTensor& src); // GM ADDR AscendC::GlobalTensor quantOutputGlobal_; @@ -163,7 +160,7 @@ private: AscendC::LocalTensor fp32Tmp_; AscendC::LocalTensor geluFp32Tmp_; - const Params *params_; + const Params* params_; int64_t n_; int64_t scaleN_; @@ -184,7 +181,7 @@ private: }; template -__aicore__ inline void BlockEpilogueGeluMxQuant::Init(Params const ¶ms) +__aicore__ inline void BlockEpilogueGeluMxQuant::Init(Params const& params) { if ASCEND_IS_AIC { return; @@ -201,10 +198,10 @@ __aicore__ inline void BlockEpilogueGeluMxQuant::Init dtypeMax_ = FP8_E5M2_MAX; } else if constexpr (AscendC::IsSameType::value) { fpEmax_ = FP4_E2M1_MAX_EXP; - dtypeMax_ = 0; // FP4不支持 + dtypeMax_ = 0; // FP4不支持 } else { fpEmax_ = FP4_E1M2_MAX_EXP; - dtypeMax_ = 0; // FP4不支持 + dtypeMax_ = 0; // FP4不支持 } // out @@ -221,43 +218,43 @@ __aicore__ inline void BlockEpilogueGeluMxQuant::Init constexpr uint32_t afterIOAndGeluExp = afterIOAndGelu + MAX_SINGLE_SCALE_NUM * sizeof(uint16_t); halfScale_ = AscendC::LocalTensor(AscendC::TPosition::VECCALC, afterIOAndGeluExp, MAX_SINGLE_SCALE_NUM); constexpr uint32_t realScaleBlockOffset = afterIOAndGeluExp + MAX_SINGLE_SCALE_NUM * sizeof(uint16_t); - quantScaleBlockOutput_ = AscendC::LocalTensor(AscendC::TPosition::VECOUT, realScaleBlockOffset, + quantScaleBlockOutput_ = AscendC::LocalTensor( + AscendC::TPosition::VECOUT, realScaleBlockOffset, params_->baseM / AscendC::GetTaskRation() * AscendC::ONE_BLK_SIZE); if (params_->geluAlg == GeluAlg::ERF) { - uint32_t ubOffset = realScaleBlockOffset - + params_->baseM / AscendC::GetTaskRation() * AscendC::ONE_BLK_SIZE * sizeof(int8_t); + uint32_t ubOffset = realScaleBlockOffset + + params_->baseM / AscendC::GetTaskRation() * AscendC::ONE_BLK_SIZE * sizeof(int8_t); if constexpr (AscendC::IsSameType::value) { erfTmp_ = AscendC::LocalTensor(AscendC::TPosition::VECCALC, ubOffset, params_->baseN); - geluFp32Tmp_ = AscendC::LocalTensor(AscendC::TPosition::VECCALC, - ubOffset + params_->baseN, params_->baseN); + geluFp32Tmp_ = AscendC::LocalTensor(AscendC::TPosition::VECCALC, ubOffset + params_->baseN, + params_->baseN); } else { fp32Tmp_ = AscendC::LocalTensor(AscendC::TPosition::VECCALC, ubOffset, params_->baseN); - erfTmp_ = AscendC::LocalTensor(AscendC::TPosition::VECCALC, - ubOffset + params_->baseN, params_->baseN); + erfTmp_ = AscendC::LocalTensor(AscendC::TPosition::VECCALC, ubOffset + params_->baseN, + params_->baseN); geluFp32Tmp_ = AscendC::LocalTensor(AscendC::TPosition::VECCALC, - ubOffset + params_->baseN + params_->baseN, params_->baseN); + ubOffset + params_->baseN + params_->baseN, params_->baseN); } } - - quantOutputGlobal_.SetGlobalBuffer((__gm__ int8_t *)params_->yGmAddr); - quantScaleGlobal_.SetGlobalBuffer((__gm__ int8_t *)params_->yScaleGmAddr); + + quantOutputGlobal_.SetGlobalBuffer((__gm__ int8_t*)params_->yGmAddr); + quantScaleGlobal_.SetGlobalBuffer((__gm__ int8_t*)params_->yScaleGmAddr); } template __aicore__ inline void BlockEpilogueGeluMxQuant::UpdateGlobalAddr( - const BlockCoord &baseOffset) + const BlockCoord& baseOffset) { if ASCEND_IS_AIV { - quantOutputGlobal_.SetGlobalBuffer( - (__gm__ int8_t *)params_->yGmAddr + AscendC::Te::Get(baseOffset)); - quantScaleGlobal_.SetGlobalBuffer( - (__gm__ int8_t *)params_->yScaleGmAddr + AscendC::Te::Get(baseOffset)); + quantOutputGlobal_.SetGlobalBuffer((__gm__ int8_t*)params_->yGmAddr + AscendC::Te::Get(baseOffset)); + quantScaleGlobal_.SetGlobalBuffer((__gm__ int8_t*)params_->yScaleGmAddr + + AscendC::Te::Get(baseOffset)); } } template __aicore__ inline void BlockEpilogueGeluMxQuant::UpdateNextProblem( - const ProblemShape &problemShape) + const ProblemShape& problemShape) { n_ = AscendC::Te::Get(problemShape); scaleN_ = Gemm::CeilDiv(static_cast(n_), static_cast(BLOCK_SIZE)); @@ -266,7 +263,7 @@ __aicore__ inline void BlockEpilogueGeluMxQuant::Upda template __aicore__ inline void BlockEpilogueGeluMxQuant::CopyOutputFromUb2Gm( - uint64_t blockCount, uint64_t offset, AscendC::LocalTensor &src) + uint64_t blockCount, uint64_t offset, AscendC::LocalTensor& src) { AscendC::DataCopyExtParams ub2GmParams{1, 0, 0, 0, 0}; ub2GmParams.blockCount = blockCount; @@ -283,7 +280,7 @@ __aicore__ inline void BlockEpilogueGeluMxQuant::Copy template __aicore__ inline void BlockEpilogueGeluMxQuant::CopyScaleFromUb2Gm( - uint64_t blockCount, uint64_t offset, AscendC::LocalTensor &src) + uint64_t blockCount, uint64_t offset, AscendC::LocalTensor& src) { AscendC::DataCopyExtParams ub2GmParams{1, 0, 0, 0, 0}; ub2GmParams.blockLen = scaleBlockN_ * sizeof(int8_t); @@ -294,7 +291,7 @@ __aicore__ inline void BlockEpilogueGeluMxQuant::Copy template __aicore__ inline void BlockEpilogueGeluMxQuant::ComputeMaxExpOCP( - __ubuf__ bfloat16_t *srcAddr, __ubuf__ uint16_t *maxExpAddr, uint16_t loopNum) + __ubuf__ bfloat16_t* srcAddr, __ubuf__ uint16_t* maxExpAddr, uint16_t loopNum) { __VEC_SCOPE__ { @@ -307,14 +304,15 @@ __aicore__ inline void BlockEpilogueGeluMxQuant::Comp AscendC::MicroAPI::Duplicate(expMaskBF16, MAX_EXP_FOR_BF16); AscendC::MicroAPI::RegTensor vdMaxExp; - AscendC::MicroAPI::MaskReg Mask = - AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg + Mask = AscendC::MicroAPI::CreateMask(); AscendC::MicroAPI::UnalignReg u1; for (uint16_t i = 0; i < loopNum; i++) { AscendC::MicroAPI::DataCopy(vdExp0, vdExp1, srcAddr, vlForHalfNumber_ * 2); - AscendC::MicroAPI::And(vdExpExtract0, (AscendC::MicroAPI::RegTensor &)vdExp0, expMaskBF16, Mask); - AscendC::MicroAPI::And(vdExpExtract1, (AscendC::MicroAPI::RegTensor &)vdExp1, expMaskBF16, Mask); + AscendC::MicroAPI::LoadDist::DIST_DINTLV_B16>(vdExp0, vdExp1, srcAddr, + vlForHalfNumber_ * 2); + AscendC::MicroAPI::And(vdExpExtract0, (AscendC::MicroAPI::RegTensor&)vdExp0, expMaskBF16, Mask); + AscendC::MicroAPI::And(vdExpExtract1, (AscendC::MicroAPI::RegTensor&)vdExp1, expMaskBF16, Mask); AscendC::MicroAPI::Max(vdMaxExp, vdExpExtract0, vdExpExtract1, Mask); AscendC::MicroAPI::ReduceMaxWithDataBlock(vdMaxExp, vdMaxExp, Mask); AscendC::MicroAPI::DataCopyUnAlign( @@ -327,7 +325,7 @@ __aicore__ inline void BlockEpilogueGeluMxQuant::Comp template __aicore__ inline void BlockEpilogueGeluMxQuant::ComputeMaxExpcuBLAS( - __ubuf__ bfloat16_t *srcAddr, __ubuf__ uint16_t *maxExpAddr, uint16_t loopNum) + __ubuf__ bfloat16_t* srcAddr, __ubuf__ uint16_t* maxExpAddr, uint16_t loopNum) { __VEC_SCOPE__ { @@ -338,21 +336,19 @@ __aicore__ inline void BlockEpilogueGeluMxQuant::Comp AscendC::MicroAPI::RegTensor absMask16Bit; AscendC::MicroAPI::Duplicate(absMask16Bit, ABS_MASK_FOR_16BIT); - AscendC::MicroAPI::MaskReg Mask = - AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg + Mask = AscendC::MicroAPI::CreateMask(); AscendC::MicroAPI::UnalignReg u1; for (uint16_t i = 0; i < loopNum; i++) { AscendC::MicroAPI::DataCopy(vdExp0, vdExp1, srcAddr, vlForHalfNumber_ * 2); - AscendC::MicroAPI::And( - (AscendC::MicroAPI::RegTensor&)vdExp0, (AscendC::MicroAPI::RegTensor&)vdExp0, - absMask16Bit, Mask); - AscendC::MicroAPI::And( - (AscendC::MicroAPI::RegTensor&)vdExp1, (AscendC::MicroAPI::RegTensor&)vdExp1, - absMask16Bit, Mask); - AscendC::MicroAPI::Max( - vdMaxExp, (AscendC::MicroAPI::RegTensor&)vdExp0, - (AscendC::MicroAPI::RegTensor&)vdExp1, Mask); + AscendC::MicroAPI::LoadDist::DIST_DINTLV_B16>(vdExp0, vdExp1, srcAddr, + vlForHalfNumber_ * 2); + AscendC::MicroAPI::And((AscendC::MicroAPI::RegTensor&)vdExp0, + (AscendC::MicroAPI::RegTensor&)vdExp0, absMask16Bit, Mask); + AscendC::MicroAPI::And((AscendC::MicroAPI::RegTensor&)vdExp1, + (AscendC::MicroAPI::RegTensor&)vdExp1, absMask16Bit, Mask); + AscendC::MicroAPI::Max(vdMaxExp, (AscendC::MicroAPI::RegTensor&)vdExp0, + (AscendC::MicroAPI::RegTensor&)vdExp1, Mask); AscendC::MicroAPI::ReduceMaxWithDataBlock(vdMaxExp, vdMaxExp, Mask); AscendC::MicroAPI::DataCopyUnAlign( maxExpAddr, vdMaxExp, u1, elementAfterReduce_); @@ -364,7 +360,7 @@ __aicore__ inline void BlockEpilogueGeluMxQuant::Comp template __aicore__ inline void BlockEpilogueGeluMxQuant::ComputeScalecuBLAS( - __ubuf__ uint16_t *maxExpAddr, __ubuf__ uint16_t *mxScaleLocalAddr, __ubuf__ uint16_t *halfScaleLocalAddr, + __ubuf__ uint16_t* maxExpAddr, __ubuf__ uint16_t* mxScaleLocalAddr, __ubuf__ uint16_t* halfScaleLocalAddr, uint32_t totalScaleInUB, uint16_t loopNumScale) { using T = bfloat16_t; @@ -410,16 +406,17 @@ __aicore__ inline void BlockEpilogueGeluMxQuant::Comp AscendC::MicroAPI::MaskMergeMode::ZEROING, AscendC::RoundMode::UNKNOWN}; for (uint16_t i = 0; i < loopNumScale; i++) { AscendC::MicroAPI::LoadAlign(vdMaxExp, maxExpAddr, vlForFloat32Number_); + AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(vdMaxExp, maxExpAddr, + vlForFloat32Number_); AscendC::MicroAPI::Cast((AscendC::MicroAPI::RegTensor&)vdMaxExp32, - (AscendC::MicroAPI::RegTensor&)vdMaxExp, mask); + (AscendC::MicroAPI::RegTensor&)vdMaxExp, mask); AscendC::MicroAPI::Compare(cmpResult, vdMaxExp32, expMask, mask); AscendC::MicroAPI::Compare(zeroMask, vdMaxExp32, zeroRegTensor32, mask); - AscendC::MicroAPI::Mul( - (AscendC::MicroAPI::RegTensor&)vdMaxExp32, (AscendC::MicroAPI::RegTensor&)vdMaxExp32, - (AscendC::MicroAPI::RegTensor&)invMax, mask); + AscendC::MicroAPI::Mul((AscendC::MicroAPI::RegTensor&)vdMaxExp32, + (AscendC::MicroAPI::RegTensor&)vdMaxExp32, + (AscendC::MicroAPI::RegTensor&)invMax, mask); AscendC::MicroAPI::ShiftRights(exp32, vdMaxExp32, SHR_NUM_FOR_FP32, mask); AscendC::MicroAPI::And(man32, vdMaxExp32, manMaskFP32, mask); @@ -447,8 +444,8 @@ __aicore__ inline void BlockEpilogueGeluMxQuant::Comp AscendC::MicroAPI::Select(halfScale, halfScale, nanRegTensor, cmpResult); AscendC::MicroAPI::Select(halfScale, halfScale, zeroRegTensor32, zeroMask); AscendC::MicroAPI::Pack(recExpOut, halfScale); - AscendC::MicroAPI::StoreAlign( - halfScaleLocalAddr + i * vlForFloat32Number_, recExpOut, dataMaskB16Half); + AscendC::MicroAPI::StoreAlign(halfScaleLocalAddr + i * vlForFloat32Number_, recExpOut, + dataMaskB16Half); } } return; @@ -456,7 +453,7 @@ __aicore__ inline void BlockEpilogueGeluMxQuant::Comp template __aicore__ inline void BlockEpilogueGeluMxQuant::ComputeScaleOCP( - __ubuf__ uint16_t *maxExpAddr, __ubuf__ uint16_t *mxScaleLocalAddr, __ubuf__ uint16_t *halfScaleLocalAddr, + __ubuf__ uint16_t* maxExpAddr, __ubuf__ uint16_t* mxScaleLocalAddr, __ubuf__ uint16_t* halfScaleLocalAddr, uint32_t totalScaleInUB, uint16_t loopNumScale) { __VEC_SCOPE__ @@ -490,8 +487,8 @@ __aicore__ inline void BlockEpilogueGeluMxQuant::Comp AscendC::MicroAPI::Select(scaleValue, scaleValue, zeroRegTensor, zeroMask); AscendC::MicroAPI::DataCopy( - mxScaleLocalAddr, scaleValue, vlForHalfNumber_ >> 1, maskScale); + AscendC::MicroAPI::StoreDist::DIST_PACK_B16>(mxScaleLocalAddr, scaleValue, + vlForHalfNumber_ >> 1, maskScale); AscendC::MicroAPI::Compare(specialDataMask, sharedExp, scaleBias, maskScale); @@ -509,7 +506,7 @@ __aicore__ inline void BlockEpilogueGeluMxQuant::Comp template __aicore__ inline void BlockEpilogueGeluMxQuant::ComputeDataForQuantTargetFp8( - __ubuf__ bfloat16_t *srcAddr, __ubuf__ uint16_t *halfScaleLocalAddr, __ubuf__ int8_t *outLocalAddr, + __ubuf__ bfloat16_t* srcAddr, __ubuf__ uint16_t* halfScaleLocalAddr, __ubuf__ int8_t* outLocalAddr, uint32_t totalCountInUB, uint16_t loopNum) { uint32_t totalCountInUB2 = totalCountInUB * 2; @@ -517,8 +514,8 @@ __aicore__ inline void BlockEpilogueGeluMxQuant::Comp __VEC_SCOPE__ { AscendC::MicroAPI::MaskReg dataMask1, dataMask2, dataMask3, dataMask4; - AscendC::MicroAPI::MaskReg maskAll = - AscendC::MicroAPI::CreateMask(); + AscendC::MicroAPI::MaskReg + maskAll = AscendC::MicroAPI::CreateMask(); AscendC::MicroAPI::RegTensor halfScaleForMul; AscendC::MicroAPI::RegTensor floatScaleForMul; AscendC::MicroAPI::RegTensor vdExp0, vdExp1, vdExp0Convert, vdExp1Convert; @@ -548,8 +545,8 @@ __aicore__ inline void BlockEpilogueGeluMxQuant::Comp AscendC::MicroAPI::LoadDist::DIST_E2B_B16>(halfScaleForMul, halfScaleLocalAddr, elementAfterReduce_); - AscendC::MicroAPI::Mul(vdExp0, vdExp0, (AscendC::MicroAPI::RegTensor &)halfScaleForMul, dataMask1); - AscendC::MicroAPI::Mul(vdExp1, vdExp1, (AscendC::MicroAPI::RegTensor &)halfScaleForMul, dataMask1); + AscendC::MicroAPI::Mul(vdExp0, vdExp0, (AscendC::MicroAPI::RegTensor&)halfScaleForMul, dataMask1); + AscendC::MicroAPI::Mul(vdExp1, vdExp1, (AscendC::MicroAPI::RegTensor&)halfScaleForMul, dataMask1); AscendC::MicroAPI::Interleave(vdExp0, vdExp1, vdExp0, vdExp1); AscendC::MicroAPI::Cast(vdExp0FP32Zero, vdExp0, dataMask1); AscendC::MicroAPI::Cast(vdExp0FP32One, vdExp0, dataMask1); @@ -563,26 +560,25 @@ __aicore__ inline void BlockEpilogueGeluMxQuant::Comp AscendC::MicroAPI::Cast(vdExp1FP8One, vdExp1FP32One, dataMask4); AscendC::MicroAPI::DataCopy( - outLocalAddr, (AscendC::MicroAPI::RegTensor &)vdExp0FP8Zero, OUT_ELE_NUM_ONE_BLK, dataMask3); + outLocalAddr, (AscendC::MicroAPI::RegTensor&)vdExp0FP8Zero, OUT_ELE_NUM_ONE_BLK, dataMask3); AscendC::MicroAPI::DataCopy( - outLocalAddr, (AscendC::MicroAPI::RegTensor &)vdExp0FP8One, OUT_ELE_NUM_ONE_BLK, dataMask3); + outLocalAddr, (AscendC::MicroAPI::RegTensor&)vdExp0FP8One, OUT_ELE_NUM_ONE_BLK, dataMask3); AscendC::MicroAPI::DataCopy( - outLocalAddr, (AscendC::MicroAPI::RegTensor &)vdExp1FP8Zero, OUT_ELE_NUM_ONE_BLK, dataMask4); + outLocalAddr, (AscendC::MicroAPI::RegTensor&)vdExp1FP8Zero, OUT_ELE_NUM_ONE_BLK, dataMask4); AscendC::MicroAPI::DataCopy( - outLocalAddr, (AscendC::MicroAPI::RegTensor &)vdExp1FP8One, OUT_ELE_NUM_ONE_BLK, dataMask4); + outLocalAddr, (AscendC::MicroAPI::RegTensor&)vdExp1FP8One, OUT_ELE_NUM_ONE_BLK, dataMask4); } } return; } - template template __aicore__ inline void BlockEpilogueGeluMxQuant::ComputeDataForQuantTargetFp4( - __ubuf__ bfloat16_t *srcAddr, __ubuf__ uint16_t *halfScaleLocalAddr, __ubuf__ int8_t *outLocalAddr, + __ubuf__ bfloat16_t* srcAddr, __ubuf__ uint16_t* halfScaleLocalAddr, __ubuf__ int8_t* outLocalAddr, uint32_t totalCountInUB, uint16_t loopNum) { using T = bfloat16_t; @@ -610,44 +606,49 @@ __aicore__ inline void BlockEpilogueGeluMxQuant::Comp dataMask1 = AscendC::MicroAPI::UpdateMask(totalCountInUB); dataMask2 = AscendC::MicroAPI::UpdateMask(totalCountInUB); AscendC::MicroAPI::DataCopy(vdExp0, vdExp1, srcAddr, vlForHalfNumber_ * 2); + AscendC::MicroAPI::LoadDist::DIST_DINTLV_B16>(vdExp0, vdExp1, srcAddr, + vlForHalfNumber_ * 2); AscendC::MicroAPI::DataCopy(halfScaleForMul, halfScaleLocalAddr, elementAfterReduce_); + AscendC::MicroAPI::LoadDist::DIST_E2B_B16>(halfScaleForMul, halfScaleLocalAddr, + elementAfterReduce_); - AscendC::MicroAPI::Mul(vdExp0, vdExp0, (AscendC::MicroAPI::RegTensor &)halfScaleForMul, dataMask1); - AscendC::MicroAPI::Mul(vdExp1, vdExp1, (AscendC::MicroAPI::RegTensor &)halfScaleForMul, dataMask1); + AscendC::MicroAPI::Mul(vdExp0, vdExp0, (AscendC::MicroAPI::RegTensor&)halfScaleForMul, dataMask1); + AscendC::MicroAPI::Mul(vdExp1, vdExp1, (AscendC::MicroAPI::RegTensor&)halfScaleForMul, dataMask1); AscendC::MicroAPI::Interleave(vdExp0, vdExp1, vdExp0, vdExp1); AscendC::MicroAPI::Cast(vdExp0FP4, vdExp0, dataMask1); AscendC::MicroAPI::Cast(vdExp1FP4, vdExp1, dataMask2); AscendC::MicroAPI::DataCopy( - outLocalAddr, (AscendC::MicroAPI::RegTensor &)vdExp0FP4, OUT_ELE_NUM_ONE_BLK, dataMask1); + outLocalAddr, (AscendC::MicroAPI::RegTensor&)vdExp0FP4, OUT_ELE_NUM_ONE_BLK, dataMask1); AscendC::MicroAPI::DataCopy( - outLocalAddr, (AscendC::MicroAPI::RegTensor &)vdExp1FP4, OUT_ELE_NUM_ONE_BLK, dataMask2); + outLocalAddr, (AscendC::MicroAPI::RegTensor&)vdExp1FP4, OUT_ELE_NUM_ONE_BLK, dataMask2); } } return; } template -__aicore__ inline void BlockEpilogueGeluMxQuant::GeluTanh(__ubuf__ bfloat16_t *geluResAddr, - uint16_t mSize, uint16_t nSize, uint32_t nAligned) +__aicore__ inline void BlockEpilogueGeluMxQuant::GeluTanh(__ubuf__ bfloat16_t* geluResAddr, + uint16_t mSize, uint16_t nSize, + uint32_t nAligned) { - constexpr uint16_t sizePerRepeat = AscendC::VECTOR_REG_WIDTH / sizeof(float); // 需要转换成float32计算 - uint16_t OneRowRepeatTimes = Gemm::CeilDiv(nSize, sizePerRepeat); // 计算为64位对齐 + constexpr uint16_t sizePerRepeat = AscendC::VECTOR_REG_WIDTH / sizeof(float); // 需要转换成float32计算 + uint16_t OneRowRepeatTimes = Gemm::CeilDiv(nSize, sizePerRepeat); // 计算为64位对齐 - __ubuf__ DataTypeIn *src = (__ubuf__ DataTypeIn *)cLocal_.GetPhyAddr(); + __ubuf__ DataTypeIn* src = (__ubuf__ DataTypeIn*)cLocal_.GetPhyAddr(); AscendC::MicroAPI::RegTensor vregInput; AscendC::MicroAPI::RegTensor vregInputSqr; AscendC::MicroAPI::RegTensor vregInputCub; AscendC::MicroAPI::RegTensor vregOutput; - AscendC::MicroAPI::RegTensor vregOutput16; // gelu总是输出bfloat16 - static constexpr AscendC::MicroAPI::CastTrait ctHalf2Fp32Zero = {AscendC::MicroAPI::RegLayout::ZERO, - AscendC::MicroAPI::SatMode::UNKNOWN, AscendC::MicroAPI::MaskMergeMode::ZEROING, AscendC::RoundMode::UNKNOWN}; - static constexpr AscendC::MicroAPI::CastTrait ctFp32toBf16 = {AscendC::MicroAPI::RegLayout::ZERO, - AscendC::MicroAPI::SatMode::NO_SAT, AscendC::MicroAPI::MaskMergeMode::ZEROING, AscendC::RoundMode::CAST_RINT}; + AscendC::MicroAPI::RegTensor vregOutput16; // gelu总是输出bfloat16 + static constexpr AscendC::MicroAPI::CastTrait ctHalf2Fp32Zero = { + AscendC::MicroAPI::RegLayout::ZERO, AscendC::MicroAPI::SatMode::UNKNOWN, + AscendC::MicroAPI::MaskMergeMode::ZEROING, AscendC::RoundMode::UNKNOWN}; + static constexpr AscendC::MicroAPI::CastTrait ctFp32toBf16 = { + AscendC::MicroAPI::RegLayout::ZERO, AscendC::MicroAPI::SatMode::NO_SAT, + AscendC::MicroAPI::MaskMergeMode::ZEROING, AscendC::RoundMode::CAST_RINT}; AscendC::MicroAPI::MaskReg mask; if constexpr (AscendC::IsSameType::value) { __VEC_SCOPE__ @@ -682,7 +683,7 @@ __aicore__ inline void BlockEpilogueGeluMxQuant::Gelu mask = AscendC::MicroAPI::UpdateMask(count); uint32_t offset = mIdx * nAligned + vfBlockIdx * sizePerRepeat; AscendC::MicroAPI::DataCopy(vregInput16, - src + offset); + src + offset); AscendC::MicroAPI::Cast(vregInput, vregInput16, mask); AscendC::MicroAPI::Mul(vregInputSqr, vregInput, vregInput, mask); AscendC::MicroAPI::Mul(vregInputCub, vregInputSqr, vregInput, mask); @@ -701,8 +702,9 @@ __aicore__ inline void BlockEpilogueGeluMxQuant::Gelu } template -__aicore__ inline void BlockEpilogueGeluMxQuant::GeluErf(__ubuf__ bfloat16_t *geluResAddr, - uint16_t mSize, uint16_t nSize, uint32_t nAligned) +__aicore__ inline void BlockEpilogueGeluMxQuant::GeluErf(__ubuf__ bfloat16_t* geluResAddr, + uint16_t mSize, uint16_t nSize, + uint32_t nAligned) { // 0.5*x*(1+erf(x/√2) constexpr uint16_t sizePerRepeat = AscendC::VECTOR_REG_WIDTH / sizeof(float); @@ -713,17 +715,18 @@ __aicore__ inline void BlockEpilogueGeluMxQuant::Gelu AscendC::MicroAPI::RegTensor vregInputAdds; AscendC::MicroAPI::RegTensor vregInputMuls; AscendC::MicroAPI::RegTensor vregOutput; - AscendC::MicroAPI::RegTensor vregOutput16; // gelu总是输出bfloat16 + AscendC::MicroAPI::RegTensor vregOutput16; // gelu总是输出bfloat16 AscendC::MicroAPI::MaskReg mask; - static constexpr AscendC::MicroAPI::CastTrait ctFp32toBf16 = {AscendC::MicroAPI::RegLayout::ZERO, - AscendC::MicroAPI::SatMode::NO_SAT, AscendC::MicroAPI::MaskMergeMode::ZEROING, AscendC::RoundMode::CAST_RINT}; - static constexpr AscendC::ErfConfig erfConfig = { AscendC::ErfAlgo::SUBSECTION_POLYNOMIAL_APPROXIMATION }; + static constexpr AscendC::MicroAPI::CastTrait ctFp32toBf16 = { + AscendC::MicroAPI::RegLayout::ZERO, AscendC::MicroAPI::SatMode::NO_SAT, + AscendC::MicroAPI::MaskMergeMode::ZEROING, AscendC::RoundMode::CAST_RINT}; + static constexpr AscendC::ErfConfig erfConfig = {AscendC::ErfAlgo::SUBSECTION_POLYNOMIAL_APPROXIMATION}; if constexpr (AscendC::IsSameType::value) { for (uint32_t mIdx = 0; mIdx < mSize; mIdx++) { AscendC::Muls(geluFp32Tmp_, cLocal_[mIdx * nAligned], ONE_OVER_SQRT_TWO, nSize); AscendC::Erf(erfTmp_, geluFp32Tmp_, nSize); - __ubuf__ float *src = (__ubuf__ float *)cLocal_.GetPhyAddr(); + __ubuf__ float* src = (__ubuf__ float*)cLocal_.GetPhyAddr(); __ubuf__ float* erfAddr = (__ubuf__ float*)erfTmp_.GetPhyAddr(); uint32_t count = nSize; __VEC_SCOPE__ @@ -738,7 +741,7 @@ __aicore__ inline void BlockEpilogueGeluMxQuant::Gelu AscendC::MicroAPI::Mul(vregOutput, vregInputAdds, vregInputMuls, mask); AscendC::MicroAPI::Cast(vregOutput16, vregOutput, mask); AscendC::MicroAPI::DataCopy( - geluResAddr + mnOffset, vregOutput16, mask); + geluResAddr + mnOffset, vregOutput16, mask); } } } @@ -747,7 +750,7 @@ __aicore__ inline void BlockEpilogueGeluMxQuant::Gelu AscendC::Cast(fp32Tmp_, cLocal_[mIdx * nAligned], AscendC::RoundMode::CAST_NONE, nSize); AscendC::Muls(geluFp32Tmp_, fp32Tmp_, ONE_OVER_SQRT_TWO, nSize); AscendC::Erf(erfTmp_, geluFp32Tmp_, nSize); - __ubuf__ float *src = (__ubuf__ float *)fp32Tmp_.GetPhyAddr(); + __ubuf__ float* src = (__ubuf__ float*)fp32Tmp_.GetPhyAddr(); __ubuf__ float* erfAddr = (__ubuf__ float*)erfTmp_.GetPhyAddr(); uint32_t count = nSize; __VEC_SCOPE__ @@ -763,7 +766,7 @@ __aicore__ inline void BlockEpilogueGeluMxQuant::Gelu AscendC::MicroAPI::Mul(vregOutput, vregInputAdds, vregInputMuls, mask); AscendC::MicroAPI::Cast(vregOutput16, vregOutput, mask); AscendC::MicroAPI::DataCopy( - geluResAddr + mnOffset, vregOutput16, mask); + geluResAddr + mnOffset, vregOutput16, mask); } } } @@ -772,12 +775,12 @@ __aicore__ inline void BlockEpilogueGeluMxQuant::Gelu template __aicore__ inline void BlockEpilogueGeluMxQuant::VFDoGeluAndQuantForMX( - __ubuf__ int8_t *outputDst, __ubuf__ uint16_t *scaleDst, uint16_t mSize, uint16_t nSize) + __ubuf__ int8_t* outputDst, __ubuf__ uint16_t* scaleDst, uint16_t mSize, uint16_t nSize) { - uint32_t nAligned = Gemm::Align32(static_cast(nSize)); // 输入为32位对齐 - AscendC::Duplicate(geluRes_, static_cast(0.0), mSize * nAligned); // 清除残余数据 + uint32_t nAligned = Gemm::Align32(static_cast(nSize)); // 输入为32位对齐 + AscendC::Duplicate(geluRes_, static_cast(0.0), mSize * nAligned); // 清除残余数据 - __ubuf__ bfloat16_t *geluResAddr = (__ubuf__ bfloat16_t *)geluRes_.GetPhyAddr(); + __ubuf__ bfloat16_t* geluResAddr = (__ubuf__ bfloat16_t*)geluRes_.GetPhyAddr(); if (params_->geluAlg == GeluAlg::ERF) { GeluErf(geluResAddr, mSize, nSize, nAligned); } else { @@ -787,18 +790,18 @@ __aicore__ inline void BlockEpilogueGeluMxQuant::VFDo uint32_t totalDataInUb = mSize * nAligned; uint32_t totalScaleInUb = totalDataInUb / BLOCK_SIZE; uint16_t loopDataNum = (totalDataInUb + vlForHalfNumber_ * 2 - 1) / (vlForHalfNumber_ * 2); - __ubuf__ uint16_t *halfScaleLocalAddr; + __ubuf__ uint16_t* halfScaleLocalAddr; if (params_->quantAlg == QuantAlg::OCP) { uint16_t loopScaleNum = (totalScaleInUb + vlForHalfNumber_ - 1) / vlForHalfNumber_; - __ubuf__ uint16_t *maxExpAddr = (__ubuf__ uint16_t *)maxExp_.GetPhyAddr(); + __ubuf__ uint16_t* maxExpAddr = (__ubuf__ uint16_t*)maxExp_.GetPhyAddr(); ComputeMaxExpOCP(geluResAddr, maxExpAddr, loopDataNum); - halfScaleLocalAddr = (__ubuf__ uint16_t *)halfScale_.GetPhyAddr(); + halfScaleLocalAddr = (__ubuf__ uint16_t*)halfScale_.GetPhyAddr(); ComputeScaleOCP(maxExpAddr, scaleDst, halfScaleLocalAddr, totalScaleInUb, loopScaleNum); } else { uint16_t loopScaleNum = (totalScaleInUb + vlForFloat32Number_ - 1) / vlForFloat32Number_; - __ubuf__ uint16_t *maxExpAddr = (__ubuf__ uint16_t *)maxExp_.GetPhyAddr(); + __ubuf__ uint16_t* maxExpAddr = (__ubuf__ uint16_t*)maxExp_.GetPhyAddr(); ComputeMaxExpcuBLAS(geluResAddr, maxExpAddr, loopDataNum); - halfScaleLocalAddr = (__ubuf__ uint16_t *)halfScale_.GetPhyAddr(); + halfScaleLocalAddr = (__ubuf__ uint16_t*)halfScale_.GetPhyAddr(); ComputeScalecuBLAS(maxExpAddr, scaleDst, halfScaleLocalAddr, totalScaleInUb, loopScaleNum); } @@ -809,14 +812,14 @@ __aicore__ inline void BlockEpilogueGeluMxQuant::VFDo if constexpr (AscendC::IsSameType::value || AscendC::IsSameType::value) { if (params_->fp4RoundMode == ROUND_MODE_FP4::FLOOR) { - ComputeDataForQuantTargetFp4(geluResAddr, halfScaleLocalAddr, - outputDst, totalDataInUb, loopDataNum); + ComputeDataForQuantTargetFp4(geluResAddr, halfScaleLocalAddr, outputDst, + totalDataInUb, loopDataNum); } else if ((params_->fp4RoundMode == ROUND_MODE_FP4::ROUND)) { - ComputeDataForQuantTargetFp4(geluResAddr, halfScaleLocalAddr, - outputDst, totalDataInUb, loopDataNum); - } else { // 默认rint - ComputeDataForQuantTargetFp4(geluResAddr, halfScaleLocalAddr, - outputDst, totalDataInUb, loopDataNum); + ComputeDataForQuantTargetFp4(geluResAddr, halfScaleLocalAddr, outputDst, + totalDataInUb, loopDataNum); + } else { // 默认rint + ComputeDataForQuantTargetFp4(geluResAddr, halfScaleLocalAddr, outputDst, + totalDataInUb, loopDataNum); } } return; @@ -825,16 +828,16 @@ __aicore__ inline void BlockEpilogueGeluMxQuant::VFDo template __aicore__ inline void BlockEpilogueGeluMxQuant::VFDoGeluForMX(uint16_t mSize) { - __ubuf__ int8_t *quantOutputInUbAddr = (__ubuf__ int8_t *)quantOutput_.GetPhyAddr(); - __ubuf__ uint16_t *quantScaleOutputInUbAddr = (__ubuf__ uint16_t *)quantScaleOutput_.GetPhyAddr(); + __ubuf__ int8_t* quantOutputInUbAddr = (__ubuf__ int8_t*)quantOutput_.GetPhyAddr(); + __ubuf__ uint16_t* quantScaleOutputInUbAddr = (__ubuf__ uint16_t*)quantScaleOutput_.GetPhyAddr(); VFDoGeluAndQuantForMX(quantOutputInUbAddr, quantScaleOutputInUbAddr, mSize, singleN_); } template __aicore__ inline void BlockEpilogueGeluMxQuant::TransMxScaleLayout(uint16_t mSize) { - __ubuf__ int8_t *quantScaleOutputInUbAddr = (__ubuf__ int8_t *)quantScaleOutput_.GetPhyAddr(); - __ubuf__ int8_t *quantScaleBlockOutputInUbAddr = (__ubuf__ int8_t *)quantScaleBlockOutput_.GetPhyAddr(); + __ubuf__ int8_t* quantScaleOutputInUbAddr = (__ubuf__ int8_t*)quantScaleOutput_.GetPhyAddr(); + __ubuf__ int8_t* quantScaleBlockOutputInUbAddr = (__ubuf__ int8_t*)quantScaleBlockOutput_.GetPhyAddr(); // scale layout: (mSize*8) -> (mSize,32) __VEC_SCOPE__ { @@ -859,8 +862,8 @@ __aicore__ inline auto BlockEpilogueGeluMxQuant::GetT } template -__aicore__ inline void BlockEpilogueGeluMxQuant::operator()(const BlockShape &blockShape, - const BlockCoord &blockCoord) +__aicore__ inline void BlockEpilogueGeluMxQuant::operator()(const BlockShape& blockShape, + const BlockCoord& blockCoord) { singleM_ = AscendC::Te::Get(blockShape); singleN_ = AscendC::Te::Get(blockShape); diff --git a/include/blaze/gemm/block/block_scheduler_qbmm.h b/include/blaze/gemm/block/block_scheduler_qbmm.h index ec8a0cb..29bbbed 100644 --- a/include/blaze/gemm/block/block_scheduler_qbmm.h +++ b/include/blaze/gemm/block/block_scheduler_qbmm.h @@ -100,18 +100,11 @@ public: endBlockIdx_ = newEndBlockIdx; } - __aicore__ inline int64_t GetTotalCnt() - { - return blockNums_; - } + __aicore__ inline int64_t GetTotalCnt() { return blockNums_; } - __aicore__ inline int64_t GetEndBlockIdx() - { - return endBlockIdx_; - } + __aicore__ inline int64_t GetEndBlockIdx() { return endBlockIdx_; } - template + template __aicore__ inline BlockShape GetBlockShape(BlockCoord blockCoord) { int64_t singleCoreM = baseM_; @@ -130,8 +123,7 @@ public: if constexpr (IsFp4() && !TRANS_B) { singleCoreNSplit = (singleCoreNSplit + 1) & ~1; } - if constexpr ((AQuantMode_ == QuantMode::PERGROUP_MODE || AQuantMode_ == QuantMode::PERBLOCK_MODE) && - TRANS_A) { + if constexpr ((AQuantMode_ == QuantMode::PERGROUP_MODE || AQuantMode_ == QuantMode::PERBLOCK_MODE) && TRANS_A) { singleCoreMSplit = PER_BLOCK_SIZE << (singleCoreMSplit > PER_BLOCK_SIZE); } else if constexpr (AQuantMode_ == QuantMode::PERBLOCK_MODE) { singleCoreMSplit = CeilPowerOfTwo(singleCoreMSplit); @@ -205,8 +197,7 @@ public: if constexpr (FullLoadMode_ == A_FULL_LOAD_MODE) { blockCoordM = blockIdx_ % mBlockNums_; int64_t curNTailTile = (roundIdx_ == roundNums_ - 1) ? nTailTile_ : 1; - blockCoordN = roundIdx_ * coreNums_ / mBlockNums_ % nBlockNums_ + - blockIdx_ / mBlockNums_ / curNTailTile; + blockCoordN = roundIdx_ * coreNums_ / mBlockNums_ % nBlockNums_ + blockIdx_ / mBlockNums_ / curNTailTile; roundIdx_++; blockCoord = BlockCoord{blockCoordM, blockCoordN, 0, 0}; return true; @@ -279,8 +270,7 @@ private: return inputValue; } - __aicore__ inline void CalSingleCoreShapeByCoord( - int64_t& singleCoreM, int64_t& singleCoreN, BlockCoord blockCoord) + __aicore__ inline void CalSingleCoreShapeByCoord(int64_t& singleCoreM, int64_t& singleCoreN, BlockCoord blockCoord) { const int64_t mIdx = AscendC::Te::Get(blockCoord); const int64_t nIdx = AscendC::Te::Get(blockCoord); diff --git a/include/blaze/gemm/kernel/kernel_qbmm_mx_activation_quant.h b/include/blaze/gemm/kernel/kernel_qbmm_mx_activation_quant.h index 8bb2fca..3e71c4e 100644 --- a/include/blaze/gemm/kernel/kernel_qbmm_mx_activation_quant.h +++ b/include/blaze/gemm/kernel/kernel_qbmm_mx_activation_quant.h @@ -31,19 +31,16 @@ namespace Gemm { namespace Kernel { #define QBMM_MX_KERNEL_CLASS_TEM_PARAMS \ template -#define QBMM_MX_KERNEL_TEM_PARAMS \ - ProblemShape, BlockMmad, BlockEpilogue, BlockScheduler, \ - AscendC::Std::enable_if_t< \ - AscendC::Std::is_same_v> +#define QBMM_MX_KERNEL_TEM_PARAMS \ + ProblemShape, BlockMmad, BlockEpilogue, BlockScheduler, \ + AscendC::Std::enable_if_t> QBMM_MX_KERNEL_CLASS_TEM_PARAMS class GemmUniversal { public: - __aicore__ inline GemmUniversal() - {} - __aicore__ inline ~GemmUniversal() - {} + __aicore__ inline GemmUniversal() {} + __aicore__ inline ~GemmUniversal() {} using BlockMmadParams = typename BlockMmad::Params; using BlockEpilogueParams = typename BlockEpilogue::Params; @@ -91,10 +88,7 @@ public: QBMMTiling qbmmParams; }; - __aicore__ inline void operator()(const Params& params) - { - Run(params); - } + __aicore__ inline void operator()(const Params& params) { Run(params); } private: static constexpr bool WEIGHT_NZ = IsWeightNz::value; @@ -116,17 +110,18 @@ private: __aicore__ inline void Init(const Params& params); __aicore__ inline void Run(const Params& params); __aicore__ inline void ResetGmAddr(const Params& params); - __aicore__ inline void ProcessSingleBatch( - const Params& params, BlockScheduler& bs, uint64_t restBatch, bool isTailRound); + __aicore__ inline void ProcessSingleBatch(const Params& params, BlockScheduler& bs, uint64_t restBatch, + bool isTailRound); __aicore__ inline void ProcessWithBatch(const Params& params, BlockScheduler& bs); - __aicore__ inline void AddBatchOffset( - const Params& params, uint64_t aBatchElementStride, uint64_t bBatchElementStride, uint64_t cBatchStride, - uint64_t scaleABatchStride, uint64_t scaleBBatchStride, uint64_t biasBatchStride); + __aicore__ inline void AddBatchOffset(const Params& params, uint64_t aBatchElementStride, + uint64_t bBatchElementStride, uint64_t cBatchStride, + uint64_t scaleABatchStride, uint64_t scaleBBatchStride, + uint64_t biasBatchStride); template - __aicore__ inline void SetL2Cache( - const ProblemShape& problemShape, uint64_t baseM, uint64_t baseN, TensorB& gmB, TensorC& gmC); + __aicore__ inline void SetL2Cache(const ProblemShape& problemShape, uint64_t baseM, uint64_t baseN, TensorB& gmB, + TensorC& gmC); __aicore__ inline void End(); @@ -184,8 +179,9 @@ __aicore__ inline void GemmUniversal::Run(const Param QBMM_MX_KERNEL_CLASS_TEM_PARAMS template -__aicore__ inline void GemmUniversal::SetL2Cache( - const ProblemShape& problemShape, uint64_t baseM, uint64_t baseN, TensorB& gmB, TensorC& gmC) +__aicore__ inline void GemmUniversal::SetL2Cache(const ProblemShape& problemShape, + uint64_t baseM, uint64_t baseN, + TensorB& gmB, TensorC& gmC) { if constexpr (IS_ATOMIC_ADD) { gmC.SetL2CacheHint(AscendC::Te::CacheMode::CACHE_MODE_DISABLE); @@ -204,16 +200,13 @@ __aicore__ inline void GemmUniversal::SetL2Cache( // 0x7f: 128 cache line alignment for FP8 weight GM streaming if constexpr (TRANS_B) { bool bAlignForL2Stream = (AscendC::Te::Get(problemShape) & cacheLineAlignMask) == 0; - gmB.SetL2CacheHint( - bAlignForL2Stream ? AscendC::Te::CacheMode::CACHE_MODE_DISABLE - : AscendC::Te::CacheMode::CACHE_MODE_NORMAL); + gmB.SetL2CacheHint(bAlignForL2Stream ? AscendC::Te::CacheMode::CACHE_MODE_DISABLE : + AscendC::Te::CacheMode::CACHE_MODE_NORMAL); } else { - bool bAlignForL2Stream = - (AscendC::Te::Get(problemShape) & cacheLineAlignMask) == 0 && - (baseN & cacheLineAlignMask) == 0; - gmB.SetL2CacheHint( - bAlignForL2Stream ? AscendC::Te::CacheMode::CACHE_MODE_DISABLE - : AscendC::Te::CacheMode::CACHE_MODE_NORMAL); + bool bAlignForL2Stream = (AscendC::Te::Get(problemShape) & cacheLineAlignMask) == 0 && + (baseN & cacheLineAlignMask) == 0; + gmB.SetL2CacheHint(bAlignForL2Stream ? AscendC::Te::CacheMode::CACHE_MODE_DISABLE : + AscendC::Te::CacheMode::CACHE_MODE_NORMAL); } } } @@ -249,8 +242,8 @@ __aicore__ inline void GemmUniversal::ResetGmAddr(con } QBMM_MX_KERNEL_CLASS_TEM_PARAMS -__aicore__ inline void GemmUniversal::ProcessWithBatch( - const Params& params, BlockScheduler& bs) +__aicore__ inline void GemmUniversal::ProcessWithBatch(const Params& params, + BlockScheduler& bs) { const auto& qbmmParams = params.qbmmParams; const auto& problemShape = params.problemShape; @@ -274,8 +267,7 @@ __aicore__ inline void GemmUniversal::ProcessWithBatc } const uint64_t cBatchStride = m * n; const uint64_t biasBatchStride = isBiasThreeDim_ ? n : 0; - const uint64_t scaleKLen = - Blaze::Gemm::CeilDiv(k, static_cast(MXFP_DIVISOR_SIZE)) * MXFP_MULTI_BASE_SIZE; + const uint64_t scaleKLen = Blaze::Gemm::CeilDiv(k, static_cast(MXFP_DIVISOR_SIZE)) * MXFP_MULTI_BASE_SIZE; const uint64_t scaleABatchStride = m * scaleKLen; const uint64_t scaleBBatchStride = n * scaleKLen; const uint64_t batchC3C4 = static_cast(qbmmParams.batchC3) * qbmmParams.batchC4; @@ -299,8 +291,8 @@ __aicore__ inline void GemmUniversal::ProcessWithBatc uint64_t curBatchC = 1UL; const uint64_t singleBatchBlockCnt = bs.GetTotalCnt(); const uint64_t batchCount = AscendC::Te::Get(params.problemShape); - const uint64_t tailRoundStart = - (singleBatchBlockCnt * batchCount / AscendC::GetBlockNum()) * AscendC::GetBlockNum(); + const uint64_t tailRoundStart = (singleBatchBlockCnt * batchCount / AscendC::GetBlockNum()) * + AscendC::GetBlockNum(); for (uint64_t b1Index = 0; b1Index < qbmmParams.batchC1; ++b1Index) { uint64_t batchC2Offset = batchC1Offset; uint64_t batchA2Offset = batchA1Offset; @@ -315,9 +307,8 @@ __aicore__ inline void GemmUniversal::ProcessWithBatc batchBOffset_ = batchB3Offset; for (uint64_t b4Index = 0; b4Index < qbmmParams.batchC4; ++b4Index) { const bool isTailRound = curBatchC * singleBatchBlockCnt > tailRoundStart; - AddBatchOffset( - params, aBatchElementStride, bBatchElementStride, cBatchStride, scaleABatchStride, - scaleBBatchStride, biasBatchStride); + AddBatchOffset(params, aBatchElementStride, bBatchElementStride, cBatchStride, scaleABatchStride, + scaleBBatchStride, biasBatchStride); ProcessSingleBatch(params, bs, batchCount - curBatchC, isTailRound); curBatchC++; batchCOffset_ += 1; @@ -356,9 +347,7 @@ __aicore__ inline void GemmUniversal::AddBatchOffset( const auto m = AscendC::Te::Get(params.problemShape); const auto n = AscendC::Te::Get(params.problemShape); const int64_t scaleN = CeilDiv(n, BLOCK_SIZE * ALIGN_NUM_2) * ALIGN_NUM_2; - epilogueOp_.UpdateGlobalAddr( - {batchCOffset_ * m * n, batchCOffset_ * m * scaleN, 0, 0, 0} - ); + epilogueOp_.UpdateGlobalAddr({batchCOffset_ * m * n, batchCOffset_ * m * scaleN, 0, 0, 0}); } QBMM_MX_KERNEL_CLASS_TEM_PARAMS @@ -372,8 +361,10 @@ __aicore__ inline void GemmUniversal::End() } QBMM_MX_KERNEL_CLASS_TEM_PARAMS -__aicore__ inline void GemmUniversal::ProcessSingleBatch( - const Params& params, BlockScheduler& bs, uint64_t restBatch, bool isTailRound) +__aicore__ inline void GemmUniversal::ProcessSingleBatch(const Params& params, + BlockScheduler& bs, + uint64_t restBatch, + bool isTailRound) { const auto& problemShape = params.problemShape; const auto m = AscendC::Te::Get(problemShape); @@ -388,11 +379,11 @@ __aicore__ inline void GemmUniversal::ProcessSingleBa auto layoutC = MakeLayoutC{}(m, n); auto gmA = AscendC::Te::MakeTensor(AscendC::Te::MakeMemPtr(aGmAddr_), layoutA); - auto gmScaleA = - AscendC::Te::MakeTensor(AscendC::Te::MakeMemPtr(scaleAGmAddr_), layoutScaleA); + auto gmScaleA = AscendC::Te::MakeTensor(AscendC::Te::MakeMemPtr(scaleAGmAddr_), + layoutScaleA); auto gmB = AscendC::Te::MakeTensor(AscendC::Te::MakeMemPtr(bGmAddr_), layoutB); - auto gmScaleB = - AscendC::Te::MakeTensor(AscendC::Te::MakeMemPtr(scaleBGmAddr_), layoutScaleB); + auto gmScaleB = AscendC::Te::MakeTensor(AscendC::Te::MakeMemPtr(scaleBGmAddr_), + layoutScaleB); auto gmBias = AscendC::Te::MakeTensor(AscendC::Te::MakeMemPtr(biasGmAddr_), layoutBias); auto gmC = AscendC::Te::MakeTensor(AscendC::Te::MakeMemPtr(cGmAddr_), layoutC); @@ -413,9 +404,8 @@ __aicore__ inline void GemmUniversal::ProcessSingleBa int64_t nPos = 0L; constexpr int64_t kPos = 0L; // K is not split, so the K coordinate is 0. while (bs.GetTileIdx(blockCoord)) { - BlockShape singleShape = - bs.template GetBlockShape( - blockCoord); + BlockShape singleShape = bs.template GetBlockShape(blockCoord); const auto baseM = AscendC::Te::Get(singleShape); const auto baseN = AscendC::Te::Get(singleShape); if (baseM <= 0 || baseN <= 0) { @@ -430,17 +420,16 @@ __aicore__ inline void GemmUniversal::ProcessSingleBa bs.GetTileCoord(blockCoord, mPos, nPos); auto gmBlockA = gmA.Slice(AscendC::Te::MakeCoord(mPos, kPos), AscendC::Te::MakeShape(baseM, k)); - auto gmBlockScaleA = - gmScaleA.Slice(AscendC::Te::MakeCoord(mPos, kPos), AscendC::Te::MakeShape(baseM, scaleKLen)); + auto gmBlockScaleA = gmScaleA.Slice(AscendC::Te::MakeCoord(mPos, kPos), + AscendC::Te::MakeShape(baseM, scaleKLen)); auto gmBlockB = gmB.Slice(AscendC::Te::MakeCoord(kPos, nPos), AscendC::Te::MakeShape(k, baseN)); - auto gmBlockScaleB = - gmScaleB.Slice(AscendC::Te::MakeCoord(kPos, nPos), AscendC::Te::MakeShape(scaleKLen, baseN)); - auto gmBlockBias = - gmBias.Slice(AscendC::Te::MakeCoord(0L, nPos), AscendC::Te::MakeShape(1L, baseN)); + auto gmBlockScaleB = gmScaleB.Slice(AscendC::Te::MakeCoord(kPos, nPos), + AscendC::Te::MakeShape(scaleKLen, baseN)); + auto gmBlockBias = gmBias.Slice(AscendC::Te::MakeCoord(0L, nPos), AscendC::Te::MakeShape(1L, baseN)); auto gmBlockC = gmC.Slice(AscendC::Te::MakeCoord(mPos, nPos), AscendC::Te::MakeShape(baseM, baseN)); - auto locOutUb = AscendC::Te::MakeTensor(ubmemPtr, - AscendC::Te::MakeFrameLayout((baseM + 1) & ~1, Align32(baseN))); - if ASCEND_IS_AIC{ + auto locOutUb = AscendC::Te::MakeTensor( + ubmemPtr, AscendC::Te::MakeFrameLayout((baseM + 1) & ~1, Align32(baseN))); + if ASCEND_IS_AIC { if (isVecSetSyncCom_) { WaitForVector(); } @@ -450,9 +439,10 @@ __aicore__ inline void GemmUniversal::ProcessSingleBa isVecSetSyncCom_ = true; if ASCEND_IS_AIV { WaitForCube(); - epilogueOp_({baseM, baseN, 0, 0}, - {mPos * n + nPos, - mPos * CeilDiv(n, BLOCK_SIZE * ALIGN_NUM_2) * ALIGN_NUM_2 + CeilDiv(nPos, BLOCK_SIZE), 0, 0, 0}); + epilogueOp_( + {baseM, baseN, 0, 0}, + {mPos * n + nPos, mPos * CeilDiv(n, BLOCK_SIZE * ALIGN_NUM_2) * ALIGN_NUM_2 + CeilDiv(nPos, BLOCK_SIZE), + 0, 0, 0}); NotifyCube(); } } diff --git a/tests/ut/op_kernel/quant_matmul_activation_quant/quant_matmul_activation_quant.h b/tests/ut/op_kernel/quant_matmul_activation_quant/quant_matmul_activation_quant.h index 3e834f5..27fc08a 100644 --- a/tests/ut/op_kernel/quant_matmul_activation_quant/quant_matmul_activation_quant.h +++ b/tests/ut/op_kernel/quant_matmul_activation_quant/quant_matmul_activation_quant.h @@ -33,10 +33,10 @@ namespace QBMMUT { template -__aicore__ inline void QBMMMxActivationQuantWrapper( - GM_ADDR x1GM, GM_ADDR x2GM, GM_ADDR x1ScaleGM, GM_ADDR x2ScaleGM, GM_ADDR biasGM, - GM_ADDR yGM, GM_ADDR yScaleGM, const QBMMV3TilingData& tilingData) + uint64_t FullLoadMode = Blaze::Gemm::NONE_FULL_LOAD_MODE> +__aicore__ inline void QBMMMxActivationQuantWrapper(GM_ADDR x1GM, GM_ADDR x2GM, GM_ADDR x1ScaleGM, GM_ADDR x2ScaleGM, + GM_ADDR biasGM, GM_ADDR yGM, GM_ADDR yScaleGM, + const QBMMV3TilingData& tilingData) { using LayoutA = AscendC::Te::NDExtLayoutPtn; using LayoutB = AscendC::Te::NZLayoutPtn; @@ -45,11 +45,12 @@ __aicore__ inline void QBMMMxActivationQuantWrapper( using ProblemShape = AscendC::Te::Shape; using DispatchPolicy = Blaze::Gemm::MatmulWithScaleMx; - using BlockScheduler = Blaze::Gemm::Block::BlockSchedulerQuantBatchMatmulV3< - ProblemShape, FullLoadMode, LayoutA, LayoutB, AType>; - using BlockMmad = Blaze::Gemm::Block::BlockMmad< - DispatchPolicy, AType, LayoutA, BType, LayoutB, CType, LayoutC, BiasType, LayoutBias>; + Blaze::Gemm::KernelMmadWithScaleMxActivationQuant, + Blaze::Gemm::L0C2UB_MODE_DUAL_DST_SPLIT_M>; + using BlockScheduler = Blaze::Gemm::Block::BlockSchedulerQuantBatchMatmulV3; + using BlockMmad = Blaze::Gemm::Block::BlockMmad; using BlockEpilogue = Blaze::Epilogue::Block::BlockEpilogueGeluMxQuant; using QBMMKernel = Blaze::Gemm::Kernel::GemmUniversal; using Params = typename QBMMKernel::Params; @@ -68,13 +69,11 @@ __aicore__ inline void QBMMMxActivationQuantWrapper( FillQbmmSchParams(params.schParams, tilingData); - params.qbmmParams = { - tilingData.batchA1, tilingData.batchA2, tilingData.batchA3, tilingData.batchA4, - tilingData.batchB1, tilingData.batchB2, tilingData.batchB3, tilingData.batchB4, - tilingData.batchC1, tilingData.batchC2, tilingData.batchC3, tilingData.batchC4, - tilingData.biasThreeDim, - tilingData.baseM_qbmm, tilingData.baseN_qbmm, tilingData.baseK_qbmm, - tilingData.isBias, tilingData.dbL0C}; + params.qbmmParams = {tilingData.batchA1, tilingData.batchA2, tilingData.batchA3, tilingData.batchA4, + tilingData.batchB1, tilingData.batchB2, tilingData.batchB3, tilingData.batchB4, + tilingData.batchC1, tilingData.batchC2, tilingData.batchC3, tilingData.batchC4, + tilingData.biasThreeDim, tilingData.baseM_qbmm, tilingData.baseN_qbmm, tilingData.baseK_qbmm, + tilingData.isBias, tilingData.dbL0C}; params.epilogueParams.yGmAddr = reinterpret_cast(yGM); params.epilogueParams.yScaleGmAddr = reinterpret_cast(yScaleGM); @@ -91,9 +90,9 @@ __aicore__ inline void QBMMMxActivationQuantWrapper( } // namespace QBMMUT template -__global__ __aicore__ void qbmm_mx_activation_quant_kernel_entry( - GM_ADDR x1GM, GM_ADDR x2GM, GM_ADDR x1ScaleGM, GM_ADDR x2ScaleGM, GM_ADDR biasGM, - GM_ADDR yGM, GM_ADDR yScaleGM, GM_ADDR tilingGM) +__global__ __aicore__ void qbmm_mx_activation_quant_kernel_entry(GM_ADDR x1GM, GM_ADDR x2GM, GM_ADDR x1ScaleGM, + GM_ADDR x2ScaleGM, GM_ADDR biasGM, GM_ADDR yGM, + GM_ADDR yScaleGM, GM_ADDR tilingGM) { const auto* tilingData = reinterpret_cast(tilingGM); QBMMUT::QBMMMxActivationQuantWrapper( diff --git a/tests/ut/op_kernel/quant_matmul_activation_quant/test_quant_matmul_activation_quant.cpp b/tests/ut/op_kernel/quant_matmul_activation_quant/test_quant_matmul_activation_quant.cpp index ca8c169..cc1ed1f 100644 --- a/tests/ut/op_kernel/quant_matmul_activation_quant/test_quant_matmul_activation_quant.cpp +++ b/tests/ut/op_kernel/quant_matmul_activation_quant/test_quant_matmul_activation_quant.cpp @@ -31,17 +31,14 @@ protected: static void TearDownTestCase() { std::string cleanCmd = std::string("cd ") + UT_KERNEL_SRC_DIR + - "/quant_matmul_activation_quant && rm -rf *.bin"; + "/quant_matmul_activation_quant && rm -rf *.bin"; system(cleanCmd.c_str()); } }; namespace { -size_t GetMxScaleKLen(int64_t k) -{ - return static_cast((k + 63) / 64) * 2UL; -} +size_t GetMxScaleKLen(int64_t k) { return static_cast((k + 63) / 64) * 2UL; } template size_t GetMxInputSize(int64_t elementCount) @@ -61,9 +58,7 @@ void FillGmBuffer(GM_ADDR addr, size_t size, uint8_t value) class GmBuffer { public: - explicit GmBuffer(size_t size) : addr_(reinterpret_cast(AscendC::GmAlloc(size))) - { - } + explicit GmBuffer(size_t size) : addr_(reinterpret_cast(AscendC::GmAlloc(size))) {} ~GmBuffer() { @@ -75,10 +70,7 @@ public: GmBuffer(const GmBuffer&) = delete; GmBuffer& operator=(const GmBuffer&) = delete; - GM_ADDR Get() const - { - return addr_; - } + GM_ADDR Get() const { return addr_; } private: GM_ADDR addr_{nullptr}; @@ -98,9 +90,8 @@ struct ActivationQuantCaseCfg { bool isBias; }; -void FillMxTiling(QBMMV3TilingData* tilingData, int64_t M, int64_t N, int64_t K, uint32_t blockNum, - uint32_t baseM, uint32_t baseN, uint32_t baseK, uint32_t kL1, uint32_t scaleKL1, - uint32_t nBufferNum, bool isBias) +void FillMxTiling(QBMMV3TilingData* tilingData, int64_t M, int64_t N, int64_t K, uint32_t blockNum, uint32_t baseM, + uint32_t baseN, uint32_t baseK, uint32_t kL1, uint32_t scaleKL1, uint32_t nBufferNum, bool isBias) { tilingData->m = M; tilingData->n = N; @@ -149,8 +140,7 @@ void FillMxTiling(QBMMV3TilingData* tilingData, int64_t M, int64_t N, int64_t K, tilingData->biasDtype = 0; } -template +template void RunMxActivationQuantSmoke(const ActivationQuantCaseCfg& cfg) { using CType = float; @@ -162,8 +152,9 @@ void RunMxActivationQuantSmoke(const ActivationQuantCaseCfg& cfg) const size_t x2ScaleSize = scaleKLen * static_cast(cfg.N) * sizeof(AscendC::fp8_e8m0_t); const size_t biasSize = static_cast(cfg.N) * sizeof(BiasType); const size_t ySize = GetMxInputSize(cfg.M * cfg.N); - const size_t yScaleSize = static_cast(cfg.M) * Blaze::Gemm::CeilDiv(static_cast(cfg.N), - static_cast(32)) * sizeof(AscendC::fp8_e8m0_t); + const size_t yScaleSize = static_cast(cfg.M) * + Blaze::Gemm::CeilDiv(static_cast(cfg.N), static_cast(32)) * + sizeof(AscendC::fp8_e8m0_t); GmBuffer x1GM(x1Size); GmBuffer x2GM(x2Size); @@ -192,15 +183,14 @@ void RunMxActivationQuantSmoke(const ActivationQuantCaseCfg& cfg) FillGmBuffer(yScaleGM.Get(), yScaleSize, 0U); auto* tilingData = reinterpret_cast(tilingGM.Get()); - FillMxTiling(tilingData, cfg.M, cfg.N, cfg.K, cfg.blockNum, cfg.baseM, cfg.baseN, cfg.baseK, - cfg.kL1, cfg.scaleKL1, cfg.nBufferNum, cfg.isBias); + FillMxTiling(tilingData, cfg.M, cfg.N, cfg.K, cfg.blockNum, cfg.baseM, cfg.baseN, cfg.baseK, cfg.kL1, cfg.scaleKL1, + cfg.nBufferNum, cfg.isBias); AscendC::SetKernelMode(KernelMode::MIX_MODE); auto kernelFunc = qbmm_mx_activation_quant_kernel_entry; - const bool ok = KERNEL_RUN_KF( - kernelFunc, cfg.blockNum, x1GM.Get(), x2GM.Get(), x1ScaleGM.Get(), x2ScaleGM.Get(), - biasGM.Get(), yGM.Get(), yScaleGM.Get(), tilingGM.Get()); + const bool ok = KERNEL_RUN_KF(kernelFunc, cfg.blockNum, x1GM.Get(), x2GM.Get(), x1ScaleGM.Get(), x2ScaleGM.Get(), + biasGM.Get(), yGM.Get(), yScaleGM.Get(), tilingGM.Get()); ASSERT_TRUE(ok) << "QBMM MX activation quant kernel execution failed"; }