diff --git a/build.sh b/build.sh index 691568b88f..f5e7dd9be4 100755 --- a/build.sh +++ b/build.sh @@ -399,13 +399,24 @@ function build_cmake_modules() ) fi - rm -rf "$BUILD_DIR" + # KEEP_BUILD_DIR=1 reuses the existing CMake cache and object files. The + # AscendC kernel targets (chunk_kda_fwd in particular) dominate a cold build, + # so keeping them lets an unrelated edit rebuild in seconds. Default stays a + # clean build; drop the variable whenever the CMake configuration changes. + if [[ "${KEEP_BUILD_DIR:-0}" == "1" && -f "$BUILD_DIR/CMakeCache.txt" ]]; then + echo "KEEP_BUILD_DIR=1: reusing $BUILD_DIR (incremental build)" + else + rm -rf "$BUILD_DIR" + fi mkdir -p "$BUILD_DIR" cmake "${cmake_args[@]}" \ -B "$BUILD_DIR" \ -S "$PROJECT_ROOT" - cmake --build "$BUILD_DIR" --target install -j 16 + # Parallelism is bounded by RAM, not cores: the CCE compiler needs ~1-2 GB per + # job and the AscendC kernel TUs are the memory-hungry ones. Keep the historical + # default and let callers tune it with BUILD_JOBS. + cmake --build "$BUILD_DIR" --target install -j "${BUILD_JOBS:-16}" } function build_deepep_kernels() diff --git a/csrc/CMakeLists.txt b/csrc/CMakeLists.txt index cf6270260f..1fcd294544 100644 --- a/csrc/CMakeLists.txt +++ b/csrc/CMakeLists.txt @@ -57,6 +57,7 @@ set(OP_SRCS ${PROJECT_OP_SRC_BASE}/sparse_attn_sharedkv_metadata/op_host/sparse_attn_sharedkv_metadata.cpp ${PROJECT_OP_SRC_BASE}/sparse_attention_score/op_host/sparse_attention_score.cpp ${PROJECT_OP_SRC_BASE}/sparse_attention_score/op_host/tiling/sparse_attention_score_tiling.cpp + ${PROJECT_OP_SRC_BASE}/chunk_kda_fwd/op_host/chunk_kda_fwd.cpp ) # A3-only host sources if(SGL_KERNEL_ENABLE_A3_ONLY_OPS) @@ -189,6 +190,32 @@ ascendc_compile_definitions(workspace_kernel PRIVATE -DHAVE_TILING ) +# chunk_kda_fwd: dedicated kernel target with its own compile options. The kernel +# manages cross-pipe synchronization manually (SetFlag/WaitFlag/SyncAll), so the +# aclnn build compiles it with --cce-auto-sync=off; keep that here instead of +# applying it to the shared workspace_kernel target. +ascendc_library(chunk_kda_fwd_kernel STATIC + ${PROJECT_OP_SRC_BASE}/chunk_kda_fwd/op_kernel/chunk_kda_fwd.cpp +) +ascendc_include_directories(chunk_kda_fwd_kernel PRIVATE + ${PROJECT_OP_SRC_BASE}/chunk_kda_fwd/op_kernel + ${PROJECT_OP_SRC_BASE}/utils/kernel + ${ASCEND_INCLUDE_DIR} +) +ascendc_compile_definitions(chunk_kda_fwd_kernel PRIVATE + -DHAVE_WORKSPACE + -DHAVE_TILING +) +ascendc_compile_options(chunk_kda_fwd_kernel PRIVATE + --cce-auto-sync=off + -Wno-deprecated-declarations +) +if(_sgl_soc_lower MATCHES "^ascend950") + ascendc_compile_definitions(chunk_kda_fwd_kernel PRIVATE + -DENABLE_CV_COMM_VIA_SSBUF=true + ) +endif() + # Build catlass kernel in its own target (no sparse_attention_score include dir), # so catlass's #include "tla/layout.hpp" resolves to its own copy and avoids # symbol redefinition (RoundDown/CeilDiv/Max/Min/...) with sparse_attention_score. @@ -258,6 +285,7 @@ target_link_libraries(${OP_PLUGIN_NAME} PRIVATE ${A3_ONLY_KERNEL_TARGETS} workspace_kernel no_workspace_kernel + chunk_kda_fwd_kernel ${CATLASS_WORKSPACE_KERNEL_TARGETS} torch_npu ascendcl @@ -289,6 +317,8 @@ target_include_directories(${OP_PLUGIN_NAME} PRIVATE ${PROJECT_OP_SRC_BASE}/sparse_attn_sharedkv/op_host ${PROJECT_OP_SRC_BASE}/sparse_attention_score/op_host ${PROJECT_OP_SRC_BASE}/sparse_attention_score/op_host/tiling + ${PROJECT_OP_SRC_BASE}/chunk_kda_fwd/op_host + ${PROJECT_OP_SRC_BASE}/chunk_kda_fwd/op_kernel ${TORCH_DIR}/include ${TORCH_DIR}/include/torch/csrc/api/include ${TORCH_NPU_DIR}/include diff --git a/csrc/chunk_kda_fwd/op_host/arch35/chunk_kda_fwd_tiling_impl.h b/csrc/chunk_kda_fwd/op_host/arch35/chunk_kda_fwd_tiling_impl.h new file mode 100644 index 0000000000..4f72a7a012 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_host/arch35/chunk_kda_fwd_tiling_impl.h @@ -0,0 +1,53 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#pragma once + +namespace optiling::arch35 { + +struct ChunkKdaFwdArch35Options { + bool computeGateInPrepare = false; + bool fusePostWu = false; + bool fusePostWuIntoFwdH = false; + bool useDenseFwdH = false; +}; + +inline ChunkKdaFwdArch35Options ConfigureChunkKdaFwdArch35( + bool isAscend950, bool qIsBf16, bool rawGIsFp32, bool hasALog, + bool useGateInKernel, bool safeGate, bool isVarLen, int64_t seqlen, + int64_t vHeads, int64_t chunkSize, int64_t kDim, int64_t vDim, + bool storeQG, bool storeVNew, bool storeH) +{ + ChunkKdaFwdArch35Options options; + const bool shapeSupported = + isAscend950 && chunkSize == 64 && kDim == 128 && vDim == 128; + if (!shapeSupported) { + return options; + } + + // Tiling keys describe shape families independently of the SoC. These + // options only enable arch35 sub-pipelines within the selected family. + options.computeGateInPrepare = + qIsBf16 && rawGIsFp32 && hasALog && + useGateInKernel && safeGate; + const bool denseAligned = !isVarLen && seqlen % chunkSize == 0; + options.useDenseFwdH = denseAligned && qIsBf16; + const bool canFusePreparePostWu = + denseAligned && qIsBf16 && safeGate && vHeads % 2 == 0; + options.fusePostWuIntoFwdH = + options.useDenseFwdH && canFusePreparePostWu && + options.computeGateInPrepare && + !storeQG && !storeVNew && !storeH; + options.fusePostWu = + canFusePreparePostWu && !options.fusePostWuIntoFwdH; + return options; +} + +} // namespace optiling::arch35 diff --git a/csrc/chunk_kda_fwd/op_host/chunk_kda_fwd.cpp b/csrc/chunk_kda_fwd/op_host/chunk_kda_fwd.cpp new file mode 100644 index 0000000000..60650632fe --- /dev/null +++ b/csrc/chunk_kda_fwd/op_host/chunk_kda_fwd.cpp @@ -0,0 +1,687 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file chunk_kda_fwd.cpp + * \brief host-side direct-launch implementation of chunk_kda_fwd + * + * Port of the aclnn ChunkKdaFwd front-end. Instead of building an aclnn executor + * graph (transpose/reshape + ADD_TO_LAUNCHER_LIST_AICORE), the layout conversions + * are done with plain torch ops and the kernel is launched directly through + * EXEC_KERNEL_CMD. The tiling computation mirrors Tiling4ChunkKdaFwd and the + * serialized tiling struct is copied GM->stack inside the kernel entry. + */ + +#include +#include +#include +#include +#include + +#include "acl/acl.h" + +#include +#include +#include + +#include "torch_npu/csrc/core/npu/NPUStream.h" +#include "torch_npu/csrc/core/npu/DeviceUtils.h" +#include "torch_npu/csrc/framework/OpCommand.h" + +#include "tiling/platform/platform_ascendc.h" + +#include "defines.h" +#include "torch_helper.h" +#include "aclrtlaunch_chunk_kda_fwd.h" + +#include "../op_kernel/chunk_kda_fwd_tiling_data.h" +#include "arch35/chunk_kda_fwd_tiling_impl.h" + +namespace sglang { +namespace npu_kernel { + +constexpr uint32_t PADDING_BYTE = 32U; + +constexpr uint64_t KDA_ALIGN = 512; +constexpr uint64_t KDA_SOLVE_SCRATCH_SLOTS = 5; +constexpr uint64_t KDA_SOLVE_PIPELINE_DEPTH = 4; +constexpr uint64_t KDA_SCORE_QUEUE_SLOTS = 4; +constexpr uint64_t KDA_SCORE_SCRATCH_PLANES = 3; +constexpr uint64_t KDA_GDN_PIPELINE_DEPTH = 2; + +constexpr int64_t MAX_KDA_K_DIM = 256; +constexpr int64_t MAX_KDA_HEAD_NUM = 128; +constexpr int64_t MAX_KDA_VARLEN_SEQUENCES = 1024; + + +enum class KdaFwdLayout { + BSND, + BNSD, + TND, + NTD, +}; + +uint64_t AlignWorkspace(uint64_t bytes) +{ + return (bytes + KDA_ALIGN - 1) / KDA_ALIGN * KDA_ALIGN; +} + +uint64_t AllocateWorkspace(uint64_t &cursor, uint64_t bytes) +{ + const uint64_t offset = AlignWorkspace(cursor); + cursor = offset + bytes; + return offset; +} + +struct KdaShapeInfo { + bool isRank3 = false; + int64_t batch = 0; + int64_t seqlen = 0; + int64_t hNum = 0; + int64_t hvNum = 0; + int64_t kDim = 0; + int64_t vDim = 0; + int64_t seqNum = 0; + int64_t totalChunks = 0; +}; + +bool ParseLayout(const std::string &layout, KdaFwdLayout &parsed) +{ + if (layout == "BSND") { + parsed = KdaFwdLayout::BSND; + } else if (layout == "BNSD") { + parsed = KdaFwdLayout::BNSD; + } else if (layout == "TND") { + parsed = KdaFwdLayout::TND; + } else if (layout == "NTD") { + parsed = KdaFwdLayout::NTD; + } else { + return false; + } + return true; +} + +int64_t Dim(const at::Tensor &tensor, size_t idx) +{ + return tensor.size(static_cast(idx)); +} + +bool SameShape(const at::Tensor &lhs, const at::Tensor &rhs) +{ + return lhs.dim() == rhs.dim() && lhs.sizes() == rhs.sizes(); +} + +bool HasShape(const at::Tensor &tensor, std::initializer_list expected) +{ + if (tensor.dim() != static_cast(expected.size())) { + return false; + } + size_t idx = 0; + for (int64_t dim : expected) { + if (tensor.size(static_cast(idx++)) != dim) { + return false; + } + } + return true; +} + +int64_t CountChunks(const c10::optional &cuSeqlens, int64_t seqlen, int64_t chunkSize) +{ + if (!cuSeqlens.has_value()) { + return (seqlen + chunkSize - 1) / chunkSize; + } + const at::Tensor cu = cuSeqlens->cpu().contiguous(); + const int64_t *cuPtr = reinterpret_cast(cu.data_ptr()); + const int64_t n = cu.numel(); + int64_t chunks = 0; + for (int64_t idx = 0; idx + 1 < n; ++idx) { + chunks += (cuPtr[idx + 1] - cuPtr[idx] + chunkSize - 1) / chunkSize; + } + return chunks; +} + +void ResolveShapeInfo(const at::Tensor &q, const at::Tensor &v, const c10::optional &beta, + const c10::optional &g, KdaFwdLayout layout, int64_t chunkSize, + const c10::optional &cuSeqlens, KdaShapeInfo &info) +{ + info.isRank3 = layout == KdaFwdLayout::TND || layout == KdaFwdLayout::NTD; + const int64_t tensorRank = info.isRank3 ? 3 : 4; + const int64_t betaRank = info.isRank3 ? 2 : 3; + TORCH_CHECK(q.dim() == tensorRank && v.dim() == tensorRank && g->dim() == tensorRank && + beta->dim() == betaRank, + "chunk_kda_fwd: q/k/v/g and beta ranks must match layout: rank3/rank2 for TND/NTD, " + "rank4/rank3 for BSND/BNSD."); + + if (layout == KdaFwdLayout::TND) { + info.batch = 1; + info.seqlen = Dim(q, 0); + info.hNum = Dim(q, 1); + info.kDim = Dim(q, 2); + info.hvNum = Dim(v, 1); + info.vDim = Dim(v, 2); + TORCH_CHECK(HasShape(v, {info.seqlen, info.hvNum, info.vDim}) && + HasShape(*g, {info.seqlen, info.hvNum, info.kDim}) && + HasShape(*beta, {info.seqlen, info.hvNum}), + "chunk_kda_fwd: TND expects v/g/beta as [T,HV,V], [T,HV,K], [T,HV]."); + } else if (layout == KdaFwdLayout::NTD) { + info.batch = 1; + info.hNum = Dim(q, 0); + info.seqlen = Dim(q, 1); + info.kDim = Dim(q, 2); + info.hvNum = Dim(v, 0); + info.vDim = Dim(v, 2); + TORCH_CHECK(HasShape(v, {info.hvNum, info.seqlen, info.vDim}) && + HasShape(*g, {info.hvNum, info.seqlen, info.kDim}) && + HasShape(*beta, {info.hvNum, info.seqlen}), + "chunk_kda_fwd: NTD expects v/g/beta as [HV,T,V], [HV,T,K], [HV,T]."); + } else if (layout == KdaFwdLayout::BSND) { + info.batch = Dim(q, 0); + info.seqlen = Dim(q, 1); + info.hNum = Dim(q, 2); + info.kDim = Dim(q, 3); + info.hvNum = Dim(v, 2); + info.vDim = Dim(v, 3); + TORCH_CHECK(HasShape(v, {info.batch, info.seqlen, info.hvNum, info.vDim}) && + HasShape(*g, {info.batch, info.seqlen, info.hvNum, info.kDim}) && + HasShape(*beta, {info.batch, info.seqlen, info.hvNum}), + "chunk_kda_fwd: BSND expects v/g/beta as [B,T,HV,V], [B,T,HV,K], [B,T,HV]."); + } else { // BNSD + info.batch = Dim(q, 0); + info.hNum = Dim(q, 1); + info.seqlen = Dim(q, 2); + info.kDim = Dim(q, 3); + info.hvNum = Dim(v, 1); + info.vDim = Dim(v, 3); + TORCH_CHECK(HasShape(v, {info.batch, info.hvNum, info.seqlen, info.vDim}) && + HasShape(*g, {info.batch, info.hvNum, info.seqlen, info.kDim}) && + HasShape(*beta, {info.batch, info.hvNum, info.seqlen}), + "chunk_kda_fwd: BNSD expects v/g/beta as [B,HV,T,V], [B,HV,T,K], [B,HV,T]."); + } + + info.seqNum = cuSeqlens.has_value() ? cuSeqlens->numel() - 1 : info.batch; + info.totalChunks = CountChunks(cuSeqlens, info.seqlen, chunkSize); +} + +void ComputeTilingData(int64_t batch, int64_t seqlen, int64_t hNum, int64_t hvNum, int64_t kDim, + int64_t vDim, int64_t chunkSize, int64_t seqNum, int64_t totalChunks, + int64_t inputRank, double scale, double lowerBound, bool hasInitialState, + bool isVarLen, bool safeGate, bool sequenceMajor, bool useGateInKernel, + bool hasALog, bool hasDtBias, int64_t gateDataType, + bool storeFinalState, bool storeGk, bool storeW, + bool storeU, bool storeQG, bool storeKg, bool storeVNew, bool storeH, + uint32_t blockDim, const optiling::arch35::ChunkKdaFwdArch35Options &arch35Options, + int64_t betaDataType, ChunkKdaFwd::ChunkKdaFwdTilingData &td) +{ + std::memset(&td, 0, sizeof(td)); + + td.batch = batch; + td.seqNum = seqNum; + td.qHeadNum = hNum; + td.vHeadNum = hvNum; + td.seqlen = seqlen; + td.kHeadDim = kDim; + td.vHeadDim = vDim; + td.chunkSize = chunkSize; + td.totalChunks = totalChunks; + td.inputRank = inputRank; + td.scale = static_cast(scale); + td.lowerBound = static_cast(lowerBound); + td.hasInitialState = hasInitialState; + td.isVarLen = isVarLen; + td.safeGate = safeGate; + td.inputSequenceMajor = sequenceMajor; + td.useGateInKernel = useGateInKernel; + td.hasALog = hasALog; + td.hasDtBias = hasDtBias; + td.computeGateInPrepare = arch35Options.computeGateInPrepare; + td.fusePostWu = arch35Options.fusePostWu; + td.fusePostWuIntoFwdH = arch35Options.fusePostWuIntoFwdH; + td.useDenseFwdH = arch35Options.useDenseFwdH; + td.storeFinalState = storeFinalState; + td.storeGk = storeGk; + td.storeW = storeW; + td.storeU = storeU; + td.storeQG = storeQG; + td.storeKg = storeKg; + td.storeVNew = storeVNew; + td.storeH = storeH; + td.gateDataType = gateDataType; + td.gateUsedCoreNum = static_cast(blockDim) * 2; + td.prepareUsedCoreNum = blockDim; + td.postWuUsedCoreNum = blockDim; + td.outputUsedCoreNum = blockDim; + td.betaDataType = betaDataType; + + const uint64_t dataBytes = 2; // direct-launch currently only supports bf16 q/k/v + const uint64_t tokenHeads = static_cast(batch) * hvNum * seqlen; + const uint64_t kTensorBytes = tokenHeads * kDim * dataBytes; + const uint64_t vTensorBytes = tokenHeads * vDim * dataBytes; + const uint64_t gkBytes = tokenHeads * kDim * sizeof(float); + const uint64_t stateElements = static_cast(seqNum) * hvNum * kDim * vDim; + const uint64_t hChunkCount = + isVarLen ? static_cast(totalChunks) : static_cast(batch) * totalChunks; + const uint64_t hBytes = hChunkCount * hvNum * kDim * vDim * dataBytes; + + uint64_t cursor = 0; + td.gkStorageOffset = storeGk ? 0 : AllocateWorkspace(cursor, gkBytes); + td.finalStateStorageOffset = + storeFinalState ? 0 : AllocateWorkspace(cursor, stateElements * sizeof(float)); + td.wStorageOffset = storeW ? 0 : AllocateWorkspace(cursor, kTensorBytes); + td.uStorageOffset = storeU ? 0 : AllocateWorkspace(cursor, vTensorBytes); + td.qgStorageOffset = storeQG ? 0 : AllocateWorkspace(cursor, kTensorBytes); + td.kgStorageOffset = storeKg ? 0 : AllocateWorkspace(cursor, kTensorBytes); + const uint64_t vNewStorageBytes = + arch35Options.useDenseFwdH && !storeVNew + ? static_cast(batch) * hvNum * chunkSize * vDim * dataBytes + : vTensorBytes; + td.vNewStorageOffset = storeVNew ? 0 : AllocateWorkspace(cursor, vNewStorageBytes); + const uint64_t hStorageBytes = + arch35Options.useDenseFwdH && !storeH + ? static_cast(batch) * hvNum * kDim * vDim * dataBytes + : hBytes; + td.hStorageOffset = storeH ? 0 : AllocateWorkspace(cursor, hStorageBytes); + td.qgScaledOffset = AllocateWorkspace(cursor, kTensorBytes); + + const uint64_t matrixBytes = tokenHeads * chunkSize * sizeof(float); + td.prepareAqkFp32Offset = AllocateWorkspace(cursor, matrixBytes); + td.prepareAkkFp32Offset = AllocateWorkspace(cursor, matrixBytes); + td.prepareScratchOffset = AlignWorkspace(cursor); + const uint64_t solveDepth = safeGate ? KDA_SOLVE_PIPELINE_DEPTH : 1; + const uint64_t solveBytes = + static_cast(blockDim) * solveDepth * KDA_SOLVE_SCRATCH_SLOTS * chunkSize * chunkSize * + sizeof(float); + const uint64_t scoreBytes = static_cast(blockDim) * KDA_SCORE_QUEUE_SLOTS * + KDA_SCORE_SCRATCH_PLANES * chunkSize * kDim * dataBytes; + cursor = td.prepareScratchOffset + AlignWorkspace(solveBytes) + scoreBytes; + + td.postWuScratchOffset = AlignWorkspace(cursor); + if (!arch35Options.fusePostWu && !arch35Options.fusePostWuIntoFwdH) { + cursor = td.postWuScratchOffset + tokenHeads * kDim * sizeof(float); + } + + td.fwdHWorkspaceBaseOffset = AlignWorkspace(cursor); + uint64_t fwdHCursor = 0; + td.vWorkspaceOffset = AllocateWorkspace( + fwdHCursor, static_cast(blockDim) * chunkSize * vDim * sizeof(float) * + KDA_GDN_PIPELINE_DEPTH); + td.vUpdateWorkspaceOffset = AllocateWorkspace( + fwdHCursor, static_cast(blockDim) * chunkSize * vDim * sizeof(float) * + KDA_GDN_PIPELINE_DEPTH); + td.kDecayWorkspaceOffset = AllocateWorkspace( + fwdHCursor, static_cast(blockDim) * chunkSize * kDim * sizeof(float) * + KDA_GDN_PIPELINE_DEPTH); + td.hWorkspaceOffset = AllocateWorkspace( + fwdHCursor, static_cast(blockDim) * kDim * vDim * sizeof(float) * + KDA_GDN_PIPELINE_DEPTH); + const uint64_t tokenBatch = isVarLen ? static_cast(seqNum) : 1; + td.numSeqWorkspaceOffset = AllocateWorkspace(fwdHCursor, (tokenBatch + 1) * sizeof(int64_t)); + td.numChunksWorkspaceOffset = AllocateWorkspace(fwdHCursor, (tokenBatch + 1) * sizeof(int64_t)); + cursor = td.fwdHWorkspaceBaseOffset + AlignWorkspace(fwdHCursor); + + td.outputScratchOffset = AllocateWorkspace(cursor, 2 * tokenHeads * vDim * sizeof(float)); +} + +std::tuple, c10::optional, at::Tensor, at::Tensor, + c10::optional, c10::optional, c10::optional, + c10::optional, c10::optional, c10::optional> +chunk_kda_fwd(const at::Tensor &q, const at::Tensor &k, const at::Tensor &v, const at::Tensor &g, + const at::Tensor &beta, const c10::optional &aLog, + const c10::optional &dtBias, const c10::optional &initialState, + const c10::optional &cuSeqlens, + const c10::optional &chunkIndices, const std::string &layout, double scale, + int64_t chunkSize, bool safeGate, double lowerBound, bool useGateInKernel, + bool stateVFirst, bool outputFinalState, bool outputGk, bool outputW, bool outputU, + bool outputQG, bool outputKg, bool outputVNew, bool outputH) +{ + TORCH_CHECK(q.defined() && k.defined() && v.defined() && g.defined() && beta.defined(), + "chunk_kda_fwd: q, k, v, g and beta must be defined"); + + TORCH_CHECK(chunkSize == 64 || chunkSize == 128, "chunk_kda_fwd: chunkSize must be 64 or 128"); + + KdaFwdLayout parsedLayout; + TORCH_CHECK(ParseLayout(layout, parsedLayout), "chunk_kda_fwd: layout must be BSND, BNSD, TND or NTD"); + + const bool isRank3 = parsedLayout == KdaFwdLayout::TND || parsedLayout == KdaFwdLayout::NTD; + + // dtype checks: q/k/v must be bf16 (direct-launch restriction), g/beta fp32 or bf16. + TORCH_CHECK(q.scalar_type() == at::kBFloat16 && k.scalar_type() == at::kBFloat16 && + v.scalar_type() == at::kBFloat16, + "chunk_kda_fwd: q, k and v must be bfloat16 (direct-launch currently supports bf16)"); + const at::ScalarType gateType = g.scalar_type(); + TORCH_CHECK(gateType == at::kFloat || gateType == at::kBFloat16, + "chunk_kda_fwd: g must be float32 or bfloat16"); + const at::ScalarType betaType = beta.scalar_type(); + TORCH_CHECK(betaType == at::kFloat || betaType == at::kBFloat16, + "chunk_kda_fwd: beta must be float32 or bfloat16"); + const bool gIsFp32 = gateType == at::kFloat; + const int64_t gateDataType = gIsFp32 ? 2 : (gateType == at::kBFloat16 ? 1 : 0); + const int64_t betaDataType = betaType == at::kFloat ? 0 : 1; + + if (aLog.has_value() && aLog->defined()) { + TORCH_CHECK(aLog->scalar_type() == at::kFloat, "chunk_kda_fwd: a_log must be float32"); + } + if (dtBias.has_value() && dtBias->defined()) { + TORCH_CHECK(dtBias->scalar_type() == at::kFloat, "chunk_kda_fwd: dt_bias must be float32"); + } + if (initialState.has_value() && initialState->defined()) { + TORCH_CHECK(initialState->scalar_type() == at::kFloat, + "chunk_kda_fwd: initial_state must be float32"); + } + if (useGateInKernel) { + TORCH_CHECK(aLog.has_value() && aLog->defined(), + "chunk_kda_fwd: a_log is required when use_gate_in_kernel is true"); + } + + KdaShapeInfo info; + ResolveShapeInfo(q, v, beta, g, parsedLayout, chunkSize, cuSeqlens, info); + + TORCH_CHECK(info.hNum > 0 && info.hvNum >= info.hNum && info.hvNum % info.hNum == 0, + "chunk_kda_fwd: H and HV must be positive, HV must be a multiple of H"); + TORCH_CHECK(info.hNum <= MAX_KDA_HEAD_NUM && info.hvNum <= MAX_KDA_HEAD_NUM, + "chunk_kda_fwd: H and HV must be <= 128"); + TORCH_CHECK(info.kDim >= 16 && info.kDim <= MAX_KDA_K_DIM && info.kDim % 16 == 0 && + info.vDim >= 16 && info.vDim <= 256 && info.vDim % 16 == 0, + "chunk_kda_fwd: K/V must be multiples of 16, K <= 256 and V <= 256"); + TORCH_CHECK(SameShape(q, k) && q.dim() == k.dim(), + "chunk_kda_fwd: q and k must have identical shape"); + TORCH_CHECK(info.seqlen > 0 && info.totalChunks > 0, "chunk_kda_fwd: invalid sequence length"); + if (cuSeqlens.has_value() && cuSeqlens->defined()) { + TORCH_CHECK(cuSeqlens->scalar_type() == at::kLong, "chunk_kda_fwd: cu_seqlens must be int64"); + TORCH_CHECK(cuSeqlens->numel() >= 2, "chunk_kda_fwd: cu_seqlens must contain at least [0, T]"); + TORCH_CHECK(!isRank3 || info.batch == 1, "chunk_kda_fwd: rank4 varlen requires B=1"); + TORCH_CHECK(info.seqNum <= MAX_KDA_VARLEN_SEQUENCES, + "chunk_kda_fwd: varlen supports at most 1024 sequences"); + const at::Tensor cuCpu = cuSeqlens->cpu().contiguous(); + const int64_t *cuPtr = reinterpret_cast(cuCpu.data_ptr()); + TORCH_CHECK(cuPtr[0] == 0, "chunk_kda_fwd: cu_seqlens[0] must be 0"); + TORCH_CHECK(cuPtr[cuCpu.numel() - 1] == info.seqlen, + "chunk_kda_fwd: cu_seqlens last element must equal the sequence length"); + for (int64_t idx = 0; idx + 1 < cuCpu.numel(); ++idx) { + TORCH_CHECK(cuPtr[idx] <= cuPtr[idx + 1], + "chunk_kda_fwd: cu_seqlens must be nondecreasing"); + } + } + if (chunkIndices.has_value() && chunkIndices->defined()) { + TORCH_CHECK(cuSeqlens.has_value() && cuSeqlens->defined(), + "chunk_kda_fwd: chunk_indices requires cu_seqlens"); + TORCH_CHECK(chunkIndices->scalar_type() == at::kLong, "chunk_kda_fwd: chunk_indices must be int64"); + TORCH_CHECK(chunkIndices->numel() == info.totalChunks * 2, + "chunk_kda_fwd: chunk_indices must contain one (seq_id, chunk_id) pair per chunk"); + } + if (useGateInKernel && safeGate) { + TORCH_CHECK(lowerBound >= -5.0 && lowerBound < 0.0, + "chunk_kda_fwd: lower_bound must be in [-5, 0) when safe_gate is true"); + } + if (initialState.has_value() && initialState->defined()) { + const bool valid = stateVFirst + ? HasShape(*initialState, {info.seqNum, info.hvNum, info.vDim, info.kDim}) + : HasShape(*initialState, {info.seqNum, info.hvNum, info.kDim, info.vDim}); + TORCH_CHECK(valid, + "chunk_kda_fwd: initial_state must be [N,HV,K,V] (state_v_first=false) or " + "[N,HV,V,K] (state_v_first=true)"); + } + + const auto ascendcPlatform = platform_ascendc::PlatformAscendCManager::GetInstance(); + const uint32_t physicalCoreNum = std::max(ascendcPlatform->GetCoreNumAic(), 1); + const bool isAscend950 = + ascendcPlatform->GetSocVersion() == platform_ascendc::SocVersion::ASCEND950; + const bool isVarLen = cuSeqlens.has_value() && cuSeqlens->defined(); + const uint64_t fwdHTaskCount = static_cast(isVarLen ? info.seqNum : info.batch) * + info.hvNum; + const uint32_t blockDim = + isAscend950 ? static_cast( + std::min(physicalCoreNum, std::max(fwdHTaskCount, 1))) : + physicalCoreNum; + + const bool hasALog = aLog.has_value() && aLog->defined(); + const bool hasDtBias = dtBias.has_value() && dtBias->defined(); + const bool hasInitialState = initialState.has_value() && initialState->defined(); + + const auto arch35Options = optiling::arch35::ConfigureChunkKdaFwdArch35( + isAscend950, true /* qIsBf16 */, gIsFp32, hasALog, useGateInKernel, safeGate, isVarLen, + info.seqlen, info.hvNum, chunkSize, info.kDim, info.vDim, outputQG, outputVNew, outputH); + + const bool storeFinalState = outputFinalState; + const bool storeGk = outputGk; + const bool storeW = outputW; + const bool storeU = outputU; + const bool storeQG = outputQG; + const bool storeKg = outputKg; + const bool storeVNew = outputVNew; + const bool storeH = outputH; + + ChunkKdaFwd::ChunkKdaFwdTilingData tilingData; + ComputeTilingData(info.batch, info.seqlen, info.hNum, info.hvNum, info.kDim, info.vDim, chunkSize, + info.seqNum, info.totalChunks, info.isRank3 ? 3 : 4, scale, lowerBound, + hasInitialState, isVarLen, safeGate, parsedLayout == KdaFwdLayout::BSND, + useGateInKernel, hasALog, hasDtBias, gateDataType, storeFinalState, + storeGk, storeW, storeU, storeQG, storeKg, storeVNew, storeH, blockDim, + arch35Options, betaDataType, tilingData); + + // ---- layout conversion to the kernel's input convention ---- + at::Tensor qHead, kHead, vHead, gHead, betaHead; + if (parsedLayout == KdaFwdLayout::BSND) { + qHead = q.contiguous(); + kHead = k.contiguous(); + vHead = v.contiguous(); + gHead = g.contiguous(); + betaHead = beta.transpose(1, 2).contiguous(); // (B, S, HV) -> (B, HV, S) + } else if (parsedLayout == KdaFwdLayout::BNSD) { + qHead = q.contiguous(); + kHead = k.contiguous(); + vHead = v.contiguous(); + gHead = g.contiguous(); + betaHead = beta.contiguous(); + } else if (parsedLayout == KdaFwdLayout::TND) { + qHead = q.transpose(0, 1).contiguous().reshape({1, info.hNum, info.seqlen, info.kDim}); + kHead = k.transpose(0, 1).contiguous().reshape({1, info.hNum, info.seqlen, info.kDim}); + vHead = v.transpose(0, 1).contiguous().reshape({1, info.hvNum, info.seqlen, info.vDim}); + gHead = g.transpose(0, 1).contiguous().reshape({1, info.hvNum, info.seqlen, info.kDim}); + betaHead = beta.transpose(0, 1).contiguous().reshape({1, info.hvNum, info.seqlen}); + } else { // NTD + qHead = q.reshape({1, info.hNum, info.seqlen, info.kDim}); + kHead = k.reshape({1, info.hNum, info.seqlen, info.kDim}); + vHead = v.reshape({1, info.hvNum, info.seqlen, info.vDim}); + gHead = g.reshape({1, info.hvNum, info.seqlen, info.kDim}); + betaHead = beta.reshape({1, info.hvNum, info.seqlen}); + } + + at::Tensor initStateContig; + void *initStatePtr = nullptr; + if (hasInitialState) { + at::Tensor init = *initialState; + if (stateVFirst) { + init = init.transpose(-1, -2).contiguous(); // (N, HV, V, K) -> (N, HV, K, V) + } + initStateContig = init.contiguous(); + initStatePtr = initStateContig.data_ptr(); + } + + void *aLogPtr = nullptr; + if (hasALog) { + aLogPtr = aLog->contiguous().data_ptr(); + } + void *dtBiasPtr = nullptr; + if (hasDtBias) { + dtBiasPtr = dtBias->contiguous().data_ptr(); + } + + at::Tensor cuSeqlensContig, chunkIndicesContig; + void *cuSeqlensPtr = nullptr; + void *chunkIndicesPtr = nullptr; + if (cuSeqlens.has_value() && cuSeqlens->defined()) { + cuSeqlensContig = cuSeqlens->contiguous(); + cuSeqlensPtr = cuSeqlensContig.data_ptr(); + } + if (chunkIndices.has_value() && chunkIndices->defined()) { + chunkIndicesContig = chunkIndices->contiguous(); + chunkIndicesPtr = chunkIndicesContig.data_ptr(); + } + + // ---- outputs (kernel internal layouts, then reshaped/transposed for the user) ---- + const int64_t b = info.batch; + const int64_t s = info.seqlen; + const int64_t hv = info.hvNum; + const int64_t kd = info.kDim; + const int64_t vd = info.vDim; + const int64_t c = chunkSize; + + at::Tensor attnOut = at::empty({b, s, hv, vd}, q.options()); + at::Tensor aqkOut = at::empty({b, hv, s, c}, q.options()); + at::Tensor akkOut = at::empty({b, hv, s, c}, q.options()); + + at::Tensor finalStateOut; + void *finalStatePtr = nullptr; + if (storeFinalState) { + finalStateOut = at::empty({info.seqNum, hv, kd, vd}, + at::TensorOptions().dtype(at::kFloat).device(q.device())); + finalStatePtr = finalStateOut.data_ptr(); + } + + at::Tensor gkOut; + void *gkPtr = nullptr; + if (storeGk) { + gkOut = at::empty({b, hv, s, kd}, + at::TensorOptions().dtype(at::kFloat).device(q.device())); + gkPtr = gkOut.data_ptr(); + } + + at::Tensor wOut; + void *wPtr = nullptr; + if (storeW) { + wOut = at::empty({b, hv, s, kd}, q.options()); + wPtr = wOut.data_ptr(); + } + + at::Tensor uOut; + void *uPtr = nullptr; + if (storeU) { + uOut = at::empty({b, hv, s, vd}, q.options()); + uPtr = uOut.data_ptr(); + } + + at::Tensor qgOut; + void *qgPtr = nullptr; + if (storeQG) { + qgOut = at::empty({b, hv, s, kd}, q.options()); + qgPtr = qgOut.data_ptr(); + } + + at::Tensor kgOut; + void *kgPtr = nullptr; + if (storeKg) { + kgOut = at::empty({b, hv, s, kd}, q.options()); + kgPtr = kgOut.data_ptr(); + } + + at::Tensor vNewOut; + void *vNewPtr = nullptr; + if (storeVNew) { + vNewOut = at::empty({b, hv, s, vd}, q.options()); + vNewPtr = vNewOut.data_ptr(); + } + + at::Tensor hOutInternal; + void *hPtr = nullptr; + if (storeH) { + // kernel-internal h layout is (B, HV, NC, K, V); reshaped below for the user. + hOutInternal = at::empty({b, hv, info.totalChunks, kd, vd}, q.options()); + hPtr = hOutInternal.data_ptr(); + } + + // ---- workspace + tiling ---- + int32_t tilingSize = (static_cast(sizeof(ChunkKdaFwd::ChunkKdaFwdTilingData)) + + static_cast(PADDING_BYTE) - 1) / + PADDING_BYTE * PADDING_BYTE; + + auto cpuTiling = at::empty({tilingSize}, at::kByte); + std::memcpy(cpuTiling.data_ptr(), &tilingData, sizeof(ChunkKdaFwd::ChunkKdaFwdTilingData)); + at::Tensor tilingTensor = TorchNpuHelper::CopyTensorHostToDevice(cpuTiling); + + const uint64_t tokenHeads = static_cast(b) * hv * s; + const uint64_t totalUserWorkspace = tilingData.outputScratchOffset + 2 * tokenHeads * vd * sizeof(float); + const uint64_t sysytemWorkspaceBytes = static_cast(ascendcPlatform->GetLibApiWorkSpaceSize()); + const int64_t totalWorkspaceBytes = + sysytemWorkspaceBytes + static_cast(AlignWorkspace(totalUserWorkspace)); + auto workspaceTensor = + at::empty({totalWorkspaceBytes}, at::TensorOptions().dtype(at::kByte).device(q.device())); + + EXEC_KERNEL_CMD(chunk_kda_fwd, blockDim, qHead, kHead, vHead, gHead, betaHead, aLogPtr, dtBiasPtr, + initStatePtr, cuSeqlensPtr, chunkIndicesPtr, attnOut, finalStatePtr, gkPtr, aqkOut, + akkOut, wPtr, uPtr, qgPtr, kgPtr, vNewPtr, hPtr, workspaceTensor, tilingTensor); + + // ---- post-process outputs to the aclnn-documented user layouts ---- + at::Tensor attn = attnOut; + at::Tensor aqk = aqkOut; + at::Tensor akk = akkOut; + if (isRank3) { + attn = attnOut.reshape({s, hv, vd}); + aqk = aqkOut.reshape({hv, s, c}); + akk = akkOut.reshape({hv, s, c}); + } + + c10::optional finalState = c10::nullopt; + if (storeFinalState) { + at::Tensor fs = finalStateOut; + if (stateVFirst) { + fs = fs.transpose(-1, -2).contiguous(); // (N, HV, K, V) -> (N, HV, V, K) + } + finalState = fs; + } + + c10::optional gk = c10::nullopt; + if (storeGk) { + gk = isRank3 ? gkOut.reshape({hv, s, kd}) : gkOut; + } + + c10::optional w = c10::nullopt; + if (storeW) { + w = isRank3 ? wOut.reshape({hv, s, kd}) : wOut; + } + + c10::optional u = c10::nullopt; + if (storeU) { + u = isRank3 ? uOut.reshape({hv, s, vd}) : uOut; + } + + c10::optional qg = c10::nullopt; + if (storeQG) { + qg = isRank3 ? qgOut.reshape({hv, s, kd}) : qgOut; + } + + c10::optional kg = c10::nullopt; + if (storeKg) { + kg = isRank3 ? kgOut.reshape({hv, s, kd}) : kgOut; + } + + c10::optional vNew = c10::nullopt; + if (storeVNew) { + vNew = isRank3 ? vNewOut.reshape({hv, s, vd}) : vNewOut; + } + + c10::optional h = c10::nullopt; + if (storeH) { + // (B, HV, NC, K, V) -> (B, NC, HV, K, V) [or state_v_first: (B, NC, HV, V, K)] + at::Tensor hView = hOutInternal.permute({0, 2, 1, 3, 4}); + if (stateVFirst) { + hView = hOutInternal.permute({0, 2, 1, 4, 3}); + } + at::Tensor hOut = hView.contiguous(); + if (isRank3) { + hOut = hOut.reshape({info.totalChunks, hv, stateVFirst ? vd : kd, stateVFirst ? kd : vd}); + } + h = hOut; + } + + return std::make_tuple(attn, finalState, gk, aqk, akk, w, u, qg, kg, vNew, h); +} + +} // namespace npu_kernel +} // namespace sglang diff --git a/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/epilogue/block/block_epilogue_kda_fwdh_update_arch22.hpp b/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/epilogue/block/block_epilogue_kda_fwdh_update_arch22.hpp new file mode 100644 index 0000000000..7d811b7cb8 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/epilogue/block/block_epilogue_kda_fwdh_update_arch22.hpp @@ -0,0 +1,526 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_KDA_FWDH_UPDATE_ARCH22_HPP +#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_KDA_FWDH_UPDATE_ARCH22_HPP +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "../kda_fwd_h_epilogue_policies.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/kda_matrix_coord.hpp" +#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" + +namespace Catlass::Epilogue::Block { + +template +class BlockEpilogue { + static constexpr bool kGated = KGatedTag::value; + static constexpr bool scalarGated = KGatedTag::scalarGated; + static constexpr bool useExp2 = KGatedTag::useExp2; + static constexpr float LN2 = 0.6931471805599453f; + +public: + using DispatchPolicy = EpilogueAtlasKdaFwdHUpdate; + using ArchTag = typename DispatchPolicy::ArchTag; + + using HElementOutput = typename HOutputType_::Element; + using GElementInput = typename GInputType_::Element; + using HElementInput = typename HInputType_::Element; + using HUpdateElementInput = typename HUpdateInputType_::Element; + using FinalStateElement = typename FinalStateType_::Element; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource) + { + constexpr uint32_t CALC_BUF_OFFSET = 0; + constexpr uint32_t PING_BUF_0_OFFSET = 32 * 1024; + constexpr uint32_t PING_BUF_1_OFFSET = 48 * 1024; + constexpr uint32_t PING_BUF_2_OFFSET = 64 * 1024; + constexpr uint32_t PING_BUF_3_OFFSET = 80 * 1024; + constexpr uint32_t PONG_BUF_0_OFFSET = 96 * 1024; + constexpr uint32_t PONG_BUF_1_OFFSET = 112 * 1024; + constexpr uint32_t PONG_BUF_2_OFFSET = 128 * 1024; + constexpr uint32_t PONG_BUF_3_OFFSET = 144 * 1024; + constexpr uint32_t PING_G_BUF_OFFSET = 160 * 1024; + constexpr uint32_t PONG_G_BUF_OFFSET = 161 * 1024; + constexpr uint32_t PING_G_SUB_BUF_OFFSET = 162 * 1024; + constexpr uint32_t PONG_G_SUB_BUF_OFFSET = 163 * 1024; + constexpr uint32_t PING_G_INPUT_BUF_OFFSET = 164 * 1024; + constexpr uint32_t PONG_G_INPUT_BUF_OFFSET = 165 * 1024; + constexpr uint32_t SHARE_BUF_OFFSET = 166 * 1024; + + calcUbTensor = resource.ubBuf.template GetBufferByByte(CALC_BUF_OFFSET); + + hUpdateUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_0_OFFSET); + gkBroadcastUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_1_OFFSET); + hUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_3_OFFSET); + finalOutputUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_3_OFFSET); + glastUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_INPUT_BUF_OFFSET); + + hUpdateUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_0_OFFSET); + gkBroadcastUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_1_OFFSET); + hUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_3_OFFSET); + finalOutputUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_3_OFFSET); + glastUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_INPUT_BUF_OFFSET); + + if constexpr (kGated) { + gkLastUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_SUB_BUF_OFFSET); + gkLastUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_SUB_BUF_OFFSET); + gkInputUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_INPUT_BUF_OFFSET); + gkInputUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_INPUT_BUF_OFFSET); + shareBufferGk_ = resource.ubBuf.template GetBufferByByte(SHARE_BUF_OFFSET); + } + } + + CATLASS_DEVICE + ~BlockEpilogue() {} + + template + CATLASS_DEVICE void CopyGmToUb(AscendC::LocalTensor dst, AscendC::GlobalTensor src, uint32_t rows, + uint32_t cols, uint32_t srcStride) + { + if (cols == srcStride) { + AscendC::DataCopy(dst, src, rows * cols); + return; + } + AscendC::DataCopyExtParams copyParams{static_cast(rows), + static_cast(cols * sizeof(Element)), + static_cast((srcStride - cols) * sizeof(Element)), 0, 0}; + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyPad(dst, src, copyParams, padParams); + } + + template + CATLASS_DEVICE void CopyUbToGm(AscendC::GlobalTensor dst, AscendC::LocalTensor src, uint32_t rows, + uint32_t cols, uint32_t dstStride) + { + if (cols == dstStride) { + AscendC::DataCopy(dst, src, rows * cols); + return; + } + AscendC::DataCopyExtParams copyParams{static_cast(rows), + static_cast(cols * sizeof(Element)), 0, + static_cast((dstStride - cols) * sizeof(Element)), 0}; + AscendC::DataCopyPad(dst, src, copyParams); + } + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor hOutput, AscendC::GlobalTensor finalState, + AscendC::GlobalTensor gInput, AscendC::GlobalTensor hInput, + AscendC::GlobalTensor hUpdateInput, AscendC::GlobalTensor gkInput, + AscendC::GlobalTensor initialState, uint32_t chunkSize, uint32_t kHeadDim, + uint32_t vBlockDim, uint32_t vHeadDim, Arch::CrossCoreFlag cube2Done, bool isInitialState, + bool isFinalState, bool storeFinalState, bool useInitialState, bool isPing, bool cube2AlreadyWaited) + { + static constexpr uint32_t ROW_TILE = 16; + uint32_t mActual = kHeadDim; + uint32_t nActual = vBlockDim; + uint32_t outputStride = vHeadDim; + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t rowsPerSubBlock = CeilDiv(mActual, subBlockNum); + uint32_t rowBegin = subBlockIdx * rowsPerSubBlock; + uint32_t rowEnd = rowBegin + rowsPerSubBlock; + if (rowEnd > mActual) { + rowEnd = mActual; + } + if (rowBegin >= mActual) { + if (!cube2AlreadyWaited) { + Arch::CrossCoreWaitFlag(cube2Done); + } + return; + } + + AscendC::ResetMask(); + + AscendC::GlobalTensor gInputThisSubBlock = gInput; + + uint32_t pingpongFlag = isPing ? 0 : pongBaseEvent; + AscendC::LocalTensor hUpdateUbTensor = isPing ? hUpdateUbTensor_ping : hUpdateUbTensor_pong; + AscendC::LocalTensor gkBroadcastUbTensor = isPing ? gkBroadcastUbTensor_ping : gkBroadcastUbTensor_pong; + AscendC::LocalTensor hUbTensor = isPing ? hUbTensor_ping : hUbTensor_pong; + AscendC::LocalTensor finalOutputUbTensor = + isPing ? finalOutputUbTensor_ping : finalOutputUbTensor_pong; + AscendC::LocalTensor glastUbTensor = isPing ? glastUbTensor_ping : glastUbTensor_pong; + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + bool useFp32StateUpdate = + storeFinalState && std::is_same::value && (!isInitialState || useInitialState); + + float muls = 1.0f; + if constexpr (scalarGated) { + GElementInput gLastVal = gInputThisSubBlock.GetValue(chunkSize - 1); + float gLastFloat = 0.0f; + if constexpr (std::is_same::value) { + gLastFloat = gLastVal; + } else if constexpr (std::is_same::value) { + gLastFloat = (float)gLastVal; + } else if constexpr (std::is_same::value) { + gLastFloat = AscendC::ToFloat(gLastVal); + } + glastUbTensor.SetValue(0, gLastFloat); + + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + if constexpr (useExp2) { + AscendC::Muls(glastUbTensor, glastUbTensor, LN2, 1); + AscendC::PipeBarrier(); + } + AscendC::Exp(glastUbTensor, glastUbTensor, 1); + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + muls = glastUbTensor.GetValue(0); + } + if constexpr (kGated) { + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + } + if constexpr (scalarGated) { + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + } + + if (nActual <= 128 && nActual == outputStride) { + uint32_t mActualThisSubBlock = rowEnd - rowBegin; + AscendC::GlobalTensor hOutputThisSubBlock = hOutput[rowBegin * outputStride]; + AscendC::GlobalTensor hInputThisSubBlock = hInput[rowBegin * outputStride]; + AscendC::GlobalTensor hUpdateInputThisSubBlock = hUpdateInput[rowBegin * nActual]; + AscendC::GlobalTensor finalStateThisSubBlock = finalState[rowBegin * outputStride]; + + if (storeFinalState && isInitialState && std::is_same::value) { + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + } else { + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + } + if constexpr (std::is_same::value) { + if (useFp32StateUpdate) { + if (isInitialState) { + AscendC::DataCopy(calcUbTensor, initialState[rowBegin * outputStride], + mActualThisSubBlock * nActual); + } else { + AscendC::DataCopy(calcUbTensor, finalStateThisSubBlock, mActualThisSubBlock * nActual); + } + } else { + AscendC::DataCopy(hUbTensor, hInputThisSubBlock, mActualThisSubBlock * nActual); + } + } else { + AscendC::DataCopy(hUbTensor, hInputThisSubBlock, mActualThisSubBlock * nActual); + } + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + + if (!useFp32StateUpdate) { + AscendC::Cast(calcUbTensor, hUbTensor, AscendC::RoundMode::CAST_NONE, mActualThisSubBlock * nActual); + AscendC::PipeBarrier(); + } + AscendC::Muls(calcUbTensor, calcUbTensor, muls, mActualThisSubBlock * nActual); + AscendC::PipeBarrier(); + + if constexpr (kGated) { + AscendC::GlobalTensor gkLastInput = gkInput[(chunkSize - 1) * kHeadDim + rowBegin]; + AscendC::LocalTensor gkLastUbTensor = isPing ? gkLastUbTensor_ping : gkLastUbTensor_pong; + AscendC::LocalTensor gkInputUbTensor = + isPing ? gkInputUbTensor_ping : gkInputUbTensor_pong; + + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + if constexpr (std::is_same::value) { + AscendC::DataCopy(gkLastUbTensor, gkLastInput, mActualThisSubBlock); + } else { + AscendC::DataCopy(gkInputUbTensor, gkLastInput, mActualThisSubBlock); + } + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + if constexpr (!std::is_same::value) { + AscendC::Cast(gkLastUbTensor, gkInputUbTensor, AscendC::RoundMode::CAST_NONE, mActualThisSubBlock); + } + AscendC::PipeBarrier(); + AscendC::Muls(gkLastUbTensor, gkLastUbTensor, LN2, mActualThisSubBlock); + AscendC::PipeBarrier(); + AscendC::Exp(gkLastUbTensor, gkLastUbTensor, mActualThisSubBlock); + AscendC::PipeBarrier(); + + uint32_t gkBrcReptime = (mActualThisSubBlock + 8 - 1) / 8; + uint32_t dstShapeGk[2] = {gkBrcReptime * 8, nActual}; + uint32_t srcShapeGk[2] = {gkBrcReptime * 8, 1}; + AscendC::Broadcast(hUpdateUbTensor, gkLastUbTensor, dstShapeGk, srcShapeGk, + shareBufferGk_); + AscendC::PipeBarrier(); + AscendC::Mul(calcUbTensor, calcUbTensor, hUpdateUbTensor, mActualThisSubBlock * nActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + } + + if (!cube2AlreadyWaited) { + Arch::CrossCoreWaitFlag(cube2Done); + } + + if constexpr (kGated) { + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + } + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::DataCopy(hUpdateUbTensor, hUpdateInputThisSubBlock, mActualThisSubBlock * nActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::Add(hUpdateUbTensor, calcUbTensor, hUpdateUbTensor, mActualThisSubBlock * nActual); + AscendC::PipeBarrier(); + if (storeFinalState && isFinalState && std::is_same::value) { + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + } + + if constexpr (std::is_same::value) { + if (storeFinalState) { + if (!isFinalState) { + AscendC::Cast(hUbTensor, hUpdateUbTensor, AscendC::RoundMode::CAST_RINT, + mActualThisSubBlock * nActual); + AscendC::PipeBarrier(); + } + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::DataCopy(finalStateThisSubBlock, hUpdateUbTensor, mActualThisSubBlock * nActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + if (!isFinalState) { + AscendC::DataCopy(hOutputThisSubBlock, hUbTensor, mActualThisSubBlock * nActual); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + } + } else { + AscendC::Cast(hUbTensor, hUpdateUbTensor, AscendC::RoundMode::CAST_RINT, + mActualThisSubBlock * nActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + AscendC::DataCopy(hOutputThisSubBlock, hUbTensor, mActualThisSubBlock * nActual); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + } + } else { + if (storeFinalState && isFinalState) { + AscendC::Cast(finalOutputUbTensor, hUpdateUbTensor, AscendC::RoundMode::CAST_RINT, + mActualThisSubBlock * nActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + AscendC::DataCopy(finalStateThisSubBlock, finalOutputUbTensor, mActualThisSubBlock * nActual); + } else { + AscendC::Cast(hUbTensor, hUpdateUbTensor, AscendC::RoundMode::CAST_RINT, + mActualThisSubBlock * nActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + AscendC::DataCopy(hOutputThisSubBlock, hUbTensor, mActualThisSubBlock * nActual); + } + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + } + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + return; + } + + if (!cube2AlreadyWaited) { + Arch::CrossCoreWaitFlag(cube2Done); + } + + bool waitHFromV = storeFinalState && isInitialState && std::is_same::value; + bool waitUpdateFromMte3 = false; + uint32_t updateReadyEvent = EVENT_ID3 + pingpongFlag; + bool waitGkFromScalar = true; + for (uint32_t rowStart = rowBegin; rowStart < rowEnd; rowStart += ROW_TILE) { + uint32_t rowsThisTile = rowEnd - rowStart; + if (rowsThisTile > ROW_TILE) { + rowsThisTile = ROW_TILE; + } + + AscendC::GlobalTensor hOutputThisTile = hOutput[rowStart * outputStride]; + AscendC::GlobalTensor hInputThisTile = hInput[rowStart * outputStride]; + AscendC::GlobalTensor hUpdateInputThisTile = hUpdateInput[rowStart * nActual]; + AscendC::GlobalTensor finalStateThisTile = finalState[rowStart * outputStride]; + + if (waitHFromV) { + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + } else { + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + } + if constexpr (std::is_same::value) { + if (useFp32StateUpdate) { + if (isInitialState) { + CopyGmToUb(calcUbTensor, initialState[rowStart * outputStride], rowsThisTile, nActual, + outputStride); + } else { + CopyGmToUb(calcUbTensor, finalStateThisTile, rowsThisTile, nActual, outputStride); + } + } else { + CopyGmToUb(hUbTensor, hInputThisTile, rowsThisTile, nActual, outputStride); + } + } else { + CopyGmToUb(hUbTensor, hInputThisTile, rowsThisTile, nActual, outputStride); + } + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + + if (!useFp32StateUpdate) { + AscendC::Cast(calcUbTensor, hUbTensor, AscendC::RoundMode::CAST_NONE, rowsThisTile * nActual); + AscendC::PipeBarrier(); + } + AscendC::Muls(calcUbTensor, calcUbTensor, muls, rowsThisTile * nActual); + AscendC::PipeBarrier(); + + if constexpr (kGated) { + AscendC::GlobalTensor gkLastInput = gkInput[(chunkSize - 1) * kHeadDim + rowStart]; + AscendC::LocalTensor gkLastUbTensor = isPing ? gkLastUbTensor_ping : gkLastUbTensor_pong; + AscendC::LocalTensor gkInputUbTensor = + isPing ? gkInputUbTensor_ping : gkInputUbTensor_pong; + + if (waitGkFromScalar) { + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + waitGkFromScalar = false; + } else { + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + } + if constexpr (std::is_same::value) { + AscendC::DataCopy(gkLastUbTensor, gkLastInput, rowsThisTile); + } else { + AscendC::DataCopy(gkInputUbTensor, gkLastInput, rowsThisTile); + } + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + if constexpr (!std::is_same::value) { + AscendC::Cast(gkLastUbTensor, gkInputUbTensor, AscendC::RoundMode::CAST_NONE, rowsThisTile); + } + AscendC::PipeBarrier(); + AscendC::Muls(gkLastUbTensor, gkLastUbTensor, LN2, rowsThisTile); + AscendC::PipeBarrier(); + AscendC::Exp(gkLastUbTensor, gkLastUbTensor, rowsThisTile); + AscendC::PipeBarrier(); + + uint32_t gkBrcReptime = (rowsThisTile + 8 - 1) / 8; + uint32_t dstShapeGk[2] = {gkBrcReptime * 8, nActual}; + uint32_t srcShapeGk[2] = {gkBrcReptime * 8, 1}; + AscendC::Broadcast(gkBroadcastUbTensor, gkLastUbTensor, dstShapeGk, srcShapeGk, + shareBufferGk_); + AscendC::PipeBarrier(); + AscendC::Mul(calcUbTensor, calcUbTensor, gkBroadcastUbTensor, rowsThisTile * nActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + } + + if (waitUpdateFromMte3) { + AscendC::WaitFlag(updateReadyEvent); + } else { + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + } + AscendC::DataCopy(hUpdateUbTensor, hUpdateInputThisTile, rowsThisTile * nActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::Add(hUpdateUbTensor, calcUbTensor, hUpdateUbTensor, rowsThisTile * nActual); + AscendC::PipeBarrier(); + if (storeFinalState && isFinalState && std::is_same::value) { + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + waitHFromV = true; + } else { + waitHFromV = false; + } + + if constexpr (std::is_same::value) { + if (storeFinalState) { + if (!isFinalState) { + AscendC::Cast(hUbTensor, hUpdateUbTensor, AscendC::RoundMode::CAST_RINT, + rowsThisTile * nActual); + AscendC::PipeBarrier(); + } + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + CopyUbToGm(finalStateThisTile, hUpdateUbTensor, rowsThisTile, nActual, outputStride); + // A2 wide-V tiles reuse this UB in the next iteration; retire MTE3 before either writer advances. + AscendC::PipeBarrier(); + AscendC::SetFlag(updateReadyEvent); + AscendC::SetFlag(updateReadyEvent); + AscendC::WaitFlag(updateReadyEvent); + waitUpdateFromMte3 = true; + if (!isFinalState) { + CopyUbToGm(hOutputThisTile, hUbTensor, rowsThisTile, nActual, outputStride); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + } + } else { + AscendC::Cast(hUbTensor, hUpdateUbTensor, AscendC::RoundMode::CAST_RINT, rowsThisTile * nActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + CopyUbToGm(hOutputThisTile, hUbTensor, rowsThisTile, nActual, outputStride); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + waitUpdateFromMte3 = false; + } + } else { + if (storeFinalState && isFinalState) { + AscendC::Cast(finalOutputUbTensor, hUpdateUbTensor, AscendC::RoundMode::CAST_RINT, + rowsThisTile * nActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + CopyUbToGm(finalStateThisTile, finalOutputUbTensor, rowsThisTile, nActual, outputStride); + } else { + AscendC::Cast(hUbTensor, hUpdateUbTensor, AscendC::RoundMode::CAST_RINT, rowsThisTile * nActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + CopyUbToGm(hOutputThisTile, hUbTensor, rowsThisTile, nActual, outputStride); + } + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + waitUpdateFromMte3 = false; + } + } + if (storeFinalState && std::is_same::value) { + AscendC::WaitFlag(updateReadyEvent); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + if (!isFinalState) { + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + } + } else { + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + } + if constexpr (kGated) { + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + } + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + } + +private: + uint32_t pongBaseEvent = 4; + + AscendC::LocalTensor calcUbTensor; + + AscendC::LocalTensor hUpdateUbTensor_ping; + AscendC::LocalTensor gkBroadcastUbTensor_ping; + AscendC::LocalTensor hUbTensor_ping; + AscendC::LocalTensor finalOutputUbTensor_ping; + AscendC::LocalTensor glastUbTensor_ping; + + AscendC::LocalTensor hUpdateUbTensor_pong; + AscendC::LocalTensor gkBroadcastUbTensor_pong; + AscendC::LocalTensor hUbTensor_pong; + AscendC::LocalTensor finalOutputUbTensor_pong; + AscendC::LocalTensor glastUbTensor_pong; + + AscendC::LocalTensor gkLastUbTensor_ping; + AscendC::LocalTensor gkLastUbTensor_pong; + AscendC::LocalTensor gkInputUbTensor_ping; + AscendC::LocalTensor gkInputUbTensor_pong; + AscendC::LocalTensor shareBufferGk_; +}; +} // namespace Catlass::Epilogue::Block + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/epilogue/block/block_epilogue_kda_fwdh_vnew_arch22.hpp b/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/epilogue/block/block_epilogue_kda_fwdh_vnew_arch22.hpp new file mode 100644 index 0000000000..7634c785ad --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/epilogue/block/block_epilogue_kda_fwdh_vnew_arch22.hpp @@ -0,0 +1,402 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_KDA_FWDH_VNEW_ARCH22_HPP +#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_KDA_FWDH_VNEW_ARCH22_HPP +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "../kda_fwd_h_epilogue_policies.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/kda_matrix_coord.hpp" +#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" + +namespace Catlass::Epilogue::Block { + +template +class BlockEpilogue { + static constexpr bool scalarGated = KGatedTag::scalarGated; + static constexpr bool useExp2 = KGatedTag::useExp2; + static constexpr float LN2 = 0.6931471805599453f; + +public: + // Type aliases + using DispatchPolicy = EpilogueAtlasKdaFwdHVnew; + using ArchTag = typename DispatchPolicy::ArchTag; + + using VElementOutput = typename VOutputType_::Element; + using GElementInput = typename GInputType_::Element; + using UElementInput = typename UInputType_::Element; + using WSElementInput = typename WSInputType_::Element; + using FinalStateElement = typename FinalStateType_::Element; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource) + { + constexpr uint32_t CALC_BUF_OFFSET = 0; + constexpr uint32_t PING_BUF_0_OFFSET = 32 * 1024; + constexpr uint32_t PING_BUF_1_OFFSET = 48 * 1024; + constexpr uint32_t PING_BUF_2_OFFSET = 64 * 1024; + constexpr uint32_t PING_BUF_3_OFFSET = 80 * 1024; + constexpr uint32_t PONG_BUF_0_OFFSET = 96 * 1024; + constexpr uint32_t PONG_BUF_1_OFFSET = 112 * 1024; + constexpr uint32_t PONG_BUF_2_OFFSET = 128 * 1024; + constexpr uint32_t PONG_BUF_3_OFFSET = 144 * 1024; + constexpr uint32_t PING_G_BUF_OFFSET = 160 * 1024; + constexpr uint32_t PONG_G_BUF_OFFSET = 161 * 1024; + constexpr uint32_t PING_G_SUB_BUF_OFFSET = 162 * 1024; + constexpr uint32_t PONG_G_SUB_BUF_OFFSET = 163 * 1024; + constexpr uint32_t PING_G_INPUT_BUF_OFFSET = 164 * 1024; + constexpr uint32_t PONG_G_INPUT_BUF_OFFSET = 165 * 1024; + constexpr uint32_t SHARE_BUF_OFFSET = 166 * 1024; + + calcUbTensor = resource.ubBuf.template GetBufferByByte(CALC_BUF_OFFSET); + + uUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_2_OFFSET); + wsUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_0_OFFSET); + gUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_BUF_OFFSET); + gLastUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_SUB_BUF_OFFSET); + gInputUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_SUB_BUF_OFFSET); + vNewOutputUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_2_OFFSET); + vNewDecayUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_2_OFFSET); + + uUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_2_OFFSET); + wsUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_0_OFFSET); + gUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_BUF_OFFSET); + gLastUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_SUB_BUF_OFFSET); + gInputUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_SUB_BUF_OFFSET); + vNewOutputUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_2_OFFSET); + vNewDecayUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_2_OFFSET); + + shareBuffer_ = resource.ubBuf.template GetBufferByByte(SHARE_BUF_OFFSET); + } + + CATLASS_DEVICE + ~BlockEpilogue() {} + + template + CATLASS_DEVICE void CopyGmToUb(AscendC::LocalTensor dst, AscendC::GlobalTensor src, uint32_t rows, + uint32_t cols, uint32_t srcStride) + { + if (cols == srcStride) { + AscendC::DataCopy(dst, src, rows * cols); + return; + } + AscendC::DataCopyExtParams copyParams{static_cast(rows), + static_cast(cols * sizeof(Element)), + static_cast((srcStride - cols) * sizeof(Element)), 0, 0}; + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyPad(dst, src, copyParams, padParams); + } + + template + CATLASS_DEVICE void CopyUbToGm(AscendC::GlobalTensor dst, AscendC::LocalTensor src, uint32_t rows, + uint32_t cols, uint32_t dstStride) + { + if (cols == dstStride) { + AscendC::DataCopy(dst, src, rows * cols); + return; + } + AscendC::DataCopyExtParams copyParams{static_cast(rows), + static_cast(cols * sizeof(Element)), 0, + static_cast((dstStride - cols) * sizeof(Element)), 0}; + AscendC::DataCopyPad(dst, src, copyParams); + } + + CATLASS_DEVICE + void PrepareG(AscendC::LocalTensor gUbTensor, AscendC::LocalTensor gLastUbTensor, + AscendC::LocalTensor gInputUbTensor, + AscendC::GlobalTensor gInputThisSubBlock, uint32_t mActual, uint32_t pingpongFlag) + { + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + if constexpr (!scalarGated) { + AscendC::Duplicate(gUbTensor, 1.0f, mActual); + AscendC::PipeBarrier(); + return; + } + if (mActual == 1) { + AscendC::Duplicate(gUbTensor, 1.0f, 1); + AscendC::PipeBarrier(); + return; + } + if constexpr (std::is_same::value) { + AscendC::DataCopyParams gUbParams{1, (uint16_t)(mActual * sizeof(float)), 0, 0}; + AscendC::DataCopyPadParams gUbPadParams{false, 0, 0, 0}; + AscendC::DataCopyPad(gUbTensor, gInputThisSubBlock, gUbParams, gUbPadParams); + } else { + AscendC::DataCopyParams gUbParams{1, (uint16_t)(mActual * sizeof(GElementInput)), 0, 0}; + AscendC::DataCopyPadParams gUbPadParams{false, 0, 0, 0}; + AscendC::DataCopyPad(gInputUbTensor, gInputThisSubBlock, gUbParams, gUbPadParams); + } + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + if constexpr (!std::is_same::value) { + AscendC::Cast(gUbTensor, gInputUbTensor, AscendC::RoundMode::CAST_NONE, mActual); + } + + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + float inputVal = gUbTensor.GetValue(mActual - 1); + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + + AscendC::PipeBarrier(); + AscendC::Duplicate(gLastUbTensor, inputVal, mActual); + AscendC::PipeBarrier(); + + AscendC::Sub(gUbTensor, gLastUbTensor, gUbTensor, mActual); + AscendC::PipeBarrier(); + if constexpr (useExp2) { + AscendC::Muls(gUbTensor, gUbTensor, LN2, mActual); + AscendC::PipeBarrier(); + } + AscendC::Exp(gUbTensor, gUbTensor, mActual); + AscendC::PipeBarrier(); + } + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor vnewOutput, + AscendC::GlobalTensor vnewdecayOutput, AscendC::GlobalTensor gInput, + AscendC::GlobalTensor uInput, AscendC::GlobalTensor wsInput, + uint32_t chunkSize, uint32_t vBlockDim, uint32_t vHeadDim, Arch::CrossCoreFlag cube1Done, + Arch::CrossCoreFlag vec1Done, bool isInitialState, bool isFinalState, bool storeFinalState, + bool waitWsFromMte3, bool isPing, bool cube1AlreadyWaited) + { + static constexpr uint32_t ROW_TILE = 16; + uint32_t mActual = chunkSize; + uint32_t nvActual = vBlockDim; + uint32_t inputStride = vHeadDim; + + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t rowsPerSubBlock = CeilDiv(mActual, subBlockNum); + uint32_t rowBegin = subBlockIdx * rowsPerSubBlock; + uint32_t rowEnd = rowBegin + rowsPerSubBlock; + if (rowEnd > mActual) { + rowEnd = mActual; + } + uint32_t pingpongFlag = isPing ? 0 : pongBaseEvent; + if (rowBegin >= mActual) { + if (!cube1AlreadyWaited) { + Arch::CrossCoreWaitFlag(cube1Done); + } + // A zero-row AIV lane still owns the EVENT0 hand-off consumed by V2. + if (waitWsFromMte3) { + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + } + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vec1Done); + return; + } + AscendC::ResetMask(); + + AscendC::GlobalTensor gInputThisSubBlock = gInput; + + AscendC::LocalTensor uUbTensor = isPing ? uUbTensor_ping : uUbTensor_pong; + AscendC::LocalTensor wsUbTensor = isPing ? wsUbTensor_ping : wsUbTensor_pong; + AscendC::LocalTensor gUbTensor = isPing ? gUbTensor_ping : gUbTensor_pong; + AscendC::LocalTensor gLastUbTensor = isPing ? gLastUbTensor_ping : gLastUbTensor_pong; + AscendC::LocalTensor gInputUbTensor = isPing ? gInputUbTensor_ping : gInputUbTensor_pong; + AscendC::LocalTensor vNewOutputUbTensor = + isPing ? vNewOutputUbTensor_ping : vNewOutputUbTensor_pong; + AscendC::LocalTensor vNewDecayUbTensor = + isPing ? vNewDecayUbTensor_ping : vNewDecayUbTensor_pong; + + if (nvActual <= 128 && nvActual == inputStride) { + uint32_t mActualThisSubBlock = rowEnd - rowBegin; + uint32_t gbrcRealStart = rowBegin & ~7; + uint32_t gbrcEffStart = rowBegin - gbrcRealStart; + uint32_t gbrcRealProcess = gbrcEffStart + mActualThisSubBlock; + uint32_t dstShape_[2] = {gbrcRealProcess, nvActual}; + uint32_t srcShape_[2] = {gbrcRealProcess, 1}; + + AscendC::GlobalTensor vnewOutputThisSubBlock = vnewOutput[rowBegin * inputStride]; + AscendC::GlobalTensor vnewdecayOutputThisSubBlock = vnewdecayOutput[rowBegin * nvActual]; + AscendC::GlobalTensor uInputThisSubBlock = uInput[rowBegin * inputStride]; + AscendC::GlobalTensor wsInputThisSubBlock = wsInput[rowBegin * nvActual]; + + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::DataCopy(uUbTensor, uInputThisSubBlock, mActualThisSubBlock * nvActual); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::Cast(calcUbTensor, uUbTensor, AscendC::RoundMode::CAST_NONE, mActualThisSubBlock * nvActual); + AscendC::PipeBarrier(); + + if constexpr (scalarGated) { + PrepareG(gUbTensor, gLastUbTensor, gInputUbTensor, gInputThisSubBlock, mActual, pingpongFlag); + } else { + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + } + if (!cube1AlreadyWaited) { + Arch::CrossCoreWaitFlag(cube1Done); + } + + if (waitWsFromMte3) { + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + } else { + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + } + AscendC::DataCopy(wsUbTensor, wsInputThisSubBlock, mActualThisSubBlock * nvActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + + AscendC::Sub(wsUbTensor, calcUbTensor, wsUbTensor, mActualThisSubBlock * nvActual); + AscendC::PipeBarrier(); + + uint32_t decayOffset = 0; + if constexpr (scalarGated) { + AscendC::Broadcast(calcUbTensor, gUbTensor[gbrcRealStart], dstShape_, srcShape_, + shareBuffer_); + AscendC::PipeBarrier(); + AscendC::Mul(calcUbTensor[gbrcEffStart * nvActual], wsUbTensor, calcUbTensor[gbrcEffStart * nvActual], + mActualThisSubBlock * nvActual); + decayOffset = gbrcEffStart * nvActual; + } else { + AscendC::Adds(calcUbTensor, wsUbTensor, 0.0f, mActualThisSubBlock * nvActual); + } + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + + AscendC::Cast(vNewDecayUbTensor, calcUbTensor[decayOffset], AscendC::RoundMode::CAST_RINT, + mActualThisSubBlock * nvActual); + AscendC::PipeBarrier(); + + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::DataCopy(vnewdecayOutputThisSubBlock, vNewDecayUbTensor, mActualThisSubBlock * nvActual); + + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::Cast(vNewOutputUbTensor, wsUbTensor, AscendC::RoundMode::CAST_RINT, + mActualThisSubBlock * nvActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::DataCopy(vnewOutputThisSubBlock, vNewOutputUbTensor, mActualThisSubBlock * nvActual); + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vec1Done); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + + return; + } + + if constexpr (scalarGated) { + PrepareG(gUbTensor, gLastUbTensor, gInputUbTensor, gInputThisSubBlock, mActual, pingpongFlag); + } else { + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + } + if (!cube1AlreadyWaited) { + Arch::CrossCoreWaitFlag(cube1Done); + } + + bool waitWsThisTileFromMte3 = waitWsFromMte3; + for (uint32_t rowStart = rowBegin; rowStart < rowEnd;) { + uint32_t alignExtra = rowStart & 7; + uint32_t maxRowsThisTile = ROW_TILE - alignExtra; + uint32_t rowsThisTile = rowEnd - rowStart; + if (rowsThisTile > maxRowsThisTile) { + rowsThisTile = maxRowsThisTile; + } + uint32_t gbrcRealStart = rowStart & ~7; + uint32_t gbrcRealProcess = alignExtra + rowsThisTile; + uint32_t dstShape_[2] = {gbrcRealProcess, nvActual}; + uint32_t srcShape_[2] = {gbrcRealProcess, 1}; + + AscendC::GlobalTensor vnewOutputThisTile = vnewOutput[rowStart * inputStride]; + AscendC::GlobalTensor vnewdecayOutputThisTile = vnewdecayOutput[rowStart * nvActual]; + AscendC::GlobalTensor uInputThisTile = uInput[rowStart * inputStride]; + AscendC::GlobalTensor wsInputThisTile = wsInput[rowStart * nvActual]; + + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + CopyGmToUb(uUbTensor, uInputThisTile, rowsThisTile, nvActual, inputStride); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::Cast(calcUbTensor, uUbTensor, AscendC::RoundMode::CAST_NONE, rowsThisTile * nvActual); + AscendC::PipeBarrier(); + + if (waitWsThisTileFromMte3) { + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + } else { + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + } + AscendC::DataCopy(wsUbTensor, wsInputThisTile, rowsThisTile * nvActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + waitWsThisTileFromMte3 = false; + + AscendC::Sub(wsUbTensor, calcUbTensor, wsUbTensor, rowsThisTile * nvActual); + AscendC::PipeBarrier(); + + uint32_t decayOffset = 0; + if constexpr (scalarGated) { + AscendC::Broadcast(calcUbTensor, gUbTensor[gbrcRealStart], dstShape_, srcShape_, + shareBuffer_); + AscendC::PipeBarrier(); + AscendC::Mul(calcUbTensor[alignExtra * nvActual], wsUbTensor, calcUbTensor[alignExtra * nvActual], + rowsThisTile * nvActual); + decayOffset = alignExtra * nvActual; + } else { + AscendC::Adds(calcUbTensor, wsUbTensor, 0.0f, rowsThisTile * nvActual); + } + AscendC::PipeBarrier(); + + AscendC::Cast(vNewDecayUbTensor, calcUbTensor[decayOffset], AscendC::RoundMode::CAST_RINT, + rowsThisTile * nvActual); + AscendC::PipeBarrier(); + + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::DataCopy(vnewdecayOutputThisTile, vNewDecayUbTensor, rowsThisTile * nvActual); + + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::Cast(vNewOutputUbTensor, wsUbTensor, AscendC::RoundMode::CAST_RINT, rowsThisTile * nvActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + CopyUbToGm(vnewOutputThisTile, vNewOutputUbTensor, rowsThisTile, nvActual, inputStride); + if (rowStart + rowsThisTile >= rowEnd) { + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vec1Done); + } + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + rowStart += rowsThisTile; + } + + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + } + +private: + uint32_t pongBaseEvent = 4; + + AscendC::LocalTensor calcUbTensor; + + AscendC::LocalTensor uUbTensor_ping; + AscendC::LocalTensor wsUbTensor_ping; + AscendC::LocalTensor gUbTensor_ping; + AscendC::LocalTensor gLastUbTensor_ping; + AscendC::LocalTensor gInputUbTensor_ping; + AscendC::LocalTensor vNewOutputUbTensor_ping; + AscendC::LocalTensor vNewDecayUbTensor_ping; + + AscendC::LocalTensor uUbTensor_pong; + AscendC::LocalTensor wsUbTensor_pong; + AscendC::LocalTensor gUbTensor_pong; + AscendC::LocalTensor gLastUbTensor_pong; + AscendC::LocalTensor gInputUbTensor_pong; + AscendC::LocalTensor vNewOutputUbTensor_pong; + AscendC::LocalTensor vNewDecayUbTensor_pong; + + AscendC::LocalTensor shareBuffer_; +}; +} // namespace Catlass::Epilogue::Block + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/epilogue/kda_fwd_h_epilogue_policies.hpp b/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/epilogue/kda_fwd_h_epilogue_policies.hpp new file mode 100644 index 0000000000..1dd37a521c --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/epilogue/kda_fwd_h_epilogue_policies.hpp @@ -0,0 +1,28 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_KDA_FWD_H_EPILOGUE_POLICIES_HPP +#define CATLASS_EPILOGUE_KDA_FWD_H_EPILOGUE_POLICIES_HPP + +#include "catlass/kda_catlass.hpp" + +namespace Catlass::Epilogue { + +struct EpilogueAtlasKdaFwdHVnew { + using ArchTag = Arch::AtlasA2; +}; + +struct EpilogueAtlasKdaFwdHUpdate { + using ArchTag = Arch::AtlasA2; +}; + +} // namespace Catlass::Epilogue + +#endif // CATLASS_EPILOGUE_KDA_FWD_H_EPILOGUE_POLICIES_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/gemm/block/block_scheduler_kda_fwd_h_arch22.hpp b/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/gemm/block/block_scheduler_kda_fwd_h_arch22.hpp new file mode 100644 index 0000000000..961ce442bc --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/gemm/block/block_scheduler_kda_fwd_h_arch22.hpp @@ -0,0 +1,445 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#include "catlass/kda_gemm_coord.hpp" +#include "../../../../chunk_kda_fwd_h_runtime_tiling.h" +using namespace Catlass; + +#ifndef CATLASS_GEMM_SCHEDULER_KDA_FWD_H_HPP +#define CATLASS_GEMM_SCHEDULER_KDA_FWD_H_HPP + +// constexpr uint32_t PING_PONG_STAGES = 1; +constexpr uint32_t PING_PONG_STAGES = 2; +constexpr uint32_t BYTE_SIZE_16_BIT = 2; +constexpr uint32_t BYTES_PER_C0 = 32; +constexpr uint32_t BYTE_SIZE_PER_REPEAT = 256; +constexpr uint32_t SIZE_16_NUM_PER_C0 = BYTES_PER_C0 / BYTE_SIZE_16_BIT; +constexpr uint32_t FLOAT_NUM_PER_REPEAT = BYTE_SIZE_PER_REPEAT / sizeof(float); +constexpr uint32_t NZ_BLOCK_SIZE = 16; + +template +CATLASS_DEVICE T AlignUp(T a, T b) +{ + return (b == 0) ? 0 : (a + b - 1) / b * b; +} + +template +CATLASS_DEVICE T Min(T a, T b) +{ + return (a > b) ? b : a; +} + +template +CATLASS_DEVICE T Max(T a, T b) +{ + return (a > b) ? a : b; +} + +namespace Catlass::Gemm::Block { + +struct KdaFwdHOffsets { + uint32_t hSrcOffset; + uint32_t hDstOffset; + uint32_t uvOffset; + uint32_t wkOffset; + uint32_t wOffset; + uint32_t gOffset; + uint32_t gkOffset; + uint32_t hWorkOffset; + uint32_t vWorkOffset; + uint32_t kDecayWorkOffset; + uint32_t vBlockOffset; + uint32_t vBlockDim; + uint32_t initialStateOffset; + uint32_t finalStateOffset; + bool isInitialState; + bool isFinalState; + uint32_t blockTokens; + uint32_t streamId; + // for debug + uint32_t batchIdx; + uint32_t headIdx; + uint32_t chunkIdx; +}; + +struct KdaFwdHStream { + uint32_t batchIdx; + uint32_t chunkIdx{0}; + uint32_t vHeadIdx; + uint32_t kHeadIdx; + uint32_t shapeBatchIdx; + uint32_t tokenBatchIdx; + + uint32_t chunkOffset; + uint32_t tokenOffset; + uint32_t batchChunks{0}; + uint32_t batchTokens; + uint32_t nextTaskIdx{0}; + bool active{false}; + + KdaFwdHOffsets offset; +}; + +struct KdaFwdHRunningQ { + KdaFwdHStream streams[PING_PONG_STAGES]; +}; + +struct BlockSchedulerKdaFwdH { + uint32_t batch; + uint32_t seqlen; + uint32_t kNumHead; + uint32_t vNumHead; + uint32_t kHeadDim; + uint32_t vHeadDim; + uint32_t chunkSize; + uint32_t vBlockSize{128}; + uint32_t isVariedLen; + uint32_t shapeBatch; + uint32_t tokenBatch; + uint32_t inputTokenBatch; + bool useInitialState; + bool storeFinalState; + uint32_t numSeqWorkspaceOffset; + uint32_t numChunksWorkspaceOffset; + + uint32_t taskIdx; + uint32_t taskStride; + uint32_t cubeCoreIdx; + uint32_t cubeCoreNum; + uint32_t taskNum; + uint32_t headGroups; + uint32_t totalChunks; + uint32_t totalTokens; + + KdaFwdHRunningQ runningQ; + + bool isRunning; + + AscendC::GlobalTensor gmSeqlen; + AscendC::GlobalTensor gmNumSeq; + AscendC::GlobalTensor gmNumChunks; + + Arch::CrossCoreFlag cube1Done[PING_PONG_STAGES] = {0, 1}; + Arch::CrossCoreFlag vec1Done[PING_PONG_STAGES] = {2, 3}; + Arch::CrossCoreFlag cube2Done[PING_PONG_STAGES] = {4, 5}; + Arch::CrossCoreFlag vec2Done[PING_PONG_STAGES] = {6, 7}; + + CATLASS_DEVICE + BlockSchedulerKdaFwdH() {} + + CATLASS_DEVICE + void Init(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR tiling, GM_ADDR user, uint32_t coreIdx, + uint32_t coreNum) + { + __gm__ ChunkKdaFwdHRuntimeTiling *__restrict kdaFwdHTilingData = + reinterpret_cast<__gm__ ChunkKdaFwdHRuntimeTiling *__restrict>(tiling); + + batch = kdaFwdHTilingData->batch; + seqlen = kdaFwdHTilingData->seqlen; + kNumHead = kdaFwdHTilingData->kNumHead; + vNumHead = kdaFwdHTilingData->vNumHead; + kHeadDim = kdaFwdHTilingData->kHeadDim; + vHeadDim = kdaFwdHTilingData->vHeadDim; + chunkSize = kdaFwdHTilingData->chunkSize; + isVariedLen = kdaFwdHTilingData->isVariedLen; + shapeBatch = kdaFwdHTilingData->shapeBatch; + tokenBatch = kdaFwdHTilingData->tokenBatch; + useInitialState = kdaFwdHTilingData->useInitialState; + storeFinalState = kdaFwdHTilingData->storeFinalState; + numSeqWorkspaceOffset = kdaFwdHTilingData->numSeqWorkspaceOffset; + numChunksWorkspaceOffset = kdaFwdHTilingData->numChunksWorkspaceOffset; + + InitRuntime(cu_seqlens, chunk_indices, user, coreIdx, coreNum); + } + + template + CATLASS_DEVICE void InitFromData(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, const TilingData &tilingData, + GM_ADDR user, uint32_t coreIdx, uint32_t coreNum) + { + batch = tilingData.batch; + seqlen = tilingData.seqlen; + kNumHead = tilingData.kNumHead; + vNumHead = tilingData.vNumHead; + kHeadDim = tilingData.kHeadDim; + vHeadDim = tilingData.vHeadDim; + chunkSize = tilingData.chunkSize; + isVariedLen = tilingData.isVariedLen; + shapeBatch = tilingData.shapeBatch; + tokenBatch = tilingData.tokenBatch; + useInitialState = tilingData.useInitialState; + storeFinalState = tilingData.storeFinalState; + numSeqWorkspaceOffset = tilingData.numSeqWorkspaceOffset; + numChunksWorkspaceOffset = tilingData.numChunksWorkspaceOffset; + + InitRuntime(cu_seqlens, chunk_indices, user, coreIdx, coreNum); + } + + CATLASS_DEVICE + void InitRuntime(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR user, uint32_t coreIdx, uint32_t coreNum) + { + gmSeqlen.SetGlobalBuffer((__gm__ int64_t *)cu_seqlens); + gmNumSeq.SetGlobalBuffer((__gm__ int64_t *)(user + numSeqWorkspaceOffset)); + gmNumChunks.SetGlobalBuffer((__gm__ int64_t *)(user + numChunksWorkspaceOffset)); + + if (isVariedLen) { + inputTokenBatch = tokenBatch; + uint32_t actualBatch = 0; + int64_t chunkPrefix = 0; + int64_t prevSeq = 0, currSeq; + for (uint32_t b = 1; b <= inputTokenBatch; b++) { + currSeq = gmSeqlen.GetValue(b); + int64_t batchSeqLen = currSeq - prevSeq; + if (batchSeqLen > 0) { + actualBatch++; + int64_t batchChunk = (batchSeqLen + chunkSize - 1) / chunkSize; + chunkPrefix += batchChunk; + } + prevSeq = currSeq; + } + tokenBatch = actualBatch; + batch = actualBatch; + totalChunks = chunkPrefix; + totalTokens = prevSeq; + } else { + totalChunks = (seqlen + chunkSize - 1) / chunkSize; + totalTokens = seqlen; + } + + cubeCoreIdx = coreIdx; + cubeCoreNum = coreNum; + vBlockSize = vHeadDim; + taskNum = batch * vNumHead; + headGroups = vNumHead / kNumHead; + InitTaskWave(0); + } + + CATLASS_DEVICE + void InitTaskWave(uint32_t waveIdx) + { + uint32_t firstTaskIdx = waveIdx * cubeCoreNum + cubeCoreIdx; + taskStride = taskNum; + for (uint32_t streamId = 0; streamId < PING_PONG_STAGES; ++streamId) { + auto &stream = runningQ.streams[streamId]; + stream.nextTaskIdx = streamId == 0 ? firstTaskIdx : taskNum; + stream.chunkIdx = 0; + stream.batchChunks = 0; + stream.active = false; + } + isRunning = firstTaskIdx < taskNum; + } + + CATLASS_DEVICE + uint32_t GetTaskWaveCount() const { return CeilDiv(taskNum, cubeCoreNum); } + + CATLASS_DEVICE + void ResolveVarlenSequence(uint32_t compactBatchIdx, KdaFwdHStream &stream) + { + uint32_t actualBatch = 0; + int64_t chunkPrefix = 0; + int64_t prevSeq = 0; + for (uint32_t b = 1; b <= inputTokenBatch; ++b) { + int64_t currSeq = gmSeqlen.GetValue(b); + int64_t batchTokens = currSeq - prevSeq; + if (batchTokens > 0) { + int64_t batchChunks = (batchTokens + chunkSize - 1) / chunkSize; + if (actualBatch == compactBatchIdx) { + stream.chunkOffset = static_cast(chunkPrefix); + stream.batchChunks = static_cast(batchChunks); + stream.tokenOffset = static_cast(prevSeq); + stream.batchTokens = static_cast(batchTokens); + return; + } + ++actualBatch; + chunkPrefix += batchChunks; + } + prevSeq = currSeq; + } + stream.chunkOffset = 0; + stream.batchChunks = 0; + stream.tokenOffset = 0; + stream.batchTokens = 0; + } + + CATLASS_DEVICE + uint32_t GetVarlenChunkOffset(uint32_t compactBatchIdx) + { + KdaFwdHStream stream; + ResolveVarlenSequence(compactBatchIdx, stream); + return stream.chunkOffset; + } + + CATLASS_DEVICE + void InitNewStream(KdaFwdHStream &newStream) + { + newStream.batchIdx = taskIdx / vNumHead; + newStream.vHeadIdx = taskIdx % vNumHead; + newStream.kHeadIdx = newStream.vHeadIdx / headGroups; + newStream.shapeBatchIdx = isVariedLen ? 0 : newStream.batchIdx; + newStream.tokenBatchIdx = isVariedLen ? newStream.batchIdx : 0; + if (isVariedLen) { + ResolveVarlenSequence(newStream.tokenBatchIdx, newStream); + } else { + newStream.chunkOffset = 0; + newStream.batchChunks = totalChunks; + newStream.tokenOffset = 0; + newStream.batchTokens = totalTokens; + } + newStream.chunkIdx = 0; + } + + CATLASS_DEVICE + void AssignNextStream(uint32_t streamId) + { + auto &stream = runningQ.streams[streamId]; + taskIdx = stream.nextTaskIdx; + if (taskIdx >= taskNum) { + stream.active = false; + stream.batchChunks = 0; + return; + } + + stream.nextTaskIdx += taskStride; + InitNewStream(stream); + stream.active = stream.batchChunks > 0; + if (stream.active) { + UpdateTask(streamId); + } + } + + CATLASS_DEVICE + void UpdateTask(uint32_t streamId) + { + auto &stream = runningQ.streams[streamId]; + auto &offset = stream.offset; + + offset.isInitialState = stream.chunkIdx == 0; + offset.isFinalState = stream.chunkIdx == (stream.batchChunks - 1); + uint32_t vBlockOffset = 0; + uint32_t vBlockDim = vBlockSize; + offset.initialStateOffset = (stream.batchIdx * vNumHead + stream.vHeadIdx) * kHeadDim * vHeadDim + vBlockOffset; + offset.finalStateOffset = (stream.batchIdx * vNumHead + stream.vHeadIdx) * kHeadDim * vHeadDim + vBlockOffset; + offset.hSrcOffset = (stream.shapeBatchIdx * vNumHead * totalChunks + stream.vHeadIdx * totalChunks + + stream.chunkOffset + stream.chunkIdx) * + kHeadDim * vHeadDim + + vBlockOffset; + offset.hDstOffset = offset.hSrcOffset + kHeadDim * vHeadDim; + if (storeFinalState && offset.isFinalState) { + offset.hDstOffset = offset.hSrcOffset; + } + offset.uvOffset = (stream.shapeBatchIdx * vNumHead * totalTokens + stream.vHeadIdx * totalTokens + + stream.tokenOffset + stream.chunkIdx * chunkSize) * + vHeadDim + + vBlockOffset; + offset.wkOffset = (stream.shapeBatchIdx * kNumHead * totalTokens + stream.kHeadIdx * totalTokens + + stream.tokenOffset + stream.chunkIdx * chunkSize) * + kHeadDim; + offset.wOffset = (stream.shapeBatchIdx * vNumHead * totalTokens + stream.vHeadIdx * totalTokens + + stream.tokenOffset + stream.chunkIdx * chunkSize) * + kHeadDim; + offset.gOffset = stream.shapeBatchIdx * vNumHead * totalTokens + stream.vHeadIdx * totalTokens + + stream.tokenOffset + stream.chunkIdx * chunkSize; + offset.gkOffset = (stream.shapeBatchIdx * vNumHead * totalTokens + stream.vHeadIdx * totalTokens + + stream.tokenOffset + stream.chunkIdx * chunkSize) * + kHeadDim; + offset.hWorkOffset = (cubeCoreIdx * PING_PONG_STAGES + streamId) * kHeadDim * vBlockSize; + offset.vWorkOffset = (cubeCoreIdx * PING_PONG_STAGES + streamId) * chunkSize * vBlockSize; + offset.kDecayWorkOffset = (cubeCoreIdx * PING_PONG_STAGES + streamId) * chunkSize * kHeadDim; + offset.vBlockOffset = vBlockOffset; + offset.vBlockDim = vBlockDim; + offset.blockTokens = offset.isFinalState ? (stream.batchTokens - stream.chunkIdx * chunkSize) : chunkSize; + offset.streamId = streamId; + offset.batchIdx = stream.batchIdx; + offset.headIdx = stream.vHeadIdx; + offset.chunkIdx = stream.chunkIdx; + } + + CATLASS_DEVICE + void InitTasks() + { + isRunning = false; + for (uint32_t streamId = 0; streamId < PING_PONG_STAGES; ++streamId) { + auto &stream = runningQ.streams[streamId]; + if (stream.active) { + stream.chunkIdx += 1; + if (stream.chunkIdx >= stream.batchChunks) { + stream.active = false; + stream.batchChunks = 0; + } + } + if (!stream.active) { + AssignNextStream(streamId); + } else { + UpdateTask(streamId); + } + if (stream.active) { + isRunning = true; + } + } + } + + CATLASS_DEVICE + const KdaFwdHStream &GetStream(uint32_t i) const { return runningQ.streams[i]; } + + CATLASS_DEVICE + uint32_t GetStreamId(uint32_t i) const { return i; } + + CATLASS_DEVICE + const KdaFwdHOffsets &GetCurTaskOffsets(const KdaFwdHStream &stream) const { return stream.offset; } + + CATLASS_DEVICE + bool StreamIsDone(const KdaFwdHStream &stream) const { return !stream.active; } + + CATLASS_DEVICE + bool NeedProcessStage2(const KdaFwdHStream &stream) { return storeFinalState || !stream.offset.isFinalState; } +}; + +struct BlockSchedulerKdaFwdHCube : public BlockSchedulerKdaFwdH { + CATLASS_DEVICE + BlockSchedulerKdaFwdHCube() {} + + CATLASS_DEVICE + void Init(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR tiling, GM_ADDR user) + { + BlockSchedulerKdaFwdH::Init(cu_seqlens, chunk_indices, tiling, user, AscendC::GetBlockIdx(), + AscendC::GetBlockNum()); + } + + template + CATLASS_DEVICE void InitFromData(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, const TilingData &tilingData, + GM_ADDR user) + { + BlockSchedulerKdaFwdH::InitFromData(cu_seqlens, chunk_indices, tilingData, user, AscendC::GetBlockIdx(), + AscendC::GetBlockNum()); + } +}; + +struct BlockSchedulerKdaFwdHVec : public BlockSchedulerKdaFwdH { + CATLASS_DEVICE + BlockSchedulerKdaFwdHVec() {} + + CATLASS_DEVICE + void Init(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR tiling, GM_ADDR user) + { + BlockSchedulerKdaFwdH::Init(cu_seqlens, chunk_indices, tiling, user, + AscendC::GetBlockIdx() / AscendC::GetSubBlockNum(), AscendC::GetBlockNum()); + } + + template + CATLASS_DEVICE void InitFromData(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, const TilingData &tilingData, + GM_ADDR user) + { + BlockSchedulerKdaFwdH::InitFromData(cu_seqlens, chunk_indices, tilingData, user, + AscendC::GetBlockIdx() / AscendC::GetSubBlockNum(), AscendC::GetBlockNum()); + } +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_SCHEDULER_KDA_FWD_H_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch22.hpp b/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch22.hpp new file mode 100644 index 0000000000..3cea7b6167 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch22/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch22.hpp @@ -0,0 +1,775 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#define CATLASS_ARCH 2201 + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/arch/kda_cross_core_sync.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/kda_debug.hpp" +#include "catlass/epilogue/block/kda_block_epilogue.hpp" +#include "../../epilogue/block/block_epilogue_kda_fwdh_update_arch22.hpp" +#include "../../epilogue/block/block_epilogue_kda_fwdh_vnew_arch22.hpp" +#include "catlass/gemm/block/kda_block_mmad.hpp" +#include "../../../../kernel_utils/block/block_mmad_pingpong_tla_multi.hpp" +#include "../block/block_scheduler_kda_fwd_h_arch22.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "tla/kda_tensor.hpp" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +#include "kernel_operator.h" +using namespace Catlass; +using namespace tla; + +namespace Catlass::Gemm::Kernel { + +struct KdaFwdHTileShapes128 { + using L1TileShape = tla::Shape<_128, _128, _128>; + using L0TileShape = L1TileShape; +}; + +struct KdaFwdHTileShapes256 { + using L1TileShape = tla::Shape<_128, _256, _128>; + using L0TileShape = tla::Shape<_128, _256, _64>; +}; + +template +struct KdaFwdHGateTag { + static constexpr bool value = KGated; + static constexpr bool scalarGated = ScalarGated; + static constexpr bool useExp2 = UseExp2; +}; + +template +class KdaFwdHKernel { +public: + using ArchTag = Arch::AtlasA2; + using CubeScheduler = typename Catlass::Gemm::Block::BlockSchedulerKdaFwdHCube; + using VecScheduler = typename Catlass::Gemm::Block::BlockSchedulerKdaFwdHVec; + + using DispatchPolicyTla = Gemm::MmadPingpongTlaMulti; + using L1TileShapeVTla = typename TileShapes::L1TileShape; + using L0TileShapeVTla = typename TileShapes::L0TileShape; + + using WType = Gemm::GemmType; + using HType = Gemm::GemmType; + using VworkType = Gemm::GemmType; + using KType = Gemm::GemmType; + using HworkType = Gemm::GemmType; + using VType = Gemm::GemmType; + using GType = Gemm::GemmType; + using UType = Gemm::GemmType; + using FinalStateType = Gemm::GemmType; + + // cube 1 + using TileCopyWH = Catlass::Gemm::Tile::PackedTileCopyTla; + using BlockMmadWH = Gemm::Block::BlockMmadTla; + + // cube 2 + using TileCopyKV = Catlass::Gemm::Tile::PackedTileCopyTla; + using BlockMmadKV = Gemm::Block::BlockMmadTla; + + // vec 1 + using DispatchPolicyKdaFwdHVnew = Epilogue::EpilogueAtlasKdaFwdHVnew; + using GateTag = KdaFwdHGateTag; + using EpilogueKdaFwdHVnew = Epilogue::Block::BlockEpilogue; + + // vec 2 + using DispatchPolicyKdaFwdHUpdate = Epilogue::EpilogueAtlasKdaFwdHUpdate; + using EpilogueKdaFwdHUpdate = Epilogue::Block::BlockEpilogue; + + using KdaFwdHOffsets = Catlass::Gemm::Block::KdaFwdHOffsets; + + using ElementK = INPUT_TYPE; + using ElementW = INPUT_TYPE; + using ElementU = INPUT_TYPE; + using ElementG = G_TYPE; + using ElementH = INPUT_TYPE; + using ElementV = INPUT_TYPE; + using ElementVWork = WORKSPACE_TYPE; + using ElementHWork = WORKSPACE_TYPE; + using ElementInitialState = STATE_TYPE; + using ElementFinalState = STATE_TYPE; + + using LayoutW = Catlass::layout::RowMajor; + using LayoutH = Catlass::layout::RowMajor; + using LayoutV = Catlass::layout::RowMajor; + using LayoutK = Catlass::layout::ColumnMajor; + + uint32_t batch; + uint32_t seqlen; + uint32_t kNumHead; + uint32_t vNumHead; + uint32_t kHeadDim; + uint32_t vHeadDim; + uint32_t chunkSize; + bool useInitialState; + bool storeFinalState; + uint32_t isVariedLen; + uint32_t shapeBatch; + uint32_t tokenBatch; + uint32_t vWorkspaceOffset; + uint32_t vUpdateWorkspaceOffset; + uint32_t hWorkspaceOffset; + uint32_t numSeqWorkspaceOffset; + uint32_t numChunksWorkspaceOffset; + uint32_t kDecayWorkspaceOffset; + + AscendC::GlobalTensor gmK; + AscendC::GlobalTensor gmW; + AscendC::GlobalTensor gmU; + AscendC::GlobalTensor gmG; + AscendC::GlobalTensor gmInitialState; + AscendC::GlobalTensor gmH; + AscendC::GlobalTensor gmV; + AscendC::GlobalTensor gmFinalState; + AscendC::GlobalTensor gmVWorkspace; + AscendC::GlobalTensor gmVUpdateWorkspace; + AscendC::GlobalTensor gmHWorkspace; + + AscendC::GlobalTensor gmGk; + AscendC::GlobalTensor gmKDecayWorkspace; + + AscendC::GlobalTensor gmSeqlen; + AscendC::GlobalTensor gmNumSeq; + AscendC::GlobalTensor gmNumChunks; + + CubeScheduler cubeBlockScheduler; + VecScheduler vecBlockScheduler; + + Arch::Resource resource; + + __aicore__ inline KdaFwdHKernel() {} + + __aicore__ inline void Init(GM_ADDR k, GM_ADDR w, GM_ADDR u, GM_ADDR g, GM_ADDR gk, GM_ADDR inital_state, + GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR h, GM_ADDR v_new, + GM_ADDR final_state, GM_ADDR tiling, GM_ADDR user) + { + __gm__ ChunkKdaFwdHRuntimeTiling *__restrict kdaFwdHTilingData = + reinterpret_cast<__gm__ ChunkKdaFwdHRuntimeTiling *__restrict>(tiling); + + batch = kdaFwdHTilingData->batch; + seqlen = kdaFwdHTilingData->seqlen; + kNumHead = kdaFwdHTilingData->kNumHead; + vNumHead = kdaFwdHTilingData->vNumHead; + kHeadDim = kdaFwdHTilingData->kHeadDim; + vHeadDim = kdaFwdHTilingData->vHeadDim; + chunkSize = kdaFwdHTilingData->chunkSize; + useInitialState = kdaFwdHTilingData->useInitialState; + storeFinalState = kdaFwdHTilingData->storeFinalState; + isVariedLen = kdaFwdHTilingData->isVariedLen; + shapeBatch = kdaFwdHTilingData->shapeBatch; + tokenBatch = kdaFwdHTilingData->tokenBatch; + vWorkspaceOffset = kdaFwdHTilingData->vWorkspaceOffset; + vUpdateWorkspaceOffset = kdaFwdHTilingData->vUpdateWorkspaceOffset; + hWorkspaceOffset = kdaFwdHTilingData->hWorkspaceOffset; + numSeqWorkspaceOffset = kdaFwdHTilingData->numSeqWorkspaceOffset; + numChunksWorkspaceOffset = kdaFwdHTilingData->numChunksWorkspaceOffset; + kDecayWorkspaceOffset = kdaFwdHTilingData->kDecayWorkspaceOffset; + + gmK.SetGlobalBuffer((__gm__ ElementK *)k); + gmW.SetGlobalBuffer((__gm__ ElementW *)w); + gmU.SetGlobalBuffer((__gm__ ElementU *)u); + gmG.SetGlobalBuffer((__gm__ ElementG *)(scalarGated ? g : gk)); + gmInitialState.SetGlobalBuffer((__gm__ ElementInitialState *)inital_state); + gmH.SetGlobalBuffer((__gm__ ElementH *)h); + gmV.SetGlobalBuffer((__gm__ ElementV *)v_new); + gmFinalState.SetGlobalBuffer((__gm__ ElementFinalState *)final_state); + gmVWorkspace.SetGlobalBuffer((__gm__ ElementVWork *)(user + vWorkspaceOffset)); + gmVUpdateWorkspace.SetGlobalBuffer((__gm__ ElementV *)(user + vUpdateWorkspaceOffset)); + gmHWorkspace.SetGlobalBuffer((__gm__ ElementHWork *)(user + hWorkspaceOffset)); + gmGk.SetGlobalBuffer((__gm__ ElementG *)(kGated ? gk : g)); + gmKDecayWorkspace.SetGlobalBuffer((__gm__ ElementK *)(user + kDecayWorkspaceOffset)); + + gmSeqlen.SetGlobalBuffer((__gm__ int64_t *)cu_seqlens); + gmNumSeq.SetGlobalBuffer((__gm__ int64_t *)(user + numSeqWorkspaceOffset)); + gmNumChunks.SetGlobalBuffer((__gm__ int64_t *)(user + numChunksWorkspaceOffset)); + + if ASCEND_IS_AIC { + cubeBlockScheduler.Init(cu_seqlens, chunk_indices, tiling, user); + } + + if ASCEND_IS_AIV { + vecBlockScheduler.Init(cu_seqlens, chunk_indices, tiling, user); + } + } + + template + __aicore__ inline void InitFromData(GM_ADDR k, GM_ADDR w, GM_ADDR u, GM_ADDR g, GM_ADDR gk, GM_ADDR inital_state, + GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR h, GM_ADDR v_new, + GM_ADDR final_state, const TilingData &tilingData, GM_ADDR user) + { + batch = tilingData.batch; + seqlen = tilingData.seqlen; + kNumHead = tilingData.kNumHead; + vNumHead = tilingData.vNumHead; + kHeadDim = tilingData.kHeadDim; + vHeadDim = tilingData.vHeadDim; + chunkSize = tilingData.chunkSize; + useInitialState = tilingData.useInitialState; + storeFinalState = tilingData.storeFinalState; + isVariedLen = tilingData.isVariedLen; + shapeBatch = tilingData.shapeBatch; + tokenBatch = tilingData.tokenBatch; + vWorkspaceOffset = tilingData.vWorkspaceOffset; + vUpdateWorkspaceOffset = tilingData.vUpdateWorkspaceOffset; + hWorkspaceOffset = tilingData.hWorkspaceOffset; + numSeqWorkspaceOffset = tilingData.numSeqWorkspaceOffset; + numChunksWorkspaceOffset = tilingData.numChunksWorkspaceOffset; + kDecayWorkspaceOffset = tilingData.kDecayWorkspaceOffset; + + gmK.SetGlobalBuffer((__gm__ ElementK *)k); + gmW.SetGlobalBuffer((__gm__ ElementW *)w); + gmU.SetGlobalBuffer((__gm__ ElementU *)u); + gmG.SetGlobalBuffer((__gm__ ElementG *)(scalarGated ? g : gk)); + gmInitialState.SetGlobalBuffer((__gm__ ElementInitialState *)inital_state); + gmH.SetGlobalBuffer((__gm__ ElementH *)h); + gmV.SetGlobalBuffer((__gm__ ElementV *)v_new); + gmFinalState.SetGlobalBuffer((__gm__ ElementFinalState *)final_state); + gmVWorkspace.SetGlobalBuffer((__gm__ ElementVWork *)(user + vWorkspaceOffset)); + gmVUpdateWorkspace.SetGlobalBuffer((__gm__ ElementV *)(user + vUpdateWorkspaceOffset)); + gmHWorkspace.SetGlobalBuffer((__gm__ ElementHWork *)(user + hWorkspaceOffset)); + gmGk.SetGlobalBuffer((__gm__ ElementG *)(kGated ? gk : g)); + gmKDecayWorkspace.SetGlobalBuffer((__gm__ ElementK *)(user + kDecayWorkspaceOffset)); + gmSeqlen.SetGlobalBuffer((__gm__ int64_t *)cu_seqlens); + gmNumSeq.SetGlobalBuffer((__gm__ int64_t *)(user + numSeqWorkspaceOffset)); + gmNumChunks.SetGlobalBuffer((__gm__ int64_t *)(user + numChunksWorkspaceOffset)); + + if ASCEND_IS_AIC { + cubeBlockScheduler.InitFromData(cu_seqlens, chunk_indices, tilingData, user); + } + if ASCEND_IS_AIV { + vecBlockScheduler.InitFromData(cu_seqlens, chunk_indices, tilingData, user); + } + } + + template + __aicore__ inline float LoadScalarAsFloat(AscendC::GlobalTensor tensor, uint32_t offset) const + { + Element value = tensor.GetValue(offset); + if constexpr (std::is_same::value) { + return AscendC::ToFloat(value); + } + return static_cast(value); + } + + // Borrow the stream's V_MTE2 free token and restore it before returning. + __aicore__ inline void ComputeTailVWorkspace(const KdaFwdHOffsets &offsets, uint32_t tailEventId) + { + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t rowsPerSubBlock = CeilDiv(offsets.blockTokens, subBlockNum); + uint32_t rowBegin = subBlockIdx * rowsPerSubBlock; + uint32_t rowEnd = Min(rowBegin + rowsPerSubBlock, offsets.blockTokens); + if (rowBegin >= rowEnd) { + return; + } + AscendC::ResetMask(); + AscendC::WaitFlag(tailEventId); + + constexpr uint32_t TAIL_INPUT_OFFSET = 166 * 1024; + constexpr uint32_t TAIL_FLOAT_OFFSET = 167 * 1024; + constexpr uint32_t TAIL_ACCUM_OFFSET = 168 * 1024; + AscendC::LocalTensor inputUb = resource.ubBuf.template GetBufferByByte(TAIL_INPUT_OFFSET); + AscendC::LocalTensor floatUb = resource.ubBuf.template GetBufferByByte(TAIL_FLOAT_OFFSET); + AscendC::LocalTensor accumUb = resource.ubBuf.template GetBufferByByte(TAIL_ACCUM_OFFSET); + + for (uint32_t tokenRow = rowBegin; tokenRow < rowEnd; ++tokenRow) { + AscendC::Duplicate(accumUb, 0.0f, offsets.vBlockDim); + AscendC::PipeBarrier(); + for (uint32_t kIdx = 0; kIdx < kHeadDim; ++kIdx) { + AscendC::DataCopy(inputUb, gmH[offsets.hSrcOffset + kIdx * vHeadDim], offsets.vBlockDim); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::Cast(floatUb, inputUb, AscendC::RoundMode::CAST_NONE, offsets.vBlockDim); + // inputUb is reused by the next kIdx iteration. PIPE_V and + // PIPE_MTE2 run independently, so a PIPE_V barrier alone does + // not stop the next DataCopy from overwriting inputUb while + // Cast is still reading it. Queue a V -> MTE2 dependency + // before that next copy, matching ComputeHWorkspaceVector(). + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::PipeBarrier(); + float weight = LoadScalarAsFloat(gmW, offsets.wOffset + tokenRow * kHeadDim + kIdx); + AscendC::Muls(floatUb, floatUb, weight, offsets.vBlockDim); + AscendC::PipeBarrier(); + AscendC::Add(accumUb, accumUb, floatUb, offsets.vBlockDim); + AscendC::PipeBarrier(); + } + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::DataCopy(gmVWorkspace[offsets.vWorkOffset + tokenRow * offsets.vBlockDim], accumUb, + offsets.vBlockDim); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + // The MTE3 copy above still reads accumUb. The next token row + // starts by clearing the same buffer on PIPE_V, so protect that + // reuse explicitly; MTE3_MTE2 does not order MTE3 against PIPE_V. + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + } + AscendC::SetFlag(tailEventId); + } + + __aicore__ inline void ComputeHWorkspaceVector(const KdaFwdHOffsets &offsets, uint32_t eventId) + { + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t rowsPerSubBlock = CeilDiv(kHeadDim, subBlockNum); + uint32_t rowBegin = subBlockIdx * rowsPerSubBlock; + uint32_t rowEnd = Min(rowBegin + rowsPerSubBlock, kHeadDim); + + constexpr uint32_t MAX_K_ROWS = 64; + constexpr uint32_t ACCUM_OFFSET = 0; + constexpr uint32_t WEIGHT_BROADCAST_OFFSET = 32 * 1024; + constexpr uint32_t WEIGHT_INPUT_OFFSET = 64 * 1024; + constexpr uint32_t WEIGHT_FLOAT_OFFSET = 65 * 1024; + constexpr uint32_t V_INPUT_OFFSET = 66 * 1024; + constexpr uint32_t V_FLOAT_OFFSET = 67 * 1024; + constexpr uint32_t PRODUCT_OFFSET = 68 * 1024; + constexpr uint32_t BROADCAST_SCRATCH_OFFSET = 96 * 1024; + + AscendC::LocalTensor accumUb = resource.ubBuf.template GetBufferByByte(ACCUM_OFFSET); + AscendC::LocalTensor weightBroadcastUb = + resource.ubBuf.template GetBufferByByte(WEIGHT_BROADCAST_OFFSET); + AscendC::LocalTensor weightInputUb = + resource.ubBuf.template GetBufferByByte(WEIGHT_INPUT_OFFSET); + AscendC::LocalTensor weightFloatUb = resource.ubBuf.template GetBufferByByte(WEIGHT_FLOAT_OFFSET); + AscendC::LocalTensor vInputUb = resource.ubBuf.template GetBufferByByte(V_INPUT_OFFSET); + AscendC::LocalTensor vFloatUb = resource.ubBuf.template GetBufferByByte(V_FLOAT_OFFSET); + AscendC::LocalTensor productUb = resource.ubBuf.template GetBufferByByte(PRODUCT_OFFSET); + AscendC::LocalTensor broadcastScratchUb = + resource.ubBuf.template GetBufferByByte(BROADCAST_SCRATCH_OFFSET); + + AscendC::ResetMask(); + AscendC::WaitFlag(eventId); + + for (uint32_t kTileStart = rowBegin; kTileStart < rowEnd; kTileStart += MAX_K_ROWS) { + uint32_t rowCount = Min(MAX_K_ROWS, rowEnd - kTileStart); + uint32_t accumElements = rowCount * offsets.vBlockDim; + AscendC::Duplicate(accumUb, 0.0f, accumElements); + AscendC::PipeBarrier(); + for (uint32_t tokenRow = 0; tokenRow < offsets.blockTokens; ++tokenRow) { + AscendC::DataCopy(weightInputUb, gmK[offsets.wkOffset + tokenRow * kHeadDim], kHeadDim); + AscendC::DataCopy(vInputUb, gmV[offsets.uvOffset + tokenRow * vHeadDim], offsets.vBlockDim); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); + AscendC::Cast(weightFloatUb, weightInputUb, AscendC::RoundMode::CAST_NONE, kHeadDim); + AscendC::Cast(vFloatUb, vInputUb, AscendC::RoundMode::CAST_NONE, offsets.vBlockDim); + AscendC::PipeBarrier(); + + uint32_t dstShape[2] = {rowCount, offsets.vBlockDim}; + uint32_t srcShape[2] = {rowCount, 1}; + AscendC::Broadcast(weightBroadcastUb, weightFloatUb[kTileStart], dstShape, srcShape, + broadcastScratchUb); + AscendC::PipeBarrier(); + + for (uint32_t localRow = 0; localRow < rowCount; ++localRow) { + uint32_t rowOffset = localRow * offsets.vBlockDim; + AscendC::Mul(productUb, vFloatUb, weightBroadcastUb[rowOffset], offsets.vBlockDim); + AscendC::PipeBarrier(); + AscendC::Add(accumUb[rowOffset], accumUb[rowOffset], productUb, offsets.vBlockDim); + AscendC::PipeBarrier(); + } + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); + } + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); + AscendC::DataCopy(gmHWorkspace[offsets.hWorkOffset + kTileStart * offsets.vBlockDim], accumUb, + accumElements); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); + } + AscendC::SetFlag(eventId); + } + + __aicore__ inline void PresetVectorPipelineEvents() + { + constexpr uint32_t pongBaseEvent = 4; + if (storeFinalState && std::is_same::value) { + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID0 + pongBaseEvent); + AscendC::SetFlag(EVENT_ID2); + AscendC::SetFlag(EVENT_ID2 + pongBaseEvent); + } else { + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID0 + pongBaseEvent); + AscendC::SetFlag(EVENT_ID2); + AscendC::SetFlag(EVENT_ID2 + pongBaseEvent); + } + AscendC::SetFlag(EVENT_ID1); + AscendC::SetFlag(EVENT_ID1 + pongBaseEvent); + AscendC::SetFlag(EVENT_ID3); + AscendC::SetFlag(EVENT_ID3 + pongBaseEvent); + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID0 + pongBaseEvent); + AscendC::SetFlag(EVENT_ID2); + AscendC::SetFlag(EVENT_ID2 + pongBaseEvent); + } + + __aicore__ inline void DrainVectorPipelineEvents(const bool (&event0FromMte3)[PING_PONG_STAGES], + const bool (&event2FromMte3)[PING_PONG_STAGES]) + { + constexpr uint32_t pongBaseEvent = 4; + if (storeFinalState && std::is_same::value) { + for (uint32_t streamId = 0; streamId < PING_PONG_STAGES; ++streamId) { + uint32_t eventOffset = streamId * pongBaseEvent; + if (event0FromMte3[streamId]) { + AscendC::WaitFlag(EVENT_ID0 + eventOffset); + } else { + AscendC::WaitFlag(EVENT_ID0 + eventOffset); + } + if (event2FromMte3[streamId]) { + AscendC::WaitFlag(EVENT_ID2 + eventOffset); + } else { + AscendC::WaitFlag(EVENT_ID2 + eventOffset); + } + } + } else { + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0 + pongBaseEvent); + AscendC::WaitFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID2 + pongBaseEvent); + } + AscendC::WaitFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID1 + pongBaseEvent); + AscendC::WaitFlag(EVENT_ID3); + AscendC::WaitFlag(EVENT_ID3 + pongBaseEvent); + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0 + pongBaseEvent); + AscendC::WaitFlag(EVENT_ID2); + AscendC::WaitFlag(EVENT_ID2 + pongBaseEvent); + } + + __aicore__ inline void Process() + { + if (isVariedLen) { + AscendC::SyncAll(); + } + + if ASCEND_IS_AIC { + uint32_t coreIdx = AscendC::GetBlockIdx(); + uint32_t coreNum = AscendC::GetBlockNum(); + + auto wLayout = + tla::MakeLayout(shapeBatch * kNumHead * cubeBlockScheduler.totalTokens, kHeadDim); + auto hLayout = tla::MakeLayout( + shapeBatch * vNumHead * cubeBlockScheduler.totalChunks * kHeadDim, vHeadDim); + auto vLayout = tla::MakeLayout(coreNum * chunkSize * PING_PONG_STAGES, + cubeBlockScheduler.vBlockSize); + auto kLayout = + tla::MakeLayout(kHeadDim, shapeBatch * kNumHead * cubeBlockScheduler.totalTokens); + auto vworkLayout = tla::MakeLayout(coreNum * chunkSize * PING_PONG_STAGES, + cubeBlockScheduler.vBlockSize); + auto hworkLayout = tla::MakeLayout(coreNum * kHeadDim * PING_PONG_STAGES, + cubeBlockScheduler.vBlockSize); + + uint32_t taskWaveCount = cubeBlockScheduler.GetTaskWaveCount(); + for (uint32_t waveIdx = 0; waveIdx < taskWaveCount; ++waveIdx) { + BlockMmadWH blockMmadWH(resource); + BlockMmadKV blockMmadKV(resource); + BlockMmadWH blockMmadWHTail(resource); + BlockMmadKV blockMmadKVTail(resource); + AscendC::SyncAll(); + cubeBlockScheduler.InitTaskWave(waveIdx); + uint32_t currStage = 0; // 0: C1, 1: C2 + while (cubeBlockScheduler.isRunning) { + if (currStage == 0) { + /* C1: v_work = w @ h[i] */ + cubeBlockScheduler.InitTasks(); + for (uint32_t i = 0; i < PING_PONG_STAGES; ++i) { + uint32_t streamId = cubeBlockScheduler.GetStreamId(i); + const auto &stream = cubeBlockScheduler.GetStream(i); + if (cubeBlockScheduler.StreamIsDone(stream)) { + continue; + } + + const KdaFwdHOffsets &cube1Offsets = cubeBlockScheduler.GetCurTaskOffsets(stream); + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec2Done[streamId]); + if (cube1Offsets.blockTokens < 16) { + Arch::CrossCoreSetFlag<0x2, PIPE_MTE2>(cubeBlockScheduler.cube1Done[streamId]); + continue; + } + int64_t cube1OffsetW = cube1Offsets.wOffset; + int64_t cube1OffsetH = cube1Offsets.hSrcOffset; + int64_t cube1OffsetVwork = cube1Offsets.vWorkOffset; + auto tensorW = tla::MakeTensor(gmW[cube1OffsetW], wLayout, Catlass::Arch::PositionGM{}); + auto tensorH = tla::MakeTensor(gmH[cube1OffsetH], hLayout, Catlass::Arch::PositionGM{}); + auto tensorV = + tla::MakeTensor(gmVWorkspace[cube1OffsetVwork], vLayout, Catlass::Arch::PositionGM{}); + GemmCoord cube1Shape{cube1Offsets.blockTokens, cube1Offsets.vBlockDim, kHeadDim}; + auto tensorBlockW = + GetTile(tensorW, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.m(), cube1Shape.k())); + auto tensorBlockH = + GetTile(tensorH, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.k(), cube1Shape.n())); + auto tensorBlockV = + GetTile(tensorV, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.m(), cube1Shape.n())); + if (cube1Offsets.blockTokens < chunkSize) { + blockMmadWHTail.preSetFlags(); + blockMmadWHTail(tensorBlockW, tensorBlockH, tensorBlockV, cube1Shape, EmptyClass{}, + true); + blockMmadWHTail.finalWaitFlags(); + } else { + blockMmadWH.preSetFlags(); + blockMmadWH(tensorBlockW, tensorBlockH, tensorBlockV, cube1Shape); + blockMmadWH.finalWaitFlags(); + } + AscendC::PipeBarrier(); + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(cubeBlockScheduler.cube1Done[streamId]); + } + } else { + /* C2: h[i+1] = k.T @ v_work */ + for (uint32_t i = 0; i < PING_PONG_STAGES; ++i) { + uint32_t streamId = cubeBlockScheduler.GetStreamId(i); + const auto &stream = cubeBlockScheduler.GetStream(i); + if (cubeBlockScheduler.StreamIsDone(stream)) { + continue; + } + const KdaFwdHOffsets &cube2Offsets = cubeBlockScheduler.GetCurTaskOffsets(stream); + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec1Done[streamId]); + + if (cubeBlockScheduler.NeedProcessStage2(stream)) { + if (cube2Offsets.blockTokens < 16) { + Arch::CrossCoreSetFlag<0x2, PIPE_MTE2>(cubeBlockScheduler.cube2Done[streamId]); + continue; + } + auto tensorK = + tla::MakeTensor(gmK[cube2Offsets.wkOffset], kLayout, Catlass::Arch::PositionGM{}); + auto tensorVwork = tla::MakeTensor(gmVUpdateWorkspace[cube2Offsets.vWorkOffset], + vworkLayout, Catlass::Arch::PositionGM{}); + auto tensorHwork = tla::MakeTensor(gmHWorkspace[cube2Offsets.hWorkOffset], hworkLayout, + Catlass::Arch::PositionGM{}); + GemmCoord cube2Shape{kHeadDim, cube2Offsets.vBlockDim, cube2Offsets.blockTokens}; + auto tensorBlockK = GetTile(tensorK, tla::MakeCoord(0, 0), + tla::MakeShape(cube2Shape.m(), cube2Shape.k())); + auto tensorBlockVwork = GetTile(tensorVwork, tla::MakeCoord(0, 0), + tla::MakeShape(cube2Shape.k(), cube2Shape.n())); + auto tensorBlockHwork = GetTile(tensorHwork, tla::MakeCoord(0, 0), + tla::MakeShape(cube2Shape.m(), cube2Shape.n())); + if (cube2Offsets.blockTokens < chunkSize) { + blockMmadKVTail.preSetFlags(); + blockMmadKVTail(tensorBlockK, tensorBlockVwork, tensorBlockHwork, cube2Shape, + EmptyClass{}, true); + blockMmadKVTail.finalWaitFlags(); + } else { + blockMmadKV.preSetFlags(); + blockMmadKV(tensorBlockK, tensorBlockVwork, tensorBlockHwork, cube2Shape); + blockMmadKV.finalWaitFlags(); + } + AscendC::PipeBarrier(); + } + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(cubeBlockScheduler.cube2Done[streamId]); + } + } + currStage ^= 0x01; + } + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec2Done[0]); + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec2Done[1]); + } + } + + if ASCEND_IS_AIV { + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t coreIdx = AscendC::GetBlockIdx() / subBlockNum; + uint32_t coreNum = AscendC::GetBlockNum(); + uint32_t taskCount = (isVariedLen ? vecBlockScheduler.tokenBatch : shapeBatch) * vNumHead; + uint32_t rowsPerSubBlock = (kHeadDim + subBlockNum - 1) / subBlockNum; + uint32_t rowBegin = subBlockIdx * rowsPerSubBlock; + uint32_t rowEnd = Min(rowBegin + rowsPerSubBlock, kHeadDim); + uint32_t hRowsPerTile = (32 * 1024) / (vHeadDim * sizeof(ElementH)); + uint32_t stateRowsPerTile = (64 * 1024) / (vHeadDim * sizeof(ElementInitialState)); + uint32_t rowsPerTile = Min(hRowsPerTile, stateRowsPerTile); + uint32_t totalChunks = isVariedLen ? vecBlockScheduler.totalChunks : ((seqlen + chunkSize - 1) / chunkSize); + uint32_t stateBlockSize = kHeadDim * vHeadDim; + AscendC::LocalTensor stateUbTensorPing = + resource.ubBuf.template GetBufferByByte(0); + AscendC::LocalTensor stateUbTensorPong = + resource.ubBuf.template GetBufferByByte(96 * 1024); + AscendC::LocalTensor hUbTensorPing = resource.ubBuf.template GetBufferByByte(64 * 1024); + AscendC::LocalTensor hUbTensorPong = + resource.ubBuf.template GetBufferByByte(160 * 1024); + uint32_t taskWaveCount = vecBlockScheduler.GetTaskWaveCount(); + for (uint32_t waveIdx = 0; waveIdx < taskWaveCount; ++waveIdx) { + EpilogueKdaFwdHVnew epilogueKdaFwdHVnew(resource); + EpilogueKdaFwdHUpdate epilogueKdaFwdHUpdate(resource); + uint32_t taskIdx = waveIdx * coreNum + coreIdx; + uint32_t pingpongFlag = 1; + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + if (taskIdx < taskCount) { + uint32_t batchIdx = taskIdx / vNumHead; + uint32_t vHeadIdx = taskIdx % vNumHead; + uint32_t chunkOffset = isVariedLen ? vecBlockScheduler.GetVarlenChunkOffset(batchIdx) : 0; + uint32_t shapeBatchIdx = isVariedLen ? 0 : batchIdx; + uint32_t hBaseOffset = + (shapeBatchIdx * vNumHead * totalChunks + vHeadIdx * totalChunks + chunkOffset) * + stateBlockSize; + uint32_t initialStateBaseOffset = taskIdx * stateBlockSize; + for (uint32_t rowOffset = rowBegin; rowOffset < rowEnd; rowOffset += rowsPerTile) { + uint32_t rowsThisTile = Min(rowsPerTile, rowEnd - rowOffset); + uint32_t stateTileElems = rowsThisTile * vHeadDim; + uint32_t hOffset = hBaseOffset + rowOffset * vHeadDim; + AscendC::LocalTensor stateUbTensor = + pingpongFlag ? stateUbTensorPing : stateUbTensorPong; + AscendC::LocalTensor hUbTensor = pingpongFlag ? hUbTensorPing : hUbTensorPong; + auto eventId = pingpongFlag ? EVENT_ID1 : EVENT_ID0; + AscendC::WaitFlag(eventId); + if (useInitialState) { + uint32_t initialStateOffset = initialStateBaseOffset + rowOffset * vHeadDim; + if constexpr (!std::is_same::value) { + AscendC::DataCopy(stateUbTensor, gmInitialState[initialStateOffset], stateTileElems); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); + AscendC::Cast(hUbTensor, stateUbTensor, AscendC::RoundMode::CAST_RINT, stateTileElems); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); + AscendC::DataCopy(gmH[hOffset], hUbTensor, stateTileElems); + } else { + AscendC::DataCopy(stateUbTensor, gmInitialState[initialStateOffset], stateTileElems); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); + AscendC::DataCopy(gmH[hOffset], stateUbTensor, stateTileElems); + } + } else { + AscendC::Duplicate(hUbTensor, static_cast(0), stateTileElems); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); + AscendC::DataCopy(gmH[hOffset], hUbTensor, stateTileElems); + } + AscendC::SetFlag(eventId); + pingpongFlag = 1 - pingpongFlag; + } + } + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + + AscendC::SyncAll(); + vecBlockScheduler.InitTaskWave(waveIdx); + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec2Done[0]); + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec2Done[1]); + PresetVectorPipelineEvents(); + uint32_t currStage = 0; // 0: V1, 1: V2 + bool waitStageFence = false; + bool event0FromMte3[PING_PONG_STAGES] = {false, false}; + bool event2FromMte3[PING_PONG_STAGES] = { + !(storeFinalState && std::is_same::value), + !(storeFinalState && std::is_same::value)}; + constexpr uint32_t pongEventOffset = 4; + while (vecBlockScheduler.isRunning) { + if (waitStageFence) { + AscendC::WaitFlag(EVENT_ID1); + } + if (currStage == 0) { + /* V1: + * gmV = gmU - gmVWorkspace + * g_buf = gmG[-1] - gmG + * g_buf = exp(g_buf) + * gmVWorkspace = g_buf * gmV + */ + vecBlockScheduler.InitTasks(); + for (uint32_t i = 0; i < PING_PONG_STAGES; ++i) { + uint32_t streamId = vecBlockScheduler.GetStreamId(i); + const auto &stream = vecBlockScheduler.GetStream(i); + if (vecBlockScheduler.StreamIsDone(stream)) { + continue; + } + const KdaFwdHOffsets &vec1Offsets = vecBlockScheduler.GetCurTaskOffsets(stream); + bool tailVectorPath = vec1Offsets.blockTokens < 16; + if (tailVectorPath) { + // For a short tail Cube does not calculate W @ H, + // but cube1Done still orders this task after the + // previous chunk's V2 update of gmH. Wait before + // the software W @ H reads gmH, not afterwards in + // the epilogue. + Arch::CrossCoreWaitFlag(vecBlockScheduler.cube1Done[streamId]); + ComputeTailVWorkspace(vec1Offsets, + EVENT_ID3 + (streamId == 0 ? 0 : pongEventOffset)); + } + bool waitWsFromMte3 = storeFinalState && std::is_same::value && + event0FromMte3[streamId]; + epilogueKdaFwdHVnew(gmV[vec1Offsets.uvOffset], gmVUpdateWorkspace[vec1Offsets.vWorkOffset], + gmG[vec1Offsets.gOffset], gmU[vec1Offsets.uvOffset], + gmVWorkspace[vec1Offsets.vWorkOffset], vec1Offsets.blockTokens, + vec1Offsets.vBlockDim, vHeadDim, vecBlockScheduler.cube1Done[streamId], + vecBlockScheduler.vec1Done[streamId], vec1Offsets.isInitialState, + vec1Offsets.isFinalState, storeFinalState, waitWsFromMte3, + (streamId == 0), tailVectorPath); + AscendC::SetFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID1); + if (storeFinalState && std::is_same::value) { + event0FromMte3[streamId] = false; + } + } + } else { + /* V2: h[i+1] += h_work if i < num_chunks - 1 else None */ + for (uint32_t i = 0; i < PING_PONG_STAGES; ++i) { + uint32_t streamId = vecBlockScheduler.GetStreamId(i); + const auto &stream = vecBlockScheduler.GetStream(i); + if (vecBlockScheduler.StreamIsDone(stream)) { + continue; + } + const KdaFwdHOffsets &vec2Offsets = vecBlockScheduler.GetCurTaskOffsets(stream); + if (vecBlockScheduler.NeedProcessStage2(stream)) { + bool vectorC2Path = vec2Offsets.blockTokens < 16; + if (vectorC2Path) { + Arch::CrossCoreWaitFlag(vecBlockScheduler.cube2Done[streamId]); + ComputeHWorkspaceVector(vec2Offsets, + EVENT_ID3 + (streamId == 0 ? 0 : pongEventOffset)); + } + if (storeFinalState && std::is_same::value) { + event0FromMte3[streamId] = true; + event2FromMte3[streamId] = !vec2Offsets.isFinalState; + } + // step 4: h[i+1] += h_work if i < num_chunks - 1 else None + epilogueKdaFwdHUpdate( + gmH[vec2Offsets.hDstOffset], gmFinalState[vec2Offsets.finalStateOffset], + gmG[vec2Offsets.gOffset], gmH[vec2Offsets.hSrcOffset], + gmHWorkspace[vec2Offsets.hWorkOffset], gmGk[vec2Offsets.gkOffset], + gmInitialState[vec2Offsets.initialStateOffset], vec2Offsets.blockTokens, kHeadDim, + vec2Offsets.vBlockDim, vHeadDim, vecBlockScheduler.cube2Done[streamId], + vec2Offsets.isInitialState, vec2Offsets.isFinalState, storeFinalState, + useInitialState, (streamId == 0), vectorC2Path); + AscendC::SetFlag(EVENT_ID1); + AscendC::WaitFlag(EVENT_ID1); + } else { + Arch::CrossCoreWaitFlag(vecBlockScheduler.cube2Done[streamId]); + } + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec2Done[streamId]); + } + } + waitStageFence = vecBlockScheduler.isRunning; + if (waitStageFence) { + AscendC::SetFlag(EVENT_ID1); + } + currStage ^= 0x01; + } + + DrainVectorPipelineEvents(event0FromMte3, event2FromMte3); + } + } + } +}; + +} // namespace Catlass::Gemm::Kernel diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_finalize_arch35.h b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_finalize_arch35.h new file mode 100644 index 0000000000..7f0fa9cb2e --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_finalize_arch35.h @@ -0,0 +1,1516 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#pragma once + +#ifndef CATLASS_ARCH +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +#define CATLASS_ARCH 3510 +#else +#define CATLASS_ARCH 2201 +#endif +#endif + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/arch/kda_cross_core_sync.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/gemm/block/kda_block_mmad.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "../kernel_utils/block/block_mmad_pingpong_tla_multi.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "kernel_operator.h" +#include "../chunk_kda_fwd_varlen.h" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +using namespace AscendC; + +namespace KdaFinalize { +namespace { +using KdaInt64 = tla::Int<64>; +using KdaInt128 = tla::Int<128>; +constexpr float LN2 = 0.69314718055994530942f; +constexpr float KDA_EXP2_CLAMP = 80.0f; +constexpr float KDA_EXP_INPUT_MAX = KDA_EXP2_CLAMP * LN2; +constexpr float KDA_EXP_INPUT_MIN = -KDA_EXP2_CLAMP * LN2; +constexpr float KDA_FP16_MAX = 65504.0f; +constexpr uint32_t EXP2_UB_ELEMENTS = 256; +constexpr uint32_t EXP2_UB_BYTES = EXP2_UB_ELEMENTS * (sizeof(float) + sizeof(uint16_t)); +constexpr uint32_t EXP2_EVENT_ID = 0; +constexpr uint32_t KDA_SOLVE_BT = 64; +constexpr uint32_t KDA_SOLVE_MATRIX_ELEMENTS = KDA_SOLVE_BT * KDA_SOLVE_BT; +constexpr uint32_t KDA_SOLVE_SCRATCH_X = 0; +constexpr uint32_t KDA_SOLVE_SCRATCH_Y0 = 1; +constexpr uint32_t KDA_SOLVE_SCRATCH_TMP = 2; +constexpr uint32_t KDA_SOLVE_SCRATCH_Y1 = 3; +constexpr uint32_t KDA_SOLVE_SCRATCH_IDENTITY = 4; +constexpr uint32_t KDA_SOLVE_SCRATCH_SLOTS = 5; +constexpr uint32_t KDA_SOLVE_DIAG_BT = 16; +constexpr uint32_t KDA_SOLVE_DIAG_BLOCKS = KDA_SOLVE_BT / KDA_SOLVE_DIAG_BT; +constexpr uint32_t KDA_SOLVE_DIAG_MCH_ITERS = 3; +constexpr uint32_t KDA_SCORE_REF_BC = 16; +constexpr uint32_t KDA_VEC_ARENA_ELEMENTS = 32768; +constexpr uint32_t KDA_BITS_PER_MASK_BYTE = 8; +constexpr uint32_t KDA_SELECT_COL_BLOCKS = 2; +constexpr uint32_t KDA_SELECT_COL_MASK_BYTES = KDA_SOLVE_MATRIX_ELEMENTS / KDA_BITS_PER_MASK_BYTE; +constexpr uint32_t KDA_SELECT_MASK_BYTES = KDA_SELECT_COL_BLOCKS * KDA_SELECT_COL_MASK_BYTES; +constexpr uint32_t KDA_SELECT_AQK_MASK_BYTE_OFFSET = 120 * 1024; +constexpr uint32_t KDA_SELECT_AKK_MASK_BYTE_OFFSET = KDA_SELECT_AQK_MASK_BYTE_OFFSET + KDA_SELECT_MASK_BYTES; +constexpr uint32_t KDA_SELECT_ZERO_BYTE_OFFSET = KDA_SELECT_AKK_MASK_BYTE_OFFSET + KDA_SELECT_MASK_BYTES; +constexpr uint32_t KDA_SELECT_ZERO_FLOAT_OFFSET = KDA_SELECT_ZERO_BYTE_OFFSET / sizeof(float); +constexpr uint8_t KDA_SCORE_DONE_FLAG0 = 2; +constexpr uint8_t KDA_SCORE_DONE_FLAG1 = 3; +constexpr uint8_t KDA_SCORE_READY_FLAG0 = 4; +constexpr uint8_t KDA_SCORE_READY_FLAG1 = 5; +constexpr uint32_t KDA_SCORE_QUEUE_DEPTH = 2; +constexpr uint32_t KDA_SYNC_REVERSE_DEPTH = 1; +constexpr uint32_t KDA_SCORE_SCRATCH_PLANES = 3; +constexpr uint32_t KDA_SCORE_SCRATCH_QG = 0; +constexpr uint32_t KDA_SCORE_SCRATCH_W = 1; +constexpr uint32_t KDA_SCORE_SCRATCH_KG = 2; +constexpr uint64_t KDA_WORKSPACE_ALIGN = 512; +constexpr uint32_t KDA_GATE_TILE_ROWS = 32; +constexpr uint32_t KDA_CUBE_MIN_REDUCTION = 16; + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +constexpr AscendC::MicroAPI::CastTrait KDA_TAIL_BF16_TO_FP32 = { + AscendC::MicroAPI::RegLayout::ZERO, + AscendC::MicroAPI::SatMode::SAT, + AscendC::MicroAPI::MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_NONE, +}; + +static __simd_vf__ inline void ComputeKdaTailOutputRegbase(__ubuf__ float *dst, __ubuf__ bfloat16_t *h, + __ubuf__ bfloat16_t *vNew, __ubuf__ bfloat16_t *qg, + __ubuf__ bfloat16_t *aqk, uint16_t rows, uint16_t curT, + float scale) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t DIM = 128; + constexpr uint16_t FP32_REG_ELEMENTS = AscendC::VECTOR_REG_WIDTH / sizeof(float); + constexpr uint16_t AQK_ROW_ELEMENTS = 16; + MaskReg fullMask = CreateMask(); + RegTensor scaleReg; + Duplicate(scaleReg, scale, fullMask); + + for (uint16_t row = 0; row < rows; ++row) { + RegTensor acc0; + RegTensor acc1; + Duplicate(acc0, 0.0f, fullMask); + Duplicate(acc1, 0.0f, fullMask); + + for (uint16_t d = 0; d < DIM; ++d) { + RegTensor matrixRaw0; + RegTensor matrixRaw1; + RegTensor weightRaw; + RegTensor matrix0; + RegTensor matrix1; + RegTensor weight; + RegTensor product0; + RegTensor product1; + DataCopy(matrixRaw0, h + static_cast(d) * DIM); + DataCopy(matrixRaw1, + h + static_cast(d) * DIM + FP32_REG_ELEMENTS); + LoadAlign(weightRaw, qg + static_cast(row) * DIM + d); + Cast(matrix0, matrixRaw0, fullMask); + Cast(matrix1, matrixRaw1, fullMask); + Cast(weight, weightRaw, fullMask); + Mul(product0, matrix0, weight, fullMask); + Mul(product1, matrix1, weight, fullMask); + Add(acc0, acc0, product0, fullMask); + Add(acc1, acc1, product1, fullMask); + } + Mul(acc0, acc0, scaleReg, fullMask); + Mul(acc1, acc1, scaleReg, fullMask); + + for (uint16_t token = 0; token < curT; ++token) { + RegTensor matrixRaw0; + RegTensor matrixRaw1; + RegTensor weightRaw; + RegTensor matrix0; + RegTensor matrix1; + RegTensor weight; + RegTensor product0; + RegTensor product1; + DataCopy(matrixRaw0, vNew + static_cast(token) * DIM); + DataCopy( + matrixRaw1, vNew + static_cast(token) * DIM + FP32_REG_ELEMENTS); + LoadAlign(weightRaw, + aqk + static_cast(row) * AQK_ROW_ELEMENTS + token); + Cast(matrix0, matrixRaw0, fullMask); + Cast(matrix1, matrixRaw1, fullMask); + Cast(weight, weightRaw, fullMask); + Mul(product0, matrix0, weight, fullMask); + Mul(product1, matrix1, weight, fullMask); + Add(acc0, acc0, product0, fullMask); + Add(acc1, acc1, product1, fullMask); + } + + DataCopy(dst + static_cast(row) * DIM, acc0, fullMask); + DataCopy(dst + static_cast(row) * DIM + FP32_REG_ELEMENTS, acc1, fullMask); + } +} +#endif + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +using KdaArchTag = Catlass::Arch::Ascend950; +#else +using KdaArchTag = Catlass::Arch::AtlasA2; +#endif +using KdaDispatchPolicy = Catlass::Gemm::MmadPingpong; +using KdaScoreDispatchPolicy = Catlass::Gemm::MmadPingpongTlaMulti; +static_assert(KdaScoreDispatchPolicy::ENABLE_L1_RESIDENT, + "KDA Aqk/Akk score MMAD must keep the shared right matrix resident in L1"); +static_assert(KdaScoreDispatchPolicy::L1B_STAGES == 1, + "KDA Aqk/Akk score MMAD needs one L1 B slot so the second MMAD reuses it"); +using KdaSolveDispatchPolicy = Catlass::Gemm::MmadPingpong; +static_assert(!KdaSolveDispatchPolicy::USE_HF32_MODE, "KDA triangular solve must use IEEE FP32 Cube mode"); +using KdaL1TileShape = tla::Shape; +using KdaL0TileShape = KdaL1TileShape; +using KdaSolveL1TileShape = tla::Shape; +using KdaSolveL0TileShape = KdaSolveL1TileShape; + +__aicore__ inline uint32_t FloatToBits(float value) +{ + union Bits { + __aicore__ Bits() {} + float f; + uint32_t u; + } bits; + bits.f = value; + return bits.u; +} + +__aicore__ inline float BitsToFloat(uint32_t value) +{ + union Bits { + __aicore__ Bits() {} + uint32_t u; + float f; + } bits; + bits.u = value; + return bits.f; +} + +__aicore__ inline uint16_t Bf16ToBits(bfloat16_t value) +{ + union Bits { + __aicore__ Bits() {} + bfloat16_t f; + uint16_t u; + } bits; + bits.f = value; + return bits.u; +} + +__aicore__ inline bfloat16_t BitsToBf16(uint16_t value) +{ + union Bits { + __aicore__ Bits() {} + uint16_t u; + bfloat16_t f; + } bits; + bits.u = value; + return bits.f; +} + +template +__aicore__ inline T FloatToType(float value) +{ + if constexpr (IsSameType::value) { + uint32_t bits = FloatToBits(value); + uint32_t bias = 0x7FFFu + ((bits >> 16) & 1u); + return BitsToBf16(static_cast((bits + bias) >> 16)); + } + return static_cast(value); +} + +template +class ChunkKdaFwdFinalizeKernel { +public: + using OUT_T = float; + using AKK_T = float; + template + __aicore__ inline void Init(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR beta, GM_ADDR initialState, + GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR preparedQG, GM_ADDR preparedAqk, + GM_ADDR propagatedVNew, GM_ADDR propagatedH, GM_ADDR o, GM_ADDR finalState, GM_ADDR aqk, + GM_ADDR akk, GM_ADDR w, GM_ADDR u, GM_ADDR qg, GM_ADDR kg, GM_ADDR vNew, GM_ADDR h, + GM_ADDR workspace, const TilingData &tiling, TPipe *pipe, bool initVecBuffers = true) + { + pipe_ = pipe; + q_.SetGlobalBuffer((__gm__ T *)q); + k_.SetGlobalBuffer((__gm__ T *)k); + v_.SetGlobalBuffer((__gm__ T *)v); + gk_.SetGlobalBuffer((__gm__ GK_T *)gk); + if (initialState != nullptr) { + initialState_.SetGlobalBuffer((__gm__ float *)initialState); + } + cuSeqlensAddr_ = reinterpret_cast<__gm__ int64_t *>(cuSeqlens); + if (preparedQG != nullptr) { + preparedQG_.SetGlobalBuffer((__gm__ T *)preparedQG); + } + if (preparedAqk != nullptr) { + preparedAqk_.SetGlobalBuffer((__gm__ T *)preparedAqk); + } + if (propagatedVNew != nullptr) { + propagatedVNew_.SetGlobalBuffer((__gm__ T *)propagatedVNew); + } + if (propagatedH != nullptr) { + propagatedH_.SetGlobalBuffer((__gm__ T *)propagatedH); + } + chunkIndicesAddr_ = reinterpret_cast<__gm__ int64_t *>(chunkIndices); + o_.SetGlobalBuffer((__gm__ OUT_T *)o); + finalState_.SetGlobalBuffer((__gm__ float *)finalState); + aqk_.SetGlobalBuffer((__gm__ float *)aqk); + akk_.SetGlobalBuffer((__gm__ AKK_T *)akk); + w_.SetGlobalBuffer((__gm__ T *)w); + u_.SetGlobalBuffer((__gm__ OUT_T *)u); + qg_.SetGlobalBuffer((__gm__ T *)qg); + kg_.SetGlobalBuffer((__gm__ T *)kg); + vNew_.SetGlobalBuffer((__gm__ T *)vNew); + h_.SetGlobalBuffer((__gm__ float *)h); + solveWorkspace_.SetGlobalBuffer((__gm__ float *)workspace); + + B_ = tiling.batch; + N_ = tiling.seqNum; + H_ = tiling.qHeadNum; + HV_ = tiling.vHeadNum; + T_ = tiling.seqlen; + K_ = tiling.kHeadDim; + V_ = tiling.vHeadDim; + BT_ = tiling.chunkSize; + NT_ = tiling.totalChunks; + scale_ = tiling.scale; + hasInitial_ = tiling.hasInitialState; + isVarLen_ = tiling.isVarLen; + usedCoreNum_ = tiling.outputUsedCoreNum; + const uint64_t outputElements = B_ * HV_ * T_ * V_; + o_.SetGlobalBuffer((__gm__ OUT_T *)workspace); + u_.SetGlobalBuffer((__gm__ OUT_T *)workspace + outputElements); + if ASCEND_IS_AIV { + uint64_t subBlockNum = static_cast(GetSubBlockNum()); + solveCoreIdx_ = subBlockNum == 0 ? 0 : static_cast(GetBlockIdx()) / subBlockNum; + } else { + solveCoreIdx_ = static_cast(GetBlockIdx()); + } + if (pipe_ != nullptr && initVecBuffers) { + pipe_->InitBuffer(exp2Buf_, EXP2_UB_BYTES); + pipe_->InitBuffer(vecBuf_, KDA_VEC_ARENA_ELEMENTS * sizeof(float)); + const uint64_t gateWritebackRows = ScoreVectorMaxRows(5 * sizeof(float) + 2 * sizeof(T) + sizeof(GK_T)); + pipe_->InitBuffer(gateWritebackBuf_, + static_cast(gateWritebackRows * K_ * (3 * sizeof(T) + sizeof(GK_T)))); + AllocVectorEvents(); + } + } + __aicore__ inline void ProcessAiv() + { + ProcessOutAiv(); + ReleaseVectorEvents(); + } + + __aicore__ inline void ProcessAic() { ProcessOutAic(); } + +private: + __aicore__ inline void AllocVectorEvents() + { + mte2ToVEvent_ = pipe_->AllocEventID(); + vToMte2Event_ = pipe_->AllocEventID(); + vToMte3Event_ = pipe_->AllocEventID(); + mte3ToVEvent_ = pipe_->AllocEventID(); + mte2ToMte3Event_ = pipe_->AllocEventID(); + mte3ToMte2Event_ = pipe_->AllocEventID(); + sToVEvent_ = pipe_->AllocEventID(); + sToMte2Event_ = pipe_->AllocEventID(); + vectorEventsAllocated_ = true; + } + + __aicore__ inline void ReleaseVectorEvents() + { + if (!vectorEventsAllocated_) { + return; + } + pipe_->ReleaseEventID(mte2ToVEvent_); + pipe_->ReleaseEventID(vToMte2Event_); + pipe_->ReleaseEventID(vToMte3Event_); + pipe_->ReleaseEventID(mte3ToVEvent_); + pipe_->ReleaseEventID(mte2ToMte3Event_); + pipe_->ReleaseEventID(mte3ToMte2Event_); + pipe_->ReleaseEventID(sToVEvent_); + pipe_->ReleaseEventID(sToMte2Event_); + vectorEventsAllocated_ = false; + } + + __aicore__ inline uint64_t QOffset(uint64_t b, uint64_t h, uint64_t t, uint64_t d) const + { + return ((b * H_ + h) * T_ + t) * K_ + d; + } + + __aicore__ inline uint64_t KVOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t d, uint64_t dim) const + { + return ((b * HV_ + hv) * T_ + t) * dim + d; + } + + __aicore__ inline uint64_t OutputOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t d) const + { + return ((b * T_ + t) * HV_ + hv) * V_ + d; + } + + __aicore__ inline uint64_t BetaOffset(uint64_t b, uint64_t hv, uint64_t t) const { return (b * HV_ + hv) * T_ + t; } + + __aicore__ inline uint64_t AOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t j) const + { + return ((b * HV_ + hv) * T_ + t) * BT_ + j; + } + + __aicore__ inline uint64_t HOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t d, uint64_t r) const + { + return (((b * HV_ + hv) * NT_ + chunkIdx) * K_ + d) * V_ + r; + } + + __aicore__ inline uint64_t WScratchOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t t, uint64_t d) const + { + return (((b * HV_ + hv) * NT_ + chunkIdx) * BT_ + t) * K_ + d; + } + + __aicore__ inline uint64_t SolveScratchOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t slot) const + { + (void)b; + (void)hv; + (void)chunkIdx; + uint64_t matrixElements = BT_ * BT_; + return solveCoreIdx_ * KDA_SOLVE_SCRATCH_SLOTS * matrixElements + slot * matrixElements; + } + + __aicore__ inline uint64_t ScoreScratchOffset(uint64_t slot, uint64_t plane, uint64_t t = 0, uint64_t d = 0) const + { + return (((solveCoreIdx_ * KDA_SCORE_QUEUE_DEPTH + slot) * KDA_SCORE_SCRATCH_PLANES + plane) * BT_ + t) * K_ + d; + } + + __aicore__ inline uint64_t ScoreRefBlockSize() const { return KDA_SCORE_REF_BC; } + + __aicore__ inline uint64_t ScoreRowBlockCount(uint64_t curT, uint64_t rowBegin) const + { + uint64_t blockSize = ScoreRefBlockSize(); + uint64_t rowCount = curT - rowBegin; + if (rowCount > blockSize) { + rowCount = blockSize; + } + return rowCount; + } + + __aicore__ inline uint64_t ScoreRefToken(uint64_t start, uint64_t curT, uint64_t rowBegin, uint64_t rowCount) const + { + uint64_t ref = rowBegin + rowCount / 2; + if (ref >= curT) { + ref = curT - 1; + } + return start + ref; + } + + __aicore__ inline void RunExp2(LocalTensor &tensor, uint32_t count) + { + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + ClampExpInput(tensor, count); + Exp(tensor, tensor, count); + PipeBarrier(); + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + } + + __aicore__ inline void ClampExpInput(LocalTensor &tensor, uint32_t count) + { + Mins(tensor, tensor, KDA_EXP_INPUT_MAX, count); + PipeBarrier(); + Maxs(tensor, tensor, KDA_EXP_INPUT_MIN, count); + PipeBarrier(); + } + + __aicore__ inline void ClampFp32ToOutputType(LocalTensor &tensor, uint32_t count) + { + if constexpr (IsSameType::value) { + Mins(tensor, tensor, KDA_FP16_MAX, count); + PipeBarrier(); + Maxs(tensor, tensor, -KDA_FP16_MAX, count); + PipeBarrier(); + } + } + + template + __aicore__ inline void CopyVectorIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset, + uint64_t count) + { + uint64_t rowBytes = count * static_cast(sizeof(CopyT)); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst, src[offset], static_cast(count)); + return; + } + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPadParams padParams{false, 0, 0, 0}; + DataCopyPad(dst, src[offset], params, padParams); + } + + template + __aicore__ inline void CopyVectorOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src, + uint64_t count) + { + uint64_t rowBytes = count * static_cast(sizeof(CopyT)); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst[offset], src, static_cast(count)); + return; + } + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPad(dst[offset], src, params); + } + + template + __aicore__ inline void CopyRowsOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src, + uint64_t rows, uint64_t cols, uint64_t dstStride) + { + if (cols == dstStride) { + CopyVectorOut(dst, offset, src, rows * cols); + return; + } + constexpr uint64_t blockBytes = 32; + const uint64_t rowBytes = cols * sizeof(CopyT); + const uint64_t gapBytes = (dstStride - cols) * sizeof(CopyT); + DataCopyParams params{ +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + 1, +#else + static_cast(rows), +#endif + static_cast(rowBytes / blockBytes), 0, +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + 0 +#else + static_cast(gapBytes / blockBytes) +#endif + }; +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + const uint64_t dstRowBytes = dstStride * sizeof(CopyT); + LoopModeParams loopParams{static_cast(rows), 1, rowBytes, dstRowBytes, 0, 0}; + // Loop-mode registers are core-local state and must not leak across DMA calls. + ResetLoopModePara(DataCopyMVType::UB_TO_OUT); + SetLoopModePara(loopParams, DataCopyMVType::UB_TO_OUT); + DataCopy(dst[offset], src, params); + ResetLoopModePara(DataCopyMVType::UB_TO_OUT); +#else + DataCopy(dst[offset], src, params); +#endif + } + + template + __aicore__ inline void CopyRowIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset) + { + CopyVectorIn(dst, src, offset, K_); + } + + template + __aicore__ inline void CopyRowOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src) + { + CopyVectorOut(dst, offset, src, K_); + } + + __aicore__ inline LocalTensor VecScratch(uint64_t slot) + { + return vecBuf_.Get()[slot * EXP2_UB_ELEMENTS]; + } + + template + __aicore__ inline void LoadAsFloatRow(GlobalTensor &src, uint64_t srcOffset, LocalTensor &dst, + uint64_t count) + { + if constexpr (IsSameType::value) { + CopyVectorIn(dst, src, srcOffset, count); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Adds(dst, dst, 0.0f, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } else { + constexpr uint32_t typedOffset = EXP2_UB_ELEMENTS * sizeof(float) / sizeof(CopyT); + LocalTensor rowLocal = exp2Buf_.Get()[typedOffset]; + CopyVectorIn(rowLocal, src, srcOffset, count); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Cast(dst, rowLocal, RoundMode::CAST_NONE, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } + PipeBarrier(); + } + + __aicore__ inline void ComputeTailLocalRows(LocalTensor &dst, uint64_t b, uint64_t hv, uint64_t start, + uint64_t curT, uint64_t rowBegin, uint64_t rows) + { + LocalTensor vRow = exp2Buf_.Get(); + LocalTensor coefficientTyped = gateWritebackBuf_.Get(); + LocalTensor coefficients = gateWritebackBuf_.Get()[BT_]; + for (uint64_t localRow = 0; localRow < rows; ++localRow) { + LocalTensor dstRow = dst[localRow * V_]; + CopyVectorIn(coefficientTyped, preparedAqk_, AOffset(b, hv, start + rowBegin + localRow, 0), curT); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Cast(coefficients, coefficientTyped, RoundMode::CAST_NONE, static_cast(curT)); + // Earlier megakernel stages reuse V_S event IDs. Drain the cast + // before the first scalar coefficient read in this tail row. + PipeBarrier(); + Duplicate(dstRow, 0.0f, static_cast(V_)); + PipeBarrier(); + for (uint64_t j = 0; j < curT; ++j) { + LoadAsFloatRow(propagatedVNew_, KVOffset(b, hv, start + j, 0, V_), vRow, V_); + float weight = coefficients.GetValue(j); + SetFlag(sToVEvent_); + WaitFlag(sToVEvent_); + Muls(vRow, vRow, weight, static_cast(V_)); + PipeBarrier(); + Add(dstRow, dstRow, vRow, static_cast(V_)); + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } + SetFlag(sToMte2Event_); + WaitFlag(sToMte2Event_); + } + } + + __aicore__ inline void ComputeTailStateRows(LocalTensor &dst, uint64_t b, uint64_t hv, uint64_t chunkIdx, + uint64_t start, uint64_t rowBegin, uint64_t rows) + { + LocalTensor hRow = exp2Buf_.Get(); + LocalTensor coefficientTyped = gateWritebackBuf_.Get(); + LocalTensor coefficients = gateWritebackBuf_.Get()[BT_]; + for (uint64_t localRow = 0; localRow < rows; ++localRow) { + LocalTensor dstRow = dst[localRow * V_]; + CopyVectorIn(coefficientTyped, preparedQG_, KVOffset(b, hv, start + rowBegin + localRow, 0, K_), K_); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Cast(coefficients, coefficientTyped, RoundMode::CAST_NONE, static_cast(K_)); + // Earlier megakernel stages reuse V_S event IDs. Drain the cast + // before the first scalar coefficient read in this tail row. + PipeBarrier(); + Duplicate(dstRow, 0.0f, static_cast(V_)); + PipeBarrier(); + for (uint64_t d = 0; d < K_; ++d) { + LoadAsFloatRow(propagatedH_, HOffset(b, hv, chunkIdx, d, 0), hRow, V_); + float weight = coefficients.GetValue(d); + SetFlag(sToVEvent_); + WaitFlag(sToVEvent_); + Muls(hRow, hRow, weight, static_cast(V_)); + PipeBarrier(); + Add(dstRow, dstRow, hRow, static_cast(V_)); + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } + SetFlag(sToMte2Event_); + WaitFlag(sToMte2Event_); + } + } + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + __aicore__ inline void ComputeTailOutputRegbaseRows(LocalTensor &dst, uint64_t b, uint64_t hv, + uint64_t chunkIdx, uint64_t start, uint64_t curT, + uint64_t rowBegin, uint64_t rows) + { + // Keep sub-16 tails on one MTE2-to-V handoff; the scalar event chain can retain stale state. + constexpr uint64_t dim = 128; + constexpr uint64_t maxTailRows = KDA_CUBE_MIN_REDUCTION - 1; + constexpr uint64_t aqkStagedCols = KDA_CUBE_MIN_REDUCTION; + constexpr uint64_t hOffset = 0; + constexpr uint64_t vNewOffset = hOffset + dim * dim; + constexpr uint64_t qgOffset = vNewOffset + maxTailRows * dim; + constexpr uint64_t aqkOffset = qgOffset + maxTailRows * dim; + LocalTensor staging = gateWritebackBuf_.Get(); + LocalTensor hLocal = staging[hOffset]; + LocalTensor vNewLocal = staging[vNewOffset]; + LocalTensor qgLocal = staging[qgOffset]; + LocalTensor aqkLocal = staging[aqkOffset]; + + CopyVectorIn(hLocal, propagatedH_, HOffset(b, hv, chunkIdx, 0, 0), dim * dim); + CopyVectorIn(vNewLocal, propagatedVNew_, KVOffset(b, hv, start, 0, V_), curT * dim); + CopyVectorIn(qgLocal, preparedQG_, KVOffset(b, hv, start + rowBegin, 0, K_), rows * dim); + for (uint64_t row = 0; row < rows; ++row) { + LocalTensor aqkRow = aqkLocal[row * aqkStagedCols]; + CopyVectorIn(aqkRow, preparedAqk_, AOffset(b, hv, start + rowBegin + row, 0), aqkStagedCols); + } + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + AscendC::VF_CALL(reinterpret_cast<__ubuf__ float *>(dst.GetPhyAddr()), + reinterpret_cast<__ubuf__ bfloat16_t *>(hLocal.GetPhyAddr()), + reinterpret_cast<__ubuf__ bfloat16_t *>(vNewLocal.GetPhyAddr()), + reinterpret_cast<__ubuf__ bfloat16_t *>(qgLocal.GetPhyAddr()), + reinterpret_cast<__ubuf__ bfloat16_t *>(aqkLocal.GetPhyAddr()), + static_cast(rows), static_cast(curT), scale_); + PipeBarrier(); + } +#endif + + template + __aicore__ inline void LoadAsFloatVector(GlobalTensor &src, uint64_t srcOffset, LocalTensor &dst, + LocalTensor &typedScratch, uint64_t count) + { + if constexpr (IsSameType::value) { + CopyVectorIn(dst, src, srcOffset, count); + } else { + CopyVectorIn(typedScratch, src, srcOffset, count); + } + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + if constexpr (!IsSameType::value) { + Cast(dst, typedScratch, RoundMode::CAST_NONE, static_cast(count)); + PipeBarrier(); + } + } + + template + __aicore__ inline void StoreFloatRow(GlobalTensor &dst, uint64_t dstOffset, LocalTensor &src, + uint64_t count) + { + if constexpr (IsSameType::value) { + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(dst, dstOffset, src, count); + } else { + constexpr uint32_t typedOffset = EXP2_UB_ELEMENTS * sizeof(float) / sizeof(CopyT); + LocalTensor rowLocal = exp2Buf_.Get()[typedOffset]; + Cast(rowLocal, src, RoundMode::CAST_RINT, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(dst, dstOffset, rowLocal, count); + } + SetFlag(mte3ToMte2Event_); + WaitFlag(mte3ToMte2Event_); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline LocalTensor Exp2NegG(uint64_t b, uint64_t hv, uint64_t t) + { + LocalTensor exp2Local = exp2Buf_.Get(); + LoadAsFloatRow(gk_, KVOffset(b, hv, t, 0, K_), exp2Local, K_); + Muls(exp2Local, exp2Local, -LN2, static_cast(K_)); + PipeBarrier(); + RunExp2(exp2Local, static_cast(K_)); + return exp2Local; + } + + __aicore__ inline uint64_t ScoreVectorMaxRows(uint64_t bytesPerElem) const + { + constexpr uint64_t arenaBytes = static_cast(KDA_VEC_ARENA_ELEMENTS) * sizeof(float); + uint64_t maxRows = (arenaBytes / bytesPerElem) / K_; + if (K_ >= 128 && maxRows > 32) { + maxRows = 32; + } + return maxRows; + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + __aicore__ inline void ComputeOutputCubeStagedArch35(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t curT) + { + SetMMLayoutTransform(true); + using ElementA = T; + using ElementB = T; + using ElementC = OUT_T; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::RowMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; + using TileMmad = Catlass::Gemm::Tile::TileMmadTla; + + constexpr uint16_t kMte2Event = 0; + constexpr uint16_t kMte1Event = 0; + constexpr uint16_t kMmadEvent = 0; + constexpr uint32_t kL1A0Offset = 0; + constexpr uint32_t kL1B0Offset = 64 * 128 * sizeof(ElementA); + constexpr uint32_t kL1A1Offset = kL1B0Offset + 128 * 128 * sizeof(ElementB); + constexpr uint32_t kL1B1Offset = kL1A1Offset + 64 * 64 * sizeof(ElementA); + + Catlass::Arch::Resource resource; + LocalTensor l1A0 = resource.l1Buf.template GetBufferByByte(kL1A0Offset); + LocalTensor l1B0 = resource.l1Buf.template GetBufferByByte(kL1B0Offset); + LocalTensor l1A1 = resource.l1Buf.template GetBufferByByte(kL1A1Offset); + LocalTensor l1B1 = resource.l1Buf.template GetBufferByByte(kL1B1Offset); + LocalTensor l0A = resource.l0ABuf.template GetBufferByByte(0); + LocalTensor l0B = resource.l0BBuf.template GetBufferByByte(0); + LocalTensor l0C = resource.l0CBuf.template GetBufferByByte(0); + + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + TileMmad tileMmad; + + auto layoutQ = tla::MakeLayout(BT_, K_); + auto layoutH = tla::MakeLayout(K_, V_); + auto layoutO = tla::MakeLayout(BT_, V_); + auto layoutAqk = tla::MakeLayout(BT_, BT_); + auto layoutV = tla::MakeLayout(BT_, V_); + + for (uint64_t nOffset = 0; nOffset < V_; nOffset += 128) { + const uint32_t curN = static_cast((V_ - nOffset) > 128 ? 128 : (V_ - nOffset)); + auto tensorH = tla::MakeTensor(propagatedH_[HOffset(b, hv, chunkIdx, 0, nOffset)], layoutH, + Catlass::Arch::PositionGM{}); + auto tensorVNew = tla::MakeTensor(propagatedVNew_[KVOffset(b, hv, start, nOffset, V_)], layoutV, + Catlass::Arch::PositionGM{}); + + for (uint64_t mOffset = 0; mOffset < curT; mOffset += 64) { + const uint32_t curM = static_cast((curT - mOffset) > 64 ? 64 : (curT - mOffset)); + auto tensorQ = tla::MakeTensor(preparedQG_[KVOffset(b, hv, start + mOffset, 0, K_)], layoutQ, + Catlass::Arch::PositionGM{}); + auto tensorAqk = tla::MakeTensor(preparedAqk_[AOffset(b, hv, start + mOffset, 0)], layoutAqk, + Catlass::Arch::PositionGM{}); + auto tensorO = tla::MakeTensor(o_[KVOffset(b, hv, start + mOffset, nOffset, V_)], layoutO, + Catlass::Arch::PositionGM{}); + auto tensorLocal = tla::MakeTensor(u_[KVOffset(b, hv, start + mOffset, nOffset, V_)], layoutO, + Catlass::Arch::PositionGM{}); + + auto blockQ = GetTile(tensorQ, tla::MakeCoord(0, 0), tla::MakeShape(curM, K_)); + auto blockH = GetTile(tensorH, tla::MakeCoord(0, 0), tla::MakeShape(K_, curN)); + auto blockAqk = GetTile(tensorAqk, tla::MakeCoord(0, 0), tla::MakeShape(curM, curT)); + auto blockVNew = GetTile(tensorVNew, tla::MakeCoord(0, 0), tla::MakeShape(curT, curN)); + auto blockO = GetTile(tensorO, tla::MakeCoord(0, 0), tla::MakeShape(curM, curN)); + auto blockLocal = GetTile(tensorLocal, tla::MakeCoord(0, 0), tla::MakeShape(curM, curN)); + + using CopyGmToL1A0 = typename TileCopy::template CopyGmToL1A; + using CopyGmToL1B0 = typename TileCopy::template CopyGmToL1B; + using CopyGmToL1A1 = typename TileCopy::template CopyGmToL1A; + using CopyGmToL1B1 = typename TileCopy::template CopyGmToL1B; + using CopyL0CToDst = typename TileCopy::template CopyL0CToDst; + CopyGmToL1A0 copyGmToL1A0; + CopyGmToL1B0 copyGmToL1B0; + CopyGmToL1A1 copyGmToL1A1; + CopyGmToL1B1 copyGmToL1B1; + CopyL0CToDst copyL0CToDst; + + auto layoutL1A0 = tla::MakeLayout(curM, K_); + auto layoutL1B0 = tla::MakeLayout(K_, curN); + auto layoutL1A1 = tla::MakeLayout(curM, curT); + auto layoutL1B1 = tla::MakeLayout(curT, curN); + auto layoutL0A0 = tla::MakeLayout(curM, K_); + auto layoutL0B0 = tla::MakeLayout(K_, curN); + auto layoutL0A1 = tla::MakeLayout(curM, curT); + auto layoutL0B1 = tla::MakeLayout(curT, curN); + auto layoutL0C = tla::MakeLayoutL0C(curM, curN); + + auto tensorL1A0 = tla::MakeTensor(l1A0, layoutL1A0, Catlass::Arch::PositionL1{}); + auto tensorL1B0 = tla::MakeTensor(l1B0, layoutL1B0, Catlass::Arch::PositionL1{}); + auto tensorL1A1 = tla::MakeTensor(l1A1, layoutL1A1, Catlass::Arch::PositionL1{}); + auto tensorL1B1 = tla::MakeTensor(l1B1, layoutL1B1, Catlass::Arch::PositionL1{}); + auto tensorL0A0 = tla::MakeTensor(l0A, layoutL0A0, Catlass::Arch::PositionL0A{}); + auto tensorL0B0 = tla::MakeTensor(l0B, layoutL0B0, Catlass::Arch::PositionL0B{}); + auto tensorL0A1 = tla::MakeTensor(l0A, layoutL0A1, Catlass::Arch::PositionL0A{}); + auto tensorL0B1 = tla::MakeTensor(l0B, layoutL0B1, Catlass::Arch::PositionL0B{}); + auto tensorL0C = tla::MakeTensor(l0C, layoutL0C, Catlass::Arch::PositionL0C{}); + uint32_t localRow = 0; + uint32_t localColumn = 0; + auto tileL1A0 = GetTile(tensorL1A0, tla::MakeCoord(localRow, localColumn), tla::MakeShape(curM, K_)); + auto tileL1B0 = GetTile(tensorL1B0, tla::MakeCoord(localRow, localColumn), tla::MakeShape(K_, curN)); + auto tileL1A1 = GetTile(tensorL1A1, tla::MakeCoord(localRow, localColumn), tla::MakeShape(curM, curT)); + auto tileL1B1 = GetTile(tensorL1B1, tla::MakeCoord(localRow, localColumn), tla::MakeShape(curT, curN)); + auto tileL0A0 = GetTile(tensorL0A0, tla::MakeCoord(localRow, localColumn), tla::MakeShape(curM, K_)); + auto tileL0B0 = GetTile(tensorL0B0, tla::MakeCoord(localRow, localColumn), tla::MakeShape(K_, curN)); + auto tileL0A1 = GetTile(tensorL0A1, tla::MakeCoord(localRow, localColumn), tla::MakeShape(curM, curT)); + auto tileL0B1 = GetTile(tensorL0B1, tla::MakeCoord(localRow, localColumn), tla::MakeShape(curT, curN)); + auto tileL0C = GetTile(tensorL0C, tla::MakeCoord(localRow, localColumn), tla::MakeShape(curM, curN)); + + copyGmToL1A0(tensorL1A0, blockQ); + copyGmToL1B0(tensorL1B0, blockH); + copyGmToL1A1(tensorL1A1, blockAqk); + copyGmToL1B1(tensorL1B1, blockVNew); + SetFlag(kMte2Event); + WaitFlag(kMte2Event); + + copyL1ToL0A(tileL0A0, tileL1A0); + copyL1ToL0B(tileL0B0, tileL1B0); + SetFlag(kMte1Event); + WaitFlag(kMte1Event); + tileMmad(tileL0C, tileL0A0, tileL0B0, curM, curN, static_cast(K_), true, 0b11); + SetFlag(kMmadEvent); + WaitFlag(kMmadEvent); + copyL0CToDst(blockO, tileL0C, 0b11); + PipeBarrier(); + + copyL1ToL0A(tileL0A1, tileL1A1); + copyL1ToL0B(tileL0B1, tileL1B1); + SetFlag(kMte1Event); + SetFlag(kMte2Event); + WaitFlag(kMte1Event); + WaitFlag(kMte2Event); + tileMmad(tileL0C, tileL0A1, tileL0B1, curM, curN, static_cast(curT), true, 0b11); + SetFlag(kMmadEvent); + WaitFlag(kMmadEvent); + copyL0CToDst(blockLocal, tileL0C, 0b11); + PipeBarrier(); + } + } + SetMMLayoutTransform(false); + } + + __aicore__ inline void PrefetchOutputTileArch35(Catlass::Arch::Resource &resource, uint32_t slot, + uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t curT, uint64_t nOffset, bool reuseSlot) + { + using ElementA = T; + using ElementB = T; + using ElementC = OUT_T; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::RowMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + + constexpr uint32_t kL1SlotBytes = 96 * 1024; + constexpr uint32_t kL1A0Offset = 0; + constexpr uint32_t kL1B0Offset = 64 * 128 * sizeof(ElementA); + constexpr uint32_t kL1A1Offset = kL1B0Offset + 128 * 128 * sizeof(ElementB); + constexpr uint32_t kL1B1Offset = kL1A1Offset + 64 * 64 * sizeof(ElementA); + const uint32_t slotBase = slot * kL1SlotBytes; + const uint32_t curM = static_cast(curT); + const uint32_t curN = static_cast((V_ - nOffset) > 128 ? 128 : (V_ - nOffset)); + + auto layoutQ = tla::MakeLayout(BT_, K_); + auto layoutH = tla::MakeLayout(K_, V_); + auto layoutAqk = tla::MakeLayout(BT_, BT_); + auto layoutV = tla::MakeLayout(BT_, V_); + auto tensorQ = + tla::MakeTensor(preparedQG_[KVOffset(b, hv, start, 0, K_)], layoutQ, Catlass::Arch::PositionGM{}); + auto tensorH = + tla::MakeTensor(propagatedH_[HOffset(b, hv, chunkIdx, 0, nOffset)], layoutH, Catlass::Arch::PositionGM{}); + auto tensorAqk = + tla::MakeTensor(preparedAqk_[AOffset(b, hv, start, 0)], layoutAqk, Catlass::Arch::PositionGM{}); + auto tensorVNew = + tla::MakeTensor(propagatedVNew_[KVOffset(b, hv, start, nOffset, V_)], layoutV, Catlass::Arch::PositionGM{}); + auto blockQ = GetTile(tensorQ, tla::MakeCoord(0, 0), tla::MakeShape(curM, K_)); + auto blockH = GetTile(tensorH, tla::MakeCoord(0, 0), tla::MakeShape(K_, curN)); + auto blockAqk = GetTile(tensorAqk, tla::MakeCoord(0, 0), tla::MakeShape(curM, curT)); + auto blockVNew = GetTile(tensorVNew, tla::MakeCoord(0, 0), tla::MakeShape(curT, curN)); + + using CopyGmToL1A0 = typename TileCopy::template CopyGmToL1A; + using CopyGmToL1B0 = typename TileCopy::template CopyGmToL1B; + using CopyGmToL1A1 = typename TileCopy::template CopyGmToL1A; + using CopyGmToL1B1 = typename TileCopy::template CopyGmToL1B; + CopyGmToL1A0 copyGmToL1A0; + CopyGmToL1B0 copyGmToL1B0; + CopyGmToL1A1 copyGmToL1A1; + CopyGmToL1B1 copyGmToL1B1; + + LocalTensor l1A0 = resource.l1Buf.template GetBufferByByte(slotBase + kL1A0Offset); + LocalTensor l1B0 = resource.l1Buf.template GetBufferByByte(slotBase + kL1B0Offset); + LocalTensor l1A1 = resource.l1Buf.template GetBufferByByte(slotBase + kL1A1Offset); + LocalTensor l1B1 = resource.l1Buf.template GetBufferByByte(slotBase + kL1B1Offset); + auto tensorL1A0 = + tla::MakeTensor(l1A0, tla::MakeLayout(curM, K_), Catlass::Arch::PositionL1{}); + auto tensorL1B0 = + tla::MakeTensor(l1B0, tla::MakeLayout(K_, curN), Catlass::Arch::PositionL1{}); + auto tensorL1A1 = + tla::MakeTensor(l1A1, tla::MakeLayout(curM, curT), Catlass::Arch::PositionL1{}); + auto tensorL1B1 = + tla::MakeTensor(l1B1, tla::MakeLayout(curT, curN), Catlass::Arch::PositionL1{}); + + if (reuseSlot) { + WaitFlag(slot); + } + copyGmToL1A0(tensorL1A0, blockQ); + copyGmToL1B0(tensorL1B0, blockH); + copyGmToL1A1(tensorL1A1, blockAqk); + copyGmToL1B1(tensorL1B1, blockVNew); + SetFlag(slot); + } + + __aicore__ inline void ComputePrefetchedOutputTileArch35(Catlass::Arch::Resource &resource, + uint32_t slot, uint64_t b, uint64_t hv, uint64_t start, + uint64_t curT, uint64_t nOffset) + { + using ElementA = T; + using ElementB = T; + using ElementC = OUT_T; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::RowMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; + using TileMmad = Catlass::Gemm::Tile::TileMmadTla; + + constexpr uint16_t kMte1Event = 0; + constexpr uint16_t kMmadEvent = 0; + constexpr uint16_t kFixEvent = 0; + constexpr uint32_t kL1SlotBytes = 96 * 1024; + constexpr uint32_t kL1A0Offset = 0; + constexpr uint32_t kL1B0Offset = 64 * 128 * sizeof(ElementA); + constexpr uint32_t kL1A1Offset = kL1B0Offset + 128 * 128 * sizeof(ElementB); + constexpr uint32_t kL1B1Offset = kL1A1Offset + 64 * 64 * sizeof(ElementA); + const uint32_t slotBase = slot * kL1SlotBytes; + const uint32_t curM = static_cast(curT); + const uint32_t curN = static_cast((V_ - nOffset) > 128 ? 128 : (V_ - nOffset)); + + LocalTensor l1A0 = resource.l1Buf.template GetBufferByByte(slotBase + kL1A0Offset); + LocalTensor l1B0 = resource.l1Buf.template GetBufferByByte(slotBase + kL1B0Offset); + LocalTensor l1A1 = resource.l1Buf.template GetBufferByByte(slotBase + kL1A1Offset); + LocalTensor l1B1 = resource.l1Buf.template GetBufferByByte(slotBase + kL1B1Offset); + LocalTensor l0A = resource.l0ABuf.template GetBufferByByte(0); + LocalTensor l0B = resource.l0BBuf.template GetBufferByByte(0); + LocalTensor l0C = resource.l0CBuf.template GetBufferByByte(0); + + auto tensorL1A0 = + tla::MakeTensor(l1A0, tla::MakeLayout(curM, K_), Catlass::Arch::PositionL1{}); + auto tensorL1B0 = + tla::MakeTensor(l1B0, tla::MakeLayout(K_, curN), Catlass::Arch::PositionL1{}); + auto tensorL1A1 = + tla::MakeTensor(l1A1, tla::MakeLayout(curM, curT), Catlass::Arch::PositionL1{}); + auto tensorL1B1 = + tla::MakeTensor(l1B1, tla::MakeLayout(curT, curN), Catlass::Arch::PositionL1{}); + auto tensorL0A0 = + tla::MakeTensor(l0A, tla::MakeLayout(curM, K_), Catlass::Arch::PositionL0A{}); + auto tensorL0B0 = + tla::MakeTensor(l0B, tla::MakeLayout(K_, curN), Catlass::Arch::PositionL0B{}); + auto tensorL0A1 = + tla::MakeTensor(l0A, tla::MakeLayout(curM, curT), Catlass::Arch::PositionL0A{}); + auto tensorL0B1 = + tla::MakeTensor(l0B, tla::MakeLayout(curT, curN), Catlass::Arch::PositionL0B{}); + auto tensorL0C = tla::MakeTensor(l0C, tla::MakeLayoutL0C(curM, curN), Catlass::Arch::PositionL0C{}); + + uint32_t localRow = 0; + uint32_t localColumn = 0; + auto tileL1A0 = GetTile(tensorL1A0, tla::MakeCoord(localRow, localColumn), tla::MakeShape(curM, K_)); + auto tileL1B0 = GetTile(tensorL1B0, tla::MakeCoord(localRow, localColumn), tla::MakeShape(K_, curN)); + auto tileL1A1 = GetTile(tensorL1A1, tla::MakeCoord(localRow, localColumn), tla::MakeShape(curM, curT)); + auto tileL1B1 = GetTile(tensorL1B1, tla::MakeCoord(localRow, localColumn), tla::MakeShape(curT, curN)); + auto tileL0A0 = GetTile(tensorL0A0, tla::MakeCoord(localRow, localColumn), tla::MakeShape(curM, K_)); + auto tileL0B0 = GetTile(tensorL0B0, tla::MakeCoord(localRow, localColumn), tla::MakeShape(K_, curN)); + auto tileL0A1 = GetTile(tensorL0A1, tla::MakeCoord(localRow, localColumn), tla::MakeShape(curM, curT)); + auto tileL0B1 = GetTile(tensorL0B1, tla::MakeCoord(localRow, localColumn), tla::MakeShape(curT, curN)); + auto tileL0C = GetTile(tensorL0C, tla::MakeCoord(localRow, localColumn), tla::MakeShape(curM, curN)); + + auto layoutO = tla::MakeLayout(BT_, V_); + auto tensorO = tla::MakeTensor(o_[KVOffset(b, hv, start, nOffset, V_)], layoutO, Catlass::Arch::PositionGM{}); + auto tensorLocal = + tla::MakeTensor(u_[KVOffset(b, hv, start, nOffset, V_)], layoutO, Catlass::Arch::PositionGM{}); + auto blockO = GetTile(tensorO, tla::MakeCoord(0, 0), tla::MakeShape(curM, curN)); + auto blockLocal = GetTile(tensorLocal, tla::MakeCoord(0, 0), tla::MakeShape(curM, curN)); + using CopyL0CToDst = typename TileCopy::template CopyL0CToDst; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToDst copyL0CToDst; + TileMmad tileMmad; + + WaitFlag(slot); + if constexpr (IsSameType::value) { + WaitFlag(kFixEvent); + } + copyL1ToL0A(tileL0A0, tileL1A0); + copyL1ToL0B(tileL0B0, tileL1B0); + SetFlag(kMte1Event); + WaitFlag(kMte1Event); + tileMmad(tileL0C, tileL0A0, tileL0B0, curM, curN, static_cast(K_), true, 0b11); + SetFlag(kMmadEvent); + WaitFlag(kMmadEvent); + if constexpr (!IsSameType::value) { + copyL0CToDst(blockO, tileL0C, 0b11); + PipeBarrier(); + } + + copyL1ToL0A(tileL0A1, tileL1A1); + copyL1ToL0B(tileL0B1, tileL1B1); + SetFlag(kMte1Event); + SetFlag(slot); + WaitFlag(kMte1Event); + tileMmad(tileL0C, tileL0A1, tileL0B1, curM, curN, static_cast(curT), + !IsSameType::value, 0b11); + if constexpr (IsSameType::value) { + SetFlag(kFixEvent); + WaitFlag(kFixEvent); + auto fixParams = FixpipeParamsV220(curN, curM, curN, static_cast(HV_ * V_), false); + fixParams.quantPre = QuantMode_t::F322BF16; + Fixpipe(vNew_[OutputOffset(b, hv, start, nOffset)], l0C, fixParams); + SetFlag(kFixEvent); + } else { + SetFlag(kMmadEvent); + WaitFlag(kMmadEvent); + copyL0CToDst(blockLocal, tileL0C, 0b11); + PipeBarrier(); + } + } + + __aicore__ inline void ProcessOutAicPipelinedArch35() + { + SetLoadDataPaddingValue(static_cast(0)); + SetMMLayoutTransform(true); + uint64_t taskNum = static_cast((isVarLen_ ? NT_ : B_ * NT_) * HV_); + uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + uint64_t currentTask = static_cast(GetBlockIdx()); + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + while (currentTask < taskNum && !ResolveFlatChunk(currentTask, seq, b, h, hv, chunkIdx, start, end)) { + currentTask += coreNum; + } + if (currentTask >= taskNum) { + SetMMLayoutTransform(false); + return; + } + + Catlass::Arch::Resource resource; + uint64_t nOffset = 0; + uint32_t slot = 0; + if constexpr (IsSameType::value) { + SetFlag(0); + } + PrefetchOutputTileArch35(resource, slot, b, hv, chunkIdx, start, end - start, nOffset, false); + uint64_t outputTileIdx = 0; + + while (true) { + uint64_t nextTask = currentTask; + uint64_t nextSeq = seq; + uint64_t nextB = b; + uint64_t nextH = h; + uint64_t nextHv = hv; + uint64_t nextChunkIdx = chunkIdx; + uint64_t nextStart = start; + uint64_t nextEnd = end; + uint64_t nextNOffset = nOffset + 128; + bool hasNext = nextNOffset < V_; + if (!hasNext) { + nextTask += coreNum; + nextNOffset = 0; + while (nextTask < taskNum && + !ResolveFlatChunk(nextTask, nextSeq, nextB, nextH, nextHv, nextChunkIdx, nextStart, nextEnd)) { + nextTask += coreNum; + } + hasNext = nextTask < taskNum; + } + + const uint32_t nextSlot = slot ^ 1U; + if (hasNext) { + PrefetchOutputTileArch35(resource, nextSlot, nextB, nextHv, nextChunkIdx, nextStart, + nextEnd - nextStart, nextNOffset, outputTileIdx + 1 >= 2); + } + ComputePrefetchedOutputTileArch35(resource, slot, b, hv, start, end - start, nOffset); + if constexpr (!IsSameType::value) { + if (nOffset + 128 >= V_) { + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(syncDoneFlag_); + } + } + if (!hasNext) { + break; + } + ++outputTileIdx; + + currentTask = nextTask; + seq = nextSeq; + b = nextB; + h = nextH; + hv = nextHv; + chunkIdx = nextChunkIdx; + start = nextStart; + end = nextEnd; + nOffset = nextNOffset; + slot = nextSlot; + } + if constexpr (IsSameType::value) { + WaitFlag(0); + } + SetMMLayoutTransform(false); + } +#endif + + __aicore__ inline void ComputeOutputCube(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t curT) + { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if (curT < KDA_CUBE_MIN_REDUCTION) { + return; + } + SetLoadDataPaddingValue(static_cast(0)); + if (BT_ == 64 && curT == BT_) { + ComputeOutputCubeStagedArch35(b, hv, chunkIdx, start, curT); + return; + } +#endif + using ElementA = T; + using ElementB = T; + using ElementC = OUT_T; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::RowMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using BlockMmad = Catlass::Gemm::Block::BlockMmadTla; + + Catlass::Arch::Resource resource; + BlockMmad blockMmad(resource); + + auto layoutQ = tla::MakeLayout(BT_, K_); + auto layoutH = tla::MakeLayout(K_, V_); + auto layoutO = tla::MakeLayout(BT_, V_); + for (uint64_t nOffset = 0; nOffset < V_; nOffset += 128) { + uint32_t curN = static_cast((V_ - nOffset) > 128 ? 128 : (V_ - nOffset)); + auto tensorH = tla::MakeTensor(propagatedH_[HOffset(b, hv, chunkIdx, 0, nOffset)], layoutH, + Catlass::Arch::PositionGM{}); + for (uint64_t mOffset = 0; mOffset < curT; mOffset += 64) { + uint32_t curM = static_cast((curT - mOffset) > 64 ? 64 : (curT - mOffset)); + Catlass::GemmCoord shapeQH{curM, curN, static_cast(K_)}; + auto tensorQ = tla::MakeTensor(preparedQG_[KVOffset(b, hv, start + mOffset, 0, K_)], layoutQ, + Catlass::Arch::PositionGM{}); + auto tensorO = tla::MakeTensor(o_[KVOffset(b, hv, start + mOffset, nOffset, V_)], layoutO, + Catlass::Arch::PositionGM{}); + auto blockQ = GetTile(tensorQ, tla::MakeCoord(0, 0), tla::MakeShape(shapeQH.m(), shapeQH.k())); + auto blockH = GetTile(tensorH, tla::MakeCoord(0, 0), tla::MakeShape(shapeQH.k(), shapeQH.n())); + auto blockO = GetTile(tensorO, tla::MakeCoord(0, 0), tla::MakeShape(shapeQH.m(), shapeQH.n())); + blockMmad(blockQ, blockH, blockO, shapeQH); + PipeBarrier(); + } + } + + if (curT < KDA_CUBE_MIN_REDUCTION) { + return; + } + + auto layoutAqk = tla::MakeLayout(BT_, BT_); + auto layoutV = tla::MakeLayout(BT_, V_); + for (uint64_t nOffset = 0; nOffset < V_; nOffset += 128) { + uint32_t curN = static_cast((V_ - nOffset) > 128 ? 128 : (V_ - nOffset)); + auto tensorVNew = tla::MakeTensor(propagatedVNew_[KVOffset(b, hv, start, nOffset, V_)], layoutV, + Catlass::Arch::PositionGM{}); + for (uint64_t mOffset = 0; mOffset < curT; mOffset += 64) { + uint32_t curM = static_cast((curT - mOffset) > 64 ? 64 : (curT - mOffset)); + Catlass::GemmCoord shapeAV{curM, curN, static_cast(curT)}; + auto tensorAqk = tla::MakeTensor(preparedAqk_[AOffset(b, hv, start + mOffset, 0)], layoutAqk, + Catlass::Arch::PositionGM{}); + auto tensorLocal = tla::MakeTensor(u_[KVOffset(b, hv, start + mOffset, nOffset, V_)], layoutO, + Catlass::Arch::PositionGM{}); + auto blockAqk = GetTile(tensorAqk, tla::MakeCoord(0, 0), tla::MakeShape(shapeAV.m(), shapeAV.k())); + auto blockVNew = GetTile(tensorVNew, tla::MakeCoord(0, 0), tla::MakeShape(shapeAV.k(), shapeAV.n())); + auto blockLocal = GetTile(tensorLocal, tla::MakeCoord(0, 0), tla::MakeShape(shapeAV.m(), shapeAV.n())); + blockMmad(blockAqk, blockVNew, blockLocal, shapeAV); + PipeBarrier(); + } + } + } + + __aicore__ inline void FinalizeOutputRows(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t curT, + uint64_t subBlockIdx, uint64_t subBlockNum) + { + if (subBlockNum == 0 || subBlockIdx >= subBlockNum || V_ == 0) { + return; + } + const uint64_t rowBegin = (curT * subBlockIdx) / subBlockNum; + const uint64_t rowEnd = (curT * (subBlockIdx + 1)) / subBlockNum; + const uint64_t gateWritebackRows = ScoreVectorMaxRows(5 * sizeof(float) + 2 * sizeof(T) + sizeof(GK_T)); + const uint64_t gateWritebackBytes = gateWritebackRows * K_ * (3 * sizeof(T) + sizeof(GK_T)); + uint64_t maxRows = KDA_VEC_ARENA_ELEMENTS / (3 * V_); + const uint64_t typedMaxRows = gateWritebackBytes / (V_ * sizeof(T)); + if (maxRows > typedMaxRows) { + maxRows = typedMaxRows; + } + if (maxRows == 0) { + return; + } + + for (uint64_t tileRow = rowBegin; tileRow < rowEnd; tileRow += maxRows) { + uint64_t tileRows = rowEnd - tileRow; + if (tileRows > maxRows) { + tileRows = maxRows; + } + const uint64_t elems = tileRows * V_; + const uint64_t ti = start + tileRow; + LocalTensor arena = vecBuf_.Get(); + LocalTensor stateLocal = arena; + LocalTensor localLocal = arena[elems]; + LocalTensor outLocal = arena[2 * elems]; + LocalTensor outTyped = gateWritebackBuf_.Get(); + + if (curT < KDA_CUBE_MIN_REDUCTION) { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (IsSameType::value) { + if (BT_ == 64 && K_ == 128 && V_ == 128) { + ComputeTailOutputRegbaseRows(outLocal, b, hv, chunkIdx, start, curT, tileRow, tileRows); + } else { + ComputeTailStateRows(stateLocal, b, hv, chunkIdx, start, tileRow, tileRows); + ComputeTailLocalRows(localLocal, b, hv, start, curT, tileRow, tileRows); + Muls(stateLocal, stateLocal, scale_, static_cast(elems)); + PipeBarrier(); + Add(outLocal, stateLocal, localLocal, static_cast(elems)); + PipeBarrier(); + } + } else { + ComputeTailStateRows(stateLocal, b, hv, chunkIdx, start, tileRow, tileRows); + ComputeTailLocalRows(localLocal, b, hv, start, curT, tileRow, tileRows); + Muls(stateLocal, stateLocal, scale_, static_cast(elems)); + PipeBarrier(); + Add(outLocal, stateLocal, localLocal, static_cast(elems)); + PipeBarrier(); + } +#else + ComputeTailStateRows(stateLocal, b, hv, chunkIdx, start, tileRow, tileRows); + ComputeTailLocalRows(localLocal, b, hv, start, curT, tileRow, tileRows); + Muls(stateLocal, stateLocal, scale_, static_cast(elems)); + PipeBarrier(); + Add(outLocal, stateLocal, localLocal, static_cast(elems)); + PipeBarrier(); +#endif + } else { + CopyVectorIn(stateLocal, o_, KVOffset(b, hv, ti, 0, V_), elems); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + CopyVectorIn(localLocal, u_, KVOffset(b, hv, ti, 0, V_), elems); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Muls(stateLocal, stateLocal, scale_, static_cast(elems)); + PipeBarrier(); + Add(outLocal, stateLocal, localLocal, static_cast(elems)); + PipeBarrier(); + } + ClampFp32ToOutputType(outLocal, static_cast(elems)); + Cast(outTyped, outLocal, RoundMode::CAST_RINT, static_cast(elems)); + PipeBarrier(); + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyRowsOut(vNew_, OutputOffset(b, hv, ti, 0), outTyped, tileRows, V_, HV_ * V_); + SetFlag(mte3ToMte2Event_); + WaitFlag(mte3ToMte2Event_); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + } + __aicore__ inline bool ResolveFlatChunk(uint64_t task, uint64_t &seq, uint64_t &b, uint64_t &h, uint64_t &hv, + uint64_t &chunkIdx, uint64_t &start, uint64_t &end) + { + hv = task % HV_; + uint64_t flatChunk = task / HV_; + if (!isVarLen_) { + seq = flatChunk / NT_; + b = seq; + chunkIdx = flatChunk % NT_; + start = chunkIdx * BT_; + end = start + BT_; + if (end > T_) { + end = T_; + } + } else { + if (!KdaVarlen::ResolveChunkRange(cuSeqlensAddr_, chunkIndicesAddr_, N_, T_, BT_, flatChunk, seq, start, + end)) { + return false; + } + b = 0; + chunkIdx = flatChunk; + } + h = hv / (HV_ / H_); + return start < end; + } + + __aicore__ inline void ProcessChunkOutAiv(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t end, + uint64_t subBlockIdx, uint64_t subBlockNum) + { + uint64_t curT = end - start; + if (curT == 0) { + return; + } + if constexpr (IsSameType::value) { + return; + } + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE2>(syncDoneFlag_); + FinalizeOutputRows(b, hv, chunkIdx, start, curT, subBlockIdx, subBlockNum); + } + + __aicore__ inline void ProcessChunkOutAic(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t end) + { + uint64_t curT = end - start; + if (curT == 0) { + return; + } + ComputeOutputCube(b, hv, chunkIdx, start, curT); + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(syncDoneFlag_); + } + + __aicore__ inline void ProcessOutAiv() + { + if constexpr (IsSameType::value) { + return; + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (IsSameType::value) { + if (!isVarLen_ && T_ % BT_ == 0 && BT_ == 64 && K_ == 128 && V_ == 128) { + return; + } + } +#endif + uint64_t subBlockNum = static_cast(GetSubBlockNum()); + if (subBlockNum == 0) { + return; + } + uint64_t subBlockIdx = static_cast(GetSubBlockIdx()); + uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + uint64_t coreIdx = static_cast(GetBlockIdx()) / subBlockNum; + uint64_t taskNum = static_cast((isVarLen_ ? NT_ : B_ * NT_) * HV_); + for (uint64_t task = coreIdx; task < taskNum; task += coreNum) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveFlatChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + (void)seq; + (void)h; + (void)chunkIdx; + ProcessChunkOutAiv(b, hv, chunkIdx, start, end, subBlockIdx, subBlockNum); + } + } + } + + __aicore__ inline void ProcessOutAic() + { + if constexpr (IsSameType::value) { + return; + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if (!isVarLen_ && T_ % BT_ == 0 && BT_ == 64 && K_ == 128 && V_ == 128) { + ProcessOutAicPipelinedArch35(); + return; + } +#endif + uint64_t taskNum = static_cast((isVarLen_ ? NT_ : B_ * NT_) * HV_); + uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + for (uint64_t task = GetBlockIdx(); task < taskNum; task += coreNum) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveFlatChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + (void)seq; + (void)h; + ProcessChunkOutAic(b, hv, chunkIdx, start, end); + } + } + } + +private: + GlobalTensor q_; + GlobalTensor k_; + GlobalTensor v_; + GlobalTensor gk_; + GlobalTensor initialState_; + GlobalTensor o_; + GlobalTensor finalState_; + GlobalTensor aqk_; + GlobalTensor akk_; + GlobalTensor w_; + GlobalTensor u_; + GlobalTensor qg_; + GlobalTensor kg_; + GlobalTensor vNew_; + GlobalTensor h_; + GlobalTensor preparedQG_; + GlobalTensor preparedAqk_; + GlobalTensor propagatedVNew_; + GlobalTensor propagatedH_; + GlobalTensor solveWorkspace_; + GlobalTensor scoreWorkspace_; + TPipe *pipe_ = nullptr; + TBuf exp2Buf_; + TBuf vecBuf_; + TBuf gateWritebackBuf_; + TEventID mte2ToVEvent_ = 0; + TEventID vToMte2Event_ = 0; + TEventID vToMte3Event_ = 0; + TEventID mte3ToVEvent_ = 0; + TEventID mte2ToMte3Event_ = 0; + TEventID mte3ToMte2Event_ = 0; + TEventID sToVEvent_ = 0; + TEventID sToMte2Event_ = 0; + bool vectorEventsAllocated_ = false; + Catlass::Arch::CrossCoreFlagWithReverse scoreReadyFlag_{KDA_SCORE_READY_FLAG0, + KDA_SCORE_READY_FLAG1}; + Catlass::Arch::CrossCoreFlagWithReverse scoreDoneFlag_{KDA_SCORE_DONE_FLAG0, + KDA_SCORE_DONE_FLAG1}; + // Score production is fully drained before solve starts, so the solve handshake can safely reuse + // the existing score flags without consuming additional hardware flag IDs. + Catlass::Arch::CrossCoreFlagWithReverse syncReadyFlag_{KDA_SCORE_READY_FLAG0, + KDA_SCORE_READY_FLAG1}; + Catlass::Arch::CrossCoreFlagWithReverse syncDoneFlag_{KDA_SCORE_DONE_FLAG0, + KDA_SCORE_DONE_FLAG1}; + uint64_t B_ = 0; + uint64_t N_ = 0; + uint64_t H_ = 0; + uint64_t HV_ = 0; + uint64_t T_ = 0; + uint64_t K_ = 0; + uint64_t V_ = 0; + uint64_t BT_ = 0; + uint64_t NT_ = 0; + float scale_ = 1.0f; + bool hasInitial_ = false; + bool isVarLen_ = false; + bool isAivOnly_ = false; + uint64_t usedCoreNum_ = 1; + uint64_t solveCoreIdx_ = 0; + __gm__ int64_t *chunkIndicesAddr_ = nullptr; + __gm__ int64_t *cuSeqlensAddr_ = nullptr; +}; +} // namespace + +template +__aicore__ inline void RunChunkKdaOutput(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR beta, + GM_ADDR initialState, GM_ADDR cuSeqlens, GM_ADDR chunkIndices, + GM_ADDR qgScaled, GM_ADDR aqk, GM_ADDR propagatedVNew, GM_ADDR propagatedH, + GM_ADDR o, GM_ADDR userWorkspace, const TilingData &tiling, TPipe &pipe) +{ + GM_ADDR outputScratch = userWorkspace + tiling.outputScratchOffset; + uint64_t outputElements = static_cast(tiling.batch) * static_cast(tiling.vHeadNum) * + static_cast(tiling.seqlen) * static_cast(tiling.vHeadDim); + GM_ADDR stateScratch = outputScratch; + GM_ADDR localScratch = outputScratch + outputElements * sizeof(float); + if ASCEND_IS_AIC { + ChunkKdaFwdFinalizeKernel op; + op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, qgScaled, aqk, propagatedVNew, propagatedH, + stateScratch, userWorkspace, aqk, userWorkspace, userWorkspace, localScratch, userWorkspace, + userWorkspace, o, propagatedH, outputScratch, tiling, &pipe, false); + op.ProcessAic(); + } + if ASCEND_IS_AIV { + ChunkKdaFwdFinalizeKernel op; + op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, qgScaled, aqk, propagatedVNew, propagatedH, + stateScratch, userWorkspace, aqk, userWorkspace, userWorkspace, localScratch, userWorkspace, + userWorkspace, o, propagatedH, outputScratch, tiling, &pipe); + op.ProcessAiv(); + } +} + +} // namespace KdaFinalize diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_fwd_h.h b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_fwd_h.h new file mode 100644 index 0000000000..fcb979cc91 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_fwd_h.h @@ -0,0 +1,1013 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CHUNK_KDA_FWD_ARCH35_FWD_H_H +#define CHUNK_KDA_FWD_ARCH35_FWD_H_H + +#include "kernel_operator.h" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" +#include "../kernel_utils/tile/copy_l0c_to_ub.hpp" + +namespace KdaForward::arch35 { + +using namespace AscendC; + +// Prepare is fully drained before FwdH starts. Keep FwdH mode-2 +// traffic outside Matmul's possible 0..7 range and SyncAll's 11..14 range. +constexpr uint64_t KDA_FWD_H_DIRECT_FREE_FLAG = 8; +constexpr uint64_t KDA_FWD_H_DIRECT_READY_FLAG = 8; +constexpr uint64_t KDA_FWD_H_L1_FREE_FLAG = 9; +constexpr uint64_t KDA_FWD_H_L1_READY_FLAG = 9; +constexpr uint64_t KDA_FWD_H_STATE_FREE_FLAG = KDA_FWD_H_L1_FREE_FLAG; +constexpr uint64_t KDA_FWD_H_VNEW_FREE_FLAG = KDA_FWD_H_L1_FREE_FLAG; +constexpr uint64_t KDA_FWD_H_STATE_READY_FLAG = KDA_FWD_H_L1_READY_FLAG; +constexpr uint64_t KDA_FWD_H_VNEW_READY_FLAG = KDA_FWD_H_L1_READY_FLAG; + +constexpr TEventID KDA_FWD_H_MTE_W_EVENT = 0; +constexpr TEventID KDA_FWD_H_MTE_Q_EVENT = 1; +constexpr TEventID KDA_FWD_H_MTE_B_EVENT = 2; +constexpr TEventID KDA_FWD_H_MTE_A_EVENT = 3; +constexpr TEventID KDA_FWD_H_M_EVENT = 4; +constexpr TEventID KDA_FWD_H_FIX_EVENT = 5; +constexpr TEventID KDA_FWD_H_IO_REUSE_EVENT = 6; + +constexpr uint32_t KDA_FWD_H_CHUNK = 64; +constexpr uint32_t KDA_FWD_H_DIM = 128; +constexpr uint32_t KDA_FWD_H_SUB_CHUNK = KDA_FWD_H_CHUNK / 2; +constexpr uint32_t KDA_FWD_H_SUB_DIM = KDA_FWD_H_DIM / 2; +constexpr uint32_t KDA_FWD_H_STATE_SUB_ELEMS = KDA_FWD_H_SUB_DIM * KDA_FWD_H_DIM; +constexpr uint32_t KDA_FWD_H_TOKEN_SUB_ELEMS = KDA_FWD_H_SUB_CHUNK * KDA_FWD_H_DIM; + +constexpr uint32_t KDA_FWD_H_L1_W_OFFSET = 0; +constexpr uint32_t KDA_FWD_H_L1_Q_OFFSET = 16 * 1024; +constexpr uint32_t KDA_FWD_H_L1_H_OFFSET = 32 * 1024; +constexpr uint32_t KDA_FWD_H_L1_KG_OFFSET = 64 * 1024; +constexpr uint32_t KDA_FWD_H_L1_AQK_OFFSET = 80 * 1024; +constexpr uint32_t KDA_FWD_H_L1_V_OFFSET = 96 * 1024; +constexpr uint32_t KDA_FWD_H_L1_W1_OFFSET = 112 * 1024; +constexpr uint32_t KDA_FWD_H_L1_Q1_OFFSET = 128 * 1024; +constexpr uint32_t KDA_FWD_H_L1_KG1_OFFSET = 144 * 1024; +constexpr uint32_t KDA_FWD_H_L1_AQK1_OFFSET = 160 * 1024; +constexpr uint32_t KDA_FWD_H_L1_W2_OFFSET = 176 * 1024; +constexpr uint32_t KDA_FWD_H_L1_Q2_OFFSET = 192 * 1024; +constexpr uint32_t KDA_FWD_H_L1_KG2_OFFSET = 208 * 1024; +constexpr uint32_t KDA_FWD_H_L1_AQK2_OFFSET = 224 * 1024; +constexpr uint32_t KDA_FWD_H_L1_W3_OFFSET = 240 * 1024; +constexpr uint32_t KDA_FWD_H_L1_Q3_OFFSET = 256 * 1024; +constexpr uint32_t KDA_FWD_H_L1_KG3_OFFSET = 272 * 1024; +constexpr uint32_t KDA_FWD_H_L1_AQK3_OFFSET = 288 * 1024; +constexpr uint32_t KDA_FWD_H_L1_AKK_OFFSET = 304 * 1024; +constexpr uint32_t KDA_FWD_H_L1_U_OFFSET = 336 * 1024; +constexpr uint32_t KDA_FWD_H_L1_STAGING_DEPTH = 4; +constexpr uint32_t KDA_FWD_H_L1_AKK_SLOT_BYTES = 8 * 1024; +constexpr uint32_t KDA_FWD_H_L1_U_SLOT_BYTES = 16 * 1024; + +constexpr uint32_t KDA_FWD_H_L0A_STATE_OFFSET = 0; +constexpr uint32_t KDA_FWD_H_L0A_VNEW_OFFSET = 16 * 1024; +constexpr uint32_t KDA_FWD_H_L0A_POST_OFFSET = 32 * 1024; +constexpr uint32_t KDA_FWD_H_L0B_STATE_OFFSET = 0; +constexpr uint32_t KDA_FWD_H_L0B_VNEW_OFFSET = 32 * 1024; +constexpr uint32_t KDA_FWD_H_L0B_POST_OFFSET = 32 * 1024; + +constexpr uint32_t KDA_FWD_H_UB_STATE_OFFSET = 0; +constexpr uint32_t KDA_FWD_H_UB_STATE_TYPED_OFFSET = 32 * 1024; +constexpr uint32_t KDA_FWD_H_UB_DIRECT_OFFSET = 48 * 1024; +constexpr uint32_t KDA_FWD_H_UB_OUT1_OFFSET = 80 * 1024; +constexpr uint32_t KDA_FWD_H_UB_VNEW_OFFSET = 96 * 1024; +constexpr uint32_t KDA_FWD_H_UB_IO_OFFSET = 112 * 1024; +constexpr uint32_t KDA_FWD_H_UB_GATE_OFFSET = 128 * 1024; + +template +class ChunkKdaFwdFwdH { +public: + using ArchTag = Catlass::Arch::Ascend950; + using LayoutRM = Catlass::layout::RowMajor; + using LayoutCM = Catlass::layout::ColumnMajor; + using TileCopyRM = Catlass::Gemm::Tile::PackedTileCopyTla< + ArchTag, T, LayoutRM, T, LayoutRM, float, LayoutRM>; + using DirectTileCopyRM = Common::Tile::PackedTileCopyTlaToUB< + ArchTag, T, LayoutRM, T, LayoutRM, float, LayoutRM, void, + Catlass::Gemm::Tile::CopyL0CToUBMode::SPLIT_M>; + using TileCopyCM = Catlass::Gemm::Tile::PackedTileCopyTla< + ArchTag, T, LayoutCM, T, LayoutRM, float, LayoutRM>; + using DirectTileCopyCM = Common::Tile::PackedTileCopyTlaToUB< + ArchTag, T, LayoutCM, T, LayoutRM, float, LayoutRM, void, + Catlass::Gemm::Tile::CopyL0CToUBMode::SPLIT_M>; + + __aicore__ inline void Init( + GM_ADDR gk, GM_ADDR initialState, GM_ADDR attnOut, GM_ADDR finalState, + GM_ADDR aqk, GM_ADDR akk, GM_ADDR w, GM_ADDR u, GM_ADDR qg, GM_ADDR kg, + GM_ADDR vNew, GM_ADDR h, const TilingData &tiling) + { + gk_.SetGlobalBuffer(reinterpret_cast<__gm__ GK_T *>(gk)); + if (initialState != nullptr) { + initialState_.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(initialState)); + } + attnOut_.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(attnOut)); + finalState_.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(finalState)); + aqk_.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(aqk)); + akk_.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(akk)); + w_.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(w)); + u_.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(u)); + qg_.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(qg)); + kg_.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(kg)); + vNew_.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(vNew)); + h_.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(h)); + batch_ = tiling.batch; + heads_ = tiling.vHeadNum; + seqlen_ = tiling.seqlen; + totalChunks_ = tiling.totalChunks; + hasInitialState_ = tiling.hasInitialState; + storeFinalState_ = tiling.storeFinalState; + storeVNew_ = tiling.storeVNew; + storeH_ = tiling.storeH; + fusePostWuIntoFwdH_ = tiling.fusePostWuIntoFwdH; + scale_ = tiling.scale; + coreNum_ = tiling.prepareUsedCoreNum; + statePublishCount_[0] = 0; + statePublishCount_[1] = 0; + vnewPublishCount_[0] = 0; + vnewPublishCount_[1] = 0; + wPublishCount_[0] = 0; + wPublishCount_[1] = 0; + } + + __aicore__ inline void Process() + { + if ASCEND_IS_AIC { + ProcessAic(); + } + if ASCEND_IS_AIV { + ProcessAiv(); + } + } + +private: + __aicore__ inline void WaitL1SlotFreeMte3( + uint64_t freeFlag, uint32_t publishCount) + { + if (publishCount != 0) { + CrossCoreWaitFlag(freeFlag); + } + } + + template + __aicore__ inline void SetL1SlotFlagAicToAiv(uint64_t flag) + { + CrossCoreSetFlag<0x2, PIPE>(flag); + } + + template + __aicore__ inline void SetL1SlotFlagAivToAic(uint64_t flag) + { + CrossCoreSetFlag<0x2, PIPE>(flag); + } + + __aicore__ inline void WaitL1SlotReadyMte1(uint64_t readyFlag) + { + CrossCoreWaitFlag(readyFlag); + } + + __aicore__ inline uint64_t MatrixOffset(uint64_t b, uint64_t hv, uint64_t t) const + { + return ((b * heads_ + hv) * seqlen_ + t) * KDA_FWD_H_DIM; + } + + __aicore__ inline uint64_t ChunkMatrixOffset( + uint64_t b, uint64_t hv, uint64_t chunk, uint64_t row = 0) const + { + return ((b * heads_ + hv) * seqlen_ + + chunk * KDA_FWD_H_CHUNK + row) * + KDA_FWD_H_DIM; + } + + __aicore__ inline uint64_t ScoreOffset(uint64_t b, uint64_t hv, uint64_t t) const + { + return ((b * heads_ + hv) * seqlen_ + t) * KDA_FWD_H_CHUNK; + } + + __aicore__ inline uint64_t ChunkScoreOffset( + uint64_t b, uint64_t hv, uint64_t chunk) const + { + return ((b * heads_ + hv) * seqlen_ + + chunk * KDA_FWD_H_CHUNK) * + KDA_FWD_H_CHUNK; + } + + __aicore__ inline uint64_t StateOffset(uint64_t b, uint64_t hv) const + { + return (b * heads_ + hv) * KDA_FWD_H_DIM * KDA_FWD_H_DIM; + } + + __aicore__ inline uint64_t HOffset(uint64_t b, uint64_t hv, uint64_t chunk) const + { + if (!storeH_) { + return StateOffset(b, hv); + } + return ((b * heads_ + hv) * totalChunks_ + chunk) * + KDA_FWD_H_DIM * KDA_FWD_H_DIM; + } + + __aicore__ inline uint64_t VNewOffset( + uint64_t b, uint64_t hv, uint64_t chunk, uint64_t row = 0) const + { + if (!storeVNew_) { + return ((b * heads_ + hv) * KDA_FWD_H_CHUNK + row) * KDA_FWD_H_DIM; + } + return ChunkMatrixOffset(b, hv, chunk, row); + } + + __aicore__ inline uint64_t OutputOffset(uint64_t b, uint64_t hv, uint64_t t) const + { + return ((b * seqlen_ + t) * heads_ + hv) * KDA_FWD_H_DIM; + } + + __aicore__ inline uint32_t L1WOffset(uint32_t slot) const + { + if (slot == 0) { + return KDA_FWD_H_L1_W_OFFSET; + } + if (slot == 1) { + return KDA_FWD_H_L1_W1_OFFSET; + } + return slot == 2 ? KDA_FWD_H_L1_W2_OFFSET : KDA_FWD_H_L1_W3_OFFSET; + } + + __aicore__ inline uint32_t L1QOffset(uint32_t slot) const + { + if (slot == 0) { + return KDA_FWD_H_L1_Q_OFFSET; + } + if (slot == 1) { + return KDA_FWD_H_L1_Q1_OFFSET; + } + return slot == 2 ? KDA_FWD_H_L1_Q2_OFFSET : KDA_FWD_H_L1_Q3_OFFSET; + } + + __aicore__ inline uint32_t L1KgOffset(uint32_t slot) const + { + if (slot == 0) { + return KDA_FWD_H_L1_KG_OFFSET; + } + if (slot == 1) { + return KDA_FWD_H_L1_KG1_OFFSET; + } + return slot == 2 ? KDA_FWD_H_L1_KG2_OFFSET : KDA_FWD_H_L1_KG3_OFFSET; + } + + __aicore__ inline uint32_t L1AqkOffset(uint32_t slot) const + { + if (slot == 0) { + return KDA_FWD_H_L1_AQK_OFFSET; + } + if (slot == 1) { + return KDA_FWD_H_L1_AQK1_OFFSET; + } + return slot == 2 ? KDA_FWD_H_L1_AQK2_OFFSET : KDA_FWD_H_L1_AQK3_OFFSET; + } + + __aicore__ inline uint32_t L1AkkOffset(uint32_t slot) const + { + return KDA_FWD_H_L1_AKK_OFFSET + slot * KDA_FWD_H_L1_AKK_SLOT_BYTES; + } + + __aicore__ inline uint32_t L1UOffset(uint32_t slot) const + { + return KDA_FWD_H_L1_U_OFFSET + slot * KDA_FWD_H_L1_U_SLOT_BYTES; + } + + template + __aicore__ inline void PublishDirectTile( + TensorL0C tensorL0C, uint32_t m, uint32_t n, + TEventID mToFixEvent, TEventID fixToMEvent) + { + auto layoutUb = tla::MakeLayout(m, n); + auto tensorUb = tla::MakeTensor( + resource_.ubBuf.template GetBufferByByte(KDA_FWD_H_UB_DIRECT_OFFSET), + layoutUb, Catlass::Arch::PositionUB{}); + using CopyL0CToDst = + typename DirectTileCopy::template CopyL0CToDst; + CopyL0CToDst copyL0CToDst; + + CrossCoreWaitFlag(KDA_FWD_H_DIRECT_FREE_FLAG); + SetFlag(mToFixEvent); + WaitFlag(mToFixEvent); + copyL0CToDst(tensorUb, tensorL0C); + CrossCoreSetFlag<0x2, PIPE_FIX>(KDA_FWD_H_DIRECT_READY_FLAG); + SetFlag(fixToMEvent); + } + + template + __aicore__ inline void PublishDirect( + LocalTensor l0C, uint32_t m, uint32_t n, + TEventID mToFixEvent, TEventID fixToMEvent) + { + auto layoutL0C = tla::MakeLayoutL0C(m, n); + auto tensorL0C = tla::MakeTensor(l0C, layoutL0C, Catlass::Arch::PositionL0C{}); + PublishDirectTile( + tensorL0C, m, n, mToFixEvent, fixToMEvent); + } + + __aicore__ inline void PrefetchIndependentProductsAic( + uint64_t b, uint64_t hv, uint64_t chunk) + { + const uint32_t slot = static_cast(chunk & 3); + WaitFlag(aicL1ReuseEvents_[slot]); + using LayoutTagL1ARm = typename TileCopyRM::LayoutTagL1A; + using LayoutTagL1BRm = typename TileCopyRM::LayoutTagL1B; + using LayoutTagL1ACm = typename TileCopyCM::LayoutTagL1A; + + auto layoutToken = tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM); + auto layoutKg = tla::MakeLayout(KDA_FWD_H_DIM, KDA_FWD_H_CHUNK); + auto layoutAqk = tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK); + auto layoutAkk = tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK); + auto tensorW = tla::MakeTensor( + w_[ChunkMatrixOffset(b, hv, chunk)], layoutToken, Catlass::Arch::PositionGM{}); + auto tensorQ = tla::MakeTensor( + qg_[ChunkMatrixOffset(b, hv, chunk)], layoutToken, + Catlass::Arch::PositionGM{}); + auto tensorKg = tla::MakeTensor( + kg_[ChunkMatrixOffset(b, hv, chunk)], layoutKg, Catlass::Arch::PositionGM{}); + auto tensorAqk = tla::MakeTensor( + aqk_[ChunkScoreOffset(b, hv, chunk)], layoutAqk, Catlass::Arch::PositionGM{}); + auto tensorAkk = tla::MakeTensor( + akk_[ChunkScoreOffset(b, hv, chunk)], layoutAkk, Catlass::Arch::PositionGM{}); + auto tensorU = tla::MakeTensor( + u_[ChunkMatrixOffset(b, hv, chunk)], layoutToken, Catlass::Arch::PositionGM{}); + auto blockW = GetTile(tensorW, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + auto blockQ = GetTile(tensorQ, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + auto blockKg = GetTile(tensorKg, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_DIM, KDA_FWD_H_CHUNK)); + auto blockAqk = GetTile(tensorAqk, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK)); + auto blockAkk = GetTile(tensorAkk, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK)); + auto blockU = GetTile(tensorU, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + using CopyGmToL1ARmW = typename TileCopyRM::template CopyGmToL1A; + using CopyGmToL1ARmQ = typename TileCopyRM::template CopyGmToL1A; + using CopyGmToL1ACm = typename TileCopyCM::template CopyGmToL1A; + using CopyGmToL1ARmAqk = typename TileCopyRM::template CopyGmToL1A; + using CopyGmToL1ARmAkk = typename TileCopyRM::template CopyGmToL1A; + using CopyGmToL1BRmU = typename TileCopyRM::template CopyGmToL1B; + + LocalTensor l1W = resource_.l1Buf.template GetBufferByByte( + L1WOffset(slot)); + LocalTensor l1Q = resource_.l1Buf.template GetBufferByByte( + L1QOffset(slot)); + LocalTensor l1Kg = resource_.l1Buf.template GetBufferByByte( + L1KgOffset(slot)); + LocalTensor l1Aqk = resource_.l1Buf.template GetBufferByByte( + L1AqkOffset(slot)); + LocalTensor l1Akk = resource_.l1Buf.template GetBufferByByte( + L1AkkOffset(slot)); + LocalTensor l1U = resource_.l1Buf.template GetBufferByByte( + L1UOffset(slot)); + auto tensorL1W = tla::MakeTensor( + l1W, tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM), + Catlass::Arch::PositionL1{}); + auto tensorL1Q = tla::MakeTensor( + l1Q, tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM), + Catlass::Arch::PositionL1{}); + auto tensorL1Kg = tla::MakeTensor( + l1Kg, tla::MakeLayout(KDA_FWD_H_DIM, KDA_FWD_H_CHUNK), + Catlass::Arch::PositionL1{}); + auto tensorL1Aqk = tla::MakeTensor( + l1Aqk, tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK), + Catlass::Arch::PositionL1{}); + auto tensorL1Akk = tla::MakeTensor( + l1Akk, tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK), + Catlass::Arch::PositionL1{}); + auto tensorL1U = tla::MakeTensor( + l1U, tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM), + Catlass::Arch::PositionL1{}); + + CopyGmToL1ARmW{}(tensorL1W, blockW); + CopyGmToL1ARmQ{}(tensorL1Q, blockQ); + CopyGmToL1ACm{}(tensorL1Kg, blockKg); + CopyGmToL1ARmAqk{}(tensorL1Aqk, blockAqk); + if (fusePostWuIntoFwdH_) { + CopyGmToL1ARmAkk{}(tensorL1Akk, blockAkk); + CopyGmToL1BRmU{}(tensorL1U, blockU); + } + SetFlag(aicMte2ToMte1Event_); + } + + __aicore__ inline void ComputePostWuAic(uint64_t chunk) + { + const uint32_t slot = static_cast(chunk & 3); + using LayoutTagL1A = typename TileCopyRM::LayoutTagL1A; + using LayoutTagL1B = typename TileCopyRM::LayoutTagL1B; + using LayoutTagL0A = typename TileCopyRM::LayoutTagL0A; + using LayoutTagL0B = typename TileCopyRM::LayoutTagL0B; + using CopyL1ToL0A = typename TileCopyRM::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopyRM::CopyL1ToL0B; + using TileMmad = Catlass::Gemm::Tile::TileMmadTla; + + LocalTensor l1Akk = resource_.l1Buf.template GetBufferByByte( + L1AkkOffset(slot)); + LocalTensor l1W = resource_.l1Buf.template GetBufferByByte( + L1WOffset(slot)); + LocalTensor l1U = resource_.l1Buf.template GetBufferByByte( + L1UOffset(slot)); + LocalTensor l0A = resource_.l0ABuf.template GetBufferByByte( + KDA_FWD_H_L0A_POST_OFFSET); + LocalTensor l0B = resource_.l0BBuf.template GetBufferByByte( + KDA_FWD_H_L0B_POST_OFFSET); + LocalTensor l0C = resource_.l0CBuf.template GetBufferByByte(0); + + auto tensorL1Akk = tla::MakeTensor( + l1Akk, tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK), + Catlass::Arch::PositionL1{}); + auto tensorL1W = tla::MakeTensor( + l1W, tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM), + Catlass::Arch::PositionL1{}); + auto tensorL1U = tla::MakeTensor( + l1U, tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM), + Catlass::Arch::PositionL1{}); + auto tensorL0A = tla::MakeTensor( + l0A, tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK), + Catlass::Arch::PositionL0A{}); + auto tensorL0B = tla::MakeTensor( + l0B, tla::MakeLayout(KDA_FWD_H_CHUNK, 2 * KDA_FWD_H_DIM), + Catlass::Arch::PositionL0B{}); + auto tensorL0C = tla::MakeTensor( + l0C, tla::MakeLayoutL0C(KDA_FWD_H_CHUNK, 2 * KDA_FWD_H_DIM), + Catlass::Arch::PositionL0C{}); + auto tileL1Akk = GetTile(tensorL1Akk, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK)); + auto tileL1W = GetTile(tensorL1W, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + auto tileL1U = GetTile(tensorL1U, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + auto tileL0A = GetTile(tensorL0A, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK)); + auto tileL0BW = GetTile(tensorL0B, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + auto tileL0BU = GetTile(tensorL0B, tla::MakeCoord(0, KDA_FWD_H_DIM), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + auto tileL0B = GetTile(tensorL0B, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, 2 * KDA_FWD_H_DIM)); + auto tileL0C = GetTile(tensorL0C, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, 2 * KDA_FWD_H_DIM)); + auto tileL0CW = GetTile(tensorL0C, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + auto tileL0CU = GetTile(tensorL0C, tla::MakeCoord(0, KDA_FWD_H_DIM), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + + WaitFlag(aicMte2ToMte1Event_); + WaitFlag(stateL0FreeEvent_); + CopyL1ToL0A{}(tileL0A, tileL1Akk); + CopyL1ToL0B{}(tileL0BW, tileL1W); + CopyL1ToL0B{}(tileL0BU, tileL1U); + SetFlag(aicMte1ToMEvent_); + WaitFlag(aicMte1ToMEvent_); + TileMmad{}(tileL0C, tileL0A, tileL0B, KDA_FWD_H_CHUNK, + 2 * KDA_FWD_H_DIM, KDA_FWD_H_CHUNK, true, 0); + SetFlag(stateL0FreeEvent_); + PublishDirectTile( + tileL0CW, KDA_FWD_H_CHUNK, KDA_FWD_H_DIM, + aicMToFixEvent_, aicFixToMEvent_); + WaitFlag(aicFixToMEvent_); + PublishDirectTile( + tileL0CU, KDA_FWD_H_CHUNK, KDA_FWD_H_DIM, + aicMToFixEvent_, aicFixToMEvent_); + WaitFlag(aicFixToMEvent_); + } + + __aicore__ inline void ComputeStateProductsAic( + uint64_t b, uint64_t hv, uint64_t chunk, bool inputsReady = false) + { + const uint32_t slot = static_cast(chunk & 3); + using LayoutTagL1A = typename TileCopyRM::LayoutTagL1A; + using LayoutTagL1B = typename TileCopyRM::LayoutTagL1B; + using LayoutTagL0A = typename TileCopyRM::LayoutTagL0A; + using LayoutTagL0B = typename TileCopyRM::LayoutTagL0B; + using CopyL1ToL0A = typename TileCopyRM::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopyRM::CopyL1ToL0B; + using TileMmad = Catlass::Gemm::Tile::TileMmadTla; + + LocalTensor l1A0 = resource_.l1Buf.template GetBufferByByte( + L1WOffset(slot)); + LocalTensor l1A1 = resource_.l1Buf.template GetBufferByByte( + L1QOffset(slot)); + LocalTensor l1B = + resource_.l1Buf.template GetBufferByByte(KDA_FWD_H_L1_H_OFFSET); + LocalTensor l0A = resource_.l0ABuf.template GetBufferByByte( + KDA_FWD_H_L0A_STATE_OFFSET); + LocalTensor l0B = resource_.l0BBuf.template GetBufferByByte( + KDA_FWD_H_L0B_STATE_OFFSET); + LocalTensor l0C = resource_.l0CBuf.template GetBufferByByte(0); + + auto layoutL1A = tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM); + auto layoutL1B = tla::MakeLayout(KDA_FWD_H_DIM, KDA_FWD_H_DIM); + auto layoutL0A = tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM); + auto layoutL0B = tla::MakeLayout(KDA_FWD_H_DIM, KDA_FWD_H_DIM); + auto layoutL0C = tla::MakeLayoutL0C(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM); + auto tensorL1A0 = tla::MakeTensor(l1A0, layoutL1A, Catlass::Arch::PositionL1{}); + auto tensorL1A1 = tla::MakeTensor(l1A1, layoutL1A, Catlass::Arch::PositionL1{}); + auto tensorL1B = tla::MakeTensor(l1B, layoutL1B, Catlass::Arch::PositionL1{}); + auto tensorL0A = tla::MakeTensor(l0A, layoutL0A, Catlass::Arch::PositionL0A{}); + auto tensorL0B = tla::MakeTensor(l0B, layoutL0B, Catlass::Arch::PositionL0B{}); + auto tensorL0C = tla::MakeTensor(l0C, layoutL0C, Catlass::Arch::PositionL0C{}); + auto tileL1B = GetTile(tensorL1B, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_DIM, KDA_FWD_H_DIM)); + auto tileL0A = GetTile(tensorL0A, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + auto tileL1W = GetTile(tensorL1A0, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + auto tileL1Q = GetTile(tensorL1A1, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + auto tileL0B = GetTile(tensorL0B, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_DIM, KDA_FWD_H_DIM)); + auto tileL0C = GetTile(tensorL0C, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + TileMmad tileMmad; + + if (!inputsReady) { + WaitFlag(aicMte2ToMte1Event_); + } + WaitFlag(stateL0FreeEvent_); + copyL1ToL0A(tileL0A, tileL1W); + copyL1ToL0B(tileL0B, tileL1B); + SetFlag(aicMte1ToMEvent_); + WaitFlag(aicMte1ToMEvent_); + tileMmad(tileL0C, tileL0A, tileL0B, KDA_FWD_H_CHUNK, + KDA_FWD_H_DIM, KDA_FWD_H_DIM, true, 0); + SetFlag(stateL0FreeEvent_); + PublishDirect( + l0C, KDA_FWD_H_CHUNK, KDA_FWD_H_DIM, + aicMToFixEvent_, aicFixToMEvent_); + SetL1SlotFlagAicToAiv(KDA_FWD_H_STATE_FREE_FLAG); + if (fusePostWuIntoFwdH_) { + SetL1SlotFlagAicToAiv(KDA_FWD_H_STATE_FREE_FLAG); + } + + WaitFlag(stateL0FreeEvent_); + WaitFlag(aicFixToMEvent_); + copyL1ToL0A(tileL0A, tileL1Q); + SetFlag(aicMte1ToMEvent_); + WaitFlag(aicMte1ToMEvent_); + tileMmad(tileL0C, tileL0A, tileL0B, KDA_FWD_H_CHUNK, + KDA_FWD_H_DIM, KDA_FWD_H_DIM, true, 0); + SetFlag(stateL0FreeEvent_); + PublishDirect( + l0C, KDA_FWD_H_CHUNK, KDA_FWD_H_DIM, + aicMToFixEvent_, aicFixToMEvent_); + WaitFlag(aicFixToMEvent_); + } + + __aicore__ inline void ComputeVnewProductsAic( + uint64_t b, uint64_t hv, uint64_t chunk, bool prefetchNext) + { + const uint32_t slot = static_cast(chunk & 3); + using LayoutTagL1AK = typename TileCopyCM::LayoutTagL1A; + using LayoutTagL1AA = typename TileCopyRM::LayoutTagL1A; + using LayoutTagL1B = typename TileCopyRM::LayoutTagL1B; + using LayoutTagL0AK = typename TileCopyCM::LayoutTagL0A; + using LayoutTagL0AA = typename TileCopyRM::LayoutTagL0A; + using LayoutTagL0B = typename TileCopyRM::LayoutTagL0B; + using CopyL1ToL0AK = typename TileCopyCM::CopyL1ToL0A; + using CopyL1ToL0AA = typename TileCopyRM::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopyRM::CopyL1ToL0B; + using TileMmadK = Catlass::Gemm::Tile::TileMmadTla; + using TileMmadA = Catlass::Gemm::Tile::TileMmadTla; + + auto layoutKg = tla::MakeLayout(KDA_FWD_H_DIM, KDA_FWD_H_CHUNK); + auto layoutAqk = tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK); + LocalTensor l1Kg = resource_.l1Buf.template GetBufferByByte( + L1KgOffset(slot)); + LocalTensor l1Aqk = resource_.l1Buf.template GetBufferByByte( + L1AqkOffset(slot)); + LocalTensor l1V = + resource_.l1Buf.template GetBufferByByte(KDA_FWD_H_L1_V_OFFSET); + LocalTensor l0A = resource_.l0ABuf.template GetBufferByByte( + KDA_FWD_H_L0A_VNEW_OFFSET); + LocalTensor l0B = resource_.l0BBuf.template GetBufferByByte( + KDA_FWD_H_L0B_VNEW_OFFSET); + LocalTensor l0C = resource_.l0CBuf.template GetBufferByByte(0); + + auto layoutL1Kg = tla::MakeLayout(KDA_FWD_H_DIM, KDA_FWD_H_CHUNK); + auto layoutL1Aqk = tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK); + auto layoutL1V = tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM); + auto layoutL0Kg = tla::MakeLayout(KDA_FWD_H_DIM, KDA_FWD_H_CHUNK); + auto layoutL0Aqk = tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK); + auto layoutL0V = tla::MakeLayout(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM); + auto baseL1Kg = tla::MakeTensor(l1Kg, layoutL1Kg, Catlass::Arch::PositionL1{}); + auto baseL1Aqk = tla::MakeTensor(l1Aqk, layoutL1Aqk, Catlass::Arch::PositionL1{}); + auto baseL1V = tla::MakeTensor(l1V, layoutL1V, Catlass::Arch::PositionL1{}); + auto baseL0Kg = tla::MakeTensor(l0A, layoutL0Kg, Catlass::Arch::PositionL0A{}); + auto baseL0Aqk = tla::MakeTensor(l0A, layoutL0Aqk, Catlass::Arch::PositionL0A{}); + auto baseL0V = tla::MakeTensor(l0B, layoutL0V, Catlass::Arch::PositionL0B{}); + auto baseL0Update = tla::MakeTensor( + l0C, tla::MakeLayoutL0C(KDA_FWD_H_DIM, KDA_FWD_H_DIM), + Catlass::Arch::PositionL0C{}); + auto baseL0Out = tla::MakeTensor( + l0C, tla::MakeLayoutL0C(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM), + Catlass::Arch::PositionL0C{}); + auto tensorL1Kg = GetTile(baseL1Kg, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_DIM, KDA_FWD_H_CHUNK)); + auto tensorL1Aqk = GetTile(baseL1Aqk, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK)); + auto tensorL1V = GetTile(baseL1V, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + auto tensorL0Kg = GetTile(baseL0Kg, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_DIM, KDA_FWD_H_CHUNK)); + auto tensorL0Aqk = GetTile(baseL0Aqk, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_CHUNK)); + auto tensorL0V = GetTile(baseL0V, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + auto tensorL0Update = GetTile(baseL0Update, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_DIM, KDA_FWD_H_DIM)); + auto tensorL0Out = GetTile(baseL0Out, tla::MakeCoord(0, 0), + tla::MakeShape(KDA_FWD_H_CHUNK, KDA_FWD_H_DIM)); + + CopyL1ToL0AK copyL1ToL0AK; + CopyL1ToL0AA copyL1ToL0AA; + CopyL1ToL0B copyL1ToL0B; + TileMmadK tileMmadK; + TileMmadA tileMmadA; + + WaitFlag(vnewL0FreeEvent_); + copyL1ToL0AK(tensorL0Kg, tensorL1Kg); + copyL1ToL0B(tensorL0V, tensorL1V); + SetFlag(aicMte1ToMEvent_); + WaitFlag(aicMte1ToMEvent_); + tileMmadK(tensorL0Update, tensorL0Kg, tensorL0V, + KDA_FWD_H_DIM, KDA_FWD_H_DIM, KDA_FWD_H_CHUNK, true, 0); + SetFlag(vnewL0FreeEvent_); + PublishDirect( + l0C, KDA_FWD_H_DIM, KDA_FWD_H_DIM, + aicMToFixEvent_, aicFixToMEvent_); + SetL1SlotFlagAicToAiv(KDA_FWD_H_VNEW_FREE_FLAG); + + WaitFlag(vnewL0FreeEvent_); + WaitFlag(aicFixToMEvent_); + copyL1ToL0AA(tensorL0Aqk, tensorL1Aqk); + SetFlag(aicL1ReuseEvents_[slot]); + if (prefetchNext) { + PrefetchIndependentProductsAic(b, hv, chunk + 1); + } + SetFlag(aicMte1ToMEvent_); + WaitFlag(aicMte1ToMEvent_); + tileMmadA(tensorL0Out, tensorL0Aqk, tensorL0V, + KDA_FWD_H_CHUNK, KDA_FWD_H_DIM, KDA_FWD_H_CHUNK, true, 0); + SetFlag(vnewL0FreeEvent_); + PublishDirect( + l0C, KDA_FWD_H_CHUNK, KDA_FWD_H_DIM, + aicMToFixEvent_, aicFixToMEvent_); + WaitFlag(aicFixToMEvent_); + } + + __aicore__ inline void ProcessAic() + { + SetLoadDataPaddingValue(static_cast(0)); + SetFlag(stateL0FreeEvent_); + for (uint32_t slot = 0; slot < KDA_FWD_H_L1_STAGING_DEPTH; ++slot) { + SetFlag(aicL1ReuseEvents_[slot]); + } + const uint64_t coreIdx = static_cast(GetBlockIdx()); + const uint64_t coreNum = coreNum_ == 0 ? 1 : coreNum_; + for (uint64_t task = coreIdx; task < batch_ * heads_; task += coreNum) { + const uint64_t b = task / heads_; + const uint64_t hv = task % heads_; + PrefetchIndependentProductsAic(b, hv, 0); + for (uint64_t chunk = 0; chunk < totalChunks_; ++chunk) { + if (fusePostWuIntoFwdH_) { + ComputePostWuAic(chunk); + } + WaitL1SlotReadyMte1(KDA_FWD_H_STATE_READY_FLAG); + if (fusePostWuIntoFwdH_) { + WaitL1SlotReadyMte1(KDA_FWD_H_STATE_READY_FLAG); + } + ComputeStateProductsAic(b, hv, chunk, fusePostWuIntoFwdH_); + WaitL1SlotReadyMte1(KDA_FWD_H_VNEW_READY_FLAG); + ComputeVnewProductsAic( + b, hv, chunk, chunk + 1 < totalChunks_); + } + CrossCoreWaitFlag(KDA_FWD_H_DIRECT_FREE_FLAG); + } + WaitFlag(stateL0FreeEvent_); + for (uint32_t slot = 0; slot < KDA_FWD_H_L1_STAGING_DEPTH; ++slot) { + WaitFlag(aicL1ReuseEvents_[slot]); + } + } + + __aicore__ inline void CopyOutputRows( + uint64_t b, uint64_t hv, uint64_t tokenStart, + LocalTensor src, uint32_t rows) + { + DataCopyExtParams params{ + static_cast(rows), + static_cast(KDA_FWD_H_DIM * sizeof(T)), + 0, + static_cast((heads_ * KDA_FWD_H_DIM - KDA_FWD_H_DIM) * sizeof(T)), + 0}; + DataCopyPad(attnOut_[OutputOffset(b, hv, tokenStart)], src, params); + } + + __aicore__ inline void InitializeStateAiv( + uint64_t b, uint64_t hv, uint32_t rowBegin, + LocalTensor state) + { + if (hasInitialState_) { + DataCopy(state, initialState_[StateOffset(b, hv) + rowBegin * KDA_FWD_H_DIM], + KDA_FWD_H_STATE_SUB_ELEMS); + SetFlag(aivMte2ToVEvent_); + WaitFlag(aivMte2ToVEvent_); + } else { + Duplicate(state, 0.0f, KDA_FWD_H_STATE_SUB_ELEMS); + PipeBarrier(); + } + } + + __aicore__ inline void StoreCurrentStateAiv( + uint64_t b, uint64_t hv, uint64_t chunk, uint32_t rowBegin, + LocalTensor state, LocalTensor stateTyped) + { + if (storeH_) { + Cast(stateTyped, state, RoundMode::CAST_RINT, KDA_FWD_H_STATE_SUB_ELEMS); + PipeBarrier(); + SetFlag(aivVToMte3Event_); + WaitFlag(aivVToMte3Event_); + DataCopy(h_[HOffset(b, hv, chunk) + rowBegin * KDA_FWD_H_DIM], + stateTyped, KDA_FWD_H_STATE_SUB_ELEMS); + SetFlag(aivMte3ToVEvent_); + WaitFlag(aivMte3ToVEvent_); + } + + constexpr uint32_t columnGroup = 64; + constexpr uint32_t columnGroups = KDA_FWD_H_DIM / columnGroup; + for (uint32_t group = 0; group < columnGroups; ++group) { + Cast(stateTyped[group * KDA_FWD_H_SUB_DIM * columnGroup], + state[group * columnGroup], RoundMode::CAST_RINT, + columnGroup, KDA_FWD_H_SUB_DIM, + {static_cast(KDA_FWD_H_SUB_DIM), 1, 1, + static_cast(columnGroups * 8)}); + } + PipeBarrier(); + SetFlag(aivVToMte3Event_); + WaitFlag(aivVToMte3Event_); + LocalTensor l1State = + resource_.l1Buf.template GetBufferByByte(KDA_FWD_H_L1_H_OFFSET); + const uint32_t subBlockIdx = rowBegin / KDA_FWD_H_SUB_DIM; + WaitL1SlotFreeMte3( + KDA_FWD_H_STATE_FREE_FLAG, statePublishCount_[subBlockIdx]); + DataCopyParams stateCopyParams; + stateCopyParams.blockCount = KDA_FWD_H_DIM / 16; + stateCopyParams.blockLen = KDA_FWD_H_SUB_DIM; + stateCopyParams.srcGap = 0; + stateCopyParams.dstGap = KDA_FWD_H_DIM - KDA_FWD_H_SUB_DIM; + DataCopy(l1State[rowBegin * 16], stateTyped, stateCopyParams); + SetL1SlotFlagAivToAic(KDA_FWD_H_STATE_READY_FLAG); + ++statePublishCount_[subBlockIdx]; + SetFlag(aivMte3ToVEvent_); + WaitFlag(aivMte3ToVEvent_); + } + + __aicore__ inline void ProcessChunkAiv( + uint64_t b, uint64_t hv, uint32_t chunk, + uint32_t subBlockIdx, LocalTensor state, + LocalTensor stateTyped, LocalTensor direct, + LocalTensor out1, LocalTensor vnew, + LocalTensor ioTyped, LocalTensor gate) + { + const uint32_t tokenBegin = subBlockIdx * KDA_FWD_H_SUB_CHUNK; + const uint32_t stateRowBegin = subBlockIdx * KDA_FWD_H_SUB_DIM; + StoreCurrentStateAiv(b, hv, chunk, stateRowBegin, state, stateTyped); + + CrossCoreWaitFlag(KDA_FWD_H_DIRECT_READY_FLAG); + WaitFlag(aivMte3ToMte2Event_); + if (fusePostWuIntoFwdH_) { + Cast(ioTyped, direct, RoundMode::CAST_RINT, KDA_FWD_H_TOKEN_SUB_ELEMS); + PipeBarrier(); + CrossCoreSetFlag<0x2, PIPE_V>(KDA_FWD_H_DIRECT_FREE_FLAG); + + SetFlag(aivVToMte3Event_); + WaitFlag(aivVToMte3Event_); + LocalTensor l1W = resource_.l1Buf.template GetBufferByByte( + L1WOffset(chunk & 3)); + WaitL1SlotFreeMte3( + KDA_FWD_H_STATE_FREE_FLAG, wPublishCount_[subBlockIdx]); + DataCopyParams wCopyParams; + wCopyParams.blockCount = KDA_FWD_H_SUB_CHUNK; + wCopyParams.blockLen = 1; + wCopyParams.srcGap = KDA_FWD_H_DIM / 16 - 1; + wCopyParams.dstGap = 0; + for (uint32_t colBlock = 0; colBlock < KDA_FWD_H_DIM / 16; ++colBlock) { + const uint32_t dstOffset = + colBlock * KDA_FWD_H_CHUNK * 16 + tokenBegin * 16; + DataCopy(l1W[dstOffset], ioTyped[colBlock * 16], wCopyParams); + } + SetFlag(aivMte3ToVEvent_); + WaitFlag(aivMte3ToVEvent_); + SetL1SlotFlagAivToAic(KDA_FWD_H_STATE_READY_FLAG); + ++wPublishCount_[subBlockIdx]; + + CrossCoreWaitFlag(KDA_FWD_H_DIRECT_READY_FLAG); + Cast(ioTyped, direct, RoundMode::CAST_RINT, KDA_FWD_H_TOKEN_SUB_ELEMS); + PipeBarrier(); + Cast(vnew, ioTyped, RoundMode::CAST_NONE, KDA_FWD_H_TOKEN_SUB_ELEMS); + PipeBarrier(); + CrossCoreSetFlag<0x2, PIPE_V>(KDA_FWD_H_DIRECT_FREE_FLAG); + CrossCoreWaitFlag(KDA_FWD_H_DIRECT_READY_FLAG); + } else { + DataCopy(ioTyped, u_[ChunkMatrixOffset(b, hv, chunk, tokenBegin)], + KDA_FWD_H_TOKEN_SUB_ELEMS); + SetFlag(aivMte2ToVEvent_); + WaitFlag(aivMte2ToVEvent_); + Cast(vnew, ioTyped, RoundMode::CAST_NONE, KDA_FWD_H_TOKEN_SUB_ELEMS); + PipeBarrier(); + } + Sub(vnew, vnew, direct, KDA_FWD_H_TOKEN_SUB_ELEMS); + PipeBarrier(); + CrossCoreSetFlag<0x2, PIPE_V>(KDA_FWD_H_DIRECT_FREE_FLAG); + if (storeVNew_) { + Cast(ioTyped, vnew, RoundMode::CAST_RINT, KDA_FWD_H_TOKEN_SUB_ELEMS); + PipeBarrier(); + SetFlag(aivVToMte3Event_); + WaitFlag(aivVToMte3Event_); + DataCopy(vNew_[VNewOffset(b, hv, chunk, tokenBegin)], + ioTyped, KDA_FWD_H_TOKEN_SUB_ELEMS); + SetFlag(aivMte3ToVEvent_); + WaitFlag(aivMte3ToVEvent_); + } + + constexpr uint32_t columnGroup = 64; + constexpr uint32_t columnGroups = KDA_FWD_H_DIM / columnGroup; + for (uint32_t group = 0; group < columnGroups; ++group) { + Cast(ioTyped[group * KDA_FWD_H_SUB_CHUNK * columnGroup], + vnew[group * columnGroup], RoundMode::CAST_RINT, + columnGroup, KDA_FWD_H_SUB_CHUNK, + {static_cast(KDA_FWD_H_SUB_CHUNK), 1, 1, + static_cast(columnGroups * 8)}); + } + PipeBarrier(); + SetFlag(aivVToMte3Event_); + WaitFlag(aivVToMte3Event_); + LocalTensor l1Vnew = + resource_.l1Buf.template GetBufferByByte(KDA_FWD_H_L1_V_OFFSET); + WaitL1SlotFreeMte3( + KDA_FWD_H_VNEW_FREE_FLAG, vnewPublishCount_[subBlockIdx]); + DataCopyParams vnewL1CopyParams; + vnewL1CopyParams.blockCount = KDA_FWD_H_DIM / 16; + vnewL1CopyParams.blockLen = KDA_FWD_H_SUB_CHUNK; + vnewL1CopyParams.srcGap = 0; + vnewL1CopyParams.dstGap = KDA_FWD_H_CHUNK - KDA_FWD_H_SUB_CHUNK; + DataCopy(l1Vnew[tokenBegin * 16], ioTyped, vnewL1CopyParams); + SetFlag(aivMte3ToVEvent_); + WaitFlag(aivMte3ToVEvent_); + SetL1SlotFlagAivToAic(KDA_FWD_H_VNEW_READY_FLAG); + ++vnewPublishCount_[subBlockIdx]; + + CrossCoreWaitFlag(KDA_FWD_H_DIRECT_READY_FLAG); + Muls(out1, direct, scale_, KDA_FWD_H_TOKEN_SUB_ELEMS); + PipeBarrier(); + CrossCoreSetFlag<0x2, PIPE_V>(KDA_FWD_H_DIRECT_FREE_FLAG); + + CrossCoreWaitFlag(KDA_FWD_H_DIRECT_READY_FLAG); + DataCopy(gate, gk_[ChunkMatrixOffset(b, hv, chunk, KDA_FWD_H_CHUNK - 1) + stateRowBegin], KDA_FWD_H_SUB_DIM); + SetFlag(aivMte2ToVEvent_); + WaitFlag(aivMte2ToVEvent_); + Muls(gate, gate, 0.6931471805599453f, KDA_FWD_H_SUB_DIM); + PipeBarrier(); + Exp(gate, gate, KDA_FWD_H_SUB_DIM); + PipeBarrier(); + AscendC::VF_CALL>( + reinterpret_cast<__ubuf__ float *>(direct.GetPhyAddr()), + reinterpret_cast<__ubuf__ float *>(state.GetPhyAddr()), + reinterpret_cast<__ubuf__ float *>(gate.GetPhyAddr()), + static_cast(KDA_FWD_H_SUB_DIM), + static_cast(KDA_FWD_H_DIM)); + PipeBarrier(); + Adds(state, direct, 0.0f, KDA_FWD_H_STATE_SUB_ELEMS); + PipeBarrier(); + + CrossCoreSetFlag<0x2, PIPE_V>(KDA_FWD_H_DIRECT_FREE_FLAG); + CrossCoreWaitFlag(KDA_FWD_H_DIRECT_READY_FLAG); + Add(direct, direct, out1, KDA_FWD_H_TOKEN_SUB_ELEMS); + PipeBarrier(); + Cast(ioTyped, direct, RoundMode::CAST_RINT, KDA_FWD_H_TOKEN_SUB_ELEMS); + PipeBarrier(); + SetFlag(aivVToMte3Event_); + WaitFlag(aivVToMte3Event_); + CopyOutputRows( + b, hv, chunk * KDA_FWD_H_CHUNK + tokenBegin, + ioTyped, KDA_FWD_H_SUB_CHUNK); + CrossCoreSetFlag<0x2, PIPE_V>(KDA_FWD_H_DIRECT_FREE_FLAG); + SetFlag(aivMte3ToMte2Event_); + } + + __aicore__ inline void ProcessAiv() + { + const uint32_t subBlockIdx = static_cast(GetSubBlockIdx()); + const uint32_t subBlockNum = static_cast(GetSubBlockNum()); + const uint64_t coreIdx = static_cast(GetBlockIdx()) / subBlockNum; + const uint64_t coreNum = coreNum_ == 0 ? 1 : coreNum_; + LocalTensor state = + resource_.ubBuf.template GetBufferByByte(KDA_FWD_H_UB_STATE_OFFSET); + LocalTensor stateTyped = + resource_.ubBuf.template GetBufferByByte(KDA_FWD_H_UB_STATE_TYPED_OFFSET); + LocalTensor direct = + resource_.ubBuf.template GetBufferByByte(KDA_FWD_H_UB_DIRECT_OFFSET); + LocalTensor out1 = + resource_.ubBuf.template GetBufferByByte(KDA_FWD_H_UB_OUT1_OFFSET); + LocalTensor vnew = + resource_.ubBuf.template GetBufferByByte(KDA_FWD_H_UB_VNEW_OFFSET); + LocalTensor ioTyped = + resource_.ubBuf.template GetBufferByByte(KDA_FWD_H_UB_IO_OFFSET); + LocalTensor gate = + resource_.ubBuf.template GetBufferByByte(KDA_FWD_H_UB_GATE_OFFSET); + + SetFlag(aivMte3ToMte2Event_); + for (uint64_t task = coreIdx; task < batch_ * heads_; task += coreNum) { + const uint64_t b = task / heads_; + const uint64_t hv = task % heads_; + const uint32_t stateRowBegin = subBlockIdx * KDA_FWD_H_SUB_DIM; + InitializeStateAiv(b, hv, stateRowBegin, state); + CrossCoreSetFlag<0x2, PIPE_V>(KDA_FWD_H_DIRECT_FREE_FLAG); + for (uint32_t chunk = 0; + chunk < static_cast(totalChunks_); ++chunk) { + ProcessChunkAiv( + b, hv, chunk, subBlockIdx, + state, stateTyped, direct, out1, vnew, ioTyped, gate); + } + if (storeFinalState_) { + SetFlag(aivVToMte3Event_); + WaitFlag(aivVToMte3Event_); + DataCopy(finalState_[StateOffset(b, hv) + + stateRowBegin * KDA_FWD_H_DIM], + state, KDA_FWD_H_STATE_SUB_ELEMS); + SetFlag(aivMte3ToVEvent_); + WaitFlag(aivMte3ToVEvent_); + } + } + WaitL1SlotFreeMte3( + KDA_FWD_H_STATE_FREE_FLAG, statePublishCount_[subBlockIdx]); + if (fusePostWuIntoFwdH_) { + WaitL1SlotFreeMte3( + KDA_FWD_H_STATE_FREE_FLAG, wPublishCount_[subBlockIdx]); + } + WaitL1SlotFreeMte3( + KDA_FWD_H_VNEW_FREE_FLAG, vnewPublishCount_[subBlockIdx]); + WaitFlag(aivMte3ToMte2Event_); + } + +private: + GlobalTensor gk_; + GlobalTensor initialState_; + GlobalTensor attnOut_; + GlobalTensor finalState_; + GlobalTensor aqk_; + GlobalTensor akk_; + GlobalTensor w_; + GlobalTensor u_; + GlobalTensor qg_; + GlobalTensor kg_; + GlobalTensor vNew_; + GlobalTensor h_; + uint32_t statePublishCount_[2] = {0, 0}; + uint32_t vnewPublishCount_[2] = {0, 0}; + uint32_t wPublishCount_[2] = {0, 0}; + TEventID aicMte2ToMte1Event_ = KDA_FWD_H_MTE_W_EVENT; + TEventID aicL1ReuseEvents_[KDA_FWD_H_L1_STAGING_DEPTH] = {0, 1, 2, 3}; + TEventID stateL0FreeEvent_ = KDA_FWD_H_M_EVENT; + TEventID vnewL0FreeEvent_ = KDA_FWD_H_M_EVENT; + TEventID aicMte1ToMEvent_ = KDA_FWD_H_M_EVENT; + TEventID aicMToFixEvent_ = KDA_FWD_H_FIX_EVENT; + TEventID aicFixToMEvent_ = KDA_FWD_H_FIX_EVENT; + TEventID aivMte2ToVEvent_ = KDA_FWD_H_MTE_W_EVENT; + TEventID aivVToMte3Event_ = KDA_FWD_H_MTE_Q_EVENT; + TEventID aivMte3ToVEvent_ = KDA_FWD_H_MTE_B_EVENT; + TEventID aivMte3ToMte2Event_ = KDA_FWD_H_IO_REUSE_EVENT; + Catlass::Arch::Resource resource_; + uint64_t batch_ = 0; + uint64_t heads_ = 0; + uint64_t seqlen_ = 0; + uint64_t totalChunks_ = 0; + uint64_t coreNum_ = 1; + bool hasInitialState_ = false; + bool storeFinalState_ = false; + bool storeVNew_ = false; + bool storeH_ = false; + bool fusePostWuIntoFwdH_ = false; + float scale_ = 1.0f; +}; + +} // namespace KdaForward::arch35 + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_impl.h b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_impl.h new file mode 100644 index 0000000000..beee398641 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_impl.h @@ -0,0 +1,61 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#pragma once + +#include "../chunk_kda_fwd_common.h" +#include "chunk_kda_fwd_fwd_h.h" + +namespace KdaForward::arch35 { + +template +__aicore__ inline void Run( + GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR g, GM_ADDR beta, + GM_ADDR aLog, GM_ADDR dtBias, GM_ADDR initialState, + GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR attnOut, + GM_ADDR finalState, GM_ADDR gk, GM_ADDR aqk, GM_ADDR akk, + GM_ADDR w, GM_ADDR u, GM_ADDR qg, GM_ADDR kg, GM_ADDR vNew, GM_ADDR h, + GM_ADDR userWorkspace, const TilingData &tiling, AscendC::TPipe &pipe) +{ + const auto addresses = ResolveAddresses( + finalState, gk, w, u, qg, kg, vNew, h, userWorkspace, tiling); + RunFrontEnd( + q, k, v, g, beta, aLog, dtBias, initialState, cuSeqlens, + chunkIndices, aqk, akk, addresses, userWorkspace, tiling, pipe); + if (tiling.useDenseFwdH) { + ChunkKdaFwdFwdH fwdH; + fwdH.Init( + addresses.gk, initialState, attnOut, addresses.finalState, + aqk, akk, addresses.w, addresses.u, addresses.qgScaled, + addresses.kg, addresses.vNew, addresses.h, tiling); + fwdH.Process(); + return; + } + + const int64_t fwdHTaskCount = + (tiling.isVarLen ? tiling.seqNum : tiling.batch) * tiling.vHeadNum; + const bool isolateGenericBackEnd = + (!tiling.isVarLen && tiling.seqlen % tiling.chunkSize == 0) || + fwdHTaskCount > tiling.prepareUsedCoreNum; + if (isolateGenericBackEnd) { + pipe.Destroy(); + RunGenericBackEnd( + q, k, v, beta, initialState, cuSeqlens, chunkIndices, aqk, + attnOut, addresses, userWorkspace, tiling); + } else { + RunGenericBackEnd( + q, k, v, beta, initialState, cuSeqlens, chunkIndices, aqk, + attnOut, addresses, userWorkspace, tiling, pipe); + } +} + +} // namespace KdaForward::arch35 diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_post_wu_arch35.h b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_post_wu_arch35.h new file mode 100644 index 0000000000..5e9eb01125 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_post_wu_arch35.h @@ -0,0 +1,2034 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#pragma once + +#ifndef CATLASS_ARCH +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +#define CATLASS_ARCH 3510 +#else +#define CATLASS_ARCH 2201 +#endif +#endif + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/arch/kda_cross_core_sync.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/gemm/block/kda_block_mmad.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "../kernel_utils/block/block_mmad_pingpong_tla_multi.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "kernel_operator.h" +#include "../chunk_kda_fwd_varlen.h" +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +#ifndef FLA_NPU_REGBASE_HPP_INCLUDED +#define FLA_NPU_REGBASE_HPP_INCLUDED +#include "../kernel_utils/vector/regbase.hpp" +#endif +#endif +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +using namespace AscendC; + +namespace KdaPostWu { +namespace { +using KdaInt64 = tla::Int<64>; +using KdaInt128 = tla::Int<128>; +constexpr float LN2 = 0.69314718055994530942f; +constexpr float KDA_EXP2_CLAMP = 80.0f; +constexpr float KDA_EXP_INPUT_MAX = KDA_EXP2_CLAMP * LN2; +constexpr float KDA_EXP_INPUT_MIN = -KDA_EXP2_CLAMP * LN2; +constexpr float KDA_FP16_MAX = 65504.0f; +constexpr uint32_t EXP2_UB_ELEMENTS = 256; +constexpr uint32_t EXP2_UB_BYTES = EXP2_UB_ELEMENTS * (sizeof(float) + sizeof(uint16_t)); +constexpr uint32_t EXP2_EVENT_ID = 0; +constexpr uint32_t KDA_SOLVE_BT = 64; +constexpr uint32_t KDA_SOLVE_MATRIX_ELEMENTS = KDA_SOLVE_BT * KDA_SOLVE_BT; +constexpr uint32_t KDA_SOLVE_SCRATCH_X = 0; +constexpr uint32_t KDA_SOLVE_SCRATCH_Y0 = 1; +constexpr uint32_t KDA_SOLVE_SCRATCH_TMP = 2; +constexpr uint32_t KDA_SOLVE_SCRATCH_Y1 = 3; +constexpr uint32_t KDA_SOLVE_SCRATCH_IDENTITY = 4; +constexpr uint32_t KDA_SOLVE_SCRATCH_SLOTS = 5; +constexpr uint32_t KDA_SOLVE_DIAG_BT = 16; +constexpr uint32_t KDA_SOLVE_DIAG_BLOCKS = KDA_SOLVE_BT / KDA_SOLVE_DIAG_BT; +constexpr uint32_t KDA_SOLVE_DIAG_MCH_ITERS = 3; +constexpr uint32_t KDA_SCORE_REF_BC = 16; +constexpr uint32_t KDA_VEC_ARENA_ELEMENTS = 32768; +constexpr uint32_t KDA_BITS_PER_MASK_BYTE = 8; +constexpr uint32_t KDA_SELECT_COL_BLOCKS = 2; +constexpr uint32_t KDA_SELECT_COL_MASK_BYTES = KDA_SOLVE_MATRIX_ELEMENTS / KDA_BITS_PER_MASK_BYTE; +constexpr uint32_t KDA_SELECT_MASK_BYTES = KDA_SELECT_COL_BLOCKS * KDA_SELECT_COL_MASK_BYTES; +constexpr uint32_t KDA_SELECT_AQK_MASK_BYTE_OFFSET = 120 * 1024; +constexpr uint32_t KDA_SELECT_AKK_MASK_BYTE_OFFSET = KDA_SELECT_AQK_MASK_BYTE_OFFSET + KDA_SELECT_MASK_BYTES; +constexpr uint32_t KDA_SELECT_ZERO_BYTE_OFFSET = KDA_SELECT_AKK_MASK_BYTE_OFFSET + KDA_SELECT_MASK_BYTES; +constexpr uint32_t KDA_SELECT_ZERO_FLOAT_OFFSET = KDA_SELECT_ZERO_BYTE_OFFSET / sizeof(float); +constexpr uint8_t KDA_SCORE_DONE_FLAG0 = 2; +constexpr uint8_t KDA_SCORE_DONE_FLAG1 = 3; +constexpr uint8_t KDA_SCORE_READY_FLAG0 = 4; +constexpr uint8_t KDA_SCORE_READY_FLAG1 = 5; +constexpr uint32_t KDA_SCORE_QUEUE_DEPTH = 2; +constexpr uint32_t KDA_SYNC_REVERSE_DEPTH = 1; +constexpr uint32_t KDA_SCORE_SCRATCH_PLANES = 3; +constexpr uint32_t KDA_SCORE_SCRATCH_QG = 0; +constexpr uint32_t KDA_SCORE_SCRATCH_W = 1; +constexpr uint32_t KDA_SCORE_SCRATCH_KG = 2; +constexpr uint64_t KDA_WORKSPACE_ALIGN = 512; +constexpr uint32_t KDA_GATE_TILE_ROWS = 32; +constexpr uint32_t KDA_TYPICAL_GATE_TILE_ROWS = 16; +constexpr uint32_t KDA_TYPICAL_GATE_PIPELINE_ROWS = 32; +constexpr uint16_t KDA_TYPICAL_GATE_PIPELINE_STAGES = 3; +constexpr uint32_t KDA_POST_EVENT = 3; +constexpr uint32_t KDA_POST_EVENT_NEXT = 4; +constexpr uint32_t KDA_POST_EVENT_FIX = 5; +constexpr uint32_t KDA_POST_PIPELINE_L1_SLOT_BYTES = 24 * 1024; +constexpr uint32_t KDA_POST_PIPELINE_L1_A_BYTES = 64 * 64 * sizeof(uint16_t); +constexpr uint32_t KDA_POST_PIPELINE_L1_B_BYTES = 64 * 128 * sizeof(uint16_t); +constexpr uint32_t KDA_POST_PIPELINE_L1_U_SLOT_BYTES = 64 * 128 * sizeof(uint16_t); +constexpr uint32_t KDA_POST_PIPELINE_L0_A_SLOT_BYTES = 64 * 64 * sizeof(uint16_t); +constexpr uint32_t KDA_POST_PIPELINE_L0_B_SLOT_BYTES = 64 * 256 * sizeof(uint16_t); +constexpr uint32_t KDA_POST_PIPELINE_L0_C_SLOT_BYTES = 64 * 256 * sizeof(float); +constexpr uint16_t KDA_POST_PIPELINE_STAGE_COUNT = 2; +constexpr uint16_t KDA_POST_FUSED_BATCH_TASKS = 4; +constexpr uint16_t KDA_POST_HEAD_PAIR_LANES = 2; +constexpr uint16_t KDA_POST_PIPELINE_U_EVENT = KDA_POST_EVENT_FIX; +constexpr uint16_t KDA_POST_REGBASE_TAIL_LIMIT = 16; +constexpr bool KDA_ENABLE_POST_AIC_PIPELINE = true; + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +static __simd_vf__ inline void ComputeKdaTailWuRegbase(__ubuf__ float *wOut, __ubuf__ float *uOut, + __ubuf__ bfloat16_t *akk, __ubuf__ bfloat16_t *wSeed, + __ubuf__ bfloat16_t *uSeed, uint16_t rows, uint16_t curT) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t DIM = 128; + constexpr uint16_t FP32_REG_ELEMENTS = AscendC::VECTOR_REG_WIDTH / sizeof(float); + constexpr uint16_t AKK_ROW_ELEMENTS = 64; + MaskReg fullMask = CreateMask(); + + for (uint16_t row = 0; row < rows; ++row) { + RegTensor wAcc0; + RegTensor wAcc1; + RegTensor uAcc0; + RegTensor uAcc1; + Duplicate(wAcc0, 0.0f, fullMask); + Duplicate(wAcc1, 0.0f, fullMask); + Duplicate(uAcc0, 0.0f, fullMask); + Duplicate(uAcc1, 0.0f, fullMask); + + for (uint16_t token = 0; token < curT; ++token) { + RegTensor wRaw0; + RegTensor wRaw1; + RegTensor uRaw0; + RegTensor uRaw1; + RegTensor weightRaw; + RegTensor w0; + RegTensor w1; + RegTensor u0; + RegTensor u1; + RegTensor weight; + RegTensor product; + DataCopy(wRaw0, wSeed + static_cast(token) * DIM); + DataCopy( + wRaw1, wSeed + static_cast(token) * DIM + FP32_REG_ELEMENTS); + DataCopy(uRaw0, uSeed + static_cast(token) * DIM); + DataCopy( + uRaw1, uSeed + static_cast(token) * DIM + FP32_REG_ELEMENTS); + LoadAlign(weightRaw, + akk + static_cast(row) * AKK_ROW_ELEMENTS + token); + Cast(w0, wRaw0, fullMask); + Cast(w1, wRaw1, fullMask); + Cast(u0, uRaw0, fullMask); + Cast(u1, uRaw1, fullMask); + Cast(weight, weightRaw, fullMask); + Mul(product, w0, weight, fullMask); + Add(wAcc0, wAcc0, product, fullMask); + Mul(product, w1, weight, fullMask); + Add(wAcc1, wAcc1, product, fullMask); + Mul(product, u0, weight, fullMask); + Add(uAcc0, uAcc0, product, fullMask); + Mul(product, u1, weight, fullMask); + Add(uAcc1, uAcc1, product, fullMask); + } + + uint32_t outputOffset = static_cast(row) * DIM; + DataCopy(wOut + outputOffset, wAcc0, fullMask); + DataCopy(wOut + outputOffset + FP32_REG_ELEMENTS, wAcc1, fullMask); + DataCopy(uOut + outputOffset, uAcc0, fullMask); + DataCopy(uOut + outputOffset + FP32_REG_ELEMENTS, uAcc1, fullMask); + } +} + +template +__simd_callee__ inline void LoadPostKdaGateRegbasePair(AscendC::MicroAPI::RegTensor &zeroReg, + AscendC::MicroAPI::RegTensor &oneReg, + __ubuf__ InputT *src, AscendC::MicroAPI::MaskReg &inputMask) +{ + using namespace AscendC::MicroAPI; + if constexpr (std::is_same()) { + LoadAlign(zeroReg, oneReg, src); + } else { + RegTensor inputReg; + LoadIn(inputReg, src); + CastHalf2Float(zeroReg, oneReg, inputReg, inputMask); + } +} + +template +__simd_callee__ inline void StorePostKdaGateRegbasePair(__ubuf__ OutputT *dst, + AscendC::MicroAPI::RegTensor &zeroReg, + AscendC::MicroAPI::RegTensor &oneReg, + AscendC::MicroAPI::MaskReg &inputMask, + AscendC::MicroAPI::MaskReg &floatMask) +{ + using namespace AscendC::MicroAPI; + if constexpr (std::is_same()) { + Mins(zeroReg, zeroReg, KDA_FP16_MAX, floatMask); + Mins(oneReg, oneReg, KDA_FP16_MAX, floatMask); + Maxs(zeroReg, zeroReg, -KDA_FP16_MAX, floatMask); + Maxs(oneReg, oneReg, -KDA_FP16_MAX, floatMask); + } + RegTensor outputReg; + CastFloat2Half(outputReg, zeroReg, oneReg, floatMask); + StoreAlign(dst, outputReg, inputMask); +} + +template +static __simd_vf__ inline void ComputePostKdaKgRegbase(__ubuf__ T *kAndKg, __ubuf__ GK_T *gate, __ubuf__ float *ref, + uint16_t rows, uint16_t cols) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t ELEMENTS_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(T); + MaskReg floatMask = CreateMask(); + for (uint16_t row = 0; row < rows; ++row) { + uint32_t rowOffset = static_cast(row) * cols; + for (uint16_t col = 0; col < cols; col += ELEMENTS_PER_REG) { + uint32_t activeCount = static_cast(cols - col); + MaskReg inputMask = UpdateMask(activeCount); + uint32_t offset = rowOffset + col; + + RegTensor gateZeroReg; + RegTensor gateOneReg; + RegTensor refZeroReg; + RegTensor refOneReg; + RegTensor expZeroReg; + RegTensor expOneReg; + RegTensor inputZeroReg; + RegTensor inputOneReg; + RegTensor outputZeroReg; + RegTensor outputOneReg; + + LoadPostKdaGateRegbasePair(gateZeroReg, gateOneReg, gate + offset, inputMask); + LoadAlign(refZeroReg, refOneReg, ref + col); + SubFloatTwoReg(expZeroReg, expOneReg, refZeroReg, refOneReg, gateZeroReg, gateOneReg, floatMask); + Muls(expZeroReg, expZeroReg, LN2, floatMask); + Muls(expOneReg, expOneReg, LN2, floatMask); + MinsFloatTwoReg(expZeroReg, expOneReg, expZeroReg, expOneReg, KDA_EXP_INPUT_MAX, floatMask); + Maxs(expZeroReg, expZeroReg, KDA_EXP_INPUT_MIN, floatMask); + Maxs(expOneReg, expOneReg, KDA_EXP_INPUT_MIN, floatMask); + ExpFloatTwoReg(expZeroReg, expOneReg, expZeroReg, expOneReg, floatMask); + + LoadPostKdaGateRegbasePair(inputZeroReg, inputOneReg, kAndKg + offset, inputMask); + MulFloatTwoReg(outputZeroReg, outputOneReg, inputZeroReg, inputOneReg, expZeroReg, expOneReg, floatMask); + StorePostKdaGateRegbasePair(kAndKg + offset, outputZeroReg, outputOneReg, inputMask, floatMask); + } + } +} +#endif + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +using KdaArchTag = Catlass::Arch::Ascend950; +#else +using KdaArchTag = Catlass::Arch::AtlasA2; +#endif +using KdaDispatchPolicy = Catlass::Gemm::MmadPingpong; +using KdaScoreDispatchPolicy = Catlass::Gemm::MmadPingpongTlaMulti; +static_assert(KdaScoreDispatchPolicy::ENABLE_L1_RESIDENT, + "KDA Aqk/Akk score MMAD must keep the shared right matrix resident in L1"); +static_assert(KdaScoreDispatchPolicy::L1B_STAGES == 1, + "KDA Aqk/Akk score MMAD needs one L1 B slot so the second MMAD reuses it"); +using KdaSolveDispatchPolicy = Catlass::Gemm::MmadPingpong; +static_assert(!KdaSolveDispatchPolicy::USE_HF32_MODE, "KDA triangular solve must use IEEE FP32 Cube mode"); +using KdaL1TileShape = tla::Shape; +using KdaL0TileShape = KdaL1TileShape; +using KdaSolveL1TileShape = tla::Shape; +using KdaSolveL0TileShape = KdaSolveL1TileShape; + +__aicore__ inline uint32_t FloatToBits(float value) +{ + union Bits { + __aicore__ Bits() {} + float f; + uint32_t u; + } bits; + bits.f = value; + return bits.u; +} + +__aicore__ inline float BitsToFloat(uint32_t value) +{ + union Bits { + __aicore__ Bits() {} + uint32_t u; + float f; + } bits; + bits.u = value; + return bits.f; +} + +__aicore__ inline uint16_t Bf16ToBits(bfloat16_t value) +{ + union Bits { + __aicore__ Bits() {} + bfloat16_t f; + uint16_t u; + } bits; + bits.f = value; + return bits.u; +} + +__aicore__ inline bfloat16_t BitsToBf16(uint16_t value) +{ + union Bits { + __aicore__ Bits() {} + uint16_t u; + bfloat16_t f; + } bits; + bits.u = value; + return bits.f; +} + +template +__aicore__ inline T FloatToType(float value) +{ + if constexpr (IsSameType::value) { + uint32_t bits = FloatToBits(value); + uint32_t bias = 0x7FFFu + ((bits >> 16) & 1u); + return BitsToBf16(static_cast((bits + bias) >> 16)); + } + return static_cast(value); +} + +template +class ChunkKdaFwdPostWuKernel { +public: + using OUT_T = T; + using AKK_T = T; + template + __aicore__ inline void Init(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR beta, GM_ADDR initialState, + GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR preparedQG, GM_ADDR preparedAqk, + GM_ADDR propagatedVNew, GM_ADDR propagatedH, GM_ADDR o, GM_ADDR finalState, GM_ADDR aqk, + GM_ADDR akk, GM_ADDR w, GM_ADDR u, GM_ADDR qg, GM_ADDR kg, GM_ADDR vNew, GM_ADDR h, + GM_ADDR workspace, const TilingData &tiling, TPipe *pipe, bool initVecBuffers = true) + { + pipe_ = pipe; + q_.SetGlobalBuffer((__gm__ T *)q); + k_.SetGlobalBuffer((__gm__ T *)k); + v_.SetGlobalBuffer((__gm__ T *)v); + gk_.SetGlobalBuffer((__gm__ GK_T *)gk); + if (initialState != nullptr) { + initialState_.SetGlobalBuffer((__gm__ float *)initialState); + } + cuSeqlensAddr_ = reinterpret_cast<__gm__ int64_t *>(cuSeqlens); + if (preparedQG != nullptr) { + preparedQG_.SetGlobalBuffer((__gm__ T *)preparedQG); + } + if (preparedAqk != nullptr) { + preparedAqk_.SetGlobalBuffer((__gm__ T *)preparedAqk); + } + if (propagatedVNew != nullptr) { + propagatedVNew_.SetGlobalBuffer((__gm__ T *)propagatedVNew); + } + if (propagatedH != nullptr) { + propagatedH_.SetGlobalBuffer((__gm__ T *)propagatedH); + } + chunkIndicesAddr_ = reinterpret_cast<__gm__ int64_t *>(chunkIndices); + o_.SetGlobalBuffer((__gm__ OUT_T *)o); + finalState_.SetGlobalBuffer((__gm__ float *)finalState); + aqk_.SetGlobalBuffer((__gm__ float *)aqk); + akk_.SetGlobalBuffer((__gm__ AKK_T *)akk); + w_.SetGlobalBuffer((__gm__ T *)w); + u_.SetGlobalBuffer((__gm__ OUT_T *)u); + qg_.SetGlobalBuffer((__gm__ T *)qg); + kg_.SetGlobalBuffer((__gm__ T *)kg); + vNew_.SetGlobalBuffer((__gm__ T *)vNew); + h_.SetGlobalBuffer((__gm__ float *)h); + solveWorkspace_.SetGlobalBuffer((__gm__ float *)workspace); + + B_ = tiling.batch; + N_ = tiling.seqNum; + H_ = tiling.qHeadNum; + HV_ = tiling.vHeadNum; + T_ = tiling.seqlen; + K_ = tiling.kHeadDim; + V_ = tiling.vHeadDim; + BT_ = tiling.chunkSize; + NT_ = tiling.totalChunks; + scale_ = tiling.scale; + hasInitial_ = tiling.hasInitialState; + isVarLen_ = tiling.isVarLen; + inputSequenceMajor_ = tiling.inputSequenceMajor; + usedCoreNum_ = tiling.postWuUsedCoreNum; + if ASCEND_IS_AIV { + uint64_t subBlockNum = static_cast(GetSubBlockNum()); + solveCoreIdx_ = subBlockNum == 0 ? 0 : static_cast(GetBlockIdx()) / subBlockNum; + } else { + solveCoreIdx_ = static_cast(GetBlockIdx()); + } + if (pipe_ != nullptr && initVecBuffers) { + pipe_->InitBuffer(exp2Buf_, EXP2_UB_BYTES); + pipe_->InitBuffer(vecBuf_, KDA_VEC_ARENA_ELEMENTS * sizeof(float)); + const uint64_t gateWritebackRows = ScoreVectorMaxRows(5 * sizeof(float) + 2 * sizeof(T) + sizeof(GK_T)); + pipe_->InitBuffer(gateWritebackBuf_, + static_cast(gateWritebackRows * K_ * (3 * sizeof(T) + sizeof(GK_T)))); + AllocVectorEvents(); + } + } + __aicore__ inline void ProcessAiv() + { + ProcessPostAiv(); + ReleaseVectorEvents(); + } + + __aicore__ inline void ProcessAic() { ProcessPostAic(); } + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + __aicore__ inline void ProcessPreparedFullHeadPairBatchArch35(const uint64_t *batchB, const uint64_t *batchHvBase, + const uint64_t *batchStart, uint16_t taskCount) + { + if (taskCount == 0) { + return; + } + SetLoadDataPaddingValue(static_cast(0)); + Catlass::Arch::Resource resource; + const uint16_t itemCount = taskCount * KDA_POST_HEAD_PAIR_LANES; + uint16_t slot = 0; + uint16_t usedSlotCount = 1; + uint16_t taskIdx = 0; + uint16_t lane = 0; + uint64_t b = batchB[taskIdx]; + uint64_t hv = batchHvBase[taskIdx] + lane; + uint64_t start = batchStart[taskIdx]; + InitializePostWuPipelineEvents(); + PrefetchPostWuPipelineArch35(resource, slot, b, hv, start, BT_, false); + PrefetchPostWuPipelineU(resource, slot, b, hv, start, BT_, false); + + for (uint16_t item = 0; item < itemCount; ++item) { + const uint16_t nextItem = item + 1; + if (nextItem < itemCount) { + const uint16_t nextTaskIdx = nextItem / KDA_POST_HEAD_PAIR_LANES; + const uint16_t nextLane = nextItem % KDA_POST_HEAD_PAIR_LANES; + const uint16_t nextSlot = slot ^ 1; + const bool reuseSlot = nextItem >= KDA_POST_PIPELINE_STAGE_COUNT; + PrefetchPostWuPipelineArch35(resource, nextSlot, batchB[nextTaskIdx], + batchHvBase[nextTaskIdx] + nextLane, batchStart[nextTaskIdx], BT_, + reuseSlot); + PrefetchPostWuPipelineU(resource, nextSlot, batchB[nextTaskIdx], batchHvBase[nextTaskIdx] + nextLane, + batchStart[nextTaskIdx], BT_, reuseSlot); + if (!reuseSlot) { + ++usedSlotCount; + } + } + + ComputePrefetchedPostWuPipelineArch35(resource, slot, b, hv, start, BT_); + if (nextItem < itemCount) { + taskIdx = nextItem / KDA_POST_HEAD_PAIR_LANES; + lane = nextItem % KDA_POST_HEAD_PAIR_LANES; + b = batchB[taskIdx]; + hv = batchHvBase[taskIdx] + lane; + start = batchStart[taskIdx]; + slot ^= 1; + } + } + FinalizePostWuPipelineEvents(usedSlotCount); + } + + __aicore__ inline void ProcessPreparedTailHeadPairArch35(uint64_t b, uint64_t hvBase, uint64_t start, uint64_t curT) + { + SetLoadDataPaddingValue(static_cast(0)); + Catlass::Arch::Resource resource; + InitializePostWuPipelineEvents(); + for (uint16_t lane = 0; lane < KDA_POST_HEAD_PAIR_LANES; ++lane) { + PrefetchPostWuPipelineArch35(resource, lane, b, hvBase + lane, start, curT, false); + PrefetchPostWuPipelineU(resource, lane, b, hvBase + lane, start, curT, false); + } + for (uint16_t lane = 0; lane < KDA_POST_HEAD_PAIR_LANES; ++lane) { + ComputePrefetchedPostWuPipelineArch35(resource, lane, b, hvBase + lane, start, curT); + } + FinalizePostWuPipelineEvents(KDA_POST_HEAD_PAIR_LANES); + } + + __aicore__ inline void ProcessPreparedTailSingleArch35(uint64_t b, uint64_t hv, uint64_t start, uint64_t curT) + { + SetLoadDataPaddingValue(static_cast(0)); + Catlass::Arch::Resource resource; + InitializePostWuPipelineSlot(0); + PrefetchPostWuPipelineArch35(resource, 0, b, hv, start, curT, false); + PrefetchPostWuPipelineU(resource, 0, b, hv, start, curT, false); + ComputePrefetchedPostWuPipelineArch35(resource, 0, b, hv, start, curT); + FinalizePostWuPipelineEvents(1); + } + + __aicore__ inline void ProcessPreparedHeadPairBatchArch35(const uint64_t *batchB, const uint64_t *batchHvBase, + const uint64_t *batchStart, const uint64_t *batchEnd, + uint16_t taskCount) + { + uint16_t fullRunBegin = 0; + for (uint16_t task = 0; task < taskCount; ++task) { + if (batchEnd[task] - batchStart[task] == BT_) { + continue; + } + if (task > fullRunBegin) { + ProcessPreparedFullHeadPairBatchArch35(batchB + fullRunBegin, batchHvBase + fullRunBegin, + batchStart + fullRunBegin, task - fullRunBegin); + } + ProcessPreparedTailHeadPairArch35(batchB[task], batchHvBase[task], batchStart[task], + batchEnd[task] - batchStart[task]); + fullRunBegin = task + 1; + } + if (fullRunBegin < taskCount) { + ProcessPreparedFullHeadPairBatchArch35(batchB + fullRunBegin, batchHvBase + fullRunBegin, + batchStart + fullRunBegin, taskCount - fullRunBegin); + } + } +#endif + +private: + __aicore__ inline void AllocVectorEvents() + { + mte2ToVEvent_ = pipe_->AllocEventID(); + vToMte2Event_ = pipe_->AllocEventID(); + vToMte3Event_ = pipe_->AllocEventID(); + mte3ToVEvent_ = pipe_->AllocEventID(); + mte2ToMte3Event_ = pipe_->AllocEventID(); + mte3ToMte2Event_ = pipe_->AllocEventID(); + vectorEventsAllocated_ = true; + } + + __aicore__ inline void ReleaseVectorEvents() + { + if (!vectorEventsAllocated_) { + return; + } + pipe_->ReleaseEventID(mte2ToVEvent_); + pipe_->ReleaseEventID(vToMte2Event_); + pipe_->ReleaseEventID(vToMte3Event_); + pipe_->ReleaseEventID(mte3ToVEvent_); + pipe_->ReleaseEventID(mte2ToMte3Event_); + pipe_->ReleaseEventID(mte3ToMte2Event_); + vectorEventsAllocated_ = false; + } + + __aicore__ inline uint64_t QOffset(uint64_t b, uint64_t h, uint64_t t, uint64_t d) const + { + if (inputSequenceMajor_) { + return ((b * T_ + t) * H_ + h) * K_ + d; + } + return ((b * H_ + h) * T_ + t) * K_ + d; + } + + __aicore__ inline uint64_t KVOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t d, uint64_t dim) const + { + return ((b * HV_ + hv) * T_ + t) * dim + d; + } + + __aicore__ inline uint64_t BetaOffset(uint64_t b, uint64_t hv, uint64_t t) const { return (b * HV_ + hv) * T_ + t; } + + __aicore__ inline uint64_t AOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t j) const + { + return ((b * HV_ + hv) * T_ + t) * BT_ + j; + } + + __aicore__ inline uint64_t HOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t d, uint64_t r) const + { + return (((b * HV_ + hv) * NT_ + chunkIdx) * K_ + d) * V_ + r; + } + + __aicore__ inline uint64_t WScratchOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t t, uint64_t d) const + { + return (((b * HV_ + hv) * NT_ + chunkIdx) * BT_ + t) * K_ + d; + } + + __aicore__ inline uint64_t SolveScratchOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t slot) const + { + (void)b; + (void)hv; + (void)chunkIdx; + uint64_t matrixElements = BT_ * BT_; + return solveCoreIdx_ * KDA_SOLVE_SCRATCH_SLOTS * matrixElements + slot * matrixElements; + } + + __aicore__ inline uint64_t ScoreScratchOffset(uint64_t slot, uint64_t plane, uint64_t t = 0, uint64_t d = 0) const + { + return (((solveCoreIdx_ * KDA_SCORE_QUEUE_DEPTH + slot) * KDA_SCORE_SCRATCH_PLANES + plane) * BT_ + t) * K_ + d; + } + + __aicore__ inline uint64_t ScoreRefBlockSize() const { return KDA_SCORE_REF_BC; } + + __aicore__ inline uint64_t ScoreRowBlockCount(uint64_t curT, uint64_t rowBegin) const + { + uint64_t blockSize = ScoreRefBlockSize(); + uint64_t rowCount = curT - rowBegin; + if (rowCount > blockSize) { + rowCount = blockSize; + } + return rowCount; + } + + __aicore__ inline uint64_t ScoreRefToken(uint64_t start, uint64_t curT, uint64_t rowBegin, uint64_t rowCount) const + { + uint64_t ref = rowBegin + rowCount / 2; + if (ref >= curT) { + ref = curT - 1; + } + return start + ref; + } + + __aicore__ inline void RunExp2(LocalTensor &tensor, uint32_t count) + { + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + ClampExpInput(tensor, count); + Exp(tensor, tensor, count); + PipeBarrier(); + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + } + + __aicore__ inline void ClampExpInput(LocalTensor &tensor, uint32_t count) + { + Mins(tensor, tensor, KDA_EXP_INPUT_MAX, count); + PipeBarrier(); + Maxs(tensor, tensor, KDA_EXP_INPUT_MIN, count); + PipeBarrier(); + } + + __aicore__ inline void ClampFp32ToOutputType(LocalTensor &tensor, uint32_t count) + { + if constexpr (IsSameType::value) { + Mins(tensor, tensor, KDA_FP16_MAX, count); + PipeBarrier(); + Maxs(tensor, tensor, -KDA_FP16_MAX, count); + PipeBarrier(); + } + } + + template + __aicore__ inline void CopyVectorIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset, + uint64_t count) + { + uint64_t rowBytes = count * static_cast(sizeof(CopyT)); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst, src[offset], static_cast(count)); + return; + } + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPadParams padParams{false, 0, 0, 0}; + DataCopyPad(dst, src[offset], params, padParams); + } + + template + __aicore__ inline void CopyRowsIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset, uint64_t rows, + uint64_t cols, uint64_t rowStride) + { + if (rows == 0 || cols == 0) { + return; + } + if (rowStride == cols) { + CopyVectorIn(dst, src, offset, rows * cols); + return; + } + DataCopyExtParams params{static_cast(rows), static_cast(cols * sizeof(CopyT)), + static_cast((rowStride - cols) * sizeof(CopyT)), 0, 0}; + DataCopyPadExtParams padParams{false, 0, 0, 0}; + DataCopyPad(dst, src[offset], params, padParams); + } + + template + __aicore__ inline void CopyVectorOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src, + uint64_t count) + { + uint64_t rowBytes = count * static_cast(sizeof(CopyT)); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst[offset], src, static_cast(count)); + return; + } + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPad(dst[offset], src, params); + } + + template + __aicore__ inline void CopyRowIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset) + { + CopyVectorIn(dst, src, offset, K_); + } + + template + __aicore__ inline void CopyRowOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src) + { + CopyVectorOut(dst, offset, src, K_); + } + + __aicore__ inline LocalTensor VecScratch(uint64_t slot) + { + return vecBuf_.Get()[slot * EXP2_UB_ELEMENTS]; + } + + template + __aicore__ inline void LoadAsFloatRow(GlobalTensor &src, uint64_t srcOffset, LocalTensor &dst, + uint64_t count) + { + if constexpr (IsSameType::value) { + CopyVectorIn(dst, src, srcOffset, count); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Adds(dst, dst, 0.0f, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } else { + constexpr uint32_t typedOffset = EXP2_UB_ELEMENTS * sizeof(float) / sizeof(CopyT); + LocalTensor rowLocal = exp2Buf_.Get()[typedOffset]; + CopyVectorIn(rowLocal, src, srcOffset, count); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Cast(dst, rowLocal, RoundMode::CAST_NONE, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } + PipeBarrier(); + } + + template + __aicore__ inline void LoadAsFloatVector(GlobalTensor &src, uint64_t srcOffset, LocalTensor &dst, + LocalTensor &typedScratch, uint64_t count) + { + if constexpr (IsSameType::value) { + CopyVectorIn(dst, src, srcOffset, count); + } else { + CopyVectorIn(typedScratch, src, srcOffset, count); + } + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + if constexpr (!IsSameType::value) { + Cast(dst, typedScratch, RoundMode::CAST_NONE, static_cast(count)); + PipeBarrier(); + } + } + + template + __aicore__ inline void StoreFloatRow(GlobalTensor &dst, uint64_t dstOffset, LocalTensor &src, + uint64_t count) + { + if constexpr (IsSameType::value) { + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(dst, dstOffset, src, count); + } else { + constexpr uint32_t typedOffset = EXP2_UB_ELEMENTS * sizeof(float) / sizeof(CopyT); + LocalTensor rowLocal = exp2Buf_.Get()[typedOffset]; + Cast(rowLocal, src, RoundMode::CAST_RINT, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(dst, dstOffset, rowLocal, count); + } + SetFlag(mte3ToMte2Event_); + WaitFlag(mte3ToMte2Event_); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline LocalTensor Exp2NegG(uint64_t b, uint64_t hv, uint64_t t) + { + LocalTensor exp2Local = exp2Buf_.Get(); + LoadAsFloatRow(gk_, KVOffset(b, hv, t, 0, K_), exp2Local, K_); + Muls(exp2Local, exp2Local, -LN2, static_cast(K_)); + PipeBarrier(); + RunExp2(exp2Local, static_cast(K_)); + return exp2Local; + } + + __aicore__ inline uint64_t ScoreVectorMaxRows(uint64_t bytesPerElem) const + { + constexpr uint64_t arenaBytes = static_cast(KDA_VEC_ARENA_ELEMENTS) * sizeof(float); + uint64_t maxRows = (arenaBytes / bytesPerElem) / K_; + if (K_ >= 128 && maxRows > 32) { + maxRows = 32; + } + return maxRows; + } + __aicore__ inline bool UsePostWuCube(uint64_t curT) const + { + return curT > 0 && curT <= BT_ && (BT_ == 64 || BT_ == 128) && K_ >= 16 && V_ >= 16 && V_ <= 256 && + K_ % 16 == 0 && V_ % 16 == 0; + } + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + __aicore__ inline bool UsePostWuCubeArch35(uint64_t curT) const + { + return curT == 64 && BT_ == 64 && K_ == 128 && V_ == 128; + } + + __aicore__ inline bool UseFullPostWuPipelineArch35(uint64_t curT) const + { + return curT == 64 && BT_ == 64 && K_ == 128 && V_ == 128; + } + + __aicore__ inline void ComputePostWuCubeFusedArch35(uint64_t b, uint64_t hv, uint64_t start, uint64_t curT) + { + using ElementA = T; + using ElementB = T; + using ElementC = T; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::RowMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + + constexpr uint32_t capacityM = 64; + constexpr uint32_t n = 128; + constexpr uint32_t capacityK = 64; + const uint32_t m = static_cast(curT); + const uint32_t k = static_cast(curT); + SetLoadDataPaddingValue(static_cast(0)); + + auto layoutA = tla::MakeLayout(capacityM, capacityK); + auto layoutB = tla::MakeLayout(capacityK, n); + auto layoutC = tla::MakeLayout(capacityM, n); + auto tensorA = tla::MakeTensor(preparedAqk_[AOffset(b, hv, start, 0)], layoutA, Catlass::Arch::PositionGM{}); + auto tensorW = + tla::MakeTensor(preparedQG_[KVOffset(b, hv, start, 0, K_)], layoutB, Catlass::Arch::PositionGM{}); + auto tensorU = + tla::MakeTensor(propagatedVNew_[KVOffset(b, hv, start, 0, V_)], layoutB, Catlass::Arch::PositionGM{}); + auto tensorWOut = tla::MakeTensor(w_[KVOffset(b, hv, start, 0, K_)], layoutC, Catlass::Arch::PositionGM{}); + auto tensorUOut = tla::MakeTensor(u_[KVOffset(b, hv, start, 0, V_)], layoutC, Catlass::Arch::PositionGM{}); + auto blockA = GetTile(tensorA, tla::MakeCoord(0, 0), tla::MakeShape(m, k)); + auto blockW = GetTile(tensorW, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + auto blockU = GetTile(tensorU, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + auto blockWOut = GetTile(tensorWOut, tla::MakeCoord(0, 0), tla::MakeShape(m, n)); + auto blockUOut = GetTile(tensorUOut, tla::MakeCoord(0, 0), tla::MakeShape(m, n)); + + Catlass::Arch::Resource resource; + constexpr uint32_t aBytes = capacityM * capacityK * sizeof(ElementA); + constexpr uint32_t bBytes = capacityK * n * sizeof(ElementB); + LocalTensor l1A = resource.l1Buf.template GetBufferByByte(0); + LocalTensor l1B0 = resource.l1Buf.template GetBufferByByte(aBytes); + LocalTensor l1B1 = resource.l1Buf.template GetBufferByByte(aBytes + bBytes); + LocalTensor l0A = resource.l0ABuf.template GetBufferByByte(0); + LocalTensor l0B = resource.l0BBuf.template GetBufferByByte(0); + LocalTensor l0C = resource.l0CBuf.template GetBufferByByte(0); + + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + using CopyGmToL1A = typename TileCopy::template CopyGmToL1A; + using CopyGmToL1B = typename TileCopy::template CopyGmToL1B; + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; + using CopyL0CToDst = typename TileCopy::template CopyL0CToDst; + using TileMmad = Catlass::Gemm::Tile::TileMmadTla; + + auto layoutL1A = tla::MakeLayout(capacityM, capacityK); + auto layoutL1B = tla::MakeLayout(capacityK, n); + auto layoutL0A = tla::MakeLayout(capacityM, capacityK); + auto layoutL0B = tla::MakeLayout(capacityK, n); + auto layoutL0C = tla::MakeLayoutL0C(capacityM, n); + auto tensorL1A = tla::MakeTensor(l1A, layoutL1A, Catlass::Arch::PositionL1{}); + auto tensorL1B0 = tla::MakeTensor(l1B0, layoutL1B, Catlass::Arch::PositionL1{}); + auto tensorL1B1 = tla::MakeTensor(l1B1, layoutL1B, Catlass::Arch::PositionL1{}); + auto tensorL0A = tla::MakeTensor(l0A, layoutL0A, Catlass::Arch::PositionL0A{}); + auto tensorL0B = tla::MakeTensor(l0B, layoutL0B, Catlass::Arch::PositionL0B{}); + auto tensorL0C = tla::MakeTensor(l0C, layoutL0C, Catlass::Arch::PositionL0C{}); + auto tileL1A = GetTile(tensorL1A, tla::MakeCoord(0, 0), tla::MakeShape(m, k)); + auto tileL1B0 = GetTile(tensorL1B0, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + auto tileL1B1 = GetTile(tensorL1B1, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + auto tileL0A = GetTile(tensorL0A, tla::MakeCoord(0, 0), tla::MakeShape(m, k)); + auto tileL0B = GetTile(tensorL0B, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + auto tileL0C = GetTile(tensorL0C, tla::MakeCoord(0, 0), tla::MakeShape(m, n)); + + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToDst copyL0CToDst; + TileMmad tileMmad; + + copyGmToL1A(tensorL1A, blockA); + copyGmToL1B(tensorL1B0, blockW); + SetFlag(KDA_POST_EVENT); + WaitFlag(KDA_POST_EVENT); + copyL1ToL0A(tileL0A, tileL1A); + copyL1ToL0B(tileL0B, tileL1B0); + SetFlag(KDA_POST_EVENT); + copyGmToL1B(tensorL1B1, blockU); + SetFlag(KDA_POST_EVENT_NEXT); + WaitFlag(KDA_POST_EVENT); + tileMmad(tileL0C, tileL0A, tileL0B, m, n, k, true, 0); + SetFlag(KDA_POST_EVENT); + SetFlag(KDA_POST_EVENT); + WaitFlag(KDA_POST_EVENT_NEXT); + WaitFlag(KDA_POST_EVENT); + WaitFlag(KDA_POST_EVENT); + copyL0CToDst(blockWOut, tileL0C); + SetFlag(KDA_POST_EVENT_FIX); + + copyL1ToL0B(tileL0B, tileL1B1); + SetFlag(KDA_POST_EVENT_NEXT); + WaitFlag(KDA_POST_EVENT_FIX); + WaitFlag(KDA_POST_EVENT_NEXT); + tileMmad(tileL0C, tileL0A, tileL0B, m, n, k, true, 0); + SetFlag(KDA_POST_EVENT); + WaitFlag(KDA_POST_EVENT); + copyL0CToDst(blockUOut, tileL0C); + SetFlag(KDA_POST_EVENT_FIX); + WaitFlag(KDA_POST_EVENT_FIX); + } + + __aicore__ inline void FinalizePostWuPipelineEvents(uint16_t usedSlotCount) + { + for (uint16_t slot = 0; slot < usedSlotCount; ++slot) { + WaitFlag(KDA_POST_EVENT + slot); + WaitFlag(KDA_POST_PIPELINE_U_EVENT + slot); + WaitFlag(KDA_POST_EVENT + slot); + WaitFlag(KDA_POST_EVENT + slot); + } + } + + __aicore__ inline void InitializePostWuPipelineEvents() + { + for (uint16_t slot = 0; slot < KDA_POST_PIPELINE_STAGE_COUNT; ++slot) { + InitializePostWuPipelineSlot(slot); + } + } + + __aicore__ inline void InitializePostWuPipelineSlot(uint16_t slot) + { + SetFlag(KDA_POST_EVENT + slot); + SetFlag(KDA_POST_EVENT + slot); + } + + __aicore__ inline void PrefetchPostWuPipelineArch35(Catlass::Arch::Resource &resource, uint16_t slot, + uint64_t b, uint64_t hv, uint64_t start, uint64_t curT, + bool reuseSlot) + { + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::RowMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = + Catlass::Gemm::Tile::PackedTileCopyTla; + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + + constexpr uint32_t capacityM = 64; + constexpr uint32_t n = 128; + constexpr uint32_t capacityK = 64; + const uint32_t m = static_cast(curT); + const uint32_t k = static_cast(curT); + auto layoutA = tla::MakeLayout(capacityM, capacityK); + auto layoutB = tla::MakeLayout(capacityK, n); + auto tensorA = tla::MakeTensor(preparedAqk_[AOffset(b, hv, start, 0)], layoutA, Catlass::Arch::PositionGM{}); + auto tensorW = + tla::MakeTensor(preparedQG_[KVOffset(b, hv, start, 0, K_)], layoutB, Catlass::Arch::PositionGM{}); + auto blockA = GetTile(tensorA, tla::MakeCoord(0, 0), tla::MakeShape(m, k)); + auto blockW = GetTile(tensorW, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + using CopyGmToL1A = typename TileCopy::template CopyGmToL1A; + using CopyGmToL1B = typename TileCopy::template CopyGmToL1B; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + + uint32_t slotBase = static_cast(slot) * KDA_POST_PIPELINE_L1_SLOT_BYTES; + LocalTensor l1A = resource.l1Buf.template GetBufferByByte(slotBase); + LocalTensor l1W = resource.l1Buf.template GetBufferByByte(slotBase + KDA_POST_PIPELINE_L1_A_BYTES); + auto layoutL1A = tla::MakeLayout(capacityM, capacityK); + auto layoutL1B = tla::MakeLayout(capacityK, n); + auto tensorL1A = tla::MakeTensor(l1A, layoutL1A, Catlass::Arch::PositionL1{}); + auto tensorL1W = tla::MakeTensor(l1W, layoutL1B, Catlass::Arch::PositionL1{}); + + uint16_t pipelineEvent = KDA_POST_EVENT + slot; + if (reuseSlot) { + WaitFlag(pipelineEvent); + } + copyGmToL1A(tensorL1A, blockA); + copyGmToL1B(tensorL1W, blockW); + SetFlag(pipelineEvent); + } + + __aicore__ inline void PrefetchPostWuPipelineU(Catlass::Arch::Resource &resource, uint16_t slot, + uint64_t b, uint64_t hv, uint64_t start, uint64_t curT, + bool reuseStage) + { + using LayoutTagB = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + + constexpr uint32_t capacityK = 64; + constexpr uint32_t n = 128; + const uint32_t k = static_cast(curT); + auto layoutB = tla::MakeLayout(capacityK, n); + auto tensorU = + tla::MakeTensor(propagatedVNew_[KVOffset(b, hv, start, 0, V_)], layoutB, Catlass::Arch::PositionGM{}); + auto blockU = GetTile(tensorU, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + using CopyGmToL1B = typename TileCopy::template CopyGmToL1B; + CopyGmToL1B copyGmToL1B; + + uint32_t uOffset = KDA_POST_PIPELINE_STAGE_COUNT * KDA_POST_PIPELINE_L1_SLOT_BYTES + + static_cast(slot) * KDA_POST_PIPELINE_L1_U_SLOT_BYTES; + LocalTensor l1U = resource.l1Buf.template GetBufferByByte(uOffset); + auto layoutL1B = tla::MakeLayout(capacityK, n); + auto tensorL1U = tla::MakeTensor(l1U, layoutL1B, Catlass::Arch::PositionL1{}); + + uint16_t pipelineEvent = KDA_POST_PIPELINE_U_EVENT + slot; + if (reuseStage) { + WaitFlag(pipelineEvent); + } + copyGmToL1B(tensorL1U, blockU); + SetFlag(pipelineEvent); + } + + __aicore__ inline void ComputePrefetchedPostWuPipelineArch35(Catlass::Arch::Resource &resource, + uint16_t slot, uint64_t b, uint64_t hv, uint64_t start, + uint64_t curT) + { + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::RowMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = + Catlass::Gemm::Tile::PackedTileCopyTla; + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; + using TileMmad = Catlass::Gemm::Tile::TileMmadTla; + + constexpr uint32_t capacityM = 64; + constexpr uint32_t n = 128; + constexpr uint32_t packedN = 256; + constexpr uint32_t capacityK = 64; + const uint32_t m = static_cast(curT); + const uint32_t k = static_cast(curT); + auto layoutC = tla::MakeLayout(capacityM, n); + auto tensorWOut = tla::MakeTensor(w_[KVOffset(b, hv, start, 0, K_)], layoutC, Catlass::Arch::PositionGM{}); + auto tensorUOut = tla::MakeTensor(u_[KVOffset(b, hv, start, 0, V_)], layoutC, Catlass::Arch::PositionGM{}); + auto blockWOut = GetTile(tensorWOut, tla::MakeCoord(0, 0), tla::MakeShape(m, n)); + auto blockUOut = GetTile(tensorUOut, tla::MakeCoord(0, 0), tla::MakeShape(m, n)); + using CopyL0CToDst = typename TileCopy::template CopyL0CToDst; + + uint32_t l1Base = static_cast(slot) * KDA_POST_PIPELINE_L1_SLOT_BYTES; + LocalTensor l1A = resource.l1Buf.template GetBufferByByte(l1Base); + LocalTensor l1W = resource.l1Buf.template GetBufferByByte(l1Base + KDA_POST_PIPELINE_L1_A_BYTES); + uint32_t uOffset = KDA_POST_PIPELINE_STAGE_COUNT * KDA_POST_PIPELINE_L1_SLOT_BYTES + + static_cast(slot) * KDA_POST_PIPELINE_L1_U_SLOT_BYTES; + LocalTensor l1U = resource.l1Buf.template GetBufferByByte(uOffset); + LocalTensor l0A = resource.l0ABuf.template GetBufferByByte(static_cast(slot) * + KDA_POST_PIPELINE_L0_A_SLOT_BYTES); + LocalTensor l0B = resource.l0BBuf.template GetBufferByByte(static_cast(slot) * + KDA_POST_PIPELINE_L0_B_SLOT_BYTES); + LocalTensor l0C = resource.l0CBuf.template GetBufferByByte(static_cast(slot) * + KDA_POST_PIPELINE_L0_C_SLOT_BYTES); + + auto layoutL1A = tla::MakeLayout(capacityM, capacityK); + auto layoutL1B = tla::MakeLayout(capacityK, n); + auto layoutL0A = tla::MakeLayout(capacityM, capacityK); + auto layoutL0B = tla::MakeLayout(capacityK, packedN); + auto layoutL0C = tla::MakeLayoutL0C(capacityM, packedN); + auto tensorL1A = tla::MakeTensor(l1A, layoutL1A, Catlass::Arch::PositionL1{}); + auto tensorL1W = tla::MakeTensor(l1W, layoutL1B, Catlass::Arch::PositionL1{}); + auto tensorL1U = tla::MakeTensor(l1U, layoutL1B, Catlass::Arch::PositionL1{}); + auto tensorL0A = tla::MakeTensor(l0A, layoutL0A, Catlass::Arch::PositionL0A{}); + auto tensorL0B = tla::MakeTensor(l0B, layoutL0B, Catlass::Arch::PositionL0B{}); + auto tensorL0C = tla::MakeTensor(l0C, layoutL0C, Catlass::Arch::PositionL0C{}); + auto tileL1A = GetTile(tensorL1A, tla::MakeCoord(0, 0), tla::MakeShape(m, k)); + auto tileL1W = GetTile(tensorL1W, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + auto tileL1U = GetTile(tensorL1U, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + auto tileL0A = GetTile(tensorL0A, tla::MakeCoord(0, 0), tla::MakeShape(m, k)); + auto tileL0BW = GetTile(tensorL0B, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + auto tileL0BU = GetTile(tensorL0B, tla::MakeCoord(0, n), tla::MakeShape(k, n)); + auto tileL0B = GetTile(tensorL0B, tla::MakeCoord(0, 0), tla::MakeShape(k, packedN)); + auto tileL0C = GetTile(tensorL0C, tla::MakeCoord(0, 0), tla::MakeShape(m, packedN)); + auto tileL0CW = GetTile(tensorL0C, tla::MakeCoord(0, 0), tla::MakeShape(m, n)); + auto tileL0CU = GetTile(tensorL0C, tla::MakeCoord(0, n), tla::MakeShape(m, n)); + + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToDst copyL0CToDst; + TileMmad tileMmad; + + uint16_t pipelineEvent = KDA_POST_EVENT + slot; + uint16_t uPipelineEvent = KDA_POST_PIPELINE_U_EVENT + slot; + WaitFlag(pipelineEvent); + WaitFlag(uPipelineEvent); + WaitFlag(pipelineEvent); + copyL1ToL0A(tileL0A, tileL1A); + copyL1ToL0B(tileL0BW, tileL1W); + copyL1ToL0B(tileL0BU, tileL1U); + SetFlag(pipelineEvent); + SetFlag(pipelineEvent); + SetFlag(uPipelineEvent); + WaitFlag(pipelineEvent); + WaitFlag(pipelineEvent); + tileMmad(tileL0C, tileL0A, tileL0B, m, packedN, k, true, 0); + SetFlag(pipelineEvent); + SetFlag(pipelineEvent); + WaitFlag(pipelineEvent); + copyL0CToDst(blockWOut, tileL0CW); + copyL0CToDst(blockUOut, tileL0CU); + SetFlag(pipelineEvent); + } +#endif + + __aicore__ inline void ComputePostWuCube(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t curT) + { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if (UsePostWuCubeArch35(curT)) { + ComputePostWuCubeFusedArch35(b, hv, start, curT); + return; + } +#endif + using ElementA = AKK_T; + using ElementB = T; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::RowMajor; + using LayoutTagC = Catlass::layout::RowMajor; +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + SetLoadDataPaddingValue(static_cast(0)); + using WTileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; +#else + using WTileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; +#endif + using UTileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using PostL1TileShape128 = tla::Shape; + using PostL0TileShape128 = tla::Shape; + using PostL1TileShape256 = tla::Shape; + using PostL0TileShape256 = tla::Shape; +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + using WBlockMmad = Catlass::Gemm::Block::BlockMmadTla; +#else + using WBlockMmad = Catlass::Gemm::Block::BlockMmadTla; +#endif + using UBlockMmad128 = + Catlass::Gemm::Block::BlockMmadTla; + using UBlockMmad256 = + Catlass::Gemm::Block::BlockMmadTla; + LayoutTagA tagA = LayoutTagA::template MakeLayout(BT_, BT_); + auto layoutA = tla::MakeLayoutFromTag(tagA); + auto tensorA = tla::MakeTensor(preparedAqk_[AOffset(b, hv, start, 0)], layoutA, Catlass::Arch::PositionGM{}); + + { + LayoutTagB tagB = LayoutTagB::template MakeLayout(BT_, K_); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + LayoutTagC tagC = LayoutTagC::template MakeLayout(BT_, K_); +#else + LayoutTagC tagC = LayoutTagC::template MakeLayout(BT_, K_); +#endif + auto layoutB = tla::MakeLayoutFromTag(tagB); + auto layoutC = tla::MakeLayoutFromTag(tagC); + Catlass::GemmCoord shape{static_cast(curT), static_cast(K_), + static_cast(curT)}; + auto tensorB = + tla::MakeTensor(preparedQG_[KVOffset(b, hv, start, 0, K_)], layoutB, Catlass::Arch::PositionGM{}); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + auto tensorC = tla::MakeTensor(w_[KVOffset(b, hv, start, 0, K_)], layoutC, Catlass::Arch::PositionGM{}); +#else + auto tensorC = + tla::MakeTensor(h_[WScratchOffset(b, hv, chunkIdx, 0, 0)], layoutC, Catlass::Arch::PositionGM{}); +#endif + auto blockA = GetTile(tensorA, tla::MakeCoord(0, 0), tla::MakeShape(shape.m(), shape.k())); + auto blockB = GetTile(tensorB, tla::MakeCoord(0, 0), tla::MakeShape(shape.k(), shape.n())); + auto blockC = GetTile(tensorC, tla::MakeCoord(0, 0), tla::MakeShape(shape.m(), shape.n())); + Catlass::Arch::Resource wResource; + WBlockMmad wBlockMmad(wResource); + wBlockMmad(blockA, blockB, blockC, shape); + PipeBarrier(); + } + + { + LayoutTagB tagB = LayoutTagB::template MakeLayout(BT_, V_); + LayoutTagC tagC = LayoutTagC::template MakeLayout(BT_, V_); + auto layoutB = tla::MakeLayoutFromTag(tagB); + auto layoutC = tla::MakeLayoutFromTag(tagC); + Catlass::GemmCoord shape{static_cast(curT), static_cast(V_), + static_cast(curT)}; + auto tensorB = + tla::MakeTensor(propagatedVNew_[KVOffset(b, hv, start, 0, V_)], layoutB, Catlass::Arch::PositionGM{}); + auto tensorC = tla::MakeTensor(u_[KVOffset(b, hv, start, 0, V_)], layoutC, Catlass::Arch::PositionGM{}); + auto blockA = GetTile(tensorA, tla::MakeCoord(0, 0), tla::MakeShape(shape.m(), shape.k())); + auto blockB = GetTile(tensorB, tla::MakeCoord(0, 0), tla::MakeShape(shape.k(), shape.n())); + auto blockC = GetTile(tensorC, tla::MakeCoord(0, 0), tla::MakeShape(shape.m(), shape.n())); + Catlass::Arch::Resource uResource; + if (V_ <= 128) { + UBlockMmad128 uBlockMmad(uResource); + uBlockMmad(blockA, blockB, blockC, shape); + } else { + UBlockMmad256 uBlockMmad(uResource); + uBlockMmad(blockA, blockB, blockC, shape); + } + PipeBarrier(); + } + } + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + __aicore__ inline bool UseTypicalPostWuGate(uint64_t curT) const + { + return curT == 64 && BT_ == 64 && K_ == 128 && V_ == 128; + } + + __aicore__ inline uint64_t TypicalGateStageElems() const + { + return static_cast(KDA_TYPICAL_GATE_TILE_ROWS) * 128; + } + + __aicore__ inline uint64_t TypicalGateStageBytes() const + { + return TypicalGateStageElems() * (sizeof(T) + sizeof(GK_T)); + } + + __aicore__ inline LocalTensor TypicalGateK(uint64_t slot) + { + return gateWritebackBuf_.Get()[slot * TypicalGateStageBytes() / sizeof(T)]; + } + + __aicore__ inline LocalTensor TypicalGateG(uint64_t slot) + { + uint64_t byteOffset = slot * TypicalGateStageBytes() + TypicalGateStageElems() * sizeof(T); + return gateWritebackBuf_.Get()[byteOffset / sizeof(GK_T)]; + } + + __aicore__ inline void PrefetchTypicalKg(uint64_t slot, uint64_t b, uint64_t h, uint64_t hv, uint64_t token, + uint64_t rows) + { + uint64_t elems = rows * K_; + LocalTensor kStage = TypicalGateK(slot); + LocalTensor gateStage = TypicalGateG(slot); + CopyRowsIn(kStage, k_, QOffset(b, h, token, 0), rows, K_, inputSequenceMajor_ ? H_ * K_ : K_); + DataCopy(gateStage, gk_[KVOffset(b, hv, token, 0, K_)], static_cast(elems)); + SetFlag(mte2ToVEvent_); + } + + __aicore__ inline void ComputeTypicalKg(uint64_t b, uint64_t h, uint64_t hv, uint64_t start, uint64_t curT, + uint64_t subBlockIdx, uint64_t subBlockNum) + { + uint64_t rowBegin = (curT * subBlockIdx) / subBlockNum; + uint64_t rowEnd = (curT * (subBlockIdx + 1)) / subBlockNum; + if (rowBegin >= rowEnd) { + return; + } + + LocalTensor gateLast = exp2Buf_.Get(); + LoadAsFloatRow(gk_, KVOffset(b, hv, start + curT - 1, 0, K_), gateLast, K_); + + uint64_t slot = 0; + uint64_t firstRows = rowEnd - rowBegin; + if (firstRows > KDA_TYPICAL_GATE_TILE_ROWS) { + firstRows = KDA_TYPICAL_GATE_TILE_ROWS; + } + PrefetchTypicalKg(slot, b, h, hv, start + rowBegin, firstRows); + + bool outputPending = false; + for (uint64_t tileRow = rowBegin; tileRow < rowEnd; tileRow += KDA_TYPICAL_GATE_TILE_ROWS) { + uint64_t tileRows = rowEnd - tileRow; + if (tileRows > KDA_TYPICAL_GATE_TILE_ROWS) { + tileRows = KDA_TYPICAL_GATE_TILE_ROWS; + } + uint64_t elems = tileRows * K_; + WaitFlag(mte2ToVEvent_); + + if (outputPending) { + WaitFlag(mte3ToMte2Event_); + } + uint64_t nextTileRow = tileRow + KDA_TYPICAL_GATE_TILE_ROWS; + if (nextTileRow < rowEnd) { + uint64_t nextRows = rowEnd - nextTileRow; + if (nextRows > KDA_TYPICAL_GATE_TILE_ROWS) { + nextRows = KDA_TYPICAL_GATE_TILE_ROWS; + } + PrefetchTypicalKg(slot ^ 1, b, h, hv, start + nextTileRow, nextRows); + } + + LocalTensor kAndKg = TypicalGateK(slot); + LocalTensor gateStage = TypicalGateG(slot); + ComputePostKdaKgRegbase((__ubuf__ T *)reinterpret_cast(kAndKg.GetPhyAddr()), + (__ubuf__ GK_T *)reinterpret_cast(gateStage.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(gateLast.GetPhyAddr()), + static_cast(tileRows), static_cast(K_)); + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(kg_[KVOffset(b, hv, start + tileRow, 0, K_)], kAndKg, static_cast(elems)); + SetFlag(mte3ToMte2Event_); + outputPending = true; + slot ^= 1; + } + if (outputPending) { + WaitFlag(mte3ToMte2Event_); + } + } + + __aicore__ inline uint64_t TypicalGatePipelineStageElems() const + { + return static_cast(KDA_TYPICAL_GATE_PIPELINE_ROWS) * 128; + } + + __aicore__ inline uint64_t TypicalGatePipelineStageBytes() const + { + return TypicalGatePipelineStageElems() * (sizeof(T) + sizeof(float)) + 128 * sizeof(float); + } + + __aicore__ inline LocalTensor TypicalGatePipelineK(uint64_t slot) + { + uint64_t byteOffset = slot * TypicalGatePipelineStageBytes(); + return vecBuf_.Get()[byteOffset / sizeof(T)]; + } + + __aicore__ inline LocalTensor TypicalGatePipelineG(uint64_t slot) + { + uint64_t byteOffset = slot * TypicalGatePipelineStageBytes() + TypicalGatePipelineStageElems() * sizeof(T); + return vecBuf_.Get()[byteOffset / sizeof(float)]; + } + + __aicore__ inline LocalTensor TypicalGatePipelineRef(uint64_t slot) + { + uint64_t byteOffset = + slot * TypicalGatePipelineStageBytes() + TypicalGatePipelineStageElems() * (sizeof(T) + sizeof(float)); + return vecBuf_.Get()[byteOffset / sizeof(float)]; + } + + __aicore__ inline bool CanPipelineTypicalKg(uint64_t curT, uint64_t subBlockIdx, uint64_t subBlockNum) const + { + if constexpr (!IsSameType::value) { + return false; + } + if (!UseTypicalPostWuGate(curT) || subBlockNum == 0) { + return false; + } + uint64_t rowBegin = (curT * subBlockIdx) / subBlockNum; + uint64_t rowEnd = (curT * (subBlockIdx + 1)) / subBlockNum; + return rowBegin < rowEnd && rowEnd - rowBegin <= KDA_TYPICAL_GATE_PIPELINE_ROWS; + } + + __aicore__ inline void PrefetchTypicalKgPipeline(uint64_t slot, uint64_t b, uint64_t h, uint64_t hv, uint64_t start, + uint64_t curT, uint64_t rowBegin, uint64_t rowEnd) + { + if constexpr (IsSameType::value) { + uint64_t elems = (rowEnd - rowBegin) * K_; + LocalTensor kStage = TypicalGatePipelineK(slot); + LocalTensor gateStage = TypicalGatePipelineG(slot); + LocalTensor refStage = TypicalGatePipelineRef(slot); + CopyRowsIn(kStage, k_, QOffset(b, h, start + rowBegin, 0), rowEnd - rowBegin, K_, + inputSequenceMajor_ ? H_ * K_ : K_); + DataCopy(gateStage, gk_[KVOffset(b, hv, start + rowBegin, 0, K_)], static_cast(elems)); + DataCopy(refStage, gk_[KVOffset(b, hv, start + curT - 1, 0, K_)], static_cast(K_)); + SetFlag(mte2ToVEvent_); + } + } + + __aicore__ inline void ComputeTypicalKgPipelineRegs(uint64_t slot, uint64_t rowBegin, uint64_t rowEnd) + { + uint64_t rows = rowEnd - rowBegin; + LocalTensor kAndKg = TypicalGatePipelineK(slot); + LocalTensor gateStage = TypicalGatePipelineG(slot); + LocalTensor refStage = TypicalGatePipelineRef(slot); + ComputePostKdaKgRegbase((__ubuf__ T *)reinterpret_cast(kAndKg.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(gateStage.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(refStage.GetPhyAddr()), + static_cast(rows), static_cast(K_)); + } + + __aicore__ inline void StoreTypicalKgPipeline(uint64_t slot, uint64_t b, uint64_t hv, uint64_t start, + uint64_t rowBegin, uint64_t rowEnd) + { + uint64_t elems = (rowEnd - rowBegin) * K_; + LocalTensor kAndKg = TypicalGatePipelineK(slot); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(kg_[KVOffset(b, hv, start + rowBegin, 0, K_)], kAndKg, static_cast(elems)); + SetFlag(mte3ToMte2Event_); + } +#endif + + __aicore__ inline void CopyScratchWAndFinalizeKg(uint64_t b, uint64_t h, uint64_t hv, uint64_t chunkIdx, + uint64_t start, uint64_t curT, uint64_t subBlockIdx, + uint64_t subBlockNum) + { + constexpr uint64_t typedOffsetFloats = 20480; + constexpr uint64_t typedOffset = typedOffsetFloats * sizeof(float) / sizeof(T); + constexpr uint64_t kgFp32Planes = 4; + uint64_t rowBegin = (curT * subBlockIdx) / subBlockNum; + uint64_t rowEnd = (curT * (subBlockIdx + 1)) / subBlockNum; + if (rowBegin >= rowEnd) { + return; + } + uint64_t maxRows = (typedOffsetFloats / kgFp32Planes) / K_; + if (maxRows > 32) { + maxRows = 32; + } + if (maxRows == 0) { + return; + } + + uint64_t last = start + curT - 1; + LocalTensor arena = vecBuf_.Get(); + LocalTensor gateLast = exp2Buf_.Get(); + LocalTensor typedLocal = vecBuf_.Get()[typedOffset]; + LoadAsFloatRow(gk_, KVOffset(b, hv, last, 0, K_), gateLast, K_); + + for (uint64_t tileRow = rowBegin; tileRow < rowEnd; tileRow += maxRows) { + uint64_t tileRows = rowEnd - tileRow; + if (tileRows > maxRows) { + tileRows = maxRows; + } + uint64_t elemCount = tileRows * K_; +#if !defined(__CCE_AICORE__) || __CCE_AICORE__ != 310 + uint64_t scratchBase = WScratchOffset(b, hv, chunkIdx, tileRow, 0); +#else + (void)chunkIdx; +#endif + uint64_t token = start + tileRow; + +#if !defined(__CCE_AICORE__) || __CCE_AICORE__ != 310 + DataCopy(arena, h_[scratchBase], static_cast(elemCount)); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Cast(typedLocal, arena, RoundMode::CAST_RINT, static_cast(elemCount)); + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(w_[KVOffset(b, hv, token, 0, K_)], typedLocal, static_cast(elemCount)); + SetFlag(mte3ToMte2Event_); + WaitFlag(mte3ToMte2Event_); +#endif + + LocalTensor kLocal = arena; + LocalTensor gLocal = arena[elemCount]; + LocalTensor expLocal = arena[2 * elemCount]; + LocalTensor outLocal = arena[3 * elemCount]; + const uint64_t gateOffsetBytes = (typedOffset + elemCount) * sizeof(T); + LocalTensor gateTyped = vecBuf_.Get()[(gateOffsetBytes + sizeof(GK_T) - 1) / sizeof(GK_T)]; + CopyRowsIn(typedLocal, k_, QOffset(b, h, token, 0), tileRows, K_, inputSequenceMajor_ ? H_ * K_ : K_); + LoadAsFloatVector(gk_, KVOffset(b, hv, token, 0, K_), gLocal, gateTyped, elemCount); + Cast(kLocal, typedLocal, RoundMode::CAST_NONE, static_cast(elemCount)); + PipeBarrier(); + + for (uint64_t row = 0; row < tileRows; ++row) { + Sub(expLocal[row * K_], gateLast, gLocal[row * K_], static_cast(K_)); + } + PipeBarrier(); + Muls(expLocal, expLocal, LN2, static_cast(elemCount)); + PipeBarrier(); + ClampExpInput(expLocal, static_cast(elemCount)); + Exp(expLocal, expLocal, static_cast(elemCount)); + PipeBarrier(); + Mul(outLocal, kLocal, expLocal, static_cast(elemCount)); + PipeBarrier(); + ClampFp32ToOutputType(outLocal, static_cast(elemCount)); + Cast(typedLocal, outLocal, RoundMode::CAST_RINT, static_cast(elemCount)); + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(kg_, KVOffset(b, hv, token, 0, K_), typedLocal, elemCount); + SetFlag(mte3ToMte2Event_); + WaitFlag(mte3ToMte2Event_); + } + if (rowEnd == curT) { + CopyVectorIn(typedLocal, k_, QOffset(b, h, last, 0), K_); + SetFlag(mte2ToMte3Event_); + WaitFlag(mte2ToMte3Event_); + CopyVectorOut(kg_, KVOffset(b, hv, last, 0, K_), typedLocal, K_); + SetFlag(mte3ToMte2Event_); + WaitFlag(mte3ToMte2Event_); + } + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + template + __aicore__ inline void ComputeTailWuRow(GlobalTensor &src, GlobalTensor &dst, + uint64_t akkBase, uint64_t srcBase, uint64_t dstBase, uint64_t curT, + uint64_t dim) + { + LocalTensor acc = vecBuf_.Get(); + LocalTensor value = vecBuf_.Get()[512]; + LocalTensor typed = vecBuf_.Get()[4096]; + LocalTensor coefficientTyped = exp2Buf_.Get(); + LocalTensor coefficients = exp2Buf_.Get()[128]; + CopyVectorIn(coefficientTyped, preparedAqk_, akkBase, curT); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Cast(coefficients, coefficientTyped, RoundMode::CAST_NONE, static_cast(curT)); + PipeBarrier(); + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + for (uint64_t j = 0; j < curT; ++j) { + LoadAsFloatVector(src, srcBase + j * dim, value, typed, dim); + float coefficient = coefficients.GetValue(j); + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + Muls(value, value, coefficient, static_cast(dim)); + PipeBarrier(); + if (j == 0) { + Adds(acc, value, 0.0f, static_cast(dim)); + } else { + Add(acc, acc, value, static_cast(dim)); + } + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + ClampFp32ToOutputType(acc, static_cast(dim)); + StoreFloatRow(dst, dstBase, acc, dim); + } + + __aicore__ inline void ComputeTailWuVector(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t curT, uint64_t subBlockIdx, uint64_t subBlockNum) + { + uint64_t rowBegin = (curT * subBlockIdx) / subBlockNum; + uint64_t rowEnd = (curT * (subBlockIdx + 1)) / subBlockNum; + for (uint64_t row = rowBegin; row < rowEnd; ++row) { + ComputeTailWuRow(preparedQG_, w_, AOffset(b, hv, start + row, 0), KVOffset(b, hv, start, 0, K_), + KVOffset(b, hv, start + row, 0, K_), curT, K_); + ComputeTailWuRow(propagatedVNew_, u_, AOffset(b, hv, start + row, 0), KVOffset(b, hv, start, 0, V_), + KVOffset(b, hv, start + row, 0, V_), curT, V_); + } + } + + __aicore__ inline void ComputeTailWuRegbaseArch35(uint64_t b, uint64_t hv, uint64_t start, uint64_t curT, + uint64_t subBlockIdx, uint64_t subBlockNum) + { + // W aliases W-seed in GM: one AIV must stage every seed row before either output is overwritten. + constexpr uint64_t dim = 128; + constexpr uint64_t maxTailRows = KDA_POST_REGBASE_TAIL_LIMIT - 1; + constexpr uint64_t tileRowsLimit = 32; + constexpr uint64_t akkStagedCols = 64; + constexpr uint64_t wSeedOffset = 0; + constexpr uint64_t uSeedOffset = wSeedOffset + maxTailRows * dim; + constexpr uint64_t akkOffset = uSeedOffset + maxTailRows * dim; + (void)subBlockNum; + if (subBlockIdx != 0 || curT == 0 || curT > maxTailRows) { + return; + } + + LocalTensor staging = gateWritebackBuf_.Get(); + LocalTensor wSeedLocal = staging[wSeedOffset]; + LocalTensor uSeedLocal = staging[uSeedOffset]; + LocalTensor akkLocal = staging[akkOffset]; + CopyVectorIn(wSeedLocal, preparedQG_, KVOffset(b, hv, start, 0, K_), curT * dim); + CopyVectorIn(uSeedLocal, propagatedVNew_, KVOffset(b, hv, start, 0, V_), curT * dim); + for (uint64_t row = 0; row < curT; ++row) { + LocalTensor akkRow = akkLocal[row * akkStagedCols]; + CopyVectorIn(akkRow, preparedAqk_, AOffset(b, hv, start + row, 0), akkStagedCols); + } + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + + for (uint64_t tileRow = 0; tileRow < curT; tileRow += tileRowsLimit) { + const uint64_t tileRows = (curT - tileRow) > tileRowsLimit ? tileRowsLimit : (curT - tileRow); + LocalTensor wFloat = vecBuf_.Get(); + LocalTensor uFloat = wFloat[tileRows * dim]; + LocalTensor akkTile = akkLocal[tileRow * akkStagedCols]; + AscendC::VF_CALL(reinterpret_cast<__ubuf__ float *>(wFloat.GetPhyAddr()), + reinterpret_cast<__ubuf__ float *>(uFloat.GetPhyAddr()), + reinterpret_cast<__ubuf__ bfloat16_t *>(akkTile.GetPhyAddr()), + reinterpret_cast<__ubuf__ bfloat16_t *>(wSeedLocal.GetPhyAddr()), + reinterpret_cast<__ubuf__ bfloat16_t *>(uSeedLocal.GetPhyAddr()), + static_cast(tileRows), static_cast(curT)); + PipeBarrier(); + + for (uint64_t row = 0; row < tileRows; ++row) { + LocalTensor wRow = wFloat[row * dim]; + LocalTensor uRow = uFloat[row * dim]; + StoreFloatRow(w_, KVOffset(b, hv, start + tileRow + row, 0, K_), wRow, dim); + StoreFloatRow(u_, KVOffset(b, hv, start + tileRow + row, 0, V_), uRow, dim); + } + } + } +#endif + __aicore__ inline bool ResolveFlatChunk(uint64_t task, uint64_t &seq, uint64_t &b, uint64_t &h, uint64_t &hv, + uint64_t &chunkIdx, uint64_t &start, uint64_t &end) + { + hv = task % HV_; + uint64_t flatChunk = task / HV_; + if (!isVarLen_) { + seq = flatChunk / NT_; + b = seq; + chunkIdx = flatChunk % NT_; + start = chunkIdx * BT_; + end = start + BT_; + if (end > T_) { + end = T_; + } + } else { + if (!KdaVarlen::ResolveChunkRange(cuSeqlensAddr_, chunkIndicesAddr_, N_, T_, BT_, flatChunk, seq, start, + end)) { + return false; + } + b = 0; + chunkIdx = flatChunk; + } + h = hv / (HV_ / H_); + return start < end; + } + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + __aicore__ inline bool ResolveHeadMajorChunk(uint64_t task, uint64_t &seq, uint64_t &b, uint64_t &h, uint64_t &hv, + uint64_t &chunkIdx, uint64_t &start, uint64_t &end) + { + uint64_t flatTask = 0; + if (isVarLen_) { + hv = task / NT_; + uint64_t flatChunk = task % NT_; + flatTask = flatChunk * HV_ + hv; + } else { + uint64_t entity = task / NT_; + uint64_t localChunk = task % NT_; + b = entity / HV_; + hv = entity % HV_; + uint64_t flatChunk = b * NT_ + localChunk; + flatTask = flatChunk * HV_ + hv; + } + return ResolveFlatChunk(flatTask, seq, b, h, hv, chunkIdx, start, end); + } + + __aicore__ inline void GetHeadMajorTaskRange(uint64_t coreIdx, uint64_t coreNum, uint64_t taskNum, + uint64_t &taskBegin, uint64_t &taskEnd) const + { + uint64_t tasksPerCore = (taskNum + coreNum - 1) / coreNum; + taskBegin = coreIdx * tasksPerCore; + taskEnd = taskBegin + tasksPerCore; + if (taskBegin > taskNum) { + taskBegin = taskNum; + } + if (taskEnd > taskNum) { + taskEnd = taskNum; + } + } + + __aicore__ inline void ProcessPostAivPipelineArch35(uint64_t taskBegin, uint64_t taskEnd, uint64_t subBlockIdx, + uint64_t subBlockNum) + { + uint64_t task = taskBegin; + while (task < taskEnd) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + bool resolved = ResolveHeadMajorChunk(task, seq, b, h, hv, chunkIdx, start, end); + uint64_t curT = resolved ? end - start : 0; + if (!resolved || !CanPipelineTypicalKg(curT, subBlockIdx, subBlockNum)) { + if (resolved) { + (void)seq; + ProcessChunkPostAiv(b, h, hv, chunkIdx, start, end, subBlockIdx, subBlockNum); + } + ++task; + continue; + } + + uint64_t rowBegin = (curT * subBlockIdx) / subBlockNum; + uint64_t rowEnd = (curT * (subBlockIdx + 1)) / subBlockNum; + uint16_t slot = 0; + bool outputPending = false; + PrefetchTypicalKgPipeline(slot, b, h, hv, start, curT, rowBegin, rowEnd); + + while (true) { + WaitFlag(mte2ToVEvent_); + + uint64_t nextTask = task + 1; + uint64_t nextSeq = 0; + uint64_t nextB = 0; + uint64_t nextH = 0; + uint64_t nextHv = 0; + uint64_t nextChunkIdx = 0; + uint64_t nextStart = 0; + uint64_t nextEnd = 0; + bool nextResolved = nextTask < taskEnd && ResolveHeadMajorChunk(nextTask, nextSeq, nextB, nextH, nextHv, + nextChunkIdx, nextStart, nextEnd); + uint64_t nextCurT = nextResolved ? nextEnd - nextStart : 0; + bool nextIsTypical = nextResolved && CanPipelineTypicalKg(nextCurT, subBlockIdx, subBlockNum); + uint64_t nextRowBegin = 0; + uint64_t nextRowEnd = 0; + if (nextIsTypical) { + nextRowBegin = (nextCurT * subBlockIdx) / subBlockNum; + nextRowEnd = (nextCurT * (subBlockIdx + 1)) / subBlockNum; + uint16_t nextSlot = (slot + 1) % KDA_TYPICAL_GATE_PIPELINE_STAGES; + PrefetchTypicalKgPipeline(nextSlot, nextB, nextH, nextHv, nextStart, nextCurT, nextRowBegin, + nextRowEnd); + } + + ComputeTypicalKgPipelineRegs(slot, rowBegin, rowEnd); + if (outputPending) { + WaitFlag(mte3ToMte2Event_); + outputPending = false; + } + StoreTypicalKgPipeline(slot, b, hv, start, rowBegin, rowEnd); + outputPending = true; + + if (!nextIsTypical) { + WaitFlag(mte3ToMte2Event_); + task = nextTask; + break; + } + + task = nextTask; + seq = nextSeq; + b = nextB; + h = nextH; + hv = nextHv; + chunkIdx = nextChunkIdx; + start = nextStart; + end = nextEnd; + curT = nextCurT; + rowBegin = nextRowBegin; + rowEnd = nextRowEnd; + slot = (slot + 1) % KDA_TYPICAL_GATE_PIPELINE_STAGES; + (void)seq; + (void)chunkIdx; + (void)end; + (void)curT; + } + } + } +#endif + + __aicore__ inline void ProcessChunkPostAiv(uint64_t b, uint64_t h, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t end, uint64_t subBlockIdx, uint64_t subBlockNum) + { + uint64_t curT = end - start; + if (curT == 0 || !UsePostWuCube(curT)) { + return; + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if (curT < BT_) { + if constexpr (IsSameType::value) { + if (BT_ == 64 && K_ == 128 && V_ == 128 && curT < KDA_POST_REGBASE_TAIL_LIMIT) { + ComputeTailWuRegbaseArch35(b, hv, start, curT, subBlockIdx, subBlockNum); + } else { + ComputeTailWuVector(b, hv, chunkIdx, start, curT, subBlockIdx, subBlockNum); + } + } else { + ComputeTailWuVector(b, hv, chunkIdx, start, curT, subBlockIdx, subBlockNum); + } + CopyScratchWAndFinalizeKg(b, h, hv, chunkIdx, start, curT, subBlockIdx, subBlockNum); + return; + } + if (UseTypicalPostWuGate(curT)) { + ComputeTypicalKg(b, h, hv, start, curT, subBlockIdx, subBlockNum); + return; + } else { + CopyScratchWAndFinalizeKg(b, h, hv, chunkIdx, start, curT, subBlockIdx, subBlockNum); + } +#else + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE2>(syncDoneFlag_); + CopyScratchWAndFinalizeKg(b, h, hv, chunkIdx, start, curT, subBlockIdx, subBlockNum); +#endif + } + + __aicore__ inline void ProcessChunkPostAic(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t end) + { + if constexpr (IsSameType::value) { + ProcessChunkPostAicTyped(b, hv, chunkIdx, start, end); + } + } + + __aicore__ inline void ProcessChunkPostAicTyped(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t end) + { + uint64_t curT = end - start; + if (curT == 0 || !UsePostWuCube(curT)) { + return; + } + ComputePostWuCube(b, hv, chunkIdx, start, curT); +#if !defined(__CCE_AICORE__) || __CCE_AICORE__ != 310 + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(syncDoneFlag_); +#endif + } + + __aicore__ inline void ProcessPostAiv() + { + if constexpr (IsSameType::value) { + return; + } + uint64_t subBlockNum = static_cast(GetSubBlockNum()); + if (subBlockNum == 0) { + return; + } + uint64_t subBlockIdx = static_cast(GetSubBlockIdx()); + uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + uint64_t coreIdx = static_cast(GetBlockIdx()) / subBlockNum; + uint64_t taskNum = static_cast((isVarLen_ ? NT_ : B_ * NT_) * HV_); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if (BT_ == 64 && K_ == 128 && V_ == 128) { + uint64_t taskBegin = 0; + uint64_t taskEnd = 0; + GetHeadMajorTaskRange(coreIdx, coreNum, taskNum, taskBegin, taskEnd); + ProcessPostAivPipelineArch35(taskBegin, taskEnd, subBlockIdx, subBlockNum); + return; + } +#endif + for (uint64_t task = coreIdx; task < taskNum; task += coreNum) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveFlatChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + (void)seq; + ProcessChunkPostAiv(b, h, hv, chunkIdx, start, end, subBlockIdx, subBlockNum); + } + } + } + + __aicore__ inline void ProcessPostAic() + { + if constexpr (IsSameType::value) { + return; + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if (KDA_ENABLE_POST_AIC_PIPELINE && BT_ == 64 && K_ == 128 && V_ == 128) { + ProcessPostAicPipelineArch35(); + return; + } +#endif + uint64_t taskNum = static_cast((isVarLen_ ? NT_ : B_ * NT_) * HV_); + uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + for (uint64_t task = GetBlockIdx(); task < taskNum; task += coreNum) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveFlatChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + (void)seq; + (void)h; + ProcessChunkPostAic(b, hv, chunkIdx, start, end); + } + } + } + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + __aicore__ inline void ProcessPostAicPipelineArch35() + { + static_assert(sizeof(T) == sizeof(uint16_t), "arch35 PostWU pipeline is specialized for fp16/bf16 inputs"); + SetLoadDataPaddingValue(static_cast(0)); + uint64_t taskNum = static_cast((isVarLen_ ? NT_ : B_ * NT_) * HV_); + uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + uint64_t coreIdx = static_cast(GetBlockIdx()); + uint64_t taskBegin = 0; + uint64_t taskEnd = 0; + GetHeadMajorTaskRange(coreIdx, coreNum, taskNum, taskBegin, taskEnd); + if (taskEnd - taskBegin < KDA_POST_PIPELINE_STAGE_COUNT) { + for (uint64_t task = taskBegin; task < taskEnd; ++task) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveHeadMajorChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + (void)seq; + (void)h; + if (end - start == BT_) { + ProcessPreparedTailSingleArch35(b, hv, start, BT_); + } + } + } + return; + } + uint64_t task = taskBegin; + Catlass::Arch::Resource resource; + + while (task < taskEnd) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + bool resolved = ResolveHeadMajorChunk(task, seq, b, h, hv, chunkIdx, start, end); + uint64_t curT = resolved ? end - start : 0; + if (!resolved || !UseFullPostWuPipelineArch35(curT)) { + (void)seq; + (void)h; + ++task; + continue; + } + + uint16_t slot = 0; + uint16_t usedSlotCount = 1; + InitializePostWuPipelineSlot(slot); + PrefetchPostWuPipelineArch35(resource, slot, b, hv, start, curT, false); + PrefetchPostWuPipelineU(resource, slot, b, hv, start, curT, false); + while (true) { + uint64_t nextTask = task + 1; + uint64_t nextSeq = 0; + uint64_t nextB = 0; + uint64_t nextH = 0; + uint64_t nextHv = 0; + uint64_t nextChunkIdx = 0; + uint64_t nextStart = 0; + uint64_t nextEnd = 0; + bool nextResolved = nextTask < taskEnd && ResolveHeadMajorChunk(nextTask, nextSeq, nextB, nextH, nextHv, + nextChunkIdx, nextStart, nextEnd); + uint64_t nextCurT = nextResolved ? nextEnd - nextStart : 0; + bool nextIsTypical = nextResolved && UseFullPostWuPipelineArch35(nextCurT); + if (nextIsTypical) { + uint16_t nextSlot = slot ^ 1; + bool reuseSlot = usedSlotCount == KDA_POST_PIPELINE_STAGE_COUNT; + if (!reuseSlot) { + InitializePostWuPipelineSlot(nextSlot); + } + PrefetchPostWuPipelineArch35(resource, nextSlot, nextB, nextHv, nextStart, nextCurT, reuseSlot); + PrefetchPostWuPipelineU(resource, nextSlot, nextB, nextHv, nextStart, nextCurT, reuseSlot); + if (!reuseSlot) { + ++usedSlotCount; + } + } + + ComputePrefetchedPostWuPipelineArch35(resource, slot, b, hv, start, curT); + if (!nextIsTypical) { + FinalizePostWuPipelineEvents(usedSlotCount); + task = nextTask; + break; + } + task = nextTask; + seq = nextSeq; + b = nextB; + h = nextH; + hv = nextHv; + chunkIdx = nextChunkIdx; + start = nextStart; + end = nextEnd; + curT = nextCurT; + slot ^= 1; + (void)seq; + (void)h; + (void)chunkIdx; + (void)end; + (void)curT; + } + } + } +#endif + +private: + GlobalTensor q_; + GlobalTensor k_; + GlobalTensor v_; + GlobalTensor gk_; + GlobalTensor initialState_; + GlobalTensor o_; + GlobalTensor finalState_; + GlobalTensor aqk_; + GlobalTensor akk_; + GlobalTensor w_; + GlobalTensor u_; + GlobalTensor qg_; + GlobalTensor kg_; + GlobalTensor vNew_; + GlobalTensor h_; + GlobalTensor preparedQG_; + GlobalTensor preparedAqk_; + GlobalTensor propagatedVNew_; + GlobalTensor propagatedH_; + GlobalTensor solveWorkspace_; + GlobalTensor scoreWorkspace_; + TPipe *pipe_ = nullptr; + TBuf exp2Buf_; + TBuf vecBuf_; + TBuf gateWritebackBuf_; + TEventID mte2ToVEvent_ = 0; + TEventID vToMte2Event_ = 0; + TEventID vToMte3Event_ = 0; + TEventID mte3ToVEvent_ = 0; + TEventID mte2ToMte3Event_ = 0; + TEventID mte3ToMte2Event_ = 0; + bool vectorEventsAllocated_ = false; + Catlass::Arch::CrossCoreFlagWithReverse scoreReadyFlag_{KDA_SCORE_READY_FLAG0, + KDA_SCORE_READY_FLAG1}; + Catlass::Arch::CrossCoreFlagWithReverse scoreDoneFlag_{KDA_SCORE_DONE_FLAG0, + KDA_SCORE_DONE_FLAG1}; + // Score production is fully drained before solve starts, so the solve handshake can safely reuse + // the existing score flags without consuming additional hardware flag IDs. + Catlass::Arch::CrossCoreFlagWithReverse syncReadyFlag_{KDA_SCORE_READY_FLAG0, + KDA_SCORE_READY_FLAG1}; + Catlass::Arch::CrossCoreFlagWithReverse syncDoneFlag_{KDA_SCORE_DONE_FLAG0, + KDA_SCORE_DONE_FLAG1}; + uint64_t B_ = 0; + uint64_t N_ = 0; + uint64_t H_ = 0; + uint64_t HV_ = 0; + uint64_t T_ = 0; + uint64_t K_ = 0; + uint64_t V_ = 0; + uint64_t BT_ = 0; + uint64_t NT_ = 0; + float scale_ = 1.0f; + bool hasInitial_ = false; + bool isVarLen_ = false; + bool isAivOnly_ = false; + bool inputSequenceMajor_ = false; + uint64_t usedCoreNum_ = 1; + uint64_t solveCoreIdx_ = 0; + __gm__ int64_t *chunkIndicesAddr_ = nullptr; + __gm__ int64_t *cuSeqlensAddr_ = nullptr; +}; +} // namespace + +template +__aicore__ inline void RunChunkKdaPostWu(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR beta, + GM_ADDR initialState, GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR wSeed, + GM_ADDR akk, GM_ADDR uSeed, GM_ADDR w, GM_ADDR u, GM_ADDR kg, GM_ADDR vNew, + GM_ADDR userWorkspace, const TilingData &tiling, TPipe &pipe) +{ + GM_ADDR postScratch = userWorkspace + tiling.postWuScratchOffset; + if ASCEND_IS_AIC { + ChunkKdaFwdPostWuKernel op; + op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, wSeed, akk, uSeed, nullptr, userWorkspace, + userWorkspace, userWorkspace, akk, w, u, userWorkspace, kg, vNew, postScratch, postScratch, tiling, + &pipe, false); + op.ProcessAic(); + } + if ASCEND_IS_AIV { + ChunkKdaFwdPostWuKernel op; + op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, wSeed, akk, uSeed, nullptr, userWorkspace, + userWorkspace, userWorkspace, akk, w, u, userWorkspace, kg, vNew, postScratch, postScratch, tiling, + &pipe); + op.ProcessAiv(); + } +} + +} // namespace KdaPostWu diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_prepare_arch35.h b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_prepare_arch35.h new file mode 100644 index 0000000000..bcb92bd14a --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/chunk_kda_fwd_prepare_arch35.h @@ -0,0 +1,4806 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#pragma once + +#ifndef CATLASS_ARCH +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +#define CATLASS_ARCH 3510 +#else +#define CATLASS_ARCH 2201 +#endif +#endif + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/arch/kda_cross_core_sync.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/gemm/block/kda_block_mmad.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "../kernel_utils/block/block_mmad_pingpong_tla_multi.hpp" +#include "../kernel_utils/tile/copy_l0c_to_ub.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "kernel_operator.h" +#include "../chunk_kda_fwd_varlen.h" +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +#ifndef FLA_NPU_REGBASE_HPP_INCLUDED +#define FLA_NPU_REGBASE_HPP_INCLUDED +#include "../kernel_utils/vector/regbase.hpp" +#endif +#endif +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" +#include "chunk_kda_fwd_post_wu_arch35.h" + +using namespace AscendC; + +namespace KdaPrepare { +namespace { +using KdaInt64 = tla::Int<64>; +using KdaInt128 = tla::Int<128>; +constexpr float LN2 = 0.69314718055994530942f; +constexpr float RCP_LN2 = 1.44269504088896340736f; +constexpr float KDA_EXP2_CLAMP = 80.0f; +constexpr float KDA_EXP_INPUT_MAX = KDA_EXP2_CLAMP * LN2; +constexpr float KDA_EXP_INPUT_MIN = -KDA_EXP2_CLAMP * LN2; +constexpr float KDA_SCORE_EXP2_CLAMP = 120.0f; +constexpr float KDA_SCORE_EXP2_MIN_CLAMP = 126.0f; +constexpr float KDA_SCORE_EXP_INPUT_MAX = KDA_SCORE_EXP2_CLAMP * LN2; +constexpr float KDA_SCORE_EXP_INPUT_MIN = -KDA_SCORE_EXP2_MIN_CLAMP * LN2; +constexpr float KDA_FP16_MAX = 65504.0f; +constexpr uint32_t EXP2_UB_ELEMENTS = 256; +constexpr uint32_t EXP2_UB_BYTES = EXP2_UB_ELEMENTS * (sizeof(float) + sizeof(uint16_t)); +constexpr uint32_t EXP2_EVENT_ID = 0; +constexpr uint32_t KDA_SOLVE_BT = 64; +constexpr uint32_t KDA_SOLVE_MATRIX_ELEMENTS = KDA_SOLVE_BT * KDA_SOLVE_BT; +constexpr uint32_t KDA_SOLVE_SCRATCH_X = 0; +constexpr uint32_t KDA_SOLVE_SCRATCH_Y0 = 1; +constexpr uint32_t KDA_SOLVE_SCRATCH_TMP = 2; +constexpr uint32_t KDA_SOLVE_SCRATCH_Y1 = 3; +constexpr uint32_t KDA_SOLVE_SCRATCH_IDENTITY = 4; +constexpr uint32_t KDA_SOLVE_SCRATCH_RAW_AKK = KDA_SOLVE_SCRATCH_Y1; +constexpr uint32_t KDA_SOLVE_SCRATCH_RAW_AQK = KDA_SOLVE_SCRATCH_IDENTITY; +constexpr uint32_t KDA_SOLVE_SCRATCH_SLOTS = 5; +constexpr uint32_t KDA_SOLVE_PIPELINE_DEPTH = 4; +constexpr uint32_t KDA_SOLVE_DIAG_BT = 16; +constexpr uint32_t KDA_SOLVE_DIAG_BLOCKS = KDA_SOLVE_BT / KDA_SOLVE_DIAG_BT; +constexpr uint32_t KDA_SOLVE_DIAG_MCH_ITERS = 3; +// Keep the local safe-gate exponent span within the BF16 score range while +// reducing repeated gate-factor work and AIV/AIC handshakes. +constexpr uint32_t KDA_SCORE_REF_BC = 32; +constexpr uint32_t KDA_SAFE_SCORE_REF_BC = 32; +constexpr uint32_t KDA_VEC_ARENA_ELEMENTS = 32768; +constexpr uint32_t KDA_BITS_PER_MASK_BYTE = 8; +constexpr uint32_t KDA_SELECT_COL_BLOCKS = 2; +constexpr uint32_t KDA_SELECT_COL_MASK_BYTES = KDA_SOLVE_MATRIX_ELEMENTS / KDA_BITS_PER_MASK_BYTE; +constexpr uint32_t KDA_SELECT_MASK_BYTES = KDA_SELECT_COL_BLOCKS * KDA_SELECT_COL_MASK_BYTES; +constexpr uint32_t KDA_SELECT_AQK_MASK_BYTE_OFFSET = 120 * 1024; +constexpr uint32_t KDA_SELECT_AKK_MASK_BYTE_OFFSET = KDA_SELECT_AQK_MASK_BYTE_OFFSET + KDA_SELECT_MASK_BYTES; +constexpr uint32_t KDA_SELECT_ZERO_BYTE_OFFSET = KDA_SELECT_AKK_MASK_BYTE_OFFSET + KDA_SELECT_MASK_BYTES; +constexpr uint32_t KDA_SELECT_ZERO_FLOAT_OFFSET = KDA_SELECT_ZERO_BYTE_OFFSET / sizeof(float); +constexpr uint8_t KDA_SCORE_DONE_FLAG0 = 2; +constexpr uint8_t KDA_SCORE_DONE_FLAG1 = 3; +constexpr uint8_t KDA_SCORE_READY_FLAG0 = 4; +constexpr uint8_t KDA_SCORE_READY_FLAG1 = 5; +constexpr uint8_t KDA_SOLVE_DONE_FLAG = 6; +constexpr uint8_t KDA_SOLVE_READY_FLAG = 7; +constexpr uint8_t KDA_POST_READY_FLAG = 0; +constexpr uint8_t KDA_POST_FREE_FLAG = 1; +constexpr uint32_t KDA_SCORE_QUEUE_DEPTH = 2; +constexpr uint32_t KDA_SCORE_LANES = 2; +constexpr uint32_t KDA_POST_QUEUE_DEPTH = 4; +constexpr uint32_t KDA_POST_QUEUE_STORAGE = KDA_POST_QUEUE_DEPTH + 1; +constexpr uint32_t KDA_DIRECT_SCORE_QUEUE_DEPTH = 2; +constexpr uint32_t KDA_DIRECT_SCORE_ROWS = 32; +constexpr uint32_t KDA_DIRECT_SCORE_MATRIX_ELEMENTS = KDA_DIRECT_SCORE_ROWS * 64; +constexpr uint32_t KDA_DIRECT_SCORE_SLOT_ELEMENTS = 3 * KDA_DIRECT_SCORE_MATRIX_ELEMENTS; +constexpr uint32_t KDA_DIRECT_SCORE_FLOAT_OFFSET = 20 * 1024; +constexpr uint32_t KDA_DIRECT_SCORE_UB_BYTE_OFFSET = EXP2_UB_BYTES + KDA_DIRECT_SCORE_FLOAT_OFFSET * sizeof(float); +constexpr uint32_t KDA_DIRECT_SCORE_L1_SLOT_ELEMENTS = 2 * KDA_DIRECT_SCORE_ROWS * 128; +constexpr uint32_t KDA_DIRECT_SCORE_L1_SLOT_BYTES = KDA_DIRECT_SCORE_L1_SLOT_ELEMENTS * sizeof(uint16_t); +constexpr uint32_t KDA_DIRECT_SCORE_L1_B_OFFSET = + KDA_SCORE_QUEUE_DEPTH * KDA_SCORE_LANES * KDA_DIRECT_SCORE_L1_SLOT_BYTES; +constexpr uint64_t KDA_DIRECT_SCORE_FREE_FLAG = 8; +constexpr uint64_t KDA_DIRECT_SCORE_READY_FLAG = 10; +constexpr uint64_t KDA_DIRECT_SCORE_SUBBLOCK_FLAG_STRIDE = 16; +constexpr uint32_t KDA_SCORE_SCRATCH_SLOTS = KDA_SCORE_QUEUE_DEPTH * KDA_SCORE_LANES; +constexpr uint32_t KDA_SYNC_REVERSE_DEPTH = 1; +constexpr uint32_t KDA_SCORE_SCRATCH_PLANES = 3; +constexpr uint32_t KDA_SCORE_SCRATCH_QG = 0; +constexpr uint32_t KDA_SCORE_SCRATCH_W = 1; +constexpr uint32_t KDA_SCORE_SCRATCH_KG = 2; +constexpr uint64_t KDA_WORKSPACE_ALIGN = 512; +constexpr uint32_t KDA_GATE_TILE_ROWS = 16; +constexpr uint32_t KDA_GATE_PIPELINE_DEPTH = 3; +// beta is staged in UB as fp32 or bf16; the slot is always sized for the wider one +// so the UB layout does not depend on the beta dtype. +constexpr uint64_t KDA_BETA_STAGE_BYTES = sizeof(float); +constexpr uint32_t KDA_AIV_UB_BUDGET_BYTES = 192 * 1024; +constexpr uint32_t KDA_LOCAL_GK_FLOAT_OFFSET = 10 * 1024; +constexpr uint32_t KDA_SCALED_QG_FLOAT_OFFSET = 18 * 1024; +constexpr bool KDA_ARCH35_ENABLE_HEAD_PAIR = true; +constexpr bool KDA_ARCH35_ENABLE_MANUAL_SCORE_PIPELINE = false; +constexpr bool KDA_ARCH35_ENABLE_DIRECT_SCORE_UB = true; +constexpr bool KDA_ARCH35_ENABLE_DIRECT_SCORE_L1 = false; +constexpr uint16_t KDA_ARCH35_SCORE_EVENT = 3; +constexpr uint16_t KDA_ARCH35_SCORE_W_EVENT = 4; +constexpr uint16_t KDA_ARCH35_SOLVE_FIX_EVENT = 7; + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +template +static __simd_vf__ inline void AccumulateRawSafeGateChunk128Regbase(__ubuf__ float *input, __ubuf__ float *bias, + __ubuf__ float *acc, uint16_t rows, float expA, + float lowerBound) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t FLOAT_ELEMENTS_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(float); + constexpr uint16_t ROW_ELEMENTS = 2 * FLOAT_ELEMENTS_PER_REG; + + MaskReg floatMask = CreateMask(); + RegTensor accZeroReg; + RegTensor accOneReg; + RegTensor oneZeroReg; + RegTensor oneOneReg; + RegTensor biasZeroReg; + RegTensor biasOneReg; + LoadAlign(accZeroReg, acc); + LoadAlign(accOneReg, acc + FLOAT_ELEMENTS_PER_REG); + Duplicate(oneZeroReg, 1.0f, floatMask); + Duplicate(oneOneReg, 1.0f, floatMask); + if constexpr (HAS_BIAS) { + LoadAlign(biasZeroReg, bias); + LoadAlign(biasOneReg, bias + FLOAT_ELEMENTS_PER_REG); + } + + const float gateScale = lowerBound * RCP_LN2; + for (uint16_t row = 0; row < rows; ++row) { + const uint32_t rowOffset = static_cast(row) * ROW_ELEMENTS; + RegTensor gateZeroReg; + RegTensor gateOneReg; + RegTensor sigmoidZeroReg; + RegTensor sigmoidOneReg; + LoadAlign(gateZeroReg, input + rowOffset); + LoadAlign(gateOneReg, input + rowOffset + FLOAT_ELEMENTS_PER_REG); + if constexpr (HAS_BIAS) { + Add(gateZeroReg, gateZeroReg, biasZeroReg, floatMask); + Add(gateOneReg, gateOneReg, biasOneReg, floatMask); + } + Muls(gateZeroReg, gateZeroReg, -expA, floatMask); + Muls(gateOneReg, gateOneReg, -expA, floatMask); + Exp(gateZeroReg, gateZeroReg, floatMask); + Exp(gateOneReg, gateOneReg, floatMask); + Adds(gateZeroReg, gateZeroReg, 1.0f, floatMask); + Adds(gateOneReg, gateOneReg, 1.0f, floatMask); + Div(sigmoidZeroReg, oneZeroReg, gateZeroReg, floatMask); + Div(sigmoidOneReg, oneOneReg, gateOneReg, floatMask); + Muls(sigmoidZeroReg, sigmoidZeroReg, gateScale, floatMask); + Muls(sigmoidOneReg, sigmoidOneReg, gateScale, floatMask); + Add(accZeroReg, accZeroReg, sigmoidZeroReg, floatMask); + Add(accOneReg, accOneReg, sigmoidOneReg, floatMask); + StoreAlign(input + rowOffset, accZeroReg, floatMask); + StoreAlign(input + rowOffset + FLOAT_ELEMENTS_PER_REG, accOneReg, floatMask); + } + StoreAlign(acc, accZeroReg, floatMask); + StoreAlign(acc + FLOAT_ELEMENTS_PER_REG, accOneReg, floatMask); +} + +template +__simd_callee__ inline void LoadKdaGateRegbasePair(AscendC::MicroAPI::RegTensor &zeroReg, + AscendC::MicroAPI::RegTensor &oneReg, __ubuf__ InputT *src, + AscendC::MicroAPI::MaskReg &inputMask) +{ + using namespace AscendC::MicroAPI; + if constexpr (std::is_same()) { + LoadAlign(zeroReg, oneReg, src); + } else { + RegTensor inputReg; + LoadIn(inputReg, src); + CastHalf2Float(zeroReg, oneReg, inputReg, inputMask); + } +} + +template +__simd_callee__ inline void ClampKdaGateRegbaseOutput(AscendC::MicroAPI::RegTensor &zeroReg, + AscendC::MicroAPI::RegTensor &oneReg, + AscendC::MicroAPI::MaskReg &floatMask) +{ + using namespace AscendC::MicroAPI; + if constexpr (std::is_same()) { + Mins(zeroReg, zeroReg, KDA_FP16_MAX, floatMask); + Mins(oneReg, oneReg, KDA_FP16_MAX, floatMask); + Maxs(zeroReg, zeroReg, -KDA_FP16_MAX, floatMask); + Maxs(oneReg, oneReg, -KDA_FP16_MAX, floatMask); + } +} + +template +__simd_callee__ inline void BuildKdaGateRegbaseExp(AscendC::MicroAPI::RegTensor &expZeroReg, + AscendC::MicroAPI::RegTensor &expOneReg, + AscendC::MicroAPI::RegTensor &gateZeroReg, + AscendC::MicroAPI::RegTensor &gateOneReg, __ubuf__ float *ref, + AscendC::MicroAPI::MaskReg &floatMask) +{ + using namespace AscendC::MicroAPI; + constexpr float expInputMax = std::is_same() ? KDA_SCORE_EXP_INPUT_MAX : KDA_EXP_INPUT_MAX; + constexpr float expInputMin = std::is_same() ? KDA_SCORE_EXP_INPUT_MIN : KDA_EXP_INPUT_MIN; + if constexpr (USE_REF) { + RegTensor refZeroReg; + RegTensor refOneReg; + LoadAlign(refZeroReg, refOneReg, ref); + if constexpr (NEGATIVE) { + SubFloatTwoReg(expZeroReg, expOneReg, refZeroReg, refOneReg, gateZeroReg, gateOneReg, floatMask); + } else { + SubFloatTwoReg(expZeroReg, expOneReg, gateZeroReg, gateOneReg, refZeroReg, refOneReg, floatMask); + } + } else if constexpr (NEGATIVE) { + Muls(expZeroReg, gateZeroReg, -1.0f, floatMask); + Muls(expOneReg, gateOneReg, -1.0f, floatMask); + } else { + Adds(expZeroReg, gateZeroReg, 0.0f, floatMask); + Adds(expOneReg, gateOneReg, 0.0f, floatMask); + } + Muls(expZeroReg, expZeroReg, LN2, floatMask); + Muls(expOneReg, expOneReg, LN2, floatMask); + MinsFloatTwoReg(expZeroReg, expOneReg, expZeroReg, expOneReg, expInputMax, floatMask); + Maxs(expZeroReg, expZeroReg, expInputMin, floatMask); + Maxs(expOneReg, expOneReg, expInputMin, floatMask); + ExpFloatTwoReg(expZeroReg, expOneReg, expZeroReg, expOneReg, floatMask); +} + +template +__simd_callee__ inline void StoreKdaGateRegbasePair(__ubuf__ OutputT *dst, AscendC::MicroAPI::RegTensor &zeroReg, + AscendC::MicroAPI::RegTensor &oneReg, + AscendC::MicroAPI::MaskReg &inputMask, + AscendC::MicroAPI::MaskReg &floatMask) +{ + using namespace AscendC::MicroAPI; + RegTensor outputReg; + ClampKdaGateRegbaseOutput(zeroReg, oneReg, floatMask); + CastFloat2Half(outputReg, zeroReg, oneReg, floatMask); + StoreAlign(dst, outputReg, inputMask); +} + +template +static __simd_vf__ inline void PrepareKdaGateQwRegbase(__ubuf__ InputT *q, __ubuf__ InputT *k, __ubuf__ OutputT *qOut, + __ubuf__ OutputT *kOut, __ubuf__ InputT *qDirect, + __ubuf__ InputT *kDirect, __ubuf__ GK_T *gate, + __ubuf__ float *ref, uint16_t rows, uint16_t cols) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t ELEMENTS_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(InputT); + + MaskReg floatMask = CreateMask(); + for (uint16_t row = 0; row < rows; ++row) { + uint32_t rowOffset = static_cast(row) * cols; + for (uint16_t col = 0; col < cols; col += ELEMENTS_PER_REG) { + uint32_t activeCount = static_cast(cols - col); + MaskReg inputMask = UpdateMask(activeCount); + uint32_t offset = rowOffset + col; + + RegTensor gateZeroReg; + RegTensor gateOneReg; + RegTensor expZeroReg; + RegTensor expOneReg; + RegTensor directZeroReg; + RegTensor directOneReg; + RegTensor inputZeroReg; + RegTensor inputOneReg; + RegTensor outputZeroReg; + RegTensor outputOneReg; + + LoadKdaGateRegbasePair(gateZeroReg, gateOneReg, gate + offset, inputMask); + BuildKdaGateRegbaseExp(expZeroReg, expOneReg, gateZeroReg, gateOneReg, ref + col, + floatMask); + BuildKdaGateRegbaseExp(directZeroReg, directOneReg, gateZeroReg, gateOneReg, + ref + col, floatMask); + + LoadKdaGateRegbasePair(inputZeroReg, inputOneReg, q + offset, inputMask); + MulFloatTwoReg(outputZeroReg, outputOneReg, inputZeroReg, inputOneReg, directZeroReg, directOneReg, + floatMask); + StoreKdaGateRegbasePair(qDirect + offset, outputZeroReg, outputOneReg, inputMask, floatMask); + MulFloatTwoReg(outputZeroReg, outputOneReg, inputZeroReg, inputOneReg, expZeroReg, expOneReg, floatMask); + StoreKdaGateRegbasePair(qOut + offset, outputZeroReg, outputOneReg, inputMask, floatMask); + + LoadKdaGateRegbasePair(inputZeroReg, inputOneReg, k + offset, inputMask); + MulFloatTwoReg(outputZeroReg, outputOneReg, inputZeroReg, inputOneReg, directZeroReg, directOneReg, + floatMask); + StoreKdaGateRegbasePair(kDirect + offset, outputZeroReg, outputOneReg, inputMask, floatMask); + MulFloatTwoReg(outputZeroReg, outputOneReg, inputZeroReg, inputOneReg, expZeroReg, expOneReg, floatMask); + StoreKdaGateRegbasePair(kOut + offset, outputZeroReg, outputOneReg, inputMask, floatMask); + } + } +} + +template +static __simd_vf__ inline void PrepareKdaGateKgRegbase(__ubuf__ OutputT *kg, __ubuf__ InputT *k, __ubuf__ GK_T *gate, + __ubuf__ float *ref, uint16_t rows, uint16_t cols, + uint16_t validRows) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t ELEMENTS_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(InputT); + + MaskReg floatMask = CreateMask(); + for (uint16_t row = 0; row < rows; ++row) { + uint32_t rowOffset = static_cast(row) * cols; + for (uint16_t col = 0; col < cols; col += ELEMENTS_PER_REG) { + uint32_t activeCount = static_cast(cols - col); + MaskReg inputMask = UpdateMask(activeCount); + uint32_t offset = rowOffset + col; + + RegTensor gateZeroReg; + RegTensor gateOneReg; + RegTensor expZeroReg; + RegTensor expOneReg; + RegTensor inputZeroReg; + RegTensor inputOneReg; + RegTensor outputZeroReg; + RegTensor outputOneReg; + + LoadKdaGateRegbasePair(gateZeroReg, gateOneReg, gate + offset, inputMask); + BuildKdaGateRegbaseExp(expZeroReg, expOneReg, gateZeroReg, gateOneReg, ref + col, + floatMask); + LoadKdaGateRegbasePair(inputZeroReg, inputOneReg, k + offset, inputMask); + MulFloatTwoReg(outputZeroReg, outputOneReg, inputZeroReg, inputOneReg, expZeroReg, expOneReg, floatMask); + if constexpr (USE_REF) { + if (row >= validRows) { + Duplicate(outputZeroReg, 0.0f, floatMask); + Duplicate(outputOneReg, 0.0f, floatMask); + } + } + StoreKdaGateRegbasePair(kg + offset, outputZeroReg, outputOneReg, inputMask, floatMask); + } + } +} + +template +static __simd_vf__ inline void PrepareKdaGateQwKgRegbase( + __ubuf__ InputT *q, __ubuf__ InputT *k, __ubuf__ OutputT *qOut, __ubuf__ OutputT *wOut, __ubuf__ OutputT *kgOut, + __ubuf__ InputT *qDirect, __ubuf__ InputT *wDirect, __ubuf__ InputT *v, __ubuf__ InputT *vDirect, + __ubuf__ InputT *finalKgOut, __ubuf__ float *beta, __ubuf__ GK_T *gate, __ubuf__ float *ref, + __ubuf__ float *finalRef, uint16_t rows, uint16_t cols, uint16_t validRows) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t ELEMENTS_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(InputT); + constexpr float scoreExpInputMax = + std::is_same() ? KDA_SCORE_EXP_INPUT_MAX : KDA_EXP_INPUT_MAX; + constexpr float scoreExpInputMin = + std::is_same() ? KDA_SCORE_EXP_INPUT_MIN : KDA_EXP_INPUT_MIN; + + MaskReg floatMask = CreateMask(); + for (uint16_t row = 0; row < rows; ++row) { + uint32_t rowOffset = static_cast(row) * cols; + for (uint16_t col = 0; col < cols; col += ELEMENTS_PER_REG) { + uint32_t activeCount = static_cast(cols - col); + MaskReg inputMask = UpdateMask(activeCount); + uint32_t offset = rowOffset + col; + + RegTensor gateZeroReg; + RegTensor gateOneReg; + RegTensor posZeroReg; + RegTensor posOneReg; + RegTensor negZeroReg; + RegTensor negOneReg; + RegTensor directZeroReg; + RegTensor directOneReg; + RegTensor finalNegZeroReg; + RegTensor finalNegOneReg; + RegTensor inputZeroReg; + RegTensor inputOneReg; + RegTensor outputZeroReg; + RegTensor outputOneReg; + RegTensor betaReg; + + LoadKdaGateRegbasePair(gateZeroReg, gateOneReg, gate + offset, inputMask); + if constexpr (USE_REF) { + RegTensor refZeroReg; + RegTensor refOneReg; + LoadAlign(refZeroReg, refOneReg, ref + col); + SubFloatTwoReg(posZeroReg, posOneReg, gateZeroReg, gateOneReg, refZeroReg, refOneReg, floatMask); + SubFloatTwoReg(negZeroReg, negOneReg, refZeroReg, refOneReg, gateZeroReg, gateOneReg, floatMask); + } else { + Adds(posZeroReg, gateZeroReg, 0.0f, floatMask); + Adds(posOneReg, gateOneReg, 0.0f, floatMask); + Muls(negZeroReg, gateZeroReg, -1.0f, floatMask); + Muls(negOneReg, gateOneReg, -1.0f, floatMask); + } + if constexpr (STORE_DIRECT) { + Adds(directZeroReg, gateZeroReg, 0.0f, floatMask); + Adds(directOneReg, gateOneReg, 0.0f, floatMask); + Muls(directZeroReg, directZeroReg, LN2, floatMask); + Muls(directOneReg, directOneReg, LN2, floatMask); + MinsFloatTwoReg(directZeroReg, directOneReg, directZeroReg, directOneReg, KDA_EXP_INPUT_MAX, floatMask); + Maxs(directZeroReg, directZeroReg, KDA_EXP_INPUT_MIN, floatMask); + Maxs(directOneReg, directOneReg, KDA_EXP_INPUT_MIN, floatMask); + ExpFloatTwoReg(directZeroReg, directOneReg, directZeroReg, directOneReg, floatMask); + } + Muls(posZeroReg, posZeroReg, LN2, floatMask); + Muls(posOneReg, posOneReg, LN2, floatMask); + Muls(negZeroReg, negZeroReg, LN2, floatMask); + Muls(negOneReg, negOneReg, LN2, floatMask); + MinsFloatTwoReg(posZeroReg, posOneReg, posZeroReg, posOneReg, scoreExpInputMax, floatMask); + MinsFloatTwoReg(negZeroReg, negOneReg, negZeroReg, negOneReg, scoreExpInputMax, floatMask); + Maxs(posZeroReg, posZeroReg, scoreExpInputMin, floatMask); + Maxs(posOneReg, posOneReg, scoreExpInputMin, floatMask); + Maxs(negZeroReg, negZeroReg, scoreExpInputMin, floatMask); + Maxs(negOneReg, negOneReg, scoreExpInputMin, floatMask); + ExpFloatTwoReg(posZeroReg, posOneReg, posZeroReg, posOneReg, floatMask); + ExpFloatTwoReg(negZeroReg, negOneReg, negZeroReg, negOneReg, floatMask); + + LoadKdaGateRegbasePair(inputZeroReg, inputOneReg, q + offset, inputMask); + if constexpr (STORE_DIRECT) { + MulFloatTwoReg(outputZeroReg, outputOneReg, inputZeroReg, inputOneReg, directZeroReg, directOneReg, + floatMask); + StoreKdaGateRegbasePair(qDirect + offset, outputZeroReg, outputOneReg, inputMask, floatMask); + } + MulFloatTwoReg(outputZeroReg, outputOneReg, inputZeroReg, inputOneReg, posZeroReg, posOneReg, floatMask); + StoreKdaGateRegbasePair(qOut + offset, outputZeroReg, outputOneReg, inputMask, floatMask); + + LoadKdaGateRegbasePair(inputZeroReg, inputOneReg, k + offset, inputMask); + if constexpr (STORE_DIRECT || SCALE_SCORE_W) { + LoadAlign(betaReg, beta + row); + } + if constexpr (STORE_DIRECT) { + MulFloatTwoReg(outputZeroReg, outputOneReg, inputZeroReg, inputOneReg, directZeroReg, directOneReg, + floatMask); + RegTensor roundedReg; + ClampKdaGateRegbaseOutput(outputZeroReg, outputOneReg, floatMask); + CastFloat2Half(roundedReg, outputZeroReg, outputOneReg, floatMask); + CastHalf2Float(outputZeroReg, outputOneReg, roundedReg, inputMask); + Mul(outputZeroReg, outputZeroReg, betaReg, floatMask); + Mul(outputOneReg, outputOneReg, betaReg, floatMask); + StoreKdaGateRegbasePair(wDirect + offset, outputZeroReg, outputOneReg, inputMask, floatMask); + } + MulFloatTwoReg(outputZeroReg, outputOneReg, inputZeroReg, inputOneReg, posZeroReg, posOneReg, floatMask); + if constexpr (SCALE_SCORE_W) { + Mul(outputZeroReg, outputZeroReg, betaReg, floatMask); + Mul(outputOneReg, outputOneReg, betaReg, floatMask); + } + StoreKdaGateRegbasePair(wOut + offset, outputZeroReg, outputOneReg, inputMask, floatMask); + MulFloatTwoReg(outputZeroReg, outputOneReg, inputZeroReg, inputOneReg, negZeroReg, negOneReg, floatMask); + if constexpr (USE_REF) { + if (row >= validRows) { + Duplicate(outputZeroReg, 0.0f, floatMask); + Duplicate(outputOneReg, 0.0f, floatMask); + } + } + StoreKdaGateRegbasePair(kgOut + offset, outputZeroReg, outputOneReg, inputMask, floatMask); + + if constexpr (STORE_DIRECT) { + LoadKdaGateRegbasePair(inputZeroReg, inputOneReg, v + offset, inputMask); + Mul(outputZeroReg, inputZeroReg, betaReg, floatMask); + Mul(outputOneReg, inputOneReg, betaReg, floatMask); + StoreKdaGateRegbasePair(vDirect + offset, outputZeroReg, outputOneReg, inputMask, floatMask); + } + if constexpr (EXPORT_FINAL_KG) { + RegTensor finalRefZeroReg; + RegTensor finalRefOneReg; + LoadAlign(finalRefZeroReg, finalRefOneReg, finalRef + col); + SubFloatTwoReg(finalNegZeroReg, finalNegOneReg, finalRefZeroReg, finalRefOneReg, gateZeroReg, + gateOneReg, floatMask); + Muls(finalNegZeroReg, finalNegZeroReg, LN2, floatMask); + Muls(finalNegOneReg, finalNegOneReg, LN2, floatMask); + MinsFloatTwoReg(finalNegZeroReg, finalNegOneReg, finalNegZeroReg, finalNegOneReg, KDA_EXP_INPUT_MAX, + floatMask); + Maxs(finalNegZeroReg, finalNegZeroReg, KDA_EXP_INPUT_MIN, floatMask); + Maxs(finalNegOneReg, finalNegOneReg, KDA_EXP_INPUT_MIN, floatMask); + ExpFloatTwoReg(finalNegZeroReg, finalNegOneReg, finalNegZeroReg, finalNegOneReg, floatMask); + LoadKdaGateRegbasePair(inputZeroReg, inputOneReg, k + offset, inputMask); + MulFloatTwoReg(outputZeroReg, outputOneReg, inputZeroReg, inputOneReg, finalNegZeroReg, finalNegOneReg, + floatMask); + StoreKdaGateRegbasePair(finalKgOut + offset, outputZeroReg, outputOneReg, inputMask, floatMask); + } + } + } +} + +static __simd_vf__ inline void ForwardSubDiag16Regbase(__ubuf__ float *diag, uint16_t valid) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t DIAG_SIZE = KDA_SOLVE_DIAG_BT; + uint32_t activeCount = DIAG_SIZE; + MaskReg rowMask = UpdateMask(activeCount); + + for (uint16_t row = 2; row < valid; ++row) { + RegTensor currentReg; + RegTensor scaleReg; + RegTensor matrixReg; + RegTensor productReg; + RegTensor sumReg; + LoadAlign(currentReg, diag + static_cast(row) * DIAG_SIZE); + Duplicate(sumReg, 0.0f, rowMask); + + for (uint16_t sourceRow = 0; sourceRow < row; ++sourceRow) { + LoadAlign(scaleReg, + diag + static_cast(row) * DIAG_SIZE + sourceRow); + LoadAlign(matrixReg, diag + static_cast(sourceRow) * DIAG_SIZE); + Mul(productReg, matrixReg, scaleReg, rowMask); + Add(sumReg, sumReg, productReg, rowMask); + } + Add(currentReg, currentReg, sumReg, rowMask); + StoreAlign(diag + static_cast(row) * DIAG_SIZE, currentReg, rowMask); + LocalMemBar(); + } +} + +static __simd_vf__ inline void SelectCausalRows64Regbase(__ubuf__ float *aqk, __ubuf__ float *akk, uint16_t rowBegin, + uint16_t rowCount) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t ROW_ELEMENTS = 64; + + MaskReg fullMask = CreateMask(); + for (uint16_t localRow = 0; localRow < rowCount; ++localRow) { + const uint16_t row = rowBegin + localRow; + const uint32_t rowOffset = static_cast(localRow) * ROW_ELEMENTS; + RegTensor zeroReg; + RegTensor aqkInputReg; + RegTensor akkInputReg; + RegTensor aqkReg; + RegTensor akkReg; + uint32_t aqkCount = static_cast(row) + 1; + uint32_t akkCount = static_cast(row); + MaskReg aqkMask = UpdateMask(aqkCount); + MaskReg akkMask = UpdateMask(akkCount); + Duplicate(zeroReg, 0.0f, fullMask); + LoadAlign(aqkInputReg, aqk + rowOffset); + LoadAlign(akkInputReg, akk + rowOffset); + Select(aqkReg, aqkInputReg, zeroReg, aqkMask); + Select(akkReg, akkInputReg, zeroReg, akkMask); + StoreAlign(aqk + rowOffset, aqkReg, fullMask); + StoreAlign(akk + rowOffset, akkReg, fullMask); + } +} + +static __simd_vf__ inline void ForwardSubDiag16StridedRegbase(__ubuf__ float *matrix, uint16_t rowStride, + uint16_t rowBegin, uint16_t colBegin, uint16_t valid) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t DIAG_SIZE = KDA_SOLVE_DIAG_BT; + uint32_t activeCount = DIAG_SIZE; + MaskReg rowMask = UpdateMask(activeCount); + + for (uint16_t row = 2; row < valid; ++row) { + uint32_t currentOffset = static_cast(rowBegin + row) * rowStride + colBegin; + RegTensor currentReg; + RegTensor scaleReg; + RegTensor matrixReg; + RegTensor productReg; + RegTensor sumReg; + LoadAlign(currentReg, matrix + currentOffset); + Duplicate(sumReg, 0.0f, rowMask); + + for (uint16_t sourceRow = 0; sourceRow < row; ++sourceRow) { + LoadAlign(scaleReg, matrix + currentOffset + sourceRow); + uint32_t sourceOffset = static_cast(rowBegin + sourceRow) * rowStride + colBegin; + LoadAlign(matrixReg, matrix + sourceOffset); + Mul(productReg, matrixReg, scaleReg, rowMask); + Add(sumReg, sumReg, productReg, rowMask); + } + Add(currentReg, currentReg, sumReg, rowMask); + StoreAlign(matrix + currentOffset, currentReg, rowMask); + LocalMemBar(); + } +} + +static __simd_vf__ inline void ForwardSubDiag16PairStridedRegbase(__ubuf__ float *matrix, uint16_t rowStride, + uint16_t colBegin, uint16_t firstValid, + uint16_t secondValid) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t DIAG_SIZE = KDA_SOLVE_DIAG_BT; + constexpr uint16_t SECOND_LOCAL_ROW = DIAG_SIZE; + uint32_t activeCount = DIAG_SIZE; + MaskReg rowMask = UpdateMask(activeCount); + + for (uint16_t row = 2; row < DIAG_SIZE; ++row) { + uint32_t firstCurrentOffset = static_cast(row) * rowStride + colBegin; + uint32_t secondCurrentOffset = static_cast(SECOND_LOCAL_ROW + row) * rowStride + colBegin + DIAG_SIZE; + RegTensor firstCurrentReg; + RegTensor secondCurrentReg; + RegTensor firstSumReg; + RegTensor secondSumReg; + LoadAlign(firstCurrentReg, matrix + firstCurrentOffset); + LoadAlign(secondCurrentReg, matrix + secondCurrentOffset); + Duplicate(firstSumReg, 0.0f, rowMask); + Duplicate(secondSumReg, 0.0f, rowMask); + + if (row < firstValid || row < secondValid) { + for (uint16_t sourceRow = 0; sourceRow < row; ++sourceRow) { + RegTensor firstScaleReg; + RegTensor secondScaleReg; + RegTensor firstMatrixReg; + RegTensor secondMatrixReg; + RegTensor firstProductReg; + RegTensor secondProductReg; + LoadAlign(firstScaleReg, matrix + firstCurrentOffset + sourceRow); + LoadAlign(secondScaleReg, matrix + secondCurrentOffset + sourceRow); + uint32_t firstSourceOffset = static_cast(sourceRow) * rowStride + colBegin; + uint32_t secondSourceOffset = + static_cast(SECOND_LOCAL_ROW + sourceRow) * rowStride + colBegin + DIAG_SIZE; + LoadAlign(firstMatrixReg, matrix + firstSourceOffset); + LoadAlign(secondMatrixReg, matrix + secondSourceOffset); + Mul(firstProductReg, firstMatrixReg, firstScaleReg, rowMask); + Mul(secondProductReg, secondMatrixReg, secondScaleReg, rowMask); + Add(firstSumReg, firstSumReg, firstProductReg, rowMask); + Add(secondSumReg, secondSumReg, secondProductReg, rowMask); + } + } + if (row < firstValid) { + Add(firstCurrentReg, firstCurrentReg, firstSumReg, rowMask); + StoreAlign(matrix + firstCurrentOffset, firstCurrentReg, rowMask); + } + if (row < secondValid) { + Add(secondCurrentReg, secondCurrentReg, secondSumReg, rowMask); + StoreAlign(matrix + secondCurrentOffset, secondCurrentReg, rowMask); + } + LocalMemBar(); + } + + RegTensor indexReg; + Arange(indexReg, 0); + for (uint16_t row = 0; row < DIAG_SIZE; ++row) { + MaskReg diagMask; + CompareScalar(diagMask, indexReg, static_cast(row), rowMask); + uint32_t firstOffset = static_cast(row) * rowStride + colBegin; + uint32_t secondOffset = static_cast(SECOND_LOCAL_ROW + row) * rowStride + colBegin + DIAG_SIZE; + RegTensor firstReg; + RegTensor secondReg; + LoadAlign(firstReg, matrix + firstOffset); + LoadAlign(secondReg, matrix + secondOffset); + Adds(firstReg, firstReg, 1.0f, diagMask); + Adds(secondReg, secondReg, 1.0f, diagMask); + StoreAlign(matrix + firstOffset, firstReg, rowMask); + StoreAlign(matrix + secondOffset, secondReg, rowMask); + } +} + +static __simd_vf__ inline void ApplyKdaRowScaleRegbase(__ubuf__ float *matrix, __ubuf__ float *rowScale, uint16_t rows, + uint16_t cols) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t FP32_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(float); + RegTensor matrixReg0; + RegTensor matrixReg1; + RegTensor scaleReg0; + RegTensor scaleReg1; + + uint16_t row = 0; + for (; row + 1 < rows; row += 2) { + LoadAlign(scaleReg0, rowScale + row); + LoadAlign(scaleReg1, rowScale + row + 1); + for (uint16_t col = 0; col < cols; col += FP32_PER_REG) { + uint32_t activeCount0 = static_cast(cols - col); + uint32_t activeCount1 = activeCount0; + MaskReg mask0 = UpdateMask(activeCount0); + MaskReg mask1 = UpdateMask(activeCount1); + uint32_t offset0 = static_cast(row) * cols + col; + uint32_t offset1 = static_cast(row + 1) * cols + col; + LoadAlign(matrixReg0, matrix + offset0); + LoadAlign(matrixReg1, matrix + offset1); + Mul(matrixReg0, matrixReg0, scaleReg0, mask0); + Mul(matrixReg1, matrixReg1, scaleReg1, mask1); + StoreAlign(matrix + offset0, matrixReg0, mask0); + StoreAlign(matrix + offset1, matrixReg1, mask1); + } + } + if (row < rows) { + LoadAlign(scaleReg0, rowScale + row); + for (uint16_t col = 0; col < cols; col += FP32_PER_REG) { + uint32_t activeCount = static_cast(cols - col); + MaskReg mask = UpdateMask(activeCount); + uint32_t offset = static_cast(row) * cols + col; + LoadAlign(matrixReg0, matrix + offset); + Mul(matrixReg0, matrixReg0, scaleReg0, mask); + StoreAlign(matrix + offset, matrixReg0, mask); + } + } +} +#endif + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +using KdaArchTag = Catlass::Arch::Ascend950; +#else +using KdaArchTag = Catlass::Arch::AtlasA2; +#endif +using KdaDispatchPolicy = Catlass::Gemm::MmadPingpong; +using KdaScoreDispatchPolicy = Catlass::Gemm::MmadPingpongTlaMulti; +static_assert(KdaScoreDispatchPolicy::ENABLE_L1_RESIDENT, + "KDA Aqk/Akk score MMAD must keep the shared right matrix resident in L1"); +static_assert(KdaScoreDispatchPolicy::L1B_STAGES == 1, + "KDA Aqk/Akk score MMAD needs one L1 B slot so the second MMAD reuses it"); +using KdaSolveDispatchPolicy = Catlass::Gemm::MmadPingpong; +static_assert(!KdaSolveDispatchPolicy::USE_HF32_MODE, "KDA triangular solve must use IEEE FP32 Cube mode"); +using KdaL1TileShape = tla::Shape; +using KdaL0TileShape = KdaL1TileShape; +using KdaSolveL1TileShape = tla::Shape; +using KdaSolveL0TileShape = KdaSolveL1TileShape; + +__aicore__ inline uint32_t FloatToBits(float value) +{ + union Bits { + __aicore__ Bits() {} + float f; + uint32_t u; + } bits; + bits.f = value; + return bits.u; +} + +__aicore__ inline float BitsToFloat(uint32_t value) +{ + union Bits { + __aicore__ Bits() {} + uint32_t u; + float f; + } bits; + bits.u = value; + return bits.f; +} + +__aicore__ inline uint16_t Bf16ToBits(bfloat16_t value) +{ + union Bits { + __aicore__ Bits() {} + bfloat16_t f; + uint16_t u; + } bits; + bits.f = value; + return bits.u; +} + +__aicore__ inline bfloat16_t BitsToBf16(uint16_t value) +{ + union Bits { + __aicore__ Bits() {} + uint16_t u; + bfloat16_t f; + } bits; + bits.u = value; + return bits.f; +} + +template +__aicore__ inline T FloatToType(float value) +{ + if constexpr (IsSameType::value) { + uint32_t bits = FloatToBits(value); + uint32_t bias = 0x7FFFu + ((bits >> 16) & 1u); + return BitsToBf16(static_cast((bits + bias) >> 16)); + } + return static_cast(value); +} + +template +class ChunkKdaFwdPrepareKernel { +public: + using OUT_T = T; + using AKK_T = float; + using SCORE_T = std::conditional_t::value, bfloat16_t, T>; + template + __aicore__ inline void Init(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR rawG, GM_ADDR aLog, GM_ADDR dtBias, + GM_ADDR beta, GM_ADDR initialState, GM_ADDR cuSeqlens, GM_ADDR chunkIndices, + GM_ADDR preparedQG, GM_ADDR preparedAqk, GM_ADDR propagatedVNew, GM_ADDR propagatedH, + GM_ADDR o, GM_ADDR finalState, GM_ADDR aqk, GM_ADDR akk, GM_ADDR w, GM_ADDR u, + GM_ADDR qg, GM_ADDR kg, GM_ADDR vNew, GM_ADDR h, GM_ADDR finalKg, GM_ADDR workspace, + const TilingData &tiling, TPipe *pipe, bool initVecBuffers = true, bool storeQG = true) + { + pipe_ = pipe; + q_.SetGlobalBuffer((__gm__ T *)q); + k_.SetGlobalBuffer((__gm__ T *)k); + v_.SetGlobalBuffer((__gm__ T *)v); + gk_.SetGlobalBuffer((__gm__ GK_T *)gk); + rawG_.SetGlobalBuffer((__gm__ float *)rawG); + aLog_.SetGlobalBuffer((__gm__ float *)aLog); + dtBias_.SetGlobalBuffer((__gm__ float *)dtBias); + betaFp32_.SetGlobalBuffer((__gm__ float *)beta); + betaBf16_.SetGlobalBuffer((__gm__ bfloat16_t *)beta); + if (initialState != nullptr) { + initialState_.SetGlobalBuffer((__gm__ float *)initialState); + } + cuSeqlensAddr_ = reinterpret_cast<__gm__ int64_t *>(cuSeqlens); + if (preparedQG != nullptr) { + preparedQG_.SetGlobalBuffer((__gm__ T *)preparedQG); + } + if (preparedAqk != nullptr) { + preparedAqk_.SetGlobalBuffer((__gm__ T *)preparedAqk); + } + if (propagatedVNew != nullptr) { + propagatedVNew_.SetGlobalBuffer((__gm__ T *)propagatedVNew); + } + if (propagatedH != nullptr) { + propagatedH_.SetGlobalBuffer((__gm__ T *)propagatedH); + } + chunkIndicesAddr_ = reinterpret_cast<__gm__ int64_t *>(chunkIndices); + o_.SetGlobalBuffer((__gm__ OUT_T *)o); + finalState_.SetGlobalBuffer((__gm__ float *)finalState); + aqk_.SetGlobalBuffer((__gm__ float *)aqk); + akk_.SetGlobalBuffer((__gm__ AKK_T *)akk); + w_.SetGlobalBuffer((__gm__ T *)w); + u_.SetGlobalBuffer((__gm__ OUT_T *)u); + qg_.SetGlobalBuffer((__gm__ T *)qg); + kg_.SetGlobalBuffer((__gm__ T *)kg); + vNew_.SetGlobalBuffer((__gm__ T *)vNew); + h_.SetGlobalBuffer((__gm__ float *)h); + finalKg_.SetGlobalBuffer((__gm__ T *)finalKg); + solveWorkspace_.SetGlobalBuffer((__gm__ float *)workspace); + + B_ = tiling.batch; + N_ = tiling.seqNum; + H_ = tiling.qHeadNum; + HV_ = tiling.vHeadNum; + T_ = tiling.seqlen; + K_ = COMPILE_K == 0 ? tiling.kHeadDim : COMPILE_K; + V_ = COMPILE_V == 0 ? tiling.vHeadDim : COMPILE_V; + BT_ = COMPILE_BT == 0 ? tiling.chunkSize : COMPILE_BT; + NT_ = tiling.totalChunks; + scale_ = tiling.scale; + hasInitial_ = tiling.hasInitialState; + isVarLen_ = tiling.isVarLen; + inputSequenceMajor_ = tiling.inputSequenceMajor; + fusePostWu_ = tiling.fusePostWu; + computeGateInPrepare_ = tiling.computeGateInPrepare; + hasALog_ = tiling.hasALog; + hasDtBias_ = tiling.hasDtBias; + lowerBound_ = tiling.lowerBound; + storeQG_ = storeQG; + usedCoreNum_ = tiling.prepareUsedCoreNum; + betaIsFp32_ = tiling.betaDataType == 0; +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (SAFE_GATE && !IsSameType::value) { + headPairMode_ = KDA_ARCH35_ENABLE_HEAD_PAIR && HV_ % KDA_SCORE_LANES == 0; + } +#endif + constexpr uint64_t solvePipelineDepth = SAFE_GATE ? KDA_SOLVE_PIPELINE_DEPTH : 1; + const uint64_t solveBytes = + usedCoreNum_ * solvePipelineDepth * KDA_SOLVE_SCRATCH_SLOTS * BT_ * BT_ * sizeof(float); + const uint64_t alignedSolveBytes = + (solveBytes + KDA_WORKSPACE_ALIGN - 1) / KDA_WORKSPACE_ALIGN * KDA_WORKSPACE_ALIGN; + scoreWorkspace_.SetGlobalBuffer((__gm__ SCORE_T *)(workspace + alignedSolveBytes)); + if ASCEND_IS_AIV { + uint64_t subBlockNum = static_cast(GetSubBlockNum()); + solveCoreIdx_ = subBlockNum == 0 ? 0 : static_cast(GetBlockIdx()) / subBlockNum; + } else { + solveCoreIdx_ = static_cast(GetBlockIdx()); + } + if (pipe_ != nullptr && initVecBuffers) { + pipe_->InitBuffer(exp2Buf_, EXP2_UB_BYTES); + pipe_->InitBuffer(vecBuf_, KDA_VEC_ARENA_ELEMENTS * sizeof(float)); + const uint64_t gateStageElems = GatePipelineRows() * K_; + const uint64_t gateInputSlotBytes = GateInputSlotBytes(); + const uint64_t gatePipelineBytes = GateBufferDepth() * (gateInputSlotBytes + gateStageElems * sizeof(T)); + pipe_->InitBuffer(gateWritebackBuf_, static_cast(gatePipelineBytes)); + AllocVectorEvents(); + } + } + __aicore__ inline void ProcessAivOnly() + { + isAivOnly_ = true; + ProcessPreAiv(); + ReleaseVectorEvents(); + } + + __aicore__ inline void ProcessAiv() + { + ProcessPreAiv(); + ReleaseVectorEvents(); + } + + __aicore__ inline void ProcessAic() { ProcessPreAic(); } + + template + __aicore__ inline void ProcessAicFused(PostWuOp &postWu) + { + ProcessPreAicHeadPairFused(postWu); + } + +private: + __aicore__ inline void AllocVectorEvents() + { + mte2ToVEvent_ = pipe_->AllocEventID(); + vToMte2Event_ = pipe_->AllocEventID(); + vToMte3Event_ = pipe_->AllocEventID(); + mte3ToVEvent_ = pipe_->AllocEventID(); + mte2ToMte3Event_ = pipe_->AllocEventID(); + vToSEvent_ = pipe_->AllocEventID(); + for (uint32_t slot = 0; slot < KDA_GATE_PIPELINE_DEPTH; ++slot) { + mte3ToMte2Events_[slot] = pipe_->AllocEventID(); + } + vectorEventsAllocated_ = true; + } + + __aicore__ inline void ReleaseVectorEvents() + { + if (!vectorEventsAllocated_) { + return; + } + pipe_->ReleaseEventID(mte2ToVEvent_); + pipe_->ReleaseEventID(vToMte2Event_); + pipe_->ReleaseEventID(vToMte3Event_); + pipe_->ReleaseEventID(mte3ToVEvent_); + pipe_->ReleaseEventID(mte2ToMte3Event_); + pipe_->ReleaseEventID(vToSEvent_); + for (uint32_t slot = 0; slot < KDA_GATE_PIPELINE_DEPTH; ++slot) { + pipe_->ReleaseEventID(mte3ToMte2Events_[slot]); + } + vectorEventsAllocated_ = false; + } + + __aicore__ inline uint64_t QOffset(uint64_t b, uint64_t h, uint64_t t, uint64_t d) const + { + if (inputSequenceMajor_) { + return ((b * T_ + t) * H_ + h) * K_ + d; + } + return ((b * H_ + h) * T_ + t) * K_ + d; + } + + __aicore__ inline uint64_t VInputOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t d) const + { + if (inputSequenceMajor_) { + return ((b * T_ + t) * HV_ + hv) * V_ + d; + } + return ((b * HV_ + hv) * T_ + t) * V_ + d; + } + + __aicore__ inline uint64_t RawGateOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t d) const + { + if (inputSequenceMajor_) { + return ((b * T_ + t) * HV_ + hv) * K_ + d; + } + return ((b * HV_ + hv) * T_ + t) * K_ + d; + } + + __aicore__ inline uint64_t KVOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t d, uint64_t dim) const + { + return ((b * HV_ + hv) * T_ + t) * dim + d; + } + + __aicore__ inline uint64_t BetaOffset(uint64_t b, uint64_t hv, uint64_t t) const { return (b * HV_ + hv) * T_ + t; } + + __aicore__ inline uint64_t AOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t j) const + { + return ((b * HV_ + hv) * T_ + t) * BT_ + j; + } + + __aicore__ inline uint64_t HOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t d, uint64_t r) const + { + return (((b * HV_ + hv) * NT_ + chunkIdx) * K_ + d) * V_ + r; + } + + __aicore__ inline uint64_t WScratchOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t t, uint64_t d) const + { + return (((b * HV_ + hv) * NT_ + chunkIdx) * BT_ + t) * K_ + d; + } + + __aicore__ inline uint64_t SolveScratchOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t slot) const + { + (void)b; + (void)hv; + (void)chunkIdx; + constexpr uint64_t solvePipelineDepth = SAFE_GATE ? KDA_SOLVE_PIPELINE_DEPTH : 1; + uint64_t matrixElements = BT_ * BT_; + return ((solveCoreIdx_ * solvePipelineDepth + activeSolveSlot_) * KDA_SOLVE_SCRATCH_SLOTS + slot) * + matrixElements; + } + + __aicore__ inline uint64_t ScoreScratchOffset(uint64_t slot, uint64_t plane, uint64_t t = 0, uint64_t d = 0) const + { + return (((solveCoreIdx_ * KDA_SCORE_SCRATCH_SLOTS + slot) * KDA_SCORE_SCRATCH_PLANES + plane) * BT_ + t) * K_ + + d; + } + + __aicore__ inline uint64_t ScoreScratchSlot(uint64_t queueSlot, uint64_t lane, bool pairHeads) const + { + return pairHeads ? queueSlot * KDA_SCORE_LANES + lane : queueSlot; + } + + __aicore__ inline uint64_t ScoreRefBlockSize() const + { + if constexpr (SAFE_GATE) { + return KDA_SAFE_SCORE_REF_BC; + } + return KDA_SCORE_REF_BC; + } + + __aicore__ inline uint64_t ScoreRowBlockCount(uint64_t curT, uint64_t rowBegin) const + { + uint64_t blockSize = ScoreRefBlockSize(); + uint64_t rowCount = curT - rowBegin; + if (rowCount > blockSize) { + rowCount = blockSize; + } + return rowCount; + } + + __aicore__ inline uint64_t ScoreRefToken(uint64_t start, uint64_t curT, uint64_t rowBegin, uint64_t rowCount) const + { + uint64_t ref = rowBegin + rowCount / 2; + if (ref >= curT) { + ref = curT - 1; + } + return start + ref; + } + + __aicore__ inline void RunExp2(LocalTensor &tensor, uint32_t count) + { + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + ClampExpInput(tensor, count); + Exp(tensor, tensor, count); + PipeBarrier(); + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + } + + __aicore__ inline void ClampExpInput(LocalTensor &tensor, uint32_t count) + { + Mins(tensor, tensor, KDA_EXP_INPUT_MAX, count); + PipeBarrier(); + Maxs(tensor, tensor, KDA_EXP_INPUT_MIN, count); + PipeBarrier(); + } + + __aicore__ inline void ClampScoreExpInput(LocalTensor &tensor, uint32_t count) + { + constexpr float expInputMax = + IsSameType::value ? KDA_SCORE_EXP_INPUT_MAX : KDA_EXP_INPUT_MAX; + constexpr float expInputMin = + IsSameType::value ? KDA_SCORE_EXP_INPUT_MIN : KDA_EXP_INPUT_MIN; + Mins(tensor, tensor, expInputMax, count); + PipeBarrier(); + Maxs(tensor, tensor, expInputMin, count); + PipeBarrier(); + } + + template + __aicore__ inline void ClampFp32ForCast(LocalTensor &tensor, uint32_t count) + { + if constexpr (IsSameType::value) { + Mins(tensor, tensor, KDA_FP16_MAX, count); + PipeBarrier(); + Maxs(tensor, tensor, -KDA_FP16_MAX, count); + PipeBarrier(); + } + } + + __aicore__ inline void ClampFp32ToOutputType(LocalTensor &tensor, uint32_t count) + { + ClampFp32ForCast(tensor, count); + } + + template + __aicore__ inline void CopyVectorIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset, + uint64_t count) + { + uint64_t rowBytes = count * static_cast(sizeof(CopyT)); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst, src[offset], static_cast(count)); + return; + } + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPadParams padParams{false, 0, 0, 0}; + DataCopyPad(dst, src[offset], params, padParams); + } + + template + __aicore__ inline void CopyRowsIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset, uint64_t rows, + uint64_t cols, uint64_t rowStride) + { + if (rows == 0 || cols == 0) { + return; + } + if (rowStride == cols) { + CopyVectorIn(dst, src, offset, rows * cols); + return; + } + DataCopyExtParams params{static_cast(rows), static_cast(cols * sizeof(CopyT)), + static_cast((rowStride - cols) * sizeof(CopyT)), 0, 0}; + DataCopyPadExtParams padParams{false, 0, 0, 0}; + DataCopyPad(dst, src[offset], params, padParams); + } + + template + __aicore__ inline void CopyVectorOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src, + uint64_t count) + { + uint64_t rowBytes = count * static_cast(sizeof(CopyT)); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst[offset], src, static_cast(count)); + return; + } + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPad(dst[offset], src, params); + } + + template + __aicore__ inline void CopyRowIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset) + { + CopyVectorIn(dst, src, offset, K_); + } + + template + __aicore__ inline void CopyRowOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src) + { + CopyVectorOut(dst, offset, src, K_); + } + + __aicore__ inline LocalTensor VecScratch(uint64_t slot) + { + return vecBuf_.Get()[slot * EXP2_UB_ELEMENTS]; + } + + __aicore__ inline uint64_t GateStageElems() const { return GatePipelineRows() * K_; } + + __aicore__ inline uint64_t GatePipelineRows() const + { + constexpr uint64_t fixedBytes = static_cast(KDA_VEC_ARENA_ELEMENTS) * sizeof(float) + EXP2_UB_BYTES; + constexpr uint64_t availableBytes = KDA_AIV_UB_BUDGET_BYTES - fixedBytes; + uint64_t bytesPerRow = K_ * KDA_GATE_PIPELINE_DEPTH * (3 * sizeof(T) + sizeof(GK_T)); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128) { + bytesPerRow = GateBufferDepth() * (K_ * (4 * sizeof(T) + sizeof(GK_T)) + KDA_BETA_STAGE_BYTES); + } +#endif + uint64_t rows = bytesPerRow == 0 ? 0 : availableBytes / bytesPerRow; + return rows < KDA_GATE_TILE_ROWS ? rows : KDA_GATE_TILE_ROWS; + } + + __aicore__ inline constexpr uint64_t GateBufferDepth() const + { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128) { + return 2; + } +#endif + return KDA_GATE_PIPELINE_DEPTH; + } + + __aicore__ inline uint64_t GateInputSlotBytes() const + { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128) { + return GateStageElems() * (3 * sizeof(T) + sizeof(GK_T)) + GatePipelineRows() * KDA_BETA_STAGE_BYTES; + } +#endif + return GateStageElems() * (2 * sizeof(T) + sizeof(GK_T)); + } + + __aicore__ inline LocalTensor GateQTyped(uint64_t slot) + { + uint64_t byteOffset = slot * GateInputSlotBytes(); + return gateWritebackBuf_.Get()[byteOffset / sizeof(T)]; + } + + __aicore__ inline LocalTensor GateKTyped(uint64_t slot) + { + uint64_t byteOffset = slot * GateInputSlotBytes() + GateStageElems() * sizeof(T); + return gateWritebackBuf_.Get()[byteOffset / sizeof(T)]; + } + + __aicore__ inline LocalTensor GateGTyped(uint64_t slot) + { + uint64_t byteOffset = slot * GateInputSlotBytes() + 2 * GateStageElems() * sizeof(T); + return gateWritebackBuf_.Get()[byteOffset / sizeof(GK_T)]; + } + + __aicore__ inline LocalTensor GateVTyped(uint64_t slot) + { + uint64_t byteOffset = slot * GateInputSlotBytes() + GateStageElems() * (2 * sizeof(T) + sizeof(GK_T)); + return gateWritebackBuf_.Get()[byteOffset / sizeof(T)]; + } + + template + __aicore__ inline LocalTensor GateBetaStaging(uint64_t slot) + { + uint64_t byteOffset = slot * GateInputSlotBytes() + GateStageElems() * (3 * sizeof(T) + sizeof(GK_T)); + return gateWritebackBuf_.Get()[byteOffset / sizeof(BetaT)]; + } + + __aicore__ inline LocalTensor GateKgTyped(uint64_t slot) + { + uint64_t byteOffset = GateBufferDepth() * GateInputSlotBytes() + slot * GateStageElems() * sizeof(T); + return gateWritebackBuf_.Get()[byteOffset / sizeof(T)]; + } + + __aicore__ inline LocalTensor LocalGateChunk() + { + constexpr uint64_t byteOffset = static_cast(KDA_LOCAL_GK_FLOAT_OFFSET) * sizeof(float); + return vecBuf_.Get()[byteOffset / sizeof(GK_T)]; + } + + __aicore__ inline LocalTensor GateScoreTyped(uint64_t slot, uint64_t tileRow) + { + if constexpr (IsSameType::value) { + if (computeGateInPrepare_) { + return LocalGateChunk()[tileRow * K_]; + } + } + return GateGTyped(slot); + } + + __aicore__ inline void LoadGateScoreRef(LocalTensor dst, uint64_t b, uint64_t hv, uint64_t token) + { + if constexpr (IsSameType::value) { + if (computeGateInPrepare_) { + const uint64_t tileRow = token - activeGateChunkStart_; + Adds(dst, LocalGateChunk()[tileRow * K_], 0.0f, static_cast(K_)); + PipeBarrier(); + return; + } + } + LoadAsFloatRow(gk_, KVOffset(b, hv, token, 0, K_), dst, K_); + } + + __aicore__ inline void PrefetchQKGate(uint64_t slot, uint64_t b, uint64_t h, uint64_t hv, uint64_t token, + uint64_t elems) + { + const uint64_t rows = elems / K_; + LocalTensor qTyped = GateQTyped(slot); + LocalTensor kTyped = GateKTyped(slot); + LocalTensor gateTyped = GateGTyped(slot); + CopyRowsIn(qTyped, q_, QOffset(b, h, token, 0), rows, K_, inputSequenceMajor_ ? H_ * K_ : K_); + CopyRowsIn(kTyped, k_, QOffset(b, h, token, 0), rows, K_, inputSequenceMajor_ ? H_ * K_ : K_); + if (!computeGateInPrepare_) { + CopyVectorIn(gateTyped, gk_, KVOffset(b, hv, token, 0, K_), elems); + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128) { + LocalTensor vTyped = GateVTyped(slot); + CopyRowsIn(vTyped, v_, VInputOffset(b, hv, token, 0), rows, V_, inputSequenceMajor_ ? HV_ * V_ : V_); + if (betaIsFp32_) { + LocalTensor betaTyped = GateBetaStaging(slot); + CopyVectorIn(betaTyped, betaFp32_, BetaOffset(b, hv, token), rows); + } else { + LocalTensor betaTyped = GateBetaStaging(slot); + CopyVectorIn(betaTyped, betaBf16_, BetaOffset(b, hv, token), rows); + } + } +#endif + SetFlag(mte2ToVEvent_); + } + + __aicore__ inline float LoadGateExpA(uint64_t hv) + { + if (!hasALog_) { + return 1.0f; + } + LocalTensor scalar = exp2Buf_.Get(); + CopyVectorIn(scalar, aLog_, hv, 1); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Exp(scalar, scalar, 1); + PipeBarrier(); + SetFlag(vToSEvent_); + WaitFlag(vToSEvent_); + __ubuf__ float *ptr = (__ubuf__ float *)scalar.GetPhyAddr(); + return ptr[0]; + } + + __aicore__ inline void PrefetchRawGateTile(uint64_t slot, uint64_t b, uint64_t hv, uint64_t token, uint64_t rows) + { + (void)slot; + const uint64_t tileRow = token - activeGateChunkStart_; + LocalTensor gate = LocalGateChunk().template ReinterpretCast()[tileRow * K_]; + CopyRowsIn(gate, rawG_, RawGateOffset(b, hv, token, 0), rows, K_, inputSequenceMajor_ ? HV_ * K_ : K_); + SetFlag(mte2ToVEvent_); + } + + __aicore__ inline void MaterializeRawGateChunkArch35(uint64_t b, uint64_t hv, uint64_t start, uint64_t rows) + { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if (!computeGateInPrepare_) { + return; + } + if constexpr (!IsSameType::value) { + return; + } else { + activeGateChunkStart_ = start; + const float expA = LoadGateExpA(hv); + LocalTensor acc = exp2Buf_.Get(); + LocalTensor bias = exp2Buf_.Get()[K_]; + if (hasDtBias_) { + CopyVectorIn(bias, dtBias_, hv * K_, K_); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + } + Duplicate(acc, 0.0f, static_cast(K_)); + PipeBarrier(); + + const uint64_t tileRows = GatePipelineRows(); + const uint64_t tileCount = (rows + tileRows - 1) / tileRows; + uint64_t currentRows = rows < tileRows ? rows : tileRows; + PrefetchRawGateTile(0, b, hv, start, currentRows); + for (uint64_t tile = 0; tile < tileCount; ++tile) { + const uint64_t slot = tile & 1; + const uint64_t tileRow = tile * tileRows; + currentRows = rows - tileRow; + if (currentRows > tileRows) { + currentRows = tileRows; + } + WaitFlag(mte2ToVEvent_); + + const uint64_t nextTile = tile + 1; + if (nextTile < tileCount) { + const uint64_t nextSlot = nextTile & 1; + uint64_t nextRows = rows - nextTile * tileRows; + if (nextRows > tileRows) { + nextRows = tileRows; + } + PrefetchRawGateTile(nextSlot, b, hv, start + nextTile * tileRows, nextRows); + } + + LocalTensor gate = LocalGateChunk().template ReinterpretCast()[tileRow * K_]; + if (hasDtBias_) { + AccumulateRawSafeGateChunk128Regbase( + (__ubuf__ float *)gate.GetPhyAddr(), (__ubuf__ float *)bias.GetPhyAddr(), + (__ubuf__ float *)acc.GetPhyAddr(), static_cast(currentRows), expA, lowerBound_); + } else { + AccumulateRawSafeGateChunk128Regbase( + (__ubuf__ float *)gate.GetPhyAddr(), (__ubuf__ float *)bias.GetPhyAddr(), + (__ubuf__ float *)acc.GetPhyAddr(), static_cast(currentRows), expA, lowerBound_); + } + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(gk_, KVOffset(b, hv, start + tileRow, 0, K_), gate, currentRows * K_); + } + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } +#else + (void)b; + (void)hv; + (void)start; + (void)rows; +#endif + } + + __aicore__ inline LocalTensor GateDirectQ(uint64_t slot) + { + return vecBuf_.Get()[slot * 3 * GateStageElems()]; + } + + __aicore__ inline LocalTensor GateDirectW(uint64_t slot) { return GateDirectQ(slot)[GateStageElems()]; } + + __aicore__ inline LocalTensor GateDirectV(uint64_t slot) { return GateDirectQ(slot)[2 * GateStageElems()]; } + + __aicore__ inline LocalTensor GateBetaFloat(uint64_t slot) + { + constexpr uint64_t directBytes = KDA_GATE_PIPELINE_DEPTH * 3 * KDA_GATE_TILE_ROWS * COMPILE_K * sizeof(T); + return vecBuf_.Get()[directBytes / sizeof(float) + slot * KDA_GATE_TILE_ROWS]; + } + + __aicore__ inline void StorePreparedQG(uint64_t b, uint64_t hv, uint64_t token, LocalTensor directQ, + uint64_t elems) + { + const uint64_t offset = KVOffset(b, hv, token, 0, K_); + if (storeQG_) { + CopyVectorOut(qg_, offset, directQ, elems); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(kg_, offset, directQ, elems); + } + + __aicore__ inline void PrefetchKGate(uint64_t slot, uint64_t b, uint64_t h, uint64_t hv, uint64_t token, + uint64_t elems) + { + const uint64_t rows = elems / K_; + LocalTensor kTyped = GateQTyped(slot); + LocalTensor gateTyped = GateGTyped(slot); + CopyRowsIn(kTyped, k_, QOffset(b, h, token, 0), rows, K_, inputSequenceMajor_ ? H_ * K_ : K_); + if (!computeGateInPrepare_) { + CopyVectorIn(gateTyped, gk_, KVOffset(b, hv, token, 0, K_), elems); + } + SetFlag(mte2ToVEvent_); + } + + __aicore__ inline void WaitGateInputReady() { WaitFlag(mte2ToVEvent_); } + + __aicore__ inline void WaitGateOutputForMte2(uint64_t slot = 0) + { + WaitFlag(mte3ToMte2Events_[slot]); + } + + __aicore__ inline void WaitGateOutputForVector() { WaitFlag(mte3ToVEvent_); } + + __aicore__ inline void SignalGateOutputDone() + { + SetFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + } + + __aicore__ inline void SignalGateOutputDoneForMte2(uint64_t slot) + { + SetFlag(mte3ToMte2Events_[slot]); + } + + template + __aicore__ inline void LoadAsFloatRow(GlobalTensor &src, uint64_t srcOffset, LocalTensor &dst, + uint64_t count) + { + if constexpr (IsSameType::value) { + CopyVectorIn(dst, src, srcOffset, count); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Adds(dst, dst, 0.0f, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } else { + constexpr uint32_t typedOffset = EXP2_UB_ELEMENTS * sizeof(float) / sizeof(CopyT); + LocalTensor rowLocal = exp2Buf_.Get()[typedOffset]; + CopyVectorIn(rowLocal, src, srcOffset, count); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Cast(dst, rowLocal, RoundMode::CAST_NONE, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } + PipeBarrier(); + } + + // beta arrives as fp32 or bf16 (tiling.betaDataType). Both are resolved at + // runtime so the whole prepare kernel stays independent of the beta dtype -- + // templating on it doubled every Prepare instantiation for one vector cast. + __aicore__ inline void LoadBetaAsFloatRow(uint64_t srcOffset, LocalTensor &dst, uint64_t count) + { + if (betaIsFp32_) { + LoadAsFloatRow(betaFp32_, srcOffset, dst, count); + } else { + LoadAsFloatRow(betaBf16_, srcOffset, dst, count); + } + } + + __aicore__ inline void ConvertStagedBeta(uint64_t slot, LocalTensor &dst, uint64_t count) + { + if (betaIsFp32_) { + LocalTensor staged = GateBetaStaging(slot); + Adds(dst, staged, 0.0f, static_cast(count)); + } else { + LocalTensor staged = GateBetaStaging(slot); + Cast(dst, staged, RoundMode::CAST_NONE, static_cast(count)); + } + } + + template + __aicore__ inline void LoadAsFloatVector(GlobalTensor &src, uint64_t srcOffset, LocalTensor &dst, + LocalTensor &typedScratch, uint64_t count) + { + if constexpr (IsSameType::value) { + CopyVectorIn(dst, src, srcOffset, count); + } else { + CopyVectorIn(typedScratch, src, srcOffset, count); + } + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + if constexpr (!IsSameType::value) { + Cast(dst, typedScratch, RoundMode::CAST_NONE, static_cast(count)); + PipeBarrier(); + } + } + + template + __aicore__ inline void StoreFloatRow(GlobalTensor &dst, uint64_t dstOffset, LocalTensor &src, + uint64_t count) + { + if constexpr (IsSameType::value) { + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(dst, dstOffset, src, count); + } else { + constexpr uint32_t typedOffset = EXP2_UB_ELEMENTS * sizeof(float) / sizeof(CopyT); + LocalTensor rowLocal = exp2Buf_.Get()[typedOffset]; + Cast(rowLocal, src, RoundMode::CAST_RINT, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(dst, dstOffset, rowLocal, count); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline LocalTensor Exp2NegG(uint64_t b, uint64_t hv, uint64_t t) + { + LocalTensor exp2Local = exp2Buf_.Get(); + LoadAsFloatRow(gk_, KVOffset(b, hv, t, 0, K_), exp2Local, K_); + Muls(exp2Local, exp2Local, -LN2, static_cast(K_)); + PipeBarrier(); + RunExp2(exp2Local, static_cast(K_)); + return exp2Local; + } + + __aicore__ inline void PrepareScoreFactorsBulk(uint64_t b, uint64_t h, uint64_t hv, uint64_t start, + uint64_t subBlockIdx, uint64_t subBlockNum, uint64_t refToken, + uint64_t scoreRowBegin, uint64_t scoreRowCount, uint64_t validColEnd, + uint64_t finalRefToken, uint64_t scoreSlot) + { + const bool useDirectScoreL1 = +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + KDA_ARCH35_ENABLE_DIRECT_SCORE_L1 && KDA_ARCH35_ENABLE_DIRECT_SCORE_UB && SAFE_GATE && BT_ == 64 && + K_ == 128 && V_ == 128 && subBlockNum == 1 && scoreRowCount == KDA_DIRECT_SCORE_ROWS && + finalRefToken == start + BT_ - 1; +#else + false; +#endif + LocalTensor refFp32 = exp2Buf_.Get(); + LoadGateScoreRef(refFp32, b, hv, refToken); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + constexpr bool exportFinalKg = SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128; + LocalTensor finalRefFp32 = exp2Buf_.Get()[K_]; + if constexpr (exportFinalKg) { + LoadGateScoreRef(finalRefFp32, b, hv, finalRefToken); + } +#endif + + uint64_t qwBegin = scoreRowBegin + (scoreRowCount * subBlockIdx) / subBlockNum; + uint64_t qwEnd = scoreRowBegin + (scoreRowCount * (subBlockIdx + 1)) / subBlockNum; + uint64_t qwMaxRows = GatePipelineRows(); + bool qwOutputPending = false; + uint64_t qwSlot = 0; + if (qwBegin < qwEnd && qwMaxRows > 0) { + uint64_t firstRows = qwEnd - qwBegin; + if (firstRows > qwMaxRows) { + firstRows = qwMaxRows; + } + PrefetchQKGate(qwSlot, b, h, hv, start + qwBegin, firstRows * K_); + } + for (uint64_t tileRow = qwBegin; tileRow < qwEnd && qwMaxRows > 0; tileRow += qwMaxRows) { + uint64_t tileRows = qwEnd - tileRow; + if (tileRows > qwMaxRows) { + tileRows = qwMaxRows; + } + uint64_t elems = tileRows * K_; + LocalTensor qTyped = GateQTyped(qwSlot); + LocalTensor kTyped = GateKTyped(qwSlot); + LocalTensor qScore = qTyped.template ReinterpretCast(); + LocalTensor kScore = kTyped.template ReinterpretCast(); + LocalTensor gateTyped = GateScoreTyped(qwSlot, tileRow); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + LocalTensor kgScore = GateKgTyped(qwSlot).template ReinterpretCast(); + LocalTensor vTyped = GateVTyped(qwSlot); + LocalTensor directQ = GateDirectQ(qwSlot); + LocalTensor directW = GateDirectW(qwSlot); + LocalTensor directV = GateDirectV(qwSlot); + LocalTensor betaFp32 = GateBetaFloat(qwSlot); +#endif +#if !defined(__CCE_AICORE__) || __CCE_AICORE__ != 310 + LocalTensor arena = vecBuf_.Get(); + LocalTensor qFp32 = arena; + LocalTensor kFp32 = arena[elems]; + LocalTensor gFp32 = arena[2 * elems]; + LocalTensor expFp32 = arena[3 * elems]; + LocalTensor outFp32 = arena[4 * elems]; +#endif + + WaitGateInputReady(); +#if !defined(__CCE_AICORE__) || __CCE_AICORE__ != 310 + Cast(qFp32, qTyped, RoundMode::CAST_NONE, static_cast(elems)); + Cast(kFp32, kTyped, RoundMode::CAST_NONE, static_cast(elems)); + if constexpr (IsSameType::value) { + gFp32 = gateTyped; + } else { + Cast(gFp32, gateTyped, RoundMode::CAST_NONE, static_cast(elems)); + } +#endif +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + ConvertStagedBeta(qwSlot, betaFp32, tileRows); + PipeBarrier(); +#endif + if (qwOutputPending) { + WaitGateOutputForMte2(); + } + uint64_t nextTileRow = tileRow + qwMaxRows; + if (nextTileRow < qwEnd) { + uint64_t nextRows = qwEnd - nextTileRow; + if (nextRows > qwMaxRows) { + nextRows = qwMaxRows; + } + PrefetchQKGate(qwSlot ^ 1, b, h, hv, start + nextTileRow, nextRows * K_); + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + bool fuseQwKg = SAFE_GATE && BT_ == 64 && K_ == 128 && V_ == 128 && subBlockNum == 1; + if (fuseQwKg) { + PrepareKdaGateQwKgRegbase( + (__ubuf__ T *)reinterpret_cast(qTyped.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(kTyped.GetPhyAddr()), + (__ubuf__ SCORE_T *)reinterpret_cast(qScore.GetPhyAddr()), + (__ubuf__ SCORE_T *)reinterpret_cast(kScore.GetPhyAddr()), + (__ubuf__ SCORE_T *)reinterpret_cast(kgScore.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(directQ.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(directW.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(vTyped.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(directV.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(vTyped.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(betaFp32.GetPhyAddr()), + (__ubuf__ GK_T *)reinterpret_cast(gateTyped.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(refFp32.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(finalRefFp32.GetPhyAddr()), + static_cast(tileRows), static_cast(K_), static_cast(tileRows)); + } else { + PrepareKdaGateQwRegbase( + (__ubuf__ T *)reinterpret_cast(qTyped.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(kTyped.GetPhyAddr()), + (__ubuf__ SCORE_T *)reinterpret_cast(qScore.GetPhyAddr()), + (__ubuf__ SCORE_T *)reinterpret_cast(kScore.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(directQ.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(directW.GetPhyAddr()), + (__ubuf__ GK_T *)reinterpret_cast(gateTyped.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(refFp32.GetPhyAddr()), static_cast(tileRows), + static_cast(K_)); + } +#else + PipeBarrier(); + for (uint64_t row = 0; row < tileRows; ++row) { + Sub(expFp32[row * K_], gFp32[row * K_], refFp32, static_cast(K_)); + } + PipeBarrier(); + Muls(expFp32, expFp32, LN2, static_cast(elems)); + PipeBarrier(); + ClampScoreExpInput(expFp32, static_cast(elems)); + Exp(expFp32, expFp32, static_cast(elems)); + PipeBarrier(); + + Mul(outFp32, qFp32, expFp32, static_cast(elems)); + PipeBarrier(); + ClampFp32ForCast(outFp32, static_cast(elems)); + Cast(qScore, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + PipeBarrier(); + Mul(outFp32, kFp32, expFp32, static_cast(elems)); + PipeBarrier(); + ClampFp32ForCast(outFp32, static_cast(elems)); + Cast(kScore, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + PipeBarrier(); +#endif + + if (qwOutputPending) { + WaitGateOutputForVector(); + } + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + if (useDirectScoreL1) { + Catlass::Arch::Resource resource; + LocalTensor scoreL1 = + resource.l1Buf.template GetBufferByByte(scoreSlot * KDA_DIRECT_SCORE_L1_SLOT_BYTES); + const uint64_t localRow = tileRow - scoreRowBegin; + constexpr uint32_t c0Elements = 16; + constexpr uint32_t c0Blocks = 128 / c0Elements; + constexpr uint32_t rowBlockElements = 16 * c0Elements; + LocalTensor nzScratch = resource.ubBuf.template GetBufferByByte( + KDA_DIRECT_SCORE_UB_BYTE_OFFSET + + (scoreSlot / KDA_SCORE_LANES) * KDA_DIRECT_SCORE_SLOT_ELEMENTS * sizeof(float) + + 2 * KDA_DIRECT_SCORE_MATRIX_ELEMENTS * sizeof(float)); + LocalTensor qNz = nzScratch; + LocalTensor kNz = qNz[tileRows * K_]; + constexpr uint8_t srcRepeatStride = 128 * sizeof(SCORE_T) / 32; + for (uint32_t colBlock = 0; colBlock < c0Blocks; ++colBlock) { + const uint64_t srcOffset = colBlock * c0Elements; + const uint64_t dstOffset = colBlock * tileRows * c0Elements; + Copy(qNz[dstOffset], qScore[srcOffset], c0Elements, static_cast(tileRows), + {1, 1, 1, srcRepeatStride}); + Copy(kNz[dstOffset], kScore[srcOffset], c0Elements, static_cast(tileRows), + {1, 1, 1, srcRepeatStride}); + } + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + const uint64_t qL1Offset = (localRow / 16) * rowBlockElements; + const uint64_t kL1Offset = ((KDA_DIRECT_SCORE_ROWS + localRow) / 16) * rowBlockElements; + DataCopyParams nzCopyParams{c0Blocks, static_cast(tileRows), 0, + static_cast(64 - tileRows)}; + DataCopy(scoreL1[qL1Offset], qNz, nzCopyParams); + DataCopy(scoreL1[kL1Offset], kNz, nzCopyParams); + } else { + CopyVectorOut(scoreWorkspace_, ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_QG, tileRow), qScore, + elems); + CopyVectorOut(scoreWorkspace_, ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_W, tileRow), kScore, + elems); + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if (fuseQwKg) { + CopyVectorOut(scoreWorkspace_, ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_KG, tileRow), kgScore, + elems); + } + StorePreparedQG(b, hv, start + tileRow, directQ, elems); + CopyVectorOut(w_, KVOffset(b, hv, start + tileRow, 0, K_), directW, elems); + CopyVectorOut(vNew_, KVOffset(b, hv, start + tileRow, 0, V_), directV, tileRows * V_); + if constexpr (exportFinalKg) { + CopyVectorOut(finalKg_, KVOffset(b, hv, start + tileRow, 0, K_), vTyped, elems); + } +#endif + SignalGateOutputDone(); + qwOutputPending = true; + qwSlot ^= 1; + } + if (qwOutputPending) { + WaitGateOutputForMte2(); + WaitGateOutputForVector(); + } + + bool fuseQwKg = false; +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + fuseQwKg = SAFE_GATE && BT_ == 64 && K_ == 128 && V_ == 128 && subBlockNum == 1; +#endif + uint64_t kgRows = fuseQwKg ? scoreRowBegin : validColEnd; + uint64_t kgBegin = (kgRows * subBlockIdx) / subBlockNum; + uint64_t kgEnd = (kgRows * (subBlockIdx + 1)) / subBlockNum; + uint64_t kgMaxRows = GatePipelineRows(); + bool kgOutputPending = false; + uint64_t kgSlot = 0; + if (kgBegin < kgEnd && kgMaxRows > 0) { + uint64_t firstRows = kgEnd - kgBegin; + if (firstRows > kgMaxRows) { + firstRows = kgMaxRows; + } + PrefetchKGate(kgSlot, b, h, hv, start + kgBegin, firstRows * K_); + } + for (uint64_t tileRow = kgBegin; tileRow < kgEnd && kgMaxRows > 0; tileRow += kgMaxRows) { + uint64_t tileRows = kgEnd - tileRow; + if (tileRows > kgMaxRows) { + tileRows = kgMaxRows; + } + uint64_t elems = tileRows * K_; + LocalTensor kTyped = GateQTyped(kgSlot); + LocalTensor kgScore = kTyped.template ReinterpretCast(); + LocalTensor gateTyped = GateScoreTyped(kgSlot, tileRow); +#if !defined(__CCE_AICORE__) || __CCE_AICORE__ != 310 + LocalTensor arena = vecBuf_.Get(); + LocalTensor kFp32 = arena; + LocalTensor gFp32 = arena[elems]; + LocalTensor expFp32 = arena[2 * elems]; + LocalTensor outFp32 = arena[3 * elems]; +#endif + + WaitGateInputReady(); +#if !defined(__CCE_AICORE__) || __CCE_AICORE__ != 310 + Cast(kFp32, kTyped, RoundMode::CAST_NONE, static_cast(elems)); + if constexpr (IsSameType::value) { + gFp32 = gateTyped; + } else { + Cast(gFp32, gateTyped, RoundMode::CAST_NONE, static_cast(elems)); + } +#endif + if (kgOutputPending) { + WaitGateOutputForMte2(); + } + uint64_t nextTileRow = tileRow + kgMaxRows; + if (nextTileRow < kgEnd) { + uint64_t nextRows = kgEnd - nextTileRow; + if (nextRows > kgMaxRows) { + nextRows = kgMaxRows; + } + PrefetchKGate(kgSlot ^ 1, b, h, hv, start + nextTileRow, nextRows * K_); + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + PrepareKdaGateKgRegbase( + (__ubuf__ SCORE_T *)reinterpret_cast(kgScore.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(kTyped.GetPhyAddr()), + (__ubuf__ GK_T *)reinterpret_cast(gateTyped.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(refFp32.GetPhyAddr()), static_cast(tileRows), + static_cast(K_), static_cast(tileRows)); +#else + PipeBarrier(); + for (uint64_t row = 0; row < tileRows; ++row) { + Sub(expFp32[row * K_], refFp32, gFp32[row * K_], static_cast(K_)); + } + PipeBarrier(); + Muls(expFp32, expFp32, LN2, static_cast(elems)); + PipeBarrier(); + ClampScoreExpInput(expFp32, static_cast(elems)); + Exp(expFp32, expFp32, static_cast(elems)); + PipeBarrier(); + Mul(outFp32, kFp32, expFp32, static_cast(elems)); + PipeBarrier(); + ClampFp32ForCast(outFp32, static_cast(elems)); + Cast(kgScore, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + PipeBarrier(); +#endif + + if (kgOutputPending) { + WaitGateOutputForVector(); + } + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(scoreWorkspace_, ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_KG, tileRow), kgScore, + elems); + SignalGateOutputDone(); + kgOutputPending = true; + kgSlot ^= 1; + } + if (kgOutputPending) { + WaitGateOutputForMte2(); + WaitGateOutputForVector(); + } + } + + __aicore__ inline void PrepareGateProductsBulk(uint64_t b, uint64_t h, uint64_t hv, uint64_t start, uint64_t curT, + uint64_t subBlockIdx, uint64_t subBlockNum, bool useRef, + uint64_t refToken, uint64_t validColEnd, bool writeScoreScratch, + uint64_t scoreSlot) + { + if constexpr (IsSameType::value) { + return; + } + if (subBlockNum == 0 || subBlockIdx >= subBlockNum || K_ == 0) { + return; + } + uint64_t rowBegin = (curT * subBlockIdx) / subBlockNum; + uint64_t rowEnd = (curT * (subBlockIdx + 1)) / subBlockNum; + if (rowBegin >= rowEnd) { + return; + } + + uint64_t maxRows = GatePipelineRows(); + if (maxRows == 0) { + return; + } + LocalTensor refFp32 = exp2Buf_.Get(); + if (useRef) { + LoadGateScoreRef(refFp32, b, hv, refToken); + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + const bool fuseScoreWriteback = writeScoreScratch && useRef && K_ * 2 <= EXP2_UB_ELEMENTS; +#endif + + bool outputPending = false; + uint64_t gateSlot = 0; + uint64_t firstRows = rowEnd - rowBegin; + if (firstRows > maxRows) { + firstRows = maxRows; + } + PrefetchQKGate(gateSlot, b, h, hv, start + rowBegin, firstRows * K_); + for (uint64_t tileRow = rowBegin; tileRow < rowEnd; tileRow += maxRows) { + uint64_t tileRows = rowEnd - tileRow; + if (tileRows > maxRows) { + tileRows = maxRows; + } + uint64_t elems = tileRows * K_; + LocalTensor qTyped = GateQTyped(gateSlot); + LocalTensor kTyped = GateKTyped(gateSlot); + LocalTensor kgTyped = GateKgTyped(gateSlot); + LocalTensor qScore = qTyped.template ReinterpretCast(); + LocalTensor wScore = kTyped.template ReinterpretCast(); + LocalTensor kgScore = kgTyped.template ReinterpretCast(); + LocalTensor gateTyped = GateScoreTyped(gateSlot, tileRow); +#if !defined(__CCE_AICORE__) || __CCE_AICORE__ != 310 + LocalTensor arena = vecBuf_.Get(); + LocalTensor qFp32 = arena; + LocalTensor kFp32 = arena[elems]; + LocalTensor gFp32 = arena[2 * elems]; + LocalTensor expFp32 = arena[3 * elems]; + LocalTensor outFp32 = arena[4 * elems]; +#else + LocalTensor vTyped = GateVTyped(gateSlot); + LocalTensor directQ = GateDirectQ(gateSlot); + LocalTensor directW = GateDirectW(gateSlot); + LocalTensor directV = GateDirectV(gateSlot); + LocalTensor betaFp32 = GateBetaFloat(gateSlot); +#endif + + uint64_t token = start + tileRow; + WaitGateInputReady(); +#if !defined(__CCE_AICORE__) || __CCE_AICORE__ != 310 + Cast(qFp32, qTyped, RoundMode::CAST_NONE, static_cast(elems)); + Cast(kFp32, kTyped, RoundMode::CAST_NONE, static_cast(elems)); + if constexpr (IsSameType::value) { + gFp32 = gateTyped; + } else { + Cast(gFp32, gateTyped, RoundMode::CAST_NONE, static_cast(elems)); + } +#endif +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + ConvertStagedBeta(gateSlot, betaFp32, tileRows); + PipeBarrier(); +#endif + uint64_t nextTileRow = tileRow + maxRows; +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + uint64_t nextGateSlot = (gateSlot + 1) % KDA_GATE_PIPELINE_DEPTH; + if (nextTileRow < rowEnd) { + uint64_t tileIndex = (tileRow - rowBegin) / maxRows; + if (tileIndex + 1 >= KDA_GATE_PIPELINE_DEPTH) { + WaitGateOutputForMte2(nextGateSlot); + } + uint64_t nextRows = rowEnd - nextTileRow; + if (nextRows > maxRows) { + nextRows = maxRows; + } + PrefetchQKGate(nextGateSlot, b, h, hv, start + nextTileRow, nextRows * K_); + } +#else + if (outputPending) { + WaitGateOutputForMte2(); + } + if (nextTileRow < rowEnd) { + uint64_t nextRows = rowEnd - nextTileRow; + if (nextRows > maxRows) { + nextRows = maxRows; + } + PrefetchQKGate(gateSlot ^ 1, b, h, hv, start + nextTileRow, nextRows * K_); + } +#endif +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + uint16_t validRows = static_cast(tileRows); + if (useRef && tileRow >= validColEnd) { + validRows = 0; + } else if (useRef && tileRow + tileRows > validColEnd) { + validRows = static_cast(validColEnd - tileRow); + } + if (writeScoreScratch) { + if (fuseScoreWriteback) { + PrepareKdaGateQwKgRegbase( + (__ubuf__ T *)reinterpret_cast(qTyped.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(kTyped.GetPhyAddr()), + (__ubuf__ SCORE_T *)reinterpret_cast(qScore.GetPhyAddr()), + (__ubuf__ SCORE_T *)reinterpret_cast(wScore.GetPhyAddr()), + (__ubuf__ SCORE_T *)reinterpret_cast(kgScore.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(directQ.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(directW.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(vTyped.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(directV.GetPhyAddr()), nullptr, + (__ubuf__ float *)reinterpret_cast(betaFp32.GetPhyAddr()), + (__ubuf__ GK_T *)reinterpret_cast(gateTyped.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(refFp32.GetPhyAddr()), nullptr, + static_cast(tileRows), static_cast(K_), validRows); + } else { + PrepareKdaGateQwKgRegbase( + (__ubuf__ T *)reinterpret_cast(qTyped.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(kTyped.GetPhyAddr()), + (__ubuf__ SCORE_T *)reinterpret_cast(qScore.GetPhyAddr()), + (__ubuf__ SCORE_T *)reinterpret_cast(wScore.GetPhyAddr()), + (__ubuf__ SCORE_T *)reinterpret_cast(kgScore.GetPhyAddr()), nullptr, nullptr, nullptr, + nullptr, nullptr, nullptr, (__ubuf__ GK_T *)reinterpret_cast(gateTyped.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(refFp32.GetPhyAddr()), nullptr, + static_cast(tileRows), static_cast(K_), validRows); + } + } else if (useRef) { + PrepareKdaGateQwKgRegbase( + (__ubuf__ T *)reinterpret_cast(qTyped.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(kTyped.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(qTyped.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(kTyped.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(kgTyped.GetPhyAddr()), nullptr, nullptr, nullptr, nullptr, + nullptr, nullptr, (__ubuf__ GK_T *)reinterpret_cast(gateTyped.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(refFp32.GetPhyAddr()), nullptr, + static_cast(tileRows), static_cast(K_), validRows); + } else { + PrepareKdaGateQwKgRegbase( + (__ubuf__ T *)reinterpret_cast(qTyped.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(kTyped.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(qTyped.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(kTyped.GetPhyAddr()), + (__ubuf__ T *)reinterpret_cast(kgTyped.GetPhyAddr()), nullptr, nullptr, nullptr, nullptr, + nullptr, nullptr, (__ubuf__ GK_T *)reinterpret_cast(gateTyped.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(refFp32.GetPhyAddr()), nullptr, + static_cast(tileRows), static_cast(K_), validRows); + } +#else + PipeBarrier(); + + if (useRef) { + for (uint64_t row = 0; row < tileRows; ++row) { + Sub(expFp32[row * K_], gFp32[row * K_], refFp32, static_cast(K_)); + } + } else { + Adds(expFp32, gFp32, 0.0f, static_cast(elems)); + } + PipeBarrier(); + Muls(expFp32, expFp32, LN2, static_cast(elems)); + PipeBarrier(); + if (writeScoreScratch) { + ClampScoreExpInput(expFp32, static_cast(elems)); + } else { + ClampExpInput(expFp32, static_cast(elems)); + } + Exp(expFp32, expFp32, static_cast(elems)); + PipeBarrier(); + + Mul(outFp32, qFp32, expFp32, static_cast(elems)); + PipeBarrier(); + if (writeScoreScratch) { + ClampFp32ForCast(outFp32, static_cast(elems)); + Cast(qScore, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + } else { + ClampFp32ToOutputType(outFp32, static_cast(elems)); + Cast(qTyped, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + } + PipeBarrier(); + + Mul(outFp32, kFp32, expFp32, static_cast(elems)); + PipeBarrier(); + if (writeScoreScratch) { + ClampFp32ForCast(outFp32, static_cast(elems)); + Cast(wScore, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + } else { + ClampFp32ToOutputType(outFp32, static_cast(elems)); + Cast(kTyped, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + } + PipeBarrier(); + + if (useRef) { + for (uint64_t row = 0; row < tileRows; ++row) { + Sub(expFp32[row * K_], refFp32, gFp32[row * K_], static_cast(K_)); + } + } else { + Muls(expFp32, gFp32, -1.0f, static_cast(elems)); + } + PipeBarrier(); + Muls(expFp32, expFp32, LN2, static_cast(elems)); + PipeBarrier(); + if (writeScoreScratch) { + ClampScoreExpInput(expFp32, static_cast(elems)); + } else { + ClampExpInput(expFp32, static_cast(elems)); + } + Exp(expFp32, expFp32, static_cast(elems)); + PipeBarrier(); + Mul(outFp32, kFp32, expFp32, static_cast(elems)); + PipeBarrier(); + if (useRef && tileRow + tileRows > validColEnd) { + for (uint64_t row = 0; row < tileRows; ++row) { + if (tileRow + row >= validColEnd) { + Duplicate(outFp32[row * K_], 0.0f, static_cast(K_)); + } + } + PipeBarrier(); + } + if (writeScoreScratch) { + ClampFp32ForCast(outFp32, static_cast(elems)); + } else { + ClampFp32ToOutputType(outFp32, static_cast(elems)); + } + if (outputPending) { + WaitGateOutputForVector(); + } + if (writeScoreScratch) { + Cast(kgScore, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + } else { + Cast(kgTyped, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + } + PipeBarrier(); +#endif + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + if (writeScoreScratch) { + CopyVectorOut(scoreWorkspace_, ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_QG, tileRow), qScore, + elems); + CopyVectorOut(scoreWorkspace_, ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_W, tileRow), wScore, + elems); + CopyVectorOut(scoreWorkspace_, ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_KG, tileRow), kgScore, + elems); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if (fuseScoreWriteback) { + StorePreparedQG(b, hv, token, directQ, elems); + CopyVectorOut(w_, KVOffset(b, hv, token, 0, K_), directW, elems); + CopyVectorOut(vNew_, KVOffset(b, hv, token, 0, V_), directV, tileRows * V_); + } +#endif + } else { + CopyVectorOut(qg_, KVOffset(b, hv, token, 0, K_), qTyped, elems); + CopyVectorOut(w_, KVOffset(b, hv, token, 0, K_), kTyped, elems); + CopyVectorOut(kg_, KVOffset(b, hv, token, 0, K_), kgTyped, elems); + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + SignalGateOutputDoneForMte2(gateSlot); +#else + SignalGateOutputDone(); +#endif + outputPending = true; +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + gateSlot = (gateSlot + 1) % KDA_GATE_PIPELINE_DEPTH; +#else + gateSlot ^= 1; +#endif + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + uint64_t tileCount = (rowEnd - rowBegin + maxRows - 1) / maxRows; + uint64_t firstPending = tileCount > KDA_GATE_PIPELINE_DEPTH ? tileCount - KDA_GATE_PIPELINE_DEPTH : 0; + for (uint64_t tile = firstPending; tile < tileCount; ++tile) { + WaitGateOutputForMte2(tile % KDA_GATE_PIPELINE_DEPTH); + } +#else + if (outputPending) { + WaitGateOutputForMte2(); + WaitGateOutputForVector(); + } +#endif + return; + } + + __aicore__ inline void ZeroScoreScratchRange(uint64_t scoreSlot, uint64_t planeBegin, uint64_t planeEnd, + uint64_t firstRow, uint64_t rowEnd) + { + if (firstRow >= rowEnd) { + return; + } + const uint64_t maxRows = GatePipelineRows(); + LocalTensor zeroLocal = GateQTyped(0).template ReinterpretCast(); + for (uint64_t row = firstRow; row < rowEnd; row += maxRows) { + uint64_t rows = rowEnd - row; + if (rows > maxRows) { + rows = maxRows; + } + const uint64_t elems = rows * K_; + Duplicate(zeroLocal, static_cast(0), static_cast(elems)); + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + for (uint64_t plane = planeBegin; plane < planeEnd; ++plane) { + CopyVectorOut(scoreWorkspace_, ScoreScratchOffset(scoreSlot, plane, row), zeroLocal, elems); + } + SignalGateOutputDone(); + WaitGateOutputForMte2(); + WaitGateOutputForVector(); + } + } + + __aicore__ inline void ZeroScoreScratchPadding(uint64_t scoreSlot, uint64_t scoreRowBegin, uint64_t scoreRowCount, + uint64_t validColEnd, uint64_t subBlockIdx, uint64_t subBlockNum) + { + if (subBlockNum == 0 || subBlockIdx + 1 != subBlockNum) { + return; + } + const uint64_t validRowEnd = scoreRowBegin + scoreRowCount; + const uint64_t paddedRowEnd = (validRowEnd + 15) / 16 * 16; + const uint64_t paddedColEnd = BT_; + ZeroScoreScratchRange(scoreSlot, KDA_SCORE_SCRATCH_QG, KDA_SCORE_SCRATCH_KG, validRowEnd, paddedRowEnd); + ZeroScoreScratchRange(scoreSlot, KDA_SCORE_SCRATCH_KG, KDA_SCORE_SCRATCH_PLANES, validColEnd, paddedColEnd); + } + + __aicore__ inline void PrepareGateProducts(uint64_t b, uint64_t h, uint64_t hv, uint64_t start, uint64_t curT, + uint64_t subBlockIdx, uint64_t subBlockNum, bool useRef = false, + uint64_t refToken = 0, uint64_t validColEnd = 0, + bool writeScoreScratch = false, uint64_t scoreSlot = 0, + uint64_t scoreRowBegin = 0, uint64_t scoreRowCount = 0) + { + if (subBlockNum == 0 || subBlockIdx >= subBlockNum) { + return; + } + if (validColEnd == 0 || validColEnd > curT) { + validColEnd = curT; + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if (writeScoreScratch && curT == BT_ && scoreRowBegin == 0 && scoreRowCount == curT && validColEnd == curT) { + PrepareGateProductsBulk(b, h, hv, start, curT, subBlockIdx, subBlockNum, useRef, refToken, validColEnd, + writeScoreScratch, scoreSlot); + ZeroScoreScratchPadding(scoreSlot, scoreRowBegin, scoreRowCount, validColEnd, subBlockIdx, subBlockNum); + return; + } +#endif + if (writeScoreScratch) { + PrepareScoreFactorsBulk(b, h, hv, start, subBlockIdx, subBlockNum, refToken, scoreRowBegin, scoreRowCount, + validColEnd, start + curT - 1, scoreSlot); + ZeroScoreScratchPadding(scoreSlot, scoreRowBegin, scoreRowCount, validColEnd, subBlockIdx, subBlockNum); + return; + } + PrepareGateProductsBulk(b, h, hv, start, curT, subBlockIdx, subBlockNum, useRef, refToken, validColEnd, + writeScoreScratch, scoreSlot); + } + + __aicore__ inline void ComputeRawAqkAkkCube(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t curT) + { + ComputeRawAqkAkkCubeBlock(b, hv, chunkIdx, start, curT, 0, curT); + } + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + template + __aicore__ inline void ComputeRawAqkAkkCubeStableBlockDirectUbArch35(uint64_t b, uint64_t hv, uint64_t start, + uint64_t rowBegin, uint64_t scoreSlot, + uint8_t subBlockIdx, uint32_t directSlot) + { + using ElementA = SCORE_T; + using ElementB = SCORE_T; + using ElementC = float; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::ColumnMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using DirectTileCopy = Common::Tile::PackedTileCopyTlaToUB; + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; + using TileMmad = Catlass::Gemm::Tile::TileMmadTla; + + constexpr uint32_t scoreRows = KDA_SAFE_SCORE_REF_BC; + constexpr uint32_t packedRows = scoreRows * 2; + constexpr uint32_t k = 128; + static_assert(scoreRows == 32 && packedRows == 64); + static_assert(N == 32 || N == 64); + + Catlass::Arch::Resource resource; + auto layoutA = tla::MakeLayout(BT_, K_); + auto layoutB = tla::MakeLayout(K_, BT_); + auto tensorQPos = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_QG)], layoutA, + Catlass::Arch::PositionGM{}); + auto tensorKPos = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_W)], layoutA, + Catlass::Arch::PositionGM{}); + auto tensorKNeg = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_KG)], layoutB, + Catlass::Arch::PositionGM{}); + + auto blockQPos = GetTile(tensorQPos, tla::MakeCoord(rowBegin, 0), tla::MakeShape(scoreRows, k)); + auto blockKPos = GetTile(tensorKPos, tla::MakeCoord(rowBegin, 0), tla::MakeShape(scoreRows, k)); + auto blockKNeg = GetTile(tensorKNeg, tla::MakeCoord(0, 0), tla::MakeShape(k, N)); + + using CopyGmToL1A = typename TileCopy::template CopyGmToL1A; + using CopyGmToL1B = typename TileCopy::template CopyGmToL1B; + + static_assert(sizeof(ElementA) == sizeof(uint16_t)); + LocalTensor l1A = + resource.l1Buf.template GetBufferByByte(scoreSlot * KDA_DIRECT_SCORE_L1_SLOT_BYTES); + LocalTensor l1B = resource.l1Buf.template GetBufferByByte(KDA_DIRECT_SCORE_L1_B_OFFSET); + LocalTensor l0A = resource.l0ABuf.template GetBufferByByte(0); + LocalTensor l0B = resource.l0BBuf.template GetBufferByByte(0); + LocalTensor l0C = resource.l0CBuf.template GetBufferByByte(0); + auto layoutL1A = tla::MakeLayout(packedRows, k); + auto layoutL1B = tla::MakeLayout(k, N); + auto layoutL0A = tla::MakeLayout(packedRows, k); + auto layoutL0B = tla::MakeLayout(k, N); + auto layoutL0C = tla::MakeLayoutL0C(packedRows, N); + auto tensorL1A = tla::MakeTensor(l1A, layoutL1A, Catlass::Arch::PositionL1{}); + auto tensorL1B = tla::MakeTensor(l1B, layoutL1B, Catlass::Arch::PositionL1{}); + auto tensorL0A = tla::MakeTensor(l0A, layoutL0A, Catlass::Arch::PositionL0A{}); + auto tensorL0B = tla::MakeTensor(l0B, layoutL0B, Catlass::Arch::PositionL0B{}); + auto tensorL0C = tla::MakeTensor(l0C, layoutL0C, Catlass::Arch::PositionL0C{}); + auto tileL1A = GetTile(tensorL1A, tla::MakeCoord(0, 0), tla::MakeShape(packedRows, k)); + auto tileL1AQ = GetTile(tensorL1A, tla::MakeCoord(0, 0), tla::MakeShape(scoreRows, k)); + auto tileL1AK = GetTile(tensorL1A, tla::MakeCoord(scoreRows, 0), tla::MakeShape(scoreRows, k)); + auto tileL1B = GetTile(tensorL1B, tla::MakeCoord(0, 0), tla::MakeShape(k, N)); + auto tileL0A = GetTile(tensorL0A, tla::MakeCoord(0, 0), tla::MakeShape(packedRows, k)); + auto tileL0B = GetTile(tensorL0B, tla::MakeCoord(0, 0), tla::MakeShape(k, N)); + auto tileL0C = GetTile(tensorL0C, tla::MakeCoord(0, 0), tla::MakeShape(packedRows, N)); + auto tileL0CTop = GetTile(tensorL0C, tla::MakeCoord(0, 0), tla::MakeShape(scoreRows, N)); + auto tileL0CBottom = GetTile(tensorL0C, tla::MakeCoord(scoreRows, 0), tla::MakeShape(scoreRows, N)); + LocalTensor directBase = resource.ubBuf.template GetBufferByByte( + KDA_DIRECT_SCORE_UB_BYTE_OFFSET + directSlot * KDA_DIRECT_SCORE_SLOT_ELEMENTS * sizeof(ElementC)); + LocalTensor directAqk = directBase; + LocalTensor directAkk = directBase[KDA_DIRECT_SCORE_MATRIX_ELEMENTS]; + auto layoutDirect = tla::MakeLayout(scoreRows, BT_); + auto tensorDirectAqk = tla::MakeTensor(directAqk, layoutDirect, Catlass::Arch::PositionUB{}); + auto tensorDirectAkk = tla::MakeTensor(directAkk, layoutDirect, Catlass::Arch::PositionUB{}); + auto blockDirectAqk = GetTile(tensorDirectAqk, tla::MakeCoord(0, 0), tla::MakeShape(scoreRows, N)); + auto blockDirectAkk = GetTile(tensorDirectAkk, tla::MakeCoord(0, 0), tla::MakeShape(scoreRows, N)); + using CopyL0CToDirectUb = typename DirectTileCopy::template CopyL0CToDst; + + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToDirectUb copyL0CToDirectUb; + TileMmad tileMmad; + + copyGmToL1A(tileL1AQ, blockQPos); + copyGmToL1A(tileL1AK, blockKPos); + copyGmToL1B(tensorL1B, blockKNeg); + SetFlag(KDA_ARCH35_SCORE_EVENT); + WaitFlag(KDA_ARCH35_SCORE_EVENT); + copyL1ToL0A(tileL0A, tileL1A); + copyL1ToL0B(tileL0B, tileL1B); + SetFlag(KDA_ARCH35_SCORE_EVENT); + WaitFlag(KDA_ARCH35_SCORE_EVENT); + tileMmad(tileL0C, tileL0A, tileL0B, packedRows, N, k, true, 0); + SetFlag(KDA_ARCH35_SCORE_EVENT); + WaitFlag(KDA_ARCH35_SCORE_EVENT); + const uint64_t flagOffset = directSlot + subBlockIdx * KDA_DIRECT_SCORE_SUBBLOCK_FLAG_STRIDE; + CrossCoreWaitFlag<0x4, PIPE_FIX>(KDA_DIRECT_SCORE_FREE_FLAG + flagOffset); + copyL0CToDirectUb(blockDirectAqk, tileL0CTop, KDA_DIRECT_SCORE_ROWS, subBlockIdx, 1, 0); + copyL0CToDirectUb(blockDirectAkk, tileL0CBottom, KDA_DIRECT_SCORE_ROWS, subBlockIdx, 1, 0); + CrossCoreSetFlag<0x4, PIPE_FIX>(KDA_DIRECT_SCORE_READY_FLAG + flagOffset); + SetFlag(KDA_ARCH35_SCORE_EVENT); + WaitFlag(KDA_ARCH35_SCORE_EVENT); + } + + template + __aicore__ inline void PrefetchRawAqkAkkHeadPairArch35(uint64_t rowBegin, uint64_t scoreSlotBase, + uint32_t l1BaseOffset, TEventID readyEvent) + { + using ElementA = SCORE_T; + using ElementB = SCORE_T; + using ElementC = float; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::ColumnMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + + constexpr uint32_t scoreRows = KDA_SAFE_SCORE_REF_BC; + constexpr uint32_t packedRows = scoreRows * 2; + constexpr uint32_t n = N; + constexpr uint32_t k = 128; + constexpr uint32_t l1ABytes = packedRows * k * sizeof(ElementA); + constexpr uint32_t l1BBytes = k * n * sizeof(ElementB); + constexpr uint32_t l1LaneBytes = l1ABytes + l1BBytes; + static_assert(scoreRows == 32 && packedRows == 64); + static_assert(N == 32 || N == 64); + + Catlass::Arch::Resource resource; + auto layoutA = tla::MakeLayout(BT_, K_); + auto layoutB = tla::MakeLayout(K_, BT_); + auto layoutL1A = tla::MakeLayout(packedRows, k); + auto layoutL1B = tla::MakeLayout(k, n); + + for (uint64_t lane = 0; lane < KDA_SCORE_LANES; ++lane) { + const uint64_t scoreSlot = scoreSlotBase + lane; + LocalTensor l1A = resource.l1Buf.template GetBufferByByte( + l1BaseOffset + static_cast(lane) * l1LaneBytes); + LocalTensor l1B = resource.l1Buf.template GetBufferByByte( + l1BaseOffset + static_cast(lane) * l1LaneBytes + l1ABytes); + auto tensorL1A = tla::MakeTensor(l1A, layoutL1A, Catlass::Arch::PositionL1{}); + auto tensorL1B = tla::MakeTensor(l1B, layoutL1B, Catlass::Arch::PositionL1{}); + auto tileL1AQ = GetTile(tensorL1A, tla::MakeCoord(0, 0), tla::MakeShape(scoreRows, k)); + auto tileL1AK = GetTile(tensorL1A, tla::MakeCoord(scoreRows, 0), tla::MakeShape(scoreRows, k)); + + auto tensorQPos = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_QG)], + layoutA, Catlass::Arch::PositionGM{}); + auto tensorKPos = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_W)], + layoutA, Catlass::Arch::PositionGM{}); + auto tensorKNeg = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_KG)], + layoutB, Catlass::Arch::PositionGM{}); + auto blockQPos = GetTile(tensorQPos, tla::MakeCoord(rowBegin, 0), tla::MakeShape(scoreRows, k)); + auto blockKPos = GetTile(tensorKPos, tla::MakeCoord(rowBegin, 0), tla::MakeShape(scoreRows, k)); + auto blockKNeg = GetTile(tensorKNeg, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + using CopyGmToL1A = typename TileCopy::template CopyGmToL1A; + using CopyGmToL1B = typename TileCopy::template CopyGmToL1B; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + copyGmToL1A(tileL1AQ, blockQPos); + copyGmToL1A(tileL1AK, blockKPos); + copyGmToL1B(tensorL1B, blockKNeg); + } + SetFlag(readyEvent); + } + + template + __aicore__ inline void ComputeRawAqkAkkCubeStableHeadPairDirectUbArch35(uint64_t rowBegin, uint64_t scoreSlotBase, + uint32_t directSlot) + { + using ElementA = SCORE_T; + using ElementB = SCORE_T; + using ElementC = float; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::ColumnMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using DirectTileCopy = Common::Tile::PackedTileCopyTlaToUB; + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; + using TileMmad = Catlass::Gemm::Tile::TileMmadTla; + + constexpr uint32_t scoreRows = KDA_SAFE_SCORE_REF_BC; + constexpr uint32_t packedRows = scoreRows * 2; + constexpr uint32_t n = N; + constexpr uint32_t k = 128; + constexpr uint32_t l1ABytes = packedRows * k * sizeof(ElementA); + constexpr uint32_t l1BBytes = k * n * sizeof(ElementB); + constexpr uint32_t l1LaneBytes = l1ABytes + l1BBytes; + static_assert(scoreRows == 32 && packedRows == 64); + static_assert(N == 32 || N == 64); + + PrefetchRawAqkAkkHeadPairArch35(rowBegin, scoreSlotBase, 0, KDA_ARCH35_SCORE_EVENT); + WaitFlag(KDA_ARCH35_SCORE_EVENT); + + Catlass::Arch::Resource resource; + auto layoutL1A = tla::MakeLayout(packedRows, k); + auto layoutL1B = tla::MakeLayout(k, n); + + constexpr uint32_t l0ABytes = packedRows * k * sizeof(ElementA); + constexpr uint32_t l0BBytes = k * n * sizeof(ElementB); + constexpr uint32_t l0CBytes = packedRows * n * sizeof(ElementC); + LocalTensor l0A0 = resource.l0ABuf.template GetBufferByByte(0); + LocalTensor l0A1 = resource.l0ABuf.template GetBufferByByte(l0ABytes); + LocalTensor l0B0 = resource.l0BBuf.template GetBufferByByte(0); + LocalTensor l0B1 = resource.l0BBuf.template GetBufferByByte(l0BBytes); + LocalTensor l0C0 = resource.l0CBuf.template GetBufferByByte(0); + LocalTensor l0C1 = resource.l0CBuf.template GetBufferByByte(l0CBytes); + auto layoutL0A = tla::MakeLayout(packedRows, k); + auto layoutL0B = tla::MakeLayout(k, n); + auto layoutL0C = tla::MakeLayoutL0C(packedRows, n); + auto tensorL0A0 = tla::MakeTensor(l0A0, layoutL0A, Catlass::Arch::PositionL0A{}); + auto tensorL0A1 = tla::MakeTensor(l0A1, layoutL0A, Catlass::Arch::PositionL0A{}); + auto tensorL0B0 = tla::MakeTensor(l0B0, layoutL0B, Catlass::Arch::PositionL0B{}); + auto tensorL0B1 = tla::MakeTensor(l0B1, layoutL0B, Catlass::Arch::PositionL0B{}); + auto tensorL0C0 = tla::MakeTensor(l0C0, layoutL0C, Catlass::Arch::PositionL0C{}); + auto tensorL0C1 = tla::MakeTensor(l0C1, layoutL0C, Catlass::Arch::PositionL0C{}); + auto tileL0A0 = GetTile(tensorL0A0, tla::MakeCoord(0, 0), tla::MakeShape(packedRows, k)); + auto tileL0A1 = GetTile(tensorL0A1, tla::MakeCoord(0, 0), tla::MakeShape(packedRows, k)); + auto tileL0B0 = GetTile(tensorL0B0, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + auto tileL0B1 = GetTile(tensorL0B1, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + auto tileL0C0 = GetTile(tensorL0C0, tla::MakeCoord(0, 0), tla::MakeShape(packedRows, n)); + auto tileL0C1 = GetTile(tensorL0C1, tla::MakeCoord(0, 0), tla::MakeShape(packedRows, n)); + auto tileL0C0Top = GetTile(tensorL0C0, tla::MakeCoord(0, 0), tla::MakeShape(scoreRows, n)); + auto tileL0C0Bottom = GetTile(tensorL0C0, tla::MakeCoord(scoreRows, 0), tla::MakeShape(scoreRows, n)); + auto tileL0C1Top = GetTile(tensorL0C1, tla::MakeCoord(0, 0), tla::MakeShape(scoreRows, n)); + auto tileL0C1Bottom = GetTile(tensorL0C1, tla::MakeCoord(scoreRows, 0), tla::MakeShape(scoreRows, n)); + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + TileMmad tileMmad; + + LocalTensor l1A0 = resource.l1Buf.template GetBufferByByte(0); + LocalTensor l1A1 = resource.l1Buf.template GetBufferByByte(l1LaneBytes); + LocalTensor l1B0 = resource.l1Buf.template GetBufferByByte(l1ABytes); + LocalTensor l1B1 = resource.l1Buf.template GetBufferByByte(l1LaneBytes + l1ABytes); + auto tensorL1A0 = tla::MakeTensor(l1A0, layoutL1A, Catlass::Arch::PositionL1{}); + auto tensorL1A1 = tla::MakeTensor(l1A1, layoutL1A, Catlass::Arch::PositionL1{}); + auto tensorL1B0 = tla::MakeTensor(l1B0, layoutL1B, Catlass::Arch::PositionL1{}); + auto tensorL1B1 = tla::MakeTensor(l1B1, layoutL1B, Catlass::Arch::PositionL1{}); + auto tileL1A0 = GetTile(tensorL1A0, tla::MakeCoord(0, 0), tla::MakeShape(packedRows, k)); + auto tileL1A1 = GetTile(tensorL1A1, tla::MakeCoord(0, 0), tla::MakeShape(packedRows, k)); + auto tileL1B0 = GetTile(tensorL1B0, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + auto tileL1B1 = GetTile(tensorL1B1, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + + LocalTensor directBase = resource.ubBuf.template GetBufferByByte( + KDA_DIRECT_SCORE_UB_BYTE_OFFSET + directSlot * KDA_DIRECT_SCORE_SLOT_ELEMENTS * sizeof(ElementC)); + LocalTensor directAqk = directBase; + LocalTensor directAkk = directBase[KDA_DIRECT_SCORE_MATRIX_ELEMENTS]; + auto layoutDirect = tla::MakeLayout(scoreRows, BT_); + auto tensorDirectAqk = tla::MakeTensor(directAqk, layoutDirect, Catlass::Arch::PositionUB{}); + auto tensorDirectAkk = tla::MakeTensor(directAkk, layoutDirect, Catlass::Arch::PositionUB{}); + auto blockDirectAqk = GetTile(tensorDirectAqk, tla::MakeCoord(0, 0), tla::MakeShape(scoreRows, n)); + auto blockDirectAkk = GetTile(tensorDirectAkk, tla::MakeCoord(0, 0), tla::MakeShape(scoreRows, n)); + using CopyL0CToDirectUb = typename DirectTileCopy::template CopyL0CToDst; + CopyL0CToDirectUb copyL0CToDirectUb; + + copyL1ToL0A(tileL0A0, tileL1A0); + copyL1ToL0B(tileL0B0, tileL1B0); + SetFlag(KDA_ARCH35_SCORE_EVENT); + copyL1ToL0A(tileL0A1, tileL1A1); + copyL1ToL0B(tileL0B1, tileL1B1); + SetFlag(KDA_ARCH35_SCORE_W_EVENT); + + WaitFlag(KDA_ARCH35_SCORE_EVENT); + tileMmad(tileL0C0, tileL0A0, tileL0B0, packedRows, n, k, true, 0); + SetFlag(KDA_ARCH35_SCORE_EVENT); + WaitFlag(KDA_ARCH35_SCORE_EVENT); + const uint64_t flagOffset0 = directSlot; + CrossCoreWaitFlag<0x4, PIPE_FIX>(KDA_DIRECT_SCORE_FREE_FLAG + flagOffset0); + copyL0CToDirectUb(blockDirectAqk, tileL0C0Top, KDA_DIRECT_SCORE_ROWS, 0, 1, 0); + copyL0CToDirectUb(blockDirectAkk, tileL0C0Bottom, KDA_DIRECT_SCORE_ROWS, 0, 1, 0); + CrossCoreSetFlag<0x4, PIPE_FIX>(KDA_DIRECT_SCORE_READY_FLAG + flagOffset0); + SetFlag(KDA_ARCH35_SCORE_EVENT); + + WaitFlag(KDA_ARCH35_SCORE_W_EVENT); + tileMmad(tileL0C1, tileL0A1, tileL0B1, packedRows, n, k, true, 0); + SetFlag(KDA_ARCH35_SCORE_W_EVENT); + WaitFlag(KDA_ARCH35_SCORE_W_EVENT); + const uint64_t flagOffset1 = directSlot + KDA_DIRECT_SCORE_SUBBLOCK_FLAG_STRIDE; + CrossCoreWaitFlag<0x4, PIPE_FIX>(KDA_DIRECT_SCORE_FREE_FLAG + flagOffset1); + copyL0CToDirectUb(blockDirectAqk, tileL0C1Top, KDA_DIRECT_SCORE_ROWS, 1, 1, 0); + copyL0CToDirectUb(blockDirectAkk, tileL0C1Bottom, KDA_DIRECT_SCORE_ROWS, 1, 1, 0); + CrossCoreSetFlag<0x4, PIPE_FIX>(KDA_DIRECT_SCORE_READY_FLAG + flagOffset1); + SetFlag(KDA_ARCH35_SCORE_W_EVENT); + + WaitFlag(KDA_ARCH35_SCORE_EVENT); + WaitFlag(KDA_ARCH35_SCORE_W_EVENT); + } + + __aicore__ inline void ComputeRawAqkAkkCubeFullArch35(uint64_t b, uint64_t hv, uint64_t start, uint64_t scoreSlot) + { + using ElementA = SCORE_T; + using ElementB = SCORE_T; + using ElementC = float; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::ColumnMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; + using TileMmad = Catlass::Gemm::Tile::TileMmadTla; + + constexpr uint32_t m = 64; + constexpr uint32_t n = 64; + constexpr uint32_t k = 128; + Catlass::Arch::Resource resource; + auto layoutA = tla::MakeLayout(m, k); + auto layoutB = tla::MakeLayout(k, n); + auto layoutC = tla::MakeLayout(m, n); + auto tensorQPos = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_QG)], layoutA, + Catlass::Arch::PositionGM{}); + auto tensorKPos = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_W)], layoutA, + Catlass::Arch::PositionGM{}); + auto tensorKNeg = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_KG)], layoutB, + Catlass::Arch::PositionGM{}); + auto tensorAqk = tla::MakeTensor(aqk_[AOffset(b, hv, start, 0)], layoutC, Catlass::Arch::PositionGM{}); + auto tensorAkk = tla::MakeTensor(akk_[AOffset(b, hv, start, 0)], layoutC, Catlass::Arch::PositionGM{}); + auto blockQPos = GetTile(tensorQPos, tla::MakeCoord(0, 0), tla::MakeShape(m, k)); + auto blockKPos = GetTile(tensorKPos, tla::MakeCoord(0, 0), tla::MakeShape(m, k)); + auto blockKNeg = GetTile(tensorKNeg, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + auto blockAqk = GetTile(tensorAqk, tla::MakeCoord(0, 0), tla::MakeShape(m, n)); + auto blockAkk = GetTile(tensorAkk, tla::MakeCoord(0, 0), tla::MakeShape(m, n)); + + using CopyGmToL1A = typename TileCopy::template CopyGmToL1A; + using CopyGmToL1B = typename TileCopy::template CopyGmToL1B; + using CopyL0CToDst = typename TileCopy::template CopyL0CToDst; + + constexpr uint32_t l1ABytes = m * k * sizeof(ElementA); + LocalTensor l1A0 = resource.l1Buf.template GetBufferByByte(0); + LocalTensor l1A1 = resource.l1Buf.template GetBufferByByte(l1ABytes); + LocalTensor l1B = resource.l1Buf.template GetBufferByByte(2 * l1ABytes); + LocalTensor l0A = resource.l0ABuf.template GetBufferByByte(0); + LocalTensor l0B = resource.l0BBuf.template GetBufferByByte(0); + LocalTensor l0C = resource.l0CBuf.template GetBufferByByte(0); + auto layoutL1A = tla::MakeLayout(m, k); + auto layoutL1B = tla::MakeLayout(k, n); + auto layoutL0A = tla::MakeLayout(m, k); + auto layoutL0B = tla::MakeLayout(k, n); + auto layoutL0C = tla::MakeLayoutL0C(m, n); + auto tensorL1A0 = tla::MakeTensor(l1A0, layoutL1A, Catlass::Arch::PositionL1{}); + auto tensorL1A1 = tla::MakeTensor(l1A1, layoutL1A, Catlass::Arch::PositionL1{}); + auto tensorL1B = tla::MakeTensor(l1B, layoutL1B, Catlass::Arch::PositionL1{}); + auto tensorL0A = tla::MakeTensor(l0A, layoutL0A, Catlass::Arch::PositionL0A{}); + auto tensorL0B = tla::MakeTensor(l0B, layoutL0B, Catlass::Arch::PositionL0B{}); + auto tensorL0C = tla::MakeTensor(l0C, layoutL0C, Catlass::Arch::PositionL0C{}); + auto tileL1A0 = GetTile(tensorL1A0, tla::MakeCoord(0, 0), tla::MakeShape(m, k)); + auto tileL1A1 = GetTile(tensorL1A1, tla::MakeCoord(0, 0), tla::MakeShape(m, k)); + auto tileL1B = GetTile(tensorL1B, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + auto tileL0A = GetTile(tensorL0A, tla::MakeCoord(0, 0), tla::MakeShape(m, k)); + auto tileL0B = GetTile(tensorL0B, tla::MakeCoord(0, 0), tla::MakeShape(k, n)); + auto tileL0C = GetTile(tensorL0C, tla::MakeCoord(0, 0), tla::MakeShape(m, n)); + + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL0CToDst copyL0CToDst; + TileMmad tileMmad; + + copyGmToL1B(tensorL1B, blockKNeg); + copyGmToL1A(tensorL1A0, blockQPos); + SetFlag(KDA_ARCH35_SCORE_EVENT); + WaitFlag(KDA_ARCH35_SCORE_EVENT); + copyL1ToL0B(tileL0B, tileL1B); + copyL1ToL0A(tileL0A, tileL1A0); + SetFlag(KDA_ARCH35_SCORE_EVENT); + copyGmToL1A(tensorL1A1, blockKPos); + SetFlag(KDA_ARCH35_SCORE_W_EVENT); + WaitFlag(KDA_ARCH35_SCORE_EVENT); + tileMmad(tileL0C, tileL0A, tileL0B, m, n, k, true, 0); + SetFlag(KDA_ARCH35_SCORE_EVENT); + SetFlag(KDA_ARCH35_SCORE_EVENT); + WaitFlag(KDA_ARCH35_SCORE_W_EVENT); + WaitFlag(KDA_ARCH35_SCORE_EVENT); + WaitFlag(KDA_ARCH35_SCORE_EVENT); + copyL0CToDst(blockAqk, tensorL0C); + SetFlag(KDA_ARCH35_SCORE_EVENT); + copyL1ToL0B(tileL0B, tileL1B); + copyL1ToL0A(tileL0A, tileL1A1); + SetFlag(KDA_ARCH35_SCORE_EVENT); + WaitFlag(KDA_ARCH35_SCORE_EVENT); + WaitFlag(KDA_ARCH35_SCORE_EVENT); + tileMmad(tileL0C, tileL0A, tileL0B, m, n, k, true, 0); + SetFlag(KDA_ARCH35_SCORE_EVENT); + WaitFlag(KDA_ARCH35_SCORE_EVENT); + copyL0CToDst(blockAkk, tensorL0C); + SetFlag(KDA_ARCH35_SCORE_EVENT); + WaitFlag(KDA_ARCH35_SCORE_EVENT); + } +#endif + + __aicore__ inline void ComputeRawAqkAkkCubeBlock(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t curT, uint64_t rowBegin, uint64_t rowCount, + bool readScoreScratch = false, uint64_t scoreSlot = 0, + uint64_t colCount = 0) + { + if (colCount == 0 || colCount > curT) { + colCount = curT; + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128) { + if (KDA_ARCH35_ENABLE_MANUAL_SCORE_PIPELINE && fusePostWu_ && rowBegin == 0 && rowCount == 64 && + colCount == 64) { + ComputeRawAqkAkkCubeFullArch35(b, hv, start, scoreSlot); + return; + } + } +#endif + using ElementA = SCORE_T; + using ElementB = SCORE_T; + using ElementC = float; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::ColumnMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using BlockMmad = Catlass::Gemm::Block::BlockMmadTla; + + Catlass::Arch::Resource resource; + BlockMmad blockMmad(resource); + auto layoutA = tla::MakeLayout(BT_, K_); + auto layoutB = tla::MakeLayout(K_, BT_); + auto layoutC = tla::MakeLayout(BT_, BT_); + const bool paddedTail = curT < BT_; + const uint64_t mmRowCount = paddedTail ? (rowCount + 15) / 16 * 16 : rowCount; + const uint64_t mmColCount = paddedTail ? BT_ : colCount; + Catlass::GemmCoord shape{static_cast(mmRowCount), static_cast(mmColCount), + static_cast(K_)}; + + (void)readScoreScratch; + auto tensorQPos = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_QG)], layoutA, + Catlass::Arch::PositionGM{}); + auto tensorKPos = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_W)], layoutA, + Catlass::Arch::PositionGM{}); + auto tensorKNeg = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_KG)], layoutB, + Catlass::Arch::PositionGM{}); + auto aqkBase = paddedTail ? solveWorkspace_[SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_RAW_AQK)] : + aqk_[AOffset(b, hv, start, 0)]; + auto akkBase = paddedTail ? solveWorkspace_[SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_RAW_AKK)] : + akk_[AOffset(b, hv, start, 0)]; + auto tensorAqk = tla::MakeTensor(aqkBase, layoutC, Catlass::Arch::PositionGM{}); + auto tensorAkk = tla::MakeTensor(akkBase, layoutC, Catlass::Arch::PositionGM{}); + + auto blockQPos = GetTile(tensorQPos, tla::MakeCoord(rowBegin, 0), tla::MakeShape(shape.m(), shape.k())); + auto blockKPos = GetTile(tensorKPos, tla::MakeCoord(rowBegin, 0), tla::MakeShape(shape.m(), shape.k())); + auto blockKNeg = GetTile(tensorKNeg, tla::MakeCoord(0, 0), tla::MakeShape(shape.k(), shape.n())); + auto blockAqk = GetTile(tensorAqk, tla::MakeCoord(rowBegin, 0), tla::MakeShape(shape.m(), shape.n())); + auto blockAkk = GetTile(tensorAkk, tla::MakeCoord(rowBegin, 0), tla::MakeShape(shape.m(), shape.n())); + + blockMmad.preSetFlags(); + blockMmad(blockQPos, blockKNeg, blockAqk, shape); + blockMmad(blockKPos, blockKNeg, blockAkk, shape); + blockMmad.finalWaitFlags(); + } + + __aicore__ inline bool UseAkkCubeSolve(uint64_t curT) const + { + return curT > 0 && curT <= BT_ && (BT_ == 64 || BT_ == 128) && K_ >= 16 && V_ >= 16 && V_ <= 256 && + K_ % 16 == 0 && V_ % 16 == 0; + } + + __aicore__ inline bool UsePostWuCube(uint64_t curT) const + { + return curT > 0 && curT <= BT_ && (BT_ == 64 || BT_ == 128) && K_ >= 16 && V_ >= 16 && V_ <= 256 && + K_ % 16 == 0 && V_ % 16 == 0; + } + + __aicore__ inline void CopyLocalFloat(LocalTensor dst, LocalTensor src, uint64_t count) + { + if (count == 0) { + return; + } + Adds(dst, src, 0.0f, static_cast(count)); + PipeBarrier(); + } + + __aicore__ inline void FillLocalFloat(LocalTensor dst, float value, uint64_t count) + { + if (count == 0) { + return; + } + Duplicate(dst, value, static_cast(count)); + PipeBarrier(); + } + + __aicore__ inline void ForwardSubDiag16(LocalTensor diag, LocalTensor row, LocalTensor prod, + LocalTensor rowBrcb, LocalTensor reduced, uint64_t valid) + { + constexpr uint32_t brcbStride = 8; + constexpr uint32_t diagSize = KDA_SOLVE_DIAG_BT; + constexpr uint8_t rowBlk = diagSize * sizeof(float) / 32; + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + ForwardSubDiag16Regbase((__ubuf__ float *)reinterpret_cast(diag.GetPhyAddr()), + static_cast(valid)); +#else + for (uint64_t i = 2; i < valid; ++i) { + uint32_t rowOffset = static_cast(i * diagSize); + DataCopy(row, diag[rowOffset], diagSize); + PipeBarrier(); + + Brcb(rowBrcb, row, diagSize / brcbStride, {1, 8}); + PipeBarrier(); + for (uint32_t col = 0; col < diagSize; col += brcbStride) { + Mul(prod[col], diag[col], rowBrcb, brcbStride, static_cast(diagSize), + {1, 1, 0, rowBlk, rowBlk, 1}); + } + PipeBarrier(); + + uint32_t remain = diagSize; + while (remain > 1) { + uint32_t calcCount = (remain / 2) * diagSize; + remain = (remain + 1) / 2; + Add(prod, prod, prod[remain * diagSize], calcCount); + PipeBarrier(); + } + DataCopy(reduced, prod, diagSize); + PipeBarrier(); + Add(row, row, reduced, diagSize); + PipeBarrier(); + DataCopy(diag[rowOffset], row, diagSize); + PipeBarrier(); + } +#endif + + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + for (uint32_t i = 0; i < diagSize; ++i) { + uint32_t diagOffset = i * diagSize + i; + if (i < valid) { + diag.SetValue(diagOffset, diag.GetValue(diagOffset) + 1.0f); + } else { + diag.SetValue(diagOffset, 1.0f); + } + } + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + } + + __aicore__ inline void SolveDiagonalBlocksInRows(LocalTensor akkMat, LocalTensor xMat, + LocalTensor arena, uint64_t scratchBase, uint64_t curT, + uint64_t rowBegin, uint64_t rowCount) + { + constexpr uint32_t diagSize = KDA_SOLVE_DIAG_BT; + constexpr uint32_t diagElements = diagSize * diagSize; + constexpr uint32_t brcbElements = diagSize * 8; + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + (void)akkMat; + (void)arena; + (void)scratchBase; + uint64_t rowEnd = rowBegin + rowCount; + for (uint64_t blockBegin = 0; blockBegin < BT_; blockBegin += diagSize) { + if (blockBegin < rowBegin || blockBegin + diagSize > rowEnd) { + continue; + } + uint64_t localBlockRow = blockBegin - rowBegin; + uint64_t valid = blockBegin < curT ? curT - blockBegin : 0; + if (valid > diagSize) { + valid = diagSize; + } + ForwardSubDiag16StridedRegbase((__ubuf__ float *)reinterpret_cast(xMat.GetPhyAddr()), + static_cast(BT_), static_cast(localBlockRow), + static_cast(blockBegin), static_cast(valid)); + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + for (uint32_t rowIdx = 0; rowIdx < diagSize; ++rowIdx) { + uint32_t diagOffset = static_cast((localBlockRow + rowIdx) * BT_ + blockBegin + rowIdx); + if (rowIdx < valid) { + xMat.SetValue(diagOffset, xMat.GetValue(diagOffset) + 1.0f); + } else { + xMat.SetValue(diagOffset, 1.0f); + } + } + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + } +#else + LocalTensor diag = arena[scratchBase]; + LocalTensor row = diag[diagElements]; + LocalTensor prod = row[diagSize]; + LocalTensor rowBrcb = prod[diagElements]; + LocalTensor reduced = rowBrcb[brcbElements]; + + uint64_t rowEnd = rowBegin + rowCount; + for (uint64_t blockBegin = 0; blockBegin < BT_; blockBegin += diagSize) { + if (blockBegin < rowBegin || blockBegin + diagSize > rowEnd) { + continue; + } + Duplicate(diag, 0.0f, diagElements); + PipeBarrier(); + + uint64_t localBlockRow = blockBegin - rowBegin; + uint64_t valid = blockBegin < curT ? curT - blockBegin : 0; + if (valid > diagSize) { + valid = diagSize; + } + for (uint32_t rowIdx = 0; rowIdx < diagSize; ++rowIdx) { + uint64_t srcOffset = (localBlockRow + rowIdx) * BT_ + blockBegin; + Muls(diag[rowIdx * diagSize], akkMat[srcOffset], -1.0f, diagSize); + } + PipeBarrier(); + + ForwardSubDiag16(diag, row, prod, rowBrcb, reduced, valid); + for (uint32_t rowIdx = 0; rowIdx < diagSize; ++rowIdx) { + uint64_t dstOffset = (localBlockRow + rowIdx) * BT_ + blockBegin; + Adds(xMat[dstOffset], diag[rowIdx * diagSize], 0.0f, diagSize); + } + PipeBarrier(); + } +#endif + } + + __aicore__ inline void BuildPrefixMask(LocalTensor dst, uint64_t prefix, uint64_t count) + { + if (prefix > count) { + prefix = count; + } + Duplicate(dst, 0.0f, static_cast(count)); + if (prefix > 0) { + Duplicate(dst, 1.0f, static_cast(prefix)); + } + PipeBarrier(); + } + + __aicore__ inline uint64_t BuildCausalMask(uint64_t threshold, uint64_t colBegin) const + { + if (threshold <= colBegin) { + return ~0ULL; + } + if (threshold >= colBegin + KDA_SOLVE_BT) { + return 0ULL; + } + return ~0ULL << (threshold - colBegin); + } + + __aicore__ inline void BuildCausalSelectMasks(LocalTensor aqkMask, LocalTensor akkMask, + uint64_t rowBegin, uint64_t rowCount, uint64_t colBegin) + { + __ubuf__ uint64_t *aqkMaskPtr = reinterpret_cast<__ubuf__ uint64_t *>(aqkMask.GetPhyAddr()); + __ubuf__ uint64_t *akkMaskPtr = reinterpret_cast<__ubuf__ uint64_t *>(akkMask.GetPhyAddr()); + for (uint32_t localRow = 0; localRow < rowCount; ++localRow) { + uint32_t row = static_cast(rowBegin + localRow); + aqkMaskPtr[localRow] = BuildCausalMask(static_cast(row) + 1, colBegin); + akkMaskPtr[localRow] = BuildCausalMask(static_cast(row), colBegin); + } + } + + __aicore__ inline void SelectCausalRows(LocalTensor aqkMat, LocalTensor akkMat, uint64_t rowBegin, + uint64_t rowCount) + { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128) { + SelectCausalRows64Regbase((__ubuf__ float *)reinterpret_cast(aqkMat.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(akkMat.GetPhyAddr()), + static_cast(rowBegin), static_cast(rowCount)); + PipeBarrier(); + return; + } +#endif + LocalTensor aqkMask = vecBuf_.Get()[KDA_SELECT_AQK_MASK_BYTE_OFFSET]; + LocalTensor akkMask = vecBuf_.Get()[KDA_SELECT_AKK_MASK_BYTE_OFFSET]; + LocalTensor zeroLocal = vecBuf_.Get()[KDA_SELECT_ZERO_FLOAT_OFFSET]; + Duplicate(zeroLocal, 0.0f, 8); + PipeBarrier(); + + uint64_t colBlockCount = (BT_ + KDA_SOLVE_BT - 1) / KDA_SOLVE_BT; + for (uint64_t colBlock = 0; colBlock < colBlockCount; ++colBlock) { + uint64_t maskOffset = colBlock * KDA_SELECT_COL_MASK_BYTES; + uint64_t colBegin = colBlock * KDA_SOLVE_BT; + BuildCausalSelectMasks(aqkMask[maskOffset], akkMask[maskOffset], rowBegin, rowCount, colBegin); + } + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + + uint8_t rowStride = static_cast(BT_ * sizeof(float) / 32); + BinaryRepeatParams repeatParams = {1, 0, 1, rowStride, 0, rowStride}; + for (uint64_t colBlock = 0; colBlock < colBlockCount; ++colBlock) { + uint64_t maskOffset = colBlock * KDA_SELECT_COL_MASK_BYTES; + uint64_t colBegin = colBlock * KDA_SOLVE_BT; + Select(aqkMat[colBegin], aqkMask[maskOffset], zeroLocal, aqkMat[colBegin], SELMODE::VSEL_TENSOR_TENSOR_MODE, + KDA_SOLVE_BT, static_cast(rowCount), repeatParams); + Select(akkMat[colBegin], akkMask[maskOffset], zeroLocal, akkMat[colBegin], SELMODE::VSEL_TENSOR_TENSOR_MODE, + KDA_SOLVE_BT, static_cast(rowCount), repeatParams); + } + PipeBarrier(); + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + } + + __aicore__ inline void PrepareAqkAkkSolveInput64(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start) + { + LocalTensor arena = vecBuf_.Get(); + LocalTensor aqkMat = arena; + LocalTensor akkMat = arena[KDA_SOLVE_MATRIX_ELEMENTS]; + LocalTensor xMat = arena[2 * KDA_SOLVE_MATRIX_ELEMENTS]; + LocalTensor betaLocal = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS]; + LocalTensor betaBrcb = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS + KDA_SOLVE_BT]; + LocalTensor maskLocal = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS + KDA_SOLVE_BT + 512]; + LocalTensor oneHotLocal = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS + KDA_SOLVE_BT + 512 + KDA_SOLVE_BT]; + + LoadBetaAsFloatRow(BetaOffset(b, hv, start), betaLocal, KDA_SOLVE_BT); + Brcb(betaBrcb, betaLocal, 8, {1, 8}); + PipeBarrier(); + + DataCopy(aqkMat, aqk_[AOffset(b, hv, start, 0)], KDA_SOLVE_MATRIX_ELEMENTS); + DataCopy(akkMat, akk_[AOffset(b, hv, start, 0)], KDA_SOLVE_MATRIX_ELEMENTS); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + + for (uint64_t col = 0; col < KDA_SOLVE_BT; col += 8) { + Mul(akkMat[col], akkMat[col], betaBrcb, 8, KDA_SOLVE_BT, {1, 1, 1, 8, 8, 1}); + PipeBarrier(); + } + SelectCausalRows(aqkMat, akkMat, 0, KDA_SOLVE_BT); + + Muls(xMat, akkMat, -1.0f, KDA_SOLVE_MATRIX_ELEMENTS); + PipeBarrier(); + for (uint64_t row = 0; row < KDA_SOLVE_BT; ++row) { + BuildPrefixMask(maskLocal, row + 1, KDA_SOLVE_BT); + BuildPrefixMask(oneHotLocal, row, KDA_SOLVE_BT); + Sub(maskLocal, maskLocal, oneHotLocal, KDA_SOLVE_BT); + PipeBarrier(); + Add(xMat[row * KDA_SOLVE_BT], xMat[row * KDA_SOLVE_BT], maskLocal, KDA_SOLVE_BT); + PipeBarrier(); + } + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(aqk_[AOffset(b, hv, start, 0)], aqkMat, KDA_SOLVE_MATRIX_ELEMENTS); + DataCopy(akk_[AOffset(b, hv, start, 0)], akkMat, KDA_SOLVE_MATRIX_ELEMENTS); + DataCopy(h_[SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X)], xMat, KDA_SOLVE_MATRIX_ELEMENTS); + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void PrepareAqkAkkSolveInputTail(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t curT) + { + uint64_t elemCount = curT * KDA_SOLVE_BT; + LocalTensor arena = vecBuf_.Get(); + LocalTensor aqkMat = arena; + LocalTensor akkMat = arena[KDA_SOLVE_MATRIX_ELEMENTS]; + LocalTensor xMat = arena[2 * KDA_SOLVE_MATRIX_ELEMENTS]; + LocalTensor betaLocal = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS]; + LocalTensor betaBrcb = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS + KDA_SOLVE_BT]; + LocalTensor maskLocal = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS + KDA_SOLVE_BT + 512]; + LocalTensor oneHotLocal = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS + KDA_SOLVE_BT + 512 + KDA_SOLVE_BT]; + + FillLocalFloat(betaLocal, 0.0f, KDA_SOLVE_BT); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + LoadBetaAsFloatRow(BetaOffset(b, hv, start), betaLocal, curT); + Brcb(betaBrcb, betaLocal, 8, {1, 8}); + PipeBarrier(); + + DataCopy(aqkMat, aqk_[AOffset(b, hv, start, 0)], static_cast(elemCount)); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + if (elemCount < KDA_SOLVE_MATRIX_ELEMENTS) { + FillLocalFloat(aqkMat[elemCount], 0.0f, KDA_SOLVE_MATRIX_ELEMENTS - elemCount); + } + DataCopy(akkMat, akk_[AOffset(b, hv, start, 0)], static_cast(elemCount)); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + if (elemCount < KDA_SOLVE_MATRIX_ELEMENTS) { + FillLocalFloat(akkMat[elemCount], 0.0f, KDA_SOLVE_MATRIX_ELEMENTS - elemCount); + } + + for (uint64_t col = 0; col < KDA_SOLVE_BT; col += 8) { + Mul(akkMat[col], akkMat[col], betaBrcb, 8, KDA_SOLVE_BT, {1, 1, 1, 8, 8, 1}); + PipeBarrier(); + } + SelectCausalRows(aqkMat, akkMat, 0, KDA_SOLVE_BT); + + Muls(xMat, akkMat, -1.0f, KDA_SOLVE_MATRIX_ELEMENTS); + PipeBarrier(); + for (uint64_t row = 0; row < KDA_SOLVE_BT; ++row) { + BuildPrefixMask(maskLocal, row + 1, KDA_SOLVE_BT); + BuildPrefixMask(oneHotLocal, row, KDA_SOLVE_BT); + Sub(maskLocal, maskLocal, oneHotLocal, KDA_SOLVE_BT); + PipeBarrier(); + Add(xMat[row * KDA_SOLVE_BT], xMat[row * KDA_SOLVE_BT], maskLocal, KDA_SOLVE_BT); + PipeBarrier(); + } + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(aqk_[AOffset(b, hv, start, 0)], aqkMat, static_cast(elemCount)); + DataCopy(h_[SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X)], xMat, KDA_SOLVE_MATRIX_ELEMENTS); + DataCopy(h_[SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_Y0)], akkMat, KDA_SOLVE_MATRIX_ELEMENTS); + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void GetSolveRowRange(uint64_t curT, uint64_t subBlockIdx, uint64_t subBlockNum, + uint64_t &rowBegin, uint64_t &rowEnd) const + { + if (subBlockNum == 0 || subBlockIdx >= subBlockNum) { + rowBegin = 0; + rowEnd = 0; + return; + } + rowBegin = (curT * subBlockIdx) / subBlockNum; + rowEnd = (curT * (subBlockIdx + 1)) / subBlockNum; + } + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + __aicore__ inline void InitializeDirectScoreUbArch35() + { + for (uint32_t slot = 0; slot < KDA_DIRECT_SCORE_QUEUE_DEPTH; ++slot) { + CrossCoreSetFlag<0x4, PIPE_V>(KDA_DIRECT_SCORE_FREE_FLAG + slot); + } + } + + __aicore__ inline void ProcessDirectScoreSolveRowsArch35(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t rowBegin, uint32_t directSlot) + { + CrossCoreWaitFlag<0x4, PIPE_V>(KDA_DIRECT_SCORE_READY_FLAG + directSlot); + + Catlass::Arch::Resource resource; + LocalTensor arena = vecBuf_.Get(); + LocalTensor directBase = resource.ubBuf.template GetBufferByByte( + KDA_DIRECT_SCORE_UB_BYTE_OFFSET + directSlot * KDA_DIRECT_SCORE_SLOT_ELEMENTS * sizeof(float)); + LocalTensor aqkMat = directBase; + LocalTensor akkMat = directBase[KDA_DIRECT_SCORE_MATRIX_ELEMENTS]; + LocalTensor xMat = directBase[2 * KDA_DIRECT_SCORE_MATRIX_ELEMENTS]; + + SelectCausalRows64Regbase((__ubuf__ float *)reinterpret_cast(aqkMat.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(akkMat.GetPhyAddr()), + static_cast(rowBegin), KDA_DIRECT_SCORE_ROWS); + PipeBarrier(); + Muls(xMat, akkMat, -1.0f, KDA_DIRECT_SCORE_MATRIX_ELEMENTS); + PipeBarrier(); + SolveDiagonalBlocksInRows(akkMat, xMat, arena, 0, BT_, rowBegin, KDA_DIRECT_SCORE_ROWS); + + LocalTensor aqkTyped = GateQTyped(0); + Muls(aqkMat, aqkMat, scale_, KDA_DIRECT_SCORE_MATRIX_ELEMENTS); + PipeBarrier(); + ClampFp32ToOutputType(aqkMat, KDA_DIRECT_SCORE_MATRIX_ELEMENTS); + Cast(aqkTyped, aqkMat, RoundMode::CAST_RINT, KDA_DIRECT_SCORE_MATRIX_ELEMENTS); + PipeBarrier(); + + const uint64_t token = start + rowBegin; + const uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X) + rowBegin * BT_; + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(o_, AOffset(b, hv, token, 0), aqkTyped, KDA_DIRECT_SCORE_MATRIX_ELEMENTS); + DataCopy(solveWorkspace_[xBase], xMat, KDA_DIRECT_SCORE_MATRIX_ELEMENTS); + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + CrossCoreSetFlag<0x4, PIPE_V>(KDA_DIRECT_SCORE_FREE_FLAG + directSlot); + } +#endif + + __aicore__ inline void PrepareAqkAkkSolveInputRows(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t curT, uint64_t rowBegin, uint64_t rowEnd, + bool storeLToAkk, bool storeLToScratch) + { + uint64_t rowCount = rowEnd - rowBegin; + if (rowCount == 0) { + return; + } + uint64_t validRowCount = rowBegin < curT ? curT - rowBegin : 0; + if (validRowCount > rowCount) { + validRowCount = rowCount; + } + uint64_t elemCount = rowCount * BT_; + uint64_t validElemCount = validRowCount * BT_; + LocalTensor arena = vecBuf_.Get(); + LocalTensor aqkMat = arena; + LocalTensor akkMat = arena[elemCount]; + LocalTensor xMat = arena[2 * elemCount]; + LocalTensor betaLocal = arena[3 * elemCount]; + LocalTensor betaBrcb = arena[3 * elemCount + BT_]; + LocalTensor maskLocal = arena[3 * elemCount + BT_ + 512]; + LocalTensor oneHotLocal = arena[3 * elemCount + BT_ + 512 + BT_]; + + uint64_t token = start + rowBegin; + + constexpr bool scoreCanIncludeBeta = +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128; +#else + false; +#endif + const bool scoreIncludesBeta = scoreCanIncludeBeta && HV_ % KDA_SCORE_LANES == 0; + if (validRowCount < rowCount) { + FillLocalFloat(aqkMat, 0.0f, elemCount); + FillLocalFloat(akkMat, 0.0f, elemCount); + if (!scoreIncludesBeta) { + FillLocalFloat(betaLocal, 0.0f, rowCount); + } + } + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + if (validRowCount > 0) { + if (!scoreIncludesBeta) { + LoadBetaAsFloatRow(BetaOffset(b, hv, token), betaLocal, validRowCount); + } + if (curT < BT_) { + DataCopy( + aqkMat, + solveWorkspace_[SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_RAW_AQK) + rowBegin * BT_], + static_cast(validElemCount)); + DataCopy( + akkMat, + solveWorkspace_[SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_RAW_AKK) + rowBegin * BT_], + static_cast(validElemCount)); + } else { + DataCopy(aqkMat, aqk_[AOffset(b, hv, token, 0)], static_cast(validElemCount)); + DataCopy(akkMat, akk_[AOffset(b, hv, token, 0)], static_cast(validElemCount)); + } + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if (!scoreIncludesBeta) { + ApplyKdaRowScaleRegbase((__ubuf__ float *)reinterpret_cast(akkMat.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(betaLocal.GetPhyAddr()), + static_cast(rowCount), static_cast(BT_)); + PipeBarrier(); + } +#else + Brcb(betaBrcb, betaLocal, static_cast((rowCount + 7) / 8), {1, 8}); + PipeBarrier(); + uint8_t rowStride = static_cast(BT_ * sizeof(float) / 32); + for (uint64_t col = 0; col < BT_; col += 8) { + Mul(akkMat[col], akkMat[col], betaBrcb, 8, static_cast(rowCount), + {1, 1, 0, rowStride, rowStride, 1}); + } + PipeBarrier(); +#endif + if (validRowCount > 0) { + SelectCausalRows(aqkMat, akkMat, rowBegin, validRowCount); + } + + Muls(xMat, akkMat, -1.0f, static_cast(elemCount)); + PipeBarrier(); + if constexpr (SAFE_GATE) { + uint64_t scratchBase = 3 * elemCount + BT_ + 512 + 2 * BT_; + SolveDiagonalBlocksInRows(akkMat, xMat, arena, scratchBase, curT, rowBegin, rowCount); + } else if (curT < BT_) { + uint64_t scratchBase = 3 * elemCount + BT_ + 512 + 2 * BT_; + SolveDiagonalBlocksInRows(akkMat, xMat, arena, scratchBase, curT, rowBegin, rowCount); + } else { + for (uint64_t localRow = 0; localRow < rowCount; ++localRow) { + uint64_t row = rowBegin + localRow; + BuildPrefixMask(maskLocal, row + 1, BT_); + BuildPrefixMask(oneHotLocal, row, BT_); + Sub(maskLocal, maskLocal, oneHotLocal, static_cast(BT_)); + PipeBarrier(); + Add(xMat[localRow * BT_], xMat[localRow * BT_], maskLocal, static_cast(BT_)); + PipeBarrier(); + } + } + + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X) + rowBegin * BT_; + uint64_t lBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_Y0) + rowBegin * BT_; +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128) { + LocalTensor aqkTyped = GateQTyped(0); + if (validElemCount > 0) { + Muls(aqkMat, aqkMat, scale_, static_cast(validElemCount)); + PipeBarrier(); + ClampFp32ToOutputType(aqkMat, static_cast(validElemCount)); + Cast(aqkTyped, aqkMat, RoundMode::CAST_RINT, static_cast(validElemCount)); + PipeBarrier(); + } + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + if (validElemCount > 0) { + CopyVectorOut(o_, AOffset(b, hv, token, 0), aqkTyped, validElemCount); + } + DataCopy(solveWorkspace_[xBase], xMat, static_cast(elemCount)); + if (storeLToScratch) { + DataCopy(solveWorkspace_[lBase], akkMat, static_cast(elemCount)); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + return; + } +#endif + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + if (validRowCount > 0) { + DataCopy(aqk_[AOffset(b, hv, token, 0)], aqkMat, static_cast(validElemCount)); + if (storeLToAkk) { + DataCopy(akk_[AOffset(b, hv, token, 0)], akkMat, static_cast(validElemCount)); + } + } + DataCopy(solveWorkspace_[xBase], xMat, static_cast(elemCount)); + if (storeLToScratch) { + DataCopy(solveWorkspace_[lBase], akkMat, static_cast(elemCount)); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void CubeGemmSolveSub(GlobalTensor &tensorA, uint64_t baseA, uint64_t rowA, uint64_t colA, + GlobalTensor &tensorB, uint64_t baseB, uint64_t rowB, uint64_t colB, + GlobalTensor &tensorC, uint64_t baseC, uint64_t rowC, uint64_t colC, + uint32_t m, uint32_t n, uint32_t k) + { + using ElementA = float; + using ElementB = float; + using ElementC = float; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::RowMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using BlockMmad = + Catlass::Gemm::Block::BlockMmadTla; + Catlass::Arch::Resource resource; + auto layoutA = tla::MakeLayout(BT_, BT_); + auto layoutB = tla::MakeLayout(BT_, BT_); + auto layoutC = tla::MakeLayout(BT_, BT_); + auto tensorLayoutA = tla::MakeTensor(tensorA[baseA], layoutA, Catlass::Arch::PositionGM{}); + auto tensorLayoutB = tla::MakeTensor(tensorB[baseB], layoutB, Catlass::Arch::PositionGM{}); + auto tensorLayoutC = tla::MakeTensor(tensorC[baseC], layoutC, Catlass::Arch::PositionGM{}); + Catlass::GemmCoord shape{m, n, k}; + auto blockA = GetTile(tensorLayoutA, tla::MakeCoord(rowA, colA), tla::MakeShape(shape.m(), shape.k())); + auto blockB = GetTile(tensorLayoutB, tla::MakeCoord(rowB, colB), tla::MakeShape(shape.k(), shape.n())); + auto blockC = GetTile(tensorLayoutC, tla::MakeCoord(rowC, colC), tla::MakeShape(shape.m(), shape.n())); + BlockMmad blockMmad(resource); + blockMmad(blockA, blockB, blockC, shape); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + SetFlag(KDA_ARCH35_SOLVE_FIX_EVENT); + WaitFlag(KDA_ARCH35_SOLVE_FIX_EVENT); +#else + PipeBarrier(); +#endif + } + + __aicore__ inline void AddSolveTmpToX(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, bool storeAkk) + { + LocalTensor arena = vecBuf_.Get(); + LocalTensor xLocal = arena; + LocalTensor tmpLocal = arena[KDA_SOLVE_MATRIX_ELEMENTS]; + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X); + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP); + + DataCopy(xLocal, h_[xBase], KDA_SOLVE_MATRIX_ELEMENTS); + DataCopy(tmpLocal, h_[tmpBase], KDA_SOLVE_MATRIX_ELEMENTS); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + + Add(xLocal, xLocal, tmpLocal, KDA_SOLVE_MATRIX_ELEMENTS); + PipeBarrier(); + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(h_[xBase], xLocal, KDA_SOLVE_MATRIX_ELEMENTS); + if (storeAkk) { + DataCopy(akk_[AOffset(b, hv, start, 0)], xLocal, KDA_SOLVE_MATRIX_ELEMENTS); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void AddSolveTmpToXTail(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t curT, + bool storeAkk) + { + uint64_t elemCount = curT * KDA_SOLVE_BT; + LocalTensor arena = vecBuf_.Get(); + LocalTensor xLocal = arena; + LocalTensor tmpLocal = arena[KDA_SOLVE_MATRIX_ELEMENTS]; + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X); + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP); + + DataCopy(xLocal, h_[xBase], KDA_SOLVE_MATRIX_ELEMENTS); + DataCopy(tmpLocal, h_[tmpBase], KDA_SOLVE_MATRIX_ELEMENTS); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + + Add(xLocal, xLocal, tmpLocal, KDA_SOLVE_MATRIX_ELEMENTS); + PipeBarrier(); + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(h_[xBase], xLocal, KDA_SOLVE_MATRIX_ELEMENTS); + if (storeAkk) { + DataCopy(akk_[AOffset(b, hv, start, 0)], xLocal, static_cast(elemCount)); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void AddSolveTmpToXRows(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t curT, + uint64_t rowBegin, uint64_t rowEnd, bool storeAkk) + { + uint64_t rowCount = rowEnd - rowBegin; + if (rowCount == 0) { + return; + } + uint64_t validRowCount = rowBegin < curT ? curT - rowBegin : 0; + if (validRowCount > rowCount) { + validRowCount = rowCount; + } + uint64_t elemCount = rowCount * BT_; + uint64_t validElemCount = validRowCount * BT_; + LocalTensor arena = vecBuf_.Get(); + LocalTensor xLocal = arena; + LocalTensor tmpLocal = arena[elemCount]; + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X) + rowBegin * BT_; + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP) + rowBegin * BT_; + uint64_t token = start + rowBegin; + + DataCopy(xLocal, solveWorkspace_[xBase], static_cast(elemCount)); + DataCopy(tmpLocal, solveWorkspace_[tmpBase], static_cast(elemCount)); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + + Add(xLocal, xLocal, tmpLocal, static_cast(elemCount)); + PipeBarrier(); + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(solveWorkspace_[xBase], xLocal, static_cast(elemCount)); + if (storeAkk && validRowCount > 0) { + DataCopy(akk_[AOffset(b, hv, token, 0)], xLocal, static_cast(validElemCount)); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void AddSolveTmpToXDiagRows(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t rowBegin, uint64_t rowEnd, bool storeAkk) + { + uint64_t rowCount = rowEnd - rowBegin; + if (rowCount == 0) { + return; + } + uint64_t elemCount = rowCount * BT_; + LocalTensor arena = vecBuf_.Get(); + LocalTensor xLocal = arena; + LocalTensor tmpLocal = arena[elemCount]; + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X) + rowBegin * BT_; + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP) + rowBegin * BT_; + uint64_t token = start + rowBegin; + + DataCopy(xLocal, solveWorkspace_[xBase], static_cast(elemCount)); + DataCopy(tmpLocal, solveWorkspace_[tmpBase], static_cast(elemCount)); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + + for (uint64_t localRow = 0; localRow < rowCount; ++localRow) { + uint64_t row = rowBegin + localRow; + uint64_t col = (row / KDA_SOLVE_DIAG_BT) * KDA_SOLVE_DIAG_BT; + uint64_t offset = localRow * BT_ + col; + Add(xLocal[offset], xLocal[offset], tmpLocal[offset], KDA_SOLVE_DIAG_BT); + PipeBarrier(); + } + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(solveWorkspace_[xBase], xLocal, static_cast(elemCount)); + if (storeAkk) { + DataCopy(akk_[AOffset(b, hv, token, 0)], xLocal, static_cast(elemCount)); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void StoreSolveXRowsToAkk(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t curT, uint64_t rowBegin, uint64_t rowEnd) + { + uint64_t validRowCount = rowBegin < curT ? curT - rowBegin : 0; + uint64_t rowCount = rowEnd - rowBegin; + if (validRowCount > rowCount) { + validRowCount = rowCount; + } + if (validRowCount == 0) { + return; + } + uint64_t elemCount = validRowCount * BT_; + LocalTensor xLocal = vecBuf_.Get(); + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X) + rowBegin * BT_; + + DataCopy(xLocal, solveWorkspace_[xBase], static_cast(elemCount)); + SetFlag(mte2ToMte3Event_); + WaitFlag(mte2ToMte3Event_); + DataCopy(akk_[AOffset(b, hv, start + rowBegin, 0)], xLocal, static_cast(elemCount)); + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + } + + __aicore__ inline void ComputeAkkMergeCube(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start) + { + uint64_t aiBase = AOffset(b, hv, start, 0); + uint64_t negABase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X); + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP); + + for (uint32_t mergeSize = 2 * KDA_SOLVE_DIAG_BT; mergeSize <= BT_; mergeSize *= 2) { + uint32_t half = mergeSize / 2; + for (uint32_t block = 0; block < BT_; block += mergeSize) { + uint32_t lower = block + half; + CubeGemmSolveSub(akk_, aiBase, lower, lower, solveWorkspace_, negABase, lower, block, solveWorkspace_, + tmpBase, 0, 0, half, half, half); + CubeGemmSolveSub(solveWorkspace_, tmpBase, 0, 0, akk_, aiBase, block, block, akk_, aiBase, lower, block, + half, half, half); + } + } + } + + __aicore__ inline void ComputeAkkMergeCubeWorkspace(uint64_t b, uint64_t hv, uint64_t chunkIdx) + { + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X); + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP); + + for (uint32_t mergeSize = 2 * KDA_SOLVE_DIAG_BT; mergeSize <= BT_; mergeSize *= 2) { + uint32_t half = mergeSize / 2; + for (uint32_t block = 0; block < BT_; block += mergeSize) { + uint32_t lower = block + half; + CubeGemmSolveSub(solveWorkspace_, xBase, lower, lower, solveWorkspace_, xBase, lower, block, + solveWorkspace_, tmpBase, 0, 0, half, half, half); + CubeGemmSolveSub(solveWorkspace_, tmpBase, 0, 0, solveWorkspace_, xBase, block, block, solveWorkspace_, + xBase, lower, block, half, half, half); + } + } + } + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + template + __aicore__ inline void LoadSolveTile(LocalTensor dst, LocalTensor src) + { + static_assert(TILE_SIZE == 16 || TILE_SIZE == 32, "arch35 solve tile must be 16 or 32 rows"); + constexpr uint32_t rowFractals = TILE_SIZE / 16; + constexpr uint32_t columnFractals = TILE_SIZE / 8; + LoadData2DParamsV2 loadParams; + loadParams.mStartPosition = 0; + loadParams.kStartPosition = 0; + loadParams.mStep = rowFractals; + loadParams.kStep = columnFractals; + loadParams.srcStride = rowFractals; + loadParams.dstStride = rowFractals; + loadParams.ifTranspose = TRANSPOSE; + LoadData(dst, src, loadParams); + } + + __aicore__ inline void ComputeAkkMergeCubeWorkspaceArch35(uint64_t b, uint64_t hv, uint64_t chunkIdx) + { + (void)b; + (void)hv; + SetMMLayoutTransform(true); + + using Element = float; + using LayoutTag = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + using TileMmad = Catlass::Gemm::Tile::TileMmadTla; + + constexpr uint32_t maxTile = 32; + constexpr uint32_t tileSlotBytes = maxTile * maxTile * sizeof(Element); + constexpr uint32_t a0Slot = 0; + constexpr uint32_t b0Slot = 1; + constexpr uint32_t a1Slot = 2; + constexpr uint32_t b1Slot = 3; + constexpr uint32_t diag0Slot = 4; + constexpr uint32_t diag1Slot = 5; + constexpr uint32_t tmp0Slot = 6; + constexpr uint32_t tmp1Slot = 7; + + Catlass::Arch::Resource resource; + LocalTensor l1A0 = resource.l1Buf.template GetBufferByByte(a0Slot * tileSlotBytes); + LocalTensor l1B0 = resource.l1Buf.template GetBufferByByte(b0Slot * tileSlotBytes); + LocalTensor l1A1 = resource.l1Buf.template GetBufferByByte(a1Slot * tileSlotBytes); + LocalTensor l1B1 = resource.l1Buf.template GetBufferByByte(b1Slot * tileSlotBytes); + LocalTensor l1Diag0 = resource.l1Buf.template GetBufferByByte(diag0Slot * tileSlotBytes); + LocalTensor l1Diag1 = resource.l1Buf.template GetBufferByByte(diag1Slot * tileSlotBytes); + LocalTensor l1Tmp0 = resource.l1Buf.template GetBufferByByte(tmp0Slot * tileSlotBytes); + LocalTensor l1Tmp1 = resource.l1Buf.template GetBufferByByte(tmp1Slot * tileSlotBytes); + + LocalTensor l0A0 = resource.l0ABuf.template GetBufferByByte(a0Slot * tileSlotBytes); + LocalTensor l0B0 = resource.l0BBuf.template GetBufferByByte(b0Slot * tileSlotBytes); + LocalTensor l0A1 = resource.l0ABuf.template GetBufferByByte(a1Slot * tileSlotBytes); + LocalTensor l0B1 = resource.l0BBuf.template GetBufferByByte(b1Slot * tileSlotBytes); + LocalTensor l0A2 = resource.l0ABuf.template GetBufferByByte(diag0Slot * tileSlotBytes); + LocalTensor l0B2 = resource.l0BBuf.template GetBufferByByte(diag0Slot * tileSlotBytes); + LocalTensor l0A3 = resource.l0ABuf.template GetBufferByByte(diag1Slot * tileSlotBytes); + LocalTensor l0B3 = resource.l0BBuf.template GetBufferByByte(diag1Slot * tileSlotBytes); + LocalTensor l0C0 = resource.l0CBuf.template GetBufferByByte(a0Slot * tileSlotBytes); + LocalTensor l0C1 = resource.l0CBuf.template GetBufferByByte(a1Slot * tileSlotBytes); + LocalTensor l0C2 = resource.l0CBuf.template GetBufferByByte(diag0Slot * tileSlotBytes); + LocalTensor l0C3 = resource.l0CBuf.template GetBufferByByte(diag1Slot * tileSlotBytes); + + TileMmad tileMmad; + const uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X); + const uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP); + auto gmLayout = tla::MakeLayout(BT_, BT_); + auto tensorX = tla::MakeTensor(solveWorkspace_[xBase], gmLayout, Catlass::Arch::PositionGM{}); + auto tensorTmp = tla::MakeTensor(solveWorkspace_[tmpBase], gmLayout, Catlass::Arch::PositionGM{}); + + constexpr uint32_t tile16 = 16; + auto shape16 = tla::MakeShape(tile16, tile16); + auto blockA0 = GetTile(tensorX, tla::MakeCoord(16, 16), shape16); + auto blockB0 = GetTile(tensorX, tla::MakeCoord(16, 0), shape16); + auto blockA1 = GetTile(tensorX, tla::MakeCoord(48, 48), shape16); + auto blockB1 = GetTile(tensorX, tla::MakeCoord(48, 32), shape16); + auto blockDiag0 = GetTile(tensorX, tla::MakeCoord(0, 0), shape16); + auto blockDiag1 = GetTile(tensorX, tla::MakeCoord(32, 32), shape16); + auto blockTmp0 = GetTile(tensorTmp, tla::MakeCoord(0, 0), shape16); + auto blockTmp1 = GetTile(tensorTmp, tla::MakeCoord(16, 0), shape16); + auto blockOut0 = GetTile(tensorX, tla::MakeCoord(16, 0), shape16); + auto blockOut1 = GetTile(tensorX, tla::MakeCoord(48, 32), shape16); + + using CopyGmToL1A16 = typename TileCopy::template CopyGmToL1A; + using CopyGmToL1B16 = typename TileCopy::template CopyGmToL1B; + using CopyL0CToDst16 = typename TileCopy::template CopyL0CToDst; + CopyGmToL1A16 copyGmToL1A16; + CopyGmToL1B16 copyGmToL1B16; + CopyL0CToDst16 copyL0CToDst16; + + auto tensorL1A0 = + tla::MakeTensor(l1A0, tla::MakeLayout(tile16, tile16), Catlass::Arch::PositionL1{}); + auto tensorL1B0 = + tla::MakeTensor(l1B0, tla::MakeLayout(tile16, tile16), Catlass::Arch::PositionL1{}); + auto tensorL1A1 = + tla::MakeTensor(l1A1, tla::MakeLayout(tile16, tile16), Catlass::Arch::PositionL1{}); + auto tensorL1B1 = + tla::MakeTensor(l1B1, tla::MakeLayout(tile16, tile16), Catlass::Arch::PositionL1{}); + auto tensorL1Diag0 = tla::MakeTensor(l1Diag0, tla::MakeLayout(tile16, tile16), + Catlass::Arch::PositionL1{}); + auto tensorL1Diag1 = tla::MakeTensor(l1Diag1, tla::MakeLayout(tile16, tile16), + Catlass::Arch::PositionL1{}); + auto tensorL1Tmp0 = tla::MakeTensor(l1Tmp0, tla::MakeLayout(tile16, tile16), + Catlass::Arch::PositionL1{}); + auto tensorL1Tmp1 = tla::MakeTensor(l1Tmp1, tla::MakeLayout(tile16, tile16), + Catlass::Arch::PositionL1{}); + + auto tensorL0A0 = + tla::MakeTensor(l0A0, tla::MakeLayout(tile16, tile16), Catlass::Arch::PositionL0A{}); + auto tensorL0B0 = + tla::MakeTensor(l0B0, tla::MakeLayout(tile16, tile16), Catlass::Arch::PositionL0B{}); + auto tensorL0A1 = + tla::MakeTensor(l0A1, tla::MakeLayout(tile16, tile16), Catlass::Arch::PositionL0A{}); + auto tensorL0B1 = + tla::MakeTensor(l0B1, tla::MakeLayout(tile16, tile16), Catlass::Arch::PositionL0B{}); + auto tensorL0A2 = + tla::MakeTensor(l0A2, tla::MakeLayout(tile16, tile16), Catlass::Arch::PositionL0A{}); + auto tensorL0B2 = + tla::MakeTensor(l0B2, tla::MakeLayout(tile16, tile16), Catlass::Arch::PositionL0B{}); + auto tensorL0A3 = + tla::MakeTensor(l0A3, tla::MakeLayout(tile16, tile16), Catlass::Arch::PositionL0A{}); + auto tensorL0B3 = + tla::MakeTensor(l0B3, tla::MakeLayout(tile16, tile16), Catlass::Arch::PositionL0B{}); + auto tensorL0C0 = tla::MakeTensor(l0C0, tla::MakeLayoutL0C(tile16, tile16), Catlass::Arch::PositionL0C{}); + auto tensorL0C1 = tla::MakeTensor(l0C1, tla::MakeLayoutL0C(tile16, tile16), Catlass::Arch::PositionL0C{}); + auto tensorL0C2 = tla::MakeTensor(l0C2, tla::MakeLayoutL0C(tile16, tile16), Catlass::Arch::PositionL0C{}); + auto tensorL0C3 = tla::MakeTensor(l0C3, tla::MakeLayoutL0C(tile16, tile16), Catlass::Arch::PositionL0C{}); + uint32_t localRow16 = 0; + uint32_t localColumn16 = 0; + auto tileL0A0 = GetTile(tensorL0A0, tla::MakeCoord(localRow16, localColumn16), shape16); + auto tileL0B0 = GetTile(tensorL0B0, tla::MakeCoord(localRow16, localColumn16), shape16); + auto tileL0A1 = GetTile(tensorL0A1, tla::MakeCoord(localRow16, localColumn16), shape16); + auto tileL0B1 = GetTile(tensorL0B1, tla::MakeCoord(localRow16, localColumn16), shape16); + auto tileL0A2 = GetTile(tensorL0A2, tla::MakeCoord(localRow16, localColumn16), shape16); + auto tileL0B2 = GetTile(tensorL0B2, tla::MakeCoord(localRow16, localColumn16), shape16); + auto tileL0A3 = GetTile(tensorL0A3, tla::MakeCoord(localRow16, localColumn16), shape16); + auto tileL0B3 = GetTile(tensorL0B3, tla::MakeCoord(localRow16, localColumn16), shape16); + auto tileL0C0 = GetTile(tensorL0C0, tla::MakeCoord(localRow16, localColumn16), shape16); + auto tileL0C1 = GetTile(tensorL0C1, tla::MakeCoord(localRow16, localColumn16), shape16); + auto tileL0C2 = GetTile(tensorL0C2, tla::MakeCoord(localRow16, localColumn16), shape16); + auto tileL0C3 = GetTile(tensorL0C3, tla::MakeCoord(localRow16, localColumn16), shape16); + + copyGmToL1A16(tensorL1A0, blockA0); + copyGmToL1B16(tensorL1B0, blockB0); + copyGmToL1A16(tensorL1A1, blockA1); + copyGmToL1B16(tensorL1B1, blockB1); + copyGmToL1B16(tensorL1Diag0, blockDiag0); + copyGmToL1B16(tensorL1Diag1, blockDiag1); + SetFlag(0); + WaitFlag(0); + + LoadSolveTile(l0A0, l1A0); + LoadSolveTile(l0B0, l1B0); + LoadSolveTile(l0A1, l1A1); + LoadSolveTile(l0B1, l1B1); + SetFlag(0); + WaitFlag(0); + tileMmad(tileL0C0, tileL0A0, tileL0B0, tile16, tile16, tile16, true, 0b11); + SetFlag(0); + tileMmad(tileL0C1, tileL0A1, tileL0B1, tile16, tile16, tile16, true, 0b11); + SetFlag(1); + WaitFlag(0); + copyL0CToDst16(blockTmp0, tileL0C0, 0b11); + SetFlag(0); + WaitFlag(1); + copyL0CToDst16(blockTmp1, tileL0C1, 0b11); + SetFlag(1); + WaitFlag(0); + copyGmToL1A16(tensorL1Tmp0, blockTmp0); + WaitFlag(1); + copyGmToL1A16(tensorL1Tmp1, blockTmp1); + SetFlag(0); + WaitFlag(0); + + LoadSolveTile(l0A2, l1Tmp0); + LoadSolveTile(l0B2, l1Diag0); + LoadSolveTile(l0A3, l1Tmp1); + LoadSolveTile(l0B3, l1Diag1); + SetFlag(0); + WaitFlag(0); + tileMmad(tileL0C2, tileL0A2, tileL0B2, tile16, tile16, tile16, true, 0b11); + SetFlag(2); + tileMmad(tileL0C3, tileL0A3, tileL0B3, tile16, tile16, tile16, true, 0b11); + SetFlag(3); + WaitFlag(2); + copyL0CToDst16(blockOut0, tileL0C2, 0b11); + SetFlag(2); + WaitFlag(3); + copyL0CToDst16(blockOut1, tileL0C3, 0b11); + SetFlag(3); + + constexpr uint32_t tile32 = 32; + auto shape32 = tla::MakeShape(tile32, tile32); + auto blockA32 = GetTile(tensorX, tla::MakeCoord(32, 32), shape32); + auto blockB32 = GetTile(tensorX, tla::MakeCoord(32, 0), shape32); + auto blockDiag32 = GetTile(tensorX, tla::MakeCoord(0, 0), shape32); + auto blockTmp32 = GetTile(tensorTmp, tla::MakeCoord(0, 0), shape32); + auto blockOut32 = GetTile(tensorX, tla::MakeCoord(32, 0), shape32); + using CopyGmToL1A32 = typename TileCopy::template CopyGmToL1A; + using CopyGmToL1B32 = typename TileCopy::template CopyGmToL1B; + using CopyL0CToDst32 = typename TileCopy::template CopyL0CToDst; + CopyGmToL1A32 copyGmToL1A32; + CopyGmToL1B32 copyGmToL1B32; + CopyL0CToDst32 copyL0CToDst32; + + auto tensorL1A32 = + tla::MakeTensor(l1A0, tla::MakeLayout(tile32, tile32), Catlass::Arch::PositionL1{}); + auto tensorL1B32 = + tla::MakeTensor(l1B0, tla::MakeLayout(tile32, tile32), Catlass::Arch::PositionL1{}); + auto tensorL1Diag32 = tla::MakeTensor(l1Diag0, tla::MakeLayout(tile32, tile32), + Catlass::Arch::PositionL1{}); + auto tensorL1Tmp32 = tla::MakeTensor(l1Tmp0, tla::MakeLayout(tile32, tile32), + Catlass::Arch::PositionL1{}); + auto tensorL0A32 = + tla::MakeTensor(l0A0, tla::MakeLayout(tile32, tile32), Catlass::Arch::PositionL0A{}); + auto tensorL0B32 = + tla::MakeTensor(l0B0, tla::MakeLayout(tile32, tile32), Catlass::Arch::PositionL0B{}); + auto tensorL0C32 = tla::MakeTensor(l0C0, tla::MakeLayoutL0C(tile32, tile32), Catlass::Arch::PositionL0C{}); + uint32_t localRow32 = 0; + uint32_t localColumn32 = 0; + auto tileL0A32 = GetTile(tensorL0A32, tla::MakeCoord(localRow32, localColumn32), shape32); + auto tileL0B32 = GetTile(tensorL0B32, tla::MakeCoord(localRow32, localColumn32), shape32); + auto tileL0C32 = GetTile(tensorL0C32, tla::MakeCoord(localRow32, localColumn32), shape32); + + WaitFlag(2); + WaitFlag(3); + copyGmToL1A32(tensorL1A32, blockA32); + copyGmToL1B32(tensorL1B32, blockB32); + copyGmToL1B32(tensorL1Diag32, blockDiag32); + SetFlag(0); + WaitFlag(0); + LoadSolveTile(l0A0, l1A0); + LoadSolveTile(l0B0, l1B0); + SetFlag(0); + WaitFlag(0); + tileMmad(tileL0C32, tileL0A32, tileL0B32, tile32, tile32, tile32, true, 0b11); + SetFlag(0); + WaitFlag(0); + copyL0CToDst32(blockTmp32, tileL0C32, 0b11); + SetFlag(0); + WaitFlag(0); + copyGmToL1A32(tensorL1Tmp32, blockTmp32); + SetFlag(0); + WaitFlag(0); + LoadSolveTile(l0A0, l1Tmp0); + LoadSolveTile(l0B0, l1Diag0); + SetFlag(0); + WaitFlag(0); + tileMmad(tileL0C32, tileL0A32, tileL0B32, tile32, tile32, tile32, true, 0b11); + SetFlag(0); + WaitFlag(0); + copyL0CToDst32(blockOut32, tileL0C32, 0b11); + SetFlag(0); + WaitFlag(0); + SetMMLayoutTransform(false); + } +#endif + + __aicore__ inline void ComputeAkkMergeCubeWorkspaceDispatch(uint64_t b, uint64_t hv, uint64_t chunkIdx) + { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128) { + ComputeAkkMergeCubeWorkspace(b, hv, chunkIdx); + return; + } +#endif + ComputeAkkMergeCubeWorkspace(b, hv, chunkIdx); + } + + __aicore__ inline void ComputeAkkInverseMchFull(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start) + { + uint64_t aBase = AOffset(b, hv, start, 0); + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X); + uint64_t yBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_Y0); + uint64_t yNextBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_Y1); + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP); + + uint32_t diagBlocks = static_cast(BT_ / KDA_SOLVE_DIAG_BT); + for (uint32_t block = 0; block < diagBlocks; ++block) { + uint32_t off = block * KDA_SOLVE_DIAG_BT; + CubeGemmSolveSub(akk_, aBase, off, off, akk_, aBase, off, off, solveWorkspace_, yBase, off, off, + KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT); + } + for (uint32_t iter = 0; iter < KDA_SOLVE_DIAG_MCH_ITERS; ++iter) { + for (uint32_t block = 0; block < diagBlocks; ++block) { + uint32_t off = block * KDA_SOLVE_DIAG_BT; + CubeGemmSolveSub(solveWorkspace_, xBase, off, off, solveWorkspace_, yBase, off, off, solveWorkspace_, + tmpBase, off, off, KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT); + } + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(mchSyncDoneFlag_); + if (iter + 1 < KDA_SOLVE_DIAG_MCH_ITERS) { + for (uint32_t block = 0; block < diagBlocks; ++block) { + uint32_t off = block * KDA_SOLVE_DIAG_BT; + CubeGemmSolveSub(solveWorkspace_, yBase, off, off, solveWorkspace_, yBase, off, off, + solveWorkspace_, yNextBase, off, off, KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT, + KDA_SOLVE_DIAG_BT); + } + } + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_FIX>(mchSyncReadyFlag_); + if (iter + 1 < KDA_SOLVE_DIAG_MCH_ITERS) { + uint64_t oldYBase = yBase; + yBase = yNextBase; + yNextBase = oldYBase; + } + } + ComputeAkkMergeCube(b, hv, chunkIdx, start); + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(mchSyncDoneFlag_); + } + + __aicore__ inline void ScaleRowsByBeta(GlobalTensor &src, GlobalTensor &dst, uint64_t b, uint64_t hv, + uint64_t start, uint64_t rowBegin, uint64_t rowCount, uint64_t dim, + LocalTensor &betaLocal, LocalTensor &betaBrcb, + LocalTensor &matrixLocal, bool sourceSequenceMajor = false) + { + constexpr uint64_t vecElemsPerRepeat = 64; + constexpr uint64_t typedOffsetFloats = 20480; + constexpr uint64_t typedOffset = typedOffsetFloats * sizeof(float) / sizeof(T); + uint64_t elemCount = rowCount * dim; + uint64_t baseOffset = KVOffset(b, hv, start + rowBegin, 0, dim); + uint64_t sourceOffset = sourceSequenceMajor ? VInputOffset(b, hv, start + rowBegin, 0) : baseOffset; + uint64_t sourceStride = sourceSequenceMajor ? HV_ * dim : dim; + + if constexpr (IsSameType::value) { + CopyRowsIn(matrixLocal, src, sourceOffset, rowCount, dim, sourceStride); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + } else { + LocalTensor matrixTyped = vecBuf_.Get()[typedOffset]; + CopyRowsIn(matrixTyped, src, sourceOffset, rowCount, dim, sourceStride); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Cast(matrixLocal, matrixTyped, RoundMode::CAST_NONE, static_cast(elemCount)); + PipeBarrier(); + } + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + ApplyKdaRowScaleRegbase((__ubuf__ float *)reinterpret_cast(matrixLocal.GetPhyAddr()), + (__ubuf__ float *)reinterpret_cast(betaLocal.GetPhyAddr()), + static_cast(rowCount), static_cast(dim)); +#else + uint8_t repeatStride = static_cast(dim * sizeof(float) / 32); + for (uint64_t col = 0; col < dim; col += vecElemsPerRepeat) { + uint64_t mask = dim - col; + if (mask > vecElemsPerRepeat) { + mask = vecElemsPerRepeat; + } + Mul(matrixLocal[col], matrixLocal[col], betaBrcb, mask, static_cast(rowCount), + {1, 1, 0, repeatStride, repeatStride, 1}); + } + PipeBarrier(); +#endif + + if constexpr (IsSameType::value) { + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(dst[baseOffset], matrixLocal, static_cast(elemCount)); + } else { + LocalTensor matrixTyped = vecBuf_.Get()[typedOffset]; + Cast(matrixTyped, matrixLocal, RoundMode::CAST_RINT, static_cast(elemCount)); + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(dst[baseOffset], matrixTyped, static_cast(elemCount)); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void PrepareWuCubeInputs(uint64_t b, uint64_t hv, uint64_t start, uint64_t curT, + uint64_t subBlockIdx, uint64_t subBlockNum) + { + uint64_t rowsPerSubBlock = (curT + subBlockNum - 1) / subBlockNum; + uint64_t rowBegin = subBlockIdx * rowsPerSubBlock; + if (rowBegin >= curT) { + return; + } + uint64_t rowCount = curT - rowBegin; + if (rowCount > rowsPerSubBlock) { + rowCount = rowsPerSubBlock; + } + LocalTensor arena = vecBuf_.Get(); + LocalTensor betaLocal = arena; + LocalTensor betaBrcb = arena[KDA_SOLVE_BT]; + LocalTensor matrixLocal = arena[KDA_SOLVE_BT + 512]; +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + // Keep each arch35 Cast/regbase panel within an 8K-element UB instruction span. + constexpr uint64_t maxScaleElements = 8192; + if (rowCount * K_ > maxScaleElements || rowCount * V_ > maxScaleElements) { + constexpr uint64_t tileRows = 16; + for (uint64_t tileRow = 0; tileRow < rowCount; tileRow += tileRows) { + uint64_t tileCount = rowCount - tileRow; + if (tileCount > tileRows) { + tileCount = tileRows; + } + LoadBetaAsFloatRow(BetaOffset(b, hv, start + rowBegin + tileRow), betaLocal, tileCount); + ScaleRowsByBeta(w_, w_, b, hv, start, rowBegin + tileRow, tileCount, K_, betaLocal, betaBrcb, + matrixLocal); + ScaleRowsByBeta(v_, vNew_, b, hv, start, rowBegin + tileRow, tileCount, V_, betaLocal, betaBrcb, + matrixLocal, inputSequenceMajor_); + } + return; + } +#endif + LoadBetaAsFloatRow(BetaOffset(b, hv, start + rowBegin), betaLocal, rowCount); +#if !defined(__CCE_AICORE__) || __CCE_AICORE__ != 310 + Brcb(betaBrcb, betaLocal, static_cast((rowCount + 7) / 8), {1, 8}); + PipeBarrier(); +#endif + ScaleRowsByBeta(w_, w_, b, hv, start, rowBegin, rowCount, K_, betaLocal, betaBrcb, matrixLocal); + ScaleRowsByBeta(v_, vNew_, b, hv, start, rowBegin, rowCount, V_, betaLocal, betaBrcb, matrixLocal, + inputSequenceMajor_); + } + + __aicore__ inline void FinalizePrepareIntermediates(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t curT, uint64_t subBlockIdx, uint64_t subBlockNum) + { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + constexpr bool qgAlreadyStored = SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128; +#else + constexpr bool qgAlreadyStored = false; +#endif + constexpr uint64_t tileRows = 32; + // Keep tail rows on the same AIV that owns their padded solve rows. Splitting by curT would + // move short-tail export to AIV1 while AIV0 is still writing the solved matrix. + const uint64_t rowBegin = (BT_ * subBlockIdx) / subBlockNum; + uint64_t rowEnd = (BT_ * (subBlockIdx + 1)) / subBlockNum; + if (rowEnd > curT) { + rowEnd = curT; + } + if (rowBegin >= rowEnd) { + return; + } + for (uint64_t tileRow = rowBegin; tileRow < rowEnd; tileRow += tileRows) { + const uint64_t rows = (rowEnd - tileRow) > tileRows ? tileRows : (rowEnd - tileRow); + const uint64_t matrixElems = rows * BT_; + const uint64_t qgElems = rows * K_; + LocalTensor arena = vecBuf_.Get(); + LocalTensor aqkLocal = arena; + LocalTensor akkLocal = arena[matrixElems]; + const uint64_t typedOffset = (2 * matrixElems + qgElems) * sizeof(float) / sizeof(T); + LocalTensor typedBase = vecBuf_.Get()[typedOffset]; + LocalTensor aqkTyped = typedBase; + LocalTensor akkTyped = typedBase[matrixElems]; + LocalTensor qgTyped = typedBase[2 * matrixElems]; + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128) { + const uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X) + tileRow * BT_; + CopyVectorIn(akkLocal, solveWorkspace_, xBase, matrixElems); + } else { + CopyVectorIn(aqkLocal, aqk_, AOffset(b, hv, start + tileRow, 0), matrixElems); + CopyVectorIn(akkLocal, akk_, AOffset(b, hv, start + tileRow, 0), matrixElems); + } +#else + CopyVectorIn(aqkLocal, aqk_, AOffset(b, hv, start + tileRow, 0), matrixElems); + CopyVectorIn(akkLocal, akk_, AOffset(b, hv, start + tileRow, 0), matrixElems); +#endif + if constexpr (!qgAlreadyStored) { + CopyVectorIn(qgTyped, qg_, KVOffset(b, hv, start + tileRow, 0, K_), qgElems); + } + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (!(SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128)) { + Muls(aqkLocal, aqkLocal, scale_, static_cast(matrixElems)); + PipeBarrier(); + } +#else + Muls(aqkLocal, aqkLocal, scale_, static_cast(matrixElems)); + PipeBarrier(); +#endif +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (!(SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128)) { + ClampFp32ToOutputType(aqkLocal, static_cast(matrixElems)); + } +#else + ClampFp32ToOutputType(aqkLocal, static_cast(matrixElems)); +#endif + ClampFp32ToOutputType(akkLocal, static_cast(matrixElems)); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (!(SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128)) { + Cast(aqkTyped, aqkLocal, RoundMode::CAST_RINT, static_cast(matrixElems)); + } +#else + Cast(aqkTyped, aqkLocal, RoundMode::CAST_RINT, static_cast(matrixElems)); +#endif + Cast(akkTyped, akkLocal, RoundMode::CAST_RINT, static_cast(matrixElems)); + PipeBarrier(); + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (!(SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128)) { + CopyVectorOut(o_, AOffset(b, hv, start + tileRow, 0), aqkTyped, matrixElems); + } +#else + CopyVectorOut(o_, AOffset(b, hv, start + tileRow, 0), aqkTyped, matrixElems); +#endif + CopyVectorOut(u_, AOffset(b, hv, start + tileRow, 0), akkTyped, matrixElems); + if constexpr (!qgAlreadyStored) { + CopyVectorOut(kg_, KVOffset(b, hv, start + tileRow, 0, K_), qgTyped, qgElems); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + } + + __aicore__ inline bool ResolveFlatChunk(uint64_t task, uint64_t &seq, uint64_t &b, uint64_t &h, uint64_t &hv, + uint64_t &chunkIdx, uint64_t &start, uint64_t &end) + { + hv = task % HV_; + uint64_t flatChunk = task / HV_; + if (!isVarLen_) { + seq = flatChunk / NT_; + b = seq; + chunkIdx = flatChunk % NT_; + start = chunkIdx * BT_; + end = start + BT_; + if (end > T_) { + end = T_; + } + } else { + if (!KdaVarlen::ResolveChunkRange(cuSeqlensAddr_, chunkIndicesAddr_, N_, T_, BT_, flatChunk, seq, start, + end)) { + return false; + } + b = 0; + chunkIdx = flatChunk; + } + h = hv / (HV_ / H_); + return start < end; + } + + __aicore__ inline void ProcessChunkPreAiv(uint64_t b, uint64_t h, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t end, uint64_t subBlockIdx, uint64_t subBlockNum) + { + if constexpr (IsSameType::value) { + ProcessChunkPreAivFp32(b, h, hv, chunkIdx, start, end, subBlockIdx, subBlockNum); + } + } + + template + __aicore__ inline void JoinAivMte3() + { + if constexpr (CORE_TYPE == AscendC::AIV) { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if (!isAivOnly_) { + if (!headPairMode_) { + Catlass::Arch::CrossCoreBarrier<0x1, PIPE_MTE3>(); + } + PipeBarrier(); + } +#endif + } + } + + template + __aicore__ inline void RunAicAfterBothAivReady(uint64_t subBlockIdx, uint64_t subBlockNum) + { + if constexpr (CORE_TYPE == AscendC::AIV) { + (void)subBlockIdx; + (void)subBlockNum; + JoinAivMte3(); + if constexpr (SAFE_GATE) { + Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(syncReadyFlag_); + Catlass::Arch::CrossCoreWaitFlag(syncDoneFlag_); + } else { + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_MTE3>(mchSyncReadyFlag_); + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE2>(mchSyncDoneFlag_); + } + } + } + + template + __aicore__ inline void SignalAicSolveReady() + { + if constexpr (CORE_TYPE == AscendC::AIV) { + JoinAivMte3(); + Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(syncReadyFlag_); + } + } + + template + __aicore__ inline void WaitAicSolveDone() + { + if constexpr (CORE_TYPE == AscendC::AIV) { + Catlass::Arch::CrossCoreWaitFlag(syncDoneFlag_); + } + } + + template + __aicore__ inline void SignalPostWuReady() + { + if constexpr (CORE_TYPE == AscendC::AIV) { + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_MTE3>(postWuReadyFlag_); + } + } + + __aicore__ inline void ProcessChunkPreAivFp32(uint64_t b, uint64_t h, uint64_t hv, uint64_t chunkIdx, + uint64_t start, uint64_t end, uint64_t subBlockIdx, + uint64_t subBlockNum, bool deferSafeSolve = false, + bool waitPendingSafeSolve = false, uint64_t scoreLane = 0, + bool pairHeads = false) + { + uint64_t curT = end - start; + if (curT == 0) { + return; + } + if constexpr (IsSameType::value) { + return; + } + + if (K_ < 16) { + return; + } + MaterializeRawGateChunkArch35(b, hv, start, curT); + bool usePostWuCube = UsePostWuCube(curT); + bool useAkkCubeSolve = UseAkkCubeSolve(curT); + uint64_t solveRowBegin = 0; + uint64_t solveRowEnd = 0; + GetSolveRowRange(BT_, subBlockIdx, subBlockNum, solveRowBegin, solveRowEnd); + // Safe-gate score factors need the bounded 16-row reference span. + // A single 64-row reference loses BF16 dynamic range for valid gates. + const bool useFullChunkScore = false; + uint64_t scoreBlockSize = useFullChunkScore ? curT : ScoreRefBlockSize(); + uint64_t scoreBlockCount = (curT + scoreBlockSize - 1) / scoreBlockSize; + uint64_t pipelineBlockCount = useFullChunkScore ? scoreBlockCount : + (scoreBlockCount + KDA_SCORE_QUEUE_DEPTH - 1) / + KDA_SCORE_QUEUE_DEPTH * KDA_SCORE_QUEUE_DEPTH; + const bool useDirectScoreUb = +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + KDA_ARCH35_ENABLE_DIRECT_SCORE_UB && pairHeads && curT == 64 && scoreBlockCount == 2; +#else + false; +#endif + bool firstSolveRowsPrepared = false; + for (uint64_t block = 0; block < pipelineBlockCount; ++block) { + if (block < scoreBlockCount) { + uint64_t rowBegin = block * scoreBlockSize; + uint64_t rowCount = useFullChunkScore ? curT - rowBegin : ScoreRowBlockCount(curT, rowBegin); + uint64_t refToken = ScoreRefToken(start, curT, rowBegin, rowCount); + uint64_t queueSlot = + useFullChunkScore ? activeSolveSlot_ / KDA_SCORE_LANES : block % KDA_SCORE_QUEUE_DEPTH; + uint64_t scoreSlot = ScoreScratchSlot(queueSlot, scoreLane, pairHeads); + PrepareGateProducts(b, h, hv, start, curT, subBlockIdx, subBlockNum, true, refToken, + rowBegin + rowCount, true, scoreSlot, rowBegin, rowCount); + } + JoinAivMte3(); + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_MTE3>(scoreReadyFlag_); + if (block > 0) { + if constexpr (SAFE_GATE) { + if (waitPendingSafeSolve) { + WaitAicSolveDone(); + waitPendingSafeSolve = false; + } + } + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE2>(scoreDoneFlag_); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128) { + if (useDirectScoreUb && useAkkCubeSolve && block == 1) { + ProcessDirectScoreSolveRowsArch35(b, hv, chunkIdx, start, 0, 0); + firstSolveRowsPrepared = true; + } else if (pairHeads && useAkkCubeSolve && block == 1) { + uint64_t firstRowBegin = 0; + uint64_t firstRowEnd = 0; + GetSolveRowRange(BT_, 0, KDA_SCORE_LANES, firstRowBegin, firstRowEnd); + PrepareAqkAkkSolveInputRows(b, hv, chunkIdx, start, curT, firstRowBegin, firstRowEnd, false, + false); + firstSolveRowsPrepared = true; + } + } +#endif + } + } + bool fusedScoreWriteback = false; +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + fusedScoreWriteback = SAFE_GATE && BT_ == 64 && K_ == 128 && V_ == 128; +#endif + if (!fusedScoreWriteback) { + // The final score MMAD only consumes scoreWorkspace_. Run the + // independent gate writeback while AIC drains its MMAD/Fixpipe path. + PrepareGateProducts(b, h, hv, start, curT, subBlockIdx, subBlockNum); + } + if (pipelineBlockCount > 0) { + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE2>(scoreDoneFlag_); + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if (useDirectScoreUb && useAkkCubeSolve) { + ProcessDirectScoreSolveRowsArch35(b, hv, chunkIdx, start, KDA_DIRECT_SCORE_ROWS, 1); + } +#endif + if constexpr (SAFE_GATE) { + if (waitPendingSafeSolve) { + WaitAicSolveDone(); + waitPendingSafeSolve = false; + } + } + if (useAkkCubeSolve) { + bool fullChunk = curT == BT_; + if constexpr (SAFE_GATE) { + if (pairHeads) { + if (!useDirectScoreUb) { + uint64_t firstRowPart = firstSolveRowsPrepared ? 1 : 0; + for (uint64_t rowPart = firstRowPart; rowPart < KDA_SCORE_LANES; ++rowPart) { + uint64_t pairRowBegin = 0; + uint64_t pairRowEnd = 0; + GetSolveRowRange(BT_, rowPart, KDA_SCORE_LANES, pairRowBegin, pairRowEnd); + PrepareAqkAkkSolveInputRows(b, hv, chunkIdx, start, curT, pairRowBegin, pairRowEnd, false, + false); + } + } + } else { + PrepareAqkAkkSolveInputRows(b, hv, chunkIdx, start, curT, solveRowBegin, solveRowEnd, false, false); + } + if (deferSafeSolve) { + SignalAicSolveReady(); + return; + } + RunAicAfterBothAivReady(subBlockIdx, subBlockNum); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (!(SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128)) { + StoreSolveXRowsToAkk(b, hv, chunkIdx, start, curT, solveRowBegin, solveRowEnd); + } +#else + StoreSolveXRowsToAkk(b, hv, chunkIdx, start, curT, solveRowBegin, solveRowEnd); +#endif + } else { + PrepareAqkAkkSolveInputRows(b, hv, chunkIdx, start, curT, solveRowBegin, solveRowEnd, fullChunk, false); + if (!fullChunk) { + RunAicAfterBothAivReady(subBlockIdx, subBlockNum); + StoreSolveXRowsToAkk(b, hv, chunkIdx, start, curT, solveRowBegin, solveRowEnd); + } else { + uint32_t solveIters = KDA_SOLVE_DIAG_MCH_ITERS; + RunAicAfterBothAivReady(subBlockIdx, subBlockNum); + for (uint32_t iter = 0; iter < solveIters; ++iter) { + AddSolveTmpToXDiagRows(b, hv, chunkIdx, start, solveRowBegin, solveRowEnd, + iter + 1 == solveIters); + RunAicAfterBothAivReady(subBlockIdx, subBlockNum); + } + } + } + } + // Host validation guarantees every accepted shape has enough workspace for this cube path. +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128) { + if (HV_ % KDA_SCORE_LANES != 0) { + PrepareWuCubeInputs(b, hv, start, curT, subBlockIdx, subBlockNum); + } + } else { + PrepareWuCubeInputs(b, hv, start, curT, subBlockIdx, subBlockNum); + } +#else + PrepareWuCubeInputs(b, hv, start, curT, subBlockIdx, subBlockNum); +#endif + FinalizePrepareIntermediates(b, hv, chunkIdx, start, curT, subBlockIdx, subBlockNum); + } + + __aicore__ inline void FinishDeferredSafeChunk(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t end, uint64_t subBlockIdx, uint64_t subBlockNum) + { + uint64_t curT = end - start; + uint64_t solveRowBegin = 0; + uint64_t solveRowEnd = 0; + GetSolveRowRange(BT_, subBlockIdx, subBlockNum, solveRowBegin, solveRowEnd); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (!(SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128)) { + StoreSolveXRowsToAkk(b, hv, chunkIdx, start, curT, solveRowBegin, solveRowEnd); + } +#else + StoreSolveXRowsToAkk(b, hv, chunkIdx, start, curT, solveRowBegin, solveRowEnd); +#endif +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128) { + if (HV_ % KDA_SCORE_LANES != 0) { + PrepareWuCubeInputs(b, hv, start, curT, subBlockIdx, subBlockNum); + } + } else { + PrepareWuCubeInputs(b, hv, start, curT, subBlockIdx, subBlockNum); + } +#else + PrepareWuCubeInputs(b, hv, start, curT, subBlockIdx, subBlockNum); +#endif + FinalizePrepareIntermediates(b, hv, chunkIdx, start, curT, subBlockIdx, subBlockNum); + } + + __aicore__ inline void FinishDeferredSafeChunkPair(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t end) + { + FinishDeferredSafeChunk(b, hv, chunkIdx, start, end, 0, 1); + } + + __aicore__ inline void ProcessChunkPreAic(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t end) + { + if constexpr (IsSameType::value) { + ProcessChunkPreAicFp32(b, hv, chunkIdx, start, end); + } + } + + __aicore__ inline void ProcessChunkPreAicFp32(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t end) + { + uint64_t curT = end - start; + if (curT == 0 || K_ < 16) { + return; + } + uint64_t scoreBlockSize = ScoreRefBlockSize(); + uint64_t scoreBlockCount = (curT + scoreBlockSize - 1) / scoreBlockSize; + uint64_t pipelineBlockCount = + (scoreBlockCount + KDA_SCORE_QUEUE_DEPTH - 1) / KDA_SCORE_QUEUE_DEPTH * KDA_SCORE_QUEUE_DEPTH; + for (uint64_t block = 0; block < pipelineBlockCount; ++block) { + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_FIX>(scoreReadyFlag_); + if (block < scoreBlockCount) { + uint64_t rowBegin = block * scoreBlockSize; + uint64_t rowCount = ScoreRowBlockCount(curT, rowBegin); + ComputeRawAqkAkkCubeBlock(b, hv, chunkIdx, start, curT, rowBegin, rowCount, true, + block % KDA_SCORE_QUEUE_DEPTH, rowBegin + rowCount); + } + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(scoreDoneFlag_); + } + bool usePostWuCube = UsePostWuCube(curT); + bool useAkkCubeSolve = UseAkkCubeSolve(curT); + if (useAkkCubeSolve) { + if constexpr (SAFE_GATE) { + Catlass::Arch::CrossCoreWaitFlag(syncReadyFlag_); + ComputeAkkMergeCubeWorkspaceDispatch(b, hv, chunkIdx); + Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(syncDoneFlag_); + } else { + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_FIX>(mchSyncReadyFlag_); + if (curT == BT_) { + ComputeAkkInverseMchFull(b, hv, chunkIdx, start); + } else { + ComputeAkkMergeCubeWorkspace(b, hv, chunkIdx); + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(mchSyncDoneFlag_); + } + } + } + (void)usePostWuCube; + (void)chunkIdx; + } + + __aicore__ inline void ProcessChunkPreAicHeadPairFp32(uint64_t b, uint64_t hvBase, uint64_t chunkIdx, + uint64_t start, uint64_t end, uint64_t localTaskIdx) + { + uint64_t curT = end - start; + if (curT == 0 || K_ < 16) { + return; + } + const bool useFullChunkScore = false; + uint64_t scoreBlockSize = useFullChunkScore ? curT : ScoreRefBlockSize(); + uint64_t scoreBlockCount = (curT + scoreBlockSize - 1) / scoreBlockSize; + uint64_t pipelineBlockCount = useFullChunkScore ? scoreBlockCount : + (scoreBlockCount + KDA_SCORE_QUEUE_DEPTH - 1) / + KDA_SCORE_QUEUE_DEPTH * KDA_SCORE_QUEUE_DEPTH; + for (uint64_t block = 0; block < pipelineBlockCount; ++block) { + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_FIX>(scoreReadyFlag_); + if (block < scoreBlockCount) { + uint64_t rowBegin = block * scoreBlockSize; + uint64_t rowCount = useFullChunkScore ? curT - rowBegin : ScoreRowBlockCount(curT, rowBegin); + uint64_t queueSlot = + useFullChunkScore ? localTaskIdx % KDA_SCORE_QUEUE_DEPTH : block % KDA_SCORE_QUEUE_DEPTH; +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + bool directScoreDispatched = false; + if constexpr (SAFE_GATE) { + if (KDA_ARCH35_ENABLE_DIRECT_SCORE_UB && curT == 64 && rowCount == 32) { + uint64_t scoreSlotBase = ScoreScratchSlot(queueSlot, 0, true); + if (rowBegin == 0) { + ComputeRawAqkAkkCubeStableHeadPairDirectUbArch35<32>(rowBegin, scoreSlotBase, + static_cast(block)); + } else { + ComputeRawAqkAkkCubeStableHeadPairDirectUbArch35<64>(rowBegin, scoreSlotBase, + static_cast(block)); + } + directScoreDispatched = true; + } + } + if (!directScoreDispatched) +#endif + { + for (uint64_t lane = 0; lane < KDA_SCORE_LANES; ++lane) { + uint64_t hv = hvBase + lane; + uint64_t scoreSlot = ScoreScratchSlot(queueSlot, lane, true); + activeSolveSlot_ = + (localTaskIdx % (KDA_SOLVE_PIPELINE_DEPTH / KDA_SCORE_LANES)) * KDA_SCORE_LANES + lane; + ComputeRawAqkAkkCubeBlock(b, hv, chunkIdx, start, curT, rowBegin, rowCount, true, scoreSlot, + rowBegin + rowCount); + } + } + } + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(scoreDoneFlag_); + } + + if (UseAkkCubeSolve(curT)) { + Catlass::Arch::CrossCoreWaitFlag(syncReadyFlag_); + for (uint64_t lane = 0; lane < KDA_SCORE_LANES; ++lane) { + activeSolveSlot_ = + (localTaskIdx % (KDA_SOLVE_PIPELINE_DEPTH / KDA_SCORE_LANES)) * KDA_SCORE_LANES + lane; + ComputeAkkMergeCubeWorkspaceDispatch(b, hvBase + lane, chunkIdx); + } + Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(syncDoneFlag_); + } + } + + __aicore__ inline bool ResolveFlatChunkForHv(uint64_t flatChunk, uint64_t hv, uint64_t &seq, uint64_t &b, + uint64_t &h, uint64_t &chunkIdx, uint64_t &start, uint64_t &end) + { + if (!isVarLen_) { + seq = flatChunk / NT_; + b = seq; + chunkIdx = flatChunk % NT_; + start = chunkIdx * BT_; + end = start + BT_; + if (end > T_) { + end = T_; + } + } else { + if (!KdaVarlen::ResolveChunkRange(cuSeqlensAddr_, chunkIndicesAddr_, N_, T_, BT_, flatChunk, seq, start, + end)) { + return false; + } + b = 0; + chunkIdx = flatChunk; + } + h = hv / (HV_ / H_); + return start < end; + } + + __aicore__ inline void ProcessPreAivHeadPair() + { + const uint64_t subBlockIdx = static_cast(GetSubBlockIdx()); + const uint64_t subBlockNum = static_cast(GetSubBlockNum()); + const uint64_t coreNum = usedCoreNum_; + const uint64_t coreIdx = static_cast(GetBlockIdx()) / subBlockNum; + const uint64_t chunkCount = isVarLen_ ? NT_ : B_ * NT_; + const uint64_t headWindows = HV_ / KDA_SCORE_LANES; + const uint64_t taskNum = chunkCount * headWindows; + bool pendingValid = false; + uint64_t pendingB = 0; + uint64_t pendingHv = 0; + uint64_t pendingChunkIdx = 0; + uint64_t pendingStart = 0; + uint64_t pendingEnd = 0; + uint64_t pendingSlot = 0; + uint64_t localTaskIdx = 0; + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (SAFE_GATE && COMPILE_BT == 64 && COMPILE_K == 128 && COMPILE_V == 128) { + if (KDA_ARCH35_ENABLE_DIRECT_SCORE_UB) { + InitializeDirectScoreUbArch35(); + } + } +#endif + + for (uint64_t task = coreIdx; task < taskNum; task += coreNum, ++localTaskIdx) { + uint64_t flatChunk = task / headWindows; + uint64_t hv = (task % headWindows) * KDA_SCORE_LANES + subBlockIdx; + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (!ResolveFlatChunkForHv(flatChunk, hv, seq, b, h, chunkIdx, start, end)) { + continue; + } + (void)seq; + uint64_t currentSlot = + (localTaskIdx % (KDA_SOLVE_PIPELINE_DEPTH / KDA_SCORE_LANES)) * KDA_SCORE_LANES + subBlockIdx; + activeSolveSlot_ = currentSlot; + bool deferSolve = UseAkkCubeSolve(end - start); + ProcessChunkPreAivFp32(b, h, hv, chunkIdx, start, end, 0, 1, deferSolve, pendingValid, subBlockIdx, true); + if (pendingValid) { + activeSolveSlot_ = pendingSlot; + FinishDeferredSafeChunkPair(pendingB, pendingHv, pendingChunkIdx, pendingStart, pendingEnd); + if (fusePostWu_) { + SignalPostWuReady(); + } + } + pendingValid = deferSolve; + if (pendingValid) { + pendingB = b; + pendingHv = hv; + pendingChunkIdx = chunkIdx; + pendingStart = start; + pendingEnd = end; + pendingSlot = currentSlot; + } + } + if (pendingValid) { + WaitAicSolveDone(); + activeSolveSlot_ = pendingSlot; + FinishDeferredSafeChunkPair(pendingB, pendingHv, pendingChunkIdx, pendingStart, pendingEnd); + if (fusePostWu_) { + SignalPostWuReady(); + } + } + } + + __aicore__ inline void ProcessPreAicHeadPair() + { + const uint64_t chunkCount = isVarLen_ ? NT_ : B_ * NT_; + const uint64_t headWindows = HV_ / KDA_SCORE_LANES; + const uint64_t taskNum = chunkCount * headWindows; + const uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + uint64_t localTaskIdx = 0; + for (uint64_t task = GetBlockIdx(); task < taskNum; task += coreNum, ++localTaskIdx) { + uint64_t flatChunk = task / headWindows; + uint64_t hvBase = (task % headWindows) * KDA_SCORE_LANES; + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveFlatChunkForHv(flatChunk, hvBase, seq, b, h, chunkIdx, start, end)) { + (void)seq; + (void)h; + ProcessChunkPreAicHeadPairFp32(b, hvBase, chunkIdx, start, end, localTaskIdx); + } + } + } + + template + __aicore__ inline void ProcessPreAicHeadPairFused(PostWuOp &postWu) + { + const uint64_t chunkCount = isVarLen_ ? NT_ : B_ * NT_; + const uint64_t headWindows = HV_ / KDA_SCORE_LANES; + const uint64_t taskNum = chunkCount * headWindows; + const uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + uint64_t batchB[KDA_POST_QUEUE_STORAGE]; + uint64_t batchHvBase[KDA_POST_QUEUE_STORAGE]; + uint64_t batchStart[KDA_POST_QUEUE_STORAGE]; + uint64_t batchEnd[KDA_POST_QUEUE_STORAGE]; + uint16_t batchCount = 0; + uint64_t localTaskIdx = 0; + + for (uint64_t task = GetBlockIdx(); task < taskNum; task += coreNum, ++localTaskIdx) { + uint64_t flatChunk = task / headWindows; + uint64_t hvBase = (task % headWindows) * KDA_SCORE_LANES; + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (!ResolveFlatChunkForHv(flatChunk, hvBase, seq, b, h, chunkIdx, start, end)) { + continue; + } + (void)seq; + (void)h; + ProcessChunkPreAicHeadPairFp32(b, hvBase, chunkIdx, start, end, localTaskIdx); + batchB[batchCount] = b; + batchHvBase[batchCount] = hvBase; + batchStart[batchCount] = start; + batchEnd[batchCount] = end; + ++batchCount; + + if (batchCount == KDA_POST_QUEUE_STORAGE) { + for (uint16_t i = 0; i < KDA_POST_QUEUE_DEPTH; ++i) { + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE2>(postWuReadyFlag_); + } + postWu.ProcessPreparedHeadPairBatchArch35(batchB, batchHvBase, batchStart, batchEnd, + KDA_POST_QUEUE_DEPTH); + batchB[0] = batchB[KDA_POST_QUEUE_DEPTH]; + batchHvBase[0] = batchHvBase[KDA_POST_QUEUE_DEPTH]; + batchStart[0] = batchStart[KDA_POST_QUEUE_DEPTH]; + batchEnd[0] = batchEnd[KDA_POST_QUEUE_DEPTH]; + batchCount = 1; + } + } + + if (batchCount > 0) { + for (uint16_t i = 0; i < batchCount; ++i) { + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE2>(postWuReadyFlag_); + } + postWu.ProcessPreparedHeadPairBatchArch35(batchB, batchHvBase, batchStart, batchEnd, batchCount); + } + } + + __aicore__ inline void ProcessPreAiv() + { + if constexpr (IsSameType::value) { + isAivOnly_ = true; + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (SAFE_GATE && !IsSameType::value) { + if (headPairMode_ && !isAivOnly_) { + ProcessPreAivHeadPair(); + return; + } + } +#endif + uint64_t subBlockNum = isAivOnly_ ? 1 : static_cast(GetSubBlockNum()); + if (subBlockNum == 0) { + return; + } + uint64_t subBlockIdx = isAivOnly_ ? 0 : static_cast(GetSubBlockIdx()); + uint64_t coreNum = isAivOnly_ ? static_cast(GetBlockNum()) : usedCoreNum_; + uint64_t coreIdx = + isAivOnly_ ? static_cast(GetBlockIdx()) : static_cast(GetBlockIdx()) / subBlockNum; + uint64_t taskNum = static_cast((isVarLen_ ? NT_ : B_ * NT_) * HV_); + if constexpr (SAFE_GATE && !IsSameType::value) { + bool pendingValid = false; + uint64_t pendingB = 0; + uint64_t pendingHv = 0; + uint64_t pendingChunkIdx = 0; + uint64_t pendingStart = 0; + uint64_t pendingEnd = 0; + uint64_t pendingSlot = 0; + uint64_t localTaskIdx = 0; + for (uint64_t task = coreIdx; task < taskNum; task += coreNum, ++localTaskIdx) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (!ResolveFlatChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + continue; + } + (void)seq; + uint64_t currentSlot = localTaskIdx % KDA_SOLVE_PIPELINE_DEPTH; + activeSolveSlot_ = currentSlot; + bool deferSolve = UseAkkCubeSolve(end - start); + if (!deferSolve && pendingValid) { + WaitAicSolveDone(); + activeSolveSlot_ = pendingSlot; + FinishDeferredSafeChunk(pendingB, pendingHv, pendingChunkIdx, pendingStart, pendingEnd, subBlockIdx, + subBlockNum); + pendingValid = false; + activeSolveSlot_ = currentSlot; + } + ProcessChunkPreAivFp32(b, h, hv, chunkIdx, start, end, subBlockIdx, subBlockNum, deferSolve, + pendingValid); + if (pendingValid) { + activeSolveSlot_ = pendingSlot; + FinishDeferredSafeChunk(pendingB, pendingHv, pendingChunkIdx, pendingStart, pendingEnd, subBlockIdx, + subBlockNum); + } + pendingValid = deferSolve; + if (pendingValid) { + pendingB = b; + pendingHv = hv; + pendingChunkIdx = chunkIdx; + pendingStart = start; + pendingEnd = end; + pendingSlot = currentSlot; + } + } + if (pendingValid) { + WaitAicSolveDone(); + activeSolveSlot_ = pendingSlot; + FinishDeferredSafeChunk(pendingB, pendingHv, pendingChunkIdx, pendingStart, pendingEnd, subBlockIdx, + subBlockNum); + } + return; + } + for (uint64_t task = coreIdx; task < taskNum; task += coreNum) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveFlatChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + (void)seq; + ProcessChunkPreAiv(b, h, hv, chunkIdx, start, end, subBlockIdx, subBlockNum); + } + } + } + + __aicore__ inline void ProcessPreAic() + { + if constexpr (IsSameType::value) { + return; + } +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (SAFE_GATE) { + if (headPairMode_) { + ProcessPreAicHeadPair(); + return; + } + } +#endif + uint64_t taskNum = static_cast((isVarLen_ ? NT_ : B_ * NT_) * HV_); + uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + uint64_t localTaskIdx = 0; + for (uint64_t task = GetBlockIdx(); task < taskNum; task += coreNum, ++localTaskIdx) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveFlatChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + if constexpr (SAFE_GATE) { + activeSolveSlot_ = localTaskIdx % KDA_SOLVE_PIPELINE_DEPTH; + } + (void)seq; + (void)h; + ProcessChunkPreAic(b, hv, chunkIdx, start, end); + } + } + } + +private: + GlobalTensor q_; + GlobalTensor k_; + GlobalTensor v_; + GlobalTensor gk_; + GlobalTensor rawG_; + GlobalTensor aLog_; + GlobalTensor dtBias_; + GlobalTensor betaFp32_; + GlobalTensor betaBf16_; + GlobalTensor initialState_; + GlobalTensor o_; + GlobalTensor finalState_; + GlobalTensor aqk_; + GlobalTensor akk_; + GlobalTensor w_; + GlobalTensor u_; + GlobalTensor qg_; + GlobalTensor kg_; + GlobalTensor vNew_; + GlobalTensor h_; + GlobalTensor finalKg_; + GlobalTensor preparedQG_; + GlobalTensor preparedAqk_; + GlobalTensor propagatedVNew_; + GlobalTensor propagatedH_; + GlobalTensor solveWorkspace_; + GlobalTensor scoreWorkspace_; + TPipe *pipe_ = nullptr; + TBuf exp2Buf_; + TBuf vecBuf_; + TBuf gateWritebackBuf_; + TEventID mte2ToVEvent_ = 0; + TEventID vToMte2Event_ = 0; + TEventID vToMte3Event_ = 0; + TEventID mte3ToVEvent_ = 0; + TEventID mte2ToMte3Event_ = 0; + TEventID vToSEvent_ = 0; + TEventID mte3ToMte2Events_[KDA_GATE_PIPELINE_DEPTH] = {0, 0, 0}; + bool vectorEventsAllocated_ = false; + Catlass::Arch::CrossCoreFlagWithReverse scoreReadyFlag_{KDA_SCORE_READY_FLAG0, + KDA_SCORE_READY_FLAG1}; + Catlass::Arch::CrossCoreFlagWithReverse scoreDoneFlag_{KDA_SCORE_DONE_FLAG0, + KDA_SCORE_DONE_FLAG1}; + // Solve has one outstanding task per core. Reuse the primary score IDs as an ordered token stream; + // score credits remain on the reverse IDs, so no additional hardware flag IDs are consumed. + Catlass::Arch::CrossCoreFlag syncReadyFlag_{KDA_SOLVE_READY_FLAG}; + Catlass::Arch::CrossCoreFlag syncDoneFlag_{KDA_SOLVE_DONE_FLAG}; + Catlass::Arch::CrossCoreFlagWithReverse postWuReadyFlag_{KDA_POST_READY_FLAG, + KDA_POST_FREE_FLAG}; + Catlass::Arch::CrossCoreFlagWithReverse mchSyncReadyFlag_{KDA_SCORE_READY_FLAG0, + KDA_SCORE_READY_FLAG1}; + Catlass::Arch::CrossCoreFlagWithReverse mchSyncDoneFlag_{KDA_SCORE_DONE_FLAG0, + KDA_SCORE_DONE_FLAG1}; + uint64_t B_ = 0; + uint64_t N_ = 0; + uint64_t H_ = 0; + uint64_t HV_ = 0; + uint64_t T_ = 0; + uint64_t K_ = 0; + uint64_t V_ = 0; + uint64_t BT_ = 0; + uint64_t NT_ = 0; + float scale_ = 1.0f; + bool hasInitial_ = false; + bool isVarLen_ = false; + bool isAivOnly_ = false; + bool headPairMode_ = false; + bool inputSequenceMajor_ = false; + bool fusePostWu_ = false; + bool computeGateInPrepare_ = false; + bool hasALog_ = false; + bool hasDtBias_ = false; + bool storeQG_ = true; + bool betaIsFp32_ = true; + float lowerBound_ = -5.0f; + uint64_t usedCoreNum_ = 1; + uint64_t solveCoreIdx_ = 0; + uint64_t activeSolveSlot_ = 0; + uint64_t activeGateChunkStart_ = 0; + __gm__ int64_t *chunkIndicesAddr_ = nullptr; + __gm__ int64_t *cuSeqlensAddr_ = nullptr; +}; +} // namespace + +template +__aicore__ inline void RunChunkKdaPrepare(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR rawG, GM_ADDR aLog, + GM_ADDR dtBias, GM_ADDR beta, GM_ADDR initialState, GM_ADDR cuSeqlens, + GM_ADDR chunkIndices, GM_ADDR aqk, GM_ADDR akk, GM_ADDR qg, GM_ADDR qgScaled, + GM_ADDR wSeed, GM_ADDR uSeed, GM_ADDR finalKg, GM_ADDR userWorkspace, + const TilingData &tiling, TPipe &pipe, bool storeQG = true) +{ + GM_ADDR aqkFp32 = userWorkspace + tiling.prepareAqkFp32Offset; + GM_ADDR akkFp32 = userWorkspace + tiling.prepareAkkFp32Offset; + GM_ADDR prepareScratch = userWorkspace + tiling.prepareScratchOffset; + + if ASCEND_IS_AIC { + ChunkKdaFwdPrepareKernel op; + op.Init(q, k, v, gk, rawG, aLog, dtBias, beta, initialState, cuSeqlens, chunkIndices, nullptr, nullptr, nullptr, + nullptr, aqk, userWorkspace, aqkFp32, akkFp32, wSeed, akk, qg, qgScaled, uSeed, userWorkspace, finalKg, + prepareScratch, tiling, &pipe, false, storeQG); + if (tiling.fusePostWu) { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + KdaPostWu::ChunkKdaFwdPostWuKernel postWu; + postWu.Init(nullptr, k, nullptr, gk, beta, initialState, cuSeqlens, chunkIndices, wSeed, akk, uSeed, + nullptr, userWorkspace, userWorkspace, userWorkspace, akk, wSeed, uSeed, userWorkspace, finalKg, + userWorkspace, prepareScratch, prepareScratch, tiling, &pipe, false); + op.ProcessAicFused(postWu); +#else + op.ProcessAic(); +#endif + } else { + op.ProcessAic(); + } + } + if ASCEND_IS_AIV { + ChunkKdaFwdPrepareKernel op; + op.Init(q, k, v, gk, rawG, aLog, dtBias, beta, initialState, cuSeqlens, chunkIndices, nullptr, nullptr, nullptr, + nullptr, aqk, userWorkspace, aqkFp32, akkFp32, wSeed, akk, qg, qgScaled, uSeed, userWorkspace, finalKg, + prepareScratch, tiling, &pipe, true, storeQG); + op.ProcessAiv(); + } +} + +} // namespace KdaPrepare diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/epilogue/block/block_epilogue_kda_fwdh_regbase.hpp b/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/epilogue/block/block_epilogue_kda_fwdh_regbase.hpp new file mode 100644 index 0000000000..be8af460d3 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/epilogue/block/block_epilogue_kda_fwdh_regbase.hpp @@ -0,0 +1,381 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef BLOCK_EPILOGUE_KDA_FWDH_REGBASE_HPP +#define BLOCK_EPILOGUE_KDA_FWDH_REGBASE_HPP + +#include "kernel_operator.h" +#ifndef FLA_NPU_REGBASE_HPP_INCLUDED +#define FLA_NPU_REGBASE_HPP_INCLUDED +#include "../../../../kernel_utils/vector/regbase.hpp" +#endif + +namespace Catlass::Epilogue::Block::detail { + +using namespace AscendC::MicroAPI; + +constexpr CastTrait KDA_B16_TO_F32_ZERO = { + RegLayout::ZERO, + SatMode::UNKNOWN, + MaskMergeMode::ZEROING, + AscendC::RoundMode::UNKNOWN, +}; + +template +__simd_callee__ inline void LoadKdaAsFloat( + RegTensor &dst, __ubuf__ T *src, MaskReg &mask) +{ + if constexpr (std::is_same()) { + DataCopy(dst, src); + } else if constexpr (std::is_same() || std::is_same()) { + RegTensor raw; + DataCopy(raw, src); + Cast(dst, raw, mask); + } else { + static_assert(!std::is_same::value, "KDA regbase only supports float/half/bfloat16_t"); + } +} + +template +__simd_callee__ inline void LoadKdaBroadcastAsFloat( + RegTensor &dst, __ubuf__ T *src, MaskReg &mask) +{ + if constexpr (std::is_same()) { + LoadAlign(dst, src); + } else if constexpr (std::is_same() || std::is_same()) { + RegTensor raw; + LoadAlign(raw, src); + Cast(dst, raw, mask); + } else { + static_assert(!std::is_same::value, "KDA regbase only supports float/half/bfloat16_t"); + } +} + +__simd_callee__ inline void LoadKdaFloat(RegTensor &dst, __ubuf__ float *src) +{ + DataCopy(dst, src); +} + +__simd_callee__ inline void StoreKdaFloat( + __ubuf__ float *dst, RegTensor &src, MaskReg &mask) +{ + DataCopy(dst, src, mask); +} + +template +static __simd_vf__ inline void PrepareKGateRegbase( + __ubuf__ float *gateOutput, __ubuf__ T *gateInput, uint16_t count) +{ + constexpr uint16_t FP32_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(float); + constexpr float LN2 = 0.6931471805599453f; + RegTensor gateReg; + MaskReg mask; + uint32_t remaining = count; + uint32_t offset = 0; + while (remaining > 0) { + mask = UpdateMask(remaining); + LoadKdaAsFloat(gateReg, gateInput + offset, mask); + if constexpr (USE_EXP2) { + Muls(gateReg, gateReg, LN2, mask); + } + Exp(gateReg, gateReg, mask); + StoreKdaFloat(gateOutput + offset, gateReg, mask); + offset += FP32_PER_REG; + } +} + +template +static __simd_vf__ inline void ComputeVNewRegbaseDualIssue( + __ubuf__ float *workspace, __ubuf__ T *uInput, uint32_t count) +{ + constexpr uint32_t FP32_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(float); + RegTensor uReg0; + RegTensor uReg1; + RegTensor wsReg0; + RegTensor wsReg1; + MaskReg mask0; + MaskReg mask1; + + uint32_t remaining = count; + uint32_t offset = 0; + while (remaining > FP32_PER_REG) { + mask0 = UpdateMask(remaining); + mask1 = UpdateMask(remaining); + LoadKdaAsFloat(uReg0, uInput + offset, mask0); + LoadKdaAsFloat(uReg1, uInput + offset + FP32_PER_REG, mask1); + LoadKdaFloat(wsReg0, workspace + offset); + LoadKdaFloat(wsReg1, workspace + offset + FP32_PER_REG); + Sub(wsReg0, uReg0, wsReg0, mask0); + Sub(wsReg1, uReg1, wsReg1, mask1); + StoreKdaFloat(workspace + offset, wsReg0, mask0); + StoreKdaFloat(workspace + offset + FP32_PER_REG, wsReg1, mask1); + offset += 2 * FP32_PER_REG; + } + if (remaining > 0) { + mask0 = UpdateMask(remaining); + LoadKdaAsFloat(uReg0, uInput + offset, mask0); + LoadKdaFloat(wsReg0, workspace + offset); + Sub(wsReg0, uReg0, wsReg0, mask0); + StoreKdaFloat(workspace + offset, wsReg0, mask0); + } +} + +template +static __simd_vf__ inline void ApplyKGateUpdateRegbaseDualIssue( + __ubuf__ float *update, __ubuf__ T *state, __ubuf__ float *rowScale, + uint16_t rows, uint16_t cols) +{ + constexpr uint16_t FP32_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(float); + RegTensor stateReg0; + RegTensor stateReg1; + RegTensor updateReg0; + RegTensor updateReg1; + RegTensor scaleReg0; + RegTensor scaleReg1; + MaskReg mask0; + MaskReg mask1; + + uint16_t row = 0; + for (; row + 1 < rows; row += 2) { + LoadAlign(scaleReg0, rowScale + row); + LoadAlign(scaleReg1, rowScale + row + 1); + uint32_t remaining0 = cols; + uint32_t remaining1 = cols; + uint16_t colLoops = static_cast((cols + FP32_PER_REG - 1) / FP32_PER_REG); + for (uint16_t colLoop = 0; colLoop < colLoops; ++colLoop) { + uint32_t col = static_cast(colLoop) * FP32_PER_REG; + mask0 = UpdateMask(remaining0); + mask1 = UpdateMask(remaining1); + LoadKdaAsFloat(stateReg0, state + static_cast(row) * cols + col, mask0); + LoadKdaAsFloat(stateReg1, state + static_cast(row + 1) * cols + col, mask1); + LoadKdaFloat(updateReg0, update + static_cast(row) * cols + col); + LoadKdaFloat(updateReg1, update + static_cast(row + 1) * cols + col); + Mul(stateReg0, stateReg0, scaleReg0, mask0); + Mul(stateReg1, stateReg1, scaleReg1, mask1); + Add(updateReg0, stateReg0, updateReg0, mask0); + Add(updateReg1, stateReg1, updateReg1, mask1); + StoreKdaFloat(update + static_cast(row) * cols + col, updateReg0, mask0); + StoreKdaFloat(update + static_cast(row + 1) * cols + col, updateReg1, mask1); + } + } + + if (row < rows) { + LoadAlign(scaleReg0, rowScale + row); + uint32_t remaining = cols; + uint16_t colLoops = static_cast((cols + FP32_PER_REG - 1) / FP32_PER_REG); + for (uint16_t colLoop = 0; colLoop < colLoops; ++colLoop) { + uint32_t col = static_cast(colLoop) * FP32_PER_REG; + mask0 = UpdateMask(remaining); + LoadKdaAsFloat(stateReg0, state + static_cast(row) * cols + col, mask0); + LoadKdaFloat(updateReg0, update + static_cast(row) * cols + col); + Mul(stateReg0, stateReg0, scaleReg0, mask0); + Add(updateReg0, stateReg0, updateReg0, mask0); + StoreKdaFloat(update + static_cast(row) * cols + col, updateReg0, mask0); + } + } +} + +template +static __simd_vf__ inline void ComputeKdaTailFinalStateRegbaseDualIssue( + __ubuf__ float *state, __ubuf__ KType *kInput, __ubuf__ VType *vInput, + __ubuf__ float *rowScale, uint16_t rows, uint16_t cols, + uint16_t tokens, uint16_t kStride, uint16_t kRowOffset) +{ + constexpr uint16_t FP32_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(float); + RegTensor stateReg0; + RegTensor stateReg1; + RegTensor valueReg; + RegTensor productReg0; + RegTensor productReg1; + RegTensor weightReg0; + RegTensor weightReg1; + RegTensor scaleReg0; + RegTensor scaleReg1; + MaskReg mask0; + MaskReg mask1; + + uint16_t row = 0; + for (; row + 1 < rows; row += 2) { + LoadAlign(scaleReg0, rowScale + row); + LoadAlign(scaleReg1, rowScale + row + 1); + uint32_t remaining0 = cols; + uint32_t remaining1 = cols; + uint16_t colLoops = static_cast((cols + FP32_PER_REG - 1) / FP32_PER_REG); + for (uint16_t colLoop = 0; colLoop < colLoops; ++colLoop) { + uint32_t col = static_cast(colLoop) * FP32_PER_REG; + mask0 = UpdateMask(remaining0); + mask1 = UpdateMask(remaining1); + LoadKdaFloat(stateReg0, state + static_cast(row) * cols + col); + LoadKdaFloat(stateReg1, state + static_cast(row + 1) * cols + col); + Mul(stateReg0, stateReg0, scaleReg0, mask0); + Mul(stateReg1, stateReg1, scaleReg1, mask1); + for (uint16_t token = 0; token < tokens; ++token) { + uint32_t kOffset = static_cast(token) * kStride + kRowOffset + row; + LoadKdaBroadcastAsFloat(weightReg0, kInput + kOffset, mask0); + LoadKdaBroadcastAsFloat(weightReg1, kInput + kOffset + 1, mask1); + LoadKdaAsFloat( + valueReg, vInput + static_cast(token) * cols + col, mask0); + Mul(productReg0, valueReg, weightReg0, mask0); + Mul(productReg1, valueReg, weightReg1, mask1); + Add(stateReg0, stateReg0, productReg0, mask0); + Add(stateReg1, stateReg1, productReg1, mask1); + } + StoreKdaFloat(state + static_cast(row) * cols + col, stateReg0, mask0); + StoreKdaFloat(state + static_cast(row + 1) * cols + col, stateReg1, mask1); + } + } + + if (row < rows) { + LoadAlign(scaleReg0, rowScale + row); + uint32_t remaining = cols; + uint16_t colLoops = static_cast((cols + FP32_PER_REG - 1) / FP32_PER_REG); + for (uint16_t colLoop = 0; colLoop < colLoops; ++colLoop) { + uint32_t col = static_cast(colLoop) * FP32_PER_REG; + mask0 = UpdateMask(remaining); + LoadKdaFloat(stateReg0, state + static_cast(row) * cols + col); + Mul(stateReg0, stateReg0, scaleReg0, mask0); + for (uint16_t token = 0; token < tokens; ++token) { + uint32_t kOffset = static_cast(token) * kStride + kRowOffset + row; + LoadKdaBroadcastAsFloat(weightReg0, kInput + kOffset, mask0); + LoadKdaAsFloat( + valueReg, vInput + static_cast(token) * cols + col, mask0); + Mul(productReg0, valueReg, weightReg0, mask0); + Add(stateReg0, stateReg0, productReg0, mask0); + } + StoreKdaFloat(state + static_cast(row) * cols + col, stateReg0, mask0); + } + } +} + +template +static __simd_vf__ inline void ComputeKdaTailVWorkspaceRegbaseDualIssue( + __ubuf__ float *output, __ubuf__ WType *wInput, __ubuf__ HType *hInput, + uint16_t rows, uint16_t cols, uint16_t reduction) +{ + constexpr uint16_t FP32_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(float); + RegTensor accumReg0; + RegTensor accumReg1; + RegTensor valueReg; + RegTensor productReg0; + RegTensor productReg1; + RegTensor weightReg0; + RegTensor weightReg1; + MaskReg mask0; + MaskReg mask1; + + uint16_t row = 0; + for (; row + 1 < rows; row += 2) { + uint32_t remaining0 = cols; + uint32_t remaining1 = cols; + uint16_t colLoops = static_cast((cols + FP32_PER_REG - 1) / FP32_PER_REG); + for (uint16_t colLoop = 0; colLoop < colLoops; ++colLoop) { + uint32_t col = static_cast(colLoop) * FP32_PER_REG; + mask0 = UpdateMask(remaining0); + mask1 = UpdateMask(remaining1); + Duplicate(accumReg0, 0.0f, mask0); + Duplicate(accumReg1, 0.0f, mask1); + for (uint16_t reductionIdx = 0; reductionIdx < reduction; ++reductionIdx) { + LoadKdaBroadcastAsFloat( + weightReg0, + wInput + static_cast(row) * reduction + reductionIdx, + mask0); + LoadKdaBroadcastAsFloat( + weightReg1, + wInput + static_cast(row + 1) * reduction + reductionIdx, + mask1); + LoadKdaAsFloat( + valueReg, + hInput + static_cast(reductionIdx) * cols + col, + mask0); + Mul(productReg0, valueReg, weightReg0, mask0); + Mul(productReg1, valueReg, weightReg1, mask1); + Add(accumReg0, accumReg0, productReg0, mask0); + Add(accumReg1, accumReg1, productReg1, mask1); + } + StoreKdaFloat(output + static_cast(row) * cols + col, accumReg0, mask0); + StoreKdaFloat(output + static_cast(row + 1) * cols + col, accumReg1, mask1); + } + } + + if (row < rows) { + uint32_t remaining = cols; + uint16_t colLoops = static_cast((cols + FP32_PER_REG - 1) / FP32_PER_REG); + for (uint16_t colLoop = 0; colLoop < colLoops; ++colLoop) { + uint32_t col = static_cast(colLoop) * FP32_PER_REG; + mask0 = UpdateMask(remaining); + Duplicate(accumReg0, 0.0f, mask0); + for (uint16_t reductionIdx = 0; reductionIdx < reduction; ++reductionIdx) { + LoadKdaBroadcastAsFloat( + weightReg0, + wInput + static_cast(row) * reduction + reductionIdx, + mask0); + LoadKdaAsFloat( + valueReg, + hInput + static_cast(reductionIdx) * cols + col, + mask0); + Mul(productReg0, valueReg, weightReg0, mask0); + Add(accumReg0, accumReg0, productReg0, mask0); + } + StoreKdaFloat(output + static_cast(row) * cols + col, accumReg0, mask0); + } + } +} + +static __simd_vf__ inline void ApplyRowScaleDualIssue( + __ubuf__ float *matrix, __ubuf__ float *rowScale, uint32_t rowScaleOffset, + uint16_t rows, uint16_t cols) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t FP32_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(float); + + RegTensor matrixReg0; + RegTensor matrixReg1; + RegTensor scaleReg0; + RegTensor scaleReg1; + MaskReg mask0; + MaskReg mask1; + + uint16_t row = 0; + for (; row + 1 < rows; row += 2) { + LoadAlign(scaleReg0, rowScale + rowScaleOffset + row); + LoadAlign(scaleReg1, rowScale + rowScaleOffset + row + 1); + uint32_t remaining0 = cols; + uint32_t remaining1 = cols; + uint16_t colLoops = static_cast((cols + FP32_PER_REG - 1) / FP32_PER_REG); + for (uint16_t colLoop = 0; colLoop < colLoops; ++colLoop) { + uint32_t col = static_cast(colLoop) * FP32_PER_REG; + mask0 = UpdateMask(remaining0); + mask1 = UpdateMask(remaining1); + LoadAlign(matrixReg0, matrix + static_cast(row) * cols + col); + LoadAlign(matrixReg1, matrix + static_cast(row + 1) * cols + col); + Mul(matrixReg0, matrixReg0, scaleReg0, mask0); + Mul(matrixReg1, matrixReg1, scaleReg1, mask1); + StoreAlign(matrix + static_cast(row) * cols + col, matrixReg0, mask0); + StoreAlign(matrix + static_cast(row + 1) * cols + col, matrixReg1, mask1); + } + } + + if (row < rows) { + LoadAlign(scaleReg0, rowScale + rowScaleOffset + row); + uint32_t remaining = cols; + uint16_t colLoops = static_cast((cols + FP32_PER_REG - 1) / FP32_PER_REG); + for (uint16_t colLoop = 0; colLoop < colLoops; ++colLoop) { + uint32_t col = static_cast(colLoop) * FP32_PER_REG; + mask0 = UpdateMask(remaining); + LoadAlign(matrixReg0, matrix + static_cast(row) * cols + col); + Mul(matrixReg0, matrixReg0, scaleReg0, mask0); + StoreAlign(matrix + static_cast(row) * cols + col, matrixReg0, mask0); + } + } +} + +} // namespace Catlass::Epilogue::Block::detail + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/epilogue/block/block_epilogue_kda_fwdh_update_arch35.hpp b/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/epilogue/block/block_epilogue_kda_fwdh_update_arch35.hpp new file mode 100644 index 0000000000..0346024e38 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/epilogue/block/block_epilogue_kda_fwdh_update_arch35.hpp @@ -0,0 +1,434 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_KDA_FWDH_UPDATE_ARCH35_HPP +#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_KDA_FWDH_UPDATE_ARCH35_HPP +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "../kda_fwd_h_epilogue_policies.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/kda_matrix_coord.hpp" +#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" +#include "block_epilogue_kda_fwdh_regbase.hpp" + +namespace Catlass::Epilogue::Block { + +template +class BlockEpilogue { + static constexpr bool kGated = KGatedTag::value; + static constexpr bool scalarGated = KGatedTag::scalarGated; + static constexpr bool useExp2 = KGatedTag::useExp2; + static constexpr float LN2 = 0.6931471805599453f; + +public: + // Type aliases + using DispatchPolicy = EpilogueAtlasKdaFwdHUpdate; + using ArchTag = typename DispatchPolicy::ArchTag; + + using HElementOutput = typename HOutputType_::Element; + using GElementInput = typename GInputType_::Element; + using HElementInput = typename HInputType_::Element; + using HUpdateElementInput = typename HUpdateInputType_::Element; + using FinalStateElement = typename FinalStateType_::Element; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource) + { + constexpr uint32_t CALC_BUF_OFFSET = 0; + constexpr uint32_t PING_BUF_0_OFFSET = 32 * 1024; + constexpr uint32_t PING_BUF_1_OFFSET = 48 * 1024; + constexpr uint32_t PING_BUF_2_OFFSET = 64 * 1024; + constexpr uint32_t PING_BUF_3_OFFSET = 80 * 1024; + constexpr uint32_t PONG_BUF_0_OFFSET = 96 * 1024; + constexpr uint32_t PONG_BUF_1_OFFSET = 112 * 1024; + constexpr uint32_t PONG_BUF_2_OFFSET = 128 * 1024; + constexpr uint32_t PONG_BUF_3_OFFSET = 144 * 1024; + constexpr uint32_t PING_G_BUF_OFFSET = 168 * 1024; + constexpr uint32_t PONG_G_BUF_OFFSET = 169 * 1024; + constexpr uint32_t PING_G_SUB_BUF_OFFSET = 170 * 1024; + constexpr uint32_t PONG_G_SUB_BUF_OFFSET = 171 * 1024; + constexpr uint32_t PING_G_INPUT_BUF_OFFSET = 172 * 1024; + constexpr uint32_t PONG_G_INPUT_BUF_OFFSET = 173 * 1024; + constexpr uint32_t UPDATE_SCRATCH_BUF_OFFSET = 160 * 1024; + constexpr uint32_t UPDATE_G_BUF_OFFSET = 176 * 1024; + + calcUbTensor = resource.ubBuf.template GetBufferByByte(CALC_BUF_OFFSET); + + hUpdateUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_0_OFFSET); + hUbTensor_ping = resource.ubBuf.template GetBufferByByte(UPDATE_SCRATCH_BUF_OFFSET); + finalOutputUbTensor_ping = + resource.ubBuf.template GetBufferByByte(UPDATE_SCRATCH_BUF_OFFSET); + glastUbTensor_ping = resource.ubBuf.template GetBufferByByte(UPDATE_G_BUF_OFFSET); + + hUpdateUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_0_OFFSET); + hUbTensor_pong = resource.ubBuf.template GetBufferByByte(UPDATE_SCRATCH_BUF_OFFSET); + finalOutputUbTensor_pong = + resource.ubBuf.template GetBufferByByte(UPDATE_SCRATCH_BUF_OFFSET); + glastUbTensor_pong = resource.ubBuf.template GetBufferByByte(UPDATE_G_BUF_OFFSET); + + if constexpr (kGated) { + gkLastUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_SUB_BUF_OFFSET); + gkLastUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_SUB_BUF_OFFSET); + gkInputUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_INPUT_BUF_OFFSET); + gkInputUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_INPUT_BUF_OFFSET); + gkBrcbUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_INPUT_BUF_OFFSET); + gkBrcbUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_INPUT_BUF_OFFSET); + } + } + + CATLASS_DEVICE + ~BlockEpilogue() {} + + template + CATLASS_DEVICE void CopyGmToUb(AscendC::LocalTensor dst, AscendC::GlobalTensor src, uint32_t rows, + uint32_t cols, uint32_t srcStride) + { + if (cols == srcStride) { + AscendC::DataCopy(dst, src, rows * cols); + return; + } + AscendC::DataCopyExtParams copyParams{static_cast(rows), + static_cast(cols * sizeof(Element)), + static_cast((srcStride - cols) * sizeof(Element)), 0, 0}; + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyPad(dst, src, copyParams, padParams); + } + + template + CATLASS_DEVICE void CopyUbToGm(AscendC::GlobalTensor dst, AscendC::LocalTensor src, uint32_t rows, + uint32_t cols, uint32_t dstStride) + { + if (cols == dstStride) { + AscendC::DataCopy(dst, src, rows * cols); + return; + } + AscendC::DataCopyExtParams copyParams{static_cast(rows), + static_cast(cols * sizeof(Element)), 0, + static_cast((dstStride - cols) * sizeof(Element)), 0}; + AscendC::DataCopyPad(dst, src, copyParams); + } + + CATLASS_DEVICE + void ApplyRowScale(AscendC::LocalTensor matrix, AscendC::LocalTensor rowScale, uint32_t rows, + uint32_t cols) + { + __ubuf__ float *matrixAddr = reinterpret_cast<__ubuf__ float *>(matrix.GetPhyAddr()); + __ubuf__ float *rowScaleAddr = reinterpret_cast<__ubuf__ float *>(rowScale.GetPhyAddr()); + AscendC::VF_CALL(matrixAddr, rowScaleAddr, 0, static_cast(rows), + static_cast(cols)); + AscendC::PipeBarrier(); + } + + CATLASS_DEVICE + void PrepareKGate(AscendC::LocalTensor gateOutput, AscendC::LocalTensor gateInput, + uint32_t count) + { + __ubuf__ float *gateOutputAddr = reinterpret_cast<__ubuf__ float *>(gateOutput.GetPhyAddr()); + __ubuf__ GElementInput *gateInputAddr = reinterpret_cast<__ubuf__ GElementInput *>(gateInput.GetPhyAddr()); + AscendC::VF_CALL>(gateOutputAddr, gateInputAddr, + static_cast(count)); + AscendC::PipeBarrier(); + } + + template + CATLASS_DEVICE void ApplyKGateUpdate(AscendC::LocalTensor update, AscendC::LocalTensor state, + AscendC::LocalTensor rowScale, uint32_t rows, uint32_t cols) + { + __ubuf__ float *updateAddr = reinterpret_cast<__ubuf__ float *>(update.GetPhyAddr()); + __ubuf__ StateElement *stateAddr = reinterpret_cast<__ubuf__ StateElement *>(state.GetPhyAddr()); + __ubuf__ float *rowScaleAddr = reinterpret_cast<__ubuf__ float *>(rowScale.GetPhyAddr()); + AscendC::VF_CALL>( + updateAddr, stateAddr, rowScaleAddr, static_cast(rows), static_cast(cols)); + AscendC::PipeBarrier(); + } + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor hOutput, AscendC::GlobalTensor finalState, + AscendC::GlobalTensor gInput, AscendC::GlobalTensor hInput, + AscendC::GlobalTensor hUpdateInput, AscendC::GlobalTensor gkInput, + AscendC::GlobalTensor initialState, uint32_t chunkSize, uint32_t kHeadDim, + uint32_t vBlockDim, uint32_t vHeadDim, Arch::CrossCoreFlag cube2Done, bool isInitialState, + bool isFinalState, bool storeFinalState, bool useInitialState, bool isPing, bool cube2AlreadyWaited, + bool useDirectFp32Ub, uint64_t directUbFreeFlagBegin, uint64_t directUbReadyFlagBegin, + uint64_t directUbFlagStride) + { + static constexpr uint32_t ROW_TILE = 16; + uint32_t mActual = kHeadDim; + uint32_t nActual = vBlockDim; + uint32_t outputStride = vHeadDim; + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t rowsPerSubBlock = CeilDiv(mActual, subBlockNum); + uint32_t rowBegin = subBlockIdx * rowsPerSubBlock; + uint32_t rowEnd = rowBegin + rowsPerSubBlock; + if (rowEnd > mActual) { + rowEnd = mActual; + } + if (rowBegin >= mActual) { + if (useDirectFp32Ub) { + uint32_t directUbSlot = isPing ? 0 : 1; + uint64_t directUbFlagOffset = directUbSlot + AscendC::GetSubBlockIdx() * directUbFlagStride; + AscendC::CrossCoreWaitFlag<0x4, PIPE_V>(directUbReadyFlagBegin + directUbFlagOffset); + AscendC::CrossCoreSetFlag<0x4, PIPE_V>(directUbFreeFlagBegin + directUbFlagOffset); + } else if (!cube2AlreadyWaited) { + Arch::CrossCoreWaitFlag(cube2Done); + } + return; + } + + AscendC::ResetMask(); + + AscendC::GlobalTensor gInputThisSubBlock = gInput; + + uint32_t pingpongFlag = isPing ? 0 : pongBaseEvent; + AscendC::LocalTensor hUpdateUbTensor = isPing ? hUpdateUbTensor_ping : hUpdateUbTensor_pong; + AscendC::LocalTensor hUbTensor = isPing ? hUbTensor_ping : hUbTensor_pong; + AscendC::LocalTensor finalOutputUbTensor = + isPing ? finalOutputUbTensor_ping : finalOutputUbTensor_pong; + AscendC::LocalTensor glastUbTensor = isPing ? glastUbTensor_ping : glastUbTensor_pong; + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + bool useFp32StateUpdate = + storeFinalState && std::is_same::value && (!isInitialState || useInitialState); + float muls = 1.0f; + if constexpr (scalarGated) { + GElementInput gLastVal = gInputThisSubBlock.GetValue(chunkSize - 1); + float gLastFloat = 0.0f; + if constexpr (std::is_same::value) { + gLastFloat = gLastVal; + } else if constexpr (std::is_same::value) { + gLastFloat = (float)gLastVal; + } else if constexpr (std::is_same::value) { + gLastFloat = AscendC::ToFloat(gLastVal); + } + glastUbTensor.SetValue(0, gLastFloat); + + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + if constexpr (useExp2) { + AscendC::Muls(glastUbTensor, glastUbTensor, LN2, 1); + AscendC::PipeBarrier(); + } + AscendC::Exp(glastUbTensor, glastUbTensor, 1); + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + muls = glastUbTensor.GetValue(0); + } + if constexpr (kGated) { + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + } + if constexpr (scalarGated) { + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + } + + if (useDirectFp32Ub) { + uint32_t directUbSlot = isPing ? 0 : 1; + uint64_t directUbFlagOffset = directUbSlot + AscendC::GetSubBlockIdx() * directUbFlagStride; + AscendC::CrossCoreWaitFlag<0x4, PIPE_V>(directUbReadyFlagBegin + directUbFlagOffset); + } else if (!cube2AlreadyWaited) { + Arch::CrossCoreWaitFlag(cube2Done); + } + // fix: need to adapt kGated. issue: A5 do not have vdim128 branch. + bool waitHFromV = storeFinalState && isInitialState && std::is_same::value; + bool waitUpdateFromMte3 = false; + uint32_t updateReadyEvent = EVENT_ID3 + pingpongFlag; + for (uint32_t rowStart = rowBegin; rowStart < rowEnd; rowStart += ROW_TILE) { + uint32_t rowsThisTile = rowEnd - rowStart; + if (rowsThisTile > ROW_TILE) { + rowsThisTile = ROW_TILE; + } + + AscendC::GlobalTensor hOutputThisTile = hOutput[rowStart * outputStride]; + AscendC::GlobalTensor hInputThisTile = hInput[rowStart * outputStride]; + AscendC::GlobalTensor hUpdateInputThisTile = hUpdateInput[rowStart * nActual]; + AscendC::GlobalTensor finalStateThisTile = finalState[rowStart * outputStride]; + AscendC::LocalTensor hUpdateUbTensorThisTile = + useDirectFp32Ub ? hUpdateUbTensor[(rowStart - rowBegin) * nActual] : hUpdateUbTensor; + + if (waitHFromV) { + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + } else { + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + } + if constexpr (std::is_same::value) { + if (useFp32StateUpdate) { + if (isInitialState) { + CopyGmToUb(calcUbTensor, initialState[rowStart * outputStride], rowsThisTile, nActual, + outputStride); + } else { + CopyGmToUb(calcUbTensor, finalStateThisTile, rowsThisTile, nActual, outputStride); + } + } else { + CopyGmToUb(hUbTensor, hInputThisTile, rowsThisTile, nActual, outputStride); + } + } else { + CopyGmToUb(hUbTensor, hInputThisTile, rowsThisTile, nActual, outputStride); + } + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + + if (!useFp32StateUpdate) { + AscendC::Cast(calcUbTensor, hUbTensor, AscendC::RoundMode::CAST_NONE, rowsThisTile * nActual); + AscendC::PipeBarrier(); + } + if constexpr (scalarGated) { + AscendC::Muls(calcUbTensor, calcUbTensor, muls, rowsThisTile * nActual); + AscendC::PipeBarrier(); + } + + if constexpr (kGated) { + AscendC::GlobalTensor gkLastInput = gkInput[(chunkSize - 1) * kHeadDim + rowStart]; + AscendC::LocalTensor gkLastUbTensor = isPing ? gkLastUbTensor_ping : gkLastUbTensor_pong; + AscendC::LocalTensor gkInputUbTensor = + isPing ? gkInputUbTensor_ping : gkInputUbTensor_pong; + if (rowStart == rowBegin) { + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + } else { + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + } + if constexpr (std::is_same::value) { + AscendC::DataCopy(gkLastUbTensor, gkLastInput, rowsThisTile); + } else { + AscendC::DataCopy(gkInputUbTensor, gkLastInput, rowsThisTile); + } + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + if constexpr (std::is_same::value) { + PrepareKGate(gkLastUbTensor, gkLastUbTensor, rowsThisTile); + } else { + PrepareKGate(gkLastUbTensor, gkInputUbTensor, rowsThisTile); + } + ApplyRowScale(calcUbTensor, gkLastUbTensor, rowsThisTile, nActual); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + } + + if (waitUpdateFromMte3) { + AscendC::WaitFlag(updateReadyEvent); + } else { + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + } + if (!useDirectFp32Ub) { + CopyGmToUb(hUpdateUbTensorThisTile, hUpdateInputThisTile, rowsThisTile, nActual, nActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + } + AscendC::Add(hUpdateUbTensorThisTile, calcUbTensor, hUpdateUbTensorThisTile, rowsThisTile * nActual); + AscendC::PipeBarrier(); + if (storeFinalState && isFinalState && std::is_same::value) { + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + waitHFromV = true; + } else { + waitHFromV = false; + } + + if constexpr (std::is_same::value) { + if (storeFinalState) { + if (!isFinalState) { + AscendC::Cast(hUbTensor, hUpdateUbTensorThisTile, AscendC::RoundMode::CAST_RINT, + rowsThisTile * nActual); + AscendC::PipeBarrier(); + } + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + CopyUbToGm(finalStateThisTile, hUpdateUbTensorThisTile, rowsThisTile, nActual, outputStride); + AscendC::PipeBarrier(); + AscendC::SetFlag(updateReadyEvent); + AscendC::SetFlag(updateReadyEvent); + AscendC::WaitFlag(updateReadyEvent); + waitUpdateFromMte3 = true; + if (!isFinalState) { + CopyUbToGm(hOutputThisTile, hUbTensor, rowsThisTile, nActual, outputStride); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + } + } else { + AscendC::Cast(hUbTensor, hUpdateUbTensorThisTile, AscendC::RoundMode::CAST_RINT, + rowsThisTile * nActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + CopyUbToGm(hOutputThisTile, hUbTensor, rowsThisTile, nActual, outputStride); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + waitUpdateFromMte3 = false; + } + } else { + if (storeFinalState && isFinalState) { + AscendC::Cast(finalOutputUbTensor, hUpdateUbTensorThisTile, AscendC::RoundMode::CAST_RINT, + rowsThisTile * nActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + CopyUbToGm(finalStateThisTile, finalOutputUbTensor, rowsThisTile, nActual, outputStride); + } else { + AscendC::Cast(hUbTensor, hUpdateUbTensorThisTile, AscendC::RoundMode::CAST_RINT, + rowsThisTile * nActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + CopyUbToGm(hOutputThisTile, hUbTensor, rowsThisTile, nActual, outputStride); + } + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + waitUpdateFromMte3 = false; + } + } + + if (storeFinalState && std::is_same::value) { + AscendC::WaitFlag(updateReadyEvent); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + if (!isFinalState) { + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + } + } else { + AscendC::WaitFlag(EVENT_ID2 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + } + if constexpr (kGated) { + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + } + if (useDirectFp32Ub) { + uint32_t directUbSlot = isPing ? 0 : 1; + uint64_t directUbFlagOffset = directUbSlot + AscendC::GetSubBlockIdx() * directUbFlagStride; + AscendC::CrossCoreSetFlag<0x4, PIPE_MTE3>(directUbFreeFlagBegin + directUbFlagOffset); + } + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID2 + pingpongFlag); + } + +private: + uint32_t pongBaseEvent = 4; + + AscendC::LocalTensor calcUbTensor; + + AscendC::LocalTensor hUpdateUbTensor_ping; + AscendC::LocalTensor hUbTensor_ping; + AscendC::LocalTensor finalOutputUbTensor_ping; + AscendC::LocalTensor glastUbTensor_ping; + + AscendC::LocalTensor hUpdateUbTensor_pong; + AscendC::LocalTensor hUbTensor_pong; + AscendC::LocalTensor finalOutputUbTensor_pong; + AscendC::LocalTensor glastUbTensor_pong; + + AscendC::LocalTensor gkLastUbTensor_ping; + AscendC::LocalTensor gkLastUbTensor_pong; + AscendC::LocalTensor gkInputUbTensor_ping; + AscendC::LocalTensor gkInputUbTensor_pong; + AscendC::LocalTensor gkBrcbUbTensor_ping; + AscendC::LocalTensor gkBrcbUbTensor_pong; +}; +} // namespace Catlass::Epilogue::Block + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/epilogue/block/block_epilogue_kda_fwdh_vnew_arch35.hpp b/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/epilogue/block/block_epilogue_kda_fwdh_vnew_arch35.hpp new file mode 100644 index 0000000000..cd908f5077 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/epilogue/block/block_epilogue_kda_fwdh_vnew_arch35.hpp @@ -0,0 +1,509 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_KDA_FWDH_VNEW_ARCH35_HPP +#define CATLASS_EPILOGUE_BLOCK_BLOCK_EPILOGUE_KDA_FWDH_VNEW_ARCH35_HPP +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "../kda_fwd_h_epilogue_policies.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/kda_matrix_coord.hpp" +#include "catlass/epilogue/tile/kda_epilogue_tile_copy.hpp" +#include "block_epilogue_kda_fwdh_regbase.hpp" + +namespace Catlass::Epilogue::Block { + +template +class BlockEpilogue { + static constexpr bool kGated = KGatedTag::value; + static constexpr bool scalarGated = KGatedTag::scalarGated; + static constexpr bool useExp2 = KGatedTag::useExp2; + static constexpr float LN2 = 0.6931471805599453f; + +public: + using DispatchPolicy = EpilogueAtlasKdaFwdHVnew; + using ArchTag = typename DispatchPolicy::ArchTag; + + using VElementOutput = typename VOutputType_::Element; + using GElementInput = typename GInputType_::Element; + using UElementInput = typename UInputType_::Element; + using WSElementInput = typename WSInputType_::Element; + using VElementUpdate = typename VUpdateType_::Element; + using VLayoutUpdate = typename VUpdateType_::Layout; + using FinalStateElement = typename FinalStateType_::Element; + + CATLASS_DEVICE + BlockEpilogue(Arch::Resource &resource) + { + constexpr uint32_t CALC_BUF_OFFSET = 0; + constexpr uint32_t PING_BUF_0_OFFSET = 32 * 1024; + constexpr uint32_t PING_BUF_1_OFFSET = 48 * 1024; + constexpr uint32_t PING_BUF_2_OFFSET = 64 * 1024; + constexpr uint32_t PING_BUF_3_OFFSET = 80 * 1024; + constexpr uint32_t PONG_BUF_0_OFFSET = 96 * 1024; + constexpr uint32_t PONG_BUF_1_OFFSET = 112 * 1024; + constexpr uint32_t PONG_BUF_2_OFFSET = 128 * 1024; + constexpr uint32_t PONG_BUF_3_OFFSET = 144 * 1024; + constexpr uint32_t PING_WIDE_IO_BUF_OFFSET = 16 * 1024; + constexpr uint32_t PONG_WIDE_IO_BUF_OFFSET = 24 * 1024; + constexpr uint32_t PING_G_BUF_OFFSET = 160 * 1024; + constexpr uint32_t PONG_G_BUF_OFFSET = 161 * 1024; + constexpr uint32_t PING_G_SUB_BUF_OFFSET = 162 * 1024; + constexpr uint32_t PONG_G_SUB_BUF_OFFSET = 163 * 1024; + constexpr uint32_t PING_G_INPUT_BUF_OFFSET = 164 * 1024; + constexpr uint32_t PONG_G_INPUT_BUF_OFFSET = 165 * 1024; + constexpr uint32_t SHARE_BUF_OFFSET = 166 * 1024; + + calcUbTensor = resource.ubBuf.template GetBufferByByte(CALC_BUF_OFFSET); + + uUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_2_OFFSET); + wsUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_0_OFFSET); + gUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_BUF_OFFSET); + gLastUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_SUB_BUF_OFFSET); + gInputUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_G_SUB_BUF_OFFSET); + vNewOutputUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_2_OFFSET); + vNewDecayUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_BUF_2_OFFSET); + wideIoUbTensor_ping = resource.ubBuf.template GetBufferByByte(PING_WIDE_IO_BUF_OFFSET); + + uUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_2_OFFSET); + wsUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_0_OFFSET); + gUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_BUF_OFFSET); + gLastUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_SUB_BUF_OFFSET); + gInputUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_G_SUB_BUF_OFFSET); + vNewOutputUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_2_OFFSET); + vNewDecayUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_BUF_2_OFFSET); + wideIoUbTensor_pong = resource.ubBuf.template GetBufferByByte(PONG_WIDE_IO_BUF_OFFSET); + + gBrcbUbTensor_ = resource.ubBuf.template GetBufferByByte(SHARE_BUF_OFFSET); + } + + CATLASS_DEVICE + ~BlockEpilogue() {} + + template + CATLASS_DEVICE void CopyGmToUb(AscendC::LocalTensor dst, AscendC::GlobalTensor src, uint32_t rows, + uint32_t cols, uint32_t srcStride) + { + if (cols == srcStride) { + AscendC::DataCopy(dst, src, rows * cols); + return; + } + AscendC::DataCopyExtParams copyParams{static_cast(rows), + static_cast(cols * sizeof(Element)), + static_cast((srcStride - cols) * sizeof(Element)), 0, 0}; + AscendC::DataCopyPadExtParams padParams{false, 0, 0, 0}; + AscendC::DataCopyPad(dst, src, copyParams, padParams); + } + + template + CATLASS_DEVICE void CopyUbToGm(AscendC::GlobalTensor dst, AscendC::LocalTensor src, uint32_t rows, + uint32_t cols, uint32_t dstStride) + { + if (cols == dstStride) { + AscendC::DataCopy(dst, src, rows * cols); + return; + } + AscendC::DataCopyExtParams copyParams{static_cast(rows), + static_cast(cols * sizeof(Element)), 0, + static_cast((dstStride - cols) * sizeof(Element)), 0}; + AscendC::DataCopyPad(dst, src, copyParams); + } + + CATLASS_DEVICE + void PrepareG(AscendC::LocalTensor gUbTensor, AscendC::LocalTensor gLastUbTensor, + AscendC::LocalTensor gInputUbTensor, + AscendC::GlobalTensor gInputThisSubBlock, uint32_t mActual, uint32_t pingpongFlag) + { + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + if constexpr (!scalarGated) { + AscendC::Duplicate(gUbTensor, 1.0f, mActual); + AscendC::PipeBarrier(); + return; + } + if (mActual == 1) { + AscendC::Duplicate(gUbTensor, 1.0f, 1); + AscendC::PipeBarrier(); + return; + } + if constexpr (std::is_same::value) { + AscendC::DataCopyParams gUbParams{1, (uint16_t)(mActual * sizeof(float)), 0, 0}; + AscendC::DataCopyPadParams gUbPadParams{false, 0, 0, 0}; + AscendC::DataCopyPad(gUbTensor, gInputThisSubBlock, gUbParams, gUbPadParams); + } else { + AscendC::DataCopyParams gUbParams{1, (uint16_t)(mActual * sizeof(GElementInput)), 0, 0}; + AscendC::DataCopyPadParams gUbPadParams{false, 0, 0, 0}; + AscendC::DataCopyPad(gInputUbTensor, gInputThisSubBlock, gUbParams, gUbPadParams); + } + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + if constexpr (!std::is_same::value) { + AscendC::Cast(gUbTensor, gInputUbTensor, AscendC::RoundMode::CAST_NONE, mActual); + AscendC::PipeBarrier(); + } + + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + float inputVal = gUbTensor.GetValue(mActual - 1); + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + + AscendC::PipeBarrier(); + AscendC::Duplicate(gLastUbTensor, inputVal, mActual); + AscendC::PipeBarrier(); + + AscendC::Sub(gUbTensor, gLastUbTensor, gUbTensor, mActual); + AscendC::PipeBarrier(); + if constexpr (useExp2) { + AscendC::Muls(gUbTensor, gUbTensor, LN2, mActual); + AscendC::PipeBarrier(); + } + AscendC::Exp(gUbTensor, gUbTensor, mActual); + AscendC::PipeBarrier(); + } + + CATLASS_DEVICE + void ApplyRowScale(AscendC::LocalTensor matrix, AscendC::LocalTensor rowScale, + uint32_t rowScaleOffset, uint32_t rows, uint32_t cols) + { + __ubuf__ float *matrixAddr = reinterpret_cast<__ubuf__ float *>(matrix.GetPhyAddr()); + __ubuf__ float *rowScaleAddr = reinterpret_cast<__ubuf__ float *>(rowScale.GetPhyAddr()); + AscendC::VF_CALL(matrixAddr, rowScaleAddr, rowScaleOffset, + static_cast(rows), static_cast(cols)); + AscendC::PipeBarrier(); + } + + CATLASS_DEVICE + void ComputeVNew(AscendC::LocalTensor workspace, AscendC::LocalTensor uInput, uint32_t count) + { + __ubuf__ float *workspaceAddr = reinterpret_cast<__ubuf__ float *>(workspace.GetPhyAddr()); + __ubuf__ UElementInput *uInputAddr = reinterpret_cast<__ubuf__ UElementInput *>(uInput.GetPhyAddr()); + AscendC::VF_CALL>(workspaceAddr, uInputAddr, count); + AscendC::PipeBarrier(); + } + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor vnewOutput, + AscendC::GlobalTensor vnewdecayOutput, + AscendC::LocalTensor l1VUpdate, AscendC::GlobalTensor gInput, + AscendC::GlobalTensor uInput, AscendC::GlobalTensor wsInput, + AscendC::GlobalTensor gkInput, AscendC::GlobalTensor kInput, + AscendC::GlobalTensor kDecayWorkspace, uint32_t chunkSize, uint32_t kHeadDim, + uint32_t vBlockDim, uint32_t vHeadDim, Arch::CrossCoreFlag cube1Done, Arch::CrossCoreFlag vec1Done, + bool isInitialState, bool isFinalState, bool storeFinalState, bool waitWsFromMte3, bool isPing, + bool cube1AlreadyWaited, bool useDirectFp32Ub, uint64_t directUbFreeFlagBegin, + uint64_t directUbReadyFlagBegin, uint64_t directUbFlagStride) + { + static constexpr uint32_t ROW_TILE = 16; + uint32_t mActual = chunkSize; + uint32_t nvActual = vBlockDim; + uint32_t nkActual = kHeadDim; + uint32_t inputStride = vHeadDim; + + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t rowsPerSubBlock = CeilDiv(mActual, subBlockNum); + uint32_t rowBegin = subBlockIdx * rowsPerSubBlock; + uint32_t rowEnd = rowBegin + rowsPerSubBlock; + if (rowEnd > mActual) { + rowEnd = mActual; + } + uint32_t pingpongFlag = isPing ? 0 : pongBaseEvent; + if (rowBegin >= mActual) { + if (useDirectFp32Ub) { + uint32_t directUbSlot = isPing ? 0 : 1; + uint64_t directUbFlagOffset = directUbSlot + AscendC::GetSubBlockIdx() * directUbFlagStride; + AscendC::CrossCoreWaitFlag<0x4, PIPE_V>(directUbReadyFlagBegin + directUbFlagOffset); + AscendC::CrossCoreSetFlag<0x4, PIPE_V>(directUbFreeFlagBegin + directUbFlagOffset); + } else if (!cube1AlreadyWaited) { + Arch::CrossCoreWaitFlag(cube1Done); + } + // A zero-row AIV lane still owns the EVENT0 hand-off consumed by V2. + if (waitWsFromMte3) { + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + } + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vec1Done); + return; + } + AscendC::ResetMask(); + + AscendC::GlobalTensor gInputThisSubBlock = gInput; + + AscendC::LocalTensor uUbTensor = isPing ? uUbTensor_ping : uUbTensor_pong; + AscendC::LocalTensor wsUbTensor = isPing ? wsUbTensor_ping : wsUbTensor_pong; + AscendC::LocalTensor gUbTensor = isPing ? gUbTensor_ping : gUbTensor_pong; + AscendC::LocalTensor gLastUbTensor = isPing ? gLastUbTensor_ping : gLastUbTensor_pong; + AscendC::LocalTensor gInputUbTensor = isPing ? gInputUbTensor_ping : gInputUbTensor_pong; + AscendC::LocalTensor vNewOutputUbTensor = + isPing ? vNewOutputUbTensor_ping : vNewOutputUbTensor_pong; + AscendC::LocalTensor vNewDecayUbTensor = + isPing ? vNewDecayUbTensor_ping : vNewDecayUbTensor_pong; + if (nvActual > 128) { + uUbTensor = isPing ? wideIoUbTensor_ping : wideIoUbTensor_pong; + vNewOutputUbTensor = isPing ? wideIoUbTensor_ping : wideIoUbTensor_pong; + vNewDecayUbTensor = isPing ? wideIoUbTensor_ping : wideIoUbTensor_pong; + } + + if (rowBegin < rowEnd && nvActual <= 128 && nvActual == inputStride) { + uint32_t mActualThisSubBlock = rowEnd - rowBegin; + AscendC::GlobalTensor vnewOutputThisSubBlock = vnewOutput[rowBegin * inputStride]; + AscendC::GlobalTensor uInputThisSubBlock = uInput[rowBegin * inputStride]; + AscendC::GlobalTensor wsInputThisSubBlock = wsInput[rowBegin * nvActual]; + + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::DataCopy(uUbTensor, uInputThisSubBlock, mActualThisSubBlock * nvActual); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + if constexpr (scalarGated) { + AscendC::Cast(calcUbTensor, uUbTensor, AscendC::RoundMode::CAST_NONE, mActualThisSubBlock * nvActual); + } + + if constexpr (scalarGated) { + PrepareG(gUbTensor, gLastUbTensor, gInputUbTensor, gInputThisSubBlock, mActual, pingpongFlag); + } else { + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + } + if (waitWsFromMte3) { + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + } else { + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + } + if (useDirectFp32Ub) { + uint32_t directUbSlot = isPing ? 0 : 1; + uint64_t directUbFlagOffset = directUbSlot + AscendC::GetSubBlockIdx() * directUbFlagStride; + AscendC::CrossCoreWaitFlag<0x4, PIPE_V>(directUbReadyFlagBegin + directUbFlagOffset); + } else { + if (!cube1AlreadyWaited) { + Arch::CrossCoreWaitFlag(cube1Done); + } + AscendC::DataCopy(wsUbTensor, wsInputThisSubBlock, mActualThisSubBlock * nvActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + } + + if constexpr (scalarGated) { + AscendC::Sub(wsUbTensor, calcUbTensor, wsUbTensor, mActualThisSubBlock * nvActual); + AscendC::PipeBarrier(); + AscendC::Copy(calcUbTensor, wsUbTensor, mActualThisSubBlock * nvActual); + AscendC::PipeBarrier(); + ApplyRowScale(calcUbTensor, gUbTensor, rowBegin, mActualThisSubBlock, nvActual); + } else { + ComputeVNew(wsUbTensor, uUbTensor, mActualThisSubBlock * nvActual); + } + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + + uint32_t nvLoops = nvActual / FLOAT_NUM_PER_REPEAT; + for (uint32_t nLoop = 0; nLoop < nvLoops; nLoop++) { + uint32_t castSrcOffset = nLoop * FLOAT_NUM_PER_REPEAT; + uint32_t castDstOffset = nLoop * mActualThisSubBlock * FLOAT_NUM_PER_REPEAT; + AscendC::LocalTensor decayInput = scalarGated ? calcUbTensor : wsUbTensor; + AscendC::Cast(vNewDecayUbTensor[castDstOffset], decayInput[castSrcOffset], + AscendC::RoundMode::CAST_RINT, FLOAT_NUM_PER_REPEAT, mActualThisSubBlock, + {(uint16_t)mActualThisSubBlock, 1, 1, (uint8_t)(nvLoops * 8)}); + } + + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + + uint32_t ubToL1Loops = nvActual / SIZE_16_NUM_PER_C0; + uint32_t mActualPadded = (mActual + NZ_BLOCK_SIZE - 1) / NZ_BLOCK_SIZE * NZ_BLOCK_SIZE; + AscendC::DataCopyParams intriParams; + intriParams.blockCount = ubToL1Loops; + intriParams.blockLen = mActualThisSubBlock; + intriParams.srcGap = 0; + intriParams.dstGap = mActualPadded - mActualThisSubBlock; + uint32_t l1Addr = rowBegin * SIZE_16_NUM_PER_C0; + AscendC::DataCopy(vnewdecayOutput[l1Addr], vNewDecayUbTensor, intriParams); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + if constexpr (!kGated) { + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vec1Done); + } + + AscendC::Cast(vNewOutputUbTensor, wsUbTensor, AscendC::RoundMode::CAST_RINT, + mActualThisSubBlock * nvActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::DataCopy(vnewOutputThisSubBlock, vNewOutputUbTensor, mActualThisSubBlock * nvActual); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + + if constexpr (kGated) { + AscendC::GlobalTensor kInputThisSubBlock = kInput[rowBegin * nkActual]; + AscendC::GlobalTensor kDecayWorkspaceThisSubBlock = + kDecayWorkspace[rowBegin * nkActual]; + // KDA passes kg = k * exp2(g_last - gk). Keep that decay exactly once. + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::DataCopy(vNewOutputUbTensor, kInputThisSubBlock, mActualThisSubBlock * nkActual); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::DataCopy(kDecayWorkspaceThisSubBlock, vNewOutputUbTensor, mActualThisSubBlock * nkActual); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vec1Done); + } + if (useDirectFp32Ub) { + uint32_t directUbSlot = isPing ? 0 : 1; + uint64_t directUbFlagOffset = directUbSlot + AscendC::GetSubBlockIdx() * directUbFlagStride; + AscendC::CrossCoreSetFlag<0x4, PIPE_V>(directUbFreeFlagBegin + directUbFlagOffset); + } + return; + } + + if constexpr (scalarGated) { + PrepareG(gUbTensor, gLastUbTensor, gInputUbTensor, gInputThisSubBlock, mActual, pingpongFlag); + } else { + AscendC::WaitFlag(EVENT_ID3 + pingpongFlag); + } + if (!cube1AlreadyWaited) { + Arch::CrossCoreWaitFlag(cube1Done); + } + + uint32_t mActualPadded = (mActual + NZ_BLOCK_SIZE - 1) / NZ_BLOCK_SIZE * NZ_BLOCK_SIZE; + bool waitWsThisTileFromMte3 = waitWsFromMte3; + for (uint32_t rowStart = rowBegin; rowStart < rowEnd;) { + uint32_t alignExtra = rowStart & 7; + uint32_t maxRowsThisTile = ROW_TILE - alignExtra; + uint32_t rowsThisTile = rowEnd - rowStart; + if (rowsThisTile > maxRowsThisTile) { + rowsThisTile = maxRowsThisTile; + } + uint32_t localRowStart = rowStart - rowBegin; + + AscendC::GlobalTensor vnewOutputThisTile = vnewOutput[rowStart * inputStride]; + AscendC::GlobalTensor uInputThisTile = uInput[rowStart * inputStride]; + AscendC::GlobalTensor wsInputThisTile = wsInput[rowStart * nvActual]; + AscendC::LocalTensor wsUbTensorThisTile = wsUbTensor[localRowStart * nvActual]; + + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + CopyGmToUb(uUbTensor, uInputThisTile, rowsThisTile, nvActual, inputStride); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + if constexpr (scalarGated) { + AscendC::Cast(calcUbTensor, uUbTensor, AscendC::RoundMode::CAST_NONE, rowsThisTile * nvActual); + AscendC::PipeBarrier(); + } + + if (waitWsThisTileFromMte3) { + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + waitWsThisTileFromMte3 = false; + } else { + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + } + CopyGmToUb(wsUbTensorThisTile, wsInputThisTile, rowsThisTile, nvActual, nvActual); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID0 + pingpongFlag); + if constexpr (scalarGated) { + AscendC::Sub(wsUbTensorThisTile, calcUbTensor, wsUbTensorThisTile, rowsThisTile * nvActual); + AscendC::PipeBarrier(); + AscendC::Copy(calcUbTensor, wsUbTensorThisTile, rowsThisTile * nvActual); + AscendC::PipeBarrier(); + ApplyRowScale(calcUbTensor, gUbTensor, rowStart, rowsThisTile, nvActual); + } else { + ComputeVNew(wsUbTensorThisTile, uUbTensor, rowsThisTile * nvActual); + } + + uint32_t nvLoops = nvActual / FLOAT_NUM_PER_REPEAT; + for (uint32_t nLoop = 0; nLoop < nvLoops; nLoop++) { + uint32_t castSrcOffset = nLoop * FLOAT_NUM_PER_REPEAT; + uint32_t castDstOffset = nLoop * rowsThisTile * FLOAT_NUM_PER_REPEAT; + AscendC::LocalTensor decayInput = scalarGated ? calcUbTensor : wsUbTensorThisTile; + AscendC::Cast(vNewDecayUbTensor[castDstOffset], decayInput[castSrcOffset], + AscendC::RoundMode::CAST_RINT, FLOAT_NUM_PER_REPEAT, rowsThisTile, + {(uint16_t)rowsThisTile, 1, 1, (uint8_t)(nvLoops * 8)}); + } + + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + + uint32_t ubToL1Loops = nvActual / SIZE_16_NUM_PER_C0; + AscendC::DataCopyParams intriParams; + intriParams.blockCount = ubToL1Loops; + intriParams.blockLen = rowsThisTile; + intriParams.srcGap = 0; + intriParams.dstGap = mActualPadded - rowsThisTile; + uint32_t l1Addr = rowStart * SIZE_16_NUM_PER_C0; + AscendC::DataCopy(vnewdecayOutput[l1Addr], vNewDecayUbTensor, intriParams); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + if constexpr (!kGated) { + if (rowStart + rowsThisTile >= rowEnd) { + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vec1Done); + } + } + + AscendC::Cast(vNewOutputUbTensor, wsUbTensorThisTile, AscendC::RoundMode::CAST_RINT, + rowsThisTile * nvActual); + AscendC::PipeBarrier(); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::SetFlag(EVENT_ID0 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + CopyUbToGm(vnewOutputThisTile, vNewOutputUbTensor, rowsThisTile, nvActual, inputStride); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + rowStart += rowsThisTile; + } + + if constexpr (kGated) { + uint32_t mActualThisSubBlock = rowEnd - rowBegin; + for (uint32_t rowOffset = 0; rowOffset < mActualThisSubBlock; rowOffset += ROW_TILE) { + uint32_t rowsThisTile = Min(ROW_TILE, mActualThisSubBlock - rowOffset); + uint32_t row = rowBegin + rowOffset; + AscendC::GlobalTensor kInputThisTile = kInput[row * nkActual]; + AscendC::GlobalTensor kDecayWorkspaceThisTile = kDecayWorkspace[row * nkActual]; + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + CopyGmToUb(vNewOutputUbTensor, kInputThisTile, rowsThisTile, nkActual, nkActual); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + CopyUbToGm(kDecayWorkspaceThisTile, vNewOutputUbTensor, rowsThisTile, nkActual, nkActual); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + } + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vec1Done); + } + + if (rowBegin < rowEnd) { + AscendC::WaitFlag(EVENT_ID1 + pingpongFlag); + AscendC::SetFlag(EVENT_ID1 + pingpongFlag); + } + AscendC::SetFlag(EVENT_ID3 + pingpongFlag); + } + +private: + uint32_t pongBaseEvent = 4; + + AscendC::LocalTensor calcUbTensor; + + AscendC::LocalTensor uUbTensor_ping; + AscendC::LocalTensor wsUbTensor_ping; + AscendC::LocalTensor gUbTensor_ping; + AscendC::LocalTensor gLastUbTensor_ping; + AscendC::LocalTensor gInputUbTensor_ping; + AscendC::LocalTensor vNewOutputUbTensor_ping; + AscendC::LocalTensor vNewDecayUbTensor_ping; + + AscendC::LocalTensor uUbTensor_pong; + AscendC::LocalTensor wsUbTensor_pong; + AscendC::LocalTensor gUbTensor_pong; + AscendC::LocalTensor gLastUbTensor_pong; + AscendC::LocalTensor gInputUbTensor_pong; + AscendC::LocalTensor vNewOutputUbTensor_pong; + AscendC::LocalTensor vNewDecayUbTensor_pong; + AscendC::LocalTensor wideIoUbTensor_ping; + AscendC::LocalTensor wideIoUbTensor_pong; + + AscendC::LocalTensor gBrcbUbTensor_; +}; +} // namespace Catlass::Epilogue::Block + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/epilogue/kda_fwd_h_epilogue_policies.hpp b/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/epilogue/kda_fwd_h_epilogue_policies.hpp new file mode 100644 index 0000000000..964e3d58ad --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/epilogue/kda_fwd_h_epilogue_policies.hpp @@ -0,0 +1,28 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_KDA_FWD_H_EPILOGUE_POLICIES_HPP +#define CATLASS_EPILOGUE_KDA_FWD_H_EPILOGUE_POLICIES_HPP + +#include "catlass/kda_catlass.hpp" + +namespace Catlass::Epilogue { + +struct EpilogueAtlasKdaFwdHVnew { + using ArchTag = Arch::Ascend950; +}; + +struct EpilogueAtlasKdaFwdHUpdate { + using ArchTag = Arch::Ascend950; +}; + +} // namespace Catlass::Epilogue + +#endif // CATLASS_EPILOGUE_KDA_FWD_H_EPILOGUE_POLICIES_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/gemm/block/block_scheduler_kda_fwd_h_arch35.hpp b/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/gemm/block/block_scheduler_kda_fwd_h_arch35.hpp new file mode 100644 index 0000000000..c0ac3c14cc --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/gemm/block/block_scheduler_kda_fwd_h_arch35.hpp @@ -0,0 +1,437 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#include "catlass/kda_gemm_coord.hpp" +#include "../../../../chunk_kda_fwd_h_runtime_tiling.h" +using namespace Catlass; + +#ifndef CATLASS_GEMM_SCHEDULER_KDA_FWD_H_HPP +#define CATLASS_GEMM_SCHEDULER_KDA_FWD_H_HPP + +// constexpr uint32_t PING_PONG_STAGES = 1; +constexpr uint32_t PING_PONG_STAGES = 2; +constexpr uint32_t BYTE_SIZE_16_BIT = 2; +constexpr uint32_t BYTES_PER_C0 = 32; +constexpr uint32_t BYTE_SIZE_PER_REPEAT = 256; +constexpr uint32_t SIZE_16_NUM_PER_C0 = BYTES_PER_C0 / BYTE_SIZE_16_BIT; +constexpr uint32_t FLOAT_NUM_PER_REPEAT = BYTE_SIZE_PER_REPEAT / sizeof(float); +constexpr uint32_t NZ_BLOCK_SIZE = 16; + +template +CATLASS_DEVICE T AlignUp(T a, T b) +{ + return (b == 0) ? 0 : (a + b - 1) / b * b; +} + +template +CATLASS_DEVICE T Min(T a, T b) +{ + return (a > b) ? b : a; +} + +template +CATLASS_DEVICE T Max(T a, T b) +{ + return (a > b) ? a : b; +} + +namespace Catlass::Gemm::Block { + +struct KdaFwdHOffsets { + uint32_t hSrcOffset; + uint32_t hDstOffset; + uint32_t uvOffset; + uint32_t wkOffset; + uint32_t wOffset; + uint32_t gOffset; + uint32_t gkOffset; + uint32_t hWorkOffset; + uint32_t vWorkOffset; + uint32_t kDecayWorkOffset; + uint32_t vBlockOffset; + uint32_t vBlockDim; + uint32_t initialStateOffset; + uint32_t finalStateOffset; + bool isInitialState; + bool isFinalState; + uint32_t blockTokens; + uint32_t streamId; + // for debug + uint32_t batchIdx; + uint32_t headIdx; + uint32_t chunkIdx; +}; + +struct KdaFwdHStream { + uint32_t batchIdx; + uint32_t chunkIdx{0}; + uint32_t vHeadIdx; + uint32_t kHeadIdx; + uint32_t shapeBatchIdx; + uint32_t tokenBatchIdx; + + uint32_t chunkOffset; + uint32_t tokenOffset; + uint32_t batchChunks{0}; + uint32_t batchTokens; + uint32_t nextTaskIdx{0}; + bool active{false}; + + KdaFwdHOffsets offset; +}; + +struct KdaFwdHRunningQ { + KdaFwdHStream streams[PING_PONG_STAGES]; +}; + +struct BlockSchedulerKdaFwdH { + uint32_t batch; + uint32_t seqlen; + uint32_t kNumHead; + uint32_t vNumHead; + uint32_t kHeadDim; + uint32_t vHeadDim; + uint32_t chunkSize; + uint32_t vBlockSize{128}; + uint32_t isVariedLen; + uint32_t shapeBatch; + uint32_t tokenBatch; + uint32_t inputTokenBatch; + bool useInitialState; + bool storeFinalState; + uint32_t numSeqWorkspaceOffset; + uint32_t numChunksWorkspaceOffset; + + uint32_t taskIdx; + uint32_t taskStride; + uint32_t cubeCoreIdx; + uint32_t cubeCoreNum; + uint32_t taskNum; + uint32_t headGroups; + uint32_t totalChunks; + uint32_t totalTokens; + + KdaFwdHRunningQ runningQ; + + bool isRunning; + + AscendC::GlobalTensor gmSeqlen; + AscendC::GlobalTensor gmNumSeq; + AscendC::GlobalTensor gmNumChunks; + + Arch::CrossCoreFlag cube1Done[PING_PONG_STAGES] = {0, 1}; + Arch::CrossCoreFlag vec1Done[PING_PONG_STAGES] = {2, 3}; + Arch::CrossCoreFlag cube2Done[PING_PONG_STAGES] = {4, 5}; + Arch::CrossCoreFlag vec2Done[PING_PONG_STAGES] = {6, 7}; + + CATLASS_DEVICE + BlockSchedulerKdaFwdH() {} + + CATLASS_DEVICE + void Init(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR tiling, GM_ADDR user, uint32_t coreIdx, + uint32_t coreNum) + { + __gm__ ChunkKdaFwdHRuntimeTiling *__restrict kdaFwdHTilingData = + reinterpret_cast<__gm__ ChunkKdaFwdHRuntimeTiling *__restrict>(tiling); + + batch = kdaFwdHTilingData->batch; + seqlen = kdaFwdHTilingData->seqlen; + kNumHead = kdaFwdHTilingData->kNumHead; + vNumHead = kdaFwdHTilingData->vNumHead; + kHeadDim = kdaFwdHTilingData->kHeadDim; + vHeadDim = kdaFwdHTilingData->vHeadDim; + chunkSize = kdaFwdHTilingData->chunkSize; + isVariedLen = kdaFwdHTilingData->isVariedLen; + shapeBatch = kdaFwdHTilingData->shapeBatch; + tokenBatch = kdaFwdHTilingData->tokenBatch; + useInitialState = kdaFwdHTilingData->useInitialState; + storeFinalState = kdaFwdHTilingData->storeFinalState; + numSeqWorkspaceOffset = kdaFwdHTilingData->numSeqWorkspaceOffset; + numChunksWorkspaceOffset = kdaFwdHTilingData->numChunksWorkspaceOffset; + + InitRuntime(cu_seqlens, chunk_indices, user, coreIdx, coreNum); + } + + template + CATLASS_DEVICE void InitFromData(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, const TilingData &tilingData, + GM_ADDR user, uint32_t coreIdx, uint32_t coreNum) + { + batch = tilingData.batch; + seqlen = tilingData.seqlen; + kNumHead = tilingData.kNumHead; + vNumHead = tilingData.vNumHead; + kHeadDim = tilingData.kHeadDim; + vHeadDim = tilingData.vHeadDim; + chunkSize = tilingData.chunkSize; + isVariedLen = tilingData.isVariedLen; + shapeBatch = tilingData.shapeBatch; + tokenBatch = tilingData.tokenBatch; + useInitialState = tilingData.useInitialState; + storeFinalState = tilingData.storeFinalState; + numSeqWorkspaceOffset = tilingData.numSeqWorkspaceOffset; + numChunksWorkspaceOffset = tilingData.numChunksWorkspaceOffset; + + InitRuntime(cu_seqlens, chunk_indices, user, coreIdx, coreNum); + } + + CATLASS_DEVICE + void InitRuntime(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR user, uint32_t coreIdx, uint32_t coreNum) + { + gmSeqlen.SetGlobalBuffer((__gm__ int64_t *)cu_seqlens); + gmNumSeq.SetGlobalBuffer((__gm__ int64_t *)(user + numSeqWorkspaceOffset)); + gmNumChunks.SetGlobalBuffer((__gm__ int64_t *)(user + numChunksWorkspaceOffset)); + + if (isVariedLen) { + inputTokenBatch = tokenBatch; + uint32_t actualBatch = 0; + int64_t chunkPrefix = 0; + int64_t prevSeq = 0, currSeq; + for (uint32_t b = 1; b <= inputTokenBatch; b++) { + currSeq = gmSeqlen.GetValue(b); + int64_t batchSeqLen = currSeq - prevSeq; + if (batchSeqLen > 0) { + actualBatch++; + int64_t batchChunk = (batchSeqLen + chunkSize - 1) / chunkSize; + chunkPrefix += batchChunk; + } + prevSeq = currSeq; + } + tokenBatch = actualBatch; + batch = actualBatch; + totalChunks = chunkPrefix; + totalTokens = prevSeq; + } else { + totalChunks = (seqlen + chunkSize - 1) / chunkSize; + totalTokens = seqlen; + } + + cubeCoreIdx = coreIdx; + cubeCoreNum = coreNum; + vBlockSize = vHeadDim; + taskNum = batch * vNumHead; + headGroups = vNumHead / kNumHead; + uint32_t maxTaskCntPerLoop = taskNum > cubeCoreNum ? PING_PONG_STAGES : 1; + taskStride = cubeCoreNum * maxTaskCntPerLoop; + for (uint32_t streamId = 0; streamId < PING_PONG_STAGES; ++streamId) { + auto &stream = runningQ.streams[streamId]; + stream.nextTaskIdx = + (streamId < maxTaskCntPerLoop) ? (cubeCoreIdx * maxTaskCntPerLoop + streamId) : taskNum; + stream.chunkIdx = 0; + stream.batchChunks = 0; + stream.active = false; + } + isRunning = cubeCoreIdx * maxTaskCntPerLoop < taskNum; + } + + CATLASS_DEVICE + void ResolveVarlenSequence(uint32_t compactBatchIdx, KdaFwdHStream &stream) + { + uint32_t actualBatch = 0; + int64_t chunkPrefix = 0; + int64_t prevSeq = 0; + for (uint32_t b = 1; b <= inputTokenBatch; ++b) { + int64_t currSeq = gmSeqlen.GetValue(b); + int64_t batchTokens = currSeq - prevSeq; + if (batchTokens > 0) { + int64_t batchChunks = (batchTokens + chunkSize - 1) / chunkSize; + if (actualBatch == compactBatchIdx) { + stream.chunkOffset = static_cast(chunkPrefix); + stream.batchChunks = static_cast(batchChunks); + stream.tokenOffset = static_cast(prevSeq); + stream.batchTokens = static_cast(batchTokens); + return; + } + ++actualBatch; + chunkPrefix += batchChunks; + } + prevSeq = currSeq; + } + stream.chunkOffset = 0; + stream.batchChunks = 0; + stream.tokenOffset = 0; + stream.batchTokens = 0; + } + + CATLASS_DEVICE + uint32_t GetVarlenChunkOffset(uint32_t compactBatchIdx) + { + KdaFwdHStream stream; + ResolveVarlenSequence(compactBatchIdx, stream); + return stream.chunkOffset; + } + + CATLASS_DEVICE + void InitNewStream(KdaFwdHStream &newStream) + { + newStream.batchIdx = taskIdx / vNumHead; + newStream.vHeadIdx = taskIdx % vNumHead; + newStream.kHeadIdx = newStream.vHeadIdx / headGroups; + newStream.shapeBatchIdx = isVariedLen ? 0 : newStream.batchIdx; + newStream.tokenBatchIdx = isVariedLen ? newStream.batchIdx : 0; + if (isVariedLen) { + ResolveVarlenSequence(newStream.tokenBatchIdx, newStream); + } else { + newStream.chunkOffset = 0; + newStream.batchChunks = totalChunks; + newStream.tokenOffset = 0; + newStream.batchTokens = totalTokens; + } + newStream.chunkIdx = 0; + } + + CATLASS_DEVICE + void AssignNextStream(uint32_t streamId) + { + auto &stream = runningQ.streams[streamId]; + taskIdx = stream.nextTaskIdx; + if (taskIdx >= taskNum) { + stream.active = false; + stream.batchChunks = 0; + return; + } + + stream.nextTaskIdx += taskStride; + InitNewStream(stream); + stream.active = stream.batchChunks > 0; + if (stream.active) { + UpdateTask(streamId); + } + } + + CATLASS_DEVICE + void UpdateTask(uint32_t streamId) + { + auto &stream = runningQ.streams[streamId]; + auto &offset = stream.offset; + + offset.isInitialState = stream.chunkIdx == 0; + offset.isFinalState = stream.chunkIdx == (stream.batchChunks - 1); + uint32_t vBlockOffset = 0; + uint32_t vBlockDim = vBlockSize; + offset.initialStateOffset = (stream.batchIdx * vNumHead + stream.vHeadIdx) * kHeadDim * vHeadDim + vBlockOffset; + offset.finalStateOffset = (stream.batchIdx * vNumHead + stream.vHeadIdx) * kHeadDim * vHeadDim + vBlockOffset; + offset.hSrcOffset = (stream.shapeBatchIdx * vNumHead * totalChunks + stream.vHeadIdx * totalChunks + + stream.chunkOffset + stream.chunkIdx) * + kHeadDim * vHeadDim + + vBlockOffset; + offset.hDstOffset = offset.hSrcOffset + kHeadDim * vHeadDim; + if (storeFinalState && offset.isFinalState) { + offset.hDstOffset = offset.hSrcOffset; + } + offset.uvOffset = (stream.shapeBatchIdx * vNumHead * totalTokens + stream.vHeadIdx * totalTokens + + stream.tokenOffset + stream.chunkIdx * chunkSize) * + vHeadDim + + vBlockOffset; + offset.wkOffset = (stream.shapeBatchIdx * kNumHead * totalTokens + stream.kHeadIdx * totalTokens + + stream.tokenOffset + stream.chunkIdx * chunkSize) * + kHeadDim; + offset.wOffset = (stream.shapeBatchIdx * vNumHead * totalTokens + stream.vHeadIdx * totalTokens + + stream.tokenOffset + stream.chunkIdx * chunkSize) * + kHeadDim; + offset.gOffset = stream.shapeBatchIdx * vNumHead * totalTokens + stream.vHeadIdx * totalTokens + + stream.tokenOffset + stream.chunkIdx * chunkSize; + offset.gkOffset = (stream.shapeBatchIdx * vNumHead * totalTokens + stream.vHeadIdx * totalTokens + + stream.tokenOffset + stream.chunkIdx * chunkSize) * + kHeadDim; + offset.hWorkOffset = (cubeCoreIdx * PING_PONG_STAGES + streamId) * kHeadDim * vBlockSize; + offset.vWorkOffset = (cubeCoreIdx * PING_PONG_STAGES + streamId) * chunkSize * vBlockSize; + offset.kDecayWorkOffset = (cubeCoreIdx * PING_PONG_STAGES + streamId) * chunkSize * kHeadDim; + offset.vBlockOffset = vBlockOffset; + offset.vBlockDim = vBlockDim; + offset.blockTokens = offset.isFinalState ? (stream.batchTokens - stream.chunkIdx * chunkSize) : chunkSize; + offset.streamId = streamId; + offset.batchIdx = stream.batchIdx; + offset.headIdx = stream.vHeadIdx; + offset.chunkIdx = stream.chunkIdx; + } + + CATLASS_DEVICE + void InitTasks() + { + isRunning = false; + for (uint32_t streamId = 0; streamId < PING_PONG_STAGES; ++streamId) { + auto &stream = runningQ.streams[streamId]; + if (stream.active) { + stream.chunkIdx += 1; + if (stream.chunkIdx >= stream.batchChunks) { + stream.active = false; + stream.batchChunks = 0; + } + } + if (!stream.active) { + AssignNextStream(streamId); + } else { + UpdateTask(streamId); + } + if (stream.active) { + isRunning = true; + } + } + } + + CATLASS_DEVICE + const KdaFwdHStream &GetStream(uint32_t i) const { return runningQ.streams[i]; } + + CATLASS_DEVICE + uint32_t GetStreamId(uint32_t i) const { return i; } + + CATLASS_DEVICE + const KdaFwdHOffsets &GetCurTaskOffsets(const KdaFwdHStream &stream) const { return stream.offset; } + + CATLASS_DEVICE + bool StreamIsDone(const KdaFwdHStream &stream) const { return !stream.active; } + + CATLASS_DEVICE + bool NeedProcessStage2(const KdaFwdHStream &stream) { return storeFinalState || !stream.offset.isFinalState; } +}; + +struct BlockSchedulerKdaFwdHCube : public BlockSchedulerKdaFwdH { + CATLASS_DEVICE + BlockSchedulerKdaFwdHCube() {} + + CATLASS_DEVICE + void Init(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR tiling, GM_ADDR user) + { + BlockSchedulerKdaFwdH::Init(cu_seqlens, chunk_indices, tiling, user, AscendC::GetBlockIdx(), + AscendC::GetBlockNum()); + } + + template + CATLASS_DEVICE void InitFromData(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, const TilingData &tilingData, + GM_ADDR user) + { + BlockSchedulerKdaFwdH::InitFromData(cu_seqlens, chunk_indices, tilingData, user, AscendC::GetBlockIdx(), + AscendC::GetBlockNum()); + } +}; + +struct BlockSchedulerKdaFwdHVec : public BlockSchedulerKdaFwdH { + CATLASS_DEVICE + BlockSchedulerKdaFwdHVec() {} + + CATLASS_DEVICE + void Init(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR tiling, GM_ADDR user) + { + BlockSchedulerKdaFwdH::Init(cu_seqlens, chunk_indices, tiling, user, + AscendC::GetBlockIdx() / AscendC::GetSubBlockNum(), AscendC::GetBlockNum()); + } + + template + CATLASS_DEVICE void InitFromData(GM_ADDR cu_seqlens, GM_ADDR chunk_indices, const TilingData &tilingData, + GM_ADDR user) + { + BlockSchedulerKdaFwdH::InitFromData(cu_seqlens, chunk_indices, tilingData, user, + AscendC::GetBlockIdx() / AscendC::GetSubBlockNum(), AscendC::GetBlockNum()); + } +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_SCHEDULER_KDA_FWD_H_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch35.hpp b/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch35.hpp new file mode 100644 index 0000000000..6039f9994e --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/arch35/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch35.hpp @@ -0,0 +1,1127 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#define CATLASS_ARCH 3510 + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/arch/kda_cross_core_sync.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/kda_debug.hpp" +#include "../block/block_scheduler_kda_fwd_h_arch35.hpp" +#include "catlass/epilogue/block/kda_block_epilogue.hpp" +#include "../../epilogue/block/block_epilogue_kda_fwdh_update_arch35.hpp" +#include "../../epilogue/block/block_epilogue_kda_fwdh_vnew_arch35.hpp" +#include "catlass/gemm/block/kda_block_mmad.hpp" +#include "../../../../kernel_utils/block/block_mmad_pingpong_tla_kernel_utils.hpp" +#include "../../../../kernel_utils/block/block_mmad_pingpong_tla_multi.hpp" +#include "../../../../kernel_utils/block/block_mmad_pingpong_tla_preloadA_l1B.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "tla/kda_tensor.hpp" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +using _0 = tla::Int<0>; +using _1 = tla::Int<1>; +using _2 = tla::Int<2>; +using _4 = tla::Int<4>; +using _8 = tla::Int<8>; +using _16 = tla::Int<16>; +using _32 = tla::Int<32>; +using _64 = tla::Int<64>; +using _128 = tla::Int<128>; +using _256 = tla::Int<256>; +using _512 = tla::Int<512>; +using _1024 = tla::Int<1024>; +using _2048 = tla::Int<2048>; +using _4096 = tla::Int<4096>; +using _8192 = tla::Int<8192>; +using _16384 = tla::Int<16384>; +using _32768 = tla::Int<32768>; +using _65536 = tla::Int<65536>; + +#include "kernel_operator.h" +using namespace Catlass; +using namespace tla; + +namespace Catlass::Gemm::Kernel { + +struct KdaFwdHTileShapes128 { + using L1TileShape = tla::Shape<_128, _128, _128>; + using L0TileShape = L1TileShape; +}; + +struct KdaFwdHTileShapes256 { + using L1TileShape = tla::Shape<_128, _256, _128>; + using L0TileShape = tla::Shape<_128, _256, _64>; +}; + +template +struct KdaFwdHGateTag { + static constexpr bool value = KGated; + static constexpr bool scalarGated = ScalarGated; + static constexpr bool useExp2 = UseExp2; +}; + +template +class KdaFwdHKernel { +public: + using ArchTag = Arch::Ascend950; + using CubeScheduler = typename Catlass::Gemm::Block::BlockSchedulerKdaFwdHCube; + using VecScheduler = typename Catlass::Gemm::Block::BlockSchedulerKdaFwdHVec; + + using DispatchPolicyTlaMulti = Gemm::MmadPingpongTlaMulti; + using DispatchPolicyTlaTail = Gemm::MmadPingpongTlaMulti; + using DispatchPolicyDirectUb = Common::MmadPingpong; + using DispatchPolicyTlaPreloadAL1B = Gemm::MmadPingpongTlaPreloadAL1B; + using L1TileShapeVTla = typename TileShapes::L1TileShape; + using L0TileShapeVTla = typename TileShapes::L0TileShape; + + using WType = Gemm::GemmType; + using HType = Gemm::GemmType; + using VworkType = Gemm::GemmType; + using KType = Gemm::GemmType; + using HworkType = Gemm::GemmType; + using VType = Gemm::GemmType; + using GType = Gemm::GemmType; + using UType = Gemm::GemmType; + using FinalStateType = Gemm::GemmType; + using VUpdateType = Gemm::GemmType; + + // cube 1 + using TileCopyWH = Catlass::Gemm::Tile::PackedTileCopyTla; + using TileCopyWHDirectUb = Common::Tile::PackedTileCopyTlaToUB; + using BlockMmadWH = Gemm::Block::BlockMmadTla; + using BlockMmadWHTail = Gemm::Block::BlockMmadTla; + using BlockMmadWHDirectUb = Common::BlockMmadTla; + + // cube 2 + using TileCopyKV = Catlass::Gemm::Tile::PackedTileCopyTla; + using TileCopyKVDirectUb = Common::Tile::PackedTileCopyTlaToUB; + using TileMmadKV = Gemm::Tile::TileMmadTla; + using BlockMmadKV = Gemm::Block::BlockMmadTla; + using BlockMmadKVTail = Gemm::Block::BlockMmadTla; + using BlockMmadKVDirectUb = Common::BlockMmadTla; + + // vec 1 + using DispatchPolicyKdaFwdHVnew = Epilogue::EpilogueAtlasKdaFwdHVnew; + using GateTag = KdaFwdHGateTag; + using EpilogueKdaFwdHVnew = Epilogue::Block::BlockEpilogue; + + // vec 2 + using DispatchPolicyKdaFwdHUpdate = Epilogue::EpilogueAtlasKdaFwdHUpdate; + using EpilogueKdaFwdHUpdate = Epilogue::Block::BlockEpilogue; + + using KdaFwdHOffsets = Catlass::Gemm::Block::KdaFwdHOffsets; + + using ElementK = INPUT_TYPE; + using ElementW = INPUT_TYPE; + using ElementU = INPUT_TYPE; + using ElementG = G_TYPE; + using ElementH = INPUT_TYPE; + using ElementV = INPUT_TYPE; + using ElementVUpdate = INPUT_TYPE; + using ElementVWork = WORKSPACE_TYPE; + using ElementHWork = WORKSPACE_TYPE; + using ElementInitialState = STATE_TYPE; + using ElementFinalState = STATE_TYPE; + + using LayoutW = Catlass::layout::RowMajor; + using LayoutH = Catlass::layout::RowMajor; + using LayoutV = Catlass::layout::RowMajor; + using LayoutK = Catlass::layout::ColumnMajor; + using LayoutVUpdate = typename VUpdateType::Layout; + + static constexpr uint64_t DIRECT_UB_FREE_FLAG_BEGIN = 1; + static constexpr uint64_t DIRECT_UB_READY_FLAG_BEGIN = 6; + static constexpr uint64_t DIRECT_UB_FLAG_STRIDE = 16; + static constexpr uint32_t DIRECT_UB_STAGES = 2; + static constexpr uint32_t DIRECT_VEC_NUM = 2; + + uint32_t batch; + uint32_t seqlen; + uint32_t kNumHead; + uint32_t vNumHead; + uint32_t kHeadDim; + uint32_t vHeadDim; + uint32_t chunkSize; + bool useInitialState; + bool storeFinalState; + uint32_t isVariedLen; + uint32_t shapeBatch; + uint32_t tokenBatch; + uint32_t vWorkspaceOffset; + uint32_t vUpdateWorkspaceOffset; + uint32_t hWorkspaceOffset; + uint32_t numSeqWorkspaceOffset; + uint32_t numChunksWorkspaceOffset; + uint32_t kDecayWorkspaceOffset; + bool useDirectFp32Ub; + + AscendC::GlobalTensor gmK; + AscendC::GlobalTensor gmW; + AscendC::GlobalTensor gmU; + AscendC::GlobalTensor gmG; + AscendC::GlobalTensor gmGk; + AscendC::GlobalTensor gmInitialState; + AscendC::GlobalTensor gmH; + AscendC::GlobalTensor gmV; + AscendC::GlobalTensor gmFinalState; + AscendC::GlobalTensor gmVWorkspace; + AscendC::GlobalTensor gmVUpdateWorkspace; + AscendC::GlobalTensor gmHWorkspace; + AscendC::GlobalTensor gmKDecayWorkspace; + + AscendC::GlobalTensor gmSeqlen; + AscendC::GlobalTensor gmNumSeq; + AscendC::GlobalTensor gmNumChunks; + + AscendC::LocalTensor ubHUpdatePing; + AscendC::LocalTensor ubHUpdatePong; + AscendC::LocalTensor ubVWorkPing; + AscendC::LocalTensor ubVWorkPong; + + AscendC::LocalTensor l1VUpdatePing; + AscendC::LocalTensor l1VUpdatePong; + + CubeScheduler cubeBlockScheduler; + VecScheduler vecBlockScheduler; + + Arch::Resource resource; + + __aicore__ inline KdaFwdHKernel() {} + + __aicore__ inline void Init(GM_ADDR k, GM_ADDR w, GM_ADDR u, GM_ADDR g, GM_ADDR gk, GM_ADDR inital_state, + GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR h, GM_ADDR v_new, + GM_ADDR final_state, GM_ADDR tiling, GM_ADDR user) + { + __gm__ ChunkKdaFwdHRuntimeTiling *__restrict kdaFwdHTilingData = + reinterpret_cast<__gm__ ChunkKdaFwdHRuntimeTiling *__restrict>(tiling); + + batch = kdaFwdHTilingData->batch; + seqlen = kdaFwdHTilingData->seqlen; + kNumHead = kdaFwdHTilingData->kNumHead; + vNumHead = kdaFwdHTilingData->vNumHead; + kHeadDim = kdaFwdHTilingData->kHeadDim; + vHeadDim = kdaFwdHTilingData->vHeadDim; + chunkSize = kdaFwdHTilingData->chunkSize; + useInitialState = kdaFwdHTilingData->useInitialState; + storeFinalState = kdaFwdHTilingData->storeFinalState; + isVariedLen = kdaFwdHTilingData->isVariedLen; + shapeBatch = kdaFwdHTilingData->shapeBatch; + tokenBatch = kdaFwdHTilingData->tokenBatch; + vWorkspaceOffset = kdaFwdHTilingData->vWorkspaceOffset; + vUpdateWorkspaceOffset = kdaFwdHTilingData->vUpdateWorkspaceOffset; + hWorkspaceOffset = kdaFwdHTilingData->hWorkspaceOffset; + numSeqWorkspaceOffset = kdaFwdHTilingData->numSeqWorkspaceOffset; + numChunksWorkspaceOffset = kdaFwdHTilingData->numChunksWorkspaceOffset; + kDecayWorkspaceOffset = kdaFwdHTilingData->kDecayWorkspaceOffset; + uint64_t effectiveTaskCount = static_cast(isVariedLen ? tokenBatch : shapeBatch) * vNumHead; + useDirectFp32Ub = std::is_same::value && chunkSize <= 64 && kHeadDim == 128 && + vHeadDim == 128 && effectiveTaskCount >= AscendC::GetBlockNum(); + + gmK.SetGlobalBuffer((__gm__ ElementK *)k); + gmW.SetGlobalBuffer((__gm__ ElementW *)w); + gmU.SetGlobalBuffer((__gm__ ElementU *)u); + gmG.SetGlobalBuffer((__gm__ ElementG *)(scalarGated ? g : gk)); + gmGk.SetGlobalBuffer((__gm__ ElementG *)(kGated ? gk : g)); + gmInitialState.SetGlobalBuffer((__gm__ ElementInitialState *)inital_state); + gmH.SetGlobalBuffer((__gm__ ElementH *)h); + gmV.SetGlobalBuffer((__gm__ ElementV *)v_new); + gmFinalState.SetGlobalBuffer((__gm__ ElementFinalState *)final_state); + gmVWorkspace.SetGlobalBuffer((__gm__ ElementVWork *)(user + vWorkspaceOffset)); + gmVUpdateWorkspace.SetGlobalBuffer((__gm__ ElementV *)(user + vUpdateWorkspaceOffset)); + gmHWorkspace.SetGlobalBuffer((__gm__ ElementHWork *)(user + hWorkspaceOffset)); + gmKDecayWorkspace.SetGlobalBuffer((__gm__ ElementK *)(user + kDecayWorkspaceOffset)); + + gmSeqlen.SetGlobalBuffer((__gm__ int64_t *)cu_seqlens); + gmNumSeq.SetGlobalBuffer((__gm__ int64_t *)(user + numSeqWorkspaceOffset)); + gmNumChunks.SetGlobalBuffer((__gm__ int64_t *)(user + numChunksWorkspaceOffset)); + + ubHUpdatePing = resource.ubBuf.template GetBufferByByte(32 * 1024); + ubHUpdatePong = resource.ubBuf.template GetBufferByByte(96 * 1024); + ubVWorkPing = resource.ubBuf.template GetBufferByByte(32 * 1024); + ubVWorkPong = resource.ubBuf.template GetBufferByByte(96 * 1024); + + l1VUpdatePing = resource.l1Buf.template GetBufferByByte(0); + l1VUpdatePong = resource.l1Buf.template GetBufferByByte(chunkSize * vHeadDim * sizeof(ElementV)); + + if ASCEND_IS_AIC { + cubeBlockScheduler.Init(cu_seqlens, chunk_indices, tiling, user); + } + + if ASCEND_IS_AIV { + vecBlockScheduler.Init(cu_seqlens, chunk_indices, tiling, user); + } + } + + template + __aicore__ inline void InitFromData(GM_ADDR k, GM_ADDR w, GM_ADDR u, GM_ADDR g, GM_ADDR gk, GM_ADDR inital_state, + GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR h, GM_ADDR v_new, + GM_ADDR final_state, const TilingData &tilingData, GM_ADDR user) + { + batch = tilingData.batch; + seqlen = tilingData.seqlen; + kNumHead = tilingData.kNumHead; + vNumHead = tilingData.vNumHead; + kHeadDim = tilingData.kHeadDim; + vHeadDim = tilingData.vHeadDim; + chunkSize = tilingData.chunkSize; + useInitialState = tilingData.useInitialState; + storeFinalState = tilingData.storeFinalState; + isVariedLen = tilingData.isVariedLen; + shapeBatch = tilingData.shapeBatch; + tokenBatch = tilingData.tokenBatch; + vWorkspaceOffset = tilingData.vWorkspaceOffset; + vUpdateWorkspaceOffset = tilingData.vUpdateWorkspaceOffset; + hWorkspaceOffset = tilingData.hWorkspaceOffset; + numSeqWorkspaceOffset = tilingData.numSeqWorkspaceOffset; + numChunksWorkspaceOffset = tilingData.numChunksWorkspaceOffset; + kDecayWorkspaceOffset = tilingData.kDecayWorkspaceOffset; + uint64_t effectiveTaskCount = static_cast(isVariedLen ? tokenBatch : shapeBatch) * vNumHead; + useDirectFp32Ub = std::is_same::value && chunkSize <= 64 && kHeadDim == 128 && + vHeadDim == 128 && effectiveTaskCount >= AscendC::GetBlockNum(); + + gmK.SetGlobalBuffer((__gm__ ElementK *)k); + gmW.SetGlobalBuffer((__gm__ ElementW *)w); + gmU.SetGlobalBuffer((__gm__ ElementU *)u); + gmG.SetGlobalBuffer((__gm__ ElementG *)(scalarGated ? g : gk)); + gmInitialState.SetGlobalBuffer((__gm__ ElementInitialState *)inital_state); + gmH.SetGlobalBuffer((__gm__ ElementH *)h); + gmV.SetGlobalBuffer((__gm__ ElementV *)v_new); + gmFinalState.SetGlobalBuffer((__gm__ ElementFinalState *)final_state); + gmVWorkspace.SetGlobalBuffer((__gm__ ElementVWork *)(user + vWorkspaceOffset)); + gmVUpdateWorkspace.SetGlobalBuffer((__gm__ ElementV *)(user + vUpdateWorkspaceOffset)); + gmHWorkspace.SetGlobalBuffer((__gm__ ElementHWork *)(user + hWorkspaceOffset)); + gmGk.SetGlobalBuffer((__gm__ ElementG *)(kGated ? gk : g)); + gmKDecayWorkspace.SetGlobalBuffer((__gm__ ElementK *)(user + kDecayWorkspaceOffset)); + gmSeqlen.SetGlobalBuffer((__gm__ int64_t *)cu_seqlens); + gmNumSeq.SetGlobalBuffer((__gm__ int64_t *)(user + numSeqWorkspaceOffset)); + gmNumChunks.SetGlobalBuffer((__gm__ int64_t *)(user + numChunksWorkspaceOffset)); + + ubHUpdatePing = resource.ubBuf.template GetBufferByByte(32 * 1024); + ubHUpdatePong = resource.ubBuf.template GetBufferByByte(96 * 1024); + ubVWorkPing = resource.ubBuf.template GetBufferByByte(32 * 1024); + ubVWorkPong = resource.ubBuf.template GetBufferByByte(96 * 1024); + + l1VUpdatePing = resource.l1Buf.template GetBufferByByte(0); + l1VUpdatePong = resource.l1Buf.template GetBufferByByte(chunkSize * vHeadDim * sizeof(ElementV)); + + if ASCEND_IS_AIC { + cubeBlockScheduler.InitFromData(cu_seqlens, chunk_indices, tilingData, user); + } + if ASCEND_IS_AIV { + vecBlockScheduler.InitFromData(cu_seqlens, chunk_indices, tilingData, user); + } + } + + // Tail helpers borrow the stream's V_MTE2 free token and restore it before returning. + __aicore__ inline void ComputeTailVWorkspace(const KdaFwdHOffsets &offsets, uint32_t tailEventId) + { + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t rowsPerSubBlock = CeilDiv(offsets.blockTokens, subBlockNum); + uint32_t rowBegin = subBlockIdx * rowsPerSubBlock; + uint32_t rowEnd = Min(rowBegin + rowsPerSubBlock, offsets.blockTokens); + if (rowBegin >= rowEnd) { + return; + } + AscendC::ResetMask(); + AscendC::WaitFlag(tailEventId); + + if (offsets.vBlockDim == vHeadDim && vHeadDim <= 256) { + constexpr uint32_t TAIL_H_INPUT_OFFSET = 0; + constexpr uint32_t TAIL_W_INPUT_OFFSET = 166 * 1024; + constexpr uint32_t TAIL_OUTPUT_OFFSET = 168 * 1024; + uint32_t rowCount = rowEnd - rowBegin; + AscendC::LocalTensor hInputUb = + resource.ubBuf.template GetBufferByByte(TAIL_H_INPUT_OFFSET); + AscendC::LocalTensor wInputUb = + resource.ubBuf.template GetBufferByByte(TAIL_W_INPUT_OFFSET); + AscendC::LocalTensor outputUb = resource.ubBuf.template GetBufferByByte(TAIL_OUTPUT_OFFSET); + AscendC::DataCopy(hInputUb, gmH[offsets.hSrcOffset], kHeadDim * offsets.vBlockDim); + AscendC::DataCopy(wInputUb, gmW[offsets.wOffset + rowBegin * kHeadDim], rowCount * kHeadDim); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + __ubuf__ float *outputAddr = reinterpret_cast<__ubuf__ float *>(outputUb.GetPhyAddr()); + __ubuf__ ElementW *wInputAddr = reinterpret_cast<__ubuf__ ElementW *>(wInputUb.GetPhyAddr()); + __ubuf__ ElementH *hInputAddr = reinterpret_cast<__ubuf__ ElementH *>(hInputUb.GetPhyAddr()); + AscendC::VF_CALL< + Catlass::Epilogue::Block::detail::ComputeKdaTailVWorkspaceRegbaseDualIssue>( + outputAddr, wInputAddr, hInputAddr, static_cast(rowCount), + static_cast(offsets.vBlockDim), static_cast(kHeadDim)); + AscendC::PipeBarrier(); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::DataCopy(gmVWorkspace[offsets.vWorkOffset + rowBegin * offsets.vBlockDim], outputUb, + rowCount * offsets.vBlockDim); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::SetFlag(tailEventId); + return; + } + + constexpr uint32_t TAIL_INPUT_OFFSET = 166 * 1024; + constexpr uint32_t TAIL_FLOAT_OFFSET = 167 * 1024; + constexpr uint32_t TAIL_ACCUM_OFFSET = 168 * 1024; + constexpr uint32_t TAIL_WEIGHT_INPUT_OFFSET = 169 * 1024; + constexpr uint32_t TAIL_WEIGHT_FLOAT_OFFSET = 170 * 1024; + AscendC::LocalTensor inputUb = resource.ubBuf.template GetBufferByByte(TAIL_INPUT_OFFSET); + AscendC::LocalTensor floatUb = resource.ubBuf.template GetBufferByByte(TAIL_FLOAT_OFFSET); + AscendC::LocalTensor accumUb = resource.ubBuf.template GetBufferByByte(TAIL_ACCUM_OFFSET); + AscendC::LocalTensor weightInputUb = + resource.ubBuf.template GetBufferByByte(TAIL_WEIGHT_INPUT_OFFSET); + AscendC::LocalTensor weightFloatUb = + resource.ubBuf.template GetBufferByByte(TAIL_WEIGHT_FLOAT_OFFSET); + + for (uint32_t tokenRow = rowBegin; tokenRow < rowEnd; ++tokenRow) { + AscendC::DataCopy(weightInputUb, gmW[offsets.wOffset + tokenRow * kHeadDim], kHeadDim); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::Cast(weightFloatUb, weightInputUb, AscendC::RoundMode::CAST_NONE, kHeadDim); + AscendC::PipeBarrier(); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + + AscendC::Duplicate(accumUb, 0.0f, offsets.vBlockDim); + AscendC::PipeBarrier(); + for (uint32_t kIdx = 0; kIdx < kHeadDim; ++kIdx) { + AscendC::DataCopy(inputUb, gmH[offsets.hSrcOffset + kIdx * vHeadDim], offsets.vBlockDim); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::Cast(floatUb, inputUb, AscendC::RoundMode::CAST_NONE, offsets.vBlockDim); + AscendC::PipeBarrier(); + float weight = weightFloatUb.GetValue(kIdx); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::Muls(floatUb, floatUb, weight, offsets.vBlockDim); + AscendC::PipeBarrier(); + AscendC::Add(accumUb, accumUb, floatUb, offsets.vBlockDim); + AscendC::PipeBarrier(); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + } + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::DataCopy(gmVWorkspace[offsets.vWorkOffset + tokenRow * offsets.vBlockDim], accumUb, + offsets.vBlockDim); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + } + AscendC::SetFlag(tailEventId); + } + + __aicore__ inline void ComputeTailHWorkspace(const KdaFwdHOffsets &offsets, uint32_t tailEventId) + { + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t rowsPerSubBlock = CeilDiv(kHeadDim, subBlockNum); + uint32_t rowBegin = subBlockIdx * rowsPerSubBlock; + uint32_t rowEnd = Min(rowBegin + rowsPerSubBlock, kHeadDim); + AscendC::ResetMask(); + AscendC::WaitFlag(tailEventId); + + constexpr uint32_t TAIL_INPUT_OFFSET = 166 * 1024; + constexpr uint32_t TAIL_FLOAT_OFFSET = 167 * 1024; + constexpr uint32_t TAIL_ACCUM_OFFSET = 168 * 1024; + constexpr uint32_t TAIL_WEIGHT_INPUT_OFFSET = 169 * 1024; + constexpr uint32_t TAIL_WEIGHT_FLOAT_OFFSET = 170 * 1024; + AscendC::LocalTensor inputUb = resource.ubBuf.template GetBufferByByte(TAIL_INPUT_OFFSET); + AscendC::LocalTensor floatUb = resource.ubBuf.template GetBufferByByte(TAIL_FLOAT_OFFSET); + AscendC::LocalTensor accumUb = resource.ubBuf.template GetBufferByByte(TAIL_ACCUM_OFFSET); + AscendC::LocalTensor weightInputUb = + resource.ubBuf.template GetBufferByByte(TAIL_WEIGHT_INPUT_OFFSET); + AscendC::LocalTensor weightFloatUb = + resource.ubBuf.template GetBufferByByte(TAIL_WEIGHT_FLOAT_OFFSET); + + for (uint32_t kRow = rowBegin; kRow < rowEnd; ++kRow) { + AscendC::Duplicate(accumUb, 0.0f, offsets.vBlockDim); + AscendC::PipeBarrier(); + for (uint32_t tokenRow = 0; tokenRow < offsets.blockTokens; ++tokenRow) { + AscendC::DataCopy(weightInputUb, gmKDecayWorkspace[offsets.kDecayWorkOffset + tokenRow * kHeadDim], + kHeadDim); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::Cast(weightFloatUb, weightInputUb, AscendC::RoundMode::CAST_NONE, kHeadDim); + AscendC::PipeBarrier(); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::DataCopy(inputUb, gmVUpdateWorkspace[offsets.vWorkOffset + tokenRow * offsets.vBlockDim], + offsets.vBlockDim); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::Cast(floatUb, inputUb, AscendC::RoundMode::CAST_NONE, offsets.vBlockDim); + AscendC::PipeBarrier(); + float weight = weightFloatUb.GetValue(kRow); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::Muls(floatUb, floatUb, weight, offsets.vBlockDim); + AscendC::PipeBarrier(); + AscendC::Add(accumUb, accumUb, floatUb, offsets.vBlockDim); + AscendC::PipeBarrier(); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + } + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::DataCopy(gmHWorkspace[offsets.hWorkOffset + kRow * offsets.vBlockDim], accumUb, offsets.vBlockDim); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + } + AscendC::SetFlag(tailEventId); + } + + __aicore__ inline void ComputeKdaTailFinalState(const KdaFwdHOffsets &offsets, uint32_t tailEventId) + { + static_assert(kGated && !scalarGated && useExp2 && std::is_same::value, + "The fused tail state path is specific to FP32 KDA state"); + constexpr uint32_t TAIL_STATE_OFFSET = 0; + constexpr uint32_t TAIL_K_INPUT_OFFSET = 166 * 1024; + constexpr uint32_t TAIL_V_INPUT_OFFSET = 170 * 1024; + constexpr uint32_t TAIL_GATE_INPUT_OFFSET = 174 * 1024; + constexpr uint32_t TAIL_GATE_SCALE_OFFSET = 175 * 1024; + + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t rowsPerSubBlock = CeilDiv(kHeadDim, subBlockNum); + uint32_t rowBegin = subBlockIdx * rowsPerSubBlock; + uint32_t rowEnd = Min(rowBegin + rowsPerSubBlock, kHeadDim); + if (rowBegin >= rowEnd) { + return; + } + + AscendC::ResetMask(); + AscendC::WaitFlag(tailEventId); + AscendC::LocalTensor stateUb = resource.ubBuf.template GetBufferByByte(TAIL_STATE_OFFSET); + AscendC::LocalTensor kInputUb = + resource.ubBuf.template GetBufferByByte(TAIL_K_INPUT_OFFSET); + AscendC::LocalTensor vInputUb = + resource.ubBuf.template GetBufferByByte(TAIL_V_INPUT_OFFSET); + AscendC::LocalTensor gateInputUb = + resource.ubBuf.template GetBufferByByte(TAIL_GATE_INPUT_OFFSET); + AscendC::LocalTensor gateScaleUb = + resource.ubBuf.template GetBufferByByte(TAIL_GATE_SCALE_OFFSET); + + uint32_t gateOffset = offsets.gkOffset + (offsets.blockTokens - 1) * kHeadDim + rowBegin; + uint32_t rowCount = rowEnd - rowBegin; + uint32_t stateElements = rowCount * offsets.vBlockDim; + uint32_t kElements = offsets.blockTokens * kHeadDim; + uint32_t vElements = offsets.blockTokens * offsets.vBlockDim; + + if (offsets.isInitialState) { + if (useInitialState) { + AscendC::DataCopy(stateUb, gmInitialState[offsets.initialStateOffset + rowBegin * offsets.vBlockDim], + stateElements); + } else { + AscendC::Duplicate(stateUb, 0.0f, stateElements); + } + } else { + AscendC::DataCopy(stateUb, gmFinalState[offsets.finalStateOffset + rowBegin * offsets.vBlockDim], + stateElements); + } + AscendC::DataCopy(kInputUb, gmK[offsets.wkOffset], kElements); + AscendC::DataCopy(vInputUb, gmV[offsets.uvOffset], vElements); + AscendC::DataCopy(gateInputUb, gmGk[gateOffset], rowCount); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + + __ubuf__ float *stateAddr = reinterpret_cast<__ubuf__ float *>(stateUb.GetPhyAddr()); + __ubuf__ ElementK *kInputAddr = reinterpret_cast<__ubuf__ ElementK *>(kInputUb.GetPhyAddr()); + __ubuf__ ElementV *vInputAddr = reinterpret_cast<__ubuf__ ElementV *>(vInputUb.GetPhyAddr()); + __ubuf__ ElementG *gateInputAddr = reinterpret_cast<__ubuf__ ElementG *>(gateInputUb.GetPhyAddr()); + __ubuf__ float *gateScaleAddr = reinterpret_cast<__ubuf__ float *>(gateScaleUb.GetPhyAddr()); + AscendC::VF_CALL>( + gateScaleAddr, gateInputAddr, static_cast(rowCount)); + AscendC::VF_CALL< + Catlass::Epilogue::Block::detail::ComputeKdaTailFinalStateRegbaseDualIssue>( + stateAddr, kInputAddr, vInputAddr, gateScaleAddr, static_cast(rowCount), + static_cast(offsets.vBlockDim), static_cast(offsets.blockTokens), + static_cast(kHeadDim), static_cast(rowBegin)); + AscendC::PipeBarrier(); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::DataCopy(gmFinalState[offsets.finalStateOffset + rowBegin * offsets.vBlockDim], stateUb, + stateElements); + AscendC::SetFlag(tailEventId); + AscendC::WaitFlag(tailEventId); + AscendC::SetFlag(tailEventId); + } + + __aicore__ inline void Process() + { + // FwdH can run after another stage in a megakernel. Start its AIC/AIV + // handshake only after every core has retired the preceding stage. + AscendC::SyncAll(); + + if ASCEND_IS_AIC { + uint32_t coreIdx = AscendC::GetBlockIdx(); + uint32_t coreNum = vecBlockScheduler.cubeCoreNum; + + BlockMmadWH blockMmadWH(resource, + chunkSize * cubeBlockScheduler.vBlockSize * sizeof(ElementV) * PING_PONG_STAGES); + BlockMmadKV blockMmadKV(resource, + chunkSize * cubeBlockScheduler.vBlockSize * sizeof(ElementV) * PING_PONG_STAGES); + BlockMmadWHTail blockMmadWHTail( + resource, chunkSize * cubeBlockScheduler.vBlockSize * sizeof(ElementV) * PING_PONG_STAGES); + BlockMmadKVTail blockMmadKVTail( + resource, chunkSize * cubeBlockScheduler.vBlockSize * sizeof(ElementV) * PING_PONG_STAGES); + bool useBoundedMmad = isVariedLen || (seqlen % chunkSize != 0); + + auto wLayout = + tla::MakeLayout(shapeBatch * kNumHead * cubeBlockScheduler.totalTokens, kHeadDim); + auto hLayout = tla::MakeLayout( + shapeBatch * vNumHead * cubeBlockScheduler.totalChunks * kHeadDim, vHeadDim); + + auto kLayout = + tla::MakeLayout(kHeadDim, shapeBatch * kNumHead * cubeBlockScheduler.totalTokens); + auto hworkLayout = tla::MakeLayout(kHeadDim, cubeBlockScheduler.vBlockSize); + + AscendC::SyncAll(); + uint32_t currStage = 0; // 0: C1, 1: C2 + while (cubeBlockScheduler.isRunning) { + if (currStage == 0) { + /* C1: v_work = w @ h[i] */ + cubeBlockScheduler.InitTasks(); + if (useDirectFp32Ub) { + BlockMmadWHDirectUb blockMmadWHDirectUb( + resource, chunkSize * cubeBlockScheduler.vBlockSize * sizeof(ElementV) * PING_PONG_STAGES); + for (uint32_t i = 0; i < PING_PONG_STAGES; ++i) { + uint32_t streamId = cubeBlockScheduler.GetStreamId(i); + const auto &stream = cubeBlockScheduler.GetStream(i); + if (cubeBlockScheduler.StreamIsDone(stream)) { + continue; + } + + const KdaFwdHOffsets &cube1Offsets = cubeBlockScheduler.GetCurTaskOffsets(stream); + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec2Done[streamId]); + if (cube1Offsets.blockTokens < 16) { + Arch::CrossCoreSetFlag<0x2, PIPE_MTE2>(cubeBlockScheduler.cube1Done[streamId]); + continue; + } + int64_t cube1OffsetW = cube1Offsets.wOffset; + int64_t cube1OffsetH = cube1Offsets.hSrcOffset; + auto tensorW = tla::MakeTensor(gmW[cube1OffsetW], wLayout, Catlass::Arch::PositionGM{}); + auto tensorH = tla::MakeTensor(gmH[cube1OffsetH], hLayout, Catlass::Arch::PositionGM{}); + GemmCoord cube1Shape{cube1Offsets.blockTokens, cube1Offsets.vBlockDim, kHeadDim}; + auto tensorBlockW = + GetTile(tensorW, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.m(), cube1Shape.k())); + auto tensorBlockH = + GetTile(tensorH, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.k(), cube1Shape.n())); + + auto ubLayout = tla::MakeLayout(cube1Shape.m(), cube1Shape.n()); + auto tensorUbPing = tla::MakeTensor(ubVWorkPing, ubLayout, Catlass::Arch::PositionUB{}); + auto tensorUbPong = tla::MakeTensor(ubVWorkPong, ubLayout, Catlass::Arch::PositionUB{}); + using UbTensor = decltype(tensorUbPing); + UbTensor tensorUbList[BlockMmadWHDirectUb::MAX_CUBE_VEC_SYNC_NUM]; + for (uint32_t ubIdx = 0; ubIdx < BlockMmadWHDirectUb::MAX_CUBE_VEC_SYNC_NUM; ++ubIdx) { + tensorUbList[ubIdx] = (ubIdx & 1U) ? tensorUbPong : tensorUbPing; + } + uint32_t ubListId = streamId; + uint32_t rowsPerSubBlock = CeilDiv(cube1Shape.m(), DIRECT_VEC_NUM); + blockMmadWHDirectUb(tensorBlockW, tensorBlockH, tensorUbList, cube1Shape, rowsPerSubBlock, + 0, DIRECT_UB_FREE_FLAG_BEGIN, DIRECT_UB_READY_FLAG_BEGIN, ubListId, + DIRECT_VEC_NUM, DIRECT_UB_STAGES); + } + } else if (useBoundedMmad) { + for (uint32_t i = 0; i < PING_PONG_STAGES; ++i) { + uint32_t streamId = cubeBlockScheduler.GetStreamId(i); + const auto &stream = cubeBlockScheduler.GetStream(i); + if (cubeBlockScheduler.StreamIsDone(stream)) { + continue; + } + + const KdaFwdHOffsets &cube1Offsets = cubeBlockScheduler.GetCurTaskOffsets(stream); + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec2Done[streamId]); + if (cube1Offsets.blockTokens < 16) { + Arch::CrossCoreSetFlag<0x2, PIPE_MTE2>(cubeBlockScheduler.cube1Done[streamId]); + continue; + } + auto vLayout = tla::MakeLayout(cube1Offsets.blockTokens, + cube1Offsets.vBlockDim); + auto tensorW = + tla::MakeTensor(gmW[cube1Offsets.wOffset], wLayout, Catlass::Arch::PositionGM{}); + auto tensorH = + tla::MakeTensor(gmH[cube1Offsets.hSrcOffset], hLayout, Catlass::Arch::PositionGM{}); + auto tensorV = tla::MakeTensor(gmVWorkspace[cube1Offsets.vWorkOffset], vLayout, + Catlass::Arch::PositionGM{}); + GemmCoord cube1Shape{cube1Offsets.blockTokens, cube1Offsets.vBlockDim, kHeadDim}; + auto tensorBlockW = + GetTile(tensorW, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.m(), cube1Shape.k())); + auto tensorBlockH = + GetTile(tensorH, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.k(), cube1Shape.n())); + auto tensorBlockV = + GetTile(tensorV, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.m(), cube1Shape.n())); + + blockMmadWHTail.preSetFlags(); + blockMmadWHTail(tensorBlockW, tensorBlockH, tensorBlockV, cube1Shape, EmptyClass{}, true); + blockMmadWHTail.finalWaitFlags(); + AscendC::PipeBarrier(); + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(cubeBlockScheduler.cube1Done[streamId]); + } + } else { + blockMmadWH.preSetFlags(); + for (uint32_t i = 0; i < PING_PONG_STAGES; ++i) { + uint32_t streamId = cubeBlockScheduler.GetStreamId(i); + const auto &stream = cubeBlockScheduler.GetStream(i); + if (cubeBlockScheduler.StreamIsDone(stream)) { + continue; + } + + const KdaFwdHOffsets &cube1Offsets = cubeBlockScheduler.GetCurTaskOffsets(stream); + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec2Done[streamId]); + if (cube1Offsets.blockTokens < 16) { + Arch::CrossCoreSetFlag<0x2, PIPE_MTE2>(cubeBlockScheduler.cube1Done[streamId]); + continue; + } + auto vLayout = tla::MakeLayout(cube1Offsets.blockTokens, + cube1Offsets.vBlockDim); + int64_t cube1OffsetW = cube1Offsets.wOffset; + int64_t cube1OffsetH = cube1Offsets.hSrcOffset; + int64_t cube1OffsetVwork = cube1Offsets.vWorkOffset; + auto tensorW = tla::MakeTensor(gmW[cube1OffsetW], wLayout, Catlass::Arch::PositionGM{}); + auto tensorH = tla::MakeTensor(gmH[cube1OffsetH], hLayout, Catlass::Arch::PositionGM{}); + auto tensorV = + tla::MakeTensor(gmVWorkspace[cube1OffsetVwork], vLayout, Catlass::Arch::PositionGM{}); + GemmCoord cube1Shape{cube1Offsets.blockTokens, cube1Offsets.vBlockDim, kHeadDim}; + auto tensorBlockW = + GetTile(tensorW, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.m(), cube1Shape.k())); + auto tensorBlockH = + GetTile(tensorH, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.k(), cube1Shape.n())); + auto tensorBlockV = + GetTile(tensorV, tla::MakeCoord(0, 0), tla::MakeShape(cube1Shape.m(), cube1Shape.n())); + + blockMmadWH(tensorBlockW, tensorBlockH, tensorBlockV, cube1Shape); + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(cubeBlockScheduler.cube1Done[streamId]); + } + blockMmadWH.finalWaitFlags(); + } + } else { + /* C2: h[i+1] = k.T @ v_work */ + if (useDirectFp32Ub) { + BlockMmadKVDirectUb blockMmadKVDirectUb( + resource, chunkSize * cubeBlockScheduler.vBlockSize * sizeof(ElementV) * PING_PONG_STAGES); + for (uint32_t i = 0; i < PING_PONG_STAGES; ++i) { + uint32_t streamId = cubeBlockScheduler.GetStreamId(i); + const auto &stream = cubeBlockScheduler.GetStream(i); + if (cubeBlockScheduler.StreamIsDone(stream)) { + continue; + } + const KdaFwdHOffsets &cube2Offsets = cubeBlockScheduler.GetCurTaskOffsets(stream); + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec1Done[streamId]); + + if (cubeBlockScheduler.NeedProcessStage2(stream)) { + if (cube2Offsets.blockTokens < 16) { + Arch::CrossCoreSetFlag<0x2, PIPE_MTE2>(cubeBlockScheduler.cube2Done[streamId]); + continue; + } + int64_t cube2OffsetK = kGated ? cube2Offsets.kDecayWorkOffset : cube2Offsets.wkOffset; + int64_t cube2OffsetVwork = cube2Offsets.vWorkOffset; + auto tensorK = + kGated ? tla::MakeTensor(gmKDecayWorkspace[cube2OffsetK], kLayout, + Catlass::Arch::PositionGM{}) : + tla::MakeTensor(gmK[cube2OffsetK], kLayout, Catlass::Arch::PositionGM{}); + auto vUpdateLayout = tla::MakeLayout( + cube2Offsets.blockTokens, cube2Offsets.vBlockDim); + auto tensorVwork = tla::MakeTensor(gmVUpdateWorkspace[cube2OffsetVwork], vUpdateLayout, + Catlass::Arch::PositionGM{}); + GemmCoord cube2Shape{kHeadDim, cube2Offsets.vBlockDim, cube2Offsets.blockTokens}; + auto tensorBlockK = GetTile(tensorK, tla::MakeCoord(0, 0), + tla::MakeShape(cube2Shape.m(), cube2Shape.k())); + auto tensorBlockVwork = GetTile(tensorVwork, tla::MakeCoord(0, 0), + tla::MakeShape(cube2Shape.k(), cube2Shape.n())); + + auto ubLayout = tla::MakeLayout(cube2Shape.m(), cube2Shape.n()); + auto tensorUbPing = + tla::MakeTensor(ubHUpdatePing, ubLayout, Catlass::Arch::PositionUB{}); + auto tensorUbPong = + tla::MakeTensor(ubHUpdatePong, ubLayout, Catlass::Arch::PositionUB{}); + using UbTensor = decltype(tensorUbPing); + UbTensor tensorUbList[BlockMmadKVDirectUb::MAX_CUBE_VEC_SYNC_NUM]; + for (uint32_t ubIdx = 0; ubIdx < BlockMmadKVDirectUb::MAX_CUBE_VEC_SYNC_NUM; ++ubIdx) { + tensorUbList[ubIdx] = (ubIdx & 1U) ? tensorUbPong : tensorUbPing; + } + uint32_t ubListId = streamId; + uint32_t rowsPerSubBlock = CeilDiv(cube2Shape.m(), DIRECT_VEC_NUM); + blockMmadKVDirectUb(tensorBlockK, tensorBlockVwork, tensorUbList, cube2Shape, + rowsPerSubBlock, 0, DIRECT_UB_FREE_FLAG_BEGIN, + DIRECT_UB_READY_FLAG_BEGIN, ubListId, DIRECT_VEC_NUM, + DIRECT_UB_STAGES); + } + } + } else if (useBoundedMmad) { + for (uint32_t i = 0; i < PING_PONG_STAGES; ++i) { + uint32_t streamId = cubeBlockScheduler.GetStreamId(i); + const auto &stream = cubeBlockScheduler.GetStream(i); + if (cubeBlockScheduler.StreamIsDone(stream)) { + continue; + } + const KdaFwdHOffsets &cube2Offsets = cubeBlockScheduler.GetCurTaskOffsets(stream); + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec1Done[streamId]); + + if (cubeBlockScheduler.NeedProcessStage2(stream)) { + if (cube2Offsets.blockTokens < 16) { + Arch::CrossCoreSetFlag<0x2, PIPE_MTE2>(cubeBlockScheduler.cube2Done[streamId]); + continue; + } + int64_t cube2OffsetK = kGated ? cube2Offsets.kDecayWorkOffset : cube2Offsets.wkOffset; + auto tensorK = + kGated ? tla::MakeTensor(gmKDecayWorkspace[cube2OffsetK], kLayout, + Catlass::Arch::PositionGM{}) : + tla::MakeTensor(gmK[cube2OffsetK], kLayout, Catlass::Arch::PositionGM{}); + auto vUpdateLayout = tla::MakeLayout( + cube2Offsets.blockTokens, cube2Offsets.vBlockDim); + auto tensorVwork = tla::MakeTensor(gmVUpdateWorkspace[cube2Offsets.vWorkOffset], + vUpdateLayout, Catlass::Arch::PositionGM{}); + auto tensorHwork = tla::MakeTensor(gmHWorkspace[cube2Offsets.hWorkOffset], hworkLayout, + Catlass::Arch::PositionGM{}); + GemmCoord cube2Shape{kHeadDim, cube2Offsets.vBlockDim, cube2Offsets.blockTokens}; + auto tensorBlockK = GetTile(tensorK, tla::MakeCoord(0, 0), + tla::MakeShape(cube2Shape.m(), cube2Shape.k())); + auto tensorBlockVwork = GetTile(tensorVwork, tla::MakeCoord(0, 0), + tla::MakeShape(cube2Shape.k(), cube2Shape.n())); + auto tensorBlockHwork = GetTile(tensorHwork, tla::MakeCoord(0, 0), + tla::MakeShape(cube2Shape.m(), cube2Shape.n())); + + blockMmadKVTail.preSetFlags(); + blockMmadKVTail(tensorBlockK, tensorBlockVwork, tensorBlockHwork, cube2Shape, + EmptyClass{}, true); + blockMmadKVTail.finalWaitFlags(); + AscendC::PipeBarrier(); + } + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(cubeBlockScheduler.cube2Done[streamId]); + } + } else { + blockMmadKV.preSetFlags(); + for (uint32_t i = 0; i < PING_PONG_STAGES; ++i) { + uint32_t streamId = cubeBlockScheduler.GetStreamId(i); + const auto &stream = cubeBlockScheduler.GetStream(i); + if (cubeBlockScheduler.StreamIsDone(stream)) { + continue; + } + const KdaFwdHOffsets &cube2Offsets = cubeBlockScheduler.GetCurTaskOffsets(stream); + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec1Done[streamId]); + + if (cubeBlockScheduler.NeedProcessStage2(stream)) { + if (cube2Offsets.blockTokens < 16) { + Arch::CrossCoreSetFlag<0x2, PIPE_MTE2>(cubeBlockScheduler.cube2Done[streamId]); + continue; + } + // step 3: h[i+1] = k.T @ v_work + int64_t cube2OffsetK = kGated ? cube2Offsets.kDecayWorkOffset : cube2Offsets.wkOffset; + int64_t cube2OffsetVwork = cube2Offsets.vWorkOffset; + auto tensorK = + kGated ? tla::MakeTensor(gmKDecayWorkspace[cube2OffsetK], kLayout, + Catlass::Arch::PositionGM{}) : + tla::MakeTensor(gmK[cube2OffsetK], kLayout, Catlass::Arch::PositionGM{}); + auto vUpdateLayout = tla::MakeLayout( + cube2Offsets.blockTokens, cube2Offsets.vBlockDim); + auto tensorVwork = tla::MakeTensor(gmVUpdateWorkspace[cube2OffsetVwork], vUpdateLayout, + Catlass::Arch::PositionGM{}); + auto tensorHwork = tla::MakeTensor(gmHWorkspace[cube2Offsets.hWorkOffset], hworkLayout, + Catlass::Arch::PositionGM{}); + GemmCoord cube2Shape{kHeadDim, cube2Offsets.vBlockDim, cube2Offsets.blockTokens}; + auto tensorBlockK = GetTile(tensorK, tla::MakeCoord(0, 0), + tla::MakeShape(cube2Shape.m(), cube2Shape.k())); + auto tensorBlockVwork = GetTile(tensorVwork, tla::MakeCoord(0, 0), + tla::MakeShape(cube2Shape.k(), cube2Shape.n())); + auto tensorBlockHwork = GetTile(tensorHwork, tla::MakeCoord(0, 0), + tla::MakeShape(cube2Shape.m(), cube2Shape.n())); + + blockMmadKV(tensorBlockK, tensorBlockVwork, tensorBlockHwork, cube2Shape); + } + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(cubeBlockScheduler.cube2Done[streamId]); + } + blockMmadKV.finalWaitFlags(); + } + } + currStage ^= 0x01; + } + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec2Done[0]); + Arch::CrossCoreWaitFlag(cubeBlockScheduler.vec2Done[1]); + if (useDirectFp32Ub) { + for (uint32_t slot = 0; slot < DIRECT_UB_STAGES; ++slot) { + AscendC::CrossCoreWaitFlag<0x4, PIPE_FIX>(DIRECT_UB_FREE_FLAG_BEGIN + slot); + AscendC::CrossCoreWaitFlag<0x4, PIPE_FIX>(DIRECT_UB_FREE_FLAG_BEGIN + DIRECT_UB_FLAG_STRIDE + slot); + } + } + } + + if ASCEND_IS_AIV { + uint32_t subBlockIdx = AscendC::GetSubBlockIdx(); + uint32_t subBlockNum = AscendC::GetSubBlockNum(); + uint32_t coreIdx = AscendC::GetBlockIdx() / subBlockNum; + uint32_t coreNum = AscendC::GetBlockNum(); + uint32_t taskCount = (isVariedLen ? vecBlockScheduler.tokenBatch : shapeBatch) * vNumHead; + uint32_t tasksPerCore = taskCount > coreNum ? PING_PONG_STAGES : 1; + uint32_t taskStride = coreNum * tasksPerCore; + uint32_t rowsPerSubBlock = (kHeadDim + subBlockNum - 1) / subBlockNum; + uint32_t rowBegin = subBlockIdx * rowsPerSubBlock; + uint32_t rowEnd = Min(rowBegin + rowsPerSubBlock, kHeadDim); + uint32_t hRowsPerTile = (32 * 1024) / (vHeadDim * sizeof(ElementH)); + uint32_t stateRowsPerTile = (64 * 1024) / (vHeadDim * sizeof(ElementInitialState)); + uint32_t rowsPerTile = Min(hRowsPerTile, stateRowsPerTile); + uint32_t totalChunks = isVariedLen ? vecBlockScheduler.totalChunks : ((seqlen + chunkSize - 1) / chunkSize); + uint32_t stateBlockSize = kHeadDim * vHeadDim; + uint32_t pingpongFlag = 1; + AscendC::LocalTensor stateUbTensorPing = + resource.ubBuf.template GetBufferByByte(0); + AscendC::LocalTensor stateUbTensorPong = + resource.ubBuf.template GetBufferByByte(96 * 1024); + AscendC::LocalTensor hUbTensorPing = resource.ubBuf.template GetBufferByByte(64 * 1024); + AscendC::LocalTensor hUbTensorPong = + resource.ubBuf.template GetBufferByByte(160 * 1024); + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID1); + for (uint32_t slot = 0; slot < tasksPerCore; ++slot) { + for (uint32_t taskIdx = coreIdx * tasksPerCore + slot; taskIdx < taskCount; taskIdx += taskStride) { + uint32_t batchIdx = taskIdx / vNumHead; + uint32_t vHeadIdx = taskIdx % vNumHead; + uint32_t chunkOffset = isVariedLen ? vecBlockScheduler.GetVarlenChunkOffset(batchIdx) : 0; + uint32_t shapeBatchIdx = isVariedLen ? 0 : batchIdx; + uint32_t hBaseOffset = + (shapeBatchIdx * vNumHead * totalChunks + vHeadIdx * totalChunks + chunkOffset) * + stateBlockSize; + uint32_t initialStateBaseOffset = taskIdx * stateBlockSize; + for (uint32_t rowOffset = rowBegin; rowOffset < rowEnd; rowOffset += rowsPerTile) { + uint32_t rowsThisTile = Min(rowsPerTile, rowEnd - rowOffset); + uint32_t stateTileElems = rowsThisTile * vHeadDim; + uint32_t hOffset = hBaseOffset + rowOffset * vHeadDim; + AscendC::LocalTensor stateUbTensor = + pingpongFlag ? stateUbTensorPing : stateUbTensorPong; + AscendC::LocalTensor hUbTensor = pingpongFlag ? hUbTensorPing : hUbTensorPong; + auto eventId = pingpongFlag ? EVENT_ID1 : EVENT_ID0; + AscendC::WaitFlag(eventId); + if (useInitialState) { + uint32_t initialStateOffset = initialStateBaseOffset + rowOffset * vHeadDim; + if constexpr (!std::is_same::value) { + AscendC::DataCopy(stateUbTensor, gmInitialState[initialStateOffset], stateTileElems); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); + AscendC::Cast(hUbTensor, stateUbTensor, AscendC::RoundMode::CAST_RINT, stateTileElems); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); + AscendC::DataCopy(gmH[hOffset], hUbTensor, stateTileElems); + } else { + AscendC::DataCopy(stateUbTensor, gmInitialState[initialStateOffset], stateTileElems); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); + AscendC::DataCopy(gmH[hOffset], stateUbTensor, stateTileElems); + } + } else { + AscendC::Duplicate(hUbTensor, static_cast(0), stateTileElems); + AscendC::SetFlag(eventId); + AscendC::WaitFlag(eventId); + AscendC::DataCopy(gmH[hOffset], hUbTensor, stateTileElems); + } + AscendC::SetFlag(eventId); + pingpongFlag = 1 - pingpongFlag; + } + } + } + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID1); + + AscendC::SyncAll(); + + if (useDirectFp32Ub) { + for (uint32_t slot = 0; slot < DIRECT_UB_STAGES; ++slot) { + uint64_t flagOffset = slot + subBlockIdx * DIRECT_UB_FLAG_STRIDE; + AscendC::CrossCoreSetFlag<0x4, PIPE_V>(DIRECT_UB_FREE_FLAG_BEGIN + flagOffset); + } + } + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec2Done[0]); + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec2Done[1]); + + EpilogueKdaFwdHVnew epilogueKdaFwdHVnew(resource); + EpilogueKdaFwdHUpdate epilogueKdaFwdHUpdate(resource); + uint32_t pongBaseEvent = 4; + + if (storeFinalState && std::is_same::value) { + AscendC::SetFlag(EVENT_ID0); // preset final_state + AscendC::SetFlag(EVENT_ID0 + pongBaseEvent); + AscendC::SetFlag(EVENT_ID2); // preset h + AscendC::SetFlag(EVENT_ID2 + pongBaseEvent); + } else { + AscendC::SetFlag(EVENT_ID0); // preset h_update + AscendC::SetFlag(EVENT_ID0 + pongBaseEvent); + AscendC::SetFlag(EVENT_ID2); // preset h + AscendC::SetFlag(EVENT_ID2 + pongBaseEvent); + } + AscendC::SetFlag(EVENT_ID1); // preset u + AscendC::SetFlag(EVENT_ID1 + pongBaseEvent); + AscendC::SetFlag(EVENT_ID3); // preset g + AscendC::SetFlag(EVENT_ID3 + pongBaseEvent); + AscendC::SetFlag(EVENT_ID0); // preset h_update + AscendC::SetFlag(EVENT_ID0 + pongBaseEvent); + AscendC::SetFlag(EVENT_ID2); // preset h + AscendC::SetFlag(EVENT_ID2 + pongBaseEvent); + uint32_t currStage = 0; // 0: V1, 1: V2 + bool event0FromMte3[PING_PONG_STAGES] = {false, false}; + bool event2FromMte3[PING_PONG_STAGES] = { + !(storeFinalState && std::is_same::value), + !(storeFinalState && std::is_same::value)}; + while (vecBlockScheduler.isRunning) { + if (currStage == 0) { + /* V1: + * gmV = gmU - gmVWorkspace + * g_buf = gmG[-1] - gmG + * g_buf = exp(g_buf) + * gmVWorkspace = g_buf * gmV + */ + vecBlockScheduler.InitTasks(); + for (uint32_t i = 0; i < PING_PONG_STAGES; ++i) { + uint32_t streamId = vecBlockScheduler.GetStreamId(i); + const auto &stream = vecBlockScheduler.GetStream(i); + if (vecBlockScheduler.StreamIsDone(stream)) { + continue; + } + const KdaFwdHOffsets &vec1Offsets = vecBlockScheduler.GetCurTaskOffsets(stream); + AscendC::LocalTensor l1VUpdate = (i == 0) ? l1VUpdatePing : l1VUpdatePong; + bool tailVectorPath = vec1Offsets.blockTokens < 16; + if (tailVectorPath) { + Arch::CrossCoreWaitFlag(vecBlockScheduler.cube1Done[streamId]); + ComputeTailVWorkspace(vec1Offsets, EVENT_ID3 + (i == 0 ? 0 : pongBaseEvent)); + } + bool waitWsFromMte3 = storeFinalState && std::is_same::value && + event0FromMte3[streamId]; + bool useDirectForTask = useDirectFp32Ub && !tailVectorPath; + epilogueKdaFwdHVnew( + gmV[vec1Offsets.uvOffset], gmVUpdateWorkspace[vec1Offsets.vWorkOffset], l1VUpdate, + gmG[vec1Offsets.gOffset], gmU[vec1Offsets.uvOffset], gmVWorkspace[vec1Offsets.vWorkOffset], + gmGk[vec1Offsets.gkOffset], gmK[vec1Offsets.wkOffset], + gmKDecayWorkspace[vec1Offsets.kDecayWorkOffset], vec1Offsets.blockTokens, kHeadDim, + vec1Offsets.vBlockDim, vHeadDim, vecBlockScheduler.cube1Done[streamId], + vecBlockScheduler.vec1Done[streamId], vec1Offsets.isInitialState, vec1Offsets.isFinalState, + storeFinalState, waitWsFromMte3, (i == 0), tailVectorPath, useDirectForTask, + DIRECT_UB_FREE_FLAG_BEGIN, DIRECT_UB_READY_FLAG_BEGIN, DIRECT_UB_FLAG_STRIDE); + if (storeFinalState && std::is_same::value) { + event0FromMte3[streamId] = false; + } + } + } else { + /* V2: h[i+1] += h_work if i < num_chunks - 1 else None */ + for (uint32_t i = 0; i < PING_PONG_STAGES; ++i) { + uint32_t streamId = vecBlockScheduler.GetStreamId(i); + const auto &stream = vecBlockScheduler.GetStream(i); + if (vecBlockScheduler.StreamIsDone(stream)) { + continue; + } + const KdaFwdHOffsets &vec2Offsets = vecBlockScheduler.GetCurTaskOffsets(stream); + if (vecBlockScheduler.NeedProcessStage2(stream)) { + bool tailVectorPath = vec2Offsets.blockTokens < 16; + bool fuseTailFinalState = false; + if constexpr (kGated && !scalarGated && useExp2 && + std::is_same::value) { + fuseTailFinalState = tailVectorPath && vec2Offsets.isFinalState && storeFinalState; + } + if (tailVectorPath) { + Arch::CrossCoreWaitFlag(vecBlockScheduler.cube2Done[streamId]); + if (fuseTailFinalState) { + ComputeKdaTailFinalState(vec2Offsets, EVENT_ID3 + (i == 0 ? 0 : pongBaseEvent)); + } else { + ComputeTailHWorkspace(vec2Offsets, EVENT_ID3 + (i == 0 ? 0 : pongBaseEvent)); + } + } + if (!fuseTailFinalState) { + if (storeFinalState && std::is_same::value) { + if (vec2Offsets.isInitialState && event2FromMte3[streamId]) { + uint32_t event2Id = EVENT_ID2 + (i == 0 ? 0 : pongBaseEvent); + AscendC::WaitFlag(event2Id); + AscendC::SetFlag(event2Id); + event2FromMte3[streamId] = false; + } + // Update always writes the FP32 state through MTE3, + // including intermediate chunks. + event0FromMte3[streamId] = true; + event2FromMte3[streamId] = !vec2Offsets.isFinalState; + } + epilogueKdaFwdHUpdate( + gmH[vec2Offsets.hDstOffset], gmFinalState[vec2Offsets.finalStateOffset], + gmG[vec2Offsets.gOffset], gmH[vec2Offsets.hSrcOffset], + gmHWorkspace[vec2Offsets.hWorkOffset], gmGk[vec2Offsets.gkOffset], + gmInitialState[vec2Offsets.initialStateOffset], vec2Offsets.blockTokens, kHeadDim, + vec2Offsets.vBlockDim, vHeadDim, vecBlockScheduler.cube2Done[streamId], + vec2Offsets.isInitialState, vec2Offsets.isFinalState, storeFinalState, + useInitialState, (i == 0), tailVectorPath, useDirectFp32Ub && !tailVectorPath, + DIRECT_UB_FREE_FLAG_BEGIN, DIRECT_UB_READY_FLAG_BEGIN, DIRECT_UB_FLAG_STRIDE); + } + } else { + if (!useDirectFp32Ub) { + Arch::CrossCoreWaitFlag(vecBlockScheduler.cube2Done[streamId]); + } + } + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(vecBlockScheduler.vec2Done[streamId]); + } + } + currStage ^= 0x01; + } + if (storeFinalState && std::is_same::value) { + if (event0FromMte3[0]) { + AscendC::WaitFlag(EVENT_ID0); + } else { + AscendC::WaitFlag(EVENT_ID0); + } + if (event0FromMte3[1]) { + AscendC::WaitFlag(EVENT_ID0 + pongBaseEvent); + } else { + AscendC::WaitFlag(EVENT_ID0 + pongBaseEvent); + } + if (event2FromMte3[0]) { + AscendC::WaitFlag(EVENT_ID2); + } else { + AscendC::WaitFlag(EVENT_ID2); + } + if (event2FromMte3[1]) { + AscendC::WaitFlag(EVENT_ID2 + pongBaseEvent); + } else { + AscendC::WaitFlag(EVENT_ID2 + pongBaseEvent); + } + } else { + AscendC::WaitFlag(EVENT_ID0); // preset h_update + AscendC::WaitFlag(EVENT_ID0 + pongBaseEvent); + AscendC::WaitFlag(EVENT_ID2); // preset h + AscendC::WaitFlag(EVENT_ID2 + pongBaseEvent); + } + AscendC::WaitFlag(EVENT_ID1); // preset u + AscendC::WaitFlag(EVENT_ID1 + pongBaseEvent); + AscendC::WaitFlag(EVENT_ID3); // preset g + AscendC::WaitFlag(EVENT_ID3 + pongBaseEvent); + AscendC::WaitFlag(EVENT_ID0); // drain h_update + AscendC::WaitFlag(EVENT_ID0 + pongBaseEvent); + AscendC::WaitFlag(EVENT_ID2); // drain h + AscendC::WaitFlag(EVENT_ID2 + pongBaseEvent); + } + } +}; + +} // namespace Catlass::Gemm::Kernel diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_arch.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_arch.hpp new file mode 100644 index 0000000000..36e70c0969 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_arch.hpp @@ -0,0 +1,52 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ARCH_HPP +#define CATLASS_KDA_ARCH_HPP + +#include "catlass/kda_catlass.hpp" + +namespace Catlass::Arch { + +struct AtlasA2 { + static constexpr uint32_t BIAS_SIZE = 1024; + static constexpr uint32_t FIXBUF_SIZE = 7 * 1024; + static constexpr uint32_t UB_SIZE = 192 * 1024; + static constexpr uint32_t L1_SIZE = 512 * 1024; + static constexpr uint32_t L0A_SIZE = 64 * 1024; + static constexpr uint32_t L0B_SIZE = 64 * 1024; + static constexpr uint32_t L0C_SIZE = 128 * 1024; +}; + +// Ascend950 +struct Ascend950 { + static constexpr uint32_t BIAS_SIZE = 4 * 1024; + static constexpr uint32_t FIXBUF_SIZE = 16 * 1024; + static constexpr uint32_t UB_SIZE = 248 * 1024; + static constexpr uint32_t L1_SIZE = 512 * 1024; + static constexpr uint32_t L0A_SIZE = 64 * 1024; + static constexpr uint32_t L0B_SIZE = 64 * 1024; + static constexpr uint32_t L0C_SIZE = 256 * 1024; +}; + +template +using PositionType = std::integral_constant; + +using PositionGM = PositionType; +using PositionL1 = PositionType; +using PositionL0A = PositionType; +using PositionL0B = PositionType; +using PositionL0C = PositionType; +using PositionBias = PositionType; +using PositionUB = PositionType; + +} // namespace Catlass::Arch + +#endif // CATLASS_KDA_ARCH_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_cross_core_sync.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_cross_core_sync.hpp new file mode 100644 index 0000000000..8b074cec46 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_cross_core_sync.hpp @@ -0,0 +1,126 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_ARCH_KDA_CROSS_CORE_SYNC_HPP +#define CATLASS_ARCH_KDA_CROSS_CORE_SYNC_HPP + +#include "catlass/kda_catlass.hpp" + +namespace Catlass::Arch { + +// A certain cross core flag can be continuously set up to 15 times without waiting. In scenarios where there is only +// one-way synchronization between cores and the synchronization operation is executed multiple times, if the core +// performing the set operation runs faster, it may result in the flag being set more than 15 times consecutively, +// leading to a system freeze. To prevent such a freeze, we need to wait for the waiting cores to perform a reverse +// synchronization operation after executing the set operation MAX_REVERSE_DEPTH times. +constexpr uint32_t MAX_REVERSE_DEPTH = 15; + +using FlagID = uint16_t; +constexpr FlagID AIV_INTER_BLOCK_BARRIER = 8; +constexpr FlagID AIC_INTER_BLOCK_BARRIER = 9; +constexpr FlagID AIV_INTER_SUBBLOCK_BARRIER = 10; +constexpr FlagID FFTS_MAX_FLAG = 7; + +struct CrossCoreFlag { + CATLASS_DEVICE + CrossCoreFlag() + : id(0) + {} + + CATLASS_DEVICE + CrossCoreFlag(FlagID id) + : id(id) + {} + + FlagID id; +}; + +template +struct CrossCoreFlagWithReverse { + CATLASS_DEVICE + CrossCoreFlagWithReverse() + : id(0), + reverseId(0) + {} + + CATLASS_DEVICE + CrossCoreFlagWithReverse(FlagID id, FlagID reverseId) + : id(id), + reverseId(reverseId) + {} + + FlagID id; + FlagID reverseId; + uint32_t count{0}; +}; + +template +struct BarrierFlag { + static_assert(MODE != 255 && MODE == 255, "Unsupported cross core barrier flag, can not find the specialization."); +}; + +template <> +struct BarrierFlag<0x0, AscendC::AIV> { + static constexpr FlagID ID = AIV_INTER_BLOCK_BARRIER; +}; + +template <> +struct BarrierFlag<0x0, AscendC::AIC> { + static constexpr FlagID ID = AIC_INTER_BLOCK_BARRIER; +}; + +template <> +struct BarrierFlag<0x1, AscendC::AIV> { + static constexpr FlagID ID = AIV_INTER_SUBBLOCK_BARRIER; +}; + +template +CATLASS_DEVICE void CrossCoreBarrier() +{ + constexpr FlagID flagId = BarrierFlag::ID; + AscendC::CrossCoreSetFlag(flagId); + AscendC::CrossCoreWaitFlag(flagId); +} + +template +CATLASS_DEVICE void CrossCoreSetFlag(CrossCoreFlag &flag) +{ + AscendC::CrossCoreSetFlag(flag.id); +} + +CATLASS_DEVICE +void CrossCoreWaitFlag(CrossCoreFlag &flag) +{ + AscendC::CrossCoreWaitFlag(flag.id); +} + +template +CATLASS_DEVICE void CrossCoreSetFlagWithReverse(CrossCoreFlagWithReverse &flag) +{ + AscendC::CrossCoreSetFlag(flag.id); + if (++flag.count >= REVERSE_DEPTH) { + AscendC::CrossCoreWaitFlag(flag.reverseId); + flag.count = 0; + } +} + +template +CATLASS_DEVICE void CrossCoreWaitFlagWithReverse(CrossCoreFlagWithReverse &flag) +{ + AscendC::CrossCoreWaitFlag(flag.id); + if (++flag.count >= REVERSE_DEPTH) { + AscendC::CrossCoreSetFlag(flag.reverseId); + flag.count = 0; + } +} + +} // namespace Catlass::Arch + +#endif // CATLASS_ARCH_KDA_CROSS_CORE_SYNC_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_local_tensor_buffer.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_local_tensor_buffer.hpp new file mode 100644 index 0000000000..8e02f4974f --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_local_tensor_buffer.hpp @@ -0,0 +1,230 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_LOCAL_TENSOR_BUFFER_HPP +#define CATLASS_KDA_LOCAL_TENSOR_BUFFER_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_arch.hpp" + +namespace Catlass::Arch { + +struct LocalTensorBufferBase { +public: + template + CATLASS_DEVICE + AscendC::LocalTensor + GetBufferByByte(const uint32_t offset) const + { + return tensor[offset].template ReinterpretCast(); + } + +protected: + CATLASS_DEVICE + LocalTensorBufferBase() = default; + + AscendC::LocalTensor tensor; +}; + +template < + class ArchTag, + AscendC::TPosition Position> +struct LocalTensorBuffer { + static_assert(DEPENDENT_FALSE, "Unsupported local tensor buffer, can not find the specialization."); +}; + +/// Partial specialization for TPosition::A1 +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::A1; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufA1; + GetTPipePtr()->InitBuffer(tbufA1, ArchTag::L1_SIZE); + tensor = tbufA1.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::A2 +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::A2; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufA2; + GetTPipePtr()->InitBuffer(tbufA2, ArchTag::L0A_SIZE); + tensor = tbufA2.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::B1 +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::B1; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufB1; + GetTPipePtr()->InitBuffer(tbufB1, ArchTag::L1_SIZE); + tensor = tbufB1.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::B2 +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::B2; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufB2; + GetTPipePtr()->InitBuffer(tbufB2, ArchTag::L0B_SIZE); + tensor = tbufB2.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::C1 +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::C1; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufC1; + GetTPipePtr()->InitBuffer(tbufC1, ArchTag::L1_SIZE); + tensor = tbufC1.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::C2 +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::C2; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufC2; + GetTPipePtr()->InitBuffer(tbufC2, ArchTag::BIAS_SIZE); + tensor = tbufC2.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::CO1 +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::CO1; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufCO1; + GetTPipePtr()->InitBuffer(tbufCO1, ArchTag::L0C_SIZE); + tensor = tbufCO1.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::C2PIPE2GM +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::C2PIPE2GM; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufC2PIPE2GM; + GetTPipePtr()->InitBuffer(tbufC2PIPE2GM, ArchTag::FIXBUF_SIZE); + tensor = tbufC2PIPE2GM.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::VECIN +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::VECIN; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufVECIN; + GetTPipePtr()->InitBuffer(tbufVECIN, ArchTag::UB_SIZE); + tensor = tbufVECIN.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::VECOUT +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::VECOUT; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufVECOUT; + GetTPipePtr()->InitBuffer(tbufVECOUT, ArchTag::UB_SIZE); + tensor = tbufVECOUT.Get(); + } +}; + +/////////////////////////////////////////////////////////// + +/// Partial specialization for TPosition::VECCALC +template +struct LocalTensorBuffer : LocalTensorBufferBase { +public: + static constexpr AscendC::TPosition Position = AscendC::TPosition::VECCALC; + + CATLASS_DEVICE + LocalTensorBuffer() + { + AscendC::TBuf tbufVECCALC; + GetTPipePtr()->InitBuffer(tbufVECCALC, ArchTag::UB_SIZE); + tensor = tbufVECCALC.Get(); + } +}; + +} // namespace Catlass::Arch + +#endif // CATLASS_KDA_LOCAL_TENSOR_BUFFER_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_resource.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_resource.hpp new file mode 100644 index 0000000000..9a521c24ce --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/arch/kda_resource.hpp @@ -0,0 +1,43 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_RESOURCE_HPP +#define CATLASS_KDA_RESOURCE_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_local_tensor_buffer.hpp" + +namespace Catlass::Arch { + +template +struct Resource { +public: + AscendC::TPipe pipe; + + LocalTensorBuffer l1Buf; + LocalTensorBuffer l0ABuf; + LocalTensorBuffer l0BBuf; + LocalTensorBuffer btBuf; + LocalTensorBuffer l0CBuf; + LocalTensorBuffer ubBuf; + LocalTensorBuffer fpBuf; + + CATLASS_DEVICE + Resource() + { + // The initialization of AscendC::Tpipe will insert some synchronization interfaces, + // which may conflict with the usage by users. Therefore, the "destroy" interface is used for releasing. + pipe.Destroy(); + } +}; + +} // namespace Catlass::Arch + +#endif // CATLASS_KDA_RESOURCE_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_alignment.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_alignment.hpp new file mode 100644 index 0000000000..f6821a9a40 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_alignment.hpp @@ -0,0 +1,103 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ALIGNMENT_HPP +#define CATLASS_KDA_ALIGNMENT_HPP + +#include "catlass/detail/kda_macros.hpp" +#include "tla/numeric/integral_constant.hpp" + +template +CATLASS_HOST_DEVICE constexpr T RoundUp(const T &val) +{ + static_assert(ALIGN != 0, "ALIGN must not be 0"); + return (val + ALIGN - 1) / ALIGN * ALIGN; +} + +template +CATLASS_HOST_DEVICE constexpr auto RoundUp(T const &val, U const &align) +{ + if constexpr (tla::is_static::value && tla::is_static::value) { // Int, Int + constexpr uint32_t res = (T::value + U::value - 1) / U::value * U::value; + return tla::Int{}; + } else if constexpr (tla::is_static::value) { // Int, int + return (T::value + align - 1) / align * align; + } else if constexpr (tla::is_static::value) { // int, Int + return (val + U::value - 1) / U::value * U::value; + } else { // int, int + return (val + align - 1) / align * align; + } +} + +template +CATLASS_HOST_DEVICE constexpr T RoundDown(const T val) +{ + static_assert(ALIGN != 0, "ALIGN must not be 0"); + return val / ALIGN * ALIGN; +} + +template +CATLASS_HOST_DEVICE constexpr auto RoundDown(T const &val, U const &align) +{ + if constexpr (tla::is_static::value && tla::is_static::value) { // Int, Int + constexpr uint32_t res = T::value / U::value * U::value; + return tla::Int{}; + } else if constexpr (tla::is_static::value) { // Int, int + return T::value / align * align; + } else if constexpr (tla::is_static::value) { // int, Int + return val / U::value * U::value; + } else { // int, int + return val / align * align; + } +} + +template +CATLASS_HOST_DEVICE constexpr T CeilDiv(const T dividend) +{ + static_assert(DIVISOR != 0, "DIVISOR must not be 0"); + return (dividend + DIVISOR - 1) / DIVISOR; +} + +template +CATLASS_HOST_DEVICE constexpr auto CeilDiv(T const ÷nd, U const &divisor) +{ + if constexpr (tla::is_static::value && tla::is_static::value) { // Int, Int + constexpr uint32_t res = (T::value + U::value - 1) / U::value; + return tla::Int{}; + } else if constexpr (tla::is_static::value) { // Int, int + return (T::value + divisor - 1) / divisor; + } else if constexpr (tla::is_static::value) { // int, Int + return (dividend + U::value - 1) / U::value; + } else { // int, int + return (dividend + divisor - 1) / divisor; + } +} + +template +CATLASS_HOST_DEVICE constexpr auto Max(T const &a, U const &b) +{ + if (a > b) { + return a; + } else { + return b; + } +} + +template +CATLASS_HOST_DEVICE constexpr auto Min(T const &a, U const &b) +{ + if (a < b) { + return a; + } else { + return b; + } +} + +#endif // CATLASS_KDA_ALIGNMENT_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_dependent_false.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_dependent_false.hpp new file mode 100644 index 0000000000..80613e908d --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_dependent_false.hpp @@ -0,0 +1,20 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_DEPENDENT_FALSE_HPP +#define CATLASS_KDA_DEPENDENT_FALSE_HPP + +template +constexpr bool DEPENDENT_BOOL_VALUE = VALUE; + +template +constexpr bool DEPENDENT_FALSE = DEPENDENT_BOOL_VALUE; + +#endif // CATLASS_KDA_DEPENDENT_FALSE_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_macros.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_macros.hpp new file mode 100644 index 0000000000..7d47093730 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_macros.hpp @@ -0,0 +1,26 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_MACROS_HPP +#define CATLASS_KDA_MACROS_HPP + +#if defined(__CCE__) +#include +#endif + +#define CATLASS_DEVICE __forceinline__ __aicore__ +#ifdef __CCE__ +#define CATLASS_HOST_DEVICE __forceinline__[host, aicore] +#else +#define CATLASS_HOST_DEVICE +#endif +#define CATLASS_GLOBAL __global__ __aicore__ + +#endif // CATLASS_KDA_MACROS_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_tag_to_layout.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_tag_to_layout.hpp new file mode 100644 index 0000000000..119cdba6fd --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/detail/kda_tag_to_layout.hpp @@ -0,0 +1,86 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_TAG_TO_LAYOUT_HPP +#define CATLASS_KDA_TAG_TO_LAYOUT_HPP + +#include "catlass/layout/kda_layout.hpp" +#include "tla/kda_tla_layout.hpp" + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +namespace Catlass::detail { +//////////////////////////////////////////////////////////////////////////////////////////////////// +// For each Catlass::layout, provides its corresponding tla layout types +template +struct TagToLayout { + using type = LayoutTag; +}; + +template +struct TagToLayout { + using type = tla::Layout, tla::Stride>>; +}; + +template +struct TagToLayout { + using type = tla::Layout, tla::Stride, int64_t>>; +}; + +template +struct TagToLayout { + using type = tla::Layout, tla::Stride>>; +}; + +template +struct TagToLayout { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + using type = tla::Layout< + tla::Shape, uint32_t>, tla::Shape, uint32_t>>, + tla::Stride, tla::Int>, + tla::Stride, int64_t>>>; +}; + +template +struct TagToLayout { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + using type = tla::Layout< + tla::Shape, uint32_t>, tla::Shape, uint32_t>>, + tla::Stride, int64_t>, + tla::Stride, tla::Int>>>; +}; + +template +struct TagToLayout { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element); + using type = tla::Layout< + tla::Shape, uint32_t>, tla::Shape, uint32_t>>, + tla::Stride, int64_t>, + tla::Stride, tla::Int>>>; +}; + +// Convenience aliases +template +using TagToLayout_t = typename TagToLayout::type; + +constexpr uint32_t ELE_NUM_PER_FRACTAL_L0C = 256; +using LayoutL0C = tla::Layout< + tla::Shape, uint32_t>, tla::Shape, uint32_t>>, + tla::Stride, tla::Int>, + tla::Stride, int64_t>>>; + +//////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::detail + +#endif // CATLASS_KDA_TAG_TO_LAYOUT_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue.hpp new file mode 100644 index 0000000000..cc072c2173 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/block/kda_block_epilogue.hpp @@ -0,0 +1,32 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_BLOCK_EPILOGUE_HPP +#define CATLASS_KDA_BLOCK_EPILOGUE_HPP + +#include "catlass/kda_catlass.hpp" + +namespace Catlass::Epilogue::Block { + +template < + class DispatchPolicy, + class... Args> +class BlockEpilogue { + static_assert(DEPENDENT_FALSE, "Could not find an epilogue specialization"); +}; + +} // namespace Catlass::Epilogue::Block + +// chunk_kda_fwd only instantiates BlockEpilogue for its own FwdH policies +// (EpilogueAtlasKdaFwdHVnew / EpilogueAtlasKdaFwdHUpdate), whose specializations +// live under arch22|arch35/fwd_h/epilogue/block/ and are included by their users. +// None of the upstream catlass epilogues were ever selected here. + +#endif // CATLASS_KDA_BLOCK_EPILOGUE_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/copy_ub_to_l1_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/copy_ub_to_l1_tla.hpp new file mode 100644 index 0000000000..90f2a76673 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/copy_ub_to_l1_tla.hpp @@ -0,0 +1,69 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_EPILOGUE_TILE_COPY_UB_TO_L1_TLA_HPP +#define CATLASS_EPILOGUE_TILE_COPY_UB_TO_L1_TLA_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "tla/kda_tensor.hpp" +#include "tla/kda_tla_layout.hpp" + +namespace Catlass::Epilogue::Tile { + +template < + class ArchTag, + class TensorSrc, + class TensorDst, + class Enable = void> +struct CopyUb2L1Tla { + static_assert(DEPENDENT_FALSE, "Unsupported CopyUb2L1Tla, can not find the specialization."); +}; + +/// Partial specialization for Ascend950, zN in and zN out. +template +struct CopyUb2L1Tla, LayoutSrc, CoordSrc, AscendC::TPosition::VECCALC>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + std::enable_if_t::value && + tla::detail::iszN::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc); + + // Methods + + CATLASS_DEVICE + CopyUb2L1Tla() = default; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::iszNUnAlign::value && + tla::detail::iszN::value && + TensorSrc::position == AscendC::TPosition::VECCALC && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be UB and zN, " + "while TensorDst must be L1 and zN"); + + int64_t srcShape = tla::get<0, 0>(srcTensor.shape()); + AscendC::DataCopyParams dataCopyParams( + tla::get<1, 1>(srcTensor.shape()), + srcShape, + (tla::get<1, 1>(srcTensor.stride()) / ELE_NUM_PER_C0 - srcShape), + (tla::get<1, 1>(dstTensor.stride()) / ELE_NUM_PER_C0 - srcShape)); + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], dataCopyParams); + } +}; + +} // namespace Catlass::Epilogue::Tile + +#endif // CATLASS_EPILOGUE_TILE_COPY_UB_TO_L1_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_gm_to_ub.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_gm_to_ub.hpp new file mode 100644 index 0000000000..683f7cefbb --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_gm_to_ub.hpp @@ -0,0 +1,181 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_COPY_GM_TO_UB_HPP +#define CATLASS_KDA_COPY_GM_TO_UB_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" + +namespace Catlass::Epilogue::Tile { + +template < + class ArchTag, + class GmType> +struct CopyGm2Ub { + static_assert(DEPENDENT_FALSE, "Unsupported copy gm to ub, can not find the specialization."); +}; + +template +struct CopyGm2Ub> { + using LayoutSrc = layout::RowMajor; + using LayoutDst = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(Element); + + CATLASS_DEVICE + CopyGm2Ub() = default; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + layout::RowMajor const &layoutDst, + layout::RowMajor const &layoutSrc) + { + AscendC::DataCopyExtParams dataCopyParams( + layoutSrc.shape(0), + layoutSrc.shape(1) * sizeof(Element), + (layoutSrc.stride(0) - layoutSrc.shape(1)) * sizeof(Element), + (layoutDst.stride(0) - layoutDst.shape(1)) / ELE_NUM_PER_BLK, + 0); + AscendC::DataCopyPadExtParams padParams(false, 0, 0, 0); + AscendC::DataCopyPad(dstTensor, srcTensor, dataCopyParams, padParams); + }; +}; + +template +struct CopyGm2Ub> { + using LayoutSrc = layout::VectorLayout; + using LayoutDst = layout::VectorLayout; + + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(Element); + + CATLASS_DEVICE + CopyGm2Ub() = default; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + layout::VectorLayout const &layoutDst, + layout::VectorLayout const &layoutSrc) + { + AscendC::DataCopyExtParams dataCopyParams( + 1, + layoutSrc.shape(0) * sizeof(Element), + 0, + 0, + 0); + AscendC::DataCopyPadExtParams padParams(false, 0, 0, 0); + AscendC::DataCopyPad(dstTensor, srcTensor, dataCopyParams, padParams); + }; +}; + +/// @brief This copy instruction used to copy per token scale from GM to UB. +/// Copy the scale of shape (m,1) on GM to the first column of shape (m,n) on UB, +/// and pad the first block of each row (i.e. pad to shape (m,8) when element type is float). +/// @tparam ArchTag: Architecture tag. +/// @tparam GmType: Type of data on GM. +template < + class ArchTag, + class GmType> +struct CopyPerTokenScale2Ub { + static_assert(std::is_same_v, + "Unsupported layout for CopyPerTokenScale2Ub."); + + using Element = typename GmType::Element; + using LayoutSrc = typename GmType::Layout; + using LayoutDst = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(Element); + + CATLASS_DEVICE + CopyPerTokenScale2Ub() = default; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, + LayoutSrc const &layoutSrc) + { + AscendC::DataCopyExtParams dataCopyParams; + AscendC::DataCopyPadExtParams padParams; + + dataCopyParams.blockCount = layoutSrc.shape(0); + dataCopyParams.blockLen = layoutSrc.shape(1) * sizeof(Element); // per token scale has only one column + dataCopyParams.srcStride = 0; + dataCopyParams.dstStride = (layoutDst.stride(0) - layoutDst.shape(1)) / ELE_NUM_PER_BLK; + // Pad the data to the complete block + padParams.isPad = true; + padParams.leftPadding = 0; + padParams.rightPadding = 0; + + AscendC::DataCopyPad(dstTensor, srcTensor, dataCopyParams, padParams); + } +}; + +template < + class ArchTag, + class GmType> +struct CopyGm2UbAligned { + static_assert(DEPENDENT_FALSE, "Unsupported copy gm to ub aligned, can not find the specialization."); +}; + +template +struct CopyGm2UbAligned> { + using LayoutSrc = layout::RowMajor; + using LayoutDst = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(Element); + static constexpr uint32_t BLOCK_LEN_LIMIT = 65536; + static constexpr uint32_t MAX_REPEAT = 4095; + static constexpr uint32_t STRIDE_LIMIT = 65536; + + CATLASS_DEVICE + CopyGm2UbAligned() = default; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + layout::RowMajor const &layoutDst, + layout::RowMajor const &layoutSrc) + { + uint32_t rows = layoutSrc.shape(0); + uint32_t cols = layoutSrc.shape(1); + uint32_t srcStride = (layoutSrc.stride(0) - layoutSrc.shape(1)) / ELE_NUM_PER_BLK; + uint32_t dstStride = (layoutDst.stride(0) - layoutDst.shape(1)) / ELE_NUM_PER_BLK; + + if ((layoutSrc.shape(1) == layoutSrc.stride(0)) && (layoutDst.shape(1) == layoutDst.stride(0))) { + DataCopy(dstTensor, srcTensor, rows * cols); + } else if (srcStride < STRIDE_LIMIT && dstStride < STRIDE_LIMIT && (cols / ELE_NUM_PER_BLK) < BLOCK_LEN_LIMIT) { + uint32_t rLoops = CeilDiv(rows, MAX_REPEAT); + for (uint32_t i = 0; i < rLoops; ++i) { + uint32_t rActual = (i < rLoops - 1) ? MAX_REPEAT : rows - i * MAX_REPEAT; + AscendC::DataCopyParams dataCopyParams( + rActual, cols / ELE_NUM_PER_BLK, srcStride, dstStride); + DataCopy(dstTensor[i * MAX_REPEAT * layoutDst.stride(0)], + srcTensor[i * MAX_REPEAT * layoutSrc.stride(0)], dataCopyParams); + } + } else { + for (uint32_t i = 0; i < rows; ++i) { + DataCopy(dstTensor[i * layoutDst.stride(0)], srcTensor[i * layoutSrc.stride(0)], cols); + } + } + }; +}; + +} // namespace Catlass::Epilogue::Tile + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_gm_to_ub_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_gm_to_ub_tla.hpp new file mode 100644 index 0000000000..380a7464f9 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_gm_to_ub_tla.hpp @@ -0,0 +1,143 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_COPY_GM_TO_UB_TLA_HPP +#define CATLASS_KDA_COPY_GM_TO_UB_TLA_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "tla/kda_tensor.hpp" +#include "tla/kda_tla_layout.hpp" + +namespace Catlass::Epilogue::Tile { + +template < + class ArchTag, + class TensorSrc, + class TensorDst, + class Enable = void> +struct CopyGm2UbTla { + static_assert(DEPENDENT_FALSE, "Unsupported CopyGm2UbTla, can not find the specialization."); +}; + +/// Partial specialization for AtlasA2, RowMajor in and RowMajor out. +template +struct CopyGm2UbTla, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::VECCALC>, + std::enable_if_t::value && + tla::detail::isRowMajor::value>> { + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(ElementSrc); + + // Methods + + CATLASS_DEVICE + CopyGm2UbTla() = default; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isRowMajor::value && + tla::detail::isRowMajor::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::VECCALC, + "The input parameters do not match. TensorSrc must be GM and RowMajor, " + "while TensorDst must be UB and RowMajor"); + + AscendC::DataCopyExtParams dataCopyParams( + tla::get<0>(srcTensor.shape()), + tla::get<1>(srcTensor.shape()) * sizeof(ElementSrc), + (tla::get<0>(srcTensor.stride()) - tla::get<1>(srcTensor.shape())) * sizeof(ElementSrc), + (tla::get<0>(dstTensor.stride()) - tla::get<1>(dstTensor.shape())) / ELE_NUM_PER_BLK, + 0); + AscendC::DataCopyPadExtParams padParams(false, 0, 0, 0); + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + AscendC::DataCopyPad(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], dataCopyParams, padParams); + }; +}; + +/// Partial specialization for Ascend950, Vector in and Vector out. +template +struct CopyGm2UbTla, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::VECCALC>, + std::enable_if_t::value && + tla::detail::isVector::value>> { + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(ElementSrc); + + // Methods + + CATLASS_DEVICE + CopyGm2UbTla() = default; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isVector::value && + tla::detail::isVector::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::VECCALC, + "The input parameters do not match. TensorSrc must be GM and Vector, " + "while TensorDst must be UB and Vector"); + + AscendC::DataCopyExtParams dataCopyParams( + 1, + tla::get<0>(srcTensor.shape()) * sizeof(ElementSrc), + 0, + 0, + 0); + AscendC::DataCopyPadExtParams padParams(false, 0, 0, 0); + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + AscendC::DataCopyPad(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], dataCopyParams, padParams); + }; +}; + +/// Partial specialization for Ascend950, RowMajor in and RowMajor out. +template +struct CopyGm2UbTla, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::VECCALC>, + std::enable_if_t::value && + tla::detail::isRowMajor::value>> { + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(ElementSrc); + + // Methods + + CATLASS_DEVICE + CopyGm2UbTla() = default; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isRowMajor::value && + tla::detail::isRowMajor::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::VECCALC, + "The input parameters do not match. TensorSrc must be GM and RowMajor, " + "while TensorDst must be UB and RowMajor"); + + AscendC::DataCopyExtParams dataCopyParams( + tla::get<0>(srcTensor.shape()), + tla::get<1>(srcTensor.shape()) * sizeof(ElementSrc), + (tla::get<0>(srcTensor.stride()) - tla::get<1>(srcTensor.shape())) * sizeof(ElementSrc), + (tla::get<0>(dstTensor.stride()) - tla::get<1>(dstTensor.shape())) / ELE_NUM_PER_BLK, + 0); + AscendC::DataCopyPadExtParams padParams(false, 0, 0, 0); + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + AscendC::DataCopyPad(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], dataCopyParams, padParams); + }; +}; + +} // namespace Catlass::Epilogue::Tile + +#endif // CATLASS_KDA_COPY_GM_TO_UB_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_ub_to_gm.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_ub_to_gm.hpp new file mode 100644 index 0000000000..8fc513fcf9 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_ub_to_gm.hpp @@ -0,0 +1,136 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_COPY_UB_TO_GM_HPP +#define CATLASS_KDA_COPY_UB_TO_GM_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" + +namespace Catlass::Epilogue::Tile { + +template < + class ArchTag, + class GmType> +struct CopyUb2Gm { + static_assert(DEPENDENT_FALSE, "Unsupported copy ub to gm, can not find the specialization."); +}; + +template +struct CopyUb2Gm> { + using LayoutDst = layout::RowMajor; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + CATLASS_DEVICE + CopyUb2Gm() = default; + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + layout::RowMajor const &layoutDst, + layout::RowMajor const &layoutSrc) + { + AscendC::DataCopyExtParams dataCopyParams( + layoutDst.shape(0), + layoutDst.shape(1) * sizeof(Element), + (layoutSrc.stride(0) - layoutSrc.shape(1)) / ELE_NUM_PER_C0, + (layoutDst.stride(0) - layoutDst.shape(1)) * sizeof(Element), + 0); + AscendC::DataCopyPad(dstTensor, srcTensor, dataCopyParams); + } +}; + +// new add vectorlayout version +template +struct CopyUb2Gm> { + using LayoutSrc = layout::VectorLayout; + using LayoutDst = layout::VectorLayout; + + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(Element); + + CATLASS_DEVICE + CopyUb2Gm() = default; + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + layout::VectorLayout const &layoutDst, + layout::VectorLayout const &layoutSrc) + { + AscendC::DataCopyExtParams dataCopyParams( + 1, + layoutDst.shape(0) * sizeof(Element), + 0, + 0, + 0); + AscendC::DataCopyPad(dstTensor, srcTensor, dataCopyParams); + }; +}; + +template < + class ArchTag, + class GmType> +struct CopyUb2GmAligned { + static_assert(DEPENDENT_FALSE, "Unsupported copy ub to gm aligned, can not find the specialization."); +}; + +template +struct CopyUb2GmAligned> { + using LayoutSrc = layout::RowMajor; + using LayoutDst = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(Element); + static constexpr uint32_t BLOCK_LEN_LIMIT = 65536; + static constexpr uint32_t MAX_REPEAT = 4095; + static constexpr uint32_t STRIDE_LIMIT = 65536; + + CATLASS_DEVICE + CopyUb2GmAligned() = default; + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + layout::RowMajor const &layoutDst, + layout::RowMajor const &layoutSrc) + { + uint32_t rows = layoutDst.shape(0); + uint32_t cols = layoutDst.shape(1); + uint32_t srcStride = (layoutSrc.stride(0) - layoutSrc.shape(1)) / ELE_NUM_PER_BLK; + uint32_t dstStride = (layoutDst.stride(0) - layoutDst.shape(1)) / ELE_NUM_PER_BLK; + + if ((layoutSrc.shape(1) == layoutSrc.stride(0)) && (layoutDst.shape(1) == layoutDst.stride(0))) { + DataCopy(dstTensor, srcTensor, rows * cols); + } else if (srcStride < STRIDE_LIMIT && dstStride < STRIDE_LIMIT && (cols / ELE_NUM_PER_BLK) < BLOCK_LEN_LIMIT) { + uint32_t rLoops = CeilDiv(rows, MAX_REPEAT); + for (uint32_t i = 0; i < rLoops; ++i) { + uint32_t rActual = (i < rLoops - 1) ? MAX_REPEAT : rows - i * MAX_REPEAT; + AscendC::DataCopyParams dataCopyParams( + rActual, cols / ELE_NUM_PER_BLK, srcStride, dstStride); + DataCopy(dstTensor[i * MAX_REPEAT * layoutDst.stride(0)], + srcTensor[i * MAX_REPEAT * layoutSrc.stride(0)], dataCopyParams); + } + } else { + for (uint32_t i = 0; i < rows; ++i) { + DataCopy(dstTensor[i * layoutDst.stride(0)], srcTensor[i * layoutSrc.stride(0)], cols); + } + } + }; +}; + +} // namespace Catlass::Epilogue::Tile + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_ub_to_gm_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_ub_to_gm_tla.hpp new file mode 100644 index 0000000000..504b79a3d9 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_copy_ub_to_gm_tla.hpp @@ -0,0 +1,104 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_COPY_UB_TO_GM_TLA_HPP +#define CATLASS_KDA_COPY_UB_TO_GM_TLA_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "tla/kda_tensor.hpp" +#include "tla/kda_tla_layout.hpp" + +namespace Catlass::Epilogue::Tile { + +template < + class ArchTag, + class TensorSrc, + class TensorDst, + class Enable = void> +struct CopyUb2GmTla { + static_assert(DEPENDENT_FALSE, "Unsupported CopyUb2GmTla, can not find the specialization."); +}; + +/// Partial specialization for AtlasA2, RowMajor in and RowMajor out. +template +struct CopyUb2GmTla, LayoutSrc, CoordSrc, AscendC::TPosition::VECCALC>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::GM>, + std::enable_if_t::value && + tla::detail::isRowMajor::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc); + + // Methods + + CATLASS_DEVICE + CopyUb2GmTla() = default; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isRowMajor::value && + tla::detail::isRowMajor::value && + TensorSrc::position == AscendC::TPosition::VECCALC && + TensorDst::position == AscendC::TPosition::GM, + "The input parameters do not match. TensorSrc must be UB and RowMajor, " + "while TensorDst must be GM and RowMajor"); + + AscendC::DataCopyExtParams dataCopyParams( + tla::get<0>(dstTensor.shape()), + tla::get<1>(dstTensor.shape()) * sizeof(ElementSrc), + (tla::get<0>(srcTensor.stride()) - tla::get<1>(srcTensor.shape())) / ELE_NUM_PER_C0, + (tla::get<0>(dstTensor.stride()) - tla::get<1>(dstTensor.shape())) * sizeof(ElementSrc), + 0); + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + AscendC::DataCopyPad(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], dataCopyParams); + }; +}; + +/// Partial specialization for Ascend950, RowMajor in and RowMajor out. +template +struct CopyUb2GmTla, LayoutSrc, CoordSrc, AscendC::TPosition::VECCALC>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::GM>, + std::enable_if_t::value && + tla::detail::isRowMajor::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc); + + // Methods + + CATLASS_DEVICE + CopyUb2GmTla() = default; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isRowMajor::value && + tla::detail::isRowMajor::value && + TensorSrc::position == AscendC::TPosition::VECCALC && + TensorDst::position == AscendC::TPosition::GM, + "The input parameters do not match. TensorSrc must be UB and RowMajor, " + "while TensorDst must be GM and RowMajor"); + + AscendC::DataCopyExtParams dataCopyParams( + tla::get<0>(dstTensor.shape()), + tla::get<1>(dstTensor.shape()) * sizeof(ElementSrc), + (tla::get<0>(srcTensor.stride()) - tla::get<1>(srcTensor.shape())) / ELE_NUM_PER_C0, + (tla::get<0>(dstTensor.stride()) - tla::get<1>(dstTensor.shape())) * sizeof(ElementSrc), + 0); + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + AscendC::DataCopyPad(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], dataCopyParams); + }; +}; + +} // namespace Catlass::Epilogue::Tile + +#endif // CATLASS_KDA_COPY_UB_TO_GM_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_epilogue_tile_copy.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_epilogue_tile_copy.hpp new file mode 100644 index 0000000000..74cd66dc2f --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/epilogue/tile/kda_epilogue_tile_copy.hpp @@ -0,0 +1,192 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_EPILOGUE_TILE_COPY_HPP +#define CATLASS_KDA_EPILOGUE_TILE_COPY_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "catlass/detail/kda_tag_to_layout.hpp" +#include "catlass/epilogue/tile/kda_copy_gm_to_ub.hpp" +#include "catlass/epilogue/tile/kda_copy_ub_to_gm.hpp" +#include "catlass/epilogue/tile/kda_copy_gm_to_ub_tla.hpp" +#include "catlass/epilogue/tile/kda_copy_ub_to_gm_tla.hpp" +#include "tla/kda_tensor.hpp" + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) +#include "catlass/epilogue/tile/copy_ub_to_l1_tla.hpp" +#endif + +namespace Catlass::Epilogue::Tile { + +template < + /// Tag indicating architecture + class ArchTag, + class... Args> +struct TileCopy { + static_assert(DEPENDENT_FALSE, "Unsupported tile copy, can not find the specialization."); +}; + +template < + class ArchTag, + /// GemmType for C matrix operand + class CType, + /// GemmType for D matrix operand + class DType> +struct TileCopy { + using ElementC = typename CType::Element; + using ElementD = typename DType::Element; + + using CopyGmToUbC = CopyGm2Ub; + using CopyUbToGmD = CopyUb2Gm; +}; + +template < + class ArchTag, + /// GemmType for C matrix operand + class CType, + /// GemmType for X matrix operand + class XType, + /// GemmType for D matrix operand + class DType> +struct TileCopy { + using ElementC = typename CType::Element; + using ElementX = typename XType::Element; + using ElementD = typename DType::Element; + + using CopyGmToUbC = CopyGm2Ub; + using CopyGmToUbX = CopyGm2Ub; + using CopyUbToGmD = CopyUb2Gm; +}; + +template < + class ArchTag, + class CType, + class XType, + class YType, + class DType> +struct TileCopy { + using ElementC = typename CType::Element; + using ElementX = typename XType::Element; + using ElementY = typename YType::Element; + using ElementD = typename DType::Element; + + using CopyGmToUbC = CopyGm2Ub; + using CopyGmToUbX = CopyGm2Ub; + using CopyGmToUbY = CopyGm2Ub; + using CopyUbToGmD = CopyUb2Gm; +}; + +template < + class ArchTag, + class CType, + class XType, + class YType, + class DType> +struct TileCopyBf16 { + using ElementC = typename CType::Element; + using ElementX = bfloat16_t; + using ElementY = bfloat16_t; + using ElementD = bfloat16_t; + + using CopyGmToUbC = CopyGm2Ub; + using CopyGmToUbX = CopyGm2Ub>; + using CopyGmToUbY = CopyGm2Ub>; + using CopyUbToGmD = CopyUb2Gm>; +}; + +template < + class ArchTag, + class CType, + class ScaleType, + class PerTokenScaleType, + class DType> +struct TileCopyPerTokenDequant { + using ElementC = typename CType::Element; + using ElementScale = typename ScaleType::Element; + using ElementPerTokenScale = typename PerTokenScaleType::Element; + using ElementD = typename DType::Element; + + using CopyGmToUbC = CopyGm2Ub; + using CopyGmToUbScale = CopyGm2Ub; + using CopyGmToUbPerTokenScale = CopyPerTokenScale2Ub; + using CopyUbToGmD = CopyUb2Gm; +}; + +template < + class ArchTag, + class CType, + class PerTokenScaleType, + class DType> +struct TileCopyW4A4Gemm { + using ElementC = typename CType::Element; + using ElementPerTokenScale = typename PerTokenScaleType::Element; + using ElementD = typename DType::Element; + + using CopyGmToUbC = CopyGm2Ub; + using CopyGmToUbPerTokenScale = CopyGm2Ub; + using CopyUbToGmD = CopyUb2Gm; +}; + +template < + class ArchTag, + /// GemmType for C matrix operand + class ElementC_, + class LayoutTagC_, + /// GemmType for X matrix operand + class ElementX_, + class LayoutTagX_, + /// GemmType for Y matrix operand + class ElementY_, + class LayoutTagY_, + /// GemmType for D matrix operand + class ElementD_, + class LayoutTagD_> +struct TileCopyDequantTla { + using ElementC = ElementC_; + using LayoutTagC = LayoutTagC_; + using LayoutC = detail::TagToLayout_t; + using TensorUbC = + tla::Tensor, LayoutC, tla::Coord, AscendC::TPosition::VECCALC>; + + using ElementX = ElementX_; + using LayoutTagX = LayoutTagX_; + using LayoutX = detail::TagToLayout_t; + using TensorUbX = + tla::Tensor, LayoutX, tla::Coord, AscendC::TPosition::VECCALC>; + + using ElementY = ElementY_; + using LayoutTagY = LayoutTagY_; + using LayoutY = detail::TagToLayout_t; + using TensorUbY = + tla::Tensor, LayoutY, tla::Coord, AscendC::TPosition::VECCALC>; + + using ElementD = ElementD_; + using LayoutTagD = LayoutTagD_; + using LayoutD = detail::TagToLayout_t; + using TensorUbD = + tla::Tensor, LayoutD, tla::Coord, AscendC::TPosition::VECCALC>; + + template + using CopyGmToUbC = CopyGm2UbTla; + + template + using CopyGmToUbX = CopyGm2UbTla; + + template + using CopyGmToUbY = CopyGm2UbTla; + + template + using CopyUbToGmD = CopyUb2GmTla; +}; + +} // namespace Catlass::Epilogue::Tile + +#endif // CATLASS_KDA_EPILOGUE_TILE_COPY_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_tla.hpp new file mode 100644 index 0000000000..f8049a3da5 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/block_mmad_pingpong_tla.hpp @@ -0,0 +1,630 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Block { + +template < + class ArchTag_, + bool ENABLE_UNIT_FLAG_, + bool USE_HF32_MODE_, + uint32_t L0C_STAGES_, + bool ENABLE_L1_RESIDENT_, + uint32_t L1A_STAGES_, + uint32_t L1B_STAGES_, + uint32_t L0A_STAGES_, + uint32_t L0B_STAGES_, + class L1TileShape_, + class L0TileShape_, + class ElementA_, + class ElementB_, + class ElementC_, + class ElementBias_, + class TileCopy_, + class TileMmad_> +struct BlockMmadTla< + MmadPingpong, + L1TileShape_, + L0TileShape_, + ElementA_, + ElementB_, + ElementC_, + ElementBias_, + TileCopy_, + TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadPingpong; + using ArchTag = typename DispatchPolicy::ArchTag; + using TileCopy = TileCopy_; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = ElementA_; + using LayoutA = typename TileCopy::LayoutA; + using ElementB = ElementB_; + using LayoutB = typename TileCopy::LayoutB; + using ElementC = ElementC_; + using LayoutC = typename TileCopy::LayoutC; + using ElementBias = ElementBias_; + + using TileMmad = TileMmad_; + + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; + using CopyL1ToBT = typename TileCopy::CopyL1ToBT; + + using ElementAccumulator = typename TileCopy::ElementAccumulator; + + static constexpr bool HAS_BIAS = TileCopy::HAS_BIAS; + + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + + using L1AAlignHelper = typename TileCopy_::L1AAlignHelper; + using L1BAlignHelper = typename TileCopy_::L1BAlignHelper; + + static_assert(tla::is_tuple::value && tla::is_static::value, + "L1TileShape must be tla::tuple and static!"); + static_assert(tla::is_tuple::value && tla::is_static::value, + "L0TileShape must be tla::tuple and static!"); + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr bool USE_HF32_MODE = DispatchPolicy::USE_HF32_MODE; + static constexpr bool ENABLE_L1_RESIDENT = DispatchPolicy::ENABLE_L1_RESIDENT; + static constexpr uint32_t L1A_STAGES = DispatchPolicy::L1A_STAGES; + static constexpr uint32_t L1B_STAGES = DispatchPolicy::L1B_STAGES; + static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES; + static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES; + static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; + static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); + static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); + static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); + static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); + static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); + static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); + + // L1 tile size + static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA); + static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); + // L0 tile size + static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L1_TILE_M * L1_TILE_N * sizeof(ElementAccumulator); + + // Check HF32_MODE + static_assert( + !USE_HF32_MODE || (USE_HF32_MODE && std::is_same_v && std::is_same_v), + "HF32 MODE only supports in float!"); + + // Check L0C_STAGES + static_assert(!(ENABLE_UNIT_FLAG && L0C_STAGES != 1), "L0C_STAGES must be 1 when UnitFlag is true!"); + + // Check LayoutC + static_assert(tla::detail::isRowMajor::value || + ((std::is_same_v || std::is_same_v || + std::is_same_v) && + tla::detail::iszN::value), + "LayoutC only supports zN in half or bfloat16 or float, RowMajor in all dtype yet!"); + + // Check L1TileShape + static_assert(L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static_assert(L0A_TILE_SIZE * L0A_STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert(L0B_TILE_SIZE * L0B_STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE * L0C_STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); + + static constexpr uint32_t _32B = 32 * 8; // in bits + static_assert(L1_TILE_M == L0_TILE_M && L1_TILE_N == L0_TILE_N, + "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); + static_assert(L0_TILE_K <= L1_TILE_K, "L0TileShape::K cannot exceed L1TileShape::K"); +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) + static_assert(L1_TILE_M * SizeOfBits::value % _32B == 0, "L1TileShape::M must be 32B aligned."); + static_assert(L1_TILE_K * SizeOfBits::value % _32B == 0, "L1TileShape::K must be 32B aligned."); + static_assert(L1_TILE_K * SizeOfBits::value % _32B == 0, "L1TileShape::K must be 32B aligned."); + static_assert(L1_TILE_N * SizeOfBits::value % _32B == 0, "L1TileShape::N must be 32B aligned."); + static_assert(L0_TILE_K * SizeOfBits::value % _32B == 0, "L0TileShape::K must be 32B aligned."); +#endif + + static_assert((!HAS_BIAS && (L1A_STAGES + L1B_STAGES) <= 8) || (HAS_BIAS && (L1A_STAGES + L1B_STAGES) <= 7), + "L1 Buffer overflow: Exceeds the supported range of EVENT(0~7)"); + + static_assert((!HAS_BIAS && (L0A_STAGES + L0B_STAGES) <= 8) || (HAS_BIAS && (L0A_STAGES + L0B_STAGES) <= 7), + "L0 Buffer overflow: Exceeds the supported range of EVENT_ID(0~7)"); + + static constexpr auto L1A_LAYOUT = + tla::MakeLayout(tla::Int{}, tla::Int{}); + static constexpr auto L1B_LAYOUT = + tla::MakeLayout(tla::Int{}, tla::Int{}); + static constexpr auto L1BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); + static constexpr auto L0BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); + + // When enabling L1 resident mode, restore the pointer and coordinates that record the last state + // to the initial state. if two blockmmad instances need to be consecutively invoked at the kernel layer, + // RestoreStatus() must be inserted between them. + CATLASS_DEVICE + void RestoreStatus() + { + for (int i = 0; i < L1A_STAGES; ++i) { + lastAddrA[i] = nullptr; + lastCoordA[i] = MatrixCoord{0U, 0U}; + } + for (int i = 0; i < L1B_STAGES; ++i) { + lastAddrB[i] = nullptr; + lastCoordB[i] = MatrixCoord{0U, 0U}; + } + } + + /// Construct + CATLASS_DEVICE + BlockMmadTla(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + if ASCEND_IS_AIC { + // use HF32 when USE_HF32_MODE is true + if constexpr (USE_HF32_MODE) { + AscendC::SetHF32Mode(true); + } else { + AscendC::SetHF32Mode(false); + } + if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { + AscendC::SetMMLayoutTransform(true); + } + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * L1A_STAGES; + // Init buffers + for (uint32_t i = 0; i < L1A_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_TILE_SIZE * i); + // Assign event ID for each stages + l1AEventList[i] = i; + // The event id that needs to be set before the loop + AscendC::SetFlag(l1AEventList[i]); + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); + // Assign event ID for each stages + l1BEventList[i] = i + L1A_STAGES; + // The event id that needs to be set before the loop + AscendC::SetFlag(l1BEventList[i]); + } + for (uint32_t i = 0; i < L0A_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); + // Assign event ID for each stages + l0AEventList[i] = i; + // The event id that needs to be set before the loop + AscendC::SetFlag(l0AEventList[i]); + } + for (uint32_t i = 0; i < L0B_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); + // Assign event ID for each stages + l0BEventList[i] = i + L0A_STAGES; + // The event id that needs to be set before the loop + AscendC::SetFlag(l0BEventList[i]); + } + if constexpr (!ENABLE_UNIT_FLAG) { + for (uint32_t i = 0; i < L0C_STAGES; i++) { + l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte(L0C_TILE_SIZE * i); + l0CEventList[i] = i; + AscendC::SetFlag(l0CEventList[i]); + } + } else { + l0CTensorList[0] = resource.l0CBuf.template GetBufferByByte(0); + } + if constexpr (HAS_BIAS) { + uint32_t l1BiasOffset = l1BOffset + L1B_TILE_SIZE * L1B_STAGES; + l1BiasTensor = resource.l1Buf.template GetBufferByByte(l1BiasOffset); + l0BiasTensor = resource.btBuf.template GetBufferByByte(0); + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + AscendC::SetFlag(L0A_STAGES + L0B_STAGES); + } + + if constexpr (ENABLE_L1_RESIDENT) { + RestoreStatus(); + } + } + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmadTla() + { + if ASCEND_IS_AIC { + if constexpr (USE_HF32_MODE) { + AscendC::SetHF32Mode(false); + } + if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { + AscendC::SetMMLayoutTransform(false); + } + for (uint32_t i = 0; i < L1A_STAGES; i++) { + AscendC::WaitFlag(l1AEventList[i]); + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + AscendC::WaitFlag(l1BEventList[i]); + } + for (uint32_t i = 0; i < L0A_STAGES; i++) { + AscendC::WaitFlag(l0AEventList[i]); + } + for (uint32_t i = 0; i < L0B_STAGES; i++) { + AscendC::WaitFlag(l0BEventList[i]); + } + if constexpr (!ENABLE_UNIT_FLAG) { + for (uint32_t i = 0; i < L0C_STAGES; i++) { + AscendC::WaitFlag(l0CEventList[i]); + } + } + if constexpr (HAS_BIAS) { + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + AscendC::WaitFlag(L0A_STAGES + L0B_STAGES); + } + } + } + + /// Perform a block-scoped matrix multiply-accumulate + template + CATLASS_DEVICE void operator()(TensorA &tensorA, TensorB &tensorB, TensorC &tensorC, GemmCoord const &actualShape, + TensorBias const &tensorBias = {}) + { + // Check L1TileShape + if constexpr (HAS_BIAS) { + static constexpr uint32_t BIAS_BUF_SIZE = L0_TILE_N * sizeof(ElementAccumulator); + static constexpr uint32_t L1BIAS_SIZE = L1_TILE_N * sizeof(ElementBias); + static_assert(BIAS_BUF_SIZE <= ArchTag::BIAS_SIZE, + "BIAS_BUF_SIZE exceeding the BT space! Reduce L0_TILE_N"); + static_assert(L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES + L1BIAS_SIZE <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + } + + using CopyGmToL1A = typename TileCopy_::template CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::template CopyGmToL1B; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) + using CopyL0CToGm = typename TileCopy_::template CopyL0CToGm; + CopyL0CToGm copyL0CToDst; +#endif +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) + using CopyL0CToDst = typename TileCopy_::template CopyL0CToDst; + CopyL0CToDst copyL0CToDst; +#endif + + uint32_t mBlockActual = actualShape.m(); + uint32_t kBlockActual = actualShape.k(); + uint32_t nBlockActual = actualShape.n(); + + uint32_t mL1Actual = mBlockActual; + if constexpr (std::is_same_v) { + // Avoid using the gemv mode in mmad + if (mL1Actual == 1) { + mL1Actual = 16; + } + } + uint32_t nL1Actual = nBlockActual; + + auto layoutInL0C = tla::MakeLayoutL0C(mL1Actual, nL1Actual); + auto tensorL0C = tla::MakeTensor(l0CTensorList[l0CListId], layoutInL0C, Arch::PositionL0C{}); + auto tensorL0Bias = tla::MakeTensor(l0BiasTensor, L0BIAS_LAYOUT, Arch::PositionBias{}); + + uint32_t kL1Actual = min(kBlockActual, L1_TILE_K); + // load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1AListId]); + auto tensorL1A = tla::MakeTensor(l1ATensorList[l1AListId], L1A_LAYOUT, Arch::PositionL1{}); + auto tensorTileA = GetTileA(tensorA, 0, 0, mBlockActual, kL1Actual); + if constexpr (ENABLE_L1_RESIDENT) { + // If the currently loaded GM pointer and block coordinates are the same as the last loaded ones, + // skip this loading. + if (lastAddrA[l1AListId] != tensorTileA.data().GetPhyAddr() || tla::get<0>(tensorTileA.coord()) != lastCoordA[l1AListId].row() || tla::get<1>(tensorTileA.coord()) != lastCoordA[l1AListId].column()) { + copyGmToL1A(tensorL1A, tensorTileA); + lastCoordA[l1AListId] = MatrixCoord{tla::get<0>(tensorTileA.coord()), tla::get<1>(tensorTileA.coord())}; + lastAddrA[l1AListId] = const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileA.data().GetPhyAddr()); + } + } else { + copyGmToL1A(tensorL1A, tensorTileA); + } + AscendC::SetFlag(l1AEventList[l1AListId]); + + // load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListId]); + auto tensorL1B = tla::MakeTensor(l1BTensorList[l1BListId], L1B_LAYOUT, Arch::PositionL1{}); + auto tensorTileB = GetTile(tensorB, tla::MakeCoord(0, 0), tla::MakeShape(kL1Actual, nBlockActual)); + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrB[l1BListId] != tensorTileB.data().GetPhyAddr() || tla::get<0>(tensorTileB.coord()) != lastCoordB[l1BListId].row() || tla::get<1>(tensorTileB.coord()) != lastCoordB[l1BListId].column()) { + copyGmToL1B(tensorL1B, tensorTileB); + lastCoordB[l1BListId] = MatrixCoord{tla::get<0>(tensorTileB.coord()), tla::get<1>(tensorTileB.coord())}; + lastAddrB[l1BListId] = const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileB.data().GetPhyAddr()); + } + } else { + copyGmToL1B(tensorL1B, tensorTileB); + } + AscendC::SetFlag(l1BEventList[l1BListId]); + + if constexpr (HAS_BIAS && !std::is_same_v) { + using CopyGmToL1Bias = typename TileCopy::template CopyGmToL1Bias; + CopyGmToL1Bias copyGmToL1Bias; + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + auto l1Bias = l1BiasTensor.template ReinterpretCast(); + auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Arch::PositionL1{}); + copyGmToL1Bias(tensorL1Bias, tensorBias); + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + } + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(l0CEventList[l0CListId]); + } + + uint32_t mL0Loop = CeilDiv(mL1Actual); + uint32_t nL0Loop = CeilDiv(nL1Actual); + + // main loop + uint32_t kL1Loop = CeilDiv(kBlockActual); + for (uint32_t kL1Idx = 0; kL1Idx < kL1Loop; kL1Idx++) { + uint32_t l1AListIdNext = (l1AListId + 1 < L1A_STAGES) ? (l1AListId + 1) : 0; + uint32_t l1BListIdNext = (l1BListId + 1 < L1B_STAGES) ? (l1BListId + 1) : 0; + uint32_t kL1ActualNext{0}; + // preload next tile from GM to L1 + if (kL1Idx < kL1Loop - 1) { + uint32_t kL1IdxNext = kL1Idx + 1; + kL1ActualNext = (kL1IdxNext < kL1Loop - 1) ? L1_TILE_K : (kBlockActual - kL1IdxNext * L1_TILE_K); + + // Get L1 tensor for next stage + auto l1ATensor = l1ATensorList[l1AListIdNext]; + auto l1BTensor = l1BTensorList[l1BListIdNext]; + auto tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Arch::PositionL1{}); + auto tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Arch::PositionL1{}); + // Get GM tile for next stage + auto tensorTileA = GetTileA(tensorA, 0, kL1IdxNext * L1_TILE_K, mBlockActual, kL1ActualNext); + auto tensorTileB = GetTile(tensorB, tla::MakeCoord(kL1IdxNext * L1_TILE_K, 0), + tla::MakeShape(kL1ActualNext, nBlockActual)); + + // load next matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1AListIdNext]); + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrA[l1AListIdNext] != tensorTileA.data().GetPhyAddr() || tla::get<0>(tensorTileA.coord()) != lastCoordA[l1AListIdNext].row() || tla::get<1>(tensorTileA.coord()) != lastCoordA[l1AListIdNext].column()) { + copyGmToL1A(tensorL1A, tensorTileA); + lastCoordA[l1AListIdNext] = + MatrixCoord{tla::get<0>(tensorTileA.coord()), tla::get<1>(tensorTileA.coord())}; + lastAddrA[l1AListIdNext] = + const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileA.data().GetPhyAddr()); + } + } else { + copyGmToL1A(tensorL1A, tensorTileA); + } + AscendC::SetFlag(l1AEventList[l1AListIdNext]); + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListIdNext]); + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrB[l1BListIdNext] != tensorTileB.data().GetPhyAddr() || tla::get<0>(tensorTileB.coord()) != lastCoordB[l1BListIdNext].row() || tla::get<1>(tensorTileB.coord()) != lastCoordB[l1BListIdNext].column()) { + copyGmToL1B(tensorL1B, tensorTileB); + lastCoordB[l1BListIdNext] = + MatrixCoord{tla::get<0>(tensorTileB.coord()), tla::get<1>(tensorTileB.coord())}; + lastAddrB[l1BListIdNext] = + const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileB.data().GetPhyAddr()); + } + } else { + copyGmToL1B(tensorL1B, tensorTileB); + } + AscendC::SetFlag(l1BEventList[l1BListIdNext]); + } + + // Get L1 tensor for current stage + auto l1ATensor = l1ATensorList[l1AListId]; + auto l1BTensor = l1BTensorList[l1BListId]; + tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Arch::PositionL1{}); + tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Arch::PositionL1{}); + // Get the loop nums on L0 + uint32_t kL0Loop = CeilDiv(kL1Actual); + + for (int mL0Idx = 0; mL0Idx < mL0Loop; mL0Idx++) { + uint32_t mL0Actual = (mL0Idx < mL0Loop - 1) ? L0_TILE_M : (mL1Actual - mL0Idx * L0_TILE_M); + + for (int kL0Idx = 0; kL0Idx < kL0Loop; kL0Idx++) { + uint32_t kL0Actual = (kL0Idx < kL0Loop - 1) ? L0_TILE_K : (kL1Actual - kL0Idx * L0_TILE_K); + + // Locate the current tile on L0A + auto l0ATile = l0ATensorList[l0AListId]; + auto layoutAInL0 = tla::MakeLayout(mL0Actual, kL0Actual); + auto tensorL0A = tla::MakeTensor(l0ATile, layoutAInL0, Arch::PositionL0A{}); + // Locate the current tile of matrix A on L1 + auto tensorTileL1A = GetTileA(tensorL1A, mL0Idx * L0_TILE_M, kL0Idx * L0_TILE_K, mL0Actual, kL0Actual); + + AscendC::WaitFlag(l0AEventList[l0AListId]); + if ((mL0Idx == 0) && (kL0Idx == 0)) { + AscendC::WaitFlag(l1AEventList[l1AListId]); + } + + // Load current tile from L1 to L0A + copyL1ToL0A(tensorL0A, tensorTileL1A); + + if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1)) { + AscendC::SetFlag(l1AEventList[l1AListId]); + } + + bool initC = ((kL1Idx == 0) && (kL0Idx == 0)); + for (int nL0Idx = 0; nL0Idx < nL0Loop; nL0Idx++) { + uint32_t nL0Actual = (nL0Idx < nL0Loop - 1) ? L0_TILE_N : (nL1Actual - nL0Idx * L0_TILE_N); + + // Locate the current tile on L0B + auto l0BTile = l0BTensorList[l0BListId]; + auto layoutBInL0 = tla::MakeLayout(kL0Actual, nL0Actual); + auto tensorL0B = tla::MakeTensor(l0BTile, layoutBInL0, Arch::PositionL0B{}); + // Locate the current tile of matrix B on L1 + auto tensorTileL1B = GetTile(tensorL1B, + tla::MakeCoord(kL0Idx * L0_TILE_K, nL0Idx * L0_TILE_N), + tla::MakeShape(kL0Actual, nL0Actual)); + + // Wait for mmad finished + AscendC::WaitFlag(l0BEventList[l0BListId]); + // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 + if ((mL0Idx == 0) && (kL0Idx == 0) && (nL0Idx == 0)) { + AscendC::WaitFlag(l1BEventList[l1BListId]); + } + + // Load current tile from L1 to L0B + copyL1ToL0B(tensorL0B, tensorTileL1B); + + // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 + if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1) && (nL0Idx == nL0Loop - 1)) { + AscendC::SetFlag(l1BEventList[l1BListId]); + } + + if constexpr (HAS_BIAS && !std::is_same_v) { + if (initC) { + if (nL0Idx == 0) { + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + } + AscendC::WaitFlag(L0A_STAGES + L0B_STAGES); + auto l1Bias = l1BiasTensor.template ReinterpretCast(); + auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Arch::PositionL1{}); + auto tensorTileL1Bias = GetTile(tensorL1Bias, + tla::MakeCoord(nL0Idx * L0_TILE_N), + tla::MakeShape(nL0Actual)); + // Load bias to l0 biasTable + copyL1ToBT(tensorL0Bias, tensorTileL1Bias); + if (nL0Idx == nL0Loop - 1) { + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + } + } + } + + // Notify to do mmad + AscendC::SetFlag(l0CEventList[l0CListId]); + + // Locate the current tile on L0C + auto tensorTileL0C = GetTile(tensorL0C, + tla::MakeCoord(mL0Idx * L0_TILE_M, nL0Idx * L0_TILE_N), + tla::MakeShape(mL0Actual, nL0Actual)); + + // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator + // Wait for loading L0B + AscendC::WaitFlag(l0CEventList[l0CListId]); + + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if ((kL1Idx == kL1Loop - 1) && (mL0Idx == mL0Loop - 1) && + (kL0Idx == kL0Loop - 1) && (nL0Idx == nL0Loop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + + if constexpr (HAS_BIAS && !std::is_same_v) { + if (initC) { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, tensorL0Bias, + mL0Actual, nL0Actual, kL0Actual, initC, unitFlag); + AscendC::SetFlag(L0A_STAGES + L0B_STAGES); + } else { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, + mL0Actual, nL0Actual, kL0Actual, initC, unitFlag); + } + } else { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, + mL0Actual, nL0Actual, kL0Actual, initC, unitFlag); + } + + // Notify to move the next L0B tile + AscendC::SetFlag(l0BEventList[l0BListId]); + l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0; + } + AscendC::SetFlag(l0AEventList[l0AListId]); + l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0; + } + } + l1AListId = l1AListIdNext; + l1BListId = l1BListIdNext; + kL1Actual = kL1ActualNext; + } + + // copy block out + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(l0CEventList[l0CListId]); + AscendC::WaitFlag(l0CEventList[l0CListId]); + copyL0CToDst(tensorC, tensorL0C); + AscendC::SetFlag(l0CEventList[l0CListId]); + l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; + } else { + copyL0CToDst(tensorC, tensorL0C, 0b11); + } + } + +protected: + template + CATLASS_DEVICE auto GetTileA(TensorA &tensorA, uint32_t mIndex, uint32_t kIndex, uint32_t mSize, uint32_t kSize) + { + if constexpr (tla::detail::isVector::value) { + return GetTile(tensorA, tla::MakeCoord(kIndex), tla::MakeShape(kSize)); + } else { + return GetTile(tensorA, tla::MakeCoord(mIndex, kIndex), tla::MakeShape(mSize, kSize)); + } + } + + // Multi-stage tensors list + AscendC::LocalTensor l1ATensorList[L1A_STAGES]; + AscendC::LocalTensor l1BTensorList[L1B_STAGES]; + AscendC::LocalTensor l0ATensorList[L0A_STAGES]; + AscendC::LocalTensor l0BTensorList[L0B_STAGES]; + AscendC::LocalTensor l0CTensorList[L0C_STAGES]; + AscendC::LocalTensor l1BiasTensor; + AscendC::LocalTensor l0BiasTensor; + + // Multi-stage event id list + int32_t l1AEventList[L1A_STAGES]; + int32_t l1BEventList[L1B_STAGES]; + int32_t l0AEventList[L0A_STAGES]; + int32_t l0BEventList[L0B_STAGES]; + int32_t l0CEventList[L0C_STAGES]; + + __gm__ typename AscendC::GlobalTensor::PrimType *lastAddrA[L1A_STAGES]; + __gm__ typename AscendC::GlobalTensor::PrimType *lastAddrB[L1B_STAGES]; + MatrixCoord lastCoordA[L1A_STAGES]; + MatrixCoord lastCoordB[L1B_STAGES]; + + // The id of current stage + uint32_t l1AListId{0}; + uint32_t l1BListId{0}; + uint32_t l0AListId{0}; + uint32_t l0BListId{0}; + uint32_t l0CListId{0}; + + TileMmad tileMmad; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL1ToBT copyL1ToBT; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad.hpp new file mode 100644 index 0000000000..290d58344b --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/block/kda_block_mmad.hpp @@ -0,0 +1,43 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_BLOCK_MMAD_HPP +#define CATLASS_KDA_BLOCK_MMAD_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" + +namespace Catlass::Gemm::Block { + +template < + class DispatchPolicy, + class L1TileShape, + class L0TileShape, + class ElementA, + class ElementB, + class ElementC, + class ElementBias = void, + class TileCopy = Gemm::Tile::PackedTileCopyTla, + class TileMmad = + Gemm::Tile::TileMmadTla > +struct BlockMmadTla { + static_assert(DEPENDENT_FALSE, "BlockMmadTla is not implemented for this DispatchPolicy"); +}; + +} // namespace Catlass::Gemm::Block + +// Only the pingpong-TLA implementation is instantiated by chunk_kda_fwd +// (Gemm::MmadPingpong). The MmadPingpongTlaMulti / MmadPingpongTlaPreloadAL1B +// specializations live under kernel_utils/block/ and are included by their users. +#include "catlass/gemm/block/block_mmad_pingpong_tla.hpp" + +#endif // CATLASS_KDA_BLOCK_MMAD_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/kda_gemm_dispatch_policy.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/kda_gemm_dispatch_policy.hpp new file mode 100644 index 0000000000..8c30bb8837 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/kda_gemm_dispatch_policy.hpp @@ -0,0 +1,373 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_GEMM_DISPATCH_POLICY_HPP +#define CATLASS_KDA_GEMM_DISPATCH_POLICY_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_arch.hpp" + +namespace Catlass::Gemm { + +// Block Mmad Policies + +template +struct MmadBase { + using ArchTag = ArchTag_; + static constexpr uint32_t ASYNC = ASYNC_; +}; + +using MmadAtlasA2 = MmadBase; +using MmadAtlasA2Async = MmadBase; + +// Now ENABLE_UNIT_FLAG_ must be false when intput element is int8 +template +struct MmadAtlasA2Pingpong : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadAtlasA2PingPongWithPrologue : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadAtlasA2SingleCoreSplitk : public MmadAtlasA2 { + static constexpr uint32_t L1A_STAGES = L1A_STAGES_; + static constexpr uint32_t L1B_STAGES = L1B_STAGES_; + static constexpr uint32_t L0AB_STAGES = 2; + static constexpr uint32_t L0C_STAGES = L0C_STAGES_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static_assert(!(ENABLE_UNIT_FLAG && (L0C_STAGES > 1)), "When L0C_STAGES > 1, can not enable unitflag"); +}; + +template +struct MmadAtlasA2PingpongSliceKWithPrologue : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadAtlasA2Preload : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool ENABLE_SHUFFLE_K = ENABLE_SHUFFLE_K_; +}; + +struct MmadAtlasA2FAQK : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; +}; + +struct MmadAtlasA2FAPV : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; +}; + +struct MmadAtlasA2MLAQK : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; +}; + +struct MmadAtlasA2MLAPV : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; +}; + +struct MmadAtlasA2MLAQKTp1Spec : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; +}; + +struct MmadAtlasA2MLAPVTp1Spec : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; +}; + +template +struct MmadAtlasA2PreloadAsync : public MmadAtlasA2Async { + static constexpr uint32_t PRELOAD_STAGES = PRELOAD_STAGES_; // Stages of emitting load instruction in advance + static constexpr uint32_t L1_STAGES = L1_STAGES_; + static constexpr uint32_t L0A_STAGES = L0A_STAGES_; + static constexpr uint32_t L0B_STAGES = L0B_STAGES_; + static constexpr uint32_t L0C_STAGES = L0C_STAGES_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool ENABLE_SHUFFLE_K = ENABLE_SHUFFLE_K_; +}; + +template +struct MmadAtlasA2PreloadAsyncWithCallback : public MmadAtlasA2PreloadAsync< + PRELOAD_STAGES_, + L1_STAGES_, + L0A_STAGES_, + L0B_STAGES_, + L0C_STAGES_, + ENABLE_UNIT_FLAG_, + ENABLE_SHUFFLE_K_> { +}; + +template +struct MmadAtlasA2DynamicPreloadAsyncWithCallback : public MmadAtlasA2PreloadAsync< + PRELOAD_STAGES_, + L1_STAGES_, + L0A_STAGES_, + L0B_STAGES_, + L0C_STAGES_, + ENABLE_UNIT_FLAG_, + ENABLE_SHUFFLE_K_> { +}; + +template +struct MmadAtlasA2W4A4MatmulPerTokenPerChannelDequant : public MmadAtlasA2PreloadAsyncWithCallback< + PRELOAD_STAGES_, + L1_STAGES_, + L0A_STAGES_, + L0B_STAGES_, + L0C_STAGES_, + ENABLE_UNIT_FLAG_, + ENABLE_SHUFFLE_K_> { +}; +//////////////////// +// new add +template +struct GemmAtlasA2 : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool ENABLE_SHUFFLE_K = ENABLE_SHUFFLE_K_; + static constexpr bool ENABLE_ABBA = ENABLE_ABBA_; +}; + +struct GemvAtlasA2 : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; +}; +//////////////////// + +template +struct MmadAtlasA2PingpongBias : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadAtlasA2FAIQK : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadAtlasA2FAIPV : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadAtlasA2FAITailQK : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadAtlasA2FAITailPV : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadAtlasA2FullLoadA : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadFullLoadA : public MmadBase { + static constexpr uint32_t L1A_STAGES = L1A_STAGES_; + static constexpr uint32_t L1B_STAGES = L1B_STAGES_; + static constexpr uint32_t L0A_STAGES = L0A_STAGES_; + static constexpr uint32_t L0B_STAGES = L0B_STAGES_; + static constexpr uint32_t L0C_STAGES = L0C_STAGES_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool ENABLE_SHUFFLE_K = ENABLE_SHUFFLE_K_; + static constexpr bool USE_HF32_MODE = USE_HF32_MODE_; + static constexpr bool ENABLE_L1_RESIDENT = ENABLE_L1_RESIDENT_; +}; + +template +struct MmadAtlasA2W8A16 : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool ENABLE_SHUFFLE_K = ENABLE_SHUFFLE_K_; +}; + +template +struct MmadAtlasA2DynamicCommon : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool ENABLE_SHUFFLE_K = ENABLE_SHUFFLE_K_; +}; + +template +struct MmadAtlasA2DynamicSmall : public MmadAtlasA2 { + static constexpr uint32_t STAGES = STAGES_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadAtlasA2Streamk : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool ENABLE_SHUFFLE_K = ENABLE_SHUFFLE_K_; +}; + +template +struct MmadAtlasA2DynamicStreamk : public MmadAtlasA2 { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool ENABLE_SHUFFLE_K = ENABLE_SHUFFLE_K_; +}; + +template +struct MmadAtlasA2DynamicSingleCoreSplitk : public MmadAtlasA2 { + static constexpr uint32_t L1A_STAGES = L1A_STAGES_; + static constexpr uint32_t L1B_STAGES = L1B_STAGES_; + static constexpr uint32_t L0AB_STAGES = 2; + static constexpr uint32_t L0C_STAGES = L0C_STAGES_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static_assert(!(ENABLE_UNIT_FLAG && (L0C_STAGES > 1)), "When L0C_STAGES > 1, can not enable unitflag"); +}; + +template +struct MmadAtlasA2DynamicAiv : public MmadAtlasA2 { + static constexpr uint32_t STAGES = STAGES_; + static constexpr uint32_t SCALAR_BUFFER_ELE_NUM = SCALAR_BUFFER_ELE_NUM_; +}; + +// UB must have a capacity of UBA(mTile) + UBB(nTile) + UBC(mTile, nTile) size. +// MmadAtlasA2AivSimple is suitable for cases: mTile <= mTileUbLimits && nTile >= params.n +template +struct MmadAtlasA2DynamicAivSimple : public MmadAtlasA2 { + static constexpr uint32_t SCALAR_BUFFER_ELE_NUM = SCALAR_BUFFER_ELE_NUM_; + static constexpr bool IS_TILE_M = IS_TILE_M_; +}; + +// mTile and nTile must be aligned to 16. +// UB must have a capacity of UBA(mTile) + UBB(nTile) + 2 * UBC(mTile, nTile) size. +// MmadAtlasA2AivTrans is suitable for cases: M > N +template +struct MmadAtlasA2DynamicAivTrans : public MmadAtlasA2 { + static constexpr uint32_t SCALAR_BUFFER_ELE_NUM = SCALAR_BUFFER_ELE_NUM_; +}; + +template +struct MmadAtlasA2Small : public MmadAtlasA2 { + static constexpr uint32_t STAGES = STAGES_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool ENABLE_SHUFFLE_K = ENABLE_SHUFFLE_K_; +}; +// Now ENABLE_UNIT_FLAG_ must be false when intput element is int8 +template +struct MmadPingpong : public MmadBase { + static constexpr uint32_t L1A_STAGES = L1A_STAGES_; + static constexpr uint32_t L1B_STAGES = L1B_STAGES_; + static constexpr uint32_t L0A_STAGES = L0A_STAGES_; + static constexpr uint32_t L0B_STAGES = L0B_STAGES_; + static constexpr uint32_t L0C_STAGES = L0C_STAGES_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool USE_HF32_MODE = USE_HF32_MODE_; + static constexpr bool ENABLE_L1_RESIDENT = ENABLE_L1_RESIDENT_; +}; + +template +struct MmadPreloadAsyncWithCallback : public MmadBase { + static constexpr uint32_t PRELOAD_STAGES = PRELOAD_STAGES_; + static constexpr uint32_t L1A_STAGES = L1A_STAGES_; + static constexpr uint32_t L1B_STAGES = L1B_STAGES_; + static constexpr uint32_t L0A_STAGES = L0A_STAGES_; + static constexpr uint32_t L0B_STAGES = L0B_STAGES_; + static constexpr uint32_t L0C_STAGES = L0C_STAGES_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool ENABLE_SHUFFLE_K = ENABLE_SHUFFLE_K_; + static constexpr bool USE_HF32_MODE = USE_HF32_MODE_; + static constexpr bool ENABLE_L1_RESIDENT = ENABLE_L1_RESIDENT_; +}; + +// 基于TLA提供block层不感知尾块逻辑，tile层感知originShape的流程。使用TileView与MakeTensorLike。 +template +struct MmadPingpongTlaV2 : public MmadBase { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct SparseMatmulMultiBlockOnKAxis : public MmadBase { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadFAIQK : public MmadBase { + static constexpr uint32_t STAGES = 2; + static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadFAITailQK : public MmadBase { + static constexpr uint32_t STAGES = 2; + static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadFAIPV : public MmadBase { + static constexpr uint32_t STAGES = 2; + static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadFAITailPV : public MmadBase { + static constexpr uint32_t STAGES = 2; + static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +template +struct MmadDequant : public MmadBase { + static constexpr uint32_t L1A_STAGES = L1A_STAGES_; + static constexpr uint32_t L1B_STAGES = L1B_STAGES_; + static constexpr uint32_t L0A_STAGES = L0A_STAGES_; + static constexpr uint32_t L0B_STAGES = L0B_STAGES_; + static constexpr uint32_t L0C_STAGES = L0C_STAGES_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool USE_HF32_MODE = USE_HF32_MODE_; + static constexpr bool ENABLE_L1_RESIDENT = ENABLE_L1_RESIDENT_; +}; + +template +struct MmadPingpongPertile : public MmadBase { + static constexpr uint32_t STAGES = 2; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +} // namespace Catlass::Gemm + +#endif // CATLASS_KDA_GEMM_DISPATCH_POLICY_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/kda_gemm_helper.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/kda_gemm_helper.hpp new file mode 100644 index 0000000000..4ae8778601 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/kda_gemm_helper.hpp @@ -0,0 +1,385 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_GEMM_HELPER_HPP +#define CATLASS_KDA_GEMM_HELPER_HPP + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/kda_numeric_size.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "tla/kda_tla_layout.hpp" + +namespace Catlass::Gemm::helper { + +template +struct L1AlignHelper { + static_assert(DEPENDENT_FALSE, "Unsupported align helper, can not find the specialization."); +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t M_ALIGNED = C0_NUM_PER_FRACTAL; + static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t N_ALIGNED = ELE_NUM_PER_C0; +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t M_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t N_ALIGNED = C0_NUM_PER_FRACTAL; +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t M_ALIGNED = C0_NUM_PER_FRACTAL; + static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t N_ALIGNED = ELE_NUM_PER_C0; +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t M_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t N_ALIGNED = C0_NUM_PER_FRACTAL; +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t M_ALIGNED = C0_NUM_PER_FRACTAL; + static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t N_ALIGNED = ELE_NUM_PER_C0; +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t M_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t N_ALIGNED = C0_NUM_PER_FRACTAL; +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t HOWO_ALIGNED = C0_NUM_PER_FRACTAL; +}; + +template +struct L1AlignHelper { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t COUT_ALIGNED = C0_NUM_PER_FRACTAL; +}; + +template +struct ElementAccumulatorSelector { + static_assert(DEPENDENT_FALSE, + "Unsupported element accumulator selector, can not find the specialization."); +}; + +template <> +struct ElementAccumulatorSelector { + using ElementAccumulator = float; +}; + +template <> +struct ElementAccumulatorSelector { + using ElementAccumulator = float; +}; + +template <> +struct ElementAccumulatorSelector { + using ElementAccumulator = int32_t; +}; + +template <> +struct ElementAccumulatorSelector { + using ElementAccumulator = float; +}; + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) +template <> +struct ElementAccumulatorSelector { + using ElementAccumulator = float; +}; + +template <> +struct ElementAccumulatorSelector { + using ElementAccumulator = float; +}; + +template <> +struct ElementAccumulatorSelector { + using ElementAccumulator = float; +}; + +template <> +struct ElementAccumulatorSelector { + using ElementAccumulator = float; +}; +#endif + +template <> +struct ElementAccumulatorSelector { + using ElementAccumulator = int32_t; +}; + +template +struct L1ATypeSelector { + static_assert(DEPENDENT_FALSE, + "Unsupported layout selector, can not find the specialization."); +}; + +template +struct L1ATypeSelector> { + using L1AType = Gemm::GemmType; +}; + +template +struct L1ATypeSelector> { + using L1AType = Gemm::GemmType; +}; + +template +struct L1ATypeSelector> { + using L1AType = Gemm::GemmType; +}; + +template +struct L1ATypeSelector> { + using L1AType = Gemm::GemmType; +}; + +template +struct L1ATypeSelector> { + using L1AType = Gemm::GemmType; +}; + +template +struct L1ATypeSelector> { + using L1AType = Gemm::GemmType; +}; + +template +struct L1ATypeSelector> { + using L1AType = Gemm::GemmType; +}; + +template +struct L1ATypeSelector> { + using L1AType = Gemm::GemmType; +}; + +template +struct L1ATypeSelector> { + using L1AType = Gemm::GemmType; +}; + +template +struct L1BTypeSelector { + static_assert(DEPENDENT_FALSE, + "Unsupported layout selector, can not find the specialization."); +}; + +template +struct L1BTypeSelector> { + using L1BType = Gemm::GemmType; +}; + +template +struct L1BTypeSelector> { + using L1BType = Gemm::GemmType; +}; + +template +struct L1BTypeSelector> { + using L1BType = Gemm::GemmType; +}; + +template +struct L1BTypeSelector> { + using L1BType = Gemm::GemmType; +}; + +template +struct L1BTypeSelector> { + using L1BType = Gemm::GemmType; +}; + +template +struct L1BTypeSelector> { + using L1BType = Gemm::GemmType; +}; + +template +struct L1BTypeSelector> { + using L1BType = Gemm::GemmType; +}; + +template +struct L1BTypeSelector> { + using L1BType = Gemm::GemmType; +}; + +template +struct L1BiasTypeSelector { + static_assert(DEPENDENT_FALSE, + "Unsupported layout selector, can not find the specialization."); +}; + +template +struct L1BiasTypeSelector { + using GMBiasType = void; + using L1BiasType = void; + using L0BiasType = void; +}; + +template +struct L1BiasTypeSelector, ElementAccumulator> { + using GMBiasType = Gemm::GemmType; + using L1BiasType = Gemm::GemmType; + using L0BiasType = Gemm::GemmType; +}; + +template +struct L0ALayoutSelector { + static_assert(DEPENDENT_FALSE, + "Unsupported layout selector, can not find the specialization."); +}; + +template <> +struct L0ALayoutSelector { + using Layout = layout::zZ; +}; + +template <> +struct L0ALayoutSelector { + using Layout = layout::zN; +}; + +template +struct L1AlignHelperTla { + static_assert(DEPENDENT_FALSE, "Unsupported align helper tla, can not find the specialization."); +}; + +template +struct L1AlignHelperTla::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t M_ALIGNED = C0_NUM_PER_FRACTAL; + static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t N_ALIGNED = ELE_NUM_PER_C0; +}; + +template +struct L1AlignHelperTla::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t M_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0; + static constexpr uint32_t N_ALIGNED = C0_NUM_PER_FRACTAL; +}; + +/////////////////////////////////////// +// new add +template <> +struct ElementAccumulatorSelector { + using ElementAccumulator = int32_t; +}; + +template +struct L1AndL0TypeSelectorGemm { + static_assert(DEPENDENT_FALSE, + "Unsupported layout selector, can not find the specialization."); + static_assert(DEPENDENT_FALSE, + "Unsupported layout selector, can not find the specialization."); +}; + +template +struct L1AndL0TypeSelectorGemm, Gemm::GemmType> { + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; + +template <> +struct L1AndL0TypeSelectorGemm, Gemm::GemmType> { + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; + +template +struct L1AndL0TypeSelectorGemm, Gemm::GemmType> { + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; + +template <> +struct L1AndL0TypeSelectorGemm, Gemm::GemmType> { + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; + +template +struct L1AndL0TypeSelectorGemm, Gemm::GemmType> { + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; + +template +struct L1AndL0TypeSelectorGemm, Gemm::GemmType> { + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; + +template <> +struct L1AndL0TypeSelectorGemm, Gemm::GemmType> { + using L1AType = Gemm::GemmType; + using L1BType = Gemm::GemmType; + using L0AType = Gemm::GemmType; + using L0BType = Gemm::GemmType; +}; + +/////////////////////////////////////// +// Check for whether the tile shape is aligned (Default: 32Bytes). +template +struct TileShapeAlignChecker { + static constexpr uint32_t _ALIGN = _aligned * 8; + + static_assert(L1TileShape::M * SizeOfBits::value % _ALIGN == 0, "L1TileShape::M is not aligned."); + static_assert(L0TileShape::M * SizeOfBits::value % _ALIGN == 0, "L0TileShape::M is not aligned."); + static_assert(L1TileShape::K * SizeOfBits::value % _ALIGN == 0, "L1TileShape::K is not aligned."); + static_assert(L1TileShape::K * SizeOfBits::value % _ALIGN == 0, "L1TileShape::K is not aligned."); + static_assert(L0TileShape::K * SizeOfBits::value % _ALIGN == 0, "L0TileShape::K is not aligned."); + static_assert(L0TileShape::K * SizeOfBits::value % _ALIGN == 0, "L0TileShape::K is not aligned."); + static_assert(L1TileShape::N * SizeOfBits::value % _ALIGN == 0, "L1TileShape::N is not aligned."); + static_assert(L0TileShape::N * SizeOfBits::value % _ALIGN == 0, "L0TileShape::N is not aligned."); +}; + +/////////////////////////////////////// +} // namespace Catlass::Gemm::helper + +#endif // CATLASS_KDA_GEMM_HELPER_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/kda_gemm_type.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/kda_gemm_type.hpp new file mode 100644 index 0000000000..9d680d795b --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/kda_gemm_type.hpp @@ -0,0 +1,29 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_GEMM_TYPE_HPP +#define CATLASS_KDA_GEMM_TYPE_HPP + +#include "catlass/kda_catlass.hpp" + +namespace Catlass::Gemm { + +//////////////////////////////////////////////////////////////////// + +template +struct GemmType { + using Element = Element_; + using Layout = Layout_; + static constexpr AscendC::TPosition POSITION = POSITION_; +}; + +} // namespace Catlass::Gemm + +#endif // CATLASS_KDA_GEMM_TYPE_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/copy_l0c_to_ub.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/copy_l0c_to_ub.hpp new file mode 100644 index 0000000000..3397404af8 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/copy_l0c_to_ub.hpp @@ -0,0 +1,166 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_TILE_ASCEND950_COPY_L0C_TO_UB_950_HPP +#define CATLASS_GEMM_TILE_ASCEND950_COPY_L0C_TO_UB_950_HPP + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/gemm/tile/ascend950/kda_ascend950_copy_l0c_to_dst.hpp" +#include "tla/kda_tensor.hpp" + +#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510) +constexpr AscendC::FixpipeConfig CFG_ROW_MAJOR_UB = {AscendC::CO2Layout::ROW_MAJOR, true}; +#endif + +namespace Catlass::Gemm::Tile { + +template +struct CopyL0CToUBTla< + Catlass::Arch::Ascend950, + TensorSrc_, + tla::Tensor, LayoutDst_, CoordDst_, AscendC::TPosition::VECCALC>, + CopyL0CToUBMode::NO_SPLIT, + ScaleGranularity::NO_QUANT, + ReluEnable_, + std::enable_if_t::value>> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = typename TensorSrc_::Element; + static constexpr auto quantPre = + CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t unitFlag = 0) + { + static_assert( + tla::detail::isRowMajor::value && TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::VECCALC, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and RowMajor"); + + AscendC::FixpipeParamsC310 intriParams; + + // Fixpipe layout information + intriParams.nSize = tla::get<1>(dstTensor.originShape()); + intriParams.mSize = tla::get<0>(dstTensor.originShape()); + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + intriParams.dstStride = tla::get<0>(dstTensor.stride()); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + // Call AscendC Fixpipe + AscendC::Fixpipe( + dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +template +struct CopyL0CToUBTla< + Catlass::Arch::Ascend950, + TensorSrc_, + tla::Tensor, LayoutDst_, CoordDst_, AscendC::TPosition::VECCALC>, + CopyL0CToUBMode::SPLIT_M, + ScaleGranularity::NO_QUANT, + ReluEnable_, + std::enable_if_t::value>> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = typename TensorSrc_::Element; + static constexpr auto quantPre = + CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t unitFlag = 0) + { + static_assert( + tla::detail::isRowMajor::value && TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::VECCALC, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and RowMajor"); + + AscendC::FixpipeParamsC310 intriParams; + + // Fixpipe layout information + intriParams.nSize = tla::get<1>(dstTensor.originShape()); + intriParams.mSize = RoundUp(tla::get<0>(dstTensor.originShape()), 2); // m must be even when spilt m + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + intriParams.dstStride = tla::get<0>(dstTensor.stride()); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + intriParams.dualDstCtl = 1; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + // Call AscendC Fixpipe + AscendC::Fixpipe( + dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +template +struct CopyL0CToUBTla< + Catlass::Arch::Ascend950, + TensorSrc_, + tla::Tensor, LayoutDst_, CoordDst_, AscendC::TPosition::VECCALC>, + CopyL0CToUBMode::SPLIT_N, + ScaleGranularity::NO_QUANT, + ReluEnable_, + std::enable_if_t::value>> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = typename TensorSrc_::Element; + static constexpr auto quantPre = + CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t unitFlag = 0) + { + static_assert( + tla::detail::isRowMajor::value && TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::VECCALC, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and RowMajor"); + + AscendC::FixpipeParamsC310 intriParams; + + // Fixpipe layout information + intriParams.nSize = RoundUp(tla::get<1>(dstTensor.originShape()), 32); + intriParams.mSize = tla::get<0>(dstTensor.originShape()); // m must be even when spilt m + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + intriParams.dstStride = tla::get<0>(dstTensor.stride()); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + intriParams.dualDstCtl = 2; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + // Call AscendC Fixpipe + AscendC::Fixpipe( + dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_GEMM_TILE_ASCEND950_COPY_L0C_TO_UB_950_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_gm_to_l1.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_gm_to_l1.hpp new file mode 100644 index 0000000000..862a427e49 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_gm_to_l1.hpp @@ -0,0 +1,727 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ASCEND950_COPY_GM_TO_L1_HPP +#define CATLASS_KDA_ASCEND950_COPY_GM_TO_L1_HPP + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/gemm/tile/kda_tile_copy_tla.hpp" +#include "catlass/kda_numeric_size.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Tile { + +/// Partial specialization for CopyGmToL1, Ascend950, RowMajor in and zN out. +template +struct TileCopyTla< + Arch::Ascend950, + tla::Tensor, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + std::enable_if_t::value && tla::detail::iszN::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()( + TensorDst const &dstTensor, + TensorSrc const &srcTensor, + uint32_t ndNum = 1, + uint32_t srcNdMatrixStride = 0, + uint32_t dstNzMatrixStride = 0) + { + static_assert( + tla::detail::isRowMajor::value && tla::detail::iszN::value && TensorSrc::position == AscendC::TPosition::GM && TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and RowMajor, while TensorDst must be L1 and zN"); + + const uint32_t nValue = tla::get<0>(srcTensor.originShape()); + const uint32_t dValue = tla::get<1>(srcTensor.originShape()); + const uint32_t srcDValue = tla::get<0>(srcTensor.stride()); + const uint32_t dstInnerStrideRow = tla::get<0, 0>(dstTensor.stride()); + const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride()); + + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = ndNum; + intriParams.nValue = nValue; + intriParams.dValue = dValue; + intriParams.srcNdMatrixStride = srcNdMatrixStride; + intriParams.srcDValue = srcDValue; + intriParams.dstNzC0Stride = dstOuterStrideCol / ELE_NUM_PER_C0; + intriParams.dstNzNStride = dstInnerStrideRow / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = dstNzMatrixStride; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +/// Partial specialization for CopyGmToL1, Ascend950, zN in and zN out. +template +struct TileCopyTla< + Arch::Ascend950, + tla::Tensor, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + std::enable_if_t::value && tla::detail::iszN::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert( + tla::detail::iszN::value && tla::detail::iszN::value && TensorSrc::position == AscendC::TPosition::GM && TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and zN, while TensorDst must be L1 and zN"); + + uint32_t blockCount = CeilDiv(tla::get<1>(srcTensor.originShape())); + uint32_t blockLen = tla::get<0>(srcTensor.originShape()); + + AscendC::DataCopyParams repeatParams; + + repeatParams.blockCount = blockCount; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / ELE_NUM_PER_C0 - blockLen; + repeatParams.dstStride = tla::get<1, 1>(dstTensor.stride()) / ELE_NUM_PER_C0 - blockLen; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], repeatParams); + } +}; + +/// Partial specialization for CopyGmToL1, Ascend950, ColumnMajor in and nZ out. +template +struct TileCopyTla< + Arch::Ascend950, + tla::Tensor, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + std::enable_if_t::value && tla::detail::isnZ::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()( + TensorDst const &dstTensor, + TensorSrc const &srcTensor, + uint32_t ndNum = 1, + uint32_t srcNdMatrixStride = 0, + uint32_t dstNzMatrixStride = 0) + { + static_assert( + tla::detail::isColumnMajor::value && tla::detail::isnZ::value && TensorSrc::position == AscendC::TPosition::GM && TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and ColumnMajor, " + "while TensorDst must be L1 and nZ"); + + const uint32_t nValue = tla::get<1>(srcTensor.originShape()); + const uint32_t dValue = tla::get<0>(srcTensor.originShape()); + const uint32_t srcDValue = tla::get<1>(srcTensor.stride()); + const uint32_t dstInnerStrideCol = tla::get<1, 0>(dstTensor.stride()); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = ndNum; + intriParams.nValue = nValue; + intriParams.dValue = dValue; + intriParams.srcNdMatrixStride = srcNdMatrixStride; + intriParams.srcDValue = srcDValue; + intriParams.dstNzC0Stride = dstOuterStrideRow / ELE_NUM_PER_C0; + intriParams.dstNzNStride = dstInnerStrideCol / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = dstNzMatrixStride; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +/// Partial specialization for CopyGmToL1, Ascend950, nZ in and nZ out. +template +struct TileCopyTla< + Arch::Ascend950, + tla::Tensor, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + std::enable_if_t::value && tla::detail::isnZ::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert( + tla::detail::isnZ::value && tla::detail::isnZ::value && TensorSrc::position == AscendC::TPosition::GM && TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and nZ, " + "while TensorDst must be L1 and nZ"); + + uint32_t blockCount = CeilDiv(tla::get<0>(srcTensor.originShape())); + uint32_t blockLen = tla::get<1>(srcTensor.originShape()); + + AscendC::DataCopyParams repeatParams; + + repeatParams.blockCount = blockCount; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = tla::get<0, 1>(srcTensor.stride()) / ELE_NUM_PER_C0 - blockLen; + repeatParams.dstStride = tla::get<0, 1>(dstTensor.stride()) / ELE_NUM_PER_C0 - blockLen; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], repeatParams); + } +}; + +/// Partial specialization for CopyGmToL1, Ascend950, VectorLayout in and VectorLayout out. +template +struct TileCopyTla< + Arch::Ascend950, + tla::Tensor, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + std::enable_if_t::value && tla::detail::isVector::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert( + tla::detail::isVector::value && tla::detail::isVector::value && TensorSrc::position == AscendC::TPosition::GM && TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and Vector, " + "while TensorDst must be L1 and Vector"); + + AscendC::DataCopyParams intriParams; + intriParams.blockCount = 1; + intriParams.blockLen = CeilDiv(tla::get<0>(srcTensor.originShape()), ELE_NUM_PER_C0); + intriParams.srcStride = 0; + intriParams.dstStride = 0; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +////////////////////////////////////CopyGmToL1(No-TLA, Ascend950)//////////////////////////////////////////////// +template < + class ArchTag, + /// GemmType for matrix operand + class GmType, + class L1Type = void, + class Enable = void> +struct CopyGmToL1 { + static_assert(DEPENDENT_FALSE, "Unsupported copy gm to l1, can not find the specialization."); +}; + +template < + class ArchTag, + /// GemmType for matrix operand + class GmType, + class L1Type = void> +struct CopyGmToL1GMMPTD { + static_assert(DEPENDENT_FALSE, "Unsupported copy gm to l1, can not find the specialization."); +}; + +template < + class ArchTag, + /// GemmType for matrix operand + class GmType, + class L1Type = void> +struct CopyGmToL1DynamicOptimized { + static_assert(DEPENDENT_FALSE, "Unsupported copy gm to l1, can not find the specialization."); +}; + +template +struct CopyGmToL1> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + const uint32_t dstInnerStrideRow = layoutDst.stride(0); + const uint32_t dstOuterStrideCol = layoutDst.stride(3); + + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.nValue = layoutSrc.shape(0); + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcDValue = layoutSrc.stride(0); + + // strideColsByFractal --> + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } + + // layoutSrc must be the layout of one of the src matrices + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc, + uint32_t ndNum, uint32_t srcNdMatrixStride, + uint32_t dstNzNStride, uint32_t dstNzMatrixStride, + uint32_t dstNzC0Stride) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = ndNum; + intriParams.nValue = layoutSrc.shape(0); + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcDValue = layoutSrc.stride(0); + + // Manually set the copy stride + intriParams.dstNzNStride = dstNzNStride; + intriParams.dstNzC0Stride = dstNzC0Stride; + intriParams.srcNdMatrixStride = srcNdMatrixStride; + intriParams.dstNzMatrixStride = dstNzMatrixStride; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +/// Partial specialization for CopyGmToL1(AtlasA5 no-tla), ColumnMajor in and nZ out. +template +struct CopyGmToL1> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::ColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + // Note: AscendC::int4b_t no longer supported on A5-platform + + // Methods + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + // dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()) --> + // tag.stride(1) --> strideRowsByFractal[语义] + // dstInnerStrideCol = tla::get<1, 0>(dstTensor.stride()) --> + // tag.stride(2) --> strideColsInFractal[语义] + + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.nValue = layoutSrc.shape(1); + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcDValue = layoutSrc.stride(1); + + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; // Outer stride -- col + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +/// Partial specialization for CopyGmToL1(AtlasA5 no-tla), zN in and zN out. +template +struct CopyGmToL1> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + // Note: AscendC::int4b_t no longer supported on A5-platform + + // Methods + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + const uint32_t blockLen = layoutSrc.shape(0) * layoutSrc.shape(1); + AscendC::DataCopyParams repeatParams; + + repeatParams.blockCount = layoutSrc.shape(3); + repeatParams.blockLen = blockLen; + + repeatParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_C0 - blockLen; + repeatParams.dstStride = layoutDst.stride(3) / ELE_NUM_PER_C0 - blockLen; + + AscendC::DataCopy(dstTensor, srcTensor, repeatParams); + } +}; + +/// Partial specialization for CopyGmToL1(AtlasA5 no-tla), nZ in and nZ out. +template +struct CopyGmToL1> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + // Note: AscendC::int4b_t no longer supported on A5-platform + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + const uint32_t blockLen = layoutSrc.shape(2) * layoutSrc.shape(3); + AscendC::DataCopyParams repeatParams; + + repeatParams.blockCount = layoutSrc.shape(1); + repeatParams.blockLen = blockLen; + + repeatParams.srcStride = layoutSrc.stride(1) / ELE_NUM_PER_C0 - blockLen; + repeatParams.dstStride = layoutDst.stride(1) / ELE_NUM_PER_C0 - blockLen; + + AscendC::DataCopy(dstTensor, srcTensor, repeatParams); + } +}; + +/// Partial specialization for CopyGmToL1(no-tla), AtlasA5, fp8_e8m0_t, MxScaleA RowMajor in and zZ out. +template +struct CopyGmToL1, + Gemm::GemmType, + std::enable_if_t>> { + using LayoutDst = layout::zZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = 2; + + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + if (layoutSrc.shape(2) != ELE_NUM_PER_C0) { + // std::cerr << "layoutSrc.shape(2) != 2" << std::endl; + return; + } + + AscendC::Dn2NzParams intriParams; + + intriParams.dnNum = 1; + intriParams.nValue = CeilDiv(layoutSrc.shape(2) * layoutSrc.shape(3)); + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcDnMatrixStride = 0; + + intriParams.srcDValue = CeilDiv(layoutSrc.stride(0)); + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = 1; + intriParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +/// Partial specialization for AtlasA2, PaddingRowMajor in and zN out. +template +struct CopyGmToL1> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::PaddingRowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Mehtods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.orgShape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + intriParams.nValue = layoutSrc.orgShape(0); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +/// Partial specialization for AtlasA2, ColumnMajor in and nZ out. +template < + class Element> +struct CopyGmToL1> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::PaddingColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Mehtods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.orgShape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + intriParams.nValue = layoutSrc.orgShape(1); + intriParams.srcDValue = layoutSrc.stride(2); + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +//////////////////////////// CopyGmToL1DynamicOptimized(Ascend950, No TLA) //////////////////////////// + +/// Partial specialization for Ascend950, zN in and zN out. +template +struct CopyGmToL1DynamicOptimized> : public CopyGmToL1> {}; + +/// Partial specialization for Ascend950, nZ in and nZ out. +template +struct CopyGmToL1DynamicOptimized> : public CopyGmToL1> {}; + +/// Partial specialization for Ascend950, RowMajor in and zN out. +template +struct CopyGmToL1DynamicOptimized> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + CATLASS_DEVICE + CopyGmToL1DynamicOptimized() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + if (layoutSrc.shape(0) <= 16) { + // If the number of matrix row is very small, call the regular interval-based data-copy + for (int i = 0; i < layoutSrc.shape(0); ++i) { + AscendC::DataCopyParams dataCopyParams( + CeilDiv(layoutSrc.shape(1), layoutDst.shape(2)), + layoutDst.shape(2) / ELE_NUM_PER_C0, + 0, + (layoutDst.stride(3) - layoutDst.shape(2)) / ELE_NUM_PER_C0); + AscendC::DataCopy(dstTensor[i * layoutDst.shape(2)], srcTensor[i * layoutSrc.stride(0)], dataCopyParams); + } + } else { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.nValue = layoutSrc.shape(0); + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } + } +}; + +/// Partial specialization for Ascend950, ColumnMajor in and nZ out. +template +struct CopyGmToL1DynamicOptimized> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::ColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + CATLASS_DEVICE + CopyGmToL1DynamicOptimized() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + if (layoutSrc.shape(1) <= 16) { + // If the number of matrix cols is 1, the regular interval-based DataCopy interface can be used instead of + // the ND2NZ DataCopy interface, resulting in higher transfer efficiency. + for (int i = 0; i < layoutSrc.shape(1); ++i) { + AscendC::DataCopyParams dataCopyParams( + CeilDiv(layoutSrc.shape(0), layoutDst.shape(0)), + layoutDst.shape(0) / ELE_NUM_PER_C0, + 0, + (layoutDst.stride(1) - layoutDst.shape(0)) / ELE_NUM_PER_C0); + AscendC::DataCopy( + dstTensor[i * layoutDst.shape(0)], srcTensor[i * layoutSrc.stride(1)], dataCopyParams); + } + } else { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.nValue = layoutSrc.shape(1); + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcDValue = layoutSrc.stride(1); + + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; // Outer stride -- col + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } + } +}; + +//////////////////////////// CopyGmToL1GMMPTD(Ascend950, No TLA) //////////////////////////// +/// Partial specialization for Ascend950, RowMajor in and zN out. +template +struct CopyGmToL1GMMPTD> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Mehtods + + CATLASS_DEVICE + CopyGmToL1GMMPTD() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.shape(0) == 1) { + // If the number of matrix rows is 1, the regular interval-based DataCopy interface can be used instead of + // the ND2NZ DataCopy interface, resulting in higher transfer efficiency. + AscendC::DataCopyParams dataCopyParams( + CeilDiv(layoutSrc.shape(1), layoutDst.shape(2)), + layoutDst.shape(2) / ELE_NUM_PER_C0, + 0, + (layoutDst.stride(3) - layoutDst.shape(2)) / ELE_NUM_PER_C0); + AscendC::DataCopy(dstTensor, srcTensor, dataCopyParams); + } else { + if (layoutSrc.shape(1) != ELE_NUM_PER_C0 || layoutSrc.stride(0) != ELE_NUM_PER_C0) { + intriParams.nValue = layoutSrc.shape(0); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + // If the matrix has ELE_NUM_PER_C0 columns and a stride of ELE_NUM_PER_C0, it follows a row-major + // layout in L1, allowing the use of the standard contiguous DataCopy interface for more efficient + // transfers. + AscendC::DataCopy(dstTensor, srcTensor, layoutSrc.shape(0) * layoutSrc.shape(1)); + } + } + } + + // layoutSrc must be the layout of one of the src matrices + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc, + uint32_t ndNum, uint32_t srcNdMatrixStride, + uint32_t dstNzNStride, uint32_t dstNzMatrixStride, + uint32_t dstNzC0Stride) + { + AscendC::Nd2NzParams intriParams; + + intriParams.nValue = layoutSrc.shape(0); + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = dstNzNStride; + intriParams.dstNzC0Stride = dstNzC0Stride; + intriParams.ndNum = ndNum; + intriParams.srcNdMatrixStride = srcNdMatrixStride; + intriParams.dstNzMatrixStride = dstNzMatrixStride; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l0c_to_dst.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l0c_to_dst.hpp new file mode 100644 index 0000000000..53d5adcc3d --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l0c_to_dst.hpp @@ -0,0 +1,274 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ASCEND950_COPY_L0C_TO_DST_HPP +#define CATLASS_KDA_ASCEND950_COPY_L0C_TO_DST_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Tile { + +enum class ScaleGranularity { + UNDEFINED = -1, + NO_QUANT = 0, + PER_TENSOR, + PER_CHANNEL, + PER_GROUP +}; + +template < + class ArchTag, + class ElementSrc, + class ElementDst, + ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT> +struct CopyL0CToDstQuantMode { + static_assert(DEPENDENT_FALSE, "Unsupported copy l0c to dst quant mode, can not find the specialization."); +}; + +// CopyL0CToDst fp32 to fp32 +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + float, float, + ScaleGranularity::NO_QUANT> { + static constexpr auto VALUE = QuantMode_t::NoQuant; +}; + +// CopyL0CToDst cast fp32 to fp16 +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + float, half, + ScaleGranularity::NO_QUANT> { + static constexpr auto VALUE = QuantMode_t::F322F16; +}; + +// CopyL0CToDst cast fp32 to bf16 +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + float, bfloat16_t, + ScaleGranularity::NO_QUANT> { + static constexpr auto VALUE = QuantMode_t::F322BF16; +}; + +// CopyL0CToDst cast float to uint8/int8 +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + float, uint8_t, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::QF322B8_PRE; +}; + +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + float, uint8_t, + ScaleGranularity::PER_CHANNEL> { + static constexpr auto VALUE = QuantMode_t::VQF322B8_PRE; +}; + +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + float, int8_t, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::QF322B8_PRE; +}; + +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + float, int8_t, + ScaleGranularity::PER_CHANNEL> { + static constexpr auto VALUE = QuantMode_t::VQF322B8_PRE; +}; + +// CopyL0CToDst output int32 +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + int32_t, int32_t, + ScaleGranularity::NO_QUANT> { + static constexpr auto VALUE = QuantMode_t::NoQuant; +}; + +// CopyL0CToDst cast int32_t to fp16 +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + int32_t, half, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::DEQF16; +}; + +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + int32_t, half, + ScaleGranularity::PER_CHANNEL> { + static constexpr auto VALUE = QuantMode_t::VDEQF16; +}; + +// CopyL0CToDst cast int32 to uint8/int8 +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + int32_t, uint8_t, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::REQ8; +}; + +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + int32_t, uint8_t, + ScaleGranularity::PER_CHANNEL> { + static constexpr auto VALUE = QuantMode_t::VREQ8; +}; + +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + int32_t, int8_t, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::REQ8; +}; + +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + int32_t, int8_t, + ScaleGranularity::PER_CHANNEL> { + static constexpr auto VALUE = QuantMode_t::VREQ8; +}; + +// CopyL0CToDst cast int32 to bf16 +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + int32_t, bfloat16_t, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::QS322BF16_PRE; +}; + +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + int32_t, bfloat16_t, + ScaleGranularity::PER_CHANNEL> { + static constexpr auto VALUE = QuantMode_t::VQS322BF16_PRE; +}; + +// CopyL0CToDst cast fp32 to fp16 +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + float, half, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::QF322F16_PRE; +}; + +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + float, half, + ScaleGranularity::PER_CHANNEL> { + static constexpr auto VALUE = QuantMode_t::VQF322F16_PRE; +}; + +// CopyL0CToDst cast fp32 to bf16 +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + float, bfloat16_t, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::QF322BF16_PRE; +}; + +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + float, bfloat16_t, + ScaleGranularity::PER_CHANNEL> { + static constexpr auto VALUE = QuantMode_t::VQF322BF16_PRE; +}; + +// CopyL0CToDst cast fp32 to fp32 +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + float, float, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::QF322F32_PRE; +}; + +template <> +struct CopyL0CToDstQuantMode< + Arch::Ascend950, + float, float, + ScaleGranularity::PER_CHANNEL> { + static constexpr auto VALUE = QuantMode_t::VQF322F32_PRE; +}; + +template < + class ArchTag, + class ElementAccumulator, + class GmType, + ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT, + bool ReluEnable = false> +struct CopyL0CToGm { + static_assert(DEPENDENT_FALSE, "Unsupported copy l0c to gm, can not find the specialization."); +}; + +///////////////////////////////////////////CopyL0CToGmTla///////////////////////////////////////////////// +// L0C copy mode +struct CopyToGM {}; +struct CopyToL1 {}; + +template < + class ArchTag, + class TensorSrc, + class TensorDst, + ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT, + bool ReluEnable = false, + class Enable = void> +struct CopyL0CToGmTla { + static_assert(DEPENDENT_FALSE, "Unsupported copy l0c to gm, can not find the specialization."); +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +enum class CopyL0CToUBMode { + NO_SPLIT = 0, + SPLIT_M, + SPLIT_N, + RESERVED +}; + +template < + class ArchTag, + class TensorSrc, + class TensorDst, + CopyL0CToUBMode CopyMode = CopyL0CToUBMode::NO_SPLIT, + ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT, + bool ReluEnable = false, + class Enable = void> +struct CopyL0CToUBTla { + static_assert(DEPENDENT_FALSE, "Unsupported copy l0c to ub, can not find the specialization."); +}; + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_ASCEND950_COPY_L0C_TO_DST_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l0c_to_gm.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l0c_to_gm.hpp new file mode 100644 index 0000000000..86266e6dd3 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l0c_to_gm.hpp @@ -0,0 +1,533 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ASCEND950_COPY_L0C_TO_GM_HPP +#define CATLASS_KDA_ASCEND950_COPY_L0C_TO_GM_HPP + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/gemm/tile/ascend950/kda_ascend950_copy_l0c_to_dst.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Tile { + +template +struct CopyL0CToGmTla< + Catlass::Arch::Ascend950, + TensorSrc_, + tla::Tensor, LayoutDst_, CoordDst_, AscendC::TPosition::GM>, + ScaleGranularity::NO_QUANT, + ReluEnable_, + std::enable_if_t::value>> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = typename TensorSrc_::Element; + static constexpr auto quantPre = + CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t unitFlag = 0) + { + static_assert( + tla::detail::isRowMajor::value && TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::GM, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be GM and RowMajor"); + + AscendC::DataCopyCO12DstParams intriParams; + + intriParams.nSize = tla::get<1>(dstTensor.originShape()); + intriParams.mSize = tla::get<0>(dstTensor.originShape()); + intriParams.dstStride = tla::get<0>(dstTensor.stride()); + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + intriParams.quantPre = quantPre; + intriParams.nz2ndEn = true; + intriParams.reluPre = reluEn; + intriParams.unitFlag = unitFlag; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::SetFixpipeNz2ndFlag(1, 1, 1); + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } + + template + CATLASS_DEVICE void + operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint32_t l0Batch, uint32_t dstNdStride) + { + static_assert( + tla::detail::isRowMajor::value && TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::GM, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be GM and RowMajor"); + + const uint32_t L0CM = tla::get<0, 0>(srcTensor.shape()) * tla::get<0, 1>(srcTensor.shape()); + const uint32_t L0CN = tla::get<1, 0>(srcTensor.shape()) * tla::get<1, 1>(srcTensor.shape()); + + // AscendC::FixpipeParamsC310 intriParams; + + // // Fixpipe layout information + // intriParams.nSize = tla::get<1>(dstTensor.shape()); + // intriParams.mSize = tla::get<0>(dstTensor.shape()); + // intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + // intriParams.dstStride = tla::get<0>(dstTensor.stride()); + + // // Fixpipe auxiliary arguments + // intriParams.quantPre = quantPre; + // intriParams.reluEn = reluEn; + // intriParams.unitFlag = 0; + + // intriParams.params.ndNum = l0Batch; + // intriParams.params.srcNdStride = L0CM * L0CN / tla::get<1, 0>(srcTensor.shape()); + // intriParams.params.dstNdStride = dstNdStride; + + // // Call AscendC Fixpipe + // AscendC::Fixpipe( + // dstTensor.data(), srcTensor.data(), intriParams); + + AscendC::DataCopyCO12DstParams intriParams; + + intriParams.nSize = tla::get<1>(dstTensor.originShape()); + intriParams.mSize = tla::get<0>(dstTensor.originShape()); + intriParams.dstStride = tla::get<0>(dstTensor.stride()); + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + intriParams.quantPre = quantPre; + intriParams.nz2ndEn = true; + intriParams.reluPre = reluEn; + intriParams.unitFlag = 0; + + AscendC::SetFixpipeNz2ndFlag(l0Batch, L0CM * L0CN / tla::get<1, 0>(srcTensor.shape()), dstNdStride); + AscendC::DataCopy(dstTensor.data(), srcTensor.data(), intriParams); + } +}; + +template +struct CopyL0CToGmTla< + Catlass::Arch::Ascend950, + TensorSrc_, + tla::Tensor, LayoutDst_, CoordDst_, AscendC::TPosition::GM>, + ScaleGranularity::NO_QUANT, + ReluEnable_, + std::enable_if_t::value>> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = typename TensorSrc_::Element; + static constexpr auto quantPre = + CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t unitFlag = 0) + { + static_assert( + tla::detail::iszN::value && TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::GM, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be GM and zN"); + + AscendC::DataCopyCO12DstParams intriParams; + + intriParams.nSize = tla::get<1>(dstTensor.originShape()); + intriParams.mSize = tla::get<0>(dstTensor.originShape()); + intriParams.dstStride = tla::get<1, 1>(dstTensor.stride()) / (BYTE_PER_C0 / sizeof(ElementDst)); + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + intriParams.quantPre = quantPre; + intriParams.nz2ndEn = false; + intriParams.reluPre = reluEn; + intriParams.unitFlag = unitFlag; + + if constexpr (std::is_same_v && std::is_same_v) { + intriParams.channelSplit = true; + } + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +template +struct CopyL0CToGmTla< + Catlass::Arch::Ascend950, + TensorSrc_, + tla::Tensor, LayoutDst_, CoordDst_, AscendC::TPosition::GM>, + ScaleGranularity::PER_TENSOR, + ReluEnable_, + std::enable_if_t::value>> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = typename TensorSrc_::Element; + static constexpr auto quantPre = + CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + struct Params { + float scale = 1.0; + + CATLASS_HOST_DEVICE + Params() = default; + + CATLASS_HOST_DEVICE + Params(float scalar) + { + scale = scalar; + } + }; + Params params; + + CATLASS_DEVICE + CopyL0CToGmTla() = default; + + CATLASS_DEVICE + CopyL0CToGmTla(Params const ¶ms_) + : params(params_) {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t unitFlag = 0) + { + static_assert( + tla::detail::isRowMajor::value && TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::GM, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be GM and RowMajor"); + + AscendC::FixpipeParamsC310 intriParams; + + // Fixpipe layout information + intriParams.nSize = tla::get<1>(dstTensor.shape()); + intriParams.mSize = tla::get<0>(dstTensor.shape()); + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + intriParams.dstStride = tla::get<0>(dstTensor.stride()); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.deqScalar = static_cast(*reinterpret_cast(¶ms.scale)); + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + // Call AscendC Fixpipe + AscendC::Fixpipe( + dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +template +struct CopyL0CToGmTla< + Catlass::Arch::Ascend950, + TensorSrc_, + tla::Tensor, LayoutDst_, CoordDst_, AscendC::TPosition::GM>, + ScaleGranularity::PER_CHANNEL, + ReluEnable_, + std::enable_if_t::value>> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = typename TensorSrc_::Element; + static constexpr auto quantPre = + CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, TensorQuant const &quantTensor, uint8_t unitFlag = 0) + { + static_assert( + tla::detail::isRowMajor::value && TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::GM, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be GM and RowMajor"); + + AscendC::FixpipeParamsC310 intriParams; + + // Fixpipe layout information + intriParams.nSize = tla::get<1>(dstTensor.shape()); + intriParams.mSize = tla::get<0>(dstTensor.shape()); + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + intriParams.dstStride = tla::get<0>(dstTensor.stride()); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + auto quantOffset = quantTensor.layout()(quantTensor.coord()); + + // Call AscendC Fixpipe + AscendC::Fixpipe( + dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], quantTensor.data()[quantOffset], intriParams); + } +}; + +////////////////////////////////////CopyL0CToGm(No-TLA, Ascend950)//////////////////////////////////////////////// +template < + class ElementAccumulator_, + class ElementDst_, + bool ReluEnable_> +struct CopyL0CToGm, + ScaleGranularity::NO_QUANT, + ReluEnable_> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = ElementAccumulator_; + using LayoutSrc = Catlass::layout::zN; + using LayoutDst = Catlass::layout::RowMajor; + static constexpr auto quantPre = CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + struct Params {}; + Params params; + + CATLASS_DEVICE + CopyL0CToGm() = default; + + CATLASS_DEVICE + CopyL0CToGm(Params const ¶ms_) + : params(params_) {}; + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor const &dst, AscendC::LocalTensor const &src, + LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0) + { + AscendC::DataCopyCO12DstParams intriParams; + + // Fixpipe layout information + intriParams.nSize = dstLayout.shape(1); + intriParams.mSize = dstLayout.shape(0); + intriParams.srcStride = srcLayout.stride(3) / srcLayout.stride(0); + intriParams.dstStride = dstLayout.stride(0); + intriParams.quantPre = quantPre; + intriParams.nz2ndEn = true; + intriParams.reluPre = reluEn; + intriParams.unitFlag = unitFlag; + + AscendC::SetFixpipeNz2ndFlag(1, 1, 1); + AscendC::DataCopy(dst, src, intriParams); + } +}; + +template < + class ElementAccumulator_, + class ElementDst_, + bool ReluEnable_> +struct CopyL0CToGm, + ScaleGranularity::PER_TENSOR, + ReluEnable_> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = ElementAccumulator_; + using LayoutSrc = Catlass::layout::zN; + using LayoutDst = Catlass::layout::RowMajor; + static constexpr auto quantPre = CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + struct Params { + float scale = 1.0; + + CATLASS_HOST_DEVICE + Params() = default; + + CATLASS_HOST_DEVICE + Params(float scalar) + { + scale = scalar; + } + }; + Params params; + + CATLASS_DEVICE + CopyL0CToGm() = default; + + CATLASS_DEVICE + CopyL0CToGm(Params const ¶ms_) + : params(params_) {}; + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor const &dst, AscendC::LocalTensor const &src, + LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0) + { + AscendC::FixpipeParamsC310 intriParams; + + // Fixpipe layout information + intriParams.nSize = dstLayout.shape(1); + intriParams.mSize = dstLayout.shape(0); + intriParams.srcStride = srcLayout.stride(3) / srcLayout.stride(0); + intriParams.dstStride = dstLayout.stride(0); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.deqScalar = static_cast(*reinterpret_cast(¶ms.scale)); + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + + // Call AscendC Fixpipe + AscendC::Fixpipe(dst, src, intriParams); + } +}; + +template < + class ElementAccumulator_, + class ElementDst_, + bool ReluEnable_> +struct CopyL0CToGm, + ScaleGranularity::PER_CHANNEL, + ReluEnable_> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = ElementAccumulator_; + using LayoutSrc = Catlass::layout::zN; + using LayoutDst = Catlass::layout::RowMajor; + static constexpr auto quantPre = CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + struct Params {}; + Params params; + + CATLASS_DEVICE + CopyL0CToGm() = default; + + CATLASS_DEVICE + CopyL0CToGm(Params const ¶ms_) + : params(params_) {}; + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor const &dst, AscendC::LocalTensor const &src, AscendC::LocalTensor const &scale, + LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0) + { + AscendC::FixpipeParamsC310 intriParams; + + // Fixpipe layout information + intriParams.nSize = dstLayout.shape(1); + intriParams.mSize = dstLayout.shape(0); + intriParams.srcStride = srcLayout.stride(3) / srcLayout.stride(0); + intriParams.dstStride = dstLayout.stride(0); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + + // Call AscendC Fixpipe + AscendC::Fixpipe(dst, src, scale, intriParams); + } +}; + +template < + class ElementAccumulator_, + class ElementDst_, + bool ReluEnable_> +struct CopyL0CToGm, + ScaleGranularity::NO_QUANT, + ReluEnable_> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = ElementAccumulator_; + using LayoutSrc = Catlass::layout::zN; + using LayoutDst = Catlass::layout::zN; + static constexpr auto quantPre = CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + struct Params {}; + Params params; + + CATLASS_DEVICE + CopyL0CToGm() = default; + + CATLASS_DEVICE + CopyL0CToGm(Params const ¶ms_) + : params(params_) {}; + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor const &dst, AscendC::LocalTensor const &src, + LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0) + { + AscendC::DataCopyCO12DstParams intriParams; + + // Fixpipe layout information + intriParams.nSize = dstLayout.shape(2) * dstLayout.shape(3); + intriParams.mSize = dstLayout.shape(0) * dstLayout.shape(1); + intriParams.srcStride = srcLayout.stride(3) / srcLayout.shape(2); + intriParams.dstStride = dstLayout.stride(3) / (BYTE_PER_C0 / sizeof(ElementDst)); + intriParams.quantPre = quantPre; + intriParams.nz2ndEn = false; + intriParams.reluPre = reluEn; + intriParams.unitFlag = unitFlag; + + if constexpr (std::is_same_v && std::is_same_v) { + intriParams.channelSplit = true; + } + + AscendC::DataCopy(dst, src, intriParams); + } +}; + +template < + class ElementAccumulator_, + class ElementDst_, + bool ReluEnable_> +struct CopyL0CToGm, + ScaleGranularity::NO_QUANT, + ReluEnable_> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = ElementAccumulator_; + using LayoutSrc = Catlass::layout::zN; + using LayoutDst = Catlass::layout::NDC1HWC0; + static constexpr auto quantPre = CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + struct Params {}; + Params params; + + CATLASS_DEVICE + CopyL0CToGm() = default; + + CATLASS_DEVICE + CopyL0CToGm(Params const ¶ms_) + : params(params_) {}; + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor const &dst, AscendC::LocalTensor const &src, + LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0) + { + AscendC::FixpipeParamsV220 intriParams; + + intriParams.nSize = srcLayout.orgShape(1); + intriParams.mSize = srcLayout.orgShape(0); + intriParams.srcStride = srcLayout.stride(3) / srcLayout.shape(2); + intriParams.dstStride = dstLayout.shape(1) * dstLayout.shape(2); + + if constexpr (AscendC::IsSameType::value && + AscendC::IsSameType::value) { + intriParams.isChannelSplit = true; + } + + intriParams.quantPre = quantPre; + intriParams.reluEn = false; + intriParams.unitFlag = unitFlag; + AscendC::Fixpipe(dst, src, intriParams); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_ASCEND950_COPY_L0C_TO_GM_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l1_to_bt.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l1_to_bt.hpp new file mode 100644 index 0000000000..97d8004b48 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l1_to_bt.hpp @@ -0,0 +1,107 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ASCEND950_COPY_L1_TO_BT_HPP +#define CATLASS_KDA_ASCEND950_COPY_L1_TO_BT_HPP + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/gemm/tile/kda_tile_copy_tla.hpp" +#include "catlass/kda_numeric_size.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Tile { + +/// Partial specialization for CopyL1ToBT, Ascend950, VectorLayout in and VectorLayout out. +template +struct TileCopyTla< + Arch::Ascend950, + tla::Tensor, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::C2>, + std::enable_if_t::value && tla::detail::isVector::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert( + tla::detail::isVector::value && tla::detail::isVector::value && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::C2, + "The input parameters do not match. TensorSrc must be L1 and Vector, " + "while TensorDst must be BT and Vector"); + + AscendC::DataCopyParams intriParams; + intriParams.blockCount = 1; + intriParams.blockLen = CeilDiv(tla::get<0>(srcTensor.originShape()), ELE_NUM_PER_C0); + if (sizeof(ElementSrc) == 4) { + // the burst length should be even when B32 + intriParams.blockLen = RoundUp(intriParams.blockLen, 2); + } + intriParams.srcStride = 0; + intriParams.dstStride = 0; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +////////////////////////////////////TileCopy(No-TLA, Ascend950)//////////////////////////////////////////////// + +template < + class ArchTag, + class L1Type, + class L0Type = void> +struct CopyL1ToBT { + static_assert(DEPENDENT_FALSE, "Unsupported copy l1 to biasTable buffer, can not find the specialization."); +}; + +template +struct CopyL1ToBT, + Catlass::Gemm::GemmType> { + using LayoutDst = layout::VectorLayout; + using LayoutSrc = layout::VectorLayout; + + static constexpr uint32_t ELE_NUM_PER_C2 = BYTE_PER_C2 / sizeof(ElementSrc); + + CATLASS_DEVICE + CopyL1ToBT() {} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc) + { + AscendC::DataCopyParams intriParams; + intriParams.blockCount = 1; + intriParams.blockLen = (layoutDst.shape(0) + ELE_NUM_PER_C2 - 1) / ELE_NUM_PER_C2; + if (sizeof(ElementSrc) == 4) { + // the burst length should be even when B32 + intriParams.blockLen = RoundUp(intriParams.blockLen, 2); + } + intriParams.srcStride = 0; + intriParams.dstStride = 0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_ASCEND950_COPY_L1_TO_BT_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l1_to_l0a.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l1_to_l0a.hpp new file mode 100644 index 0000000000..bb2424a8cd --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l1_to_l0a.hpp @@ -0,0 +1,467 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ASCEND950_COPY_L1_TO_L0A_HPP +#define CATLASS_KDA_ASCEND950_COPY_L1_TO_L0A_HPP + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/kda_numeric_size.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "catlass/gemm/tile/kda_tile_copy_tla.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Tile { + +/// Partial specialization for CopyL1ToL0A, Ascend950, zN in and zN out. +template +struct TileCopyTla< + Arch::Ascend950, + tla::Tensor, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A2>, + std::enable_if_t::value && tla::detail::iszN::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert( + tla::detail::iszN::value && tla::detail::iszN::value && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::A2, + "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0A and zN"); + + const uint32_t dstOuterShapeRow = CeilDiv(tla::get<0>(dstTensor.originShape()), tla::get<0, 0>(dstTensor.shape())); + const uint32_t dstOuterShapeCol = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape())); + const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride()); + auto srcCoord = srcTensor.coord(); + + AscendC::LoadData2DParamsV2 loadDataParams; + loadDataParams.mStartPosition = CeilDiv(tla::get<0>(srcCoord)); + loadDataParams.kStartPosition = CeilDiv(tla::get<1>(srcCoord)); + loadDataParams.mStep = dstOuterShapeRow; + loadDataParams.kStep = dstOuterShapeCol; + loadDataParams.srcStride = CeilDiv(srcOuterStrideCol); + loadDataParams.dstStride = CeilDiv(dstOuterStrideCol); + loadDataParams.ifTranspose = false; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data(), loadDataParams); + } + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint32_t l0Batch) + { + static_assert( + tla::detail::iszN::value && tla::detail::iszN::value && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::A2, + "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0A and zN"); + + const uint32_t dstOuterShapeRow = CeilDiv(tla::get<0>(dstTensor.originShape()), tla::get<0, 0>(dstTensor.shape())); + const uint32_t dstOuterShapeCol = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape())); + const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride()); + + AscendC::LoadData2DParamsV2 loadDataParams; + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = dstOuterShapeRow; + loadDataParams.kStep = dstOuterShapeCol * l0Batch; + loadDataParams.srcStride = CeilDiv(srcOuterStrideCol); + loadDataParams.dstStride = CeilDiv(dstOuterStrideCol); + loadDataParams.ifTranspose = false; + + AscendC::LoadData(dstTensor.data(), srcTensor.data(), loadDataParams); + } +}; + +/// Partial specialization for CopyL1ToL0A, Ascend950, not B8, nZ in and zN out. (Transpose A) +template +struct TileCopyTla< + Arch::Ascend950, + tla::Tensor, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A2>, + std::enable_if_t< + !AscendC::Std::is_one_of_v && + !AscendC::Std::is_one_of_v && + tla::detail::isnZ::value && tla::detail::iszN::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert( + !AscendC::Std::is_one_of_v && + !AscendC::Std::is_one_of_v && + tla::detail::isnZ::value && tla::detail::iszN::value && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::A2, + "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0A and zN"); + + const uint32_t L0M = tla::get<0>(dstTensor.originShape()); + const uint32_t L0K = tla::get<1>(dstTensor.originShape()); + const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride()); + auto srcCoord = srcTensor.coord(); + + AscendC::LoadData2DParamsV2 loadDataParams; + loadDataParams.mStartPosition = CeilDiv(tla::get<1>(srcCoord)); + loadDataParams.kStartPosition = CeilDiv(tla::get<0>(srcCoord)); + loadDataParams.mStep = CeilDiv(L0K); + loadDataParams.kStep = CeilDiv(L0M); + if constexpr (AscendC::Std::is_one_of_v) { + loadDataParams.kStep = RoundUp<2>(loadDataParams.kStep); // for b32 data types, ensure kStep is even + } + loadDataParams.srcStride = CeilDiv(srcOuterStrideRow); + loadDataParams.dstStride = CeilDiv(dstOuterStrideCol); + loadDataParams.ifTranspose = true; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data(), loadDataParams); + } + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint32_t l0Batch) + { + static_assert( + !AscendC::Std::is_one_of_v && + !AscendC::Std::is_one_of_v && + tla::detail::isnZ::value && tla::detail::iszN::value && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::A2, + "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0A and zN"); + + const uint32_t L1M = tla::get<0, 0>(srcTensor.shape()) * tla::get<0, 1>(srcTensor.shape()); + const uint32_t L1K = tla::get<1, 0>(srcTensor.shape()) * tla::get<1, 1>(srcTensor.shape()); + const uint32_t L0M = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape()); + const uint32_t L0K = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape()); + const uint32_t L0MOrigin = tla::get<0>(dstTensor.originShape()); + const uint32_t L0KOrigin = tla::get<1>(dstTensor.originShape()); + const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride()); + + AscendC::LoadData2DParamsV2 loadDataParams; + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = CeilDiv(L0KOrigin); + loadDataParams.kStep = CeilDiv(L0MOrigin); + if constexpr (AscendC::Std::is_one_of_v) { + loadDataParams.kStep = RoundUp<2>(loadDataParams.kStep); // for b32 data types, ensure kStep is even + } + loadDataParams.srcStride = CeilDiv(srcOuterStrideRow); + loadDataParams.dstStride = CeilDiv(dstOuterStrideCol); + loadDataParams.ifTranspose = true; + + for (uint32_t l0BatchIdx = 0; l0BatchIdx < l0Batch; l0BatchIdx++) { + AscendC::LoadData( + dstTensor.data()[l0BatchIdx * L0M * L0K], srcTensor.data()[l0BatchIdx * L1M * L1K], loadDataParams); + } + } +}; + +/// Partial specialization for CopyL1ToL0A, Ascend950, B8, nZ in and zN out. (Transpose A) +template +struct TileCopyTla< + Arch::Ascend950, + tla::Tensor, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A2>, + std::enable_if_t< + AscendC::Std::is_one_of_v && + AscendC::Std::is_one_of_v && + tla::detail::isnZ::value && tla::detail::iszN::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert( + AscendC::Std::is_one_of_v && + AscendC::Std::is_one_of_v && + tla::detail::isnZ::value && tla::detail::iszN::value && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::A2, + "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0A and zN"); + + const uint32_t L0MPadded = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape()); + const uint32_t L0MOrigin = tla::get<0>(dstTensor.originShape()); + const uint32_t L0KOrigin = tla::get<1>(dstTensor.originShape()); + const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride()); + auto srcCoord = srcTensor.coord(); + + AscendC::LoadData2DParamsV2 loadDataParams; + if (RoundUp(L0MOrigin) % ELE_NUM_PER_C0 == 0) { + loadDataParams.mStartPosition = CeilDiv(tla::get<1>(srcCoord)); + loadDataParams.kStartPosition = CeilDiv(tla::get<0>(srcCoord)); + loadDataParams.mStep = RoundUp<2>(CeilDiv(L0KOrigin)); + loadDataParams.kStep = CeilDiv(L0MOrigin); + loadDataParams.srcStride = CeilDiv(srcOuterStrideRow); + loadDataParams.dstStride = CeilDiv(dstOuterStrideCol); + loadDataParams.ifTranspose = true; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data(), loadDataParams); + } else { + for (uint32_t kIdx = 0; kIdx < CeilDiv(L0KOrigin); kIdx++) { + loadDataParams.mStartPosition = CeilDiv(tla::get<1>(srcCoord)) + kIdx * 2; + loadDataParams.kStartPosition = CeilDiv(tla::get<0>(srcCoord)); + loadDataParams.mStep = 2; + loadDataParams.kStep = CeilDiv(L0MOrigin); + loadDataParams.srcStride = CeilDiv(srcOuterStrideRow); + loadDataParams.dstStride = CeilDiv(dstOuterStrideCol); + loadDataParams.ifTranspose = true; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + AscendC::LoadData( + dstTensor.data()[dstOffset + kIdx * L0MPadded * ELE_NUM_PER_C0], srcTensor.data(), loadDataParams); + } + } + } + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint32_t l0Batch) + { + static_assert( + AscendC::Std::is_one_of_v && + AscendC::Std::is_one_of_v && + tla::detail::isnZ::value && tla::detail::iszN::value && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::A2, + "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0A and zN"); + + const uint32_t L1M = tla::get<0, 0>(srcTensor.shape()) * tla::get<0, 1>(srcTensor.shape()); + const uint32_t L1K = tla::get<1, 0>(srcTensor.shape()) * tla::get<1, 1>(srcTensor.shape()); + const uint32_t L0M = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape()); + const uint32_t L0K = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape()); + const uint32_t L0MPadded = L0M; + const uint32_t L0KPadded = L0K; + const uint32_t L0MOrigin = tla::get<0>(dstTensor.originShape()); + const uint32_t L0KOrigin = tla::get<1>(dstTensor.originShape()); + const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride()); + + AscendC::LoadData2DParamsV2 loadDataParams; + if (RoundUp(L0MOrigin) % ELE_NUM_PER_C0 == 0) { + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = RoundUp<2>(CeilDiv(L0KOrigin)); + loadDataParams.kStep = CeilDiv(L0MOrigin); + loadDataParams.srcStride = CeilDiv(srcOuterStrideRow); + loadDataParams.dstStride = CeilDiv(dstOuterStrideCol); + loadDataParams.ifTranspose = true; + + for (uint32_t l0BatchIdx = 0; l0BatchIdx < l0Batch; l0BatchIdx++) { + AscendC::LoadData( + dstTensor.data()[l0BatchIdx * L0MPadded * L0KPadded], srcTensor.data()[l0BatchIdx * L1M * L1K], loadDataParams); + } + } else { + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = 2; + loadDataParams.kStep = CeilDiv(L0MOrigin); + loadDataParams.srcStride = CeilDiv(srcOuterStrideRow); + loadDataParams.dstStride = CeilDiv(dstOuterStrideCol); + loadDataParams.ifTranspose = true; + for (uint32_t l0BatchIdx = 0; l0BatchIdx < l0Batch; l0BatchIdx++) { + for (uint32_t kIdx = 0; kIdx < CeilDiv(L0KOrigin); kIdx++) { + AscendC::LoadData( + dstTensor.data()[l0BatchIdx * L0MPadded * L0KPadded + kIdx * L0MPadded * ELE_NUM_PER_C0], + srcTensor.data()[l0BatchIdx * L1M * L1K + kIdx * ELE_NUM_PER_FRACTAL * 2], loadDataParams); + } + } + } + } +}; + +template +struct TileCopyTla< + Arch::Ascend950, + tla::Tensor, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A2>, + std::enable_if_t::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + uint16_t aL1M = tla::get<0, 0>(srcTensor.stride()); + uint16_t madM = tla::get<1, 1>(dstTensor.stride()); + uint16_t madK = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape())); + + AscendC::LoadData2DParamsV2 loadDataParams; + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = CeilDiv(madM); + loadDataParams.kStep = CeilDiv(madK); + loadDataParams.srcStride = CeilDiv(aL1M); + loadDataParams.dstStride = CeilDiv(madM); + + loadDataParams.ifTranspose = false; + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + AscendC::LoadData(dstTensor.data(), srcTensor.data()[srcOffset], loadDataParams); + } +}; + +////////////////////////////////////CopyL1ToL0A(No-TLA, Ascend950)//////////////////////////////////////////////// +template < + class ArchTag, + class L1Type, + class L0Type = void, + class Enable = void> +struct CopyL1ToL0A { + static_assert(DEPENDENT_FALSE, "Unsupported copy l1 to l0, can not find the specialization."); +}; + +/// Partial specialization for CopyL1ToL0A, AtlasA5, zN in and zN out. +template +struct CopyL1ToL0A> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0A() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2DParamsV2 loadDataParams; + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = layoutDst.shape(1); + loadDataParams.kStep = layoutDst.shape(3); + loadDataParams.srcStride = CeilDiv(layoutSrc.stride(3)); + loadDataParams.dstStride = CeilDiv(layoutDst.stride(3)); + loadDataParams.ifTranspose = false; + + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + } +}; + +/// Partial specialization for CopyL1ToL0A, AtlasA5, not B8 or B4, nZ in and zN out. (Transpose A) +template +struct CopyL1ToL0A, void, + std::enable_if_t< + !AscendC::Std::is_one_of_v>> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0A() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + const uint32_t L0M = layoutDst.shape(0) * layoutDst.shape(1); + const uint32_t L0K = layoutDst.shape(2) * layoutDst.shape(3); + const uint32_t srcOuterStrideRow = layoutSrc.stride(1); + const uint32_t dstOuterStrideCol = layoutDst.stride(3); + + AscendC::LoadData2DParamsV2 loadDataParams; + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = CeilDiv(L0K); + loadDataParams.kStep = CeilDiv(L0M); + loadDataParams.srcStride = CeilDiv(srcOuterStrideRow); + loadDataParams.dstStride = CeilDiv(dstOuterStrideCol); + loadDataParams.ifTranspose = true; + + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + } +}; + +/// Partial specialization for CopyL1ToL0A, AtlasA5, B8 or B4, nZ in and zN out. (Transpose A) +template +struct CopyL1ToL0A, void, + std::enable_if_t< + AscendC::Std::is_one_of_v>> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0A() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + const uint32_t L0M = layoutDst.shape(0) * layoutDst.shape(1); + const uint32_t L0K = layoutDst.shape(2) * layoutDst.shape(3); + const uint32_t srcOuterStrideRow = layoutSrc.stride(1); + const uint32_t dstOuterStrideCol = layoutDst.stride(3); + + AscendC::LoadData2DParamsV2 loadDataParams; + if (L0M % ELE_NUM_PER_C0 == 0) { + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = CeilDiv(L0K); + loadDataParams.kStep = CeilDiv(L0M); + loadDataParams.srcStride = CeilDiv(srcOuterStrideRow); + loadDataParams.dstStride = CeilDiv(dstOuterStrideCol); + loadDataParams.ifTranspose = true; + + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + } else { + for (uint32_t kIdx = 0; kIdx < L0K / ELE_NUM_PER_C0; kIdx++) { + loadDataParams.mStartPosition = kIdx * 2; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = 2; + loadDataParams.kStep = CeilDiv(L0M); + loadDataParams.srcStride = CeilDiv(srcOuterStrideRow); + loadDataParams.dstStride = CeilDiv(dstOuterStrideCol); + loadDataParams.ifTranspose = true; + + AscendC::LoadData( + dstTensor[kIdx * L0M * ELE_NUM_PER_C0], srcTensor, loadDataParams); + } + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_ASCEND950_COPY_L1_TO_L0A_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l1_to_l0b.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l1_to_l0b.hpp new file mode 100644 index 0000000000..d4a59a1439 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/ascend950/kda_ascend950_copy_l1_to_l0b.hpp @@ -0,0 +1,437 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ASCEND950_COPY_L1_TO_L0B_HPP +#define CATLASS_KDA_ASCEND950_COPY_L1_TO_L0B_HPP + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/kda_numeric_size.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "catlass/gemm/tile/kda_tile_copy_tla.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Tile { + +/// Partial specialization for CopyL1ToL0B, Ascend950, not B8, zN in and nZ out. +template +struct TileCopyTla< + Arch::Ascend950, + tla::Tensor, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::B2>, + std::enable_if_t< + !AscendC::Std::is_one_of_v && + !AscendC::Std::is_one_of_v && + tla::detail::iszN::value && tla::detail::isnZ::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert( + !AscendC::Std::is_one_of_v && + !AscendC::Std::is_one_of_v && + tla::detail::iszN::value && tla::detail::isnZ::value && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::B2, + "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0B and nZ"); + + const uint32_t L0KOrigin = tla::get<0>(dstTensor.originShape()); + const uint32_t L0NOrigin = tla::get<1>(dstTensor.originShape()); + const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + auto srcCoord = srcTensor.coord(); + + AscendC::LoadData2DParamsV2 loadDataParams; + loadDataParams.mStartPosition = CeilDiv(tla::get<0>(srcCoord)); + loadDataParams.kStartPosition = CeilDiv(tla::get<1>(srcCoord)); + loadDataParams.mStep = CeilDiv(L0KOrigin); + loadDataParams.kStep = CeilDiv(L0NOrigin); + if constexpr (AscendC::Std::is_one_of_v) { + loadDataParams.kStep = RoundUp<2>(loadDataParams.kStep); // for b32 data types, ensure kStep is even + } + loadDataParams.srcStride = CeilDiv(srcOuterStrideCol); + loadDataParams.dstStride = CeilDiv(dstOuterStrideRow); + loadDataParams.ifTranspose = true; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data(), loadDataParams); + } + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint32_t l0Batch) + { + static_assert( + !AscendC::Std::is_one_of_v && + !AscendC::Std::is_one_of_v && + tla::detail::iszN::value && tla::detail::isnZ::value && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::B2, + "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0B and nZ"); + + const uint32_t L1K = tla::get<0, 0>(srcTensor.shape()) * tla::get<0, 1>(srcTensor.shape()); + const uint32_t L1N = tla::get<1, 0>(srcTensor.shape()) * tla::get<1, 1>(srcTensor.shape()); + const uint32_t L0K = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape()); + const uint32_t L0N = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape()); + const uint32_t L0KOrigin = tla::get<0>(dstTensor.originShape()); + const uint32_t L0NOrigin = tla::get<1>(dstTensor.originShape()); + const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + + AscendC::LoadData2DParamsV2 loadDataParams; + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = CeilDiv(L0KOrigin); + loadDataParams.kStep = CeilDiv(L0NOrigin); + if constexpr (AscendC::Std::is_one_of_v) { + loadDataParams.kStep = RoundUp<2>(loadDataParams.kStep); // for b32 data types, ensure kStep is even + } + loadDataParams.srcStride = CeilDiv(srcOuterStrideCol); + loadDataParams.dstStride = CeilDiv(dstOuterStrideRow); + loadDataParams.ifTranspose = true; + + for (uint32_t l0BatchIdx = 0; l0BatchIdx < l0Batch; l0BatchIdx++) { + AscendC::LoadData( + dstTensor.data()[l0BatchIdx * L0N * L0K], srcTensor.data()[l0BatchIdx * L1N * L1K], loadDataParams); + } + } +}; + +/// Partial specialization for CopyL1ToL0B, Ascend950, B8, zN in and nZ out. +template +struct TileCopyTla< + Arch::Ascend950, + tla::Tensor, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::B2>, + std::enable_if_t< + AscendC::Std::is_one_of_v && + AscendC::Std::is_one_of_v && + tla::detail::iszN::value && tla::detail::isnZ::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert( + AscendC::Std::is_one_of_v && + AscendC::Std::is_one_of_v && + tla::detail::iszN::value && tla::detail::isnZ::value && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::B2, + "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0B and nZ"); + + const uint32_t L0NPadded = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape()); + const uint32_t L0KOrigin = tla::get<0>(dstTensor.originShape()); + const uint32_t L0NOrigin = tla::get<1>(dstTensor.originShape()); + const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + auto srcCoord = srcTensor.coord(); + + AscendC::LoadData2DParamsV2 loadDataParams; + if (RoundUp(L0NOrigin) % ELE_NUM_PER_C0 == 0) { + loadDataParams.mStartPosition = CeilDiv(tla::get<0>(srcCoord)); + loadDataParams.kStartPosition = CeilDiv(tla::get<1>(srcCoord)); + loadDataParams.mStep = RoundUp<2>(CeilDiv(L0KOrigin)); + loadDataParams.kStep = CeilDiv(L0NOrigin); + loadDataParams.srcStride = CeilDiv(srcOuterStrideCol); + loadDataParams.dstStride = CeilDiv(dstOuterStrideRow); + loadDataParams.ifTranspose = true; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data(), loadDataParams); + } else { + for (uint32_t kIdx = 0; kIdx < CeilDiv(L0KOrigin); kIdx++) { + loadDataParams.mStartPosition = CeilDiv(tla::get<0>(srcCoord)) + kIdx * 2; + loadDataParams.kStartPosition = CeilDiv(tla::get<1>(srcCoord)); + loadDataParams.mStep = 2; + loadDataParams.kStep = CeilDiv(L0NOrigin); + loadDataParams.srcStride = CeilDiv(srcOuterStrideCol); + loadDataParams.dstStride = CeilDiv(dstOuterStrideRow); + loadDataParams.ifTranspose = true; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + AscendC::LoadData( + dstTensor.data()[dstOffset + kIdx * L0NPadded * ELE_NUM_PER_C0], srcTensor.data(), loadDataParams); + } + } + } + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint32_t l0Batch) + { + static_assert( + AscendC::Std::is_one_of_v && + AscendC::Std::is_one_of_v && + tla::detail::iszN::value && tla::detail::isnZ::value && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::B2, + "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0B and nZ"); + + const uint32_t L1K = tla::get<0, 0>(srcTensor.shape()) * tla::get<0, 1>(srcTensor.shape()); + const uint32_t L1N = tla::get<1, 0>(srcTensor.shape()) * tla::get<1, 1>(srcTensor.shape()); + const uint32_t L0K = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape()); + const uint32_t L0N = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape()); + const uint32_t L0KOrigin = tla::get<0>(dstTensor.originShape()); + const uint32_t L0NOrigin = tla::get<1>(dstTensor.originShape()); + const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + + AscendC::LoadData2DParamsV2 loadDataParams; + if (RoundUp(L0NOrigin) % ELE_NUM_PER_C0 == 0) { + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = RoundUp<2>(CeilDiv(L0KOrigin)); + loadDataParams.kStep = CeilDiv(L0NOrigin); + loadDataParams.srcStride = CeilDiv(srcOuterStrideCol); + loadDataParams.dstStride = CeilDiv(dstOuterStrideRow); + loadDataParams.ifTranspose = true; + + for (uint32_t l0BatchIdx = 0; l0BatchIdx < l0Batch; l0BatchIdx++) { + AscendC::LoadData( + dstTensor.data()[l0BatchIdx * L0N * L0K], srcTensor.data()[l0BatchIdx * L1N * L1K], loadDataParams); + } + } else { + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = 2; + loadDataParams.kStep = CeilDiv(L0NOrigin); + loadDataParams.srcStride = CeilDiv(srcOuterStrideCol); + loadDataParams.dstStride = CeilDiv(dstOuterStrideRow); + loadDataParams.ifTranspose = true; + for (uint32_t l0BatchIdx = 0; l0BatchIdx < l0Batch; l0BatchIdx++) { + for (uint32_t kIdx = 0; kIdx < CeilDiv(L0KOrigin); kIdx++) { + AscendC::LoadData( + dstTensor.data()[l0BatchIdx * L0N * L0K + kIdx * L0N * ELE_NUM_PER_C0], + srcTensor.data()[l0BatchIdx * L1N * L1K + kIdx * ELE_NUM_PER_FRACTAL * 2], loadDataParams); + } + } + } + } +}; + +/// Partial specialization for CopyL1ToL0B, Ascend950, nZ in and nZ out. (Transpose B) +template +struct TileCopyTla< + Arch::Ascend950, + tla::Tensor, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::B2>, + std::enable_if_t::value && tla::detail::isnZ::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert( + tla::detail::isnZ::value && tla::detail::isnZ::value && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::B2, + "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0B and nZ"); + + const uint32_t dstOuterShapeRow = CeilDiv(tla::get<0>(dstTensor.originShape()), tla::get<0, 0>(dstTensor.shape())); + const uint32_t dstOuterShapeCol = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape())); + const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + auto srcCoord = srcTensor.coord(); + + AscendC::LoadData2DParamsV2 loadDataParams; + loadDataParams.mStartPosition = CeilDiv(tla::get<1>(srcCoord)); + loadDataParams.kStartPosition = CeilDiv(tla::get<0>(srcCoord)); + loadDataParams.mStep = dstOuterShapeCol; + loadDataParams.kStep = dstOuterShapeRow; + loadDataParams.srcStride = CeilDiv(srcOuterStrideRow); + loadDataParams.dstStride = CeilDiv(dstOuterStrideRow); + loadDataParams.ifTranspose = false; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data(), loadDataParams); + } + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint32_t l0Batch) + { + static_assert( + tla::detail::isnZ::value && tla::detail::isnZ::value && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::B2, + "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0B and nZ"); + + const uint32_t dstOuterShapeRow = CeilDiv(tla::get<0>(dstTensor.originShape()), tla::get<0, 0>(dstTensor.shape())); + const uint32_t dstOuterShapeCol = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape())); + const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + + AscendC::LoadData2DParamsV2 loadDataParams; + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = dstOuterShapeCol; + loadDataParams.kStep = dstOuterShapeRow * l0Batch; + loadDataParams.srcStride = CeilDiv(srcOuterStrideRow); + loadDataParams.dstStride = CeilDiv(dstOuterStrideRow); + loadDataParams.ifTranspose = false; + + AscendC::LoadData(dstTensor.data(), srcTensor.data(), loadDataParams); + } +}; + +////////////////////////////////////CopyL1ToL0B(No-TLA, Ascend950)//////////////////////////////////////////////// +template < + class ArchTag, + class L1Type, + class L0Type = void, + class Enable = void> +struct CopyL1ToL0B { + static_assert(DEPENDENT_FALSE, "Unsupported copy l1 to l0, can not find the specialization."); +}; + +/// Partial specialization for CopyL1ToL0B, Ascend950, nZ in and nZ out. (Transpose B) +template +struct CopyL1ToL0B> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + const uint32_t dstOuterShapeRow = layoutDst.shape(1); + const uint32_t dstOuterShapeCol = layoutDst.shape(3); + const uint32_t srcOuterStrideRow = layoutSrc.stride(1); + const uint32_t dstOuterStrideRow = layoutDst.stride(1); + + AscendC::LoadData2DParamsV2 loadDataParams; + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = dstOuterShapeCol; + loadDataParams.kStep = dstOuterShapeRow; + loadDataParams.srcStride = CeilDiv(srcOuterStrideRow); + loadDataParams.dstStride = CeilDiv(dstOuterStrideRow); + loadDataParams.ifTranspose = false; + + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + } +}; + +/// Partial specialization for CopyL1ToL0B, AtlasA5, not B8 or B4, zN in and nZ out. +template +struct CopyL1ToL0B, void, + std::enable_if_t< + !AscendC::Std::is_one_of_v>> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + const uint32_t L0K = layoutDst.shape(0) * layoutDst.shape(1); + const uint32_t L0N = layoutDst.shape(2) * layoutDst.shape(3); + const uint32_t srcOuterStrideCol = layoutSrc.stride(3); + const uint32_t dstOuterStrideRow = layoutDst.stride(1); + + AscendC::LoadData2DParamsV2 loadDataParams; + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = CeilDiv(L0K); + loadDataParams.kStep = CeilDiv(L0N); + loadDataParams.srcStride = CeilDiv(srcOuterStrideCol); + loadDataParams.dstStride = CeilDiv(dstOuterStrideRow); + loadDataParams.ifTranspose = true; + + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + } +}; + +/// Partial specialization for CopyL1ToL0B, AtlasA5, B8 or B4, zN in and nZ out. +template +struct CopyL1ToL0B, void, + std::enable_if_t< + AscendC::Std::is_one_of_v>> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + const uint32_t L0K = layoutDst.shape(0) * layoutDst.shape(1); + const uint32_t L0N = layoutDst.shape(2) * layoutDst.shape(3); + const uint32_t srcOuterStrideCol = layoutSrc.stride(3); + const uint32_t dstOuterStrideRow = layoutDst.stride(1); + + AscendC::LoadData2DParamsV2 loadDataParams; + if (L0N % ELE_NUM_PER_C0 == 0) { + loadDataParams.mStartPosition = 0; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = CeilDiv(L0K); + loadDataParams.kStep = CeilDiv(L0N); + loadDataParams.srcStride = CeilDiv(srcOuterStrideCol); + loadDataParams.dstStride = CeilDiv(dstOuterStrideRow); + loadDataParams.ifTranspose = true; + + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + } else { + for (uint32_t kIdx = 0; kIdx < L0K / ELE_NUM_PER_C0; kIdx++) { + loadDataParams.mStartPosition = kIdx * 2; + loadDataParams.kStartPosition = 0; + loadDataParams.mStep = 2; + loadDataParams.kStep = CeilDiv(L0N); + loadDataParams.srcStride = CeilDiv(srcOuterStrideCol); + loadDataParams.dstStride = CeilDiv(dstOuterStrideRow); + loadDataParams.ifTranspose = true; + + AscendC::LoadData( + dstTensor[kIdx * L0N * ELE_NUM_PER_C0], srcTensor, loadDataParams); + } + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_ASCEND950_COPY_L1_TO_L0B_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_cast_int4_to_int8.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_cast_int4_to_int8.hpp new file mode 100644 index 0000000000..b11b3644fb --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_cast_int4_to_int8.hpp @@ -0,0 +1,180 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ATLASA2_CAST_INT4_TO_INT8_HPP +#define CATLASS_KDA_ATLASA2_CAST_INT4_TO_INT8_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" + +namespace Catlass::Gemm::Tile { + +template < + class ArchTag, + class SrcType_, + class DstType_, + // Length of the compute elements + uint32_t COMPUTE_LEN_, + uint32_t STAGES = 2> +struct TileCastInt4ToInt8 { + using ElementSrc = typename SrcType_::Element; + using ElementDst = typename DstType_::Element; + using LayoutSrc = typename SrcType_::Layout; + using LayoutDst = typename DstType_::Layout; + + static constexpr uint32_t ELE_NUM_PER_BLK_INT8 = BYTE_PER_BLK / sizeof(ElementDst); + static constexpr uint32_t ELE_NUM_PER_BLK_INT4 = BYTE_PER_BLK / sizeof(ElementDst) * 2; + + static_assert( + sizeof(ElementSrc) == sizeof(ElementDst), + "Error: Src and Dst element sizes are equal (forbidden)"); + + static_assert( + std::is_same_v || + std::is_same_v, + "Unsupported layout, only can be RowMajor, ColumnMajor or RowMajorInt4 or ColumnMajorInt4"); + + static constexpr uint32_t COMPUTE_LEN = COMPUTE_LEN_; + static_assert(COMPUTE_LEN <= 24 * 1024, "COMPUTE_LEN cannot exceed 24 * 1024"); + + struct Params {}; + + /// Construct + CATLASS_DEVICE + TileCastInt4ToInt8(Arch::Resource const &resource, Params const ¶ms) + { + if constexpr (g_coreType == AscendC::AIV) { + uint32_t ubOffset = 0; + uint32_t ubInSize = COMPUTE_LEN * sizeof(int8_t) / 2; + uint32_t ubOutSize = COMPUTE_LEN * sizeof(ElementDst); + uint32_t ubWorkspaceSize = COMPUTE_LEN * sizeof(half); + // Init buffers + for (uint32_t i = 0; i < STAGES; i++) { + ubInTensorList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += ubInSize; + ubOutTensorList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += ubOutSize; + ubWorkspaceList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += ubWorkspaceSize; + + ubEventList[i] = i; + AscendC::SetFlag(ubEventList[i]); + AscendC::SetFlag(ubEventList[i]); + } + } + } + + /// Destructor + CATLASS_DEVICE + ~TileCastInt4ToInt8() + { + if constexpr (g_coreType == AscendC::AIV) { + for (uint32_t i = 0; i < STAGES; i++) { + AscendC::WaitFlag(ubEventList[i]); + AscendC::WaitFlag(ubEventList[i]); + } + } + } + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &gmDst, LayoutDst const &layoutDst, + AscendC::GlobalTensor const &gmSrc, LayoutSrc const &layoutSrc) + { + uint32_t tilesNum = layoutSrc.shape(0); + uint32_t tileLen = layoutSrc.shape(1); + uint32_t tileLenRoundInt8 = RoundUp(layoutSrc.shape(1), ELE_NUM_PER_BLK_INT8); + uint32_t tileLenRoundInt4 = RoundUp(layoutSrc.shape(1), ELE_NUM_PER_BLK_INT4); + uint64_t tileStrideSrc = layoutSrc.stride(0); + uint64_t tileStrideDst = layoutDst.stride(0); + if constexpr (std::is_same_v) { + tilesNum = layoutSrc.shape(1); + tileLen = layoutSrc.shape(0); + tileLenRoundInt8 = RoundUp(layoutSrc.shape(0), ELE_NUM_PER_BLK_INT8); + tileLenRoundInt4 = RoundUp(layoutSrc.shape(0), ELE_NUM_PER_BLK_INT4); + tileStrideSrc = layoutSrc.stride(1); + tileStrideDst = layoutDst.stride(1); + } + uint32_t tilesPerAiv = tilesNum / AscendC::GetSubBlockNum(); + if (AscendC::GetSubBlockIdx() < (tilesNum % AscendC::GetSubBlockNum())) { + tilesPerAiv++; + } + uint64_t taskOffsetSrc = AscendC::GetSubBlockIdx() * tilesPerAiv * tileStrideSrc; + uint64_t taskOffsetDst = AscendC::GetSubBlockIdx() * tilesPerAiv * tileStrideDst; + if (AscendC::GetSubBlockIdx() >= (tilesNum % AscendC::GetSubBlockNum())) { + taskOffsetSrc += (tilesNum % AscendC::GetSubBlockNum()) * tileStrideSrc; + taskOffsetDst += (tilesNum % AscendC::GetSubBlockNum()) * tileStrideDst; + } + uint32_t tilesPerLoop = 32; + uint32_t loops = CeilDiv(tilesPerAiv, tilesPerLoop); + uint32_t pingpong = 0; + for (uint32_t loopIdx = 0; loopIdx < loops; loopIdx++) { + uint32_t actualTiles = tilesPerLoop; + if (loopIdx == loops - 1) { + actualTiles = tilesPerAiv - loopIdx * tilesPerLoop; + } + uint64_t tileOffsetSrc = loopIdx * tilesPerLoop * tileStrideSrc; + AscendC::DataCopyExtParams dataCopyParamsIn( + actualTiles, (tileLen + 1) / 2 * sizeof(int8_t), + (tileStrideSrc - tileLen) * sizeof(int8_t) / 2, + (tileLenRoundInt4 - tileLen) / ELE_NUM_PER_BLK_INT4, 0); + AscendC::DataCopyPadExtParams padParams(false, 0, 0, 0); + + AscendC::WaitFlag(ubEventList[pingpong]); + AscendC::DataCopyPad(ubInTensorList[pingpong], + gmSrc[taskOffsetSrc + tileOffsetSrc], dataCopyParamsIn, padParams); + + AscendC::SetFlag(ubEventList[pingpong]); + AscendC::WaitFlag(ubEventList[pingpong]); + + AscendC::WaitFlag(ubEventList[pingpong]); + + AscendC::Cast(ubWorkspaceList[pingpong], ubInTensorList[pingpong], + AscendC::RoundMode::CAST_NONE, actualTiles * tileLenRoundInt4); + + AscendC::PipeBarrier(); + + AscendC::SetFlag(ubEventList[pingpong]); + + AscendC::Cast(ubOutTensorList[pingpong], ubWorkspaceList[pingpong], + AscendC::RoundMode::CAST_NONE, actualTiles * tileLenRoundInt4); + + AscendC::SetFlag(ubEventList[pingpong]); + AscendC::WaitFlag(ubEventList[pingpong]); + + uint64_t tileOffsetDst = loopIdx * tilesPerLoop * tileStrideDst; + AscendC::DataCopyExtParams dataCopyParamsOut( + actualTiles, tileLen * sizeof(ElementDst), + (tileLenRoundInt4 - tileLen) / ELE_NUM_PER_BLK_INT8, + (tileStrideDst - tileLen) * sizeof(ElementDst), 0); + AscendC::DataCopyPad(gmDst[taskOffsetDst + tileOffsetDst], + ubOutTensorList[pingpong], dataCopyParamsOut); + AscendC::SetFlag(ubEventList[pingpong]); + + pingpong = (pingpong + 1) % STAGES; + } + } + +protected: + /// Data members + AscendC::LocalTensor ubInTensorList[STAGES]; + AscendC::LocalTensor ubOutTensorList[STAGES]; + AscendC::LocalTensor ubWorkspaceList[STAGES]; + + int32_t ubEventList[STAGES]; +}; + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_ATLASA2_CAST_INT4_TO_INT8_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_gm_to_l1.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_gm_to_l1.hpp new file mode 100644 index 0000000000..cb35cafa6b --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_gm_to_l1.hpp @@ -0,0 +1,2194 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ATLASA2_COPY_GM_TO_L1_HPP +#define CATLASS_KDA_ATLASA2_COPY_GM_TO_L1_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/kda_numeric_size.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "catlass/gemm/tile/kda_tile_copy_tla.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Tile { + +template < + class ArchTag, + /// GemmType for matrix operand + class GmType, + class L1Type = void> +struct CopyGmToL1 { + static_assert(DEPENDENT_FALSE, "Unsupported copy gm to l1, can not find the specialization."); +}; + +template < + class ArchTag, + /// GemmType for matrix operand + class GmType, + class L1Type = void> +struct CopyGmToL1IntervalDataCopy { + static_assert(DEPENDENT_FALSE, "Unsupported copy gm to l1, can not find the specialization."); +}; + +template < + class ArchTag, + /// GemmType for matrix operand + class GmType, + class L1Type = void> +struct CopyGmToL1GMMPTD { + static_assert(DEPENDENT_FALSE, "Unsupported copy gm to l1, can not find the specialization."); +}; + +template < + class ArchTag, + /// GemmType for matrix operand + class GmType, + class L1Type = void> +struct CopyGmToL1DynamicOptimized { + static_assert(DEPENDENT_FALSE, "Unsupported copy gm to l1, can not find the specialization."); +}; + +/// Partial specialization for AtlasA2, RowMajor in and zN out. +template +struct CopyGmToL1DynamicOptimized> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1DynamicOptimized() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.shape(0) <= 16) { + // If the number of matrix rows is 1, the regular interval-based DataCopy interface can be used instead of + // the ND2NZ DataCopy interface, resulting in higher transfer efficiency. + for (int i = 0; i < layoutSrc.shape(0); ++i) { + AscendC::DataCopyParams dataCopyParams( + CeilDiv(layoutSrc.shape(1), layoutDst.shape(2)), + layoutDst.shape(2) / ELE_NUM_PER_C0, + 0, + (layoutDst.stride(3) - layoutDst.shape(2)) / ELE_NUM_PER_C0); + AscendC::DataCopy( + dstTensor[i * layoutDst.shape(2)], srcTensor[i * layoutSrc.stride(0)], dataCopyParams); + } + } else { + if (layoutSrc.stride(0) < STRIDE_LIMIT) { + if (layoutSrc.shape(1) != ELE_NUM_PER_C0 || layoutSrc.stride(0) != ELE_NUM_PER_C0) { + intriParams.nValue = layoutSrc.shape(0); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + // If the matrix has ELE_NUM_PER_C0, columns and a stride of ELE_NUM_PER_C0, it follows a row-major + // layout in L1, allowing the use of standard contiguous DataCopy interface for more efficient + // transfers. + AscendC::DataCopy(dstTensor, srcTensor, layoutSrc.shape(0) * layoutSrc.shape(1)); + } + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(0); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(0)], intriParams); + } + } + } + } +}; + +/// Partial specialization for AtlasA2, ColumnMajor in and nZ out. +template +struct CopyGmToL1DynamicOptimized> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::ColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1DynamicOptimized() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.shape(1) <= 16) { + // If the number of matrix cols is 1, the regular interval-based DataCopy interface can be used instead of + // the ND2NZ DataCopy interface, resulting in higher transfer efficiency. + for (int i = 0; i < layoutSrc.shape(1); ++i) { + AscendC::DataCopyParams dataCopyParams( + CeilDiv(layoutSrc.shape(0), layoutDst.shape(0)), + layoutDst.shape(0) / ELE_NUM_PER_C0, + 0, + (layoutDst.stride(1) - layoutDst.shape(0)) / ELE_NUM_PER_C0); + AscendC::DataCopy( + dstTensor[i * layoutDst.shape(0)], srcTensor[i * layoutSrc.stride(1)], dataCopyParams); + } + } else { + if (layoutSrc.stride(1) < STRIDE_LIMIT) { + if (layoutSrc.shape(0) != ELE_NUM_PER_C0 || layoutSrc.stride(1) != ELE_NUM_PER_C0) { + intriParams.nValue = layoutSrc.shape(1); + intriParams.srcDValue = layoutSrc.stride(1); + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + // If the matrix has ELE_NUM_PER_C0, rows and a stride of ELE_NUM_PER_C0, it follows a col-major + // layout in L1, allowing the use of standard contiguous DataCopy interface for more efficient + // transfers. + AscendC::DataCopy(dstTensor, srcTensor, layoutSrc.shape(0) * layoutSrc.shape(1)); + } + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(1); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(1)], intriParams); + } + } + } + } +}; + +/// Partial specialization for AtlasA2, zN in and zN out. +template +struct CopyGmToL1DynamicOptimized> : public CopyGmToL1> {}; + +/// Partial specialization for AtlasA2, nZ in and nZ out. +template +struct CopyGmToL1DynamicOptimized> : public CopyGmToL1> {}; + +/// Partial specialization for AtlasA2, PaddingRowMajor in and zN out. +template +struct CopyGmToL1DynamicOptimized> : public CopyGmToL1> {}; + +/// Partial specialization for AtlasA2, PaddingColumnMajor in and nZ out. +template +struct CopyGmToL1DynamicOptimized> : public CopyGmToL1> {}; + +/// Partial specialization for AtlasA2, RowMajor in and zN out. +template +struct CopyGmToL1GMMPTD> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1GMMPTD() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.shape(0) == 1) { + // If the number of matrix rows is 1, the regular interval-based DataCopy interface can be used instead of + // the ND2NZ DataCopy interface, resulting in higher transfer efficiency. + AscendC::DataCopyParams dataCopyParams( + CeilDiv(layoutSrc.shape(1), layoutDst.shape(2)), + layoutDst.shape(2) / ELE_NUM_PER_C0, + 0, + (layoutDst.stride(3) - layoutDst.shape(2)) / ELE_NUM_PER_C0); + AscendC::DataCopy(dstTensor, srcTensor, dataCopyParams); + } else { + if (layoutSrc.stride(0) < STRIDE_LIMIT) { + if (layoutSrc.shape(1) != ELE_NUM_PER_C0 || layoutSrc.stride(0) != ELE_NUM_PER_C0) { + intriParams.nValue = layoutSrc.shape(0); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + // If the matrix has ELE_NUM_PER_C0 columns and a stride of ELE_NUM_PER_C0, it follows a row-major + // layout in L1, allowing the use of the standard contiguous DataCopy interface for more efficient + // transfers. + AscendC::DataCopy(dstTensor, srcTensor, layoutSrc.shape(0) * layoutSrc.shape(1)); + } + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(0); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(0)], intriParams); + } + } + } + } + + // layoutSrc must be the layout of one of the src matrices + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc, + uint32_t ndNum, uint32_t srcNdMatrixStride, + uint32_t dstNzNStride, uint32_t dstNzMatrixStride, + uint32_t dstNzC0Stride) + { + AscendC::Nd2NzParams intriParams; + + intriParams.nValue = layoutSrc.shape(0); + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = dstNzNStride; + intriParams.dstNzC0Stride = dstNzC0Stride; + if (srcNdMatrixStride < STRIDE_LIMIT) { + intriParams.ndNum = ndNum; + intriParams.srcNdMatrixStride = srcNdMatrixStride; + intriParams.dstNzMatrixStride = dstNzMatrixStride; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.ndNum = 1; + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzMatrixStride = 0; + for (uint32_t i = 0; i < ndNum; i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * srcNdMatrixStride], intriParams); + } + } + } +}; + +//////////////////////////////////////// +/// Using the standard strided DataCopy interface to implement nd2nz +/// transfer may achieve higher data transfer efficiency when the data block shape is short and wide +/// Partial specialization for AtlasA2, half, RowMajor in and zN out. +template <> +struct CopyGmToL1IntervalDataCopy> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::RowMajor; + using Element = half; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1IntervalDataCopy() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + for (int i = 0; i < layoutSrc.shape(0); ++i) { + AscendC::DataCopyParams dataCopyParams( + CeilDiv(layoutSrc.shape(1), layoutDst.shape(2)), + layoutDst.shape(2) / ELE_NUM_PER_C0, + 0, + (layoutDst.stride(3) - layoutDst.shape(2)) / ELE_NUM_PER_C0); + AscendC::DataCopy(dstTensor[i * layoutDst.shape(2)], srcTensor[i * layoutSrc.stride(0)], dataCopyParams); + } + } +}; + +/// Partial specialization for AtlasA2, half, PaddingRowMajor in and zN out. +/// Using the standard strided DataCopy interface to implement nd2nz +/// transfer may achieve higher data transfer efficiency when the data block shape is short and wide +template <> +struct CopyGmToL1IntervalDataCopy> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::PaddingRowMajor; + using Element = half; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1IntervalDataCopy() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + for (int i = 0; i < layoutSrc.orgShape(0); ++i) { + AscendC::DataCopyParams dataCopyParams( + CeilDiv(layoutSrc.orgShape(1), layoutDst.shape(2)), + layoutDst.shape(2) / ELE_NUM_PER_C0, + 0, + (layoutDst.stride(3) - layoutDst.shape(2)) / ELE_NUM_PER_C0); + AscendC::DataCopy(dstTensor[i * layoutDst.shape(2)], srcTensor[i * layoutSrc.stride(0)], dataCopyParams); + } + } +}; + +/// Partial specialization for AtlasA2, half, ColumnMajor in and zN out. +/// Using the standard strided DataCopy interface to implement nd2nz +/// transfer may achieve higher data transfer efficiency when the data block shape is tall and narrow +template <> +struct CopyGmToL1IntervalDataCopy> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::ColumnMajor; + using Element = half; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1IntervalDataCopy() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + for (int i = 0; i < layoutSrc.shape(1); ++i) { + AscendC::DataCopyParams dataCopyParams( + CeilDiv(layoutSrc.shape(0), layoutDst.shape(0)), + layoutDst.shape(0) / ELE_NUM_PER_C0, + 0, + (layoutDst.stride(1) - layoutDst.shape(0)) / ELE_NUM_PER_C0); + AscendC::DataCopy(dstTensor[i * layoutDst.shape(0)], srcTensor[i * layoutSrc.stride(1)], dataCopyParams); + } + } +}; + +/// Partial specialization for AtlasA2, half, PaddingColumnMajor in and zN out. +/// Using the standard strided DataCopy interface to implement nd2nz +/// transfer may achieve higher data transfer efficiency when the data block shape is tall and narrow +template <> +struct CopyGmToL1IntervalDataCopy> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::PaddingColumnMajor; + using Element = half; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1IntervalDataCopy() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + for (int i = 0; i < layoutSrc.orgShape(1); ++i) { + AscendC::DataCopyParams dataCopyParams( + CeilDiv(layoutSrc.orgShape(0), layoutDst.shape(0)), + layoutDst.shape(0) / ELE_NUM_PER_C0, + 0, + (layoutDst.stride(1) - layoutDst.shape(0)) / ELE_NUM_PER_C0); + AscendC::DataCopy(dstTensor[i * layoutDst.shape(0)], srcTensor[i * layoutSrc.stride(2)], dataCopyParams); + } + } +}; + +/// new add gemm +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.stride(0) < STRIDE_LIMIT) { + intriParams.nValue = layoutSrc.shape(0); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(0); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(0)], intriParams); + } + } + } +}; + +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::zZ; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + uint32_t srcNdStride = C0_NUM_PER_FRACTAL * layoutSrc.stride(0); + uint32_t ndNum = layoutSrc.shape(0) / C0_NUM_PER_FRACTAL; + uint32_t remains = layoutSrc.shape(0) % C0_NUM_PER_FRACTAL; + if (srcNdStride < STRIDE_LIMIT) { + if (ndNum) { + intriParams.ndNum = ndNum; + intriParams.nValue = C0_NUM_PER_FRACTAL; + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcNdMatrixStride = srcNdStride; + intriParams.srcDValue = layoutSrc.stride(0); + + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + + intriParams.dstNzMatrixStride = layoutDst.stride(1); + + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } + + if (remains) { + AscendC::Nd2NzParams tailParams; + tailParams.ndNum = 1; + tailParams.nValue = remains; + tailParams.dValue = layoutSrc.shape(1); + tailParams.srcNdMatrixStride = srcNdStride; + tailParams.srcDValue = layoutSrc.stride(0); + + tailParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + tailParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + tailParams.dstNzMatrixStride = 0; //` + + AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(1)], srcTensor[ndNum * srcNdStride], tailParams); + } + } else if (layoutSrc.stride(0) < STRIDE_LIMIT) { + for (uint32_t i = 0; i < ndNum; i++) { + AscendC::Nd2NzParams intriParams; + intriParams.ndNum = 1; + intriParams.nValue = C0_NUM_PER_FRACTAL; + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.srcDValue = layoutSrc.stride(0); + + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[i * layoutDst.stride(1)], srcTensor[i * srcNdStride], intriParams); + } + if (remains) { + AscendC::Nd2NzParams tailParams; + tailParams.ndNum = 1; + tailParams.nValue = remains; + tailParams.dValue = layoutSrc.shape(1); + tailParams.srcNdMatrixStride = 0; + tailParams.srcDValue = layoutSrc.stride(0); + + tailParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + tailParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + tailParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(1)], srcTensor[ndNum * srcNdStride], tailParams); + } + } else { + for (uint32_t i = 0; i < layoutSrc.shape(0); i++) { + uint32_t idxR0 = i / C0_NUM_PER_FRACTAL; + uint32_t idxInR0 = i % C0_NUM_PER_FRACTAL; + + AscendC::Nd2NzParams intriParams; + intriParams.ndNum = 1; + intriParams.nValue = 1; + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.srcDValue = 0; + + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = 0; + intriParams.dstNzMatrixStride = 0; + + uint32_t offsetDst = i * idxR0 * layoutDst.stride(1) + idxInR0 * ELE_NUM_PER_C0; + uint32_t offsetSrc = i * layoutSrc.stride(0); + AscendC::DataCopy(dstTensor[offsetDst], srcTensor[offsetSrc], intriParams); + } + } + } +}; + +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::nN; + using LayoutSrc = layout::ColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + uint32_t srcNdStride = C0_NUM_PER_FRACTAL * layoutSrc.stride(1); + uint32_t ndNum = layoutSrc.shape(1) / C0_NUM_PER_FRACTAL; + uint32_t remains = layoutSrc.shape(1) % C0_NUM_PER_FRACTAL; + if (srcNdStride < STRIDE_LIMIT) { + if (ndNum) { + intriParams.ndNum = ndNum; + intriParams.nValue = C0_NUM_PER_FRACTAL; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = srcNdStride; + intriParams.srcDValue = layoutSrc.stride(1); + + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + + intriParams.dstNzMatrixStride = layoutDst.stride(3); + + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } + + if (remains) { + AscendC::Nd2NzParams tailParams; + tailParams.ndNum = 1; + tailParams.nValue = remains; + tailParams.dValue = layoutSrc.shape(0); + tailParams.srcNdMatrixStride = srcNdStride; + tailParams.srcDValue = layoutSrc.stride(1); + + tailParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + tailParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + tailParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(3)], srcTensor[ndNum * srcNdStride], tailParams); + } + } else if (layoutSrc.stride(1) < STRIDE_LIMIT) { + for (uint32_t i = 0; i < ndNum; i++) { + AscendC::Nd2NzParams intriParams; + intriParams.ndNum = 1; + intriParams.nValue = C0_NUM_PER_FRACTAL; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.srcDValue = layoutSrc.stride(1); + + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[i * layoutDst.stride(3)], srcTensor[i * srcNdStride], intriParams); + } + if (remains) { + AscendC::Nd2NzParams tailParams; + tailParams.ndNum = 1; + tailParams.nValue = remains; + tailParams.dValue = layoutSrc.shape(0); + tailParams.srcNdMatrixStride = 0; + tailParams.srcDValue = layoutSrc.stride(1); + + tailParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + tailParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + tailParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(3)], srcTensor[ndNum * srcNdStride], tailParams); + } + } else { + for (uint32_t i = 0; i < layoutSrc.shape(1); i++) { + uint32_t idxR0 = i / C0_NUM_PER_FRACTAL; + uint32_t idxInR0 = i % C0_NUM_PER_FRACTAL; + + AscendC::Nd2NzParams intriParams; + intriParams.ndNum = 1; + intriParams.nValue = 1; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.srcDValue = 0; + + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = 0; + intriParams.dstNzMatrixStride = 0; + + uint32_t offsetDst = i * idxR0 * layoutDst.stride(3) + idxInR0 * ELE_NUM_PER_C0; + uint32_t offsetSrc = i * layoutSrc.stride(1); + AscendC::DataCopy(dstTensor[offsetDst], srcTensor[offsetSrc], intriParams); + } + } + } +}; + +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::ColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.stride(1) < STRIDE_LIMIT) { + intriParams.nValue = layoutSrc.shape(1); + intriParams.srcDValue = layoutSrc.stride(1); + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(1); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(1)], intriParams); + } + } + } +}; + +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::ColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.stride(1) < STRIDE_LIMIT) { + intriParams.nValue = layoutSrc.shape(1); + intriParams.srcDValue = layoutSrc.stride(1); + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(1); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(1)], intriParams); + } + } + } +}; +//////////////////////////////////////// + +/////////////////////////////////////// +/// new add gemv, VectorLayout -> zN +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::VectorLayout; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + intriParams.nValue = 1; + intriParams.srcDValue = layoutSrc.shape(0); + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +template +struct CopyGmToL1> { + using LayoutDst = layout::NDC1HWC0; + using LayoutSrc = layout::NDC1HWC0; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + const static uint64_t MAX_UINT16 = 65535; + + uint32_t cin1LoadL1 = layoutDst.orgShape(2); + uint32_t hiLoadL1 = layoutDst.orgShape(3); + + uint32_t dilationD = layoutSrc.orgShape(1); + uint32_t OriC1 = layoutSrc.orgShape(2); + uint32_t OriH = layoutSrc.orgShape(3); + uint32_t OriW = layoutSrc.orgShape(4); + uint32_t OriK0 = layoutSrc.orgShape(5); + + uint64_t dataCopyLoop = CeilDiv(cin1LoadL1, OriC1); + uint64_t dataCopySubLoop = 0; + uint64_t blockCount = OriC1; + bool srcStrideBeyondMaxU16 = false; + if (OriH * OriW - hiLoadL1 * OriW > MAX_UINT16) { + dataCopySubLoop = dataCopyLoop > 0 ? cin1LoadL1 / dataCopyLoop : 0; + blockCount = 1; + srcStrideBeyondMaxU16 = true; + } + + uint64_t aL1GmOffset = 0; + uint64_t aL1Offset = 0; + if (cin1LoadL1 > OriC1 || srcStrideBeyondMaxU16) { + repeatParams.blockCount = blockCount; + repeatParams.blockLen = hiLoadL1 * OriW; + repeatParams.srcStride = OriW * OriH - repeatParams.blockLen; + repeatParams.dstStride = 0; + for (uint64_t i = 0; i < dataCopyLoop; i++) { + if (srcStrideBeyondMaxU16) { + uint64_t aL1GmSubOffset = aL1GmOffset; + uint64_t aL1SubOffset = aL1Offset; + for (uint64_t j = 0; j < dataCopySubLoop; j++) { + AscendC::DataCopy(dstTensor[aL1SubOffset], srcTensor[aL1GmSubOffset], repeatParams); + aL1GmSubOffset += OriH * OriW * OriK0; + aL1SubOffset += hiLoadL1 * OriW * OriK0; + } + } else { + AscendC::DataCopy(dstTensor[aL1Offset], srcTensor[aL1GmOffset], repeatParams); + } + aL1GmOffset += dilationD * OriC1 * OriH * OriW * OriK0; + aL1Offset += OriC1 * hiLoadL1 * OriW * OriK0; + } + } else { + repeatParams.blockCount = cin1LoadL1; + repeatParams.blockLen = hiLoadL1 * OriW; + repeatParams.srcStride = OriW * OriH - repeatParams.blockLen; + repeatParams.dstStride = 0; + AscendC::DataCopy(dstTensor[aL1Offset], srcTensor[aL1GmOffset], repeatParams); + aL1Offset += cin1LoadL1 * hiLoadL1 * OriW * OriK0; + } + } + +private: + uint64_t aL1GmOffset = 0; + AscendC::DataCopyParams repeatParams; +}; + +template +struct CopyGmToL1> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::KDC1KHKWN1N0C0; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + uint32_t currentNBL1 = layoutDst.orgShape(1); + uint32_t currentKBL1 = layoutDst.orgShape(0); + + uint32_t N1 = layoutSrc.shape(3); + uint32_t N0 = layoutSrc.shape(2); + uint32_t C0 = layoutSrc.shape(0); + uint32_t OriCoAlign = N1 * N0; + + const static uint32_t LOAD2D_MAX_REPEAT_TIMES = 255; + + if (currentNBL1 >= OriCoAlign) { + uint32_t repeatTimes = (currentKBL1 * currentNBL1) / (N0 * C0); + if (repeatTimes > LOAD2D_MAX_REPEAT_TIMES) { + repeatParams.blockCount = 1; + repeatParams.srcStride = 0; + repeatParams.blockLen = CeilDiv(currentKBL1 * currentNBL1, C0); + AscendC::DataCopy(dstTensor, srcTensor, repeatParams); + } else { + AscendC::LoadData2DParams loadData2dParams; + loadData2dParams.srcStride = 1; + loadData2dParams.repeatTimes = repeatTimes; + AscendC::LoadData(dstTensor, srcTensor, loadData2dParams); + } + } else { + repeatParams.blockCount = currentKBL1 / C0; + repeatParams.blockLen = currentNBL1; + repeatParams.srcStride = N1 * N0 - currentNBL1; + AscendC::DataCopy(dstTensor, srcTensor, repeatParams); + } + } + +private: + AscendC::DataCopyParams repeatParams; +}; + +/////////////////////////////////////// +/// new add gemv, ColumnMajor -> nN +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::nN; + using LayoutSrc = layout::ColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + uint32_t srcNdStride = C0_NUM_PER_FRACTAL * layoutSrc.stride(1); + uint32_t ndNum = layoutSrc.shape(1) / C0_NUM_PER_FRACTAL; + uint32_t remains = layoutSrc.shape(1) % C0_NUM_PER_FRACTAL; + if (srcNdStride < STRIDE_LIMIT) { + if (ndNum) { + intriParams.ndNum = ndNum; + intriParams.nValue = C0_NUM_PER_FRACTAL; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = srcNdStride; + intriParams.srcDValue = layoutSrc.stride(1); + + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + + intriParams.dstNzMatrixStride = layoutDst.stride(3); + + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } + + if (remains) { + AscendC::Nd2NzParams tailParams; + tailParams.ndNum = 1; + tailParams.nValue = remains; + tailParams.dValue = layoutSrc.shape(0); + tailParams.srcNdMatrixStride = srcNdStride; + tailParams.srcDValue = layoutSrc.stride(1); + + tailParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + tailParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + tailParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(3)], srcTensor[ndNum * srcNdStride], tailParams); + } + } else if (layoutSrc.stride(1) < STRIDE_LIMIT) { + for (uint32_t i = 0; i < ndNum; i++) { + AscendC::Nd2NzParams intriParams; + intriParams.ndNum = 1; + intriParams.nValue = C0_NUM_PER_FRACTAL; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.srcDValue = layoutSrc.stride(1); + + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[i * layoutDst.stride(3)], srcTensor[i * srcNdStride], intriParams); + } + if (remains) { + AscendC::Nd2NzParams tailParams; + tailParams.ndNum = 1; + tailParams.nValue = remains; + tailParams.dValue = layoutSrc.shape(0); + tailParams.srcNdMatrixStride = 0; + tailParams.srcDValue = layoutSrc.stride(1); + + tailParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + tailParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + tailParams.dstNzMatrixStride = 0; + + AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(3)], srcTensor[ndNum * srcNdStride], tailParams); + } + } else { + for (uint32_t i = 0; i < layoutSrc.shape(1); i++) { + uint32_t idxR0 = i / C0_NUM_PER_FRACTAL; + uint32_t idxInR0 = i % C0_NUM_PER_FRACTAL; + + AscendC::Nd2NzParams intriParams; + intriParams.ndNum = 1; + intriParams.nValue = 1; + intriParams.dValue = layoutSrc.shape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.srcDValue = 0; + + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzNStride = 0; + intriParams.dstNzMatrixStride = 0; + + uint32_t offsetDst = i * idxR0 * layoutDst.stride(3) + idxInR0 * ELE_NUM_PER_C0; + uint32_t offsetSrc = i * layoutSrc.stride(1); + AscendC::DataCopy(dstTensor[offsetDst], srcTensor[offsetSrc], intriParams); + } + } + } +}; + +template +struct CopyGmToL1, Gemm::GemmType> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.stride(0) < STRIDE_LIMIT) { + intriParams.nValue = layoutSrc.shape(0); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(0); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(0)], intriParams); + } + } + } +}; +///////////////////////////////// + +/// Partial specialization for AtlasA2, RowMajor in and zN out. +template +struct CopyGmToL1> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(1); + // There is an issue with AscenC in the data migration of nd2nz of type int4b_t: dValue needs to be passed + // with the length of int8, and it is uncertain whether this will be fixed in subsequent versions. + if constexpr (std::is_same_v) { + intriParams.dValue = CeilDiv(intriParams.dValue, 2); + } + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.stride(0) < STRIDE_LIMIT) { + intriParams.nValue = layoutSrc.shape(0); + intriParams.srcDValue = layoutSrc.stride(0); + if constexpr (std::is_same_v) { + intriParams.srcDValue = CeilDiv(intriParams.srcDValue, 2); + } + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(0); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(0)], intriParams); + } + } + } + + // layoutSrc must be the layout of one of the src matrices + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc, + uint32_t ndNum, uint32_t srcNdMatrixStride, + uint32_t dstNzNStride, uint32_t dstNzMatrixStride, + uint32_t dstNzC0Stride) + { + AscendC::Nd2NzParams intriParams; + + intriParams.nValue = layoutSrc.shape(0); + intriParams.dValue = layoutSrc.shape(1); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = dstNzNStride; + intriParams.dstNzC0Stride = dstNzC0Stride; + if (srcNdMatrixStride < STRIDE_LIMIT) { + intriParams.ndNum = ndNum; + intriParams.srcNdMatrixStride = srcNdMatrixStride; + intriParams.dstNzMatrixStride = dstNzMatrixStride; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.ndNum = 1; + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzMatrixStride = 0; + for (uint32_t i = 0; i < ndNum; i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * srcNdMatrixStride], intriParams); + } + } + } +}; +///////////////////////////////// + +/// Partial specialization for AtlasA2, ColumnMajor in and nZ out. +template < + class Element> +struct CopyGmToL1> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::ColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.shape(0); + // There is an issue with AscenC in the data migration of nd2nz of type int4b_t: dValue needs to be passed + // with the length of int8, and it is uncertain whether this will be fixed in subsequent versions. + if constexpr (std::is_same_v) { + intriParams.dValue = CeilDiv(intriParams.dValue, 2); + } + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + if (layoutSrc.stride(1) < STRIDE_LIMIT) { + intriParams.nValue = layoutSrc.shape(1); + intriParams.srcDValue = layoutSrc.stride(1); + if constexpr (std::is_same_v) { + intriParams.srcDValue = CeilDiv(intriParams.srcDValue, 2); + } + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < layoutSrc.shape(1); i++) { + AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(1)], intriParams); + } + } + } +}; + +/// Partial specialization for zN in and zN out. +template < + class ArchTag, + class Element> +struct CopyGmToL1> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + uint32_t blockCount = CeilDiv(layoutSrc.orgShape(1)); + uint32_t blockLen = layoutSrc.orgShape(0); + + AscendC::DataCopyParams repeatParams; + + if (layoutSrc.stride(3) / ELE_NUM_PER_C0 < STRIDE_LIMIT) { + repeatParams.blockCount = blockCount; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_C0 - blockLen; + repeatParams.dstStride = layoutDst.stride(3) / ELE_NUM_PER_C0 - blockLen; + AscendC::DataCopy(dstTensor, srcTensor, repeatParams); + } else { + repeatParams.blockCount = 1; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = 0; + repeatParams.dstStride = 0; + for (uint32_t i = 0; i < blockCount; i++) { + uint64_t dstOffset = i * layoutDst.stride(3); + uint64_t srcOffset = i * layoutSrc.stride(3); + AscendC::DataCopy(dstTensor[dstOffset], srcTensor[srcOffset], repeatParams); + } + } + } +}; + +/// Partial specialization for nZ in and nZ out. +template < + class ArchTag, + class Element> +struct CopyGmToL1> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + uint32_t blockCount = CeilDiv(layoutSrc.orgShape(0)); + uint32_t blockLen = layoutSrc.orgShape(1); + + AscendC::DataCopyParams repeatParams; + + if (layoutSrc.stride(1) / ELE_NUM_PER_C0 < STRIDE_LIMIT) { + repeatParams.blockCount = blockCount; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = layoutSrc.stride(1) / ELE_NUM_PER_C0 - blockLen; + repeatParams.dstStride = layoutDst.stride(1) / ELE_NUM_PER_C0 - blockLen; + AscendC::DataCopy(dstTensor, srcTensor, repeatParams); + } else { + repeatParams.blockCount = 1; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = 0; + repeatParams.dstStride = 0; + for (uint32_t i = 0; i < blockCount; i++) { + uint64_t dstOffset = i * layoutDst.stride(1); + uint64_t srcOffset = i * layoutSrc.stride(1); + AscendC::DataCopy(dstTensor[dstOffset], srcTensor[srcOffset], repeatParams); + } + } + } +}; + +/// Partial specialization for AtlasA2, PaddingRowMajor in and zN out. +template +struct CopyGmToL1> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::PaddingRowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.orgShape(1); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + intriParams.nValue = layoutSrc.orgShape(0); + intriParams.srcDValue = layoutSrc.stride(0); + intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +/// Partial specialization for AtlasA2, ColumnMajor in and nZ out. +template < + class Element> +struct CopyGmToL1> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::PaddingColumnMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = layoutSrc.orgShape(0); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + intriParams.nValue = layoutSrc.orgShape(1); + intriParams.srcDValue = layoutSrc.stride(2); + intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +/// Partial specialization for AtlasA2, RowMajor in and RowMajor out. +template +struct CopyGmToL1, + Gemm::GemmType> { + using LayoutDst = layout::RowMajor; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_BLK = BytesToBits(BYTE_PER_BLK) / SizeOfBits::value; + static constexpr uint32_t BLOCK_LEN_LIMIT = 65536; + static constexpr uint32_t MAX_REPEAT = 4095; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + uint32_t rows = layoutSrc.shape(0); + uint32_t cols = layoutSrc.shape(1); + uint32_t srcStride = (layoutSrc.stride(0) - layoutSrc.shape(1)) / ELE_NUM_PER_BLK; + uint32_t dstStride = (layoutDst.stride(0) - layoutDst.shape(1)) / ELE_NUM_PER_BLK; + + if ((layoutSrc.shape(1) == layoutSrc.stride(0)) && (layoutDst.shape(1) == layoutDst.stride(0))) { + DataCopy(dstTensor, srcTensor, rows * cols); + } else if (srcStride < STRIDE_LIMIT && dstStride < STRIDE_LIMIT && (cols / ELE_NUM_PER_BLK) < BLOCK_LEN_LIMIT) { + uint32_t rLoops = CeilDiv(rows, MAX_REPEAT); + for (uint32_t i = 0; i < rLoops; ++i) { + uint32_t rActual = (i < rLoops - 1) ? MAX_REPEAT : rows - i * MAX_REPEAT; + AscendC::DataCopyParams dataCopyParams( + rActual, cols / ELE_NUM_PER_BLK, srcStride, dstStride); + DataCopy(dstTensor[i * MAX_REPEAT * layoutDst.stride(0)], + srcTensor[i * MAX_REPEAT * layoutSrc.stride(0)], dataCopyParams); + } + } else { + for (uint32_t i = 0; i < rows; ++i) { + DataCopy(dstTensor[i * layoutDst.stride(0)], srcTensor[i * layoutSrc.stride(0)], cols); + } + } + } +}; + +///////////////////////////////////////////TileCopyTla////////////////////////////////////////////////////// +/// Partial specialization for CopyGmToL1, AtlasA2, RowMajor in and zN out. +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + std::enable_if_t::value && + tla::detail::iszN::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isRowMajor::value && + tla::detail::iszN::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and RowMajor, while TensorDst must be L1 and zN"); + + const uint32_t nValue = tla::get<0>(srcTensor.originShape()); + const uint32_t dValue = tla::get<1>(srcTensor.originShape()); + const uint32_t srcDValue = tla::get<0>(srcTensor.stride()); + const uint32_t dstInnerStrideRow = tla::get<0, 0>(dstTensor.stride()); + const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride()); + + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = dValue; + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = dstOuterStrideCol / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + if (srcDValue < STRIDE_LIMIT) { + intriParams.nValue = nValue; + intriParams.srcDValue = srcDValue; + intriParams.dstNzNStride = dstInnerStrideRow / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < nValue; i++) { + AscendC::DataCopy(dstTensor.data()[dstOffset + i * ELE_NUM_PER_C0], + srcTensor.data()[srcOffset + i * srcDValue], + intriParams); + } + } + } +}; + +/// Partial specialization for CopyGmToL1, AtlasA2, ColumnMajor in and nZ out. +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + std::enable_if_t::value && + tla::detail::isnZ::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isColumnMajor::value && + tla::detail::isnZ::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and ColumnMajor, " + "while TensorDst must be L1 and nZ"); + + const uint32_t nValue = tla::get<1>(srcTensor.originShape()); + const uint32_t dValue = tla::get<0>(srcTensor.originShape()); + const uint32_t srcDValue = tla::get<1>(srcTensor.stride()); + const uint32_t dstInnerStrideRow = tla::get<1, 0>(dstTensor.stride()); + const uint32_t dstOuterStrideCol = tla::get<0, 1>(dstTensor.stride()); + + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = dValue; + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = dstOuterStrideCol / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + if (srcDValue < STRIDE_LIMIT) { + intriParams.nValue = nValue; + intriParams.srcDValue = srcDValue; + intriParams.dstNzNStride = dstInnerStrideRow / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < nValue; i++) { + AscendC::DataCopy(dstTensor.data()[dstOffset + i * ELE_NUM_PER_C0], + srcTensor.data()[srcOffset + i * srcDValue], + intriParams); + } + } + } +}; + +/// Partial specialization for CopyGmToL1, AtlasA2, zN in and zN out. +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + std::enable_if_t::value && + tla::detail::iszN::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::iszN::value && + tla::detail::iszN::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and zN, while TensorDst must be L1 and zN"); + + const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride()); + + uint32_t blockCount = CeilDiv(tla::get<1>(srcTensor.originShape())); + uint32_t blockLen = tla::get<0>(srcTensor.originShape()); + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::DataCopyParams repeatParams; + if (srcOuterStrideCol / ELE_NUM_PER_C0 <= STRIDE_LIMIT) { + repeatParams.blockCount = blockCount; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = srcOuterStrideCol / ELE_NUM_PER_C0 - blockLen; + repeatParams.dstStride = dstOuterStrideCol / ELE_NUM_PER_C0 - blockLen; + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], repeatParams); + } else { + repeatParams.blockCount = 1; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = 0; + repeatParams.dstStride = 0; + for (uint32_t i = 0; i < blockCount; i++) { + AscendC::DataCopy(dstTensor.data()[dstOffset + i * dstOuterStrideCol], + srcTensor.data()[srcOffset + i * srcOuterStrideCol], + repeatParams); + } + } + } +}; + +/// Partial specialization for CopyGmToL1, AtlasA2, nZ in and nZ out. +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + std::enable_if_t::value && + tla::detail::isnZ::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isnZ::value && + tla::detail::isnZ::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and nZ, while TensorDst must be L1 and nZ"); + + const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + + uint32_t blockCount = CeilDiv(tla::get<0>(srcTensor.originShape())); + uint32_t blockLen = tla::get<1>(srcTensor.originShape()); + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::DataCopyParams repeatParams; + if (srcOuterStrideRow / ELE_NUM_PER_C0 <= STRIDE_LIMIT) { + repeatParams.blockCount = blockCount; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = srcOuterStrideRow / ELE_NUM_PER_C0 - blockLen; + repeatParams.dstStride = dstOuterStrideRow / ELE_NUM_PER_C0 - blockLen; + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], repeatParams); + } else { + repeatParams.blockCount = 1; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = 0; + repeatParams.dstStride = 0; + for (uint32_t i = 0; i < blockCount; i++) { + AscendC::DataCopy(dstTensor.data()[dstOffset + i * dstOuterStrideRow], + srcTensor.data()[srcOffset + i * srcOuterStrideRow], + repeatParams); + } + } + } +}; + +///////////////////////////////////////////TileCopyTlaExt////////////////////////////////////////////////////// +/// Partial specialization for TileCopyTlaExt, CopyGmToL1, AtlasA2, PaddingRowMajor in and zN out. +template +struct TileCopyTlaExt, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + layout::RowMajor, layout::zN> { + using ActualShape = tla::Shape; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTlaExt() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, ActualShape actualShape) + { + static_assert(tla::detail::isRowMajor::value && + tla::detail::iszN::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and RowMajor, while TensorDst must be L1 and zN"); + + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = tla::get<1>(actualShape); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = tla::get<1, 1>(dstTensor.stride()) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + intriParams.nValue = tla::get<0>(actualShape); + intriParams.srcDValue = tla::get<0>(srcTensor.stride()); + intriParams.dstNzNStride = tla::get<0, 0>(dstTensor.stride()) / ELE_NUM_PER_C0; + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +/// Partial specialization for TileCopyTlaExt, CopyGmToL1, AtlasA2, PaddingRowMajor in and zN out. +template +struct TileCopyTlaExt, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + layout::PaddingRowMajor, layout::zN> { + using ActualShape = tla::Shape; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTlaExt() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, ActualShape actualShape) + { + static_assert(tla::detail::iszN::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and PaddingRowMajor, " + "while TensorDst must be L1 and zN"); + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = tla::get<1>(actualShape); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = tla::get<1, 1>(dstTensor.stride()) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + intriParams.nValue = tla::get<0>(actualShape); + intriParams.srcDValue = tla::get<0, 0>(srcTensor.stride()); + intriParams.dstNzNStride = tla::get<0, 0>(dstTensor.stride()) / ELE_NUM_PER_C0; + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +/// Partial specialization for TileCopyTlaExt, CopyGmToL1, AtlasA2, ColumnMajor in and nZ out. +template +struct TileCopyTlaExt, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + layout::ColumnMajor, layout::nZ> { + using ActualShape = tla::Shape; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTlaExt() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, ActualShape actualShape) + { + static_assert(tla::detail::isColumnMajor::value && + tla::detail::isnZ::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and ColumnMajor, " + "while TensorDst must be L1 and nZ"); + + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = tla::get<0>(actualShape); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = tla::get<0, 1>(dstTensor.stride()) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + intriParams.nValue = tla::get<1>(actualShape); + intriParams.srcDValue = tla::get<1>(srcTensor.stride()); + intriParams.dstNzNStride = tla::get<1, 0>(dstTensor.stride()) / ELE_NUM_PER_C0; + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +/// Partial specialization for TileCopyTlaExt, CopyGmToL1, AtlasA2, PaddingColumnMajor in and nZ out. +template +struct TileCopyTlaExt, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + layout::PaddingColumnMajor, layout::nZ> { + using ActualShape = tla::Shape; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTlaExt() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, ActualShape actualShape) + { + static_assert(tla::detail::isnZ::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and PaddingColumnMajor, " + "while TensorDst must be L1 and nZ"); + + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = tla::get<0>(actualShape); + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = tla::get<0, 1>(dstTensor.stride()) / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + intriParams.nValue = tla::get<1>(actualShape); + intriParams.srcDValue = tla::get<1, 0>(srcTensor.stride()); + intriParams.dstNzNStride = tla::get<1, 0>(dstTensor.stride()) / ELE_NUM_PER_C0; + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +/// Partial specialization for TileCopyTlaExt, CopyGmToL1, AtlasA2, zN in and zN out. +template +struct TileCopyTlaExt, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + layout::zN, layout::zN> { + using ActualShape = tla::Shape; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTlaExt() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, ActualShape actualShape) + { + static_assert(tla::detail::iszN::value && + tla::detail::iszN::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and zN, while TensorDst must be L1 and zN"); + + const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride()); + + uint32_t blockCount = CeilDiv(tla::get<1>(actualShape)); + uint32_t blockLen = tla::get<0>(actualShape); + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::DataCopyParams repeatParams; + if (srcOuterStrideCol / ELE_NUM_PER_C0 <= STRIDE_LIMIT) { + repeatParams.blockCount = blockCount; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = srcOuterStrideCol / ELE_NUM_PER_C0 - blockLen; + repeatParams.dstStride = dstOuterStrideCol / ELE_NUM_PER_C0 - blockLen; + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], repeatParams); + } else { + repeatParams.blockCount = 1; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = 0; + repeatParams.dstStride = 0; + for (uint32_t i = 0; i < blockCount; i++) { + AscendC::DataCopy(dstTensor.data()[dstOffset + i * dstOuterStrideCol], + srcTensor.data()[srcOffset + i * srcOuterStrideCol], + repeatParams); + } + } + } +}; + +/// Partial specialization for TileCopyTlaExt, CopyGmToL1, AtlasA2, nZ in and nZ out. +template +struct TileCopyTlaExt, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + layout::nZ, layout::nZ> { + using ActualShape = tla::Shape; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTlaExt() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, ActualShape actualShape) + { + static_assert(tla::detail::isnZ::value && + tla::detail::isnZ::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and nZ, while TensorDst must be L1 and nZ"); + + const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride()); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + + uint32_t blockCount = CeilDiv(tla::get<0>(actualShape)); + uint32_t blockLen = tla::get<1>(actualShape); + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::DataCopyParams repeatParams; + if (srcOuterStrideRow / ELE_NUM_PER_C0 <= STRIDE_LIMIT) { + repeatParams.blockCount = blockCount; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = srcOuterStrideRow / ELE_NUM_PER_C0 - blockLen; + repeatParams.dstStride = dstOuterStrideRow / ELE_NUM_PER_C0 - blockLen; + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], repeatParams); + } else { + repeatParams.blockCount = 1; + repeatParams.blockLen = blockLen; + repeatParams.srcStride = 0; + repeatParams.dstStride = 0; + for (uint32_t i = 0; i < blockCount; i++) { + AscendC::DataCopy(dstTensor.data()[dstOffset + i * dstOuterStrideRow], + srcTensor.data()[srcOffset + i * srcOuterStrideRow], + repeatParams); + } + } + } +}; + +template +struct CopyGmToL1, + Gemm::GemmType> { + using LayoutDst = layout::VectorLayout; + using LayoutSrc = layout::VectorLayout; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyGmToL1() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::DataCopyParams intriParams; + intriParams.blockCount = 1; + intriParams.blockLen = layoutDst.shape(0) / ELE_NUM_PER_C0; + intriParams.srcStride = 0; + intriParams.dstStride = 0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +/// Partial specialization for CopyGmToL1, AtlasA2, RowMajor or ColumnMajor in and zN out. +template +struct TileCopySparseTla, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + std::enable_if_t<(tla::detail::isRowMajor::value || + tla::detail::isColumnMajor::value) && + tla::detail::iszN::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc); + + // Methods + + CATLASS_DEVICE + TileCopySparseTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert((tla::detail::isRowMajor::value || + tla::detail::isColumnMajor::value) && + tla::detail::iszN::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and RowMajor or ColumnMajor, while TensorDst must be L1 and zN"); + + auto srcShape = srcTensor.shape(); + auto srcStride = srcTensor.stride(); + auto dstShape = dstTensor.shape(); + auto coord = srcTensor.coord(); + + AscendC::Nd2NzParams nd2nzParams; + nd2nzParams.ndNum = 1; + nd2nzParams.nValue = + tla::min(tla::get<0, 0>(dstShape) * tla::get<0, 1>(dstShape), tla::get<0>(srcShape) - tla::get<0>(coord)); + nd2nzParams.dValue = + tla::min(tla::get<1, 0>(dstShape) * tla::get<1, 1>(dstShape), tla::get<1>(srcShape) - tla::get<1>(coord)); + nd2nzParams.srcNdMatrixStride = 0; + nd2nzParams.srcDValue = tla::get<0>(srcStride); + nd2nzParams.dstNzC0Stride = tla::get<0, 0>(dstShape) * tla::get<0, 1>(dstShape); + nd2nzParams.dstNzNStride = 1; + nd2nzParams.dstNzMatrixStride = 0; + + auto offset = srcTensor.layout()(coord); + + AscendC::DataCopy(dstTensor.data(), srcTensor.data()[offset], nd2nzParams); + } +}; + +template +struct TileCopySparseTla, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + std::enable_if_t::value && + tla::detail::isnZ::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc); + + // Methods + + CATLASS_DEVICE + TileCopySparseTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert((tla::detail::isColumnMajor::value) && + tla::detail::isnZ::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and ColumnMajor, while TensorDst must be L1 and nZ"); + const uint32_t nValue = tla::get<1>(srcTensor.shape()); + const uint32_t dValue = tla::get<0>(srcTensor.shape()); + const uint32_t srcDValue = tla::get<1>(srcTensor.stride()); + const uint32_t dstInnerStrideRow = tla::get<1, 0>(dstTensor.stride()); + const uint32_t dstOuterStrideCol = tla::get<0, 1>(dstTensor.stride()); + + AscendC::Nd2NzParams intriParams; + + intriParams.ndNum = 1; + intriParams.dValue = dValue; + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzC0Stride = dstOuterStrideCol / ELE_NUM_PER_C0; + intriParams.dstNzMatrixStride = 0; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + if (srcDValue < STRIDE_LIMIT) { + intriParams.nValue = nValue; + intriParams.srcDValue = srcDValue; + intriParams.dstNzNStride = dstInnerStrideRow / ELE_NUM_PER_C0; + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } else { + intriParams.nValue = 1; + intriParams.srcDValue = 0; + intriParams.dstNzNStride = 0; + for (uint32_t i = 0; i < nValue; i++) { + AscendC::DataCopy(dstTensor.data()[dstOffset + i * ELE_NUM_PER_C0], + srcTensor.data()[srcOffset + i * srcDValue], + intriParams); + } + } + } +}; + +/// Partial specialization for CopyGmToL1, AtlasA2, sparse zn in and zN out. +template +struct TileCopySparseTla, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + std::enable_if_t::value && + tla::detail::iszN::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc); + + // Methods + + CATLASS_DEVICE + TileCopySparseTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::iszN::value && + tla::detail::iszN::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and zN, while TensorDst must be L1 and zN"); + + using TransT = typename TensorSrc::Element; + + auto srcShape = srcTensor.shape(); + auto dstShape = dstTensor.shape(); + auto coord = srcTensor.coord(); + int32_t alignedGRow = tla::get<0, 0>(srcShape) * tla::get<0, 1>(srcShape); + // height direction need to be 16 aligned + auto alignHeight = tla::get<0, 0>(dstShape) * tla::get<0, 1>(dstShape); + int32_t blockLen = alignHeight * tla::get<1, 0>(dstShape) * sizeof(TransT) / BYTE_PER_C0; + int32_t srcStride = (alignedGRow - alignHeight) * tla::get<1, 0>(srcShape) * sizeof(TransT) / BYTE_PER_C0; + uint16_t nburst = tla::get<1, 1>(dstShape); + auto offset = srcTensor.layout()(coord); + AscendC::DataCopy( + dstTensor.data(), srcTensor.data()[offset], + {nburst, static_cast(blockLen), static_cast(srcStride), static_cast(0)}); + } +}; + +/// Partial specialization for CopyGmToL1, AtlasA2, sparse nZ in and nZ out. +template +struct TileCopySparseTla, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>, + std::enable_if_t::value && + tla::detail::isnZ::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc); + + // Methods + + CATLASS_DEVICE + TileCopySparseTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isnZ::value && + tla::detail::isnZ::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM and nZ, while TensorDst must be L1 and nZ"); + + using TransT = typename TensorSrc::Element; + auto srcShape = srcTensor.shape(); + auto dstShape = dstTensor.shape(); + auto coord = srcTensor.coord(); + int32_t alignedGCol = tla::get<1, 0>(srcShape) * tla::get<1, 1>(srcShape); + auto alignCol = tla::get<1, 0>(dstShape) * tla::get<1, 1>(dstShape); + int32_t blockLen = alignCol * tla::get<0, 0>(dstShape) * sizeof(TransT) / BYTE_PER_C0; + int32_t srcStride = (alignedGCol - alignCol) * tla::get<0, 0>(srcShape) * sizeof(TransT) / BYTE_PER_C0; + uint16_t nburst = tla::get<0, 1>(dstShape); + auto offset = srcTensor.layout()(coord); + AscendC::DataCopy( + dstTensor.data(), srcTensor.data()[offset], + {nburst, static_cast(blockLen), static_cast(srcStride), static_cast(0)}); + } +}; + +///////////////////////////////////////////TileCopyFAQTla////////////////////////////////////////////////////// +/// used for FA LoadQ, multi matrix datacopy ND->NZ +template +struct TileCopyFAQTla, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A1>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyFAQTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::iszN::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::A1, + "The input parameters do not match. TensorSrc must be GM, while TensorDst must be L1 and zN"); + + const uint32_t nValue = tla::get<1>(srcTensor.shape()); + const uint32_t dValue = tla::get<2>(srcTensor.shape()); + const uint32_t srcDValue = tla::get<1>(srcTensor.stride()); + const uint32_t ndNum = tla::get<0>(srcTensor.shape()); + const uint32_t srcNdMatrixStride = tla::get<0>(srcTensor.stride()); + const uint32_t dstNzNStride = tla::get<0>(srcTensor.shape()); + const uint32_t dstNzMatrixStride = ELE_NUM_PER_C0; + const uint32_t dstNzC0Stride = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape()); + const uint32_t dstInnerStrideRow = tla::get<0, 0>(dstTensor.stride()); + const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride()); + + AscendC::Nd2NzParams intriParams; + + intriParams.nValue = nValue; + intriParams.dValue = dValue; + intriParams.srcDValue = srcDValue; + intriParams.dstNzNStride = dstNzNStride; + intriParams.dstNzC0Stride = dstNzC0Stride; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + if (srcNdMatrixStride < STRIDE_LIMIT) { + intriParams.ndNum = ndNum; + intriParams.srcNdMatrixStride = srcNdMatrixStride; + intriParams.dstNzMatrixStride = dstNzMatrixStride; + AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } else { + intriParams.ndNum = 1; + intriParams.srcNdMatrixStride = 0; + intriParams.dstNzMatrixStride = 0; + for (uint32_t i = 0; i < nValue; i++) { + AscendC::DataCopy(dstTensor.data()[dstOffset + i * ELE_NUM_PER_C0], + srcTensor.data()[srcOffset + i * srcDValue], + intriParams); + } + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_ATLASA2_COPY_GM_TO_L1_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_gm_to_ub.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_gm_to_ub.hpp new file mode 100644 index 0000000000..35b1fb4002 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_gm_to_ub.hpp @@ -0,0 +1,99 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ATLASA2_COPY_GM_TO_UB_HPP +#define CATLASS_KDA_ATLASA2_COPY_GM_TO_UB_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "catlass/gemm/tile/kda_tile_copy_tla.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Tile { + +/// Partial specialization for AtlasA2, RowMajor in and RowMajor out. +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::GM>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::VECCALC>, + std::enable_if_t::value && + tla::detail::isRowMajor::value>> { + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(ElementSrc); + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isRowMajor::value && + tla::detail::isRowMajor::value && + TensorSrc::position == AscendC::TPosition::GM && + TensorDst::position == AscendC::TPosition::VECCALC, + "The input parameters do not match. TensorSrc must be GM and RowMajor, " + "while TensorDst must be UB and RowMajor"); + + const uint16_t row = tla::get<0>(srcTensor.originShape()); + const uint16_t col = tla::get<1>(srcTensor.originShape()); + + AscendC::DataCopyExtParams dataCopyParams( + row, + col * sizeof(ElementSrc), + (tla::get<0>(srcTensor.stride()) - col) * sizeof(ElementSrc), + (tla::get<0>(dstTensor.stride()) - col) / ELE_NUM_PER_BLK, + 0); + AscendC::DataCopyPadExtParams padParams(false, 0, 0, 0); + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + AscendC::DataCopyPad(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], dataCopyParams, padParams); + }; +}; + +template < + class ArchTag, + class GmType> +struct CopyGm2Ub { + static_assert(DEPENDENT_FALSE, "Unsupported copy gm to ub, can not find the specialization."); +}; + +template +struct CopyGm2Ub> { + using LayoutSrc = layout::VectorLayout; + using LayoutDst = layout::VectorLayout; + + static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(Element); + + CATLASS_DEVICE + CopyGm2Ub() = default; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::GlobalTensor const &srcTensor, + layout::VectorLayout const &layoutDst, + layout::VectorLayout const &layoutSrc) + { + AscendC::DataCopyExtParams dataCopyParams( + 1, + layoutSrc.shape(0) * sizeof(Element), + 0, + 0, + 0); + AscendC::DataCopyPadExtParams padParams(false, 0, 0, 0); + AscendC::DataCopyPad(dstTensor, srcTensor, dataCopyParams, padParams); + }; +}; + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_ATLASA2_COPY_GM_TO_UB_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l0c_to_gm.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l0c_to_gm.hpp new file mode 100644 index 0000000000..baa55940e6 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l0c_to_gm.hpp @@ -0,0 +1,632 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ATLASA2_COPY_L0C_TO_GM_HPP +#define CATLASS_KDA_ATLASA2_COPY_L0C_TO_GM_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Tile { + +enum class ScaleGranularity { + UNDEFINED = -1, + NO_QUANT = 0, + PER_TENSOR, + PER_CHANNEL, + PER_GROUP +}; + +template < + class ArchTag, + class ElementSrc, + class ElementDst, + ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT> +struct CopyL0CToGmQuantMode { + static_assert(DEPENDENT_FALSE, "Unsupported copy l0c to gm, can not find the specialization."); +}; + +// CopyL0CToGm fp32 to fp32 +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + float, float, + ScaleGranularity::NO_QUANT> { + static constexpr auto VALUE = QuantMode_t::NoQuant; +}; + +// CopyL0CToGm cast fp32 to fp16 +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + float, half, + ScaleGranularity::NO_QUANT> { + static constexpr auto VALUE = QuantMode_t::F322F16; +}; + +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + float, half, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::F322F16; +}; + +// CopyL0CToGm cast fp32 to bf16 +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + float, bfloat16_t, + ScaleGranularity::NO_QUANT> { + static constexpr auto VALUE = QuantMode_t::F322BF16; +}; + +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + float, bfloat16_t, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::F322BF16; +}; + +// CopyL0CToGm cast float to uint8/int8 +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + float, uint8_t, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::QF322B8_PRE; +}; + +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + float, uint8_t, + ScaleGranularity::PER_CHANNEL> { + static constexpr auto VALUE = QuantMode_t::VQF322B8_PRE; +}; + +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + float, int8_t, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::QF322B8_PRE; +}; + +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + float, int8_t, + ScaleGranularity::PER_CHANNEL> { + static constexpr auto VALUE = QuantMode_t::VQF322B8_PRE; +}; + +// CopyL0CToGm output int32 +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + int32_t, int32_t, + ScaleGranularity::NO_QUANT> { + static constexpr auto VALUE = QuantMode_t::NoQuant; +}; + +// CopyL0CToGm cast int32_t to fp16 +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + int32_t, half, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::DEQF16; +}; + +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + int32_t, half, + ScaleGranularity::PER_CHANNEL> { + static constexpr auto VALUE = QuantMode_t::VDEQF16; +}; + +// CopyL0CToGm cast int32 to uint8/int8 +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + int32_t, uint8_t, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::REQ8; +}; + +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + int32_t, uint8_t, + ScaleGranularity::PER_CHANNEL> { + static constexpr auto VALUE = QuantMode_t::VREQ8; +}; + +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + int32_t, int8_t, + ScaleGranularity::PER_TENSOR> { + static constexpr auto VALUE = QuantMode_t::REQ8; +}; + +template <> +struct CopyL0CToGmQuantMode< + Catlass::Arch::AtlasA2, + int32_t, int8_t, + ScaleGranularity::PER_CHANNEL> { + static constexpr auto VALUE = QuantMode_t::VREQ8; +}; + +template < + class ArchTag, + class ElementAccumulator, + class GmType, + ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT, + bool ReluEnable = false> +struct CopyL0CToGm { + static_assert(DEPENDENT_FALSE, "Unsupported copy l0c to gm, can not find the specialization."); +}; + +template < + class ElementAccumulator_, + class ElementDst_, + bool ReluEnable_> +struct CopyL0CToGm, + ScaleGranularity::NO_QUANT, + ReluEnable_> { + using ArchTag = Catlass::Arch::AtlasA2; + using ElementDst = ElementDst_; + using ElementSrc = ElementAccumulator_; + using LayoutSrc = Catlass::layout::zN; + using LayoutDst = Catlass::layout::RowMajor; + static constexpr auto quantPre = CopyL0CToGmQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + struct Params {}; + Params params; + + CATLASS_DEVICE + CopyL0CToGm() = default; + + CATLASS_DEVICE + CopyL0CToGm(Params const ¶ms_) + : params(params_) {}; + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor const &dst, AscendC::LocalTensor const &src, + LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0) + { + AscendC::FixpipeParamsV220 intriParams; + + // Fixpipe layout information + intriParams.nSize = dstLayout.shape(1); + intriParams.mSize = dstLayout.shape(0); + intriParams.srcStride = srcLayout.stride(3) / srcLayout.stride(0); + intriParams.dstStride = dstLayout.stride(0); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + + // Call AscendC Fixpipe + AscendC::Fixpipe(dst, src, intriParams); + } +}; + +template < + class ElementAccumulator_, + class ElementDst_, + bool ReluEnable_> +struct CopyL0CToGm, + ScaleGranularity::PER_TENSOR, + ReluEnable_> { + using ArchTag = Catlass::Arch::AtlasA2; + using ElementDst = ElementDst_; + using ElementSrc = ElementAccumulator_; + using LayoutSrc = Catlass::layout::zN; + using LayoutDst = Catlass::layout::RowMajor; + static constexpr auto quantPre = CopyL0CToGmQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + struct Params { + float scale = 1.0; + + CATLASS_HOST_DEVICE + Params() = default; + + CATLASS_HOST_DEVICE + Params(float scalar) + { + scale = scalar; + } + }; + Params params; + + CATLASS_DEVICE + CopyL0CToGm() = default; + + CATLASS_DEVICE + CopyL0CToGm(Params const ¶ms_) + : params(params_) {}; + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor const &dst, AscendC::LocalTensor const &src, + LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0) + { + AscendC::FixpipeParamsV220 intriParams; + + // Fixpipe layout information + intriParams.nSize = dstLayout.shape(1); + intriParams.mSize = dstLayout.shape(0); + intriParams.srcStride = srcLayout.stride(3) / srcLayout.stride(0); + intriParams.dstStride = dstLayout.stride(0); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.deqScalar = static_cast(*reinterpret_cast(¶ms.scale)); + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + + // Call AscendC Fixpipe + AscendC::Fixpipe(dst, src, intriParams); + } +}; + +template < + class ElementAccumulator_, + class ElementDst_, + bool ReluEnable_> +struct CopyL0CToGm, + ScaleGranularity::PER_CHANNEL, + ReluEnable_> { + using ArchTag = Catlass::Arch::AtlasA2; + using ElementDst = ElementDst_; + using ElementSrc = ElementAccumulator_; + using LayoutSrc = Catlass::layout::zN; + using LayoutDst = Catlass::layout::RowMajor; + static constexpr auto quantPre = CopyL0CToGmQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + struct Params {}; + Params params; + + CATLASS_DEVICE + CopyL0CToGm() = default; + + CATLASS_DEVICE + CopyL0CToGm(Params const ¶ms_) + : params(params_) {}; + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor const &dst, AscendC::LocalTensor const &src, AscendC::LocalTensor const &scale, + LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0) + { + AscendC::FixpipeParamsV220 intriParams; + + // Fixpipe layout information + intriParams.nSize = dstLayout.shape(1); + intriParams.mSize = dstLayout.shape(0); + intriParams.srcStride = srcLayout.stride(3) / srcLayout.stride(0); + intriParams.dstStride = dstLayout.stride(0); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + + // Call AscendC Fixpipe + AscendC::SetFixPipeConfig(scale, false); + AscendC::PipeBarrier(); + AscendC::Fixpipe(dst, src, intriParams); + } +}; + +template < + class ElementAccumulator_, + class ElementDst_, + bool ReluEnable_> +struct CopyL0CToGm, + ScaleGranularity::NO_QUANT, + ReluEnable_> { + using ArchTag = Catlass::Arch::AtlasA2; + using ElementDst = ElementDst_; + using ElementSrc = ElementAccumulator_; + using LayoutSrc = Catlass::layout::zN; + using LayoutDst = Catlass::layout::zN; + static constexpr auto quantPre = CopyL0CToGmQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + struct Params {}; + Params params; + + CATLASS_DEVICE + CopyL0CToGm() = default; + + CATLASS_DEVICE + CopyL0CToGm(Params const ¶ms_) + : params(params_) {}; + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor const &dst, AscendC::LocalTensor const &src, + LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0) + { + AscendC::FixpipeParamsV220 intriParams; + + // Fixpipe layout information + intriParams.nSize = dstLayout.shape(2) * dstLayout.shape(3); + intriParams.mSize = dstLayout.shape(0) * dstLayout.shape(1); + intriParams.srcStride = srcLayout.stride(3) / srcLayout.shape(2); + intriParams.dstStride = dstLayout.stride(3) / (BYTE_PER_C0 / sizeof(ElementDst)); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + + if constexpr (std::is_same_v && std::is_same_v) { + intriParams.isChannelSplit = true; + } + + // Call AscendC Fixpipe + AscendC::Fixpipe(dst, src, intriParams); + } +}; + +template < + class ElementAccumulator_, + class ElementDst_, + bool ReluEnable_> +struct CopyL0CToGm, + ScaleGranularity::NO_QUANT, + ReluEnable_> { + using ArchTag = Catlass::Arch::AtlasA2; + using ElementDst = ElementDst_; + using ElementSrc = ElementAccumulator_; + using LayoutSrc = Catlass::layout::zN; + using LayoutDst = Catlass::layout::NDC1HWC0; + static constexpr auto quantPre = CopyL0CToGmQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + struct Params {}; + Params params; + + CATLASS_DEVICE + CopyL0CToGm() = default; + + CATLASS_DEVICE + CopyL0CToGm(Params const ¶ms_) + : params(params_) {}; + + CATLASS_DEVICE + void operator()(AscendC::GlobalTensor const &dst, AscendC::LocalTensor const &src, + LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0) + { + AscendC::FixpipeParamsV220 intriParams; + + intriParams.nSize = srcLayout.orgShape(1); + intriParams.mSize = srcLayout.orgShape(0); + intriParams.srcStride = srcLayout.stride(3) / srcLayout.shape(2); + intriParams.dstStride = dstLayout.shape(1) * dstLayout.shape(2); + + if constexpr (AscendC::IsSameType::value && + AscendC::IsSameType::value) { + intriParams.isChannelSplit = true; + } + + intriParams.quantPre = quantPre; + intriParams.reluEn = false; + intriParams.unitFlag = unitFlag; + AscendC::Fixpipe(dst, src, intriParams); + } +}; + +///////////////////////////////////////////CopyL0CToGmTla///////////////////////////////////////////////// +// L0C copy mode +struct CopyToGM {}; +struct CopyToL1 {}; + +template < + class ArchTag, + class TensorSrc, + class TensorDst, + ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT, + bool ReluEnable = false, + class Enable = void> +struct CopyL0CToGmTla { + static_assert(DEPENDENT_FALSE, "Unsupported copy l0c to gm, can not find the specialization."); +}; + +template < + class TensorSrc_, + class ElementDst_, + class LayoutDst_, + class CoordDst_, + bool ReluEnable_> +struct CopyL0CToGmTla, LayoutDst_, CoordDst_, AscendC::TPosition::GM>, + ScaleGranularity::NO_QUANT, + ReluEnable_, + std::enable_if_t::value>> { + using ArchTag = Catlass::Arch::AtlasA2; + using ElementDst = ElementDst_; + using ElementSrc = typename TensorSrc_::Element; + static constexpr auto quantPre = CopyL0CToGmQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + struct Params {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t unitFlag = 0) + { + static_assert(tla::detail::isRowMajor::value && + TensorSrc::position == AscendC::TPosition::CO1 && + TensorDst::position == AscendC::TPosition::GM, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be GM and RowMajor"); + + AscendC::FixpipeParamsV220 intriParams; + + // Fixpipe layout information + intriParams.nSize = tla::get<1>(dstTensor.originShape()); + intriParams.mSize = tla::get<0>(dstTensor.originShape()); + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + intriParams.dstStride = tla::get<0>(dstTensor.stride()); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + // Call AscendC Fixpipe + AscendC::Fixpipe( + dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +template < + class TensorSrc_, + class ElementDst_, + class LayoutDst_, + class CoordDst_, + bool ReluEnable_> +struct CopyL0CToGmTla, LayoutDst_, CoordDst_, AscendC::TPosition::GM>, + ScaleGranularity::NO_QUANT, + ReluEnable_, + std::enable_if_t::value>> { + using ArchTag = Catlass::Arch::AtlasA2; + using ElementDst = ElementDst_; + using ElementSrc = typename TensorSrc_::Element; + static constexpr auto quantPre = CopyL0CToGmQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t unitFlag = 0) + { + static_assert(tla::detail::iszN::value && + TensorSrc::position == AscendC::TPosition::CO1 && + TensorDst::position == AscendC::TPosition::GM, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be GM and zN"); + + AscendC::FixpipeParamsV220 intriParams; + + // Fixpipe layout information + intriParams.nSize = tla::get<1>(dstTensor.originShape()); + intriParams.mSize = tla::get<0>(dstTensor.originShape()); + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<1, 0>(srcTensor.shape()); + intriParams.dstStride = tla::get<1, 1>(dstTensor.stride()) / (BYTE_PER_C0 / sizeof(ElementDst)); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + + if constexpr (std::is_same_v && std::is_same_v) { + intriParams.isChannelSplit = true; + } + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + // Call AscendC Fixpipe + AscendC::Fixpipe( + dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams); + } +}; + +template < + class ArchTag, + class TensorSrc, + class TensorDst, + ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT, + class Enable = void> +struct CopyL0CToGmSparseTla { + static_assert(DEPENDENT_FALSE, "Unsupported copy l0c to gm, can not find the specialization."); +}; + +template < + class TensorSrc_, + class ElementDst_, + class LayoutDst_, + class CoordDst_> +struct CopyL0CToGmSparseTla, LayoutDst_, CoordDst_, AscendC::TPosition::GM>, + ScaleGranularity::NO_QUANT, + std::enable_if_t::value>> { + using ArchTag = Catlass::Arch::AtlasA2; + using ElementDst = ElementDst_; + using ElementSrc = typename TensorSrc_::Element; + static constexpr auto quantPre = CopyL0CToGmQuantMode::VALUE; + + struct Params {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isRowMajor::value && + TensorSrc::position == AscendC::TPosition::CO1 && + TensorDst::position == AscendC::TPosition::GM, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be GM and RowMajor"); + + auto dstShape = dstTensor.shape(); + auto srcShape = srcTensor.shape(); + auto dstStride = dstTensor.stride(); + auto coord = dstTensor.coord(); + + AscendC::FixpipeParamsV220 params; + params.nSize = + tla::min(tla::get<1, 0>(srcShape) * tla::get<1, 1>(srcShape), tla::get<1>(dstShape) - tla::get<1>(coord)); + params.mSize = + tla::min(tla::get<0, 0>(srcShape) * tla::get<0, 1>(srcShape), tla::get<0>(dstShape) - tla::get<0>(coord)); + params.srcStride = tla::get<0, 0>(srcShape) * tla::get<0, 1>(srcShape); + params.dstStride = tla::get<0>(dstStride); + params.quantPre = quantPre; + + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + auto dstOffset = dstTensor.layout()(coord); + AscendC::Fixpipe(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], params); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_ATLASA2_COPY_L0C_TO_GM_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_bt.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_bt.hpp new file mode 100644 index 0000000000..ebe5bc4fbd --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_bt.hpp @@ -0,0 +1,62 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ATLASA2_COPY_L1_TO_BT_HPP +#define CATLASS_KDA_ATLASA2_COPY_L1_TO_BT_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "tla/kda_tensor.hpp" + +using namespace tla; + +namespace Catlass::Gemm::Tile { + +template < + class ArchTag, + class L1Type, + class L0Type = void> +struct CopyL1ToBT { + static_assert(DEPENDENT_FALSE, "Unsupported copy l1 to biasTable buffer, can not find the specialization."); +}; + +template +struct CopyL1ToBT, + Catlass::Gemm::GemmType> { + using LayoutDst = layout::VectorLayout; + using LayoutSrc = layout::VectorLayout; + + static constexpr uint32_t ELE_NUM_PER_C2 = BYTE_PER_C2 / sizeof(ElementSrc); + + CATLASS_DEVICE + CopyL1ToBT() {} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc) + { + AscendC::DataCopyParams intriParams; + intriParams.blockCount = 1; + intriParams.blockLen = (layoutDst.shape(0) + ELE_NUM_PER_C2 - 1) / ELE_NUM_PER_C2; + intriParams.srcStride = 0; + intriParams.dstStride = 0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_ATLASA2_COPY_L1_TO_BT_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_fp.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_fp.hpp new file mode 100644 index 0000000000..52bc44969a --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_fp.hpp @@ -0,0 +1,62 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ATLASA2_COPY_L1_TO_FP_HPP +#define CATLASS_KDA_ATLASA2_COPY_L1_TO_FP_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "tla/kda_tensor.hpp" + +using namespace tla; + +namespace Catlass::Gemm::Tile { + +template < + class ArchTag, + class L1Type, + class L0Type = void> +struct CopyL1ToFP { + static_assert(DEPENDENT_FALSE, "Unsupported copy l1 to fixpipe buffer, can not find the specialization."); +}; + +template +struct CopyL1ToFP, + Catlass::Gemm::GemmType> { + using LayoutDst = layout::VectorLayout; + using LayoutSrc = layout::VectorLayout; + + static constexpr uint32_t ELE_NUM_PER_FP = BYTE_PER_BLK_FP / sizeof(ElementSrc); + + CATLASS_DEVICE + CopyL1ToFP() {} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc) + { + AscendC::DataCopyParams intriParams; + intriParams.blockCount = 1; + intriParams.blockLen = (layoutDst.shape(0) + ELE_NUM_PER_FP - 1) / ELE_NUM_PER_FP; + intriParams.srcStride = 0; + intriParams.dstStride = 0; + AscendC::DataCopy(dstTensor, srcTensor, intriParams); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_ATLASA2_COPY_L1_TO_FP_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_l0a.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_l0a.hpp new file mode 100644 index 0000000000..89b0637c43 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_l0a.hpp @@ -0,0 +1,719 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ATLASA2_COPY_L1_TO_L0A_HPP +#define CATLASS_KDA_ATLASA2_COPY_L1_TO_L0A_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/kda_numeric_size.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "catlass/gemm/tile/kda_tile_copy_tla.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Tile { + +template < + class ArchTag, + class L1Type, + class L0Type = void> +struct CopyL1ToL0A { + static_assert(DEPENDENT_FALSE, "Unsupported copy l1 to l0, can not find the specialization."); +}; + +//////////////////////////////// +/// new add gemm +template +struct CopyL1ToL0A, Catlass::Gemm::GemmType> { + using LayoutDst = layout::zZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + CATLASS_DEVICE + CopyL1ToL0A() {} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(layoutDst.shape(3)); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(3) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + + for (uint32_t i = 0; i < layoutDst.shape(1); i++) { + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; + +template +struct CopyL1ToL0A, Catlass::Gemm::GemmType> { + using LayoutDst = layout::zZ; + using LayoutSrc = layout::nN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + CATLASS_DEVICE + CopyL1ToL0A() {} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = static_cast(CeilDiv(layoutSrc.orgShape(0))); + ; + loadDataParams.sid = 0; + loadDataParams.dstGap = 0; + loadDataParams.ifTranspose = true; + loadDataParams.addrMode = 0; + for (uint32_t i = 0; i < CeilDiv(layoutSrc.orgShape(0)); i++) { + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; + +template +struct CopyL1ToL0A, Catlass::Gemm::GemmType> { + using Element = float; + using LayoutDst = layout::zZ; + using LayoutSrc = layout::nN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + CATLASS_DEVICE + CopyL1ToL0A() {} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc) + { + AscendC::LoadData2dTransposeParams loadDataParams; + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = static_cast(CeilDiv(layoutSrc.orgShape(0))); + loadDataParams.dstGap = 1; + loadDataParams.dstFracGap = 0; + for (uint32_t i = 0; i < CeilDiv(layoutSrc.orgShape(0)); i++) { + AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1) * 2], loadDataParams); + } + } +}; + +template +struct CopyL1ToL0A, Catlass::Gemm::GemmType> { + using Element = int8_t; + using LayoutDst = layout::zZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + CATLASS_DEVICE + CopyL1ToL0A() {} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc) + { + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = 1; + loadDataParams.dstGap = 0; + loadDataParams.dstFracGap = CeilDiv(layoutDst.orgShape(1)) - 1; + + for (uint32_t i = 0; i < CeilDiv(layoutDst.orgShape(0)); i++) { + AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1) * 2], + srcTensor[i * layoutSrc.stride(1)], + loadDataParams); + } + } +}; + +template +struct CopyL1ToL0A> { + using LayoutDst = layout::zZ; + using LayoutSrc = layout::NDC1HWC0; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint8_t RIGHT_MOVE_8 = 8; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0A(uint32_t strideW = 0, + uint32_t strideH = 0, + uint32_t filterW = 0, + uint32_t filterH = 0, + uint32_t dilationFilterW = 0, + uint32_t dilationFilterH = 0) + { + loadData3Dv2Params.strideW = strideW; + loadData3Dv2Params.strideH = strideH; + loadData3Dv2Params.filterW = filterW; + loadData3Dv2Params.filterSizeW = filterW >> RIGHT_MOVE_8; + loadData3Dv2Params.filterH = filterH; + loadData3Dv2Params.filterSizeH = filterH >> RIGHT_MOVE_8; + loadData3Dv2Params.dilationFilterW = dilationFilterW; + loadData3Dv2Params.dilationFilterH = dilationFilterH; + } + + CATLASS_DEVICE + static CopyL1ToL0A MakeCopyL1ToL0A( + uint32_t strideW = 0, + uint32_t strideH = 0, + uint32_t filterW = 0, + uint32_t filterH = 0, + uint32_t dilationFilterW = 0, + uint32_t dilationFilterH = 0) + { + return CopyL1ToL0A(strideW, strideH, filterW, filterH, dilationFilterW, dilationFilterH); + } + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc, + uint32_t kStartPt, uint32_t mStartPt, + uint32_t l1H, uint32_t l1W, uint8_t *padList) + { + loadData3Dv2Params.l1H = l1H; + loadData3Dv2Params.l1W = l1W; + loadData3Dv2Params.padList[0] = padList[0]; + loadData3Dv2Params.padList[1] = padList[1]; + loadData3Dv2Params.padList[2] = padList[2]; + loadData3Dv2Params.padList[3] = padList[3]; + loadData3Dv2Params.kStartPt = kStartPt; + loadData3Dv2Params.mStartPt = mStartPt; + loadData3Dv2Params.kExtension = layoutDst.orgShape(1); + loadData3Dv2Params.mExtension = layoutDst.orgShape(0); + loadData3Dv2Params.channelSize = layoutSrc.orgShape(1) * layoutSrc.orgShape(2) * layoutSrc.orgShape(5); + static constexpr AscendC::IsResetLoad3dConfig CONV3D_LOAD3DV2_DEFAULT_CONFIG = {true, true}; + AscendC::LoadData(dstTensor, srcTensor, loadData3Dv2Params); + } + +private: + AscendC::LoadData3DParamsV2 loadData3Dv2Params; +}; + +////////////////////////////////////////// + +/// Partial specialization for zN in and zZ out. +template +struct CopyL1ToL0A> { + using LayoutDst = layout::zZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0A() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(layoutDst.shape(3)); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(3) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + + for (uint32_t i = 0; i < layoutDst.shape(1); i++) { + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; + +/// Partial specialization for float, zN in and zZ out. +template +struct CopyL1ToL0A> { + using Element = float; + using LayoutDst = layout::zZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0A() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + constexpr uint8_t PAD_LIST[4] = {0, 0, 0, 0}; + uint16_t l1M = layoutSrc.shape(0) * layoutSrc.shape(1); + uint16_t l1K = layoutSrc.shape(2) * layoutSrc.shape(3); + uint16_t l0M = layoutDst.shape(0) * layoutDst.shape(1); + uint16_t l0K = layoutDst.shape(2) * layoutDst.shape(3); + AscendC::SetFmatrix(1, l1M, PAD_LIST, AscendC::FmatrixMode::FMATRIX_LEFT); + static constexpr AscendC::IsResetLoad3dConfig config = {false, false}; + AscendC::LoadData3DParamsV2 loadDataParams; + loadDataParams.kExtension = l0K; + loadDataParams.mExtension = l0M; + loadDataParams.channelSize = l1K; + + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + } +}; + +template +struct CopyL1ToL0A> { + using LayoutDst = layout::zZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + CATLASS_DEVICE + CopyL1ToL0A() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(3) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = true; + loadDataParams.addrMode = 0; + + for (uint32_t i = 0; i < CeilDiv(layoutDst.orgShape(0)); i++) { + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; + +/// Partial specialization for int8_t, nZ in and zZ out. (Transpose A) +template +struct CopyL1ToL0A> { + using Element = int8_t; + using LayoutDst = layout::zZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0A() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = 1; + loadDataParams.dstGap = 0; + loadDataParams.dstFracGap = CeilDiv(layoutDst.orgShape(1)) - 1; + + for (uint32_t i = 0; i < CeilDiv(layoutDst.orgShape(0)); i++) { + AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1) * 2], + srcTensor[i * layoutSrc.stride(1)], + loadDataParams); + } + } +}; + +/// Partial specialization for float, nZ in and zZ out. (Transpose A) +template +struct CopyL1ToL0A> { + using Element = float; + using LayoutDst = layout::zZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0A() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + constexpr uint8_t PAD_LIST[] = {0, 0, 0, 0}; + uint16_t l1M = layoutSrc.shape(0) * layoutSrc.shape(1); + uint16_t l1K = layoutSrc.shape(2) * layoutSrc.shape(3); + uint16_t l0M = layoutDst.shape(0) * layoutDst.shape(1); + uint16_t l0K = layoutDst.shape(2) * layoutDst.shape(3); + // K, M need to be 16 aligned for f32 + uint16_t l1MAlign = RoundUp(l1M); + uint16_t l1KAlign = RoundUp(l1K); + uint16_t l0MAlign = RoundUp(l0M); + uint16_t l0KAlign = RoundUp(l0K); + AscendC::SetFmatrix(1, l1KAlign, PAD_LIST, AscendC::FmatrixMode::FMATRIX_LEFT); + static constexpr AscendC::IsResetLoad3dConfig config = {false, false}; + AscendC::LoadData3DParamsV2 loadDataParams; + loadDataParams.kExtension = l0MAlign; + loadDataParams.mExtension = l0KAlign; + loadDataParams.enTranspose = true; + loadDataParams.channelSize = l1MAlign; + + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + } +}; + +///////////////////////////////////////////TileCopyTla////////////////////////////////////////////////////// + +/// Partial specialization for CopyL1ToL0A, AtlasA2, zN in and zZ out. +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A2>, + std::enable_if_t::value && + tla::detail::iszZ::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::iszN::value && + tla::detail::iszZ::value && + TensorSrc::position == AscendC::TPosition::A1 && + TensorDst::position == AscendC::TPosition::A2, + "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0A and zZ"); + + const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride()); + const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride()); + const uint32_t dstOuterShapeRow = CeilDiv(tla::get<0>(dstTensor.originShape()), tla::get<0, 0>(dstTensor.shape())); + const uint32_t dstOuterShapeCol = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape())); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + + AscendC::LoadData2DParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = dstOuterShapeCol; + loadDataParams.srcStride = srcOuterStrideCol / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = 0; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + for (uint32_t i = 0; i < dstOuterShapeRow; i++) { + AscendC::LoadData(dstTensor.data()[dstOffset + i * dstOuterStrideRow], + srcTensor.data()[srcOffset + i * srcOuterStrideRow], + loadDataParams); + } + } +}; + +/// Partial specialization for CopyL1ToL0A, AtlasA2, float, zN in and zZ out. +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A2>, + std::enable_if_t::value && tla::detail::iszZ::value>> { + using Element = float; + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(std::is_same_v && + std::is_same_v && + tla::detail::iszN::value && + tla::detail::iszZ::value && + TensorSrc::position == AscendC::TPosition::A1 && + TensorDst::position == AscendC::TPosition::A2, + "The input parameters do not match. TensorSrc must be float, L1 and zN, " + "while TensorDst must be float, L0A and zZ"); + + constexpr uint8_t PAD_LIST[4] = {0, 0, 0, 0}; + uint16_t l1M = tla::get<1, 1>(srcTensor.stride()) / tla::get<1, 0>(srcTensor.shape()); + uint16_t l1K = RoundUp(tla::get<1>(srcTensor.originShape())); + uint16_t l0M = RoundUp(tla::get<0>(dstTensor.originShape())); + uint16_t l0K = RoundUp(tla::get<1>(dstTensor.originShape())); + AscendC::SetFmatrix(1, l1M, PAD_LIST, AscendC::FmatrixMode::FMATRIX_LEFT); + static constexpr AscendC::IsResetLoad3dConfig config = {false, false}; + AscendC::LoadData3DParamsV2 loadDataParams; + loadDataParams.kExtension = l0K; + loadDataParams.mExtension = l0M; + loadDataParams.channelSize = l1K; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], loadDataParams); + } +}; + +/// Partial specialization for CopyL1ToL0A, AtlasA2, nZ in and zZ out. (Transpose A) +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A2>, + std::enable_if_t::value && + tla::detail::iszZ::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isnZ::value && + tla::detail::iszZ::value && + TensorSrc::position == AscendC::TPosition::A1 && + TensorDst::position == AscendC::TPosition::A2, + "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0A and zZ"); + + const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride()); + const uint32_t dstOuterShapeRow = CeilDiv(tla::get<0>(dstTensor.originShape()), tla::get<0, 0>(dstTensor.shape())); + const uint32_t dstOuterShapeCol = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape())); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + + AscendC::LoadData2DParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = dstOuterShapeCol; + loadDataParams.srcStride = 1; + loadDataParams.sid = 0; + loadDataParams.dstGap = 0; + loadDataParams.ifTranspose = true; + loadDataParams.addrMode = 0; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + for (uint32_t i = 0; i < dstOuterShapeRow; i++) { + AscendC::LoadData(dstTensor.data()[dstOffset + i * dstOuterStrideRow], + srcTensor.data()[srcOffset + i * srcOuterStrideRow], + loadDataParams); + } + } +}; + +/// Partial specialization for CopyL1ToL0A, AtlasA2, int8_t, nZ in and zZ out. (Transpose A) +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A2>, + std::enable_if_t::value && tla::detail::iszZ::value>> { + using Element = int8_t; + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(std::is_same_v && + std::is_same_v && + tla::detail::isnZ::value && + tla::detail::iszZ::value && + TensorSrc::position == AscendC::TPosition::A1 && + TensorDst::position == AscendC::TPosition::A2, + "The input parameters do not match. TensorSrc must be int8_t, L1 and nZ, " + "while TensorDst must be int8_t, L0A and zZ"); + + const uint32_t srcOuterShapeRow = CeilDiv(tla::get<0>(srcTensor.originShape()), tla::get<0, 0>(srcTensor.shape())); + const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride()); + const uint32_t dstOuterShapeCol = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape())); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = dstOuterShapeCol; + loadDataParams.srcStride = 1; + loadDataParams.dstGap = 0; + loadDataParams.dstFracGap = dstOuterShapeCol - 1; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + for (uint32_t i = 0; i < srcOuterShapeRow; i++) { + AscendC::LoadDataWithTranspose(dstTensor.data()[dstOffset + i * dstOuterStrideRow * 2], + srcTensor.data()[srcOffset + i * srcOuterStrideRow], + loadDataParams); + } + } +}; + +/// Partial specialization for CopyL1ToL0A, AtlasA2, float, nZ in and zZ out. (Transpose A) +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A2>, + std::enable_if_t::value && tla::detail::iszZ::value>> { + using Element = float; + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(std::is_same_v && + std::is_same_v && + tla::detail::isnZ::value && + tla::detail::iszZ::value && + TensorSrc::position == AscendC::TPosition::A1 && + TensorDst::position == AscendC::TPosition::A2, + "The input parameters do not match. TensorSrc must be float, L1 and nZ, " + "while TensorDst must be float, L0A and zZ"); + + constexpr uint8_t PAD_LIST[4] = {0, 0, 0, 0}; + uint16_t l1M = tla::get<0>(srcTensor.originShape()); + uint16_t l1K = tla::get<0, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.shape()); + uint16_t l0M = tla::get<0>(dstTensor.originShape()); + uint16_t l0K = tla::get<1>(dstTensor.originShape()); + // K, M need to be 16 aligned for f32 + uint16_t l1MAlign = RoundUp(l1M); + uint16_t l1KAlign = RoundUp(l1K); + uint16_t l0MAlign = RoundUp(l0M); + uint16_t l0KAlign = RoundUp(l0K); + AscendC::SetFmatrix(1, l1KAlign, PAD_LIST, AscendC::FmatrixMode::FMATRIX_LEFT); + static constexpr AscendC::IsResetLoad3dConfig config = {false, false}; + AscendC::LoadData3DParamsV2 loadDataParams; + loadDataParams.kExtension = l0MAlign; + loadDataParams.mExtension = l0KAlign; + loadDataParams.enTranspose = true; + loadDataParams.channelSize = l1MAlign; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], loadDataParams); + } +}; + +/// Partial specialization for CopyL1ToL0A, AtlasA2, zN in and zZ out. +template +struct TileCopySparseTla, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::A2>, + std::enable_if_t::value && + tla::detail::iszZ::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(ElementSrc); + static constexpr uint16_t HW_N0 = 16; ///< Hardware configuration parameter for N0, reserved for future hardware configuration + static constexpr uint16_t HW_M0 = 16; ///< Hardware configuration parameter for M0 + static constexpr uint64_t M_POS_BIT = 48; ///< Bit position for M index + static constexpr uint64_t K_POS_BIT = 32; ///< Bit position for K index + static constexpr uint64_t M_STEP_BIT = 16; ///< Bit step for M index + static constexpr uint8_t PAD_LIST[4] = {0, 0, 0, 0}; ///< Padding list for data alignment + + // Methods + + CATLASS_DEVICE + TileCopySparseTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::iszN::value && + tla::detail::iszZ::value && + TensorSrc::position == AscendC::TPosition::A1 && + TensorDst::position == AscendC::TPosition::A2, + "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0A and zZ"); + + auto srcShape = srcTensor.shape(); + auto dstShape = dstTensor.shape(); + auto coord = srcTensor.coord(); + uint16_t aL1K = tla::get<1, 0>(srcShape) * tla::get<1, 1>(srcShape); + uint16_t aL1M = tla::get<0, 0>(srcShape) * tla::get<0, 1>(srcShape); + uint16_t wAlign = RoundUp(aL1M, HW_M0); + AscendC::Load3DSetFMatrixCal(1, wAlign, PAD_LIST); + // format(M, K), K_axis is k direction, and M_axis is m direction in load3d intrin + uint16_t madMAlign = tla::get<0, 0>(dstShape) * tla::get<0, 1>(dstShape); + uint16_t kStep = tla::get<1, 0>(dstShape) * tla::get<1, 1>(dstShape); + AscendC::LoadData3DParamsV2Pro loadData3DV2; + loadData3DV2.channelSize = aL1K; + loadData3DV2.extConfig = ((uint64_t)tla::get<0>(coord) << M_POS_BIT) | ((uint64_t)tla::get<1>(coord) << K_POS_BIT) | + ((uint64_t)madMAlign << M_STEP_BIT) | (uint64_t)kStep; + + AscendC::LoadData(dstTensor.data(), srcTensor.data(), loadData3DV2); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_ATLASA2_COPY_L1_TO_L0A_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_l0b.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_l0b.hpp new file mode 100644 index 0000000000..b94c228834 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_l0b.hpp @@ -0,0 +1,881 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ATLASA2_COPY_L1_TO_L0B_HPP +#define CATLASS_KDA_ATLASA2_COPY_L1_TO_L0B_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/kda_numeric_size.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "catlass/gemm/tile/kda_tile_copy_tla.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Tile { + +template < + class ArchTag, + class L1Type, + class L0Type = void> +struct CopyL1ToL0B { + static_assert(DEPENDENT_FALSE, "Unsupported copy l1 to l0, can not find the specialization."); +}; + +//////////////////////////////////////// +/// new add gemm +template +struct CopyL1ToL0B, Catlass::Gemm::GemmType> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + CATLASS_DEVICE + CopyL1ToL0B() {} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutSrc.orgShape(1))); + loadDataParams.srcStride = 1; + loadDataParams.sid = 0; + loadDataParams.dstGap = 0; + loadDataParams.ifTranspose = true; + loadDataParams.addrMode = 0; + for (uint32_t i = 0; i < CeilDiv(layoutDst.orgShape(0)); i++) { // K N + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; + +template +struct CopyL1ToL0B, Catlass::Gemm::GemmType> { + using Element = float; + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + + CATLASS_DEVICE + CopyL1ToL0B() {} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc) + { + AscendC::LoadData2dTransposeParams loadDataParams; + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutSrc.orgShape(1))); + loadDataParams.srcStride = 1; + loadDataParams.dstGap = 0; + loadDataParams.dstFracGap = static_cast(CeilDiv(layoutDst.orgShape(1))) - 1; + for (uint32_t i = 0; i < CeilDiv(layoutDst.orgShape(0)); i++) { // K N + AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1) * 2], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; + +template +struct CopyL1ToL0B, Catlass::Gemm::GemmType> { + using Element = int8_t; + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + CATLASS_DEVICE + CopyL1ToL0B() {} + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor dstTensor, + AscendC::LocalTensor srcTensor, + LayoutDst layoutDst, LayoutSrc layoutSrc) + { + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL / 2; + loadDataParams.dstGap = 1; + loadDataParams.dstFracGap = 0; + + for (uint32_t i = 0; i < CeilDiv(layoutDst.orgShape(0)); i++) { + AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1)], + srcTensor[i * layoutSrc.stride(1) * 2], + loadDataParams); + } + } +}; + +template +struct CopyL1ToL0B, Catlass::Gemm::GemmType> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(layoutDst.shape(3)); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(3) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + + for (uint32_t i = 0; i < layoutDst.shape(1); i++) { + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; +///////////////////////////////////////////// + +//////////////////////////////////////////// +/// new add gemv +template +struct CopyL1ToL0B, Catlass::Gemm::GemmType> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(layoutDst.shape(1)); + loadDataParams.srcStride = layoutSrc.stride(1) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(1) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + + for (uint32_t i = 0; i < layoutDst.shape(3); i++) { + AscendC::LoadData(dstTensor[i * layoutDst.stride(3)], srcTensor[i * layoutSrc.stride(3)], loadDataParams); + } + } +}; + +template +struct CopyL1ToL0B, Catlass::Gemm::GemmType> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::nN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = layoutDst.shape(1) * layoutDst.shape(3); + loadDataParams.srcStride = layoutSrc.stride(1) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(1) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = true; + loadDataParams.addrMode = 0; + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + }; +}; + +template +struct CopyL1ToL0B, Catlass::Gemm::GemmType> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::nN; + using Element = float; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(0))); + loadDataParams.srcStride = 1; + loadDataParams.dstGap = 0; + loadDataParams.dstFracGap = CeilDiv(layoutDst.orgShape(0)) - 1; + + for (uint32_t i = 0; i < CeilDiv<2 * ELE_NUM_PER_C0>(layoutDst.orgShape(1)); i++) { + AscendC::LoadDataWithTranspose( + dstTensor[i * layoutDst.stride(3) * 2], + srcTensor[i * layoutSrc.stride(3)], + loadDataParams); + } + }; +}; + +template +struct CopyL1ToL0B, Catlass::Gemm::GemmType> { + using LayoutDst = layout::zN; + using LayoutSrc = layout::nZ; + using Element = int8_t; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(0))); + loadDataParams.srcStride = layoutSrc.stride(1) / ELE_NUM_PER_FRACTAL / 2; + loadDataParams.dstGap = 1; + loadDataParams.dstFracGap = 0; + + for (uint32_t i = 0; i < CeilDiv(layoutDst.orgShape(1)); i++) { + AscendC::LoadDataWithTranspose( + dstTensor[i * layoutDst.stride(3)], + srcTensor[i * layoutSrc.stride(3) * 2], + loadDataParams); + } + } +}; +//////////////////////////////////////////// + +/// Partial specialization for int8_t, zN in and nZ out. +template +struct CopyL1ToL0B> { + using Element = int8_t; + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL / 2; + loadDataParams.dstGap = 1; + loadDataParams.dstFracGap = 0; + + for (uint32_t i = 0; i < CeilDiv(layoutDst.orgShape(0)); i++) { + AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1)], + srcTensor[i * layoutSrc.stride(1) * 2], + loadDataParams); + } + } +}; + +/// Partial specialization for float, zN in and nZ out. +template +struct CopyL1ToL0B> { + using Element = float; + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + constexpr uint8_t PAD_LIST[4] = {0, 0, 0, 0}; + uint16_t l1K = layoutSrc.shape(0) * layoutSrc.shape(1); + uint16_t l1N = layoutSrc.shape(2) * layoutSrc.shape(3); + uint16_t l0K = layoutDst.shape(0) * layoutDst.shape(1); + uint16_t l0N = layoutDst.shape(2) * layoutDst.shape(3); + // K, N need to be 16 aligned for f32 + uint16_t l1KAlign = RoundUp(l1K); + uint16_t l1NAlign = RoundUp(l1N); + uint16_t l0KAlign = RoundUp(l0K); + uint16_t l0NAlign = RoundUp(l0N); + AscendC::SetFmatrix(1, l1KAlign, PAD_LIST, AscendC::FmatrixMode::FMATRIX_RIGHT); + static constexpr AscendC::IsResetLoad3dConfig config = {false, false}; + AscendC::LoadData3DParamsV2 loadDataParams; + loadDataParams.kExtension = l0NAlign; + loadDataParams.mExtension = l0KAlign; + loadDataParams.channelSize = l1NAlign; + loadDataParams.fMatrixCtrl = true; + + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + } +}; + +/// Partial specialization for zN in and nZ out. +template +struct CopyL1ToL0B> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(3) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = true; + loadDataParams.addrMode = 0; + + for (uint32_t i = 0; i < CeilDiv(layoutDst.orgShape(0)); i++) { + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } +}; + +template +struct CopyL1ToL0B> { + using Element = AscendC::int4b_t; + using LayoutDst = layout::nZ; + using LayoutSrc = layout::zN; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(CeilDiv(layoutDst.orgShape(1))); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL / 4; + loadDataParams.dstGap = 3; + loadDataParams.dstFracGap = 0; + + for (uint32_t i = 0; i < CeilDiv(layoutDst.orgShape(0)); i++) { + AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1)], + srcTensor[i * layoutSrc.stride(1) * 4], + loadDataParams); + } + } +}; + +/// Partial specialization for nZ in and nZ out. (Transpose B) +template +struct CopyL1ToL0B> { + using LayoutDst = layout::nZ; + using LayoutSrc = layout::nZ; + + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + CopyL1ToL0B() {}; + + CATLASS_DEVICE + void operator()( + AscendC::LocalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + LayoutDst const &layoutDst, LayoutSrc const &layoutSrc) + { + AscendC::LoadData2DParams loadDataParams; + if (layoutSrc.shape(3) == layoutDst.shape(3)) { + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(layoutDst.shape(1) * layoutDst.shape(3)); + loadDataParams.srcStride = 1; + loadDataParams.sid = 0; + loadDataParams.dstGap = 0; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + + AscendC::LoadData(dstTensor, srcTensor, loadDataParams); + } else { + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(layoutDst.shape(3)); + loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = layoutDst.stride(3) / ELE_NUM_PER_FRACTAL - 1; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + + for (uint32_t i = 0; i < layoutDst.shape(1); i++) { + AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams); + } + } + } +}; + +///////////////////////////////////////////TileCopyTla////////////////////////////////////////////////////// +/// Partial specialization for CopyL1ToL0B, AtlasA2, zN in and nZ out. +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::B2>, + std::enable_if_t::value && + tla::detail::isnZ::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::iszN::value && + tla::detail::isnZ::value && + TensorSrc::position == AscendC::TPosition::A1 && + TensorDst::position == AscendC::TPosition::B2, + "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0B and nZ"); + + const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride()); + const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride()); + const uint32_t dstOuterShapeRow = CeilDiv(tla::get<0>(dstTensor.originShape()), tla::get<0, 0>(dstTensor.shape())); + const uint32_t dstOuterShapeCol = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape())); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + + AscendC::LoadData2DParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = dstOuterShapeCol; + loadDataParams.srcStride = srcOuterStrideCol / ELE_NUM_PER_FRACTAL; + loadDataParams.sid = 0; + loadDataParams.dstGap = 0; + loadDataParams.ifTranspose = true; + loadDataParams.addrMode = 0; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + for (uint32_t i = 0; i < dstOuterShapeRow; i++) { + AscendC::LoadData(dstTensor.data()[dstOffset + i * dstOuterStrideRow], + srcTensor.data()[srcOffset + i * srcOuterStrideRow], + loadDataParams); + } + } +}; + +/// Partial specialization for CopyL1ToL0B, AtlasA2, nZ in and nZ out. (Transpose B) +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::B2>, + std::enable_if_t::value && + tla::detail::isnZ::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isnZ::value && + tla::detail::isnZ::value && + TensorSrc::position == AscendC::TPosition::A1 && + TensorDst::position == AscendC::TPosition::B2, + "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0B and nZ"); + + const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride()); + const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride()); + const uint32_t dstOuterShapeRow = CeilDiv(tla::get<0>(dstTensor.originShape()), tla::get<0, 0>(dstTensor.shape())); + const uint32_t dstOuterShapeCol = CeilDiv(tla::get<1>(dstTensor.originShape()), tla::get<1, 0>(dstTensor.shape())); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::LoadData2DParams loadDataParams; + if (srcOuterStrideRow == dstOuterStrideRow) { + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = dstOuterShapeRow * dstOuterShapeCol; + loadDataParams.srcStride = 1; + loadDataParams.sid = 0; + loadDataParams.dstGap = 0; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + + AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], loadDataParams); + } else { + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = dstOuterShapeCol; + loadDataParams.srcStride = 1; + loadDataParams.sid = 0; + loadDataParams.dstGap = 0; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + + for (uint32_t i = 0; i < dstOuterShapeRow; i++) { + AscendC::LoadData(dstTensor.data()[dstOffset + i * dstOuterStrideRow], + srcTensor.data()[srcOffset + i * srcOuterStrideRow], + loadDataParams); + } + } + } +}; + +/// Partial specialization for CopyL1ToL0B, AtlasA2, int8_t, zN in and nZ out. +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::B2>, + std::enable_if_t::value && tla::detail::isnZ::value>> { + using Element = int8_t; + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(std::is_same_v && + std::is_same_v && + tla::detail::iszN::value && + tla::detail::isnZ::value && + TensorSrc::position == AscendC::TPosition::A1 && + TensorDst::position == AscendC::TPosition::B2, + "The input parameters do not match. TensorSrc must be int8_t, L1 and zN, " + "while TensorDst must be int8_t, L0B and nZ"); + + const uint32_t srcOuterShapeCol = CeilDiv(tla::get<1>(srcTensor.originShape()), tla::get<1, 0>(srcTensor.shape())); + const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride()); + const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride()); + const uint32_t dstOuterShapeRow = CeilDiv(tla::get<0>(dstTensor.originShape()), tla::get<0, 0>(dstTensor.shape())); + const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride()); + + AscendC::LoadData2dTransposeParams loadDataParams; + + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = srcOuterShapeCol; + loadDataParams.srcStride = srcOuterStrideCol / ELE_NUM_PER_FRACTAL / 2; + loadDataParams.dstGap = 1; + loadDataParams.dstFracGap = 0; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + for (uint32_t i = 0; i < dstOuterShapeRow; i++) { + AscendC::LoadDataWithTranspose(dstTensor.data()[dstOffset + i * dstOuterStrideRow], + srcTensor.data()[srcOffset + i * srcOuterStrideRow * 2], + loadDataParams); + } + } +}; + +/// Partial specialization for CopyL1ToL0B, AtlasA2, float, zN in and nZ out. +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::B2>, + std::enable_if_t::value && tla::detail::isnZ::value>> { + using Element = float; + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(std::is_same_v && + std::is_same_v && + tla::detail::iszN::value && + tla::detail::isnZ::value && + TensorSrc::position == AscendC::TPosition::A1 && + TensorDst::position == AscendC::TPosition::B2, + "The input parameters do not match. TensorSrc must be float, L1 and zN, " + "while TensorDst must be float, L0B and nZ"); + + constexpr uint8_t PAD_LIST[4] = {0, 0, 0, 0}; + uint16_t l1K = tla::get<1, 1>(srcTensor.stride()) / tla::get<1, 0>(srcTensor.shape()); + uint16_t l1N = tla::get<1>(srcTensor.originShape()); + uint16_t l0K = tla::get<0>(dstTensor.originShape()); + uint16_t l0N = tla::get<1>(dstTensor.originShape()); + // K, N need to be 16 aligned for f32 + uint16_t l1KAlign = RoundUp(l1K); + uint16_t l1NAlign = RoundUp(l1N); + uint16_t l0KAlign = RoundUp(l0K); + uint16_t l0NAlign = RoundUp(l0N); + AscendC::SetFmatrix(1, l1KAlign, PAD_LIST, AscendC::FmatrixMode::FMATRIX_RIGHT); + static constexpr AscendC::IsResetLoad3dConfig config = {false, false}; + AscendC::LoadData3DParamsV2 loadDataParams; + loadDataParams.kExtension = l0NAlign; + loadDataParams.mExtension = l0KAlign; + loadDataParams.channelSize = l1NAlign; + loadDataParams.fMatrixCtrl = true; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], loadDataParams); + } +}; + +/// Partial specialization for CopyL1ToL0B, AtlasA2, zN in and nZ out. +template +struct CopyL1ToL0BSparseTla, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::B2>, + tla::Tensor, LayoutIdx, CoordIdx, AscendC::TPosition::A1>, + std::enable_if_t::value && + tla::detail::isnZ::value && + tla::detail::iszN::value>> { + static constexpr uint8_t INDEX_SHIFT = 2; ///< Shift value for index manipulation + + // Methods + + CATLASS_DEVICE + CopyL1ToL0BSparseTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, TensorIdx const &idxTensor) + { + static_assert(tla::detail::iszN::value && + tla::detail::isnZ::value && + TensorSrc::position == AscendC::TPosition::A1 && + TensorDst::position == AscendC::TPosition::B2, + "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0B and nZ"); + + auto dstShape = dstTensor.shape(); + auto dstStride = dstTensor.stride(); + auto srcShape = srcTensor.shape(); + auto srcStride = srcTensor.stride(); + auto coord = srcTensor.coord(); + + // SET LOAD2D parameters , loop axis: K or M, or 1 + uint16_t madK = tla::get<1, 0>(dstShape) * tla::get<1, 1>(dstShape); + + // k is c0Size_ aligned for f32 + uint16_t kC0 = CeilDiv(madK, tla::get<1, 0>(dstShape)); + uint16_t nFraC0 = tla::get<0, 1>(dstShape); + uint16_t l0bLoop = 1; + uint64_t l0bDstAddrStride = 0; + uint64_t l0bSrcAddrStride = 0; + uint8_t l0bRepeat = kC0 * nFraC0; + uint16_t l0bDststride = 0; + uint16_t l0bSrcstride = 1; + + uint16_t bL1N = tla::get<0, 0>(srcShape) * tla::get<0, 1>(srcShape); + if (nFraC0 * tla::get<0, 0>(dstShape) == bL1N) { // loop=1 + l0bLoop = 1; + } else if (nFraC0 >= kC0) { // LOOP is K and repeat is n axis + l0bLoop = kC0; + l0bDstAddrStride = nFraC0 * tla::get<1, 1>(dstStride); + l0bSrcAddrStride = bL1N * tla::get<1, 0>(srcShape); + l0bRepeat = nFraC0; + + l0bDststride = 0; + l0bSrcstride = 1; + } else { // LOOP is N and repeat is K axis + l0bLoop = nFraC0; + l0bDstAddrStride = tla::get<1, 1>(dstStride); + l0bSrcAddrStride = tla::get<0, 1>(srcStride); + l0bRepeat = kC0; + + l0bDststride = nFraC0 - 1; + l0bSrcstride = bL1N; + } + // use load2d for L1_2_L0B + // startIndex, repeatTimes, srcStride, sid, dstGap, ifTranspose, addrmode + AscendC::LoadData2dParams loadDataParams{0, l0bRepeat, l0bSrcstride, 0, l0bDststride, 0, 0}; + uint16_t bL1KOffset = tla::get<1>(coord); + uint64_t l1bOffset = tla::get<0>(coord) * tla::get<1, 0>(srcShape) + bL1KOffset * bL1N; + uint64_t l0bOffset = 0; + AscendC::LocalTensor dstLocal; + AscendC::LocalTensor srcLocal; + AscendC::LocalTensor l1BIndex; + for (uint64_t i = 0; i < l0bLoop; i++) { + dstLocal.SetAddr(dstTensor.data()[l0bOffset].address_); + srcLocal.SetAddr(srcTensor.data()[l1bOffset].address_); + l1BIndex.SetAddr(idxTensor.data()[l1bOffset >> INDEX_SHIFT].address_); + AscendC::LoadDataWithSparse(dstLocal, srcLocal, l1BIndex, loadDataParams); + l1bOffset += l0bSrcAddrStride; + l0bOffset += l0bDstAddrStride; + } + } +}; + +// Partial specialization for CopyL1ToL0B, AtlasA2, nZ in and nZ out. (Transpose B) +template +struct CopyL1ToL0BSparseTla, LayoutSrc, CoordSrc, AscendC::TPosition::A1>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::B2>, + tla::Tensor, LayoutIdx, CoordIdx, AscendC::TPosition::A1>, + std::enable_if_t::value && + tla::detail::isnZ::value && + tla::detail::isnZ::value>> { + static constexpr uint8_t INDEX_SHIFT = 2; ///< Shift value for index manipulation + + // Methods + + CATLASS_DEVICE + CopyL1ToL0BSparseTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, TensorIdx const &idxTensor) + { + static_assert(tla::detail::isnZ::value && + tla::detail::isnZ::value && + TensorSrc::position == AscendC::TPosition::A1 && + TensorDst::position == AscendC::TPosition::B2, + "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0B and nZ"); + + auto dstShape = dstTensor.shape(); + auto dstStride = dstTensor.stride(); + auto srcShape = srcTensor.shape(); + auto srcStride = srcTensor.stride(); + auto coord = srcTensor.coord(); + + // SET LOAD2D parameters , loop axis: K or M, or 1 + uint16_t madK = tla::get<0, 0>(dstShape) * tla::get<0, 1>(dstShape); + + // k is c0Size_ aligned for f32 + uint16_t kC0 = CeilDiv(madK, tla::get<0, 0>(dstShape)); + uint16_t nFraC0 = tla::get<1, 1>(dstShape); + uint16_t l0bLoop = 1; + uint64_t l0bDstAddrStride = 0; + uint64_t l0bSrcAddrStride = 0; + uint8_t l0bRepeat = kC0 * nFraC0; + + AscendC::LoadData2DParams loadDataParams; + loadDataParams.startIndex = 0; + loadDataParams.repeatTimes = static_cast(l0bRepeat); + loadDataParams.srcStride = 1; + loadDataParams.sid = 0; + loadDataParams.dstGap = 0; + loadDataParams.ifTranspose = false; + loadDataParams.addrMode = 0; + uint16_t bL1N = tla::get<1, 0>(srcShape) * tla::get<1, 1>(srcShape); + uint16_t bL1KOffset = tla::get<0>(coord); + uint64_t l1bOffset = tla::get<1>(coord) * tla::get<0, 0>(srcShape) + bL1KOffset * bL1N; + + AscendC::LocalTensor dstLocal; + AscendC::LocalTensor srcLocal; + AscendC::LocalTensor l1BIndex; + dstLocal.SetAddr(dstTensor.data().address_); + srcLocal.SetAddr(srcTensor.data()[l1bOffset].address_); + l1BIndex.SetAddr(idxTensor.data()[l1bOffset >> INDEX_SHIFT].address_); + + AscendC::LoadDataWithSparse(dstLocal, srcLocal, l1BIndex, loadDataParams); + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_ATLASA2_COPY_L1_TO_L0B_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_ub_to_gm.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_ub_to_gm.hpp new file mode 100644 index 0000000000..4d558811d3 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_atlasa2_copy_ub_to_gm.hpp @@ -0,0 +1,131 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_ATLASA2_COPY_UB_TO_GM_HPP +#define CATLASS_KDA_ATLASA2_COPY_UB_TO_GM_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_arch.hpp" +#include "catlass/gemm/tile/kda_tile_copy_tla.hpp" +#include "catlass/gemm/kda_gemm_type.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Tile { + +/// Partial specialization for AtlasA2, RowMajor in and RowMajor out. +template +struct TileCopyTla, LayoutSrc, CoordSrc, AscendC::TPosition::VECCALC>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::GM>, + std::enable_if_t::value && + tla::detail::isRowMajor::value>> { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc); + + // Methods + + CATLASS_DEVICE + TileCopyTla() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isRowMajor::value && + tla::detail::isRowMajor::value && + TensorSrc::position == AscendC::TPosition::VECCALC && + TensorDst::position == AscendC::TPosition::GM, + "The input parameters do not match. TensorSrc must be GM and RowMajor, " + "while TensorDst must be UB and RowMajor"); + + const uint16_t row = tla::get<0>(dstTensor.originShape()); + const uint16_t col = tla::get<1>(dstTensor.originShape()); + + AscendC::DataCopyExtParams dataCopyParams( + row, + col * sizeof(ElementSrc), + (tla::get<0>(srcTensor.stride()) - col) / ELE_NUM_PER_C0, + (tla::get<0>(dstTensor.stride()) - col) * sizeof(ElementSrc), + 0); + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + AscendC::DataCopyPad(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], dataCopyParams); + }; +}; + +/// Partial specialization for AtlasA2, RowMajor in and PaddingRowMajor out. +template +struct TileCopyTlaExt, LayoutSrc, CoordSrc, AscendC::TPosition::VECCALC>, + tla::Tensor, LayoutDst, CoordDst, AscendC::TPosition::GM>, + layout::RowMajor, layout::PaddingRowMajor> { + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc); + + // Methods + + CATLASS_DEVICE + TileCopyTlaExt() {}; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor) + { + static_assert(tla::detail::isRowMajor::value && + TensorSrc::position == AscendC::TPosition::VECCALC && + TensorDst::position == AscendC::TPosition::GM, + "The input parameters do not match. TensorSrc must be GM and PaddingRowMajor, " + "while TensorDst must be UB and RowMajor"); + + AscendC::DataCopyExtParams dataCopyParams( + tla::get<1, 1>(dstTensor.shape()), + tla::get<1, 0>(dstTensor.shape()) * sizeof(ElementSrc), + (tla::get<0>(srcTensor.stride()) - tla::get<1>(srcTensor.shape())) / ELE_NUM_PER_C0, + (tla::get<1, 1>(dstTensor.stride()) - tla::get<1, 0>(dstTensor.shape())) * sizeof(ElementSrc), + 0); + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + AscendC::DataCopyPad(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], dataCopyParams); + }; +}; + +template < + class ArchTag, + class GmType> +struct CopyUb2Gm { + static_assert(DEPENDENT_FALSE, "Unsupported copy ub to gm, can not find the specialization."); +}; + +template +struct CopyUb2Gm> { + using LayoutDst = layout::RowMajor; + using LayoutSrc = layout::RowMajor; + + static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element); + + CATLASS_DEVICE + CopyUb2Gm() = default; + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &dstTensor, + AscendC::LocalTensor const &srcTensor, + layout::RowMajor const &layoutDst, + layout::RowMajor const &layoutSrc) + { + AscendC::DataCopyExtParams dataCopyParams( + layoutDst.shape(0), + layoutDst.shape(1) * sizeof(Element), + (layoutSrc.stride(0) - layoutSrc.shape(1)) / ELE_NUM_PER_C0, + (layoutDst.stride(0) - layoutDst.shape(1)) * sizeof(Element), + 0); + AscendC::DataCopyPad(dstTensor, srcTensor, dataCopyParams); + } +}; + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_ATLASA2_COPY_UB_TO_GM_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_cast_fp8_to_fp16.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_cast_fp8_to_fp16.hpp new file mode 100644 index 0000000000..9ffb0a69ff --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_cast_fp8_to_fp16.hpp @@ -0,0 +1,415 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_TILE_ATLASA2_KDA_CAST_FP8_TO_FP16_HPP +#define CATLASS_GEMM_TILE_ATLASA2_KDA_CAST_FP8_TO_FP16_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" + +namespace Catlass::Gemm::Tile { + +template < + class ArchTag, + class SrcType_, + class DstType_, + uint32_t COMPUTE_LENGTH> +struct TileCastFp8ToFp16Dequant { + using ElementSrc = typename SrcType_::Element; + using ElementDst = typename DstType_::Element; + using LayoutSrc = typename SrcType_::Layout; + using LayoutDst = typename DstType_::Layout; + using LayoutRowMajor = Catlass::layout::RowMajor; + + static_assert(std::is_same_v || std::is_same_v, + "Unsupported layout, only can be Row/Column Major."); + static_assert(std::is_same_v, "layout src and dst must be the same."); + + static const uint32_t Alignment = 256; + // static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(int8_t); + + struct Params { + half scalar; + half zeroPoint; + + CATLASS_HOST_DEVICE + Params() = default; + + CATLASS_HOST_DEVICE + Params(half scalar_, half zeroPoint_) + { + scalar = scalar_; + zeroPoint = zeroPoint_; + } + }; + + CATLASS_DEVICE + TileCastFp8ToFp16Dequant() + {} + + CATLASS_DEVICE + TileCastFp8ToFp16Dequant(Arch::Resource &resource, Params const ¶ms_) + : params(params_) + { + int64_t bufferOffset = 0; + for (uint32_t i = 0; i < BUFFER_NUM; i++) { + inputBuffer[i] = resource.ubBuf.template GetBufferByByte(bufferOffset * sizeof(ElementSrc)); + bufferOffset += COMPUTE_LENGTH; + } + for (uint32_t i = 0; i < BUFFER_NUM; i++) { + outputBuffer[i] = (resource.ubBuf.template GetBufferByByte(bufferOffset * sizeof(ElementSrc))) + .template ReinterpretCast(); + bufferOffset += COMPUTE_LENGTH * 2; + } + for (uint32_t i = 0; i < BUFFER_NUM; i++) { + workspace[i] = (resource.ubBuf.template GetBufferByByte(bufferOffset * sizeof(ElementSrc))) + .template ReinterpretCast(); + bufferOffset += COMPUTE_LENGTH * 2; + } + InitLocalMaskVec(resource, bufferOffset); + } + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor gmDst, LayoutDst const &layoutDst, + AscendC::GlobalTensor gmSrc, LayoutSrc const &layoutSrc, + uint32_t &bufferIndex) + { + uint32_t tilesNum, tileLen, srcStride, dstStride; + if constexpr (std::is_same_v) { + tilesNum = layoutSrc.shape(0); + tileLen = layoutSrc.shape(1); + srcStride = layoutSrc.stride(0); + dstStride = layoutDst.stride(0); + } else if constexpr (std::is_same_v) { + tilesNum = layoutSrc.shape(1); + tileLen = layoutSrc.shape(0); + srcStride = layoutSrc.stride(1); + dstStride = layoutDst.stride(1); + } + + uint32_t tilesPerAiv = tilesNum / AscendC::GetSubBlockNum(); + uint32_t tilesRemain = tilesNum % AscendC::GetSubBlockNum(); + if (AscendC::GetSubBlockIdx() < tilesRemain) { + tilesPerAiv++; + } + + uint64_t taskSrcOffset = AscendC::GetSubBlockIdx() * tilesPerAiv * srcStride; + uint64_t taskDstOffset = AscendC::GetSubBlockIdx() * tilesPerAiv * dstStride; + if (AscendC::GetSubBlockIdx() >= tilesRemain) { + taskSrcOffset += tilesRemain * srcStride; + taskDstOffset += tilesRemain * dstStride; + } + + uint32_t totalLoops = 0; + uint32_t loopsPerTile, tilesInALoop; + uint32_t tileLenRoundFp8 = RoundUp(tileLen); + if (tileLenRoundFp8 > COMPUTE_LENGTH / 2) { + // One signle tile length is bigger than COMPUTE_LENGTH, which should be clipped. + loopsPerTile = CeilDiv(tileLen, COMPUTE_LENGTH); + totalLoops = tilesPerAiv * loopsPerTile; + } else if (tileLenRoundFp8 != 0) { + // COMPUTE_LENGTH is bigger than tile length, such that more than one tiles can be arranged together. + tilesInALoop = COMPUTE_LENGTH / tileLenRoundFp8; + totalLoops = CeilDiv(tilesPerAiv, tilesInALoop); + } // tileLenRoundFp8 == 0 --> totalLoops = 0 + + uint32_t tileTailLen = tileLen % COMPUTE_LENGTH; + uint64_t srcProcessOffset, dstProcessOffset; + uint32_t loadLen = COMPUTE_LENGTH, storeLen, loadRepeat = 1, storeRepeat = 1; + // uint32_t srcLoadStride = srcStride, dstLoadStride = tileLenRoundFp8, + // srcStoreStride = tileLenRoundFp8, dstStoreStride = dstStride; + for (int ldx = 0; ldx < totalLoops; ldx++) { + // Dynamic compute length + if (tileLenRoundFp8 > COMPUTE_LENGTH / 2) { + uint32_t fullTileRounds = ldx / loopsPerTile; + uint32_t residueTileRounds = ldx % loopsPerTile; + srcProcessOffset = taskSrcOffset + fullTileRounds * srcStride + residueTileRounds * COMPUTE_LENGTH; + dstProcessOffset = taskDstOffset + fullTileRounds * dstStride + residueTileRounds * COMPUTE_LENGTH; + + if ((residueTileRounds == loopsPerTile - 1) && (tileTailLen != 0)) { + loadLen = tileTailLen; + } + } else { + uint32_t fullTileRounds = ldx * tilesInALoop; + srcProcessOffset = taskSrcOffset + fullTileRounds * srcStride; + dstProcessOffset = taskDstOffset + fullTileRounds * dstStride; + + loadLen = tileLen; + storeLen = tileLen; + loadRepeat = tilesInALoop; + storeRepeat = tilesInALoop; + + if ((ldx == totalLoops - 1) && (tilesPerAiv % tilesInALoop != 0)) { + loadRepeat = tilesPerAiv % tilesInALoop; + storeRepeat = loadRepeat; + } + } + // uint32_t srcLoadStride = srcStride; + // uint32_t dstLoadStride = tileLenRoundFp8; + // uint32_t srcStoreStride = tileLenRoundFp8; + + // GM -> UB + AscendC::DataCopyExtParams dataCopyParamsIn( + loadRepeat, + loadLen * sizeof(ElementSrc), + (srcStride - loadLen) * sizeof(ElementSrc), // + (tileLenRoundFp8 - loadLen) * sizeof(ElementSrc) / BYTE_PER_BLK, + 0); + AscendC::DataCopyPadExtParams padParams(false, 0, 0, 0); + + AscendC::WaitFlag(EventIdBuffer[bufferIndex]); + AscendC::DataCopyPad(inputBuffer[bufferIndex], gmSrc[srcProcessOffset], + dataCopyParamsIn, padParams); + + AscendC::SetFlag(EventIdBuffer[bufferIndex]); + AscendC::WaitFlag(EventIdBuffer[bufferIndex]); + AscendC::WaitFlag(EventIdBuffer[bufferIndex]); + + Dequant(inputBuffer[bufferIndex], + outputBuffer[bufferIndex], + value_vector1, + value_vector2, + workspace[bufferIndex]); + AscendC::SetFlag(EventIdBuffer[bufferIndex]); + AscendC::WaitFlag(EventIdBuffer[bufferIndex]); + AscendC::SetFlag(EventIdBuffer[bufferIndex]); + + AscendC::DataCopyExtParams dataCopyParams( + storeRepeat, + storeLen * sizeof(ElementDst), + (tileLenRoundFp8 - storeLen) * sizeof(ElementDst) / BYTE_PER_C0, + (dstStride - storeLen) * sizeof(ElementDst), + 0); + AscendC::DataCopyPad(gmDst[dstProcessOffset], outputBuffer[bufferIndex], dataCopyParams); + AscendC::SetFlag(EventIdBuffer[bufferIndex]); + bufferIndex = (bufferIndex + 1) % BUFFER_NUM; + } + } + + /* + * DATAFLOW: gmSrc -> ubInTensor -> ubOutTensor -> gmDst + */ + CATLASS_DEVICE + void EpCastFp32ToFp16( + AscendC::GlobalTensor gmDst, LayoutRowMajor layoutDst, + AscendC::GlobalTensor gmSrc, LayoutRowMajor layoutSrc) + { + AscendC::LocalTensor ubInTensor[BUFFER_NUM]; + AscendC::LocalTensor ubOutTensor[BUFFER_NUM]; + + Arch::Resource resource; + int64_t bufferOffset = 0; + const int64_t CAST_LENGTH = 32 * 1024 / sizeof(half); // 一次处理16K个数据 + for (int i = 0; i < BUFFER_NUM; i++) { + ubInTensor[i] = resource.ubBuf.template GetBufferByByte(bufferOffset); + bufferOffset += CAST_LENGTH * 4; // float 4字节 + } + for (int i = 0; i < BUFFER_NUM; i++) { + ubOutTensor[i] = resource.ubBuf.template GetBufferByByte(bufferOffset); + bufferOffset += CAST_LENGTH * 2; // half 2字节 + } + + uint32_t tilesNum, tileLen, srcStride, dstStride; + tilesNum = layoutSrc.shape(0); + tileLen = layoutSrc.shape(1); + srcStride = layoutSrc.stride(0); + dstStride = layoutDst.stride(0); // Always RowMajor + + uint32_t tilesPerAiv = tilesNum / AscendC::GetSubBlockNum(); + uint32_t tilesRemain = tilesNum % AscendC::GetSubBlockNum(); + if (AscendC::GetSubBlockIdx() < tilesRemain) { + tilesPerAiv++; + } + + uint64_t taskSrcOffset = AscendC::GetSubBlockIdx() * tilesPerAiv * srcStride; + uint64_t taskDstOffset = AscendC::GetSubBlockIdx() * tilesPerAiv * dstStride; + if (AscendC::GetSubBlockIdx() >= tilesRemain) { + taskSrcOffset += tilesRemain * srcStride; + taskDstOffset += tilesRemain * dstStride; + } + + uint32_t totalLoops = 0; + uint32_t loopsPerTile, tilesInALoop; + uint32_t tileLenRoundFp8 = RoundUp(tileLen); + if (tileLenRoundFp8 > COMPUTE_LENGTH / 2) { + // One signle tile length is bigger than COMPUTE_LENGTH, which should be clipped. + loopsPerTile = CeilDiv(tileLen, COMPUTE_LENGTH); + totalLoops = tilesPerAiv * loopsPerTile; + } else if (tileLenRoundFp8 != 0) { + // COMPUTE_LENGTH is bigger than tile length, such that more than one tiles can be arranged together. + tilesInALoop = COMPUTE_LENGTH / tileLenRoundFp8; + totalLoops = CeilDiv(tilesPerAiv, tilesInALoop); + } // tileLenRoundFp8 == 0 --> totalLoops = 0 + + uint32_t tileTailLen = tileLen % COMPUTE_LENGTH; + uint64_t srcProcessOffset, dstProcessOffset; + uint32_t loadLen = COMPUTE_LENGTH, storeLen, loadRepeat = 1, storeRepeat = 1; + // uint32_t srcLoadStride = srcStride, dstLoadStride = tileLenRoundFp8, + // srcStoreStride = tileLenRoundFp8, dstStoreStride = dstStride; + for (int ldx = 0; ldx < totalLoops; ldx++) { + // Dynamic compute length + if (tileLenRoundFp8 > COMPUTE_LENGTH / 2) { + uint32_t fullTileRounds = ldx / loopsPerTile; + uint32_t residueTileRounds = ldx % loopsPerTile; + srcProcessOffset = taskSrcOffset + fullTileRounds * srcStride + residueTileRounds * COMPUTE_LENGTH; + dstProcessOffset = taskDstOffset + fullTileRounds * dstStride + residueTileRounds * COMPUTE_LENGTH; + + if ((residueTileRounds == loopsPerTile - 1) && (tileTailLen != 0)) { + loadLen = tileTailLen; + } + } else { + uint32_t fullTileRounds = ldx * tilesInALoop; + srcProcessOffset = taskSrcOffset + fullTileRounds * srcStride; + dstProcessOffset = taskDstOffset + fullTileRounds * dstStride; + + loadLen = tileLen; + storeLen = tileLen; + loadRepeat = tilesInALoop; + storeRepeat = tilesInALoop; + + if ((ldx == totalLoops - 1) && (tilesPerAiv % tilesInALoop != 0)) { + loadRepeat = tilesPerAiv % tilesInALoop; + storeRepeat = loadRepeat; + } + } + + // copy GM -> UB + AscendC::WaitFlag(EventIdBufferForCast[bufferIndexForCast]); + AscendC::DataCopyExtParams dataCopyParamsIn( + loadRepeat, + loadLen * sizeof(float), // float loaded + (srcStride - loadLen) * sizeof(float), + (tileLenRoundFp8 - loadLen) * sizeof(float) / BYTE_PER_BLK, + 0); + AscendC::DataCopyPadExtParams padParams(false, 0, 0, 0); + AscendC::DataCopyPad(ubInTensor[bufferIndexForCast], gmSrc[srcProcessOffset], + dataCopyParamsIn, padParams); + // Begin casting ... + AscendC::SetFlag(EventIdBufferForCast[bufferIndexForCast]); + AscendC::WaitFlag(EventIdBufferForCast[bufferIndexForCast]); + AscendC::WaitFlag(EventIdBufferForCast[bufferIndexForCast]); + AscendC::Cast(ubOutTensor[bufferIndexForCast], ubInTensor[bufferIndexForCast], + AscendC::RoundMode::CAST_RINT, CAST_LENGTH); + AscendC::SetFlag(EventIdBufferForCast[bufferIndexForCast]); + AscendC::WaitFlag(EventIdBufferForCast[bufferIndexForCast]); + AscendC::SetFlag(EventIdBufferForCast[bufferIndexForCast]); + // End casting ... + + AscendC::DataCopyExtParams dataCopyParams( + storeRepeat, + storeLen * sizeof(half), + (tileLenRoundFp8 - storeLen) * sizeof(half) / BYTE_PER_C0, + (dstStride - storeLen) * sizeof(half), + 0); + AscendC::DataCopyPad(gmDst[dstProcessOffset], ubOutTensor[bufferIndexForCast], dataCopyParams); + AscendC::SetFlag(EventIdBufferForCast[bufferIndexForCast]); + bufferIndexForCast = (bufferIndexForCast + 1) % BUFFER_NUM; + } + } + +private: + CATLASS_DEVICE + void InitLocalMaskVec(Arch::Resource &resource, int64_t &bufferOffset) + { + int16_t value_uint = 0x4000; + value_vector1 = (resource.ubBuf.template GetBufferByByte(bufferOffset * sizeof(ElementSrc))) + .template ReinterpretCast(); + bufferOffset += 256; + AscendC::Duplicate(value_vector1, value_uint, 128); + AscendC::PipeBarrier(); + // AscendC::PipeBarrier(); + value_uint = 0x3FFF; + value_vector2 = (resource.ubBuf.template GetBufferByByte(bufferOffset * sizeof(ElementSrc))) + .template ReinterpretCast(); + bufferOffset += 256; + AscendC::Duplicate(value_vector2, value_uint, 128); + AscendC::PipeBarrier(); + } + + CATLASS_DEVICE + void Dequant(AscendC::LocalTensor &src, AscendC::LocalTensor &dst, + AscendC::LocalTensor &value_vector1, AscendC::LocalTensor &value_vector2, + AscendC::LocalTensor &workspace) + { + pipe_barrier(PIPE_V); + uint32_t num = COMPUTE_LENGTH; + num = (num + 128 - 1) / 128 * 128; + AscendC::Cast(dst.template ReinterpretCast(), + src.template ReinterpretCast(), + AscendC::RoundMode::CAST_NONE, + num); + pipe_barrier(PIPE_V); + + AscendC::Adds(dst, dst, 1024, num); + pipe_barrier(PIPE_V); + + AscendC::ShiftLeft( + dst.template ReinterpretCast(), dst.template ReinterpretCast(), 7, num); + pipe_barrier(PIPE_V); + + uint64_t mask = 128; + AscendC::And(workspace.template ReinterpretCast(), + dst.template ReinterpretCast(), + value_vector1, + mask, + num / 128, + {1, 1, 1, 8, 8, 0}); + pipe_barrier(PIPE_V); + + AscendC::ShiftLeft( + workspace.template ReinterpretCast(), workspace.template ReinterpretCast(), 1, num); + pipe_barrier(PIPE_V); + + AscendC::And(dst.template ReinterpretCast(), + dst.template ReinterpretCast(), + value_vector2, + mask, + num / 128, + {1, 1, 1, 8, 8, 0}); + pipe_barrier(PIPE_V); + + AscendC::Or(dst.template ReinterpretCast(), + dst.template ReinterpretCast(), + workspace.template ReinterpretCast(), + num); + pipe_barrier(PIPE_V); + + AscendC::Muls(dst.template ReinterpretCast(), dst.template ReinterpretCast(), 1 << 8, num); + pipe_barrier(PIPE_V); + + AscendC::Adds(dst, dst, params.zeroPoint, num); + pipe_barrier(PIPE_V); + + AscendC::Muls(dst, dst, params.scalar, num); + pipe_barrier(PIPE_V); + } + +private: + static const uint32_t BUFFER_NUM = 2; + AscendC::LocalTensor inputBuffer[BUFFER_NUM]; + AscendC::LocalTensor value_vector1; + AscendC::LocalTensor value_vector2; + AscendC::LocalTensor outputBuffer[BUFFER_NUM]; + AscendC::LocalTensor workspace[BUFFER_NUM]; + AscendC::TEventID EventIdBuffer[BUFFER_NUM] = {EVENT_ID0, EVENT_ID1}; + AscendC::TEventID EventIdBufferForCast[BUFFER_NUM] = {EVENT_ID2, EVENT_ID3}; + uint32_t bufferIndexForCast{0}; + + Params params; +}; + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_GEMM_TILE_ATLASA2_KDA_CAST_FP8_TO_FP16_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_cast_int8_to_fp16.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_cast_int8_to_fp16.hpp new file mode 100644 index 0000000000..e68d4bf327 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/atlasa2/kda_cast_int8_to_fp16.hpp @@ -0,0 +1,188 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_TILE_ATLASA2_KDA_CAST_INT8_TO_FP16_HPP +#define CATLASS_GEMM_TILE_ATLASA2_KDA_CAST_INT8_TO_FP16_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" + +namespace Catlass::Gemm::Tile { + +template < + class ArchTag, + class SrcType_, + class DstType_, + // Length of the compute elements + uint32_t COMPUTE_LEN_, + uint32_t STAGES = 2> +struct TileCastInt8ToFp16Dequant { + using ElementSrc = typename SrcType_::Element; + using ElementDst = typename DstType_::Element; + using LayoutSrc = typename SrcType_::Layout; + using LayoutDst = typename DstType_::Layout; + + static_assert(std::is_same_v || std::is_same_v, + "Unsupported layout, only can be RowMajor and ColumnMajor"); + static_assert(std::is_same_v, "layoutDst and layoutSrc must be the same"); + + static constexpr uint32_t ELE_NUM_PER_BLK_INT8 = BYTE_PER_BLK / sizeof(ElementSrc); + static constexpr uint32_t ELE_NUM_PER_BLK_HALF = BYTE_PER_BLK / sizeof(ElementDst); + + static constexpr uint32_t COMPUTE_LEN = COMPUTE_LEN_; + + static_assert(COMPUTE_LEN <= 32 * 1024, "COMPUTE_LEN cannot exceed 32 * 1024"); + + struct Params { + half deqScalar; + half deqZeroPoint; + + CATLASS_HOST_DEVICE + Params() = default; + + CATLASS_HOST_DEVICE + Params(half deqScalar_, half deqZeroPoint_) + { + deqScalar = deqScalar_; + deqZeroPoint = deqZeroPoint_; + } + }; + + /// Construct + CATLASS_DEVICE + TileCastInt8ToFp16Dequant(Arch::Resource const &resource, Params const ¶ms_) + : params(params_) + { + if constexpr (g_coreType == AscendC::AIV) { + uint32_t ubOffset = 0; + uint32_t ubInSize = COMPUTE_LEN * sizeof(ElementSrc); + uint32_t ubOutSize = COMPUTE_LEN * sizeof(ElementDst); + // Init buffers + for (uint32_t i = 0; i < STAGES; i++) { + ubInTensorList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += ubInSize; + ubOutTensorList[i] = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += ubOutSize; + + ubEventList[i] = i; + AscendC::SetFlag(ubEventList[i]); + AscendC::SetFlag(ubEventList[i]); + } + } + } + + /// Destructor + CATLASS_DEVICE + ~TileCastInt8ToFp16Dequant() + { + if constexpr (g_coreType == AscendC::AIV) { + for (uint32_t i = 0; i < STAGES; i++) { + AscendC::WaitFlag(ubEventList[i]); + AscendC::WaitFlag(ubEventList[i]); + } + } + } + + CATLASS_DEVICE + void operator()( + AscendC::GlobalTensor const &gmDst, LayoutDst const &layoutDst, + AscendC::GlobalTensor const &gmSrc, LayoutSrc const &layoutSrc) + { + uint32_t tilesNum = layoutSrc.shape(0); + uint32_t tileLen = layoutSrc.shape(1); + uint32_t tileLenRoundInt8 = RoundUp(layoutSrc.shape(1), ELE_NUM_PER_BLK_INT8); + uint64_t tileStrideSrc = layoutSrc.stride(0); + uint64_t tileStrideDst = layoutDst.stride(0); + if constexpr (std::is_same_v) { + tilesNum = layoutSrc.shape(1); + tileLen = layoutSrc.shape(0); + tileLenRoundInt8 = RoundUp(layoutSrc.shape(0), ELE_NUM_PER_BLK_INT8); + tileStrideSrc = layoutSrc.stride(1); + tileStrideDst = layoutDst.stride(1); + } + uint32_t tilesPerAiv = tilesNum / AscendC::GetSubBlockNum(); + if (AscendC::GetSubBlockIdx() < (tilesNum % AscendC::GetSubBlockNum())) { + tilesPerAiv++; + } + uint64_t taskOffsetSrc = AscendC::GetSubBlockIdx() * tilesPerAiv * tileStrideSrc; + uint64_t taskOffsetDst = AscendC::GetSubBlockIdx() * tilesPerAiv * tileStrideDst; + if (AscendC::GetSubBlockIdx() >= (tilesNum % AscendC::GetSubBlockNum())) { + taskOffsetSrc += (tilesNum % AscendC::GetSubBlockNum()) * tileStrideSrc; + taskOffsetDst += (tilesNum % AscendC::GetSubBlockNum()) * tileStrideDst; + } + uint32_t tilesPerLoop = COMPUTE_LEN / tileLenRoundInt8; + uint32_t loops = CeilDiv(tilesPerAiv, tilesPerLoop); + uint32_t pingpong = 0; + for (uint32_t loopIdx = 0; loopIdx < loops; loopIdx++) { + uint32_t actualTiles = tilesPerLoop; + if (loopIdx == loops - 1) { + actualTiles = tilesPerAiv - loopIdx * tilesPerLoop; + } + uint64_t tileOffsetSrc = loopIdx * tilesPerLoop * tileStrideSrc; + AscendC::DataCopyExtParams dataCopyParamsIn( + actualTiles, tileLen * sizeof(ElementSrc), + (tileStrideSrc - tileLen) * sizeof(ElementSrc), + (tileLenRoundInt8 - tileLen) / ELE_NUM_PER_BLK_INT8, 0); + AscendC::DataCopyPadExtParams padParams(false, 0, 0, 0); + + AscendC::WaitFlag(ubEventList[pingpong]); + AscendC::DataCopyPad(ubInTensorList[pingpong], + gmSrc[taskOffsetSrc + tileOffsetSrc], dataCopyParamsIn, padParams); + + AscendC::SetFlag(ubEventList[pingpong]); + AscendC::WaitFlag(ubEventList[pingpong]); + + AscendC::WaitFlag(ubEventList[pingpong]); + + AscendC::Cast(ubOutTensorList[pingpong], ubInTensorList[pingpong], + AscendC::RoundMode::CAST_NONE, actualTiles * tileLenRoundInt8); + AscendC::PipeBarrier(); + AscendC::SetFlag(ubEventList[pingpong]); + + AscendC::Adds(ubOutTensorList[pingpong], ubOutTensorList[pingpong], + params.deqZeroPoint, actualTiles * tileLenRoundInt8); + AscendC::PipeBarrier(); + + AscendC::Muls(ubOutTensorList[pingpong], ubOutTensorList[pingpong], + params.deqScalar, actualTiles * tileLenRoundInt8); + AscendC::PipeBarrier(); + + AscendC::SetFlag(ubEventList[pingpong]); + AscendC::WaitFlag(ubEventList[pingpong]); + + uint64_t tileOffsetDst = loopIdx * tilesPerLoop * tileStrideDst; + AscendC::DataCopyExtParams dataCopyParamsOut( + actualTiles, tileLen * sizeof(ElementDst), + (tileLenRoundInt8 - tileLen) / ELE_NUM_PER_BLK_HALF, + (tileStrideDst - tileLen) * sizeof(ElementDst), 0); + AscendC::DataCopyPad(gmDst[taskOffsetDst + tileOffsetDst], + ubOutTensorList[pingpong], dataCopyParamsOut); + AscendC::SetFlag(ubEventList[pingpong]); + + pingpong = (pingpong + 1) % STAGES; + } + } + +protected: + /// Data members + AscendC::LocalTensor ubInTensorList[STAGES]; + AscendC::LocalTensor ubOutTensorList[STAGES]; + int32_t ubEventList[STAGES]; + + Params params; +}; + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_GEMM_TILE_ATLASA2_KDA_CAST_INT8_TO_FP16_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/copy_l0c_to_ub.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/copy_l0c_to_ub.hpp new file mode 100644 index 0000000000..785a5c0408 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/copy_l0c_to_ub.hpp @@ -0,0 +1,18 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_TILE_COPY_L0C_TO_UB_HPP +#define CATLASS_GEMM_TILE_COPY_L0C_TO_UB_HPP + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) +#include "catlass/gemm/tile/ascend950/copy_l0c_to_ub.hpp" +#endif + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_cast_fp8_to_fp16.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_cast_fp8_to_fp16.hpp new file mode 100644 index 0000000000..95e6fa7197 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_cast_fp8_to_fp16.hpp @@ -0,0 +1,18 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_CAST_FP8_TO_FP16_HPP +#define CATLASS_KDA_CAST_FP8_TO_FP16_HPP + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) +#include "catlass/gemm/tile/atlasa2/kda_cast_fp8_to_fp16.hpp" +#endif + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_cast_int4_to_int8.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_cast_int4_to_int8.hpp new file mode 100644 index 0000000000..75c265b5af --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_cast_int4_to_int8.hpp @@ -0,0 +1,18 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_CAST_INT4_TO_INT8_HPP +#define CATLASS_KDA_CAST_INT4_TO_INT8_HPP + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) +#include "catlass/gemm/tile/atlasa2/kda_atlasa2_cast_int4_to_int8.hpp" +#endif + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_cast_int8_to_fp16.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_cast_int8_to_fp16.hpp new file mode 100644 index 0000000000..e99e12754d --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_cast_int8_to_fp16.hpp @@ -0,0 +1,18 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_CAST_INT8_TO_FP16_HPP +#define CATLASS_KDA_CAST_INT8_TO_FP16_HPP + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) +#include "catlass/gemm/tile/atlasa2/kda_cast_int8_to_fp16.hpp" +#endif + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_gm_to_l1.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_gm_to_l1.hpp new file mode 100644 index 0000000000..6097bc8cb6 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_gm_to_l1.hpp @@ -0,0 +1,21 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_COPY_GM_TO_L1_HPP +#define CATLASS_KDA_COPY_GM_TO_L1_HPP + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) +#include "catlass/gemm/tile/atlasa2/kda_atlasa2_copy_gm_to_l1.hpp" +#endif +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) +#include "catlass/gemm/tile/ascend950/kda_ascend950_copy_gm_to_l1.hpp" +#endif + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_gm_to_ub.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_gm_to_ub.hpp new file mode 100644 index 0000000000..013d24d116 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_gm_to_ub.hpp @@ -0,0 +1,18 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_GEMM_COPY_GM_TO_UB_HPP +#define CATLASS_KDA_GEMM_COPY_GM_TO_UB_HPP + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) +#include "catlass/gemm/tile/atlasa2/kda_atlasa2_copy_gm_to_ub.hpp" +#endif + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l0c_to_gm.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l0c_to_gm.hpp new file mode 100644 index 0000000000..6bfa7b1945 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l0c_to_gm.hpp @@ -0,0 +1,21 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_COPY_L0C_TO_GM_HPP +#define CATLASS_KDA_COPY_L0C_TO_GM_HPP + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) +#include "catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l0c_to_gm.hpp" +#endif +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) +#include "catlass/gemm/tile/ascend950/kda_ascend950_copy_l0c_to_gm.hpp" +#endif + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_bt.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_bt.hpp new file mode 100644 index 0000000000..ba44dbd032 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_bt.hpp @@ -0,0 +1,21 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_COPY_L1_TO_BT_HPP +#define CATLASS_KDA_COPY_L1_TO_BT_HPP + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) +#include "catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_bt.hpp" +#endif +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) +#include "catlass/gemm/tile/ascend950/kda_ascend950_copy_l1_to_bt.hpp" +#endif + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_fp.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_fp.hpp new file mode 100644 index 0000000000..547871d161 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_fp.hpp @@ -0,0 +1,18 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_COPY_L1_TO_FP_HPP +#define CATLASS_KDA_COPY_L1_TO_FP_HPP + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) +#include "catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_fp.hpp" +#endif + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_l0a.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_l0a.hpp new file mode 100644 index 0000000000..327897533a --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_l0a.hpp @@ -0,0 +1,21 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_COPY_L1_TO_L0A_HPP +#define CATLASS_KDA_COPY_L1_TO_L0A_HPP + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) +#include "catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_l0a.hpp" +#endif +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) +#include "catlass/gemm/tile/ascend950/kda_ascend950_copy_l1_to_l0a.hpp" +#endif + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_l0b.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_l0b.hpp new file mode 100644 index 0000000000..97ca958a9a --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_l1_to_l0b.hpp @@ -0,0 +1,21 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_COPY_L1_TO_L0B_HPP +#define CATLASS_KDA_COPY_L1_TO_L0B_HPP + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) +#include "catlass/gemm/tile/atlasa2/kda_atlasa2_copy_l1_to_l0b.hpp" +#endif +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) +#include "catlass/gemm/tile/ascend950/kda_ascend950_copy_l1_to_l0b.hpp" +#endif + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_ub_to_gm.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_ub_to_gm.hpp new file mode 100644 index 0000000000..02d9303a06 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_copy_ub_to_gm.hpp @@ -0,0 +1,18 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_GEMM_COPY_UB_TO_GM_HPP +#define CATLASS_KDA_GEMM_COPY_UB_TO_GM_HPP + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) +#include "catlass/gemm/tile/atlasa2/kda_atlasa2_copy_ub_to_gm.hpp" +#endif + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_gemm_tile_copy.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_gemm_tile_copy.hpp new file mode 100644 index 0000000000..9d463a5a6c --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_gemm_tile_copy.hpp @@ -0,0 +1,520 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_GEMM_TILE_COPY_HPP +#define CATLASS_KDA_GEMM_TILE_COPY_HPP + +#include + +#include "catlass/kda_catlass.hpp" +#include "catlass/detail/kda_tag_to_layout.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_cast_fp8_to_fp16.hpp" +#include "catlass/gemm/tile/kda_cast_int4_to_int8.hpp" +#include "catlass/gemm/tile/kda_cast_int8_to_fp16.hpp" +#include "catlass/gemm/tile/kda_copy_gm_to_l1.hpp" +#include "catlass/gemm/tile/kda_copy_gm_to_ub.hpp" +#include "catlass/gemm/tile/kda_copy_l0c_to_gm.hpp" +#include "catlass/gemm/tile/copy_l0c_to_ub.hpp" +#include "catlass/gemm/tile/kda_copy_l1_to_bt.hpp" +#include "catlass/gemm/tile/kda_copy_l1_to_fp.hpp" +#include "catlass/gemm/tile/kda_copy_l1_to_l0a.hpp" +#include "catlass/gemm/tile/kda_copy_l1_to_l0b.hpp" +#include "catlass/gemm/tile/kda_copy_ub_to_gm.hpp" +#include "catlass/gemm/tile/kda_tile_copy_tla.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm::Tile { + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) +template < + /// Tag indicating architecture + class ArchTag, + /// GemmType for A matrix operand + class AType, + /// GemmType type for B matrix operand + class BType, + /// GemmType type for C matrix operand + class CType, + /// GemmType type for Bias operand + class BiasType = void> +struct TileCopy { + using ElementA = typename AType::Element; + using ElementB = typename BType::Element; + using ElementAccumulator = typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + + using CopyGmToL1A = Gemm::Tile::CopyGmToL1; + using CopyGmToL1B = Gemm::Tile::CopyGmToL1; + using CopyL1ToL0A = Gemm::Tile::CopyL1ToL0A::L1AType>; + using CopyL1ToL0B = Gemm::Tile::CopyL1ToL0B::L1BType>; + using CopyL0CToGm = Gemm::Tile::CopyL0CToGm; + using BiasTypeSelector = helper::L1BiasTypeSelector; + using CopyGmToL1Bias = std::conditional_t< + std::is_same_v, + void, + Gemm::Tile::CopyGmToL1>; + using CopyL1ToBT = std::conditional_t< + std::is_same_v, + void, + Gemm::Tile::CopyL1ToBT>; +}; + +template +struct TileCopyWithPrologueDeqPerTensor { + using ElementA = typename AType::Element; + using ElementB = typename BType::Element; + + using ElementAccumulator = typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + + using CopyGmToL1A = Gemm::Tile::CopyGmToL1; + using CopyGmToL1B = Gemm::Tile::CopyGmToL1; + + using PrologueA = PrologueA_; + using PrologueB = PrologueB_; + + using CopyL1ToL0A = Gemm::Tile::CopyL1ToL0A::L1AType>; + using CopyL1ToL0B = Gemm::Tile::CopyL1ToL0B::L1BType>; + using CopyL0CToGm = Gemm::Tile::CopyL0CToGm; + using BiasTypeSelector = helper::L1BiasTypeSelector; + using CopyGmToL1Bias = std::conditional_t< + std::is_same_v, + void, + Gemm::Tile::CopyGmToL1>; + using CopyL1ToBT = std::conditional_t< + std::is_same_v, + void, + Gemm::Tile::CopyL1ToBT>; +}; + +template +struct TileCopyWithPrologue { + using ElementA = typename AType::Element; + using ElementB = typename BType::Element; + + using ElementAccumulator = typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + + using CopyGmToL1A = Gemm::Tile::CopyGmToL1; + using CopyGmToL1B = Gemm::Tile::CopyGmToL1; + + using PrologueA = PrologueA_; + using PrologueB = PrologueB_; + + using CopyL1ToL0A = Gemm::Tile::CopyL1ToL0A::L1AType>; + using CopyL1ToL0B = Gemm::Tile::CopyL1ToL0B::L1BType>; + using CopyL0CToGm = Gemm::Tile::CopyL0CToGm; + using BiasTypeSelector = helper::L1BiasTypeSelector; + using CopyGmToL1Bias = std::conditional_t< + std::is_same_v, + void, + Gemm::Tile::CopyGmToL1>; + using CopyL1ToBT = std::conditional_t< + std::is_same_v, + void, + Gemm::Tile::CopyL1ToBT>; +}; + +/////////////////////////////////// +/// new add +template < + /// Tag indicating architecture + class ArchTag, + /// GemmType for A matrix operand + class AType, + /// GemmType type for B matrix operand + class BType, + /// GemmType type for C matrix operand + class CType, + /// GemmTpe type for Bias operand + class BiasType = void> +struct TileCopyGemm { + using ElementA = typename AType::Element; + using ElementB = typename BType::Element; + using ElementAccumulator = typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + // change structual + using L1AType = typename Gemm::helper::L1AndL0TypeSelectorGemm::L1AType; + using L1BType = typename Gemm::helper::L1AndL0TypeSelectorGemm::L1BType; + using L0AType = typename Gemm::helper::L1AndL0TypeSelectorGemm::L0AType; + using L0BType = typename Gemm::helper::L1AndL0TypeSelectorGemm::L0BType; + + using CopyGmToL1A = Gemm::Tile::CopyGmToL1; + using CopyGmToL1B = Gemm::Tile::CopyGmToL1; + using CopyL1ToL0A = Gemm::Tile::CopyL1ToL0A; + using CopyL1ToL0B = Gemm::Tile::CopyL1ToL0B; + using CopyL0CToGm = Gemm::Tile::CopyL0CToGm; +}; +////////////////////////////// +template < + /// Tag indicating architecture + class ArchTag, + /// GemmType for A matrix operand + class AType, + /// GemmType type for B matrix operand + class BType, + /// GemmType type for C matrix operand + class CType, + /// GemmType type for Bias operand + class BiasType> +struct ConvTileCopy { + using ElementA = typename AType::Element; + using ElementB = typename BType::Element; + using ElementAccumulator = typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + + using CopyGmToL1A = Gemm::Tile::CopyGmToL1; + using CopyGmToL1B = Gemm::Tile::CopyGmToL1; + using CopyL1ToL0A = Gemm::Tile::CopyL1ToL0A::L1AType>; + using CopyL1ToL0B = Gemm::Tile::CopyL1ToL0B::L1BType>; + using CopyL0CToGm = Gemm::Tile::CopyL0CToGm; + using BiasTypeSelector = helper::L1BiasTypeSelector; + using CopyGmToL1Bias = std::conditional_t< + std::is_same_v, + void, + Gemm::Tile::CopyGmToL1>; + using CopyL1ToBT = std::conditional_t< + std::is_same_v, + void, + Gemm::Tile::CopyL1ToBT>; +}; + +// fixpipe开启relu开关 +template < + /// Tag indicating architecture + class ArchTag, + /// GemmType for A matrix operand + class AType, + /// GemmType type for B matrix operand + class BType, + /// GemmType type for C matrix operand + class CType, + /// GemmType type for Bias operand + class BiasType = void> +struct ReluTileCopy : public TileCopy { + // 重写 CopyL0CToGm + using ElementAccumulator = typename TileCopy::ElementAccumulator; + using CopyL0CToGm = Gemm::Tile:: + CopyL0CToGm; +}; + +// fixpipe开启随路量化 +template < + /// Tag indicating architecture + class ArchTag, + /// GemmType for A matrix operand + class AType, + /// GemmType type for B matrix operand + class BType, + /// GemmType type for C matrix operand + class CType, + /// GemmType type for Bias operand + class BiasType = void, + /// GemmType type for Bias operand + ScaleGranularity SCALE_GRANU = ScaleGranularity::PER_TENSOR> +struct QuantTileCopy : public TileCopy { + // 重写 CopyL0CToGm + using ElementAccumulator = typename TileCopy::ElementAccumulator; + using CopyL0CToGm = Gemm::Tile::CopyL0CToGm; + + using CopyGmToL1Scale = Gemm::Tile::CopyGmToL1< + ArchTag, + Gemm::GemmType, + Gemm::GemmType>; + + using CopyL1ToFP = Gemm::Tile::CopyL1ToFP< + ArchTag, + Gemm::GemmType, + Gemm::GemmType>; +}; + +// sparse Tile +template < + /// Tag indicating architecture + class ArchTag, + class ElementA_, + class LayoutTagA, + class ElementB_, + class LayoutTagB, + class ElementC_, + class LayoutTagC> +struct SparseTileCopyTla { + using ElementA = ElementA_; + using ElementB = ElementB_; + using ElementAccumulator = typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + + using LayoutTagL1A = typename helper::L1ATypeSelector>::L1AType::Layout; + using LayoutTagL1B = typename helper::L1BTypeSelector>::L1BType::Layout; + using LayoutTagL0A = layout::zZ; + using LayoutTagL0B = layout::nZ; + using LayoutTagL1BIdx = typename helper::L1BTypeSelector>::L1BType::Layout; + + using LayoutA = detail::TagToLayout_t; + using LayoutB = detail::TagToLayout_t; + using LayoutC = detail::TagToLayout_t; + + using LayoutL1A = detail::TagToLayout_t; + using LayoutL1B = detail::TagToLayout_t; + using LayoutL1BIdx = detail::TagToLayout_t; + using LayoutL0A = detail::TagToLayout_t; + using LayoutL0B = detail::TagToLayout_t; + using LayoutL0C = typename detail::LayoutL0C; + + using TensorL1A = + tla::Tensor, LayoutL1A, tla::Coord, AscendC::TPosition::A1>; + using TensorL1B = + tla::Tensor, LayoutL1B, tla::Coord, AscendC::TPosition::A1>; + using TensorL1BIdx = + tla::Tensor, LayoutL1BIdx, tla::Coord, AscendC::TPosition::A1>; + using TensorL0A = + tla::Tensor, LayoutL0A, tla::Coord, AscendC::TPosition::A2>; + using TensorL0B = + tla::Tensor, LayoutL0B, tla::Coord, AscendC::TPosition::B2>; + using TensorL0C = tla:: + Tensor, LayoutL0C, tla::Coord, AscendC::TPosition::CO1>; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + template + using CopyGmToL1A = Gemm::Tile::TileCopySparseTla; + + template + using CopyGmToL1B = Gemm::Tile::TileCopySparseTla; + + template + using CopyGmToL1BIdx = Gemm::Tile::TileCopySparseTla; + + using CopyL1ToL0A = Gemm::Tile::TileCopySparseTla; + + using CopyL1ToL0B = Gemm::Tile::CopyL1ToL0BSparseTla; + + template + using CopyL0CToGm = Gemm::Tile::CopyL0CToGmSparseTla; +}; + +#endif + +template < + /// Tag indicating architecture + class ArchTag, + class ElementA_, + class LayoutTagA_, + class ElementB_, + class LayoutTagB_, + class ElementC_, + class LayoutTagC_, + class ElementBias = void, + bool ReluEnable_ = false, + ScaleGranularity DEQUANT_GRANULARITY_ = ScaleGranularity::NO_QUANT, + class L0CCopyMode = CopyToGM> +struct PackedTileCopyTla { + using ElementA = ElementA_; + using ElementB = ElementB_; + using LayoutTagA = LayoutTagA_; + using LayoutTagB = LayoutTagB_; + using LayoutTagC = LayoutTagC_; + using ElementAccumulator = typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + static constexpr bool ReluEnable = ReluEnable_; + static constexpr ScaleGranularity DEQUANT_GRANULARITY = DEQUANT_GRANULARITY_; + + static constexpr bool HAS_BIAS = !std::is_void_v; + static constexpr bool HAS_QUANT_TENSOR = (DEQUANT_GRANULARITY == ScaleGranularity::PER_CHANNEL); + + using LayoutTagL1A = typename helper::L1ATypeSelector>::L1AType::Layout; + using LayoutTagL1B = typename helper::L1BTypeSelector>::L1BType::Layout; + using LayoutTagL0A = typename helper::L0ALayoutSelector::Layout; + using LayoutTagL0B = layout::nZ; + using LayoutTagL0C = layout::L0C; + + using LayoutA = detail::TagToLayout_t; + using LayoutB = detail::TagToLayout_t; + using LayoutC = detail::TagToLayout_t; + + using LayoutL1A = detail::TagToLayout_t; + using LayoutL1B = detail::TagToLayout_t; + using LayoutL0A = detail::TagToLayout_t; + using LayoutL0B = detail::TagToLayout_t; + using LayoutL0C = typename detail::LayoutL0C; + + using TensorL1AVectorLayout = + tla::Tensor, LayoutL1A, tla::Coord, AscendC::TPosition::A1>; + using TensorL1ALayout = + tla::Tensor, LayoutL1A, tla::Coord, AscendC::TPosition::A1>; + + using TensorL1A = + std::conditional_t::value, TensorL1AVectorLayout, TensorL1ALayout>; + using TensorL1B = + tla::Tensor, LayoutL1B, tla::Coord, AscendC::TPosition::A1>; + using TensorL0A = + tla::Tensor, LayoutL0A, tla::Coord, AscendC::TPosition::A2>; + using TensorL0B = + tla::Tensor, LayoutL0B, tla::Coord, AscendC::TPosition::B2>; + using TensorL0C = tla:: + Tensor, LayoutL0C, tla::Coord, AscendC::TPosition::CO1>; + using TensorL1Bias = std::conditional_t< + HAS_BIAS, + tla::Tensor< + AscendC::LocalTensor, + detail::TagToLayout_t, + tla::Coord, + AscendC::TPosition::A1>, + EmptyClass>; + using TensorL0Bias = tla::Tensor< + AscendC::LocalTensor, + detail::TagToLayout_t, + tla::Coord, + AscendC::TPosition::C2>; + using TensorL1Quant = std::conditional_t< + HAS_QUANT_TENSOR, + tla::Tensor, + detail::TagToLayout_t, + tla::Coord, + AscendC::TPosition::A1>, + EmptyClass>; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + template + using CopyGmToL1A = Gemm::Tile::TileCopyTla; + + template + using CopyGmToL1B = Gemm::Tile::TileCopyTla; + + template + using CopyGmToL1Bias = + std::conditional_t, EmptyClass>; + + template + using CopyGmToL1Scale = + std::conditional_t, EmptyClass>; + + using CopyL1ToL0A = Gemm::Tile::TileCopyTla; + using CopyL1ToL0B = Gemm::Tile::TileCopyTla; + using CopyL1ToBT = + std::conditional_t, EmptyClass>; + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) + template + using CopyL0CToDst = Gemm::Tile::CopyL0CToGmTla; +#endif +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) + template + using CopyL0CToGm = Gemm::Tile::CopyL0CToGmTla; +#endif +}; + +template < + /// Tag indicating architecture + class ArchTag, + class TensorA, + class LayoutTagA, + class TensorB, + class LayoutTagB, + class TensorC, + class LayoutTagC, + class TensorBias = void, + class LayoutTagBias = void, + bool IS_PADDING_A = false, + bool IS_PADDING_B = false> +struct PaddingPackedTileCopyTla { + static_assert( + std::is_same_v || std::is_same_v, + "Unsupported layout, only can be RowMajor and ColumnMajor"); + static_assert( + std::is_same_v || std::is_same_v, + "Unsupported layout, only can be RowMajor and ColumnMajor"); + using ElementA = typename TensorA::Element; + using ElementB = typename TensorB::Element; + using ElementAccumulator = typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + + using LayoutTagL1A = typename helper::L1ATypeSelector>::L1AType::Layout; + using LayoutTagL1B = typename helper::L1BTypeSelector>::L1BType::Layout; + using LayoutTagL0A = layout::zZ; + using LayoutTagL0B = layout::nZ; + + using LayoutL1A = detail::TagToLayout_t; + using LayoutL1B = detail::TagToLayout_t; + using LayoutL0A = detail::TagToLayout_t; + using LayoutL0B = detail::TagToLayout_t; + using LayoutL0C = typename detail::LayoutL0C; + + using TensorL1A = + tla::Tensor, LayoutL1A, tla::Coord, AscendC::TPosition::A1>; + using TensorL1B = + tla::Tensor, LayoutL1B, tla::Coord, AscendC::TPosition::A1>; + using TensorL0A = + tla::Tensor, LayoutL0A, tla::Coord, AscendC::TPosition::A2>; + using TensorL0B = + tla::Tensor, LayoutL0B, tla::Coord, AscendC::TPosition::B2>; + using TensorL0C = tla:: + Tensor, LayoutL0C, tla::Coord, AscendC::TPosition::CO1>; + + using L1AAlignHelper = Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Gemm::helper::L1AlignHelper; + + using LayoutPaddingTagA = std:: + conditional_t, layout::PaddingRowMajor, layout::PaddingColumnMajor>; + using LayoutPaddingTagB = std:: + conditional_t, layout::PaddingRowMajor, layout::PaddingColumnMajor>; + + using CopyGmToL1A = std::conditional_t< + IS_PADDING_A, + Gemm::Tile::TileCopyTlaExt, + Gemm::Tile::TileCopyTlaExt>; + using CopyGmToL1B = std::conditional_t< + IS_PADDING_B, + Gemm::Tile::TileCopyTlaExt, + Gemm::Tile::TileCopyTlaExt>; + + using CopyL1ToL0A = Gemm::Tile::TileCopyTla; + using CopyL1ToL0B = Gemm::Tile::TileCopyTla; +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) + using CopyL0CToDst = Gemm::Tile::CopyL0CToGmTla; +#endif +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) + using CopyL0CToGm = Gemm::Tile::CopyL0CToGmTla; +#endif +}; + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) +template < + /// Tag indicating architecture + class ArchTag, + class ElementA_, + class LayoutTagA, + class ElementB_, + class LayoutTagB, + class ElementC_, + class LayoutTagC, + class ElementBias = void, + CopyL0CToUBMode CopyMode_ = CopyL0CToUBMode::NO_SPLIT, + bool ReluEnable = false, + ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT> +struct PackedTileCopyTlaToUB + : public PackedTileCopyTla { + static constexpr CopyL0CToUBMode CopyMode = CopyMode_; + // 重写 CopyL0CToDst + using TensorL0C = typename PackedTileCopyTla< + ArchTag, + ElementA_, + LayoutTagA, + ElementB_, + LayoutTagB, + ElementC_, + LayoutTagC, + ElementBias>::TensorL0C; + + template + using CopyL0CToDst = + Gemm::Tile::CopyL0CToUBTla; +}; +#endif + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_GEMM_TILE_COPY_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_copy_tla.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_copy_tla.hpp new file mode 100644 index 0000000000..46e2406361 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_copy_tla.hpp @@ -0,0 +1,70 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_TILE_COPY_TLA_HPP +#define CATLASS_KDA_TILE_COPY_TLA_HPP + +#include "catlass/kda_catlass.hpp" + +namespace Catlass::Gemm::Tile { + +template < + class ArchTag, + class TensorSrc, + class TensorDst, + class Enable = void> +struct TileCopyTla { + static_assert(DEPENDENT_FALSE, "Unsupported TileCopyTla, can not find the specialization."); +}; + +// Extended template for TileCopyTla that supports manually specifying LayoutTagSrc and LayoutTagDst. +// Users can specialize the copy class by LayoutTagSrc and LayoutTagDst. +template < + class ArchTag, + class TensorSrc, + class TensorDst, + class LayoutTagSrc, + class LayoutTagDst> +struct TileCopyTlaExt { + static_assert(DEPENDENT_FALSE, "Unsupported TileCopyTlaExt, can not find the specialization."); +}; + +template < + class ArchTag, + class TensorSrc, + class TensorDst, + class Enable = void> +struct TileCopySparseTla { + static_assert(DEPENDENT_FALSE, "Unsupported TileCopySparseTla, can not find the specialization."); +}; + +template < + class ArchTag, + class Element, + class TensorSrc, + class TensorDst, + class TensorIdx, + class Enable = void> +struct CopyL1ToL0BSparseTla { + static_assert(DEPENDENT_FALSE, "Unsupported CopyL1ToL0BSparseTla, can not find the specialization."); +}; + +// used for FA LoadQ, multi matrix datacopy ND->NZ +template < + class ArchTag, + class TensorSrc, + class TensorDst> +struct TileCopyFAQTla { + static_assert(DEPENDENT_FALSE, "Unsupported TileCopyFAQTla, can not find the specialization."); +}; + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_TILE_COPY_TLA_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_mmad.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_mmad.hpp new file mode 100644 index 0000000000..cca5eca621 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/gemm/tile/kda_tile_mmad.hpp @@ -0,0 +1,285 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_TILE_MMAD_HPP +#define CATLASS_KDA_TILE_MMAD_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "tla/kda_tensor.hpp" +namespace Catlass::Gemm::Tile { + +/////////////////////////////////////////////////////////// + +template < + /// Tag indicating architecture + class ArchTag_, + /// GemmType for A matrix operand + class AType_, + /// GemmType type for B matrix operand + class BType_, + /// GemmType type for Bias operand + class BiasType_> +struct TileMmad { + using ElementA = typename AType_::Element; + using ElementB = typename BType_::Element; + using ElementAccumulator = + typename Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + + // Methods + + CATLASS_DEVICE + TileMmad() {} + + CATLASS_DEVICE + void operator()(AscendC::LocalTensor const &l0CTensor, + AscendC::LocalTensor const &l0ATensor, + AscendC::LocalTensor const &l0BTensor, + uint32_t m, uint32_t n, uint32_t k, + bool initC = true, uint8_t unitFlag = 0) + { + AscendC::MmadParams mmadParams; + mmadParams.m = m; + mmadParams.n = n; + mmadParams.k = k; + mmadParams.unitFlag = unitFlag; + mmadParams.cmatrixInitVal = initC; +#if (defined(__NPU_ARCH__) && __NPU_ARCH__ == 2201) + if constexpr (std::is_same_v && + (std::is_same_v || + std::is_same_v)) { + mmadParams.kDirectionAlign = true; + } +#endif +#if (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3510) + if constexpr (std::is_same_v) { + mmadParams.disableGemv = false; + } else { + mmadParams.disableGemv = true; + } +#endif + + AscendC::Mmad(l0CTensor, + l0ATensor, + l0BTensor, + mmadParams); + + const uint32_t PIPE_M_BARRIER_THRESHOLD = 10; + if ((m / C0_NUM_PER_FRACTAL) * (n / C0_NUM_PER_FRACTAL) < PIPE_M_BARRIER_THRESHOLD) { + AscendC::PipeBarrier(); + } + } + + CATLASS_DEVICE + void operator()(AscendC::LocalTensor const &l0CTensor, + AscendC::LocalTensor const &l0ATensor, + AscendC::LocalTensor const &l0BTensor, + AscendC::LocalTensor const &l0BiasTensor, + uint32_t m, uint32_t n, uint32_t k, + bool initC = true, uint8_t unitFlag = 0) + { + AscendC::MmadParams mmadParams; + mmadParams.m = m; + mmadParams.n = n; + mmadParams.k = k; + mmadParams.unitFlag = unitFlag; + mmadParams.cmatrixInitVal = false; +#if (defined(__NPU_ARCH__) && __NPU_ARCH__ == 2201) + using LayoutL1A = typename helper::L1ATypeSelector::L1AType::Layout; + if constexpr (std::is_same_v && std::is_same_v) { + mmadParams.kDirectionAlign = true; + } +#endif +#if (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3510) + mmadParams.disableGemv = true; +#endif + + AscendC::Mmad(l0CTensor, + l0ATensor, + l0BTensor, + l0BiasTensor, + mmadParams); + + const uint32_t PIPE_M_BARRIER_THRESHOLD = 10; + if ((m / C0_NUM_PER_FRACTAL) * (n / C0_NUM_PER_FRACTAL) < PIPE_M_BARRIER_THRESHOLD) { + AscendC::PipeBarrier(); + } + } +}; + +///////////////////////////////////////////TileMmadTla///////////////////////////////////////////////// + +template < + /// Tag indicating architecture + class ArchTag, + /// Element for A matrix operand + class ElementA, + /// LayoutTag for A matrix operand in L1 + class LayoutTagL1A> +struct TileMmadTla { + // Methods + + CATLASS_DEVICE + TileMmadTla() {} + + template + CATLASS_DEVICE void operator()(TensorC const &l0CTensor, + TensorA const &l0ATensor, + TensorB const &l0BTensor, + uint32_t m, uint32_t n, uint32_t k, + bool initC = true, uint8_t unitFlag = 0) + { + AscendC::MmadParams mmadParams; + mmadParams.m = m; + mmadParams.n = n; + mmadParams.k = k; + mmadParams.unitFlag = unitFlag; + mmadParams.cmatrixInitVal = initC; +#if (defined(__NPU_ARCH__) && __NPU_ARCH__ == 2201) + if constexpr (std::is_same_v && std::is_same_v) { + mmadParams.kDirectionAlign = true; + } +#endif +#if (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3510) + if constexpr (std::is_same_v) { + mmadParams.disableGemv = false; + } else { + mmadParams.disableGemv = true; + } +#endif + + AscendC::Mmad(l0CTensor.data(), + l0ATensor.data(), + l0BTensor.data(), + mmadParams); + + const uint32_t PIPE_M_BARRIER_THRESHOLD = 10; + if ((m / C0_NUM_PER_FRACTAL) * (n / C0_NUM_PER_FRACTAL) < PIPE_M_BARRIER_THRESHOLD) { + AscendC::PipeBarrier(); + } + } + + template + CATLASS_DEVICE void operator()(TensorC const &l0CTensor, + TensorA const &l0ATensor, + TensorB const &l0BTensor, + TensorBias const &l0BiasTensor, + uint32_t m, uint32_t n, uint32_t k, + bool initC = true, uint8_t unitFlag = 0) + { + AscendC::MmadParams mmadParams; + mmadParams.m = m; + mmadParams.n = n; + mmadParams.k = k; + mmadParams.unitFlag = unitFlag; + mmadParams.cmatrixInitVal = false; +#if (defined(__NPU_ARCH__) && __NPU_ARCH__ == 2201) + if constexpr (std::is_same_v && std::is_same_v) { + mmadParams.kDirectionAlign = true; + } +#endif +#if (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3510) + mmadParams.disableGemv = true; +#endif + + AscendC::Mmad(l0CTensor.data(), + l0ATensor.data(), + l0BTensor.data(), + l0BiasTensor.data(), + mmadParams); + + const uint32_t PIPE_M_BARRIER_THRESHOLD = 10; + if ((m / C0_NUM_PER_FRACTAL) * (n / C0_NUM_PER_FRACTAL) < PIPE_M_BARRIER_THRESHOLD) { + AscendC::PipeBarrier(); + } + } + + template + CATLASS_DEVICE void operator()(TensorC const &l0CTensor, + TensorA const &l0ATensor, + TensorB const &l0BTensor, + uint32_t m, uint32_t n, uint32_t k, + uint32_t l0Batch) + { + const uint32_t L0AM = tla::get<0, 0>(l0ATensor.shape()) * tla::get<0, 1>(l0ATensor.shape()); + const uint32_t L0AK = tla::get<1, 0>(l0ATensor.shape()) * tla::get<1, 1>(l0ATensor.shape()); + const uint32_t L0BK = tla::get<0, 0>(l0BTensor.shape()) * tla::get<0, 1>(l0BTensor.shape()); + const uint32_t L0BN = tla::get<1, 0>(l0BTensor.shape()) * tla::get<1, 1>(l0BTensor.shape()); + const uint32_t L0CM = tla::get<0, 0>(l0CTensor.shape()) * tla::get<0, 1>(l0CTensor.shape()); + const uint32_t L0CN = tla::get<1, 0>(l0CTensor.shape()) * tla::get<1, 1>(l0CTensor.shape()); + + AscendC::MmadParams mmadParams; + mmadParams.m = m; + mmadParams.n = n; + mmadParams.k = k; + mmadParams.unitFlag = 0; + mmadParams.cmatrixInitVal = true; +#if (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3510) + mmadParams.disableGemv = true; +#endif + for (uint32_t l0BatchIdx = 0; l0BatchIdx < l0Batch; l0BatchIdx++) { + AscendC::Mmad(l0CTensor.data()[l0BatchIdx * L0CM * L0CN], + l0ATensor.data()[l0BatchIdx * L0AM * L0AK], + l0BTensor.data()[l0BatchIdx * L0BK * L0BN], + mmadParams); + } + } + + // automatically extracts actual sizes from tensor originShape + template + CATLASS_DEVICE void operator()(TensorC const &l0CTensor, + TensorA const &l0ATensor, + TensorB const &l0BTensor, + bool initC = true, uint8_t unitFlag = 0) + { + // Get actual sizes from tensor originShape + uint32_t m = tla::get<0>(l0CTensor.layout().originShape()); + uint32_t n = tla::get<1>(l0CTensor.layout().originShape()); + uint32_t k = tla::get<1>(l0ATensor.layout().originShape()); + AscendC::MmadParams mmadParams; +#if (defined(__NPU_ARCH__) && __NPU_ARCH__ == 2201) + if constexpr (std::is_same_v && std::is_same_v) { + mmadParams.kDirectionAlign = true; + } + if constexpr (!std::is_same_v) { + if (m == 1) + m = 16; // avoid gemv mode + } +#endif +#if (defined(__NPU_ARCH__) && __NPU_ARCH__ == 3510) + if constexpr (std::is_same_v) { + mmadParams.disableGemv = false; + } else { + mmadParams.disableGemv = true; + } +#endif + mmadParams.m = m; + mmadParams.n = n; + mmadParams.k = k; + mmadParams.unitFlag = unitFlag; + mmadParams.cmatrixInitVal = initC; + + AscendC::Mmad(l0CTensor.data(), + l0ATensor.data(), + l0BTensor.data(), + mmadParams); + + const uint32_t PIPE_M_BARRIER_THRESHOLD = 10; + if ((m / C0_NUM_PER_FRACTAL) * (n / C0_NUM_PER_FRACTAL) < PIPE_M_BARRIER_THRESHOLD) { + AscendC::PipeBarrier(); + } + } +}; + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Catlass::Gemm::Tile + +#endif // CATLASS_KDA_TILE_MMAD_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/kda_catlass.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_catlass.hpp new file mode 100644 index 0000000000..fd8c3b6845 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_catlass.hpp @@ -0,0 +1,55 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_CATLASS_HPP +#define CATLASS_KDA_CATLASS_HPP + +#include + +#if defined(__CCE__) +#include +#endif + +#include "catlass/detail/kda_alignment.hpp" +#include "catlass/detail/kda_dependent_false.hpp" +#include "catlass/detail/kda_macros.hpp" + +namespace Catlass { + +constexpr uint32_t BYTE_PER_C0 = 32; +constexpr uint32_t BYTE_PER_C2 = 64; +constexpr uint32_t C0_NUM_PER_FRACTAL = 16; +constexpr uint32_t BYTE_PER_FRACTAL = BYTE_PER_C0 * C0_NUM_PER_FRACTAL; + +constexpr uint32_t BYTE_PER_BLK = 32; +constexpr uint32_t BLK_NUM_PER_VECTOR_FRACTAL = 8; +constexpr uint32_t BYTE_PER_VECTOR_FRACTAL = BYTE_PER_BLK * BLK_NUM_PER_VECTOR_FRACTAL; + +constexpr uint64_t L2_OFFSET = 0; +constexpr uint32_t STRIDE_LIMIT = 65536; + +constexpr uint32_t BYTE_PER_BLK_FP = 128; /// datablock size of A1->C2PiPE2GM + +class EmptyClass {}; + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) +constexpr uint32_t MX_SCALE_COPY_GROUP_NUM = 2; // Mx-scale matrix 2-byte aligned +constexpr uint32_t MX_SCALE_GROUP_NUM = 32; // Data count for one MX-scale factor per group +constexpr uint32_t MX_BASEK_FACTOR = 64; // Data matrix alignment at K-dimension +#endif +} // namespace Catlass + +#if defined(__CCE__) && defined(L2_CACHE_HINT) && defined(CATLASS_BUILD_LEGACY) +inline __gm__ struct OpSystemRunCfg g_opSystemRunCfg { + Catlass::L2_OFFSET +}; +#endif + +#endif // CATLASS_KDA_CATLASS_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/kda_conv_coord.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_conv_coord.hpp new file mode 100644 index 0000000000..d179ef0510 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_conv_coord.hpp @@ -0,0 +1,1320 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_CONV_COORD_HPP +#define CATLASS_KDA_CONV_COORD_HPP + +#include "catlass/kda_coord.hpp" + +namespace Catlass { + +/// Shape of conv3d operation +struct Conv3dParams { +public: + typedef uint32_t Index; + static constexpr uint32_t N0 = 16; + using Fmap6HDShape = Coord<6, Index>; // {batch, di, cin1, hi, wi, cin0} + using FilterFracZ3DShape = Coord<7, Index>; // {kd, cin1, kh, kw, n1, n0, cin0} + using Out6HDShape = Coord<6, Index>; // {batch, do, cout1, ho, wo, cout0} + using Strides = Coord<3, Index>; + using Pads = Coord<3, Index>; + using Dilations = Coord<3, Index>; + +private: + Fmap6HDShape fmap6HDShape_; + FilterFracZ3DShape filterFracZ3DShape_; + Out6HDShape out6HDShape_; + Strides strides_; + Pads pads_; + Dilations dilations_; + Index cout_; + +public: + CATLASS_HOST_DEVICE + Conv3dParams( + Index BATCH = 1, + Index Di = 1, + Index Cin1 = 1, + Index Hi = 1, + Index Wi = 1, + Index C0 = 16, + Index Kd = 1, + Index Kh = 1, + Index Kw = 1, + Index N1 = 1, + Index Do = 1, + Index Ho = 1, + Index Wo = 1, + Index Cout1 = 1, + Index Cout = 1, + Index padHead = 0, + Index padTop = 0, + Index padLeft = 0, + Index strideD = 1, + Index strideH = 1, + Index strideW = 1, + Index dilationD = 1, + Index dilationH = 1, + Index dilationW = 1) + : fmap6HDShape_(MakeCoord(BATCH, Di, Cin1, Hi, Wi, C0)), + filterFracZ3DShape_(MakeCoord(Kd, Cin1, Kh, Kw, N1, N0, C0)), + out6HDShape_(MakeCoord(BATCH, Do, Cout1, Ho, Wo, C0)), + cout_(Cout), + pads_(MakeCoord(padHead, padTop, padLeft)), + strides_(MakeCoord(strideD, strideH, strideW)), + dilations_(MakeCoord(dilationD, dilationH, dilationW)) + {} + + CATLASS_HOST_DEVICE + static Conv3dParams MakeConvCoord( + const uint32_t *fmapShape, + const uint32_t *filterShape, + const uint32_t *paddings, + const uint32_t *strides, + const uint32_t *dilations) + { + return Conv3dParams( + fmapShape[0], + fmapShape[1], + fmapShape[2], + fmapShape[3], + fmapShape[4], + fmapShape[5], + filterShape[0], + filterShape[1], + filterShape[2], + CeilDiv(filterShape[3], N0), + (fmapShape[1] + paddings[0] * 2 - dilations[0] * (filterShape[0] - 1) - 1) / strides[0] + 1, // Do + (fmapShape[3] + paddings[1] * 2 - dilations[1] * (filterShape[1] - 1) - 1) / strides[1] + 1, // Ho + (fmapShape[4] + paddings[2] * 2 - dilations[2] * (filterShape[2] - 1) - 1) / strides[2] + 1, // Wo + CeilDiv(filterShape[3], fmapShape[5]), + filterShape[3], + paddings[0], + paddings[1], + paddings[2], + strides[0], + strides[1], + strides[2], + dilations[0], + dilations[1], + dilations[2]); + } + + // fmapShape + CATLASS_HOST_DEVICE + Index const &batch() const { return fmap6HDShape_[0]; } + CATLASS_HOST_DEVICE + Index const &cin1() const { return fmap6HDShape_[2]; } + CATLASS_HOST_DEVICE + Index const &di() const { return fmap6HDShape_[1]; } + CATLASS_HOST_DEVICE + Index const &hi() const { return fmap6HDShape_[3]; } + CATLASS_HOST_DEVICE + Index const &wi() const { return fmap6HDShape_[4]; } + CATLASS_HOST_DEVICE + Index const &cin0() const { return fmap6HDShape_[5]; } + CATLASS_HOST_DEVICE + Index const hiwi() const { return fmap6HDShape_[3] * fmap6HDShape_[4]; } + + // filterShape + CATLASS_HOST_DEVICE + Index const &kd() const { return filterFracZ3DShape_[0]; } + CATLASS_HOST_DEVICE + Index const &kh() const { return filterFracZ3DShape_[2]; } + CATLASS_HOST_DEVICE + Index const &kw() const { return filterFracZ3DShape_[3]; } + CATLASS_HOST_DEVICE + Index const khkw() const + { + return filterFracZ3DShape_[2] * filterFracZ3DShape_[3]; + } + CATLASS_HOST_DEVICE + Index const kdc1khkw() const + { + return filterFracZ3DShape_[0] * filterFracZ3DShape_[1] * filterFracZ3DShape_[2] * filterFracZ3DShape_[3]; + } + CATLASS_HOST_DEVICE + Index const &n1() const { return filterFracZ3DShape_[4]; } + CATLASS_HOST_DEVICE + Index const &n0() const { return filterFracZ3DShape_[5]; } + + // outShape + CATLASS_HOST_DEVICE + Index const &dout() const { return out6HDShape_[1]; } + CATLASS_HOST_DEVICE + Index const &ho() const { return out6HDShape_[3]; } + CATLASS_HOST_DEVICE + Index const &wo() const { return out6HDShape_[4]; } + CATLASS_HOST_DEVICE + Index const &cout1() const { return out6HDShape_[2]; } + CATLASS_HOST_DEVICE + Index const &cout0() const { return out6HDShape_[5]; } + CATLASS_HOST_DEVICE + Index const &cout() const { return cout_; } + + /// paddings + CATLASS_HOST_DEVICE + Index const &padhead() const { return pads_[0]; } + CATLASS_HOST_DEVICE + Index const &padtail() const { return pads_[0]; } + CATLASS_HOST_DEVICE + Index const &padtop() const { return pads_[1]; } + CATLASS_HOST_DEVICE + Index const &padbottom() const { return pads_[1]; } + CATLASS_HOST_DEVICE + Index const &padleft() const { return pads_[2]; } + CATLASS_HOST_DEVICE + Index const &padright() const { return pads_[2]; } + + /// strideSize + CATLASS_HOST_DEVICE + Index const &sD() const { return strides_[0]; } + CATLASS_HOST_DEVICE + Index const &sH() const { return strides_[1]; } + CATLASS_HOST_DEVICE + Index const &sW() const { return strides_[2]; } + + /// dilationSize + CATLASS_HOST_DEVICE + Index const &dD() const { return dilations_[0]; } + CATLASS_HOST_DEVICE + Index const dilatedKernelD() const { return 1 + (filterFracZ3DShape_[0] - 1) * dilations_[0]; } + CATLASS_HOST_DEVICE + Index const &dH() const { return dilations_[1]; } + CATLASS_HOST_DEVICE + Index const dilatedKernelH() const { return 1 + (filterFracZ3DShape_[2] - 1) * dilations_[1]; } + CATLASS_HOST_DEVICE + Index const &dW() const { return dilations_[2]; } + CATLASS_HOST_DEVICE + Index const dilatedKernelW() const { return 1 + (filterFracZ3DShape_[3] - 1) * dilations_[2]; } + + ///// used in block + CATLASS_HOST_DEVICE + Index const howo() const { return out6HDShape_[3] * out6HDShape_[4]; } + CATLASS_HOST_DEVICE + Index const alignCout() const { return out6HDShape_[2] * out6HDShape_[5]; } + CATLASS_HOST_DEVICE + Index const wicin0() const { return fmap6HDShape_[4] * fmap6HDShape_[5]; } + CATLASS_HOST_DEVICE + Index const khkwcin0() const { return filterFracZ3DShape_[2] * filterFracZ3DShape_[3] * filterFracZ3DShape_[6]; } + CATLASS_HOST_DEVICE + Index const alignCinKhKwKd() const + { + return filterFracZ3DShape_[0] * filterFracZ3DShape_[1] * filterFracZ3DShape_[2] * filterFracZ3DShape_[3] * filterFracZ3DShape_[6]; + } + CATLASS_HOST_DEVICE + Index const kdcin1() const { return filterFracZ3DShape_[0] * filterFracZ3DShape_[1]; } + CATLASS_HOST_DEVICE + Index const fmapOneBatchSize() const + { + return fmap6HDShape_[1] * fmap6HDShape_[2] * fmap6HDShape_[3] * fmap6HDShape_[4] * fmap6HDShape_[5]; + } + CATLASS_HOST_DEVICE + Index const outputOneBatchSize() const + { + return out6HDShape_[1] * out6HDShape_[2] * out6HDShape_[3] * out6HDShape_[4] * out6HDShape_[5]; + } +}; + +template < + uint32_t noCnt_ = 1, + uint32_t doCnt_ = 1, + uint32_t co1Cnt_ = 1, + uint32_t howoCnt_ = 1> +struct ConvCoreShape { + static uint32_t const noCnt = noCnt_; + static uint32_t const doCnt = doCnt_; + static uint32_t const co1Cnt = co1Cnt_; + static uint32_t const howoCnt = howoCnt_; + + /// Returns a Coord object + CATLASS_HOST_DEVICE + static Coord<4> ToCoord() + { + return MakeCoord(noCnt, doCnt, co1Cnt, howoCnt); + } +}; + +template < + uint32_t mAL1_ = 1, + uint32_t Kd_ = 1, + uint32_t Ci1_ = 1> +struct ConvFmapL1Shape { + static uint32_t constexpr mAL1 = mAL1_; + static uint32_t constexpr Kd = Kd_; + static uint32_t constexpr Ci1 = Ci1_; + + /// Returns a Coord object + CATLASS_HOST_DEVICE + static Coord<3> ToCoord() + { + return MakeCoord(mAL1, Kd, Ci1); + } +}; + +template < + uint32_t Kd_ = 1, + uint32_t Ci1_ = 1, + uint32_t nBL1_ = 1> +struct ConvFilterL1Shape { + static uint32_t constexpr Kd = Kd_; + static uint32_t constexpr Ci1 = Ci1_; + static uint32_t constexpr nBL1 = nBL1_; + + /// Returns a Coord object + CATLASS_HOST_DEVICE + static Coord<3> ToCoord() + { + return MakeCoord(Kd, Ci1, nBL1); + } +}; + +template < + uint32_t mL0_ = 1, + uint32_t kL0_ = 1, + uint32_t nL0_ = 1> +struct ConvL0Shape { + static uint32_t constexpr mL0 = mL0_; + static uint32_t constexpr kL0 = kL0_; + static uint32_t constexpr nL0 = nL0_; + + /// Returns a Coord object + CATLASS_HOST_DEVICE + static Coord<3> ToCoord() + { + return MakeCoord(mL0, kL0, nL0); + } +}; + +struct Conv3d6HdCoord : public Coord<4, uint32_t> { + using Index = uint32_t; + + using Base = Coord<4, Index>; + + static constexpr int N_INDEX = 0; + static constexpr int D_INDEX = 1; + static constexpr int C1_INDEX = 2; + static constexpr int HW_INDEX = 3; + + /// Default ctor + CATLASS_HOST_DEVICE + Conv3d6HdCoord() {} + + CATLASS_HOST_DEVICE + Conv3d6HdCoord(Coord<4, Index> const &coord) + : Base(coord) + {} + + CATLASS_HOST_DEVICE + Conv3d6HdCoord(Index n, Index d, Index c1, Index hw) + : Base(MakeCoord(n, d, c1, hw)) + {} + + CATLASS_HOST_DEVICE + Index const &n() const + { + return this->At(N_INDEX); + } + CATLASS_HOST_DEVICE + Index &n() + { + return this->At(N_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &d() const + { + return this->At(D_INDEX); + } + CATLASS_HOST_DEVICE + Index &d() + { + return this->At(D_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &c1() const + { + return this->At(C1_INDEX); + } + CATLASS_HOST_DEVICE + Index &c1() + { + return this->At(C1_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &hw() const + { + return this->At(HW_INDEX); + } + CATLASS_HOST_DEVICE + Index &hw() + { + return this->At(HW_INDEX); + } +}; + +struct Conv3dFracZ3dCoord : public Coord<2, uint32_t> { + using Index = uint32_t; + + using Base = Coord<2, Index>; + + static constexpr int KDC1KHKW_INDEX = 0; + static constexpr int N1_INDEX = 1; + + /// Default ctor + CATLASS_HOST_DEVICE + Conv3dFracZ3dCoord() {} + + CATLASS_HOST_DEVICE + Conv3dFracZ3dCoord(Index kdc1khkw, Index n1) + : Base(MakeCoord(kdc1khkw, n1)) + {} + + CATLASS_HOST_DEVICE + Index const &kdc1khkw() const + { + return this->At(KDC1KHKW_INDEX); + } + CATLASS_HOST_DEVICE + Index &kdc1khkw() + { + return this->At(KDC1KHKW_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &n1() const + { + return this->At(N1_INDEX); + } + CATLASS_HOST_DEVICE + Index &n1() + { + return this->At(N1_INDEX); + } +}; + +/////////////////// Shapes and Coords for Conv2d /////////////////// + +template < + uint32_t Ho_ = 1, + uint32_t Wo_ = 1, + uint32_t Cin1_ = 1> +struct Conv2dFmapL1Shape { // (Ho, Wo, Cin1) + static constexpr uint32_t Ho = Ho_; + static constexpr uint32_t Wo = Wo_; + static constexpr uint32_t Cin1 = Cin1_; + + /// Returns a Coord object + CATLASS_HOST_DEVICE + static Coord<3> ToCoord() + { + return MakeCoord(Ho, Wo, Cin1); + } +}; + +template < + uint32_t Cout_ = 16, + uint32_t Cin1_ = 1> +struct Conv2dFilterL1Shape { // (Cout, Cin1) + static constexpr uint32_t Cout = Cout_; + static constexpr uint32_t Cin1 = Cin1_; + + /// Returns a Coord object + CATLASS_HOST_DEVICE + static Coord<2> ToCoord() + { + return MakeCoord(Cout, Cin1); + } +}; + +template < + uint32_t M_ = 16, + uint32_t N_ = 16, + uint32_t K_ = 16> +struct Conv2dL0Shape { + static constexpr uint32_t M = M_; + static constexpr uint32_t N = N_; + static constexpr uint32_t K = K_; + + /// Returns a Coord object + CATLASS_HOST_DEVICE + static Coord<3> ToCoord() + { + return MakeCoord(M, N, K); + } +}; + +struct Conv2dFmapCoord : public Coord<5, uint32_t> { // (Batch, C1, H, W, C0) +public: + /// Integer-valued index + using Index = uint32_t; + + /// Base type is a Coord of rank=5 + using Base = Coord<5, Index>; + + /// LongIndex type + using LongIndex = typename Base::LongIndex; + + /// dimensions + static constexpr uint32_t BATCH_INDEX = 0; + static constexpr uint32_t C1_INDEX = 1; + static constexpr uint32_t H_INDEX = 2; + static constexpr uint32_t W_INDEX = 3; + static constexpr uint32_t C0_INDEX = 4; + + /// Default ctor + CATLASS_HOST_DEVICE + Conv2dFmapCoord() {} + + /// Constructs from Coord<5> + CATLASS_HOST_DEVICE + Conv2dFmapCoord(Coord<5, Index> const &coord) + : Base(coord) + {} + + /// Helper to construct from C1, H, W, C0 + CATLASS_HOST_DEVICE + Conv2dFmapCoord(Index batch, Index c1, Index h, Index w, Index c0) + : Base(MakeCoord(batch, c1, h, w, c0)) + {} + + CATLASS_HOST_DEVICE + Conv2dFmapCoord(LongIndex batch, LongIndex c1, LongIndex h, LongIndex w, LongIndex c0) + : Base(MakeCoord(Index(batch), Index(c1), Index(h), Index(w), Index(c0))) + {} + + CATLASS_HOST_DEVICE + Index const &c1() const + { + return this->At(C1_INDEX); + } + CATLASS_HOST_DEVICE + Index &c1() + { + return this->At(C1_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &batch() const + { + return this->At(BATCH_INDEX); + } + CATLASS_HOST_DEVICE + Index &batch() + { + return this->At(BATCH_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &h() const + { + return this->At(H_INDEX); + } + CATLASS_HOST_DEVICE + Index &h() + { + return this->At(H_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &w() const + { + return this->At(W_INDEX); + } + CATLASS_HOST_DEVICE + Index &w() + { + return this->At(W_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &c0() const + { + return this->At(C0_INDEX); + } + CATLASS_HOST_DEVICE + Index &c0() + { + return this->At(C0_INDEX); + } + + /// Element-wise addition + CATLASS_HOST_DEVICE + Conv2dFmapCoord operator+(Base const &b) const + { + return Conv2dFmapCoord(Base::operator+(b)); + } + + /// In-place addition + CATLASS_HOST_DEVICE + Conv2dFmapCoord &operator+=(Base const &b) + { + Base::operator+=(b); + return *this; + } +}; + +struct Conv2dFilterCoord : public Coord<5, uint32_t> { // (Cin1, Kh, Kw, Cout, C0) +public: + /// Integer-valued index + using Index = uint32_t; + + /// Base type is a Coord of rank=5 + using Base = Coord<5, Index>; + + /// LongIndex type + using LongIndex = typename Base::LongIndex; + + /// dimensions + static constexpr uint32_t CIN1_INDEX = 0; + static constexpr uint32_t KH_INDEX = 1; + static constexpr uint32_t KW_INDEX = 2; + static constexpr uint32_t COUT_INDEX = 3; + static constexpr uint32_t C0_INDEX = 4; + + /// Default ctor + CATLASS_HOST_DEVICE + Conv2dFilterCoord() {} + + /// Constructs from Coord<5> + CATLASS_HOST_DEVICE + Conv2dFilterCoord(Coord<5, Index> const &coord) + : Base(coord) + {} + + /// Helper to construct from Cin1, Kh, Kw, Cout, C0 + CATLASS_HOST_DEVICE + Conv2dFilterCoord(Index cin1, Index kh, Index kw, Index cout, Index c0) + : Base(MakeCoord(cin1, kh, kw, cout, c0)) + {} + + CATLASS_HOST_DEVICE + Conv2dFilterCoord(LongIndex cin1, LongIndex kh, LongIndex kw, LongIndex cout, LongIndex c0) + : Base(MakeCoord(Index(cin1), Index(kh), Index(kw), Index(cout), Index(c0))) + {} + + CATLASS_HOST_DEVICE + Index const &cin1() const + { + return this->At(CIN1_INDEX); + } + CATLASS_HOST_DEVICE + Index &cin1() + { + return this->At(CIN1_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &kh() const + { + return this->At(KH_INDEX); + } + CATLASS_HOST_DEVICE + Index &kh() + { + return this->At(KH_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &kw() const + { + return this->At(KW_INDEX); + } + CATLASS_HOST_DEVICE + Index &kw() + { + return this->At(KW_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &cout() const + { + return this->At(COUT_INDEX); + } + CATLASS_HOST_DEVICE + Index &cout() + { + return this->At(COUT_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &c0() const + { + return this->At(C0_INDEX); + } + CATLASS_HOST_DEVICE + Index &c0() + { + return this->At(C0_INDEX); + } + + /// Element-wise addition + CATLASS_HOST_DEVICE + Conv2dFilterCoord operator+(Base const &b) const + { + return Conv2dFilterCoord(Base::operator+(b)); + } + + /// In-place addition + CATLASS_HOST_DEVICE + Conv2dFilterCoord &operator+=(Base const &b) + { + Base::operator+=(b); + return *this; + } +}; + +struct Conv2dHoWoCoCoord : public Coord<3, uint32_t> { // (Ho, Wo, Cout) +public: + /// Integer-valued index + using Index = uint32_t; + + /// Base type is a Coord of rank=3 + using Base = Coord<3, Index>; + + /// LongIndex type + using LongIndex = typename Base::LongIndex; + + /// dimensions + static constexpr uint32_t HO_INDEX = 0; + static constexpr uint32_t WO_INDEX = 1; + static constexpr uint32_t COUT_INDEX = 2; + + /// Default ctor + CATLASS_HOST_DEVICE + Conv2dHoWoCoCoord() {} + + /// Constructs from Coord<3> + CATLASS_HOST_DEVICE + Conv2dHoWoCoCoord(Coord<3, Index> const &coord) + : Base(coord) + {} + + /// Helper to construct from Ho, Wo, Cout + CATLASS_HOST_DEVICE + Conv2dHoWoCoCoord(Index ho, Index wo, Index cout) + : Base(MakeCoord(ho, wo, cout)) + {} + + CATLASS_HOST_DEVICE + Conv2dHoWoCoCoord(LongIndex ho, LongIndex wo, LongIndex cout) + : Base(MakeCoord(Index(ho), Index(wo), Index(cout))) + {} + + CATLASS_HOST_DEVICE + Index const &ho() const + { + return this->At(HO_INDEX); + } + + CATLASS_HOST_DEVICE + Index &ho() + { + return this->At(HO_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &wo() const + { + return this->At(WO_INDEX); + } + + CATLASS_HOST_DEVICE + Index &wo() + { + return this->At(WO_INDEX); + } + + CATLASS_HOST_DEVICE + Index howo() const + { + return this->At(HO_INDEX) * this->At(WO_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &cout() const + { + return this->At(COUT_INDEX); + } + + CATLASS_HOST_DEVICE + Index &cout() + { + return this->At(COUT_INDEX); + } + + /// Element-wise addition + CATLASS_HOST_DEVICE + Conv2dHoWoCoCoord operator+(Base const &b) const + { + return Conv2dHoWoCoCoord(Base::operator+(b)); + } + + /// In-place addition + CATLASS_HOST_DEVICE + Conv2dHoWoCoCoord &operator+=(Base const &b) + { + Base::operator+=(b); + return *this; + } +}; + +struct Conv2dCoord : public Coord<5, uint32_t> { // (Batch, H, W, Cout, Cin1) +public: + /// Integer-valued index + using Index = uint32_t; + + /// Base type is a Coord of rank=5 + using Base = Coord<5, Index>; + + /// LongIndex type + using LongIndex = typename Base::LongIndex; + + /// dimensions + static constexpr uint32_t BATCH_INDEX = 0; + static constexpr uint32_t H_INDEX = 1; + static constexpr uint32_t W_INDEX = 2; + static constexpr uint32_t COUT_INDEX = 3; + static constexpr uint32_t CIN1_INDEX = 4; + + /// Default ctor + CATLASS_HOST_DEVICE + Conv2dCoord() {} + + /// Constructs from Coord<5> + CATLASS_HOST_DEVICE + Conv2dCoord(Coord<5, Index> const &coord) + : Base(coord) + {} + + /// Helper to construct from Batch, H, W, Cout, Cin1 + CATLASS_HOST_DEVICE + Conv2dCoord(Index batch, Index h, Index w, Index cout, Index cin1) + : Base(MakeCoord(batch, h, w, cout, cin1)) + {} + + CATLASS_HOST_DEVICE + Conv2dCoord(LongIndex batch, LongIndex h, LongIndex w, LongIndex cout, LongIndex cin1) + : Base(MakeCoord(Index(batch), Index(h), Index(w), Index(cout), Index(cin1))) + {} + + CATLASS_HOST_DEVICE + Index const &batch() const + { + return this->At(BATCH_INDEX); + } + CATLASS_HOST_DEVICE + Index &batch() + { + return this->At(BATCH_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &h() const + { + return this->At(H_INDEX); + } + CATLASS_HOST_DEVICE + Index &h() + { + return this->At(H_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &w() const + { + return this->At(W_INDEX); + } + CATLASS_HOST_DEVICE + Index &w() + { + return this->At(W_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &cout() const + { + return this->At(COUT_INDEX); + } + CATLASS_HOST_DEVICE + Index &cout() + { + return this->At(COUT_INDEX); + } + + CATLASS_HOST_DEVICE + Index const &cin1() const + { + return this->At(CIN1_INDEX); + } + CATLASS_HOST_DEVICE + Index &cin1() + { + return this->At(CIN1_INDEX); + } + + /// Element-wise addition + CATLASS_HOST_DEVICE + Conv2dCoord operator+(Base const &b) const + { + return Conv2dCoord(Base::operator+(b)); + } + + /// In-place addition + CATLASS_HOST_DEVICE + Conv2dCoord &operator+=(Base const &b) + { + Base::operator+=(b); + return *this; + } + + CATLASS_HOST_DEVICE + auto GetHoWoCoCoord() const + { + return this->GetCoordByAxis(); + } +}; + +class Conv2dFilterParams { +public: + typedef uint8_t ShortIndex; + typedef uint32_t Index; + static constexpr uint32_t C0 = 16; + using Ks = Coord<2, ShortIndex>; + using Pads = Coord<4, ShortIndex>; + using Strides = Coord<2, ShortIndex>; + using Dilations = Coord<2, ShortIndex>; + +private: + Ks ks; + Pads pads; + Strides strides; + Dilations dilations; + +public: + Conv2dFilterParams(ShortIndex kh = 0, ShortIndex kw = 0, + ShortIndex padLeft = 0, ShortIndex padRight = 0, ShortIndex padTop = 0, ShortIndex padBottom = 0, + ShortIndex strideH = 0, ShortIndex strideW = 0, ShortIndex dilationH = 0, ShortIndex dilationW = 0) + : ks(MakeCoord(kh, kw)), + pads(MakeCoord(padLeft, padRight, padTop, padBottom)), + strides(MakeCoord(strideH, strideW)), + dilations(MakeCoord(dilationH, dilationW)) + {} + + CATLASS_HOST_DEVICE + ShortIndex const &kh() const + { + return this->ks[0]; + } + CATLASS_HOST_DEVICE + ShortIndex &kh() + { + return this->ks[0]; + } + + CATLASS_HOST_DEVICE + ShortIndex const &kw() const + { + return this->ks[1]; + } + CATLASS_HOST_DEVICE + ShortIndex &kw() + { + return this->ks[1]; + } + + CATLASS_HOST_DEVICE + ShortIndex const &padLeft() const + { + return this->pads[0]; + } + CATLASS_HOST_DEVICE + ShortIndex const &padLeft() + { + return this->pads[0]; + } + + CATLASS_HOST_DEVICE + ShortIndex const &padRight() const + { + return this->pads[1]; + } + CATLASS_HOST_DEVICE + ShortIndex const &padRight() + { + return this->pads[1]; + } + + CATLASS_HOST_DEVICE + ShortIndex const &padTop() const + { + return this->pads[2]; + } + CATLASS_HOST_DEVICE + ShortIndex const &padTop() + { + return this->pads[2]; + } + + CATLASS_HOST_DEVICE + ShortIndex const &padBottom() const + { + return this->pads[3]; + } + CATLASS_HOST_DEVICE + ShortIndex const &padBottom() + { + return this->pads[3]; + } + + CATLASS_HOST_DEVICE + ShortIndex const &strideH() const + { + return this->strides[0]; + } + CATLASS_HOST_DEVICE + ShortIndex const &strideH() + { + return this->strides[0]; + } + + CATLASS_HOST_DEVICE + ShortIndex const &strideW() const + { + return this->strides[1]; + } + CATLASS_HOST_DEVICE + ShortIndex const &strideW() + { + return this->strides[1]; + } + + CATLASS_HOST_DEVICE + ShortIndex const &dilationH() const + { + return this->dilations[0]; + } + CATLASS_HOST_DEVICE + ShortIndex const &dilationH() + { + return this->dilations[0]; + } + + CATLASS_HOST_DEVICE + ShortIndex const &dilationW() const + { + return this->dilations[1]; + } + CATLASS_HOST_DEVICE + ShortIndex const &dilationW() + { + return this->dilations[1]; + } +}; + +struct Conv2dParams { +public: + typedef uint8_t ShortIndex; + typedef uint32_t Index; + static constexpr uint32_t C0 = 16; + using Pads = Coord<4, ShortIndex>; + using Strides = Coord<2, ShortIndex>; + using Dilations = Coord<2, ShortIndex>; + +private: + // Batch, Hi, Wi, Cin, Cout, Kh, Kw + Conv2dFmapCoord fmapShape; // {Batch, Cin1, Hi, Wi, C0} + Conv2dFilterCoord filterShape; // {Cin1, Kh, Kw, Cout, C0} + Conv2dFmapCoord outputShape; // {Batch, Cout1, Ho, Wo, C0} + Conv2dFilterParams configs; // {Ks, Pads, Strides, Dilations} + Conv2dCoord postIm2colShape; // {Batch, Ho, Wo, Cout, Cin1} +public: + /// Default ctor + CATLASS_HOST_DEVICE + Conv2dParams() {} + + CATLASS_HOST_DEVICE + Conv2dParams(Index batch, Index hi, Index wi, Index cin, Index cout, + ShortIndex kh, ShortIndex kw, + ShortIndex padLeft, ShortIndex padRight, ShortIndex padTop, ShortIndex padBottom, + ShortIndex strideH, ShortIndex strideW, + ShortIndex dilationH, ShortIndex dilationW) + : fmapShape(MakeCoord(batch, CeilDiv(cin, C0), hi, wi, C0)), + filterShape(MakeCoord(CeilDiv(cin, C0), (Index)kh, (Index)kw, cout, C0)), + configs(kh, kw, padLeft, padRight, padTop, padBottom, strideH, strideW, dilationH, dilationW) + { + Index cout1 = CeilDiv(cout, C0); + Index ho = (hi + padTop + padBottom - dilationH * (kh - 1) - 1) / strideH + 1; + Index wo = (wi + padLeft + padRight - dilationW * (kw - 1) - 1) / strideW + 1; + outputShape = MakeCoord(batch, cout1, ho, wo, C0); + postIm2colShape = MakeCoord(batch, ho, wo, cout, filterShape.cin1()); + } + + CATLASS_HOST_DEVICE + static Conv2dParams MakeConv2dParams( + const Index *dataSizes, // [Batch, Hi, Wi, Cin, Cout] + const ShortIndex *filterSizes, // [Kh, Kw] + const ShortIndex *pads, // [padLeft, padRight, padTop, padBottom] + const ShortIndex *strides, // [strideH, strideW] + const ShortIndex *dilations) // [dilationH, dilationW] + { + return Conv2dParams( + dataSizes[0], dataSizes[1], dataSizes[2], dataSizes[3], dataSizes[4], + filterSizes[0], filterSizes[1], + pads[0], pads[1], pads[2], pads[3], + strides[0], strides[1], strides[2], strides[3]); + } + + CATLASS_HOST_DEVICE + Conv2dFilterParams const &getFilterParams() const + { + return this->configs; + } + + CATLASS_HOST_DEVICE + Conv2dFmapCoord const &getOutputShape() const + { + return this->outputShape; + } + + CATLASS_HOST_DEVICE + Conv2dCoord const &getPostIm2colShape() const + { + return this->postIm2colShape; + } + + CATLASS_HOST_DEVICE + Index const &batch() const + { + return this->fmapShape.batch(); + } + CATLASS_HOST_DEVICE + Index &batch() + { + return this->fmapShape.batch(); + } + + CATLASS_HOST_DEVICE + Index const &hi() const + { + return this->fmapShape.h(); + } + CATLASS_HOST_DEVICE + Index &hi() + { + return this->fmapShape.h(); + } + + CATLASS_HOST_DEVICE + Index const &wi() const + { + return this->fmapShape.w(); + } + CATLASS_HOST_DEVICE + Index &wi() + { + return this->fmapShape.w(); + } + + CATLASS_HOST_DEVICE + Index cin() const + { + return this->fmapShape.c1() * C0; + } + + CATLASS_HOST_DEVICE + Index const &cin1() const + { + return this->fmapShape.c1(); + } + CATLASS_HOST_DEVICE + Index &cin1() + { + return this->fmapShape.c1(); + } + + CATLASS_HOST_DEVICE + Index const &cout() const + { + return this->filterShape.cout(); + } + CATLASS_HOST_DEVICE + Index &cout() + { + return this->filterShape.cout(); + } + + CATLASS_HOST_DEVICE + Index const &cout1() const + { + return this->outputShape.c1(); + } + CATLASS_HOST_DEVICE + Index &cout1() + { + return this->outputShape.c1(); + } + + CATLASS_HOST_DEVICE + Index coutRound() const + { + return (this->filterShape.cout() + C0 - 1) / C0 * C0; + } + + CATLASS_HOST_DEVICE + Index const &ho() const + { + return this->outputShape.h(); + } + CATLASS_HOST_DEVICE + Index &ho() + { + return this->outputShape.h(); + } + + CATLASS_HOST_DEVICE + Index const &wo() const + { + return this->outputShape.w(); + } + CATLASS_HOST_DEVICE + Index &wo() + { + return this->outputShape.w(); + } + + CATLASS_HOST_DEVICE + Index howo() const + { + return this->outputShape.h() * this->outputShape.w(); + } + + CATLASS_HOST_DEVICE + Index howoRound() const + { + return (this->howo() + C0 - 1) / C0 * C0; + } + + CATLASS_HOST_DEVICE + ShortIndex const &kh() const + { + return this->configs.kh(); + } + CATLASS_HOST_DEVICE + ShortIndex &kh() + { + return this->configs.kh(); + } + + CATLASS_HOST_DEVICE + ShortIndex const &kw() const + { + return this->configs.kw(); + } + CATLASS_HOST_DEVICE + ShortIndex &kw() + { + return this->configs.kw(); + } + + CATLASS_HOST_DEVICE + ShortIndex const &padTop() const + { + return this->configs.padTop(); + } + CATLASS_HOST_DEVICE + ShortIndex const &padTop() + { + return this->configs.padTop(); + } + + CATLASS_HOST_DEVICE + ShortIndex const &padBottom() const + { + return this->configs.padBottom(); + } + CATLASS_HOST_DEVICE + ShortIndex const &padBottom() + { + return this->configs.padBottom(); + } + + CATLASS_HOST_DEVICE + ShortIndex const &padLeft() const + { + return this->configs.padLeft(); + } + CATLASS_HOST_DEVICE + ShortIndex const &padLeft() + { + return this->configs.padLeft(); + } + + CATLASS_HOST_DEVICE + ShortIndex const &padRight() const + { + return this->configs.padRight(); + } + CATLASS_HOST_DEVICE + ShortIndex const &padRight() + { + return this->configs.padRight(); + } + + CATLASS_HOST_DEVICE + ShortIndex const &strideH() const + { + return this->configs.strideH(); + } + CATLASS_HOST_DEVICE + ShortIndex const &strideH() + { + return this->configs.strideH(); + } + + CATLASS_HOST_DEVICE + ShortIndex const &strideW() const + { + return this->configs.strideW(); + } + CATLASS_HOST_DEVICE + ShortIndex const &strideW() + { + return this->configs.strideW(); + } + + CATLASS_HOST_DEVICE + ShortIndex const &dilationH() const + { + return this->configs.dilationH(); + } + CATLASS_HOST_DEVICE + ShortIndex const &dilationH() + { + return this->configs.dilationH(); + } + + CATLASS_HOST_DEVICE + ShortIndex const &dilationW() const + { + return this->configs.dilationW(); + } + CATLASS_HOST_DEVICE + ShortIndex const &dilationW() + { + return this->configs.dilationW(); + } +}; + +} // namespace Catlass + +#endif // CATLASS_KDA_CONV_COORD_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/kda_coord.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_coord.hpp new file mode 100644 index 0000000000..d68fdb1b28 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_coord.hpp @@ -0,0 +1,336 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_COORD_HPP +#define CATLASS_KDA_COORD_HPP + +#include "catlass/kda_catlass.hpp" + +namespace Catlass { + +/// Statically-sized array specifying Coords within a tensor +template < + int RANK_, ///< Logical rank of coordinate + class Index_ = uint32_t, ///< Index type used for each dimension + class LongIndex_ = int64_t ///< Long index type used for linear offsets + > +struct Coord { +public: + // Number of elements in Coord + static const int RANK = RANK_; + + // Index typen used to store elements + using Index = Index_; + + // Type used to represent linear offsets + using LongIndex = LongIndex_; + + // Default ctor initializes uniformly + CATLASS_HOST_DEVICE constexpr explicit Coord(Index value = Index(0)) + { + for (int i = 0; i < RANK; ++i) { + idx[i] = value; + } + } + + // Constructs from an array of integers + CATLASS_HOST_DEVICE constexpr Coord(Index const (&idx_)[RANK]) + { + for (int i = 0; i < RANK; ++i) { + idx[i] = idx_[i]; + } + } + + // Constructs frrom an array of integers + CATLASS_HOST_DEVICE + int Argmin() const + { + int i = 0; + for (int j = 1; j < RANK; ++j) { + if (idx[j] < idx[i]) { + i = j; + } + } + return i; + } + + // Returns the index of the dimension with greatest value + CATLASS_HOST_DEVICE + int Argmax() const + { + int i = 0; + for (int j = 1; j < RANK; ++j) { + if (idx[j] > idx[i]) { + i = j; + } + } + return i; + } + + // Returns true if Coord is non-zero + CATLASS_HOST_DEVICE + explicit operator bool() const + { + for (int i = 0; i < RANK; ++i) { + if (idx[i]) { + return true; + } + } + return false; + } + + // Return true if Coord is uniformly zero. + CATLASS_HOST_DEVICE + bool operator!() const + { + for (int i = 0; i < RANK; ++i) { + if (idx[i]) { + return false; + } + } + return true; + } + + // Element-wise addition + CATLASS_HOST_DEVICE + Coord operator+(Coord const &b) const + { + Coord c; + for (int i = 0; i < RANK; ++i) { + c.idx[i] = idx[i] + b.idx[i]; + } + return c; + } + + // Add a scalar to each element + CATLASS_HOST_DEVICE + Coord operator+(const Index val) const + { + Coord c; + for (int i = 0; i < RANK; ++i) { + c.idx[i] = idx[i] + val; + } + return c; + } + + // Element-wise subtraction + CATLASS_HOST_DEVICE + Coord operator-(Coord const &b) const + { + Coord c; + for (int i = 0; i < RANK; i++) { + c.idx[i] = idx[i] - b.idx[i]; + } + return c; + } + + // Subtract a scalar from each element + CATLASS_HOST_DEVICE + Coord operator-(Index const val) const + { + Coord c; + for (int i = 0; i < RANK; ++i) { + c.idx[i] = idx[i] - val; + } + return c; + } + + // Element-wise multiply + CATLASS_HOST_DEVICE + Coord operator*(Coord const &b) const + { + Coord c; + for (int i = 0; i < RANK; i++) { + c.idx[i] = idx[i] * b.idx[i]; + } + return c; + } + + // Element-wise division + CATLASS_HOST_DEVICE + Coord operator/(Coord const &b) const + { + Coord c; + for (int i = 0; i < RANK; i++) { + c.idx[i] = idx[i] / b.idx[i]; + } + return c; + } + + // Element-wise mod + CATLASS_HOST_DEVICE + Coord operator%(Coord const &b) const + { + Coord c; + for (int i = 0; i < RANK; i++) { + c.idx[i] = idx[i] % b.idx[i]; + } + return c; + } + + // In-place addition + CATLASS_HOST_DEVICE + Coord &operator+=(Coord const &b) + { + for (int i = 0; i < RANK; ++i) { + idx[i] += b.idx[i]; + } + return *this; + } + + // In-place equal + CATLASS_HOST_DEVICE + bool operator==(Coord const &b) const + { + for (int i = 0; i < RANK; ++i) { + if (idx[i] != b.idx[i]) { + return false; + } + } + return true; + } + + // In-place equal + CATLASS_HOST_DEVICE + bool operator==(Index const val) const + { + for (int i = 0; i < RANK; ++i) { + if (idx[i] != val) { + return false; + } + } + return true; + } + + // Member acces operator + CATLASS_HOST_DEVICE + Index &operator[](int dim) + { + return idx[dim]; + } + + // Member access operator + CATLASS_HOST_DEVICE + Index const &operator[](int dim) const + { + return idx[dim]; + } + + // Gets the index of a given Coord element + template + CATLASS_HOST_DEVICE Index &At() + { + return idx[DIM]; + } + + // Access via index; may limit unrolling potential + CATLASS_HOST_DEVICE + Index &At(int dim) + { + return idx[dim]; + } + + // Gets the index of a given Coord element + template + CATLASS_HOST_DEVICE + Index const & + At() const + { + return idx[DIM]; + } + + // Access via index; may limit unrolling potential + CATLASS_HOST_DEVICE + Index const &At(int dim) const + { + return idx[dim]; + } + + template + CATLASS_HOST_DEVICE auto GetCoordByAxis() const + { + Index idx_[sizeof...(Is)]{idx[Is]...}; + return Coord{idx_}; + } + + CATLASS_HOST_DEVICE + static Coord Min(Coord const &a, Coord const &b) + { + Coord res; + for (int i = 0; i < RANK; ++i) { + res[i] = a[i] < b[i] ? a[i] : b[i]; + } + return res; + } + +private: + // Indices + Index idx[RANK]; +}; + +// Helper to make a 1-element coordinate +template +CATLASS_HOST_DEVICE constexpr Coord<1, T> MakeCoord(T dim0) +{ + T values[1] = {dim0}; + return Coord<1, T>(values); +} + +/// Helper to make a 2-element coordinate +template +CATLASS_HOST_DEVICE constexpr Coord<2, T> MakeCoord(T dim0, T dim1) +{ + T values[2] = {dim0, dim1}; + return Coord<2, T>(values); +} + +/// Helper to make a 3-element coordinate +template +CATLASS_HOST_DEVICE constexpr Coord<3, T> MakeCoord(T dim0, T dim1, T dim2) +{ + T values[3] = {dim0, dim1, dim2}; + return Coord<3, T>(values); +} + +/// Helper to make a 4-element coordinate +template +CATLASS_HOST_DEVICE constexpr Coord<4, T> MakeCoord(T dim0, T dim1, T dim2, T dim3) +{ + T values[4] = {dim0, dim1, dim2, dim3}; + return Coord<4, T>(values); +} + +/// Helper to make a 5-element coordinate +template +CATLASS_HOST_DEVICE constexpr Coord<5, T> MakeCoord(T dim0, T dim1, T dim2, T dim3, T dim4) +{ + T values[5] = {dim0, dim1, dim2, dim3, dim4}; + return Coord<5, T>(values); +} + +/// Helper to make a 6-element coordinate +template +CATLASS_HOST_DEVICE constexpr Coord<6, T> MakeCoord(T dim0, T dim1, T dim2, T dim3, T dim4, T dim5) +{ + T values[6] = {dim0, dim1, dim2, dim3, dim4, dim5}; + return Coord<6, T>(values); +} + +/// Helper to make a 7-element coordinate +template +CATLASS_HOST_DEVICE constexpr Coord<7, T> MakeCoord(T dim0, T dim1, T dim2, T dim3, T dim4, T dim5, T dim6) +{ + T values[7] = {dim0, dim1, dim2, dim3, dim4, dim5, dim6}; + return Coord<7, T>(values); +} + +} // namespace Catlass + +#endif // CATLASS_KDA_COORD_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/kda_debug.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_debug.hpp new file mode 100644 index 0000000000..6e0ecc9ecd --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_debug.hpp @@ -0,0 +1,73 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_DEBUG_HPP +#define CATLASS_KDA_DEBUG_HPP + +#include +#include +#include + +#include + +#define SINGLE_CORE_DUMPSIZE (1024 * 1024) +#ifdef CATLASS_ARCH +#if CATLASS_ARCH == 2201 +#define DUMP_CORE_NUM 75 +#endif +#if CATLASS_ARCH == 3510 +#define DUMP_CORE_NUM 108 +#endif +#endif +#ifndef DUMP_CORE_NUM +#define DUMP_CORE_NUM 0 +#endif +// DUMP_CORE_NUM is from AscendC host stub +#define ALL_DUMPSIZE (DUMP_CORE_NUM * SINGLE_CORE_DUMPSIZE) + +using LogFuncType = std::function; +/** + * @brief Check acl api status code. + * @param status The return code of acl api. + * @param logFunc Log function, which receives a C-Style string. + * @return + */ +inline void aclCheck( + aclError status, + LogFuncType logFunc = [](const char *logStrPtr) { std::cerr << logStrPtr; }) +{ + if (status != ACL_SUCCESS) { + std::stringstream ss; + ss << "AclError: " << status; + logFunc(ss.str().c_str()); + } +} +/** + * @brief Check rt api status code. + * @param status The return code of rt api. + * @param logFunc Log function, which receives a C-Style string. + * @return + */ +inline void rtCheck( + int status, + LogFuncType logFunc = [](const char *logStrPtr) { std::cerr << logStrPtr; }) +{ + if (status != 0) { + std::stringstream ss; + ss << "RtError: " << status; + logFunc(ss.str().c_str()); + } +} + +namespace Adx { +void AdumpPrintWorkSpace(const void *dumpBufferAddr, const size_t dumpBufferSize, aclrtStream stream, const char *opType); +} // namespace Adx + +#endif // CATLASS_KDA_DEBUG_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/kda_gemm_coord.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_gemm_coord.hpp new file mode 100644 index 0000000000..0c23e4ac8e --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_gemm_coord.hpp @@ -0,0 +1,161 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_GEMM_COORD_HPP +#define CATLASS_KDA_GEMM_COORD_HPP + +#include "catlass/kda_coord.hpp" + +namespace Catlass { + +/// Shape of a matrix multiply-add operation +template < + /// Rows of matrix product + uint32_t M_ = 1, + /// Columns of matrix product + uint32_t N_ = 1, + /// Inner dimension of matrix product + uint32_t K_ = 1> +struct GemmShape { + static constexpr uint32_t M = M_; + static constexpr uint32_t N = N_; + static constexpr uint32_t K = K_; + + static constexpr int64_t MN = M * N; + static constexpr int64_t MK = M * K; + static constexpr int64_t KN = N * K; + static constexpr int64_t MNK = M * N * K; + + static constexpr int64_t COUNT = MNK; + + /// Returns a Coord object + CATLASS_HOST_DEVICE + static Coord<3> ToCoord() + { + return MakeCoord(M, N, K); + } + + CATLASS_HOST_DEVICE + static Coord<2> ToCoordMN() + { + return MakeCoord(M, N); + } + + CATLASS_HOST_DEVICE + static Coord<2> ToCoordMK() + { + return MakeCoord(M, K); + } + + CATLASS_HOST_DEVICE + static Coord<2> ToCoordKN() + { + return MakeCoord(K, N); + } +}; + +/// GemmCoord is a structure derived from Coord<3> that specifies a location within the +/// coordinate space of a Gemm problem. +struct GemmCoord : public Coord<3, uint32_t> { + /// Integer-valued index + using Index = uint32_t; + + /// Base type is a Coord of rank=3 + using Base = Coord<3, Index>; + + /// Gemm M dimension - rows of the output C matrix + static constexpr int M_INDEX = 0; + + /// Gemm N dimension - columns of the output C matrix + static constexpr int N_INDEX = 1; + + /// Gemm K dimension - inner dimension of the Gemm problem + static constexpr int K_INDEX = 2; + + /// Default ctor + CATLASS_HOST_DEVICE + GemmCoord() {} + + /// Constructs from Coord<3> and a batch + CATLASS_HOST_DEVICE + GemmCoord(Coord<3, Index> const &coord) + : Base(coord) + {} + + /// Helper to construct from a K, N, M, batch variables + CATLASS_HOST_DEVICE + GemmCoord(Index m, Index n, Index k) + : Base(MakeCoord(m, n, k)) + {} + + /// Returns the Gemm M coordinate + CATLASS_HOST_DEVICE + Index const &m() const + { + return this->At(M_INDEX); + } + + /// Returns reference to the Gemm M coordinate + CATLASS_HOST_DEVICE + Index &m() + { + return this->At(M_INDEX); + } + + /// Returns the Gemm N coordinate + CATLASS_HOST_DEVICE + Index const &n() const + { + return this->At(N_INDEX); + } + + /// Returns reference to the Gemm N coordinate + CATLASS_HOST_DEVICE + Index &n() + { + return this->At(N_INDEX); + } + + /// Returns the Gemm K coordinate + CATLASS_HOST_DEVICE + Index const &k() const + { + return this->At(K_INDEX); + } + + /// Returns reference to the Gemm K coordinate + CATLASS_HOST_DEVICE + Index &k() + { + return this->At(K_INDEX); + } + + CATLASS_HOST_DEVICE + auto GetCoordMN() const + { + return this->GetCoordByAxis(); + } + + CATLASS_HOST_DEVICE + auto GetCoordMK() const + { + return this->GetCoordByAxis(); + } + + CATLASS_HOST_DEVICE + auto GetCoordKN() const + { + return this->GetCoordByAxis(); + } +}; + +} // namespace Catlass + +#endif // CATLASS_KDA_GEMM_COORD_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/kda_matrix_coord.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_matrix_coord.hpp new file mode 100644 index 0000000000..fccd956347 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_matrix_coord.hpp @@ -0,0 +1,108 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_MATRIX_COORD_HPP +#define CATLASS_KDA_MATRIX_COORD_HPP + +#include "catlass/kda_coord.hpp" + +namespace Catlass { + +template < + uint32_t ROW_ = 1, + uint32_t COLUMN_ = 1> +struct MatrixShape { + static constexpr uint32_t ROW = ROW_; + static constexpr uint32_t COLUMN = COLUMN_; + + static constexpr int64_t COUNT = ROW * COLUMN; + + CATLASS_HOST_DEVICE + static Coord<2> ToCoord() + { + return MakeCoord(ROW, COLUMN); + } +}; + +/// MatrixCoord wraps Coord<2, uint32_t> to provide a helper for accessing named dimensions. Classes +/// expecting a coordinate in the rank=2 index space of a matrix should use MatrixCoord. +struct MatrixCoord : public Coord<2, uint32_t> { + /// Integer-valued index + using Index = uint32_t; + + /// Base type is a Coord of rank=2 + using Base = Coord<2, Index>; + + /// LongIndex type + using LongIndex = typename Base::LongIndex; + + /// Rows dimension + static constexpr uint32_t ROW_INDEX = 0; + + /// Columns dimension + static constexpr uint32_t COLUMN_INDEX = 1; + + /// Default ctor + CATLASS_HOST_DEVICE + MatrixCoord() {} + + /// Constructs from Coord<2> + CATLASS_HOST_DEVICE + MatrixCoord(Coord<2, Index> const &coord) + : Base(coord) + {} + + /// Helper to construct from a row and column + CATLASS_HOST_DEVICE + MatrixCoord(Index row, Index column) + : Base(MakeCoord(row, column)) + {} + + /// Helper to construct from a row and column, which are LongIndex based + CATLASS_HOST_DEVICE + MatrixCoord(LongIndex row, LongIndex column) + : Base(MakeCoord(Index(row), Index(column))) + {} + + /// Returns the row of the coordinate + CATLASS_HOST_DEVICE + Index const &row() const { return this->At(ROW_INDEX); } + + /// Returns the row of the coordinate + CATLASS_HOST_DEVICE + Index &row() { return this->At(ROW_INDEX); } + + /// Returns the column of the coordinate + CATLASS_HOST_DEVICE + Index const &column() const { return this->At(COLUMN_INDEX); } + + /// Returns the column of the coordinate + CATLASS_HOST_DEVICE + Index &column() { return this->At(COLUMN_INDEX); } + + /// Element-wise addition + CATLASS_HOST_DEVICE + MatrixCoord operator+(Base const &b) const + { + return MatrixCoord(Base::operator+(b)); + } + + /// In-place addition + CATLASS_HOST_DEVICE + MatrixCoord &operator+=(Base const &b) + { + Base::operator+=(b); + return *this; + } +}; + +} // namespace Catlass + +#endif diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/kda_numeric_size.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_numeric_size.hpp new file mode 100644 index 0000000000..11b5a05cf2 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/kda_numeric_size.hpp @@ -0,0 +1,42 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef CATLASS_KDA_NUMERIC_SIZE_HPP +#define CATLASS_KDA_NUMERIC_SIZE_HPP + +#include "catlass/kda_catlass.hpp" + +namespace Catlass { + +#if defined(__CCE__) +using AscendC::SizeOfBits; +#else +template +struct SizeOfBits { + static constexpr size_t value = sizeof(T) * 8; +}; +#endif + +/// Returns the number of bytes required to hold a specified number of bits +template +CATLASS_HOST_DEVICE constexpr ReturnType BitsToBytes(T bits) +{ + return (static_cast(bits) + static_cast(7)) / static_cast(8); +} + +/// Returns the number of bits required to hold a specified number of bytes +template +CATLASS_HOST_DEVICE constexpr ReturnType BytesToBits(T bytes) +{ + return static_cast(bytes) * static_cast(8); +} + +} // namespace Catlass + +#endif // CATLASS_KDA_NUMERIC_SIZE_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_layout.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_layout.hpp new file mode 100644 index 0000000000..1f0c85e2c1 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_layout.hpp @@ -0,0 +1,19 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_LAYOUT_HPP +#define CATLASS_KDA_LAYOUT_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/layout/kda_matrix.hpp" +#include "catlass/layout/kda_vector.hpp" +#include "catlass/layout/kda_tensor.hpp" + +#endif // CATLASS_KDA_LAYOUT_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_matrix.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_matrix.hpp new file mode 100644 index 0000000000..6e2efbeeeb --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_matrix.hpp @@ -0,0 +1,1818 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_MATRIX_HPP +#define CATLASS_KDA_MATRIX_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_numeric_size.hpp" +#include "catlass/detail/kda_alignment.hpp" +#include "catlass/kda_matrix_coord.hpp" +#include "catlass/kda_conv_coord.hpp" + +namespace Catlass::layout { + +/// Mapping function for row-major matrices +struct RowMajor { +public: + /// Logical rank of tensor + static constexpr int RANK = 2; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + /// Constructor + CATLASS_HOST_DEVICE + RowMajor(Index rows = 0, Index cols = 0) + : shape_(MakeCoord(rows, cols)), + stride_(MakeCoord(LongIndex(cols), LongIndex(1))) + {} + + /// Constructor + CATLASS_HOST_DEVICE + RowMajor(Index rows, Index cols, LongIndex ldm) + : shape_(MakeCoord(rows, cols)), + stride_(MakeCoord(ldm, LongIndex(1))) + {} + + /// Ctor + CATLASS_HOST_DEVICE + RowMajor(Shape shape, Stride stride) + : shape_(shape), + stride_(stride) + {} + + template + CATLASS_HOST_DEVICE static RowMajor MakeLayout(Index rows, Index cols) + { + return RowMajor(rows, cols); + } + + template + CATLASS_HOST_DEVICE static RowMajor MakeLayoutInUb(MatrixCoord const &shape) + { + constexpr uint32_t ELE_NUM_PER_BLK = BytesToBits(BYTE_PER_BLK) / SizeOfBits::value; + return RowMajor(shape.row(), shape.column(), RoundUp(shape.column())); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + CATLASS_HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + return LongIndex(coord.row()) * stride_[0] + LongIndex(coord.column()); + } + + /// Returns the layout of a tile. + CATLASS_HOST_DEVICE + RowMajor GetTileLayout(MatrixCoord const &tileShape) const + { + return RowMajor(tileShape, stride()); + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + + /// Returns the length of the layout + CATLASS_HOST_DEVICE + LongIndex Capacity() const + { + return static_cast(shape_[0]) * stride_[0]; + } + +protected: + // + // Data members + // + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/// Mapping function for col-major matrices +struct ColumnMajor { +public: + /// Logical rank of tensor + static constexpr int RANK = 2; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + // Methods + + /// Constructor + CATLASS_HOST_DEVICE + ColumnMajor(Index rows = 0, Index cols = 0) + : shape_(MakeCoord(rows, cols)), + stride_(MakeCoord(LongIndex(1), LongIndex(rows))) + {} + + /// Constructor + CATLASS_HOST_DEVICE + ColumnMajor(Index rows, Index cols, LongIndex ldm) + : shape_(MakeCoord(rows, cols)), + stride_(MakeCoord(LongIndex(1), ldm)) + {} + + /// Ctor + CATLASS_HOST_DEVICE + ColumnMajor(Shape shape, Stride stride) + : shape_(shape), + stride_(stride) + {} + + template + CATLASS_HOST_DEVICE static ColumnMajor MakeLayout(Index rows, Index cols) + { + return ColumnMajor(rows, cols); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + CATLASS_HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + return LongIndex(coord.row()) + LongIndex(coord.column()) * stride_[1]; + } + + /// Returns the layout of a tile. + CATLASS_HOST_DEVICE + ColumnMajor GetTileLayout(MatrixCoord const &tileShape) const + { + return ColumnMajor(tileShape, stride()); + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + + /// Returns the length of the layout + CATLASS_HOST_DEVICE + LongIndex Capacity() const + { + return static_cast(shape_[1]) * stride_[1]; + } + +protected: + // + // Data members + // + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/// Mapping function for nZ matrices which is col-major inside fractal and row-major between fractal +struct nZ { +public: + /// Logical rank of tensor + static constexpr int RANK = 4; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical rank of orgshape + static constexpr int ORG_SHAPE_RANK = 2; + + /// Logical coordinate + using OrgShape = Coord; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + // Methods + + /// Constructor + CATLASS_HOST_DEVICE constexpr nZ(Index orgRows = 0, /// Number of rows of origin matrices + Index orgCols = 0, /// Number of cols of origin matrices + Index rowsInFractal = 0, /// Number of rows inside the fractal + Index rowsByFractal = 0, /// number of rows by the fractal + Index colsInFractal = 0, /// number of cols inside the fractal + Index colsByFractal = 0, /// number of cols by the fractal + LongIndex strideRowsInFractal = 0, /// number of elements between adjacent rows inside the fractal + LongIndex strideRowsByFractal = 0, /// number of elements between adjacent fractal rows + LongIndex strideColsInFractal = 0, /// number of elements between adjacent cols inside the fractal + LongIndex strideColsByFractal = 0) /// number of elements between adjacent fractal cols + : orgShape_(MakeCoord(orgRows, orgCols)), + shape_(MakeCoord(rowsInFractal, rowsByFractal, colsInFractal, colsByFractal)), + stride_(MakeCoord(strideRowsInFractal, strideRowsByFractal, strideColsInFractal, strideColsByFractal)) + {} + + /// Ctor + CATLASS_HOST_DEVICE constexpr nZ(OrgShape orgShape, Shape shape, Stride stride) + : orgShape_(orgShape), + shape_(shape), + stride_(stride) + {} + + /// Make the layout of a coordinate (row, column) + template + CATLASS_HOST_DEVICE constexpr static nZ MakeLayout(Index orgRows, Index orgCols) + { + constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + Index rowsRound = RoundUp(orgRows); + Index colsRound = RoundUp(orgCols); + return nZ(orgRows, + orgCols, + ELE_NUM_PER_C0, + rowsRound / ELE_NUM_PER_C0, + C0_NUM_PER_FRACTAL, + colsRound / C0_NUM_PER_FRACTAL, + 1, + colsRound * ELE_NUM_PER_C0, + ELE_NUM_PER_C0, + ELE_NUM_PER_FRACTAL); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + CATLASS_HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + return LongIndex(coord.row()) / shape_[0] * stride_[1] + LongIndex(coord.column()) / shape_[2] * stride_[3] + + (LongIndex(coord.row()) % shape_[0]) * stride_[0] + (LongIndex(coord.column()) % shape_[2]) * stride_[2]; + } + + /// Returns the layout of a tile. + CATLASS_HOST_DEVICE + nZ GetTileLayout(MatrixCoord const &tileOriShape) const + { + auto tileShape = MakeCoord( + shape(0), CeilDiv(tileOriShape.row(), shape(0)), + shape(2), CeilDiv(tileOriShape.column(), shape(2))); + return nZ(tileOriShape, tileShape, stride()); + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const + { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index &orgShape(int idx) + { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + + /// Returns the length of the layout + CATLASS_HOST_DEVICE + LongIndex Capacity() const + { + return static_cast(stride_[1]) * shape_[1]; + } + +private: + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/// Mapping function for zN matrices which is row-major inside fractal and col-major between fractal +struct zN { +public: + /// Logical rank of tensor + static constexpr int RANK = 4; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical rank of orgshape + static constexpr int ORG_SHAPE_RANK = 2; + + /// Logical coordinate + using OrgShape = Coord; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + // Methods + + /// Constructor + CATLASS_HOST_DEVICE constexpr zN(Index orgRows = 0, /// Number of rows of origin matrices + Index orgCols = 0, /// Number of cols of origin matrices + Index rowsInFractal = 0, /// Number of rows inside the fractal + Index rowsByFractal = 0, /// number of rows by the fractal + Index colsInFractal = 0, /// number of cols inside the fractal + Index colsByFractal = 0, /// number of cols by the fractal + LongIndex strideRowsInFractal = 0, /// number of elements between adjacent rows inside the fractal + LongIndex strideRowsByFractal = 0, /// number of elements between adjacent fractal rows + LongIndex strideColsInFractal = 0, /// number of elements between adjacent cols inside the fractal + LongIndex strideColsByFractal = 0) /// number of elements between adjacent fractal cols + : orgShape_(MakeCoord(orgRows, orgCols)), + shape_(MakeCoord(rowsInFractal, rowsByFractal, colsInFractal, colsByFractal)), + stride_(MakeCoord(strideRowsInFractal, strideRowsByFractal, strideColsInFractal, strideColsByFractal)) + {} + + /// Ctor + CATLASS_HOST_DEVICE constexpr zN(OrgShape orgShape, Shape shape, Stride stride) + : orgShape_(orgShape), + shape_(shape), + stride_(stride) + {} + + /// Make the layout of a coordinate (row, column) + template + CATLASS_HOST_DEVICE constexpr static zN MakeLayout(Index orgRows, Index orgCols) + { + constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + Index rowsRound = RoundUp(orgRows); + Index colsRound = RoundUp(orgCols); + return zN(orgRows, + orgCols, + C0_NUM_PER_FRACTAL, + rowsRound / C0_NUM_PER_FRACTAL, + ELE_NUM_PER_C0, + colsRound / ELE_NUM_PER_C0, + ELE_NUM_PER_C0, + ELE_NUM_PER_FRACTAL, + 1, + rowsRound * ELE_NUM_PER_C0); + } + + CATLASS_HOST_DEVICE + static zN MakeLayoutInL0C(MatrixCoord const &shape) + { + return zN(shape.row(), + shape.column(), + C0_NUM_PER_FRACTAL, + CeilDiv(shape.row()), + C0_NUM_PER_FRACTAL, + CeilDiv(shape.column()), + C0_NUM_PER_FRACTAL, + C0_NUM_PER_FRACTAL * C0_NUM_PER_FRACTAL, + 1, + RoundUp(shape.row()) * C0_NUM_PER_FRACTAL); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + CATLASS_HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + return LongIndex(coord.row()) / shape_[0] * stride_[1] + LongIndex(coord.column()) / shape_[2] * stride_[3] + + (LongIndex(coord.row()) % shape_[0]) * stride_[0] + (LongIndex(coord.column()) % shape_[2]) * stride_[2]; + } + + /// Returns the layout of a tile. + CATLASS_HOST_DEVICE + zN GetTileLayout(MatrixCoord const &tileOriShape) const + { + auto tileShape = MakeCoord( + shape(0), CeilDiv(tileOriShape.row(), shape(0)), + shape(2), CeilDiv(tileOriShape.column(), shape(2))); + return zN(tileOriShape, tileShape, stride()); + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const + { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index &orgShape(int idx) + { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + + /// Returns the length of the layout + CATLASS_HOST_DEVICE + LongIndex Capacity() const + { + return static_cast(stride_[3]) * shape_[3]; + } + +private: + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/// Mapping function for zN matrices which is row-major inside fractal and row-major between fractal +struct zZ { +public: + /// Logical rank of tensor + static constexpr int RANK = 4; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical rank of orgshape + static constexpr int ORG_SHAPE_RANK = 2; + + /// Logical coordinate + using OrgShape = Coord; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + // Methods + + /// Constructor + CATLASS_HOST_DEVICE constexpr zZ(Index orgRows = 0, /// Number of rows of origin matrices + Index orgCols = 0, /// Number of cols of origin matrices + Index rowsInFractal = 0, /// Number of rows inside the fractal + Index rowsByFractal = 0, /// number of rows by the fractal + Index colsInFractal = 0, /// number of cols inside the fractal + Index colsByFractal = 0, /// number of cols by the fractal + LongIndex strideRowsInFractal = 0, /// number of elements between adjacent rows inside the fractal + LongIndex strideRowsByFractal = 0, /// number of elements between adjacent fractal rows + LongIndex strideColsInFractal = 0, /// number of elements between adjacent cols inside the fractal + LongIndex strideColsByFractal = 0) /// number of elements between adjacent fractal cols + : orgShape_(MakeCoord(orgRows, orgCols)), + shape_(MakeCoord(rowsInFractal, rowsByFractal, colsInFractal, colsByFractal)), + stride_(MakeCoord(strideRowsInFractal, strideRowsByFractal, strideColsInFractal, strideColsByFractal)) + {} + + /// Ctor + CATLASS_HOST_DEVICE constexpr zZ(OrgShape orgShape, Shape shape, Stride stride) + : orgShape_(orgShape), + shape_(shape), + stride_(stride) + {} + + /// Make the layout of a coordinate (row, column) + template + CATLASS_HOST_DEVICE constexpr static zZ MakeLayout(Index orgRows, Index orgCols) + { + constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + Index rowsRound = RoundUp(orgRows); + Index colsRound = RoundUp(orgCols); + return zZ(orgRows, + orgCols, + C0_NUM_PER_FRACTAL, + rowsRound / C0_NUM_PER_FRACTAL, + ELE_NUM_PER_C0, + colsRound / ELE_NUM_PER_C0, + ELE_NUM_PER_C0, + colsRound * C0_NUM_PER_FRACTAL, + 1, + ELE_NUM_PER_FRACTAL); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + CATLASS_HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + return LongIndex(coord.row()) / shape_[0] * stride_[1] + LongIndex(coord.column()) / shape_[2] * stride_[3]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const + { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index &orgShape(int idx) + { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + +private: + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/// Mapping function for L0C matrices +/// A special data layout for L0C memory, used for accumulator storage. +/// This layout similar to zN layout, but uses C0_NUM_PER_FRACTAL for both row and column inner dimensions. +struct L0C { +public: + /// Logical rank of tensor + static constexpr int RANK = 4; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical rank of orgshape + static constexpr int ORG_SHAPE_RANK = 2; + + /// Logical coordinate + using OrgShape = Coord; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + // Methods + + /// Constructor + CATLASS_HOST_DEVICE constexpr L0C(Index orgRows = 0, /// Number of rows of origin matrices + Index orgCols = 0, /// Number of cols of origin matrices + Index rowsInFractal = 0, /// Number of rows inside the fractal + Index rowsByFractal = 0, /// number of rows by the fractal + Index colsInFractal = 0, /// number of cols inside the fractal + Index colsByFractal = 0, /// number of cols by the fractal + LongIndex strideRowsInFractal = 0, /// number of elements between adjacent rows inside the fractal + LongIndex strideRowsByFractal = 0, /// number of elements between adjacent fractal rows + LongIndex strideColsInFractal = 0, /// number of elements between adjacent cols inside the fractal + LongIndex strideColsByFractal = 0) /// number of elements between adjacent fractal cols + : orgShape_(MakeCoord(orgRows, orgCols)), + shape_(MakeCoord(rowsInFractal, rowsByFractal, colsInFractal, colsByFractal)), + stride_(MakeCoord(strideRowsInFractal, strideRowsByFractal, strideColsInFractal, strideColsByFractal)) + {} + + /// Ctor + CATLASS_HOST_DEVICE constexpr L0C(OrgShape orgShape, Shape shape, Stride stride) + : orgShape_(orgShape), + shape_(shape), + stride_(stride) + {} + + /// Make the layout of a coordinate (row, column) + template + CATLASS_HOST_DEVICE constexpr static L0C MakeLayout(Index orgRows, Index orgCols) + { + constexpr uint32_t ELE_NUM_PER_FRACTAL = 256; + Index rowsRound = RoundUp(orgRows); + Index colsRound = RoundUp(orgCols); + return L0C(orgRows, + orgCols, + C0_NUM_PER_FRACTAL, + rowsRound / C0_NUM_PER_FRACTAL, + C0_NUM_PER_FRACTAL, + colsRound / C0_NUM_PER_FRACTAL, + C0_NUM_PER_FRACTAL, + ELE_NUM_PER_FRACTAL, + 1, + rowsRound * C0_NUM_PER_FRACTAL); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + CATLASS_HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + return LongIndex(coord.row()) / shape_[0] * stride_[1] + LongIndex(coord.column()) / shape_[2] * stride_[3] + + (LongIndex(coord.row()) % shape_[0]) * stride_[0] + (LongIndex(coord.column()) % shape_[2]) * stride_[2]; + } + + /// Returns the layout of a tile. + CATLASS_HOST_DEVICE + L0C GetTileLayout(MatrixCoord const &tileOriShape) const + { + auto tileShape = MakeCoord( + shape(0), CeilDiv(tileOriShape.row(), shape(0)), + shape(2), CeilDiv(tileOriShape.column(), shape(2))); + return L0C(tileOriShape, tileShape, stride()); + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const + { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index &orgShape(int idx) + { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + + /// Returns the length of the layout + CATLASS_HOST_DEVICE + LongIndex Capacity() const + { + return static_cast(stride_[3]) * shape_[3]; + } + +private: + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/// Mapping function for padding rowmajor matrices +/// A special data layout designed to improve the efficiency of matrix operations in non-512B aligned scenarios. +/// This layout is row-major within blocks and also row-major between blocks. +struct PaddingRowMajor { +public: + /// Logical rank of tensor + static constexpr int RANK = 4; + + /// Logical rank of orgshape + static constexpr int ORG_SHAPE_RANK = 2; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using OrgShape = Coord; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + /// Constructor + CATLASS_HOST_DEVICE + PaddingRowMajor(Index orgRows = 0, Index orgCols = 0, Index blockRows = 0, Index blockCols = 0) + : orgShape_(MakeCoord(orgRows, orgCols)), + shape_(MakeCoord(blockRows, CeilDiv(orgRows, blockRows), blockCols, CeilDiv(orgCols, blockCols))), + stride_(MakeCoord((LongIndex)blockCols, (LongIndex)blockRows * (LongIndex)RoundUp(orgCols, blockCols), + (LongIndex)1, (LongIndex)blockRows * (LongIndex)blockCols)) + {} + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + CATLASS_HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + LongIndex blockRows = (LongIndex)shape_[0]; + LongIndex blockCols = (LongIndex)shape_[2]; + return (LongIndex)coord.row() / blockRows * stride_[1] + (LongIndex)coord.column() / blockCols * stride_[3] + (LongIndex)coord.row() % blockRows * stride_[0] + (LongIndex)coord.column() % blockCols; + } + + CATLASS_HOST_DEVICE + PaddingRowMajor GetTileLayout(MatrixCoord const &tileShape) const + { + return PaddingRowMajor(tileShape.row(), tileShape.column(), shape_[0], shape_[2]); + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const + { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index &orgShape(int idx) + { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + +private: + // + // Data members + // + + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/// Mapping function for padding columnmajor matrices +/// A special data layout designed to improve the efficiency of matrix operations in non-512B aligned scenarios. +/// This layout is column-major within blocks and also column-major between blocks. +struct PaddingColumnMajor { +public: + /// Logical rank of tensor + static constexpr int RANK = 4; + + /// Logical rank of orgshape + static constexpr int ORG_SHAPE_RANK = 2; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using OrgShape = Coord; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + /// Constructor + CATLASS_HOST_DEVICE + PaddingColumnMajor(Index orgRows = 0, Index orgCols = 0, Index blockRows = 0, Index blockCols = 0) + : orgShape_(MakeCoord(orgRows, orgCols)), + shape_(MakeCoord(blockRows, CeilDiv(orgRows, blockRows), blockCols, CeilDiv(orgCols, blockCols))), + stride_(MakeCoord((LongIndex)1, (LongIndex)blockRows * (LongIndex)blockCols, (LongIndex)blockRows, + (LongIndex)RoundUp(orgRows, blockRows) * (LongIndex)blockCols)) + {} + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + CATLASS_HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + LongIndex blockRows = (LongIndex)shape_[0]; + LongIndex blockCols = (LongIndex)shape_[2]; + return (LongIndex)coord.row() / blockRows * stride_[1] + (LongIndex)coord.column() / blockCols * stride_[3] + (LongIndex)coord.row() % blockRows + (LongIndex)coord.column() % blockCols * stride_[2]; + } + + CATLASS_HOST_DEVICE + PaddingColumnMajor GetTileLayout(MatrixCoord const &tileShape) const + { + return PaddingColumnMajor(tileShape.row(), tileShape.column(), shape_[0], shape_[2]); + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const + { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index &orgShape(int idx) + { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + +private: + // + // Data members + // + + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +/////////////////////// +// new add layout nN +// nN layout +struct nN { +public: + /// Logical rank of tensor + static constexpr int RANK = 4; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical rank of orgshape + static constexpr int ORG_SHAPE_RANK = 2; + + /// Logical coordinate + using OrgShape = Coord; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + // Methods + + /// Constructor + CATLASS_HOST_DEVICE + nN(Index orgRows = 0, /// Number of rows of origin matrices + Index orgCols = 0, /// Number of cols of origin matrices + + Index rowsInFractal = 0, /// Number of rows inside the fractal + Index rowsByFractal = 0, /// number of rows by the fractal + Index colsInFractal = 0, /// number of cols inside the fractal + Index colsByFractal = 0, /// number of cols by the fractal + + LongIndex strideRowsInFractal = 0, /// number of elements between adjacent rows inside the fractal + LongIndex strideRowsByFractal = 0, /// number of elements between adjacent fractal rows + LongIndex strideColsInFractal = 0, /// number of elements between adjacent cols inside the fractal + LongIndex strideColsByFractal = 0) /// number of elements between adjacent fractal cols + : orgShape_(MakeCoord(orgRows, orgCols)), + shape_(MakeCoord(rowsInFractal, rowsByFractal, colsInFractal, colsByFractal)), + stride_(MakeCoord(strideRowsInFractal, strideRowsByFractal, strideColsInFractal, strideColsByFractal)) + { + } + + /// Ctor + CATLASS_HOST_DEVICE + nN(OrgShape orgShape, Shape shape, Stride stride) + : orgShape_(orgShape), + shape_(shape), + stride_(stride) + {} + + /// Make the layout of a coordinate (row, column) + template + CATLASS_HOST_DEVICE static nN MakeLayout(Index orgRows, Index orgCols) + { + static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits::value; + static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits::value; + Index rowsRound = RoundUp(orgRows); + Index colsRound = RoundUp(orgCols); + return nN(orgRows, + orgCols, + + ELE_NUM_PER_C0, + rowsRound / ELE_NUM_PER_C0, + C0_NUM_PER_FRACTAL, + colsRound / C0_NUM_PER_FRACTAL, + + 1, + ELE_NUM_PER_FRACTAL, + ELE_NUM_PER_C0, + rowsRound * C0_NUM_PER_FRACTAL); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (row, column) + CATLASS_HOST_DEVICE + LongIndex GetOffset(MatrixCoord const &coord) const + { + return LongIndex(coord.row()) / shape_[0] * stride_[1] + LongIndex(coord.column()) / shape_[2] * stride_[3]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const + { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index &orgShape(int idx) + { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + +private: + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +struct NDC1HWC0 { +public: + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical rank of orgshape + /// (N,D,C1,H,W,C0) + static constexpr int ORG_SHAPE_RANK = 6; + + static constexpr int RANK = 5; + using OrgShape = Coord; + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + // Methods + + /// Constructor + CATLASS_HOST_DEVICE constexpr NDC1HWC0( + Index batch = 0, + Index D = 0, + Index C1 = 0, + Index H = 0, + Index W = 0, + Index C0 = 0, + + Index rowsInFractal = 0, /// Number of rows inside the fractal + Index rowsByFractal = 0, /// number of rows by the fractal + Index colsInFractal = 0, /// number of cols inside the fractal + Index colsByFractal = 0, /// number of cols by the fractal + + LongIndex strideC0 = 0, /// number of elements between adjacent C0 cols + LongIndex strideHW = 0, /// number of elements between adjacent W rows + LongIndex StrideC1 = 0, /// number of elements between adjacent C1 cols + LongIndex StrideD = 0, /// number of elements between adjacent D batchCols + LongIndex StrideN = 0 /// number of elements between adjacent batch + ) + : orgShape_(MakeCoord(batch, D, C1, H, W, C0)), + shape_(MakeCoord(batch, rowsInFractal, rowsByFractal, colsInFractal, colsByFractal)), + stride_(MakeCoord(strideC0, strideHW, StrideC1, StrideD, StrideN)) + {} + + /// Ctor + CATLASS_HOST_DEVICE constexpr NDC1HWC0(OrgShape orgshape, Shape shape, Stride stride) + : orgShape_(orgshape), + shape_(shape), + stride_(stride) + {} + + CATLASS_HOST_DEVICE constexpr static NDC1HWC0 MakeLayout(Index Batch, Index D, Index C1, Index H, Index W, Index C0) + { + return NDC1HWC0(Batch, + D, + C1, + H, + W, + C0, + + W, + H, + C0, + D * C1, + + 1, /// StrideC0 + C0, /// StrideHW + H * W * C0, /// StrideC1 + H * W * C0 * C1, /// StrideD + H * W * C0 * C1 * D /// StrideN + ); + } + + // CATLASS_HOST_DEVICE + /// Returns the offset of a coordinate in linear memory. + CATLASS_HOST_DEVICE + LongIndex GetOffset(Conv3d6HdCoord const &coord) const + { + return LongIndex(coord.n()) * stride_[4] + LongIndex(coord.d()) * stride_[3] + + LongIndex(coord.c1()) * stride_[2] + LongIndex(coord.hw()) * stride_[1]; + } + + /// Returns the layout of a tile. + CATLASS_HOST_DEVICE + NDC1HWC0 GetTileLayout(OrgShape const &tileOriShape) const + { + Shape tileShape = MakeCoord( + tileOriShape[0], + tileOriShape[4], + tileOriShape[3], + shape(3), + tileOriShape[1] * tileOriShape[2]); + + Stride tileStride = MakeCoord(stride(0), + stride(1), + (LongIndex)(tileOriShape[3] * tileOriShape[4] * shape(3)), + (LongIndex)(tileOriShape[2] * tileOriShape[3] * tileOriShape[4] * shape(3)), + (LongIndex)(tileOriShape[1] * tileOriShape[2] * tileOriShape[3] * tileOriShape[4] * shape(3))); + return NDC1HWC0(tileOriShape, tileShape, tileStride); + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const + { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index &orgShape(int idx) + { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + +private: + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +struct KDC1KHKWN1N0C0 { +public: + static constexpr int RANK = 4; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical rank of orgshape + static constexpr int ORG_SHAPE_RANK = 4; + + /// Logical coordinate + using OrgShape = Coord; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + // Methods + + /// Constructor + CATLASS_HOST_DEVICE constexpr KDC1KHKWN1N0C0(Index KdC1KhKw = 0, /// Merging Kd,Kh,Kw,C1 axes of KDC1KHKWN1N0C0 + Index N1 = 0, /// Cout = N1*N0 + Index N0 = 0, + Index C0 = 0, + + Index rowsInFractal = 0, /// Number of rows inside the fractal + Index rowsByFractal = 0, /// number of rows by the fractal + Index colsInFractal = 0, /// number of cols inside the fractal + Index colsByFractal = 0, /// number of cols by the fractal + + LongIndex strideC0 = 0, /// number of elements between adjacent rows inside the fractal + LongIndex StrideDC1HW = 0, /// number of elements between adjacent fractal rows + LongIndex strideN0 = 0, /// number of elements between adjacent cols inside the fractal + LongIndex strideN1 = 0 /// number of elements between adjacent fractal cols + ) + : orgShape_(MakeCoord(KdC1KhKw, N1, N0, C0)), + shape_(MakeCoord(rowsInFractal, rowsByFractal, colsInFractal, colsByFractal)), + stride_(MakeCoord(strideC0, strideN0, strideN1, StrideDC1HW)) + {} + + /// Ctor + CATLASS_HOST_DEVICE constexpr KDC1KHKWN1N0C0(OrgShape orgShape, Shape shape, Stride stride) + : orgShape_(orgShape), + shape_(shape), + stride_(stride) + {} + + /// Make the layout of a coordinate (Kd*C1*Kh*Kw,N1,N0,C0) + CATLASS_HOST_DEVICE constexpr static KDC1KHKWN1N0C0 MakeLayout(Index KdC1KhKw, Index N1, Index N0, Index C0) + { + return KDC1KHKWN1N0C0(KdC1KhKw, + N1, + N0, + C0, + + C0, + KdC1KhKw, + N0, + N1, + + 1, /// StrideC0 + C0 * N0 * N1, /// StrideDC1HW + C0, /// StrideN0 + C0 * N0 /// StrideN1 + ); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (KdC1KhKw_idx, N1_idx) + CATLASS_HOST_DEVICE + LongIndex GetOffset(Conv3dFracZ3dCoord const &coord) const + { + return LongIndex(coord.kdc1khkw()) * stride_[3] + LongIndex(coord.n1()) * stride_[1]; + } + + /// Returns the layout of a tile. + CATLASS_HOST_DEVICE + KDC1KHKWN1N0C0 GetTileLayout(OrgShape const &tileOriShape) const + { + Shape tileShape = MakeCoord(shape(0), /// C0 + tileOriShape[0], /// Kd*C1*Kh*Kw + shape(2), /// N0 + tileOriShape[1] /// N1 + ); + Stride tileStride = MakeCoord(stride(0), /// TileStrideC0 + stride(2) * tileOriShape[1] * tileOriShape[2], /// TileStrideDC1HW + (LongIndex)shape(0), /// TileStrideN0 + stride(2) * tileOriShape[2] /// TileStrideN1 + ); + return KDC1KHKWN1N0C0(tileOriShape, tileShape, tileStride); + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index orgShape(int idx) const + { + return orgShape_[idx]; + } + + /// Returns the origin shape of the layout + CATLASS_HOST_DEVICE + typename OrgShape::Index &orgShape(int idx) + { + return orgShape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + +private: + /// Origin Shape data member + OrgShape orgShape_; + + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; +} // namespace Catlass::layout + +#endif // CATLASS_KDA_MATRIX_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_tensor.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_tensor.hpp new file mode 100644 index 0000000000..3e8d607104 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_tensor.hpp @@ -0,0 +1,315 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ +#ifndef CATLASS_KDA_TENSOR_HPP +#define CATLASS_KDA_TENSOR_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/kda_conv_coord.hpp" +#include "catlass/detail/kda_alignment.hpp" + +namespace Catlass::layout { + +struct NC1HWC0 { // (Batch, C1, H, W, C0) + /// Logical rank of tensor + static constexpr int RANK = 5; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + /// Ctor + CATLASS_HOST_DEVICE + NC1HWC0(Index batch = 0, Index c1 = 0, Index h = 0, Index w = 0, Index c0 = 0) + : shape_(MakeCoord(batch, c1, h, w, c0)) + { + LongIndex strideBatch = c1 * h * w * c0; + LongIndex strideC1 = h * w * c0; + LongIndex strideH = w * c0; + LongIndex strideW = c0; + LongIndex strideC0 = 1; + stride_ = MakeCoord(strideBatch, strideC1, strideH, strideW, strideC0); + } + + CATLASS_HOST_DEVICE + NC1HWC0( + Index batch, + Index c1, + Index h, + Index w, + Index c0, + LongIndex strideBatch, + LongIndex strideC1, + LongIndex strideH, + LongIndex strideW, + LongIndex strideC0) + : shape_(MakeCoord(batch, c1, h, w, c0)), + stride_(MakeCoord(strideBatch, strideC1, strideH, strideW, strideC0)) + { + } + + CATLASS_HOST_DEVICE + NC1HWC0(Shape shape, Stride stride) + : shape_(shape), + stride_(stride) + { + } + + /// Make the layout of a coordinate (batch, c1, h, w, c0) + template + CATLASS_HOST_DEVICE static NC1HWC0 MakeLayout(Index batch, Index c1, Index h, Index w, Index c0) + { + return NC1HWC0(batch, c1, h, w, c0); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (batch, c1, h, w, c0) + CATLASS_HOST_DEVICE + LongIndex GetOffset(Conv2dFmapCoord const &coord) const + { + return LongIndex(coord.batch()) * stride_[0] + LongIndex(coord.c1()) * stride_[1] + LongIndex(coord.h()) * stride_[2] + LongIndex(coord.w()) * stride_[3] + LongIndex(coord.c0()) * stride_[4]; + } + + /// Returns the layout of a tile. + CATLASS_HOST_DEVICE + NC1HWC0 GetTileLayout(Conv2dFmapCoord const &tileShape) const + { + return NC1HWC0(tileShape, stride()); + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + + /// Returns the length of the layout + CATLASS_HOST_DEVICE + size_t Capacity() + { + return static_cast(shape_[0]) * stride_[0]; + } + +private: + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +struct CI1KHKWCOCI0 { // (Cin1, Kh, Kw, Cout, C0) + /// Logical rank of tensor + static constexpr int RANK = 5; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Logical coordinate + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + +public: + /// Ctor + CATLASS_HOST_DEVICE + CI1KHKWCOCI0(Index cin1 = 0, Index kh = 0, Index kw = 0, Index cout = 0, Index c0 = 0) + : shape_(MakeCoord(cin1, kh, kw, cout, c0)) + { + LongIndex strideCin1 = kh * kw * cout * c0; + LongIndex strideKh = kw * cout * c0; + LongIndex strideKw = cout * c0; + LongIndex strideCout = c0; + LongIndex strideC0 = 1; + stride_ = MakeCoord(strideCin1, strideKh, strideKw, strideCout, strideC0); + } + + CATLASS_HOST_DEVICE + CI1KHKWCOCI0( + Index cin1, + Index kh, + Index kw, + Index cout, + Index c0, + LongIndex strideCin1, + LongIndex strideKh, + LongIndex strideKw, + LongIndex strideCout, + LongIndex strideC0) + : shape_(MakeCoord(cin1, kh, kw, cout, c0)), + stride_(MakeCoord(strideCin1, strideKh, strideKw, strideCout, strideC0)) + { + } + + CATLASS_HOST_DEVICE + CI1KHKWCOCI0(Shape shape, Stride stride) + : shape_(shape), + stride_(stride) + { + } + + /// Make the layout of a coordinate (cin1, h, w, c0) + template + CATLASS_HOST_DEVICE static CI1KHKWCOCI0 MakeLayout(Index cin1, Index kh, Index kw, Index cout, Index c0) + { + return CI1KHKWCOCI0(cin1, kh, kw, cout, c0); + } + + /// Returns the offset of a coordinate in linear memory. + /// Assumes coordinate has convention (cin1, kh, kw, cout, c0) + CATLASS_HOST_DEVICE + LongIndex GetOffset(Conv2dFilterCoord const &coord) const + { + return LongIndex(coord.cin1()) * stride_[0] + LongIndex(coord.kh()) * stride_[1] + LongIndex(coord.kw()) * stride_[2] + LongIndex(coord.cout()) * stride_[3] + LongIndex(coord.c0()) * stride_[4]; + } + + /// Returns the layout of a tile. + CATLASS_HOST_DEVICE + CI1KHKWCOCI0 GetTileLayout(Conv2dFilterCoord const &tileShape) const + { + return CI1KHKWCOCI0(tileShape, stride()); + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + + /// Returns the length of the layout + CATLASS_HOST_DEVICE + size_t Capacity() + { + return static_cast(shape_[0]) * stride_[0]; + } + +private: + /// Shape data member + Shape shape_; + + /// Stride data member + Stride stride_; +}; + +} // namespace Catlass::layout + +#endif // CATLASS_KDA_TENSOR_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_vector.hpp b/csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_vector.hpp new file mode 100644 index 0000000000..2d7750b0c5 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/catlass/layout/kda_vector.hpp @@ -0,0 +1,139 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_KDA_VECTOR_HPP +#define CATLASS_KDA_VECTOR_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_numeric_size.hpp" + +namespace Catlass::layout { + +struct VectorLayout { +public: + /// Logical rank of tensor + static constexpr int RANK = 1; + + /// Index type used for coordinates + using Index = uint32_t; + + /// Long index type used for offsets + using LongIndex = int64_t; + + /// Shape vector + using Shape = Coord; + + /// Stride vector + using Stride = Coord; + + /// Logical coordinate + using TensorCoord = Coord; + +public: + // Methods + + CATLASS_HOST_DEVICE + VectorLayout(Index size = 0) + : shape_(MakeCoord(size)), + stride_(MakeCoord(LongIndex(1))) + {} + + CATLASS_HOST_DEVICE + VectorLayout(Shape shape, Stride stride) + : shape_(shape), + stride_(stride) + {} + + template + CATLASS_HOST_DEVICE static VectorLayout MakeLayoutInUb(TensorCoord const &tileShape) + { + constexpr uint32_t ELE_NUM_PER_BLK = BytesToBits(BYTE_PER_BLK) / SizeOfBits::value; + return VectorLayout{ELE_NUM_PER_BLK > (tileShape[0])}; + } + + CATLASS_HOST_DEVICE + LongIndex GetOffset(TensorCoord const &coord) const + { + return stride_[0] * coord[0]; + } + + /// Returns the layout of a tile. + CATLASS_HOST_DEVICE + VectorLayout GetTileLayout(TensorCoord const &tileShape) const + { + return VectorLayout(tileShape, stride()); + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape shape() const + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + Shape &shape() + { + return shape_; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index shape(int idx) const + { + return shape_[idx]; + } + + /// Returns the shape of the layout + CATLASS_HOST_DEVICE + typename Shape::Index &shape(int idx) + { + return shape_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride stride() const + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + Stride &stride() + { + return stride_; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index stride(int idx) const + { + return stride_[idx]; + } + + /// Returns the stride of the layout + CATLASS_HOST_DEVICE + typename Stride::Index &stride(int idx) + { + return stride_[idx]; + } + +private: + /// Stride data member + Shape shape_; + Stride stride_; +}; + +} // namespace Catlass::layout + +#endif // CATLASS_KDA_VECTOR_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd.cpp b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd.cpp new file mode 100644 index 0000000000..1b4c07461b --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd.cpp @@ -0,0 +1,179 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#include "kernel_operator.h" +#include "lib/matmul_intf.h" + +#include "chunk_kda_fwd_common.h" +#include "chunk_kda_fwd_tiling_data.h" +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +#define KDA_COMPILE_ARCH35_FAST_PATH 1 +#include "arch35/chunk_kda_fwd_impl.h" +#else +#define KDA_COMPILE_ARCH35_FAST_PATH 0 +#endif + +namespace KdaForward { + +template +__aicore__ inline void DispatchGeneric( + GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR g, GM_ADDR beta, + GM_ADDR aLog, GM_ADDR dtBias, GM_ADDR initialState, + GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR attnOut, + GM_ADDR finalState, GM_ADDR gk, GM_ADDR aqk, GM_ADDR akk, + GM_ADDR w, GM_ADDR u, GM_ADDR qg, GM_ADDR kg, GM_ADDR vNew, GM_ADDR h, + GM_ADDR userWorkspace, const TilingData &tiling) +{ + RunGeneric( + q, k, v, g, beta, aLog, dtBias, initialState, cuSeqlens, + chunkIndices, attnOut, finalState, gk, aqk, akk, w, u, qg, kg, + vNew, h, userWorkspace, tiling); +} + +#if KDA_COMPILE_ARCH35_FAST_PATH +template +__aicore__ inline void DispatchArch35SafeGate( + GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR g, GM_ADDR beta, + GM_ADDR aLog, GM_ADDR dtBias, GM_ADDR initialState, + GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR attnOut, + GM_ADDR finalState, GM_ADDR gk, GM_ADDR aqk, GM_ADDR akk, + GM_ADDR w, GM_ADDR u, GM_ADDR qg, GM_ADDR kg, GM_ADDR vNew, GM_ADDR h, + GM_ADDR userWorkspace, const TilingData &tiling) +{ + AscendC::TPipe pipe; + if (tiling.safeGate) { + arch35::Run( + q, k, v, g, beta, aLog, dtBias, initialState, cuSeqlens, + chunkIndices, attnOut, finalState, gk, aqk, akk, w, u, qg, + kg, vNew, h, userWorkspace, tiling, pipe); + } else { + arch35::Run( + q, k, v, g, beta, aLog, dtBias, initialState, cuSeqlens, + chunkIndices, attnOut, finalState, gk, aqk, akk, w, u, qg, + kg, vNew, h, userWorkspace, tiling, pipe); + } +} +#elif defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +template +__aicore__ inline void DispatchArch35SafeGate( + GM_ADDR, GM_ADDR, GM_ADDR, GM_ADDR, GM_ADDR, + GM_ADDR, GM_ADDR, GM_ADDR, + GM_ADDR, GM_ADDR, + GM_ADDR, GM_ADDR, GM_ADDR, GM_ADDR, GM_ADDR, + GM_ADDR, GM_ADDR, GM_ADDR, GM_ADDR, GM_ADDR, GM_ADDR, + GM_ADDR, + const TilingData &) +{ +} +#endif + +template +__aicore__ inline void DispatchGenericSafeGate( + GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR g, GM_ADDR beta, + GM_ADDR aLog, GM_ADDR dtBias, GM_ADDR initialState, + GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR attnOut, + GM_ADDR finalState, GM_ADDR gk, GM_ADDR aqk, GM_ADDR akk, + GM_ADDR w, GM_ADDR u, GM_ADDR qg, GM_ADDR kg, GM_ADDR vNew, GM_ADDR h, + GM_ADDR userWorkspace, const TilingData &tiling) +{ + if (tiling.safeGate) { + DispatchGeneric( + q, k, v, g, beta, aLog, dtBias, initialState, cuSeqlens, + chunkIndices, attnOut, finalState, gk, aqk, akk, w, u, qg, + kg, vNew, h, userWorkspace, tiling); + } else { + DispatchGeneric( + q, k, v, g, beta, aLog, dtBias, initialState, cuSeqlens, + chunkIndices, attnOut, finalState, gk, aqk, akk, w, u, qg, + kg, vNew, h, userWorkspace, tiling); + } +} + +} // namespace KdaForward + +using ChunkKdaFwd::ChunkKdaFwdTilingData; + +// FwdH uses SyncAll() and cross-core flags. The original framework +// operator requests batch scheduling with TilingContext::SetScheduleMode(1). +// Kernel-launch operators must express the same requirement on the entrypoint +// so all requested cores are available together instead of being scheduled in +// partial waves. +// +// ascendc_library's generated device translation unit defines the guard below, +// temporarily rewrites __global__ to inline, and includes this source as the +// implementation body. Schedule-mode attributes are invalid on that inline +// copy, so attach the attribute only to the real kernel declaration processed +// outside the generated inline-body pass. +#if defined(__CHUNK_KDA_FWD__KERNEL_FUN_H__) +#define KDA_FWD_SCHEDULE_MODE +#else +#define KDA_FWD_SCHEDULE_MODE __schedmode__(1) +#endif +extern "C" KDA_FWD_SCHEDULE_MODE __global__ __aicore__ void chunk_kda_fwd( + GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR g, GM_ADDR beta, + GM_ADDR a_log, GM_ADDR dt_bias, GM_ADDR initial_state, + GM_ADDR cu_seqlens, GM_ADDR chunk_indices, GM_ADDR attn_out, + GM_ADDR final_state, GM_ADDR gk, GM_ADDR aqk, GM_ADDR akk, + GM_ADDR w, GM_ADDR u, GM_ADDR qg, GM_ADDR kg, GM_ADDR v_new, GM_ADDR h, + GM_ADDR workspace, GM_ADDR tiling) +{ + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2); + REGISTER_TILING_DEFAULT(ChunkKdaFwdTilingData); + GM_ADDR userWorkspace = AscendC::GetUserWorkspace(workspace); + + // Direct-launch: copy the tiling blob from GM to the stack. The aclnn + // GET_TILING_DATA_WITH_STRUCT macro is unavailable under the plain AscendC + // direct-launch compilation, so the host-serialized struct is rehydrated + // here exactly like chunk_gated_delta_rule does. + constexpr size_t TILING_WORDS = (sizeof(ChunkKdaFwdTilingData) + 7) / 8; + alignas(8) uint64_t tilingBuf[TILING_WORDS]; + const __gm__ uint64_t *gmTilingWords = reinterpret_cast(tiling); + for (size_t i = 0; i < TILING_WORDS; ++i) { + tilingBuf[i] = gmTilingWords[i]; + } + ChunkKdaFwdTilingData *tilingData = reinterpret_cast(tilingBuf); + + // Tiling key 2 is selected when chunk_size == 64 and K == V == 128; the + // host encodes it implicitly through these shape fields, so dispatch on + // them instead of the aclnn TILING_KEY_IS compile-time key. + const bool useFastPath = tilingData->chunkSize == 64 && tilingData->kHeadDim == 128 && + tilingData->vHeadDim == 128; + if (useFastPath) { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + KdaForward::DispatchArch35SafeGate( + q, k, v, g, beta, a_log, dt_bias, initial_state, cu_seqlens, + chunk_indices, attn_out, final_state, gk, aqk, akk, w, u, qg, + kg, v_new, h, userWorkspace, *tilingData); +#else + KdaForward::DispatchGenericSafeGate( + q, k, v, g, beta, a_log, dt_bias, initial_state, cu_seqlens, + chunk_indices, attn_out, final_state, gk, aqk, akk, w, u, qg, + kg, v_new, h, userWorkspace, *tilingData); +#endif + } else { + KdaForward::DispatchGenericSafeGate( + q, k, v, g, beta, a_log, dt_bias, initial_state, cu_seqlens, + chunk_indices, attn_out, final_state, gk, aqk, akk, w, u, qg, + kg, v_new, h, userWorkspace, *tilingData); + } +} + +#undef KDA_FWD_SCHEDULE_MODE + +#undef KDA_COMPILE_ARCH35_FAST_PATH diff --git a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_common.h b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_common.h new file mode 100644 index 0000000000..ec0cfcf6b0 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_common.h @@ -0,0 +1,340 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#pragma once + +#include "kernel_operator.h" +#include "kda_gate_cumsum_kernel.h" +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +#include "arch35/chunk_kda_fwd_prepare_arch35.h" +#include "arch35/chunk_kda_fwd_post_wu_arch35.h" +#include "arch35/chunk_kda_fwd_finalize_arch35.h" +#else +#include "chunk_kda_fwd_prepare.h" +#include "chunk_kda_fwd_post_wu.h" +#include "chunk_kda_fwd_finalize.h" +#endif + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +#include "arch35/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch35.hpp" +#else +#include "arch22/fwd_h/gemm/kernel/kda_fwd_h_kernel_arch22.hpp" +#endif + +namespace KdaForward { + +using namespace AscendC; + +struct GateRuntimeTiling { + int64_t batch; + int64_t t; + int64_t hv; + int64_t k; + int64_t rank; + int64_t chunkSize; + int64_t seqNum; + int64_t hasCuSeqlens; + int64_t hasALog; + int64_t hasDtBias; + int64_t dataType; + int64_t useGateInKernel; + int64_t safeGate; + int64_t inputSequenceMajor; + float lowerBound; + int64_t usedCoreNum; +}; + +struct ChunkKdaFwdAddresses { + GM_ADDR gk; + GM_ADDR finalState; + GM_ADDR w; + GM_ADDR u; + GM_ADDR qg; + GM_ADDR kg; + GM_ADDR vNew; + GM_ADDR h; + GM_ADDR qgScaled; + GM_ADDR uSeed; +}; + +struct FwdHTilingView { + int64_t batch; + int64_t seqlen; + int64_t kNumHead; + int64_t vNumHead; + int64_t kHeadDim; + int64_t vHeadDim; + int64_t chunkSize; + bool useInitialState; + bool storeFinalState; + int64_t isVariedLen; + int64_t shapeBatch; + int64_t tokenBatch; + int64_t vWorkspaceOffset; + int64_t vUpdateWorkspaceOffset; + int64_t kDecayWorkspaceOffset; + int64_t hWorkspaceOffset; + int64_t numSeqWorkspaceOffset; + int64_t numChunksWorkspaceOffset; +}; + +template +__aicore__ inline FwdHTilingView MakeFwdHTiling(const TilingData &tiling) +{ + constexpr bool persistFwdHStateInFp32 = true; + return { + tiling.isVarLen ? tiling.seqNum : tiling.batch, + tiling.seqlen, + tiling.vHeadNum, + tiling.vHeadNum, + tiling.kHeadDim, + tiling.vHeadDim, + tiling.chunkSize, + tiling.hasInitialState, + persistFwdHStateInFp32, + tiling.isVarLen ? 1 : 0, + tiling.isVarLen ? 1 : tiling.batch, + tiling.isVarLen ? tiling.seqNum : 1, + tiling.vWorkspaceOffset, + tiling.vUpdateWorkspaceOffset, + tiling.kDecayWorkspaceOffset, + tiling.hWorkspaceOffset, + tiling.numSeqWorkspaceOffset, + tiling.numChunksWorkspaceOffset, + }; +} + +__aicore__ inline GM_ADDR ResolveStorage( + GM_ADDR output, GM_ADDR userWorkspace, int64_t offset, bool storeOutput) +{ + return storeOutput ? output : userWorkspace + offset; +} + +template +__aicore__ inline ChunkKdaFwdAddresses ResolveAddresses( + GM_ADDR finalState, GM_ADDR gk, GM_ADDR w, GM_ADDR u, GM_ADDR qg, + GM_ADDR kg, GM_ADDR vNew, GM_ADDR h, GM_ADDR userWorkspace, + const TilingData &tiling) +{ + return { + ResolveStorage(gk, userWorkspace, tiling.gkStorageOffset, tiling.storeGk), + ResolveStorage(finalState, userWorkspace, tiling.finalStateStorageOffset, + tiling.storeFinalState), + ResolveStorage(w, userWorkspace, tiling.wStorageOffset, tiling.storeW), + ResolveStorage(u, userWorkspace, tiling.uStorageOffset, tiling.storeU), + ResolveStorage(qg, userWorkspace, tiling.qgStorageOffset, tiling.storeQG), + ResolveStorage(kg, userWorkspace, tiling.kgStorageOffset, tiling.storeKg), + ResolveStorage(vNew, userWorkspace, tiling.vNewStorageOffset, tiling.storeVNew), + ResolveStorage(h, userWorkspace, tiling.hStorageOffset, tiling.storeH), + userWorkspace + tiling.qgScaledOffset, + userWorkspace + tiling.outputScratchOffset, + }; +} + +template +__aicore__ inline GateRuntimeTiling MakeGateTiling(const TilingData &tiling) +{ + return { + tiling.batch, + tiling.seqlen, + tiling.vHeadNum, + tiling.kHeadDim, + tiling.inputRank, + tiling.chunkSize, + tiling.seqNum, + tiling.isVarLen ? 1 : 0, + tiling.hasALog ? 1 : 0, + tiling.hasDtBias ? 1 : 0, + tiling.gateDataType, + tiling.useGateInKernel ? 1 : 0, + tiling.safeGate ? 1 : 0, + tiling.inputSequenceMajor ? 1 : 0, + tiling.lowerBound, + tiling.gateUsedCoreNum, + }; +} + +template +__aicore__ inline void RunGateCumsum( + GM_ADDR g, GM_ADDR aLog, GM_ADDR dtBias, GM_ADDR cuSeqlens, + GM_ADDR gk, const TilingData &tiling) +{ + if (tiling.computeGateInPrepare) { + return; + } + if ASCEND_IS_AIV { + GateRuntimeTiling gateTiling = MakeGateTiling(tiling); + TPipe gatePipe; + if (gateTiling.dataType == 2) { + KdaGateCumsum::DispatchKdaGateCumsum( + g, aLog, dtBias, cuSeqlens, gk, gateTiling, &gatePipe); + } else if (gateTiling.dataType == 1) { + KdaGateCumsum::DispatchKdaGateCumsum( + g, aLog, dtBias, cuSeqlens, gk, gateTiling, &gatePipe); + } else { + KdaGateCumsum::DispatchKdaGateCumsum( + g, aLog, dtBias, cuSeqlens, gk, gateTiling, &gatePipe); + } + } +} + +template +__aicore__ inline void RunFrontEnd( + GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR g, GM_ADDR beta, + GM_ADDR aLog, GM_ADDR dtBias, GM_ADDR initialState, + GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR aqk, GM_ADDR akk, + const ChunkKdaFwdAddresses &addresses, GM_ADDR userWorkspace, + const TilingData &tiling, TPipe &pipe) +{ + RunGateCumsum(g, aLog, dtBias, cuSeqlens, addresses.gk, tiling); + if (!tiling.computeGateInPrepare) { + SyncAll(); + } + GM_ADDR uSeed = (tiling.fusePostWu || tiling.fusePostWuIntoFwdH) ? addresses.u : addresses.uSeed; + GM_ADDR wSeed = addresses.w; +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if (!tiling.fusePostWu && !tiling.fusePostWuIntoFwdH) { + // Arch35 PostWU reads the prepared W seed while producing W. Keep the + // seed in the already-reserved PostWU workspace so output stores can + // never overwrite seed rows that have not been consumed yet. + wSeed = userWorkspace + tiling.postWuScratchOffset; + } +#endif + + KdaPrepare::RunChunkKdaPrepare( + q, k, v, addresses.gk, g, aLog, dtBias, beta, initialState, + cuSeqlens, chunkIndices, aqk, akk, addresses.qg, + addresses.qgScaled, wSeed, uSeed, addresses.kg, + userWorkspace, tiling, pipe, tiling.storeQG); + SyncAll(); + pipe.Reset(); + + if (!tiling.fusePostWu && !tiling.fusePostWuIntoFwdH) { + KdaPostWu::RunChunkKdaPostWu( + q, k, v, addresses.gk, beta, initialState, cuSeqlens, + chunkIndices, wSeed, akk, uSeed, + addresses.w, addresses.u, addresses.kg, addresses.vNew, + userWorkspace, tiling, pipe); + SyncAll(); + pipe.Reset(); + } +} + +template +__aicore__ inline void RunFwdH( + GM_ADDR initialState, GM_ADDR cuSeqlens, GM_ADDR chunkIndices, + const ChunkKdaFwdAddresses &addresses, GM_ADDR userWorkspace, + const TilingData &tiling) +{ + using FwdHKernel = Catlass::Gemm::Kernel::KdaFwdHKernel< + T, float, float, float, TileShapes, true, false, true>; + const auto fwdHTiling = MakeFwdHTiling(tiling); + FwdHKernel stateOp; + stateOp.InitFromData( + addresses.kg, addresses.w, addresses.u, addresses.gk, addresses.gk, + initialState, cuSeqlens, chunkIndices, addresses.h, addresses.vNew, + addresses.finalState, fwdHTiling, + userWorkspace + tiling.fwdHWorkspaceBaseOffset); + stateOp.Process(); +} + +template +__aicore__ inline void RunGenericBackEnd( + GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR beta, GM_ADDR initialState, + GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR aqk, GM_ADDR attnOut, + const ChunkKdaFwdAddresses &addresses, GM_ADDR userWorkspace, + const TilingData &tiling) +{ + if (tiling.vHeadDim > 128) { + RunFwdH( + initialState, cuSeqlens, chunkIndices, addresses, + userWorkspace, tiling); + } else { + RunFwdH( + initialState, cuSeqlens, chunkIndices, addresses, + userWorkspace, tiling); + } + SyncAll(); + TPipe pipe; + KdaFinalize::RunChunkKdaOutput( + q, k, v, addresses.gk, beta, initialState, cuSeqlens, + chunkIndices, addresses.qgScaled, aqk, + addresses.vNew, addresses.h, attnOut, userWorkspace, tiling, pipe); +} + +template +__aicore__ inline void RunGenericBackEnd( + GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR beta, GM_ADDR initialState, + GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR aqk, GM_ADDR attnOut, + const ChunkKdaFwdAddresses &addresses, GM_ADDR userWorkspace, + const TilingData &tiling, TPipe &pipe) +{ + if (tiling.vHeadDim > 128) { + RunFwdH( + initialState, cuSeqlens, chunkIndices, addresses, + userWorkspace, tiling); + } else { + RunFwdH( + initialState, cuSeqlens, chunkIndices, addresses, + userWorkspace, tiling); + } + SyncAll(); + KdaFinalize::RunChunkKdaOutput( + q, k, v, addresses.gk, beta, initialState, cuSeqlens, + chunkIndices, addresses.qgScaled, aqk, + addresses.vNew, addresses.h, attnOut, userWorkspace, tiling, pipe); +} + +template +__aicore__ inline void RunGeneric( + GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR g, GM_ADDR beta, + GM_ADDR aLog, GM_ADDR dtBias, GM_ADDR initialState, + GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR attnOut, + GM_ADDR finalState, GM_ADDR gk, GM_ADDR aqk, GM_ADDR akk, + GM_ADDR w, GM_ADDR u, GM_ADDR qg, GM_ADDR kg, GM_ADDR vNew, GM_ADDR h, + GM_ADDR userWorkspace, const TilingData &tiling) +{ + const auto addresses = ResolveAddresses( + finalState, gk, w, u, qg, kg, vNew, h, userWorkspace, tiling); +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + TPipe pipe; + RunFrontEnd( + q, k, v, g, beta, aLog, dtBias, initialState, cuSeqlens, + chunkIndices, aqk, akk, addresses, userWorkspace, tiling, pipe); + if (!tiling.isVarLen && tiling.seqlen % tiling.chunkSize == 0) { + pipe.Destroy(); + RunGenericBackEnd( + q, k, v, beta, initialState, cuSeqlens, chunkIndices, aqk, + attnOut, addresses, userWorkspace, tiling); + } else { + RunGenericBackEnd( + q, k, v, beta, initialState, cuSeqlens, chunkIndices, aqk, + attnOut, addresses, userWorkspace, tiling, pipe); + } +#else + { + TPipe pipe; + RunFrontEnd( + q, k, v, g, beta, aLog, dtBias, initialState, cuSeqlens, + chunkIndices, aqk, akk, addresses, userWorkspace, tiling, pipe); + } + RunGenericBackEnd( + q, k, v, beta, initialState, cuSeqlens, chunkIndices, aqk, + attnOut, addresses, userWorkspace, tiling); +#endif +} + +} // namespace KdaForward diff --git a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_finalize.h b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_finalize.h new file mode 100644 index 0000000000..6654d36203 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_finalize.h @@ -0,0 +1,808 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#pragma once + +#ifndef CATLASS_ARCH +#define CATLASS_ARCH 2201 +#endif + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/arch/kda_cross_core_sync.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/gemm/block/kda_block_mmad.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "kernel_utils/block/block_mmad_pingpong_tla_multi.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "kernel_operator.h" +#include "chunk_kda_fwd_varlen.h" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +using namespace AscendC; + +namespace KdaFinalize { +namespace { +using KdaInt64 = tla::Int<64>; +using KdaInt128 = tla::Int<128>; +constexpr float LN2 = 0.69314718055994530942f; +constexpr float KDA_EXP2_CLAMP = 80.0f; +constexpr float KDA_EXP_INPUT_MAX = KDA_EXP2_CLAMP * LN2; +constexpr float KDA_EXP_INPUT_MIN = -KDA_EXP2_CLAMP * LN2; +constexpr float KDA_FP16_MAX = 65504.0f; +constexpr uint32_t EXP2_UB_ELEMENTS = 256; +constexpr uint32_t EXP2_UB_BYTES = EXP2_UB_ELEMENTS * (sizeof(float) + sizeof(uint16_t)); +constexpr uint32_t EXP2_EVENT_ID = 0; +constexpr uint32_t KDA_SOLVE_BT = 64; +constexpr uint32_t KDA_SOLVE_MATRIX_ELEMENTS = KDA_SOLVE_BT * KDA_SOLVE_BT; +constexpr uint32_t KDA_SOLVE_SCRATCH_X = 0; +constexpr uint32_t KDA_SOLVE_SCRATCH_Y0 = 1; +constexpr uint32_t KDA_SOLVE_SCRATCH_TMP = 2; +constexpr uint32_t KDA_SOLVE_SCRATCH_Y1 = 3; +constexpr uint32_t KDA_SOLVE_SCRATCH_IDENTITY = 4; +constexpr uint32_t KDA_SOLVE_SCRATCH_SLOTS = 5; +constexpr uint32_t KDA_SOLVE_DIAG_BT = 16; +constexpr uint32_t KDA_SOLVE_DIAG_BLOCKS = KDA_SOLVE_BT / KDA_SOLVE_DIAG_BT; +constexpr uint32_t KDA_SOLVE_DIAG_MCH_ITERS = 3; +constexpr uint32_t KDA_SCORE_REF_BC = 16; +constexpr uint32_t KDA_VEC_ARENA_ELEMENTS = 32768; +constexpr uint32_t KDA_BITS_PER_MASK_BYTE = 8; +constexpr uint32_t KDA_SELECT_COL_BLOCKS = 2; +constexpr uint32_t KDA_SELECT_COL_MASK_BYTES = KDA_SOLVE_MATRIX_ELEMENTS / KDA_BITS_PER_MASK_BYTE; +constexpr uint32_t KDA_SELECT_MASK_BYTES = KDA_SELECT_COL_BLOCKS * KDA_SELECT_COL_MASK_BYTES; +constexpr uint32_t KDA_SELECT_AQK_MASK_BYTE_OFFSET = 120 * 1024; +constexpr uint32_t KDA_SELECT_AKK_MASK_BYTE_OFFSET = KDA_SELECT_AQK_MASK_BYTE_OFFSET + KDA_SELECT_MASK_BYTES; +constexpr uint32_t KDA_SELECT_ZERO_BYTE_OFFSET = KDA_SELECT_AKK_MASK_BYTE_OFFSET + KDA_SELECT_MASK_BYTES; +constexpr uint32_t KDA_SELECT_ZERO_FLOAT_OFFSET = KDA_SELECT_ZERO_BYTE_OFFSET / sizeof(float); +constexpr uint8_t KDA_SCORE_DONE_FLAG0 = 2; +constexpr uint8_t KDA_SCORE_DONE_FLAG1 = 3; +constexpr uint8_t KDA_SCORE_READY_FLAG0 = 4; +constexpr uint8_t KDA_SCORE_READY_FLAG1 = 5; +constexpr uint32_t KDA_SCORE_QUEUE_DEPTH = 2; +constexpr uint32_t KDA_SYNC_REVERSE_DEPTH = 1; +constexpr uint32_t KDA_SCORE_SCRATCH_PLANES = 3; +constexpr uint32_t KDA_SCORE_SCRATCH_QG = 0; +constexpr uint32_t KDA_SCORE_SCRATCH_W = 1; +constexpr uint32_t KDA_SCORE_SCRATCH_KG = 2; +constexpr uint64_t KDA_WORKSPACE_ALIGN = 512; +constexpr uint32_t KDA_GATE_TILE_ROWS = 32; + +using KdaArchTag = Catlass::Arch::AtlasA2; +using KdaDispatchPolicy = Catlass::Gemm::MmadPingpong; +using KdaScoreDispatchPolicy = Catlass::Gemm::MmadPingpongTlaMulti; +static_assert(KdaScoreDispatchPolicy::ENABLE_L1_RESIDENT, + "KDA Aqk/Akk score MMAD must keep the shared right matrix resident in L1"); +static_assert(KdaScoreDispatchPolicy::L1B_STAGES == 1, + "KDA Aqk/Akk score MMAD needs one L1 B slot so the second MMAD reuses it"); +using KdaSolveDispatchPolicy = Catlass::Gemm::MmadPingpong; +static_assert(!KdaSolveDispatchPolicy::USE_HF32_MODE, "KDA triangular solve must use IEEE FP32 Cube mode"); +using KdaL1TileShape = tla::Shape; +using KdaL0TileShape = KdaL1TileShape; +using KdaSolveL1TileShape = tla::Shape; +using KdaSolveL0TileShape = KdaSolveL1TileShape; + +__aicore__ inline uint32_t FloatToBits(float value) +{ + union Bits { + __aicore__ Bits() {} + float f; + uint32_t u; + } bits; + bits.f = value; + return bits.u; +} + +__aicore__ inline float BitsToFloat(uint32_t value) +{ + union Bits { + __aicore__ Bits() {} + uint32_t u; + float f; + } bits; + bits.u = value; + return bits.f; +} + +__aicore__ inline uint16_t Bf16ToBits(bfloat16_t value) +{ + union Bits { + __aicore__ Bits() {} + bfloat16_t f; + uint16_t u; + } bits; + bits.f = value; + return bits.u; +} + +__aicore__ inline bfloat16_t BitsToBf16(uint16_t value) +{ + union Bits { + __aicore__ Bits() {} + uint16_t u; + bfloat16_t f; + } bits; + bits.u = value; + return bits.f; +} + +template +__aicore__ inline T FloatToType(float value) +{ + if constexpr (IsSameType::value) { + uint32_t bits = FloatToBits(value); + uint32_t bias = 0x7FFFu + ((bits >> 16) & 1u); + return BitsToBf16(static_cast((bits + bias) >> 16)); + } + return static_cast(value); +} + +template +class ChunkKdaFwdFinalizeKernel { +public: + using OUT_T = float; + using AKK_T = float; + template + __aicore__ inline void Init(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR beta, GM_ADDR initialState, + GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR preparedQG, GM_ADDR preparedAqk, + GM_ADDR propagatedVNew, GM_ADDR propagatedH, GM_ADDR o, GM_ADDR finalState, GM_ADDR aqk, + GM_ADDR akk, GM_ADDR w, GM_ADDR u, GM_ADDR qg, GM_ADDR kg, GM_ADDR vNew, GM_ADDR h, + GM_ADDR workspace, const TilingData &tiling, TPipe *pipe, bool initVecBuffers = true) + { + pipe_ = pipe; + q_.SetGlobalBuffer((__gm__ T *)q); + k_.SetGlobalBuffer((__gm__ T *)k); + v_.SetGlobalBuffer((__gm__ T *)v); + gk_.SetGlobalBuffer((__gm__ GK_T *)gk); + if (initialState != nullptr) { + initialState_.SetGlobalBuffer((__gm__ float *)initialState); + } + cuSeqlensAddr_ = reinterpret_cast<__gm__ int64_t *>(cuSeqlens); + if (preparedQG != nullptr) { + preparedQG_.SetGlobalBuffer((__gm__ T *)preparedQG); + } + if (preparedAqk != nullptr) { + preparedAqk_.SetGlobalBuffer((__gm__ T *)preparedAqk); + } + if (propagatedVNew != nullptr) { + propagatedVNew_.SetGlobalBuffer((__gm__ T *)propagatedVNew); + } + if (propagatedH != nullptr) { + propagatedH_.SetGlobalBuffer((__gm__ T *)propagatedH); + } + chunkIndicesAddr_ = reinterpret_cast<__gm__ int64_t *>(chunkIndices); + o_.SetGlobalBuffer((__gm__ OUT_T *)o); + finalState_.SetGlobalBuffer((__gm__ float *)finalState); + aqk_.SetGlobalBuffer((__gm__ float *)aqk); + akk_.SetGlobalBuffer((__gm__ AKK_T *)akk); + w_.SetGlobalBuffer((__gm__ T *)w); + u_.SetGlobalBuffer((__gm__ OUT_T *)u); + qg_.SetGlobalBuffer((__gm__ T *)qg); + kg_.SetGlobalBuffer((__gm__ T *)kg); + vNew_.SetGlobalBuffer((__gm__ T *)vNew); + h_.SetGlobalBuffer((__gm__ float *)h); + solveWorkspace_.SetGlobalBuffer((__gm__ float *)workspace); + + B_ = tiling.batch; + N_ = tiling.seqNum; + H_ = tiling.qHeadNum; + HV_ = tiling.vHeadNum; + T_ = tiling.seqlen; + K_ = tiling.kHeadDim; + V_ = tiling.vHeadDim; + BT_ = tiling.chunkSize; + NT_ = tiling.totalChunks; + scale_ = tiling.scale; + hasInitial_ = tiling.hasInitialState; + isVarLen_ = tiling.isVarLen; + usedCoreNum_ = tiling.outputUsedCoreNum; + const uint64_t outputElements = B_ * HV_ * T_ * V_; + o_.SetGlobalBuffer((__gm__ OUT_T *)workspace); + u_.SetGlobalBuffer((__gm__ OUT_T *)workspace + outputElements); + if ASCEND_IS_AIV { + uint64_t subBlockNum = static_cast(GetSubBlockNum()); + solveCoreIdx_ = subBlockNum == 0 ? 0 : static_cast(GetBlockIdx()) / subBlockNum; + } else { + solveCoreIdx_ = static_cast(GetBlockIdx()); + } + if (pipe_ != nullptr && initVecBuffers) { + pipe_->InitBuffer(exp2Buf_, EXP2_UB_BYTES); + pipe_->InitBuffer(vecBuf_, KDA_VEC_ARENA_ELEMENTS * sizeof(float)); + const uint64_t gateWritebackRows = ScoreVectorMaxRows(5 * sizeof(float) + 2 * sizeof(T) + sizeof(GK_T)); + pipe_->InitBuffer(gateWritebackBuf_, + static_cast(gateWritebackRows * K_ * (3 * sizeof(T) + sizeof(GK_T)))); + AllocVectorEvents(); + } + } + __aicore__ inline void ProcessAiv() + { + ProcessOutAiv(); + ReleaseVectorEvents(); + } + + __aicore__ inline void ProcessAic() { ProcessOutAic(); } + +private: + __aicore__ inline void AllocVectorEvents() + { + mte2ToVEvent_ = pipe_->AllocEventID(); + vToMte2Event_ = pipe_->AllocEventID(); + vToMte3Event_ = pipe_->AllocEventID(); + mte3ToVEvent_ = pipe_->AllocEventID(); + mte2ToMte3Event_ = pipe_->AllocEventID(); + mte3ToMte2Event_ = pipe_->AllocEventID(); + vectorEventsAllocated_ = true; + } + + __aicore__ inline void ReleaseVectorEvents() + { + if (!vectorEventsAllocated_) { + return; + } + pipe_->ReleaseEventID(mte2ToVEvent_); + pipe_->ReleaseEventID(vToMte2Event_); + pipe_->ReleaseEventID(vToMte3Event_); + pipe_->ReleaseEventID(mte3ToVEvent_); + pipe_->ReleaseEventID(mte2ToMte3Event_); + pipe_->ReleaseEventID(mte3ToMte2Event_); + vectorEventsAllocated_ = false; + } + + __aicore__ inline uint64_t QOffset(uint64_t b, uint64_t h, uint64_t t, uint64_t d) const + { + return ((b * H_ + h) * T_ + t) * K_ + d; + } + + __aicore__ inline uint64_t KVOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t d, uint64_t dim) const + { + return ((b * HV_ + hv) * T_ + t) * dim + d; + } + + __aicore__ inline uint64_t OutputOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t d) const + { + return ((b * T_ + t) * HV_ + hv) * V_ + d; + } + + __aicore__ inline uint64_t BetaOffset(uint64_t b, uint64_t hv, uint64_t t) const { return (b * HV_ + hv) * T_ + t; } + + __aicore__ inline uint64_t AOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t j) const + { + return ((b * HV_ + hv) * T_ + t) * BT_ + j; + } + + __aicore__ inline uint64_t HOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t d, uint64_t r) const + { + return (((b * HV_ + hv) * NT_ + chunkIdx) * K_ + d) * V_ + r; + } + + __aicore__ inline uint64_t WScratchOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t t, uint64_t d) const + { + return (((b * HV_ + hv) * NT_ + chunkIdx) * BT_ + t) * K_ + d; + } + + __aicore__ inline uint64_t SolveScratchOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t slot) const + { + (void)b; + (void)hv; + (void)chunkIdx; + uint64_t matrixElements = BT_ * BT_; + return solveCoreIdx_ * KDA_SOLVE_SCRATCH_SLOTS * matrixElements + slot * matrixElements; + } + + __aicore__ inline uint64_t ScoreScratchOffset(uint64_t slot, uint64_t plane, uint64_t t = 0, uint64_t d = 0) const + { + return (((solveCoreIdx_ * KDA_SCORE_QUEUE_DEPTH + slot) * KDA_SCORE_SCRATCH_PLANES + plane) * BT_ + t) * K_ + d; + } + + __aicore__ inline uint64_t ScoreRefBlockSize() const { return KDA_SCORE_REF_BC; } + + __aicore__ inline uint64_t ScoreRowBlockCount(uint64_t curT, uint64_t rowBegin) const + { + uint64_t blockSize = ScoreRefBlockSize(); + uint64_t rowCount = curT - rowBegin; + if (rowCount > blockSize) { + rowCount = blockSize; + } + return rowCount; + } + + __aicore__ inline uint64_t ScoreRefToken(uint64_t start, uint64_t curT, uint64_t rowBegin, uint64_t rowCount) const + { + uint64_t ref = rowBegin + rowCount / 2; + if (ref >= curT) { + ref = curT - 1; + } + return start + ref; + } + + __aicore__ inline void RunExp2(LocalTensor &tensor, uint32_t count) + { + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + ClampExpInput(tensor, count); + Exp(tensor, tensor, count); + PipeBarrier(); + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + } + + __aicore__ inline void ClampExpInput(LocalTensor &tensor, uint32_t count) + { + Mins(tensor, tensor, KDA_EXP_INPUT_MAX, count); + PipeBarrier(); + Maxs(tensor, tensor, KDA_EXP_INPUT_MIN, count); + PipeBarrier(); + } + + __aicore__ inline void ClampFp32ToOutputType(LocalTensor &tensor, uint32_t count) + { + if constexpr (IsSameType::value) { + Mins(tensor, tensor, KDA_FP16_MAX, count); + PipeBarrier(); + Maxs(tensor, tensor, -KDA_FP16_MAX, count); + PipeBarrier(); + } + } + + template + __aicore__ inline void CopyVectorIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset, + uint64_t count) + { + uint64_t rowBytes = count * static_cast(sizeof(CopyT)); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst, src[offset], static_cast(count)); + return; + } + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPadParams padParams{false, 0, 0, 0}; + DataCopyPad(dst, src[offset], params, padParams); + } + + template + __aicore__ inline void CopyVectorOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src, + uint64_t count) + { + uint64_t rowBytes = count * static_cast(sizeof(CopyT)); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst[offset], src, static_cast(count)); + return; + } + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPad(dst[offset], src, params); + } + + template + __aicore__ inline void CopyRowsOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src, + uint64_t rows, uint64_t cols, uint64_t dstStride) + { + if (cols == dstStride) { + CopyVectorOut(dst, offset, src, rows * cols); + return; + } + constexpr uint64_t blockBytes = 32; + const uint64_t rowBytes = cols * sizeof(CopyT); + const uint64_t gapBytes = (dstStride - cols) * sizeof(CopyT); + DataCopyParams params{static_cast(rows), static_cast(rowBytes / blockBytes), 0, + static_cast(gapBytes / blockBytes)}; + DataCopy(dst[offset], src, params); + } + + template + __aicore__ inline void CopyRowIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset) + { + CopyVectorIn(dst, src, offset, K_); + } + + template + __aicore__ inline void CopyRowOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src) + { + CopyVectorOut(dst, offset, src, K_); + } + + __aicore__ inline LocalTensor VecScratch(uint64_t slot) + { + return vecBuf_.Get()[slot * EXP2_UB_ELEMENTS]; + } + + template + __aicore__ inline void LoadAsFloatRow(GlobalTensor &src, uint64_t srcOffset, LocalTensor &dst, + uint64_t count) + { + if constexpr (IsSameType::value) { + CopyVectorIn(dst, src, srcOffset, count); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Adds(dst, dst, 0.0f, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } else { + constexpr uint32_t typedOffset = EXP2_UB_ELEMENTS * sizeof(float) / sizeof(CopyT); + LocalTensor rowLocal = exp2Buf_.Get()[typedOffset]; + CopyVectorIn(rowLocal, src, srcOffset, count); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Cast(dst, rowLocal, RoundMode::CAST_NONE, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } + PipeBarrier(); + } + + template + __aicore__ inline void LoadAsFloatVector(GlobalTensor &src, uint64_t srcOffset, LocalTensor &dst, + LocalTensor &typedScratch, uint64_t count) + { + if constexpr (IsSameType::value) { + CopyVectorIn(dst, src, srcOffset, count); + } else { + CopyVectorIn(typedScratch, src, srcOffset, count); + } + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + if constexpr (!IsSameType::value) { + Cast(dst, typedScratch, RoundMode::CAST_NONE, static_cast(count)); + PipeBarrier(); + } + } + + template + __aicore__ inline void StoreFloatRow(GlobalTensor &dst, uint64_t dstOffset, LocalTensor &src, + uint64_t count) + { + if constexpr (IsSameType::value) { + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(dst, dstOffset, src, count); + } else { + constexpr uint32_t typedOffset = EXP2_UB_ELEMENTS * sizeof(float) / sizeof(CopyT); + LocalTensor rowLocal = exp2Buf_.Get()[typedOffset]; + Cast(rowLocal, src, RoundMode::CAST_RINT, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(dst, dstOffset, rowLocal, count); + } + SetFlag(mte3ToMte2Event_); + WaitFlag(mte3ToMte2Event_); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline LocalTensor Exp2NegG(uint64_t b, uint64_t hv, uint64_t t) + { + LocalTensor exp2Local = exp2Buf_.Get(); + LoadAsFloatRow(gk_, KVOffset(b, hv, t, 0, K_), exp2Local, K_); + Muls(exp2Local, exp2Local, -LN2, static_cast(K_)); + PipeBarrier(); + RunExp2(exp2Local, static_cast(K_)); + return exp2Local; + } + + __aicore__ inline uint64_t ScoreVectorMaxRows(uint64_t bytesPerElem) const + { + constexpr uint64_t arenaBytes = static_cast(KDA_VEC_ARENA_ELEMENTS) * sizeof(float); + uint64_t maxRows = (arenaBytes / bytesPerElem) / K_; + if (K_ >= 128 && maxRows > 32) { + maxRows = 32; + } + return maxRows; + } + + __aicore__ inline void ComputeOutputCube(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t curT) + { + using ElementA = T; + using ElementB = T; + using ElementC = OUT_T; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::RowMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using BlockMmad = Catlass::Gemm::Block::BlockMmadTla; + + Catlass::Arch::Resource resource; + BlockMmad blockMmad(resource); + + auto layoutQ = tla::MakeLayout(BT_, K_); + auto layoutH = tla::MakeLayout(K_, V_); + auto layoutO = tla::MakeLayout(BT_, V_); + for (uint64_t nOffset = 0; nOffset < V_; nOffset += 128) { + uint32_t curN = static_cast((V_ - nOffset) > 128 ? 128 : (V_ - nOffset)); + auto tensorH = tla::MakeTensor(propagatedH_[HOffset(b, hv, chunkIdx, 0, nOffset)], layoutH, + Catlass::Arch::PositionGM{}); + for (uint64_t mOffset = 0; mOffset < curT; mOffset += 64) { + uint32_t curM = static_cast((curT - mOffset) > 64 ? 64 : (curT - mOffset)); + Catlass::GemmCoord shapeQH{curM, curN, static_cast(K_)}; + auto tensorQ = tla::MakeTensor(preparedQG_[KVOffset(b, hv, start + mOffset, 0, K_)], layoutQ, + Catlass::Arch::PositionGM{}); + auto tensorO = tla::MakeTensor(o_[KVOffset(b, hv, start + mOffset, nOffset, V_)], layoutO, + Catlass::Arch::PositionGM{}); + auto blockQ = GetTile(tensorQ, tla::MakeCoord(0, 0), tla::MakeShape(shapeQH.m(), shapeQH.k())); + auto blockH = GetTile(tensorH, tla::MakeCoord(0, 0), tla::MakeShape(shapeQH.k(), shapeQH.n())); + auto blockO = GetTile(tensorO, tla::MakeCoord(0, 0), tla::MakeShape(shapeQH.m(), shapeQH.n())); + blockMmad(blockQ, blockH, blockO, shapeQH); + PipeBarrier(); + } + } + + auto layoutAqk = tla::MakeLayout(BT_, BT_); + auto layoutV = tla::MakeLayout(BT_, V_); + for (uint64_t nOffset = 0; nOffset < V_; nOffset += 128) { + uint32_t curN = static_cast((V_ - nOffset) > 128 ? 128 : (V_ - nOffset)); + auto tensorVNew = tla::MakeTensor(propagatedVNew_[KVOffset(b, hv, start, nOffset, V_)], layoutV, + Catlass::Arch::PositionGM{}); + for (uint64_t mOffset = 0; mOffset < curT; mOffset += 64) { + uint32_t curM = static_cast((curT - mOffset) > 64 ? 64 : (curT - mOffset)); + Catlass::GemmCoord shapeAV{curM, curN, static_cast(curT)}; + auto tensorAqk = tla::MakeTensor(preparedAqk_[AOffset(b, hv, start + mOffset, 0)], layoutAqk, + Catlass::Arch::PositionGM{}); + auto tensorLocal = tla::MakeTensor(u_[KVOffset(b, hv, start + mOffset, nOffset, V_)], layoutO, + Catlass::Arch::PositionGM{}); + auto blockAqk = GetTile(tensorAqk, tla::MakeCoord(0, 0), tla::MakeShape(shapeAV.m(), shapeAV.k())); + auto blockVNew = GetTile(tensorVNew, tla::MakeCoord(0, 0), tla::MakeShape(shapeAV.k(), shapeAV.n())); + auto blockLocal = GetTile(tensorLocal, tla::MakeCoord(0, 0), tla::MakeShape(shapeAV.m(), shapeAV.n())); + blockMmad(blockAqk, blockVNew, blockLocal, shapeAV); + PipeBarrier(); + } + } + } + + __aicore__ inline void FinalizeOutputRows(uint64_t b, uint64_t hv, uint64_t start, uint64_t curT, + uint64_t subBlockIdx, uint64_t subBlockNum) + { + if (subBlockNum == 0 || subBlockIdx >= subBlockNum || V_ == 0) { + return; + } + const uint64_t rowBegin = (curT * subBlockIdx) / subBlockNum; + const uint64_t rowEnd = (curT * (subBlockIdx + 1)) / subBlockNum; + const uint64_t gateWritebackRows = ScoreVectorMaxRows(5 * sizeof(float) + 2 * sizeof(T) + sizeof(GK_T)); + const uint64_t gateWritebackBytes = gateWritebackRows * K_ * (3 * sizeof(T) + sizeof(GK_T)); + uint64_t maxRows = KDA_VEC_ARENA_ELEMENTS / (3 * V_); + const uint64_t typedMaxRows = gateWritebackBytes / (V_ * sizeof(T)); + if (maxRows > typedMaxRows) { + maxRows = typedMaxRows; + } + if (maxRows == 0) { + return; + } + + for (uint64_t tileRow = rowBegin; tileRow < rowEnd; tileRow += maxRows) { + uint64_t tileRows = rowEnd - tileRow; + if (tileRows > maxRows) { + tileRows = maxRows; + } + const uint64_t elems = tileRows * V_; + const uint64_t ti = start + tileRow; + LocalTensor arena = vecBuf_.Get(); + LocalTensor stateLocal = arena; + LocalTensor localLocal = arena[elems]; + LocalTensor outLocal = arena[2 * elems]; + LocalTensor outTyped = gateWritebackBuf_.Get(); + + CopyVectorIn(stateLocal, o_, KVOffset(b, hv, ti, 0, V_), elems); + CopyVectorIn(localLocal, u_, KVOffset(b, hv, ti, 0, V_), elems); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Muls(stateLocal, stateLocal, scale_, static_cast(elems)); + PipeBarrier(); + Add(outLocal, stateLocal, localLocal, static_cast(elems)); + PipeBarrier(); + ClampFp32ToOutputType(outLocal, static_cast(elems)); + Cast(outTyped, outLocal, RoundMode::CAST_RINT, static_cast(elems)); + PipeBarrier(); + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyRowsOut(vNew_, OutputOffset(b, hv, ti, 0), outTyped, tileRows, V_, HV_ * V_); + SetFlag(mte3ToMte2Event_); + WaitFlag(mte3ToMte2Event_); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + } + __aicore__ inline bool ResolveFlatChunk(uint64_t task, uint64_t &seq, uint64_t &b, uint64_t &h, uint64_t &hv, + uint64_t &chunkIdx, uint64_t &start, uint64_t &end) + { + hv = task % HV_; + uint64_t flatChunk = task / HV_; + if (!isVarLen_) { + seq = flatChunk / NT_; + b = seq; + chunkIdx = flatChunk % NT_; + start = chunkIdx * BT_; + end = start + BT_; + if (end > T_) { + end = T_; + } + } else { + if (!KdaVarlen::ResolveChunkRange(cuSeqlensAddr_, chunkIndicesAddr_, N_, T_, BT_, flatChunk, seq, start, + end)) { + return false; + } + b = 0; + chunkIdx = flatChunk; + } + h = hv / (HV_ / H_); + return start < end; + } + + __aicore__ inline void ProcessChunkOutAiv(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t end, + uint64_t subBlockIdx, uint64_t subBlockNum) + { + uint64_t curT = end - start; + if (curT == 0) { + return; + } + if constexpr (IsSameType::value) { + return; + } + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE2>(syncDoneFlag_); + FinalizeOutputRows(b, hv, start, curT, subBlockIdx, subBlockNum); + } + + __aicore__ inline void ProcessChunkOutAic(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t end) + { + uint64_t curT = end - start; + if (curT == 0) { + return; + } + ComputeOutputCube(b, hv, chunkIdx, start, curT); + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(syncDoneFlag_); + } + + __aicore__ inline void ProcessOutAiv() + { + if constexpr (IsSameType::value) { + return; + } + uint64_t subBlockNum = static_cast(GetSubBlockNum()); + if (subBlockNum == 0) { + return; + } + uint64_t subBlockIdx = static_cast(GetSubBlockIdx()); + uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + uint64_t coreIdx = static_cast(GetBlockIdx()) / subBlockNum; + uint64_t taskNum = static_cast((isVarLen_ ? NT_ : B_ * NT_) * HV_); + for (uint64_t task = coreIdx; task < taskNum; task += coreNum) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveFlatChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + (void)seq; + (void)h; + (void)chunkIdx; + ProcessChunkOutAiv(b, hv, chunkIdx, start, end, subBlockIdx, subBlockNum); + } + } + } + + __aicore__ inline void ProcessOutAic() + { + if constexpr (IsSameType::value) { + return; + } + uint64_t taskNum = static_cast((isVarLen_ ? NT_ : B_ * NT_) * HV_); + uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + for (uint64_t task = GetBlockIdx(); task < taskNum; task += coreNum) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveFlatChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + (void)seq; + (void)h; + ProcessChunkOutAic(b, hv, chunkIdx, start, end); + } + } + } + +private: + GlobalTensor q_; + GlobalTensor k_; + GlobalTensor v_; + GlobalTensor gk_; + GlobalTensor initialState_; + GlobalTensor o_; + GlobalTensor finalState_; + GlobalTensor aqk_; + GlobalTensor akk_; + GlobalTensor w_; + GlobalTensor u_; + GlobalTensor qg_; + GlobalTensor kg_; + GlobalTensor vNew_; + GlobalTensor h_; + GlobalTensor preparedQG_; + GlobalTensor preparedAqk_; + GlobalTensor propagatedVNew_; + GlobalTensor propagatedH_; + GlobalTensor solveWorkspace_; + GlobalTensor scoreWorkspace_; + TPipe *pipe_ = nullptr; + TBuf exp2Buf_; + TBuf vecBuf_; + TBuf gateWritebackBuf_; + TEventID mte2ToVEvent_ = 0; + TEventID vToMte2Event_ = 0; + TEventID vToMte3Event_ = 0; + TEventID mte3ToVEvent_ = 0; + TEventID mte2ToMte3Event_ = 0; + TEventID mte3ToMte2Event_ = 0; + bool vectorEventsAllocated_ = false; + Catlass::Arch::CrossCoreFlagWithReverse scoreReadyFlag_{KDA_SCORE_READY_FLAG0, + KDA_SCORE_READY_FLAG1}; + Catlass::Arch::CrossCoreFlagWithReverse scoreDoneFlag_{KDA_SCORE_DONE_FLAG0, + KDA_SCORE_DONE_FLAG1}; + // Score production is fully drained before solve starts, so the solve handshake can safely reuse + // the existing score flags without consuming additional hardware flag IDs. + Catlass::Arch::CrossCoreFlagWithReverse syncReadyFlag_{KDA_SCORE_READY_FLAG0, + KDA_SCORE_READY_FLAG1}; + Catlass::Arch::CrossCoreFlagWithReverse syncDoneFlag_{KDA_SCORE_DONE_FLAG0, + KDA_SCORE_DONE_FLAG1}; + uint64_t B_ = 0; + uint64_t N_ = 0; + uint64_t H_ = 0; + uint64_t HV_ = 0; + uint64_t T_ = 0; + uint64_t K_ = 0; + uint64_t V_ = 0; + uint64_t BT_ = 0; + uint64_t NT_ = 0; + float scale_ = 1.0f; + bool hasInitial_ = false; + bool isVarLen_ = false; + bool isAivOnly_ = false; + uint64_t usedCoreNum_ = 1; + uint64_t solveCoreIdx_ = 0; + __gm__ int64_t *chunkIndicesAddr_ = nullptr; + __gm__ int64_t *cuSeqlensAddr_ = nullptr; +}; +} // namespace + +template +__aicore__ inline void RunChunkKdaOutput(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR beta, + GM_ADDR initialState, GM_ADDR cuSeqlens, GM_ADDR chunkIndices, + GM_ADDR qgScaled, GM_ADDR aqk, GM_ADDR propagatedVNew, GM_ADDR propagatedH, + GM_ADDR o, GM_ADDR userWorkspace, const TilingData &tiling, TPipe &pipe) +{ + GM_ADDR outputScratch = userWorkspace + tiling.outputScratchOffset; + uint64_t outputElements = static_cast(tiling.batch) * static_cast(tiling.vHeadNum) * + static_cast(tiling.seqlen) * static_cast(tiling.vHeadDim); + GM_ADDR stateScratch = outputScratch; + GM_ADDR localScratch = outputScratch + outputElements * sizeof(float); + if ASCEND_IS_AIC { + ChunkKdaFwdFinalizeKernel op; + op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, qgScaled, aqk, propagatedVNew, propagatedH, + stateScratch, userWorkspace, aqk, userWorkspace, userWorkspace, localScratch, userWorkspace, + userWorkspace, o, propagatedH, outputScratch, tiling, &pipe, false); + op.ProcessAic(); + } + if ASCEND_IS_AIV { + ChunkKdaFwdFinalizeKernel op; + op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, qgScaled, aqk, propagatedVNew, propagatedH, + stateScratch, userWorkspace, aqk, userWorkspace, userWorkspace, localScratch, userWorkspace, + userWorkspace, o, propagatedH, outputScratch, tiling, &pipe); + op.ProcessAiv(); + } +} + +} // namespace KdaFinalize diff --git a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_h_runtime_tiling.h b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_h_runtime_tiling.h new file mode 100644 index 0000000000..9ef1bd15fe --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_h_runtime_tiling.h @@ -0,0 +1,42 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CHUNK_KDA_FWD_H_RUNTIME_TILING_H +#define CHUNK_KDA_FWD_H_RUNTIME_TILING_H + +#include + +struct ChunkKdaFwdHRuntimeTiling { + int64_t batch; + int64_t seqlen; + int64_t kNumHead; + int64_t vNumHead; + int64_t kHeadDim; + int64_t vHeadDim; + int64_t chunkSize; + bool useInitialState; + bool storeFinalState; + int64_t dataType; + int64_t gDataType; + int64_t stateDataType; + int64_t isVariedLen; + int64_t shapeBatch; + int64_t tokenBatch; + bool useG; + bool useGk; + int64_t vWorkspaceOffset; + int64_t vUpdateWorkspaceOffset; + int64_t kDecayWorkspaceOffset; + int64_t hWorkspaceOffset; + int64_t numSeqWorkspaceOffset; + int64_t numChunksWorkspaceOffset; +}; + +#endif // CHUNK_KDA_FWD_H_RUNTIME_TILING_H diff --git a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_post_wu.h b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_post_wu.h new file mode 100644 index 0000000000..1146108de2 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_post_wu.h @@ -0,0 +1,924 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#pragma once + +#ifndef CATLASS_ARCH +#define CATLASS_ARCH 2201 +#endif + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/arch/kda_cross_core_sync.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/gemm/block/kda_block_mmad.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "kernel_utils/block/block_mmad_pingpong_tla_multi.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "kernel_operator.h" +#include "chunk_kda_fwd_varlen.h" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +using namespace AscendC; + +namespace KdaPostWu { +namespace { +using KdaInt64 = tla::Int<64>; +using KdaInt128 = tla::Int<128>; +constexpr float LN2 = 0.69314718055994530942f; +constexpr float KDA_EXP2_CLAMP = 80.0f; +constexpr float KDA_EXP_INPUT_MAX = KDA_EXP2_CLAMP * LN2; +constexpr float KDA_EXP_INPUT_MIN = -KDA_EXP2_CLAMP * LN2; +constexpr float KDA_FP16_MAX = 65504.0f; +constexpr uint32_t EXP2_UB_ELEMENTS = 256; +constexpr uint32_t EXP2_UB_BYTES = EXP2_UB_ELEMENTS * (sizeof(float) + sizeof(uint16_t)); +constexpr uint32_t EXP2_EVENT_ID = 0; +constexpr uint32_t KDA_SOLVE_BT = 64; +constexpr uint32_t KDA_SOLVE_MATRIX_ELEMENTS = KDA_SOLVE_BT * KDA_SOLVE_BT; +constexpr uint32_t KDA_SOLVE_SCRATCH_X = 0; +constexpr uint32_t KDA_SOLVE_SCRATCH_Y0 = 1; +constexpr uint32_t KDA_SOLVE_SCRATCH_TMP = 2; +constexpr uint32_t KDA_SOLVE_SCRATCH_Y1 = 3; +constexpr uint32_t KDA_SOLVE_SCRATCH_IDENTITY = 4; +constexpr uint32_t KDA_SOLVE_SCRATCH_SLOTS = 5; +constexpr uint32_t KDA_SOLVE_DIAG_BT = 16; +constexpr uint32_t KDA_SOLVE_DIAG_BLOCKS = KDA_SOLVE_BT / KDA_SOLVE_DIAG_BT; +constexpr uint32_t KDA_SOLVE_DIAG_MCH_ITERS = 3; +constexpr uint32_t KDA_SCORE_REF_BC = 16; +constexpr uint32_t KDA_VEC_ARENA_ELEMENTS = 32768; +constexpr uint32_t KDA_BITS_PER_MASK_BYTE = 8; +constexpr uint32_t KDA_SELECT_COL_BLOCKS = 2; +constexpr uint32_t KDA_SELECT_COL_MASK_BYTES = KDA_SOLVE_MATRIX_ELEMENTS / KDA_BITS_PER_MASK_BYTE; +constexpr uint32_t KDA_SELECT_MASK_BYTES = KDA_SELECT_COL_BLOCKS * KDA_SELECT_COL_MASK_BYTES; +constexpr uint32_t KDA_SELECT_AQK_MASK_BYTE_OFFSET = 120 * 1024; +constexpr uint32_t KDA_SELECT_AKK_MASK_BYTE_OFFSET = KDA_SELECT_AQK_MASK_BYTE_OFFSET + KDA_SELECT_MASK_BYTES; +constexpr uint32_t KDA_SELECT_ZERO_BYTE_OFFSET = KDA_SELECT_AKK_MASK_BYTE_OFFSET + KDA_SELECT_MASK_BYTES; +constexpr uint32_t KDA_SELECT_ZERO_FLOAT_OFFSET = KDA_SELECT_ZERO_BYTE_OFFSET / sizeof(float); +constexpr uint8_t KDA_SCORE_DONE_FLAG0 = 2; +constexpr uint8_t KDA_SCORE_DONE_FLAG1 = 3; +constexpr uint8_t KDA_SCORE_READY_FLAG0 = 4; +constexpr uint8_t KDA_SCORE_READY_FLAG1 = 5; +constexpr uint32_t KDA_SCORE_QUEUE_DEPTH = 2; +constexpr uint32_t KDA_SYNC_REVERSE_DEPTH = 1; +constexpr uint32_t KDA_SCORE_SCRATCH_PLANES = 3; +constexpr uint32_t KDA_SCORE_SCRATCH_QG = 0; +constexpr uint32_t KDA_SCORE_SCRATCH_W = 1; +constexpr uint32_t KDA_SCORE_SCRATCH_KG = 2; +constexpr uint64_t KDA_WORKSPACE_ALIGN = 512; +constexpr uint32_t KDA_GATE_TILE_ROWS = 32; +constexpr uint32_t KDA_POST_EVENT = 3; +constexpr uint32_t KDA_POST_EVENT_NEXT = 4; +constexpr uint32_t KDA_POST_EVENT_FIX = 5; + +using KdaArchTag = Catlass::Arch::AtlasA2; +using KdaDispatchPolicy = Catlass::Gemm::MmadPingpong; +using KdaScoreDispatchPolicy = Catlass::Gemm::MmadPingpongTlaMulti; +static_assert(KdaScoreDispatchPolicy::ENABLE_L1_RESIDENT, + "KDA Aqk/Akk score MMAD must keep the shared right matrix resident in L1"); +static_assert(KdaScoreDispatchPolicy::L1B_STAGES == 1, + "KDA Aqk/Akk score MMAD needs one L1 B slot so the second MMAD reuses it"); +using KdaSolveDispatchPolicy = Catlass::Gemm::MmadPingpong; +static_assert(!KdaSolveDispatchPolicy::USE_HF32_MODE, "KDA triangular solve must use IEEE FP32 Cube mode"); +using KdaL1TileShape = tla::Shape; +using KdaL0TileShape = KdaL1TileShape; +using KdaSolveL1TileShape = tla::Shape; +using KdaSolveL0TileShape = KdaSolveL1TileShape; + +__aicore__ inline uint32_t FloatToBits(float value) +{ + union Bits { + __aicore__ Bits() {} + float f; + uint32_t u; + } bits; + bits.f = value; + return bits.u; +} + +__aicore__ inline float BitsToFloat(uint32_t value) +{ + union Bits { + __aicore__ Bits() {} + uint32_t u; + float f; + } bits; + bits.u = value; + return bits.f; +} + +__aicore__ inline uint16_t Bf16ToBits(bfloat16_t value) +{ + union Bits { + __aicore__ Bits() {} + bfloat16_t f; + uint16_t u; + } bits; + bits.f = value; + return bits.u; +} + +__aicore__ inline bfloat16_t BitsToBf16(uint16_t value) +{ + union Bits { + __aicore__ Bits() {} + uint16_t u; + bfloat16_t f; + } bits; + bits.u = value; + return bits.f; +} + +template +__aicore__ inline T FloatToType(float value) +{ + if constexpr (IsSameType::value) { + uint32_t bits = FloatToBits(value); + uint32_t bias = 0x7FFFu + ((bits >> 16) & 1u); + return BitsToBf16(static_cast((bits + bias) >> 16)); + } + return static_cast(value); +} + +template +class ChunkKdaFwdPostWuKernel { +public: + using OUT_T = T; + using AKK_T = T; + template + __aicore__ inline void Init(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR beta, GM_ADDR initialState, + GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR preparedQG, GM_ADDR preparedAqk, + GM_ADDR propagatedVNew, GM_ADDR propagatedH, GM_ADDR o, GM_ADDR finalState, GM_ADDR aqk, + GM_ADDR akk, GM_ADDR w, GM_ADDR u, GM_ADDR qg, GM_ADDR kg, GM_ADDR vNew, GM_ADDR h, + GM_ADDR workspace, const TilingData &tiling, TPipe *pipe, bool initVecBuffers = true) + { + pipe_ = pipe; + q_.SetGlobalBuffer((__gm__ T *)q); + k_.SetGlobalBuffer((__gm__ T *)k); + v_.SetGlobalBuffer((__gm__ T *)v); + gk_.SetGlobalBuffer((__gm__ GK_T *)gk); + if (initialState != nullptr) { + initialState_.SetGlobalBuffer((__gm__ float *)initialState); + } + cuSeqlensAddr_ = reinterpret_cast<__gm__ int64_t *>(cuSeqlens); + if (preparedQG != nullptr) { + preparedQG_.SetGlobalBuffer((__gm__ T *)preparedQG); + } + if (preparedAqk != nullptr) { + preparedAqk_.SetGlobalBuffer((__gm__ T *)preparedAqk); + } + if (propagatedVNew != nullptr) { + propagatedVNew_.SetGlobalBuffer((__gm__ T *)propagatedVNew); + } + if (propagatedH != nullptr) { + propagatedH_.SetGlobalBuffer((__gm__ T *)propagatedH); + } + chunkIndicesAddr_ = reinterpret_cast<__gm__ int64_t *>(chunkIndices); + o_.SetGlobalBuffer((__gm__ OUT_T *)o); + finalState_.SetGlobalBuffer((__gm__ float *)finalState); + aqk_.SetGlobalBuffer((__gm__ float *)aqk); + akk_.SetGlobalBuffer((__gm__ AKK_T *)akk); + w_.SetGlobalBuffer((__gm__ T *)w); + u_.SetGlobalBuffer((__gm__ OUT_T *)u); + qg_.SetGlobalBuffer((__gm__ T *)qg); + kg_.SetGlobalBuffer((__gm__ T *)kg); + vNew_.SetGlobalBuffer((__gm__ T *)vNew); + h_.SetGlobalBuffer((__gm__ float *)h); + solveWorkspace_.SetGlobalBuffer((__gm__ float *)workspace); + + B_ = tiling.batch; + N_ = tiling.seqNum; + H_ = tiling.qHeadNum; + HV_ = tiling.vHeadNum; + T_ = tiling.seqlen; + K_ = tiling.kHeadDim; + V_ = tiling.vHeadDim; + BT_ = tiling.chunkSize; + NT_ = tiling.totalChunks; + scale_ = tiling.scale; + hasInitial_ = tiling.hasInitialState; + isVarLen_ = tiling.isVarLen; + inputSequenceMajor_ = tiling.inputSequenceMajor; + usedCoreNum_ = tiling.postWuUsedCoreNum; + if ASCEND_IS_AIV { + uint64_t subBlockNum = static_cast(GetSubBlockNum()); + solveCoreIdx_ = subBlockNum == 0 ? 0 : static_cast(GetBlockIdx()) / subBlockNum; + } else { + solveCoreIdx_ = static_cast(GetBlockIdx()); + } + if (pipe_ != nullptr && initVecBuffers) { + pipe_->InitBuffer(exp2Buf_, EXP2_UB_BYTES); + pipe_->InitBuffer(vecBuf_, KDA_VEC_ARENA_ELEMENTS * sizeof(float)); + const uint64_t gateWritebackRows = ScoreVectorMaxRows(5 * sizeof(float) + 2 * sizeof(T) + sizeof(GK_T)); + pipe_->InitBuffer(gateWritebackBuf_, + static_cast(gateWritebackRows * K_ * (3 * sizeof(T) + sizeof(GK_T)))); + AllocVectorEvents(); + } + } + __aicore__ inline void ProcessAiv() + { + ProcessPostAiv(); + ReleaseVectorEvents(); + } + + __aicore__ inline void ProcessAic() { ProcessPostAic(); } + +private: + __aicore__ inline void AllocVectorEvents() + { + mte2ToVEvent_ = pipe_->AllocEventID(); + vToMte2Event_ = pipe_->AllocEventID(); + vToMte3Event_ = pipe_->AllocEventID(); + mte3ToVEvent_ = pipe_->AllocEventID(); + mte2ToMte3Event_ = pipe_->AllocEventID(); + mte3ToMte2Event_ = pipe_->AllocEventID(); + vectorEventsAllocated_ = true; + } + + __aicore__ inline void ReleaseVectorEvents() + { + if (!vectorEventsAllocated_) { + return; + } + pipe_->ReleaseEventID(mte2ToVEvent_); + pipe_->ReleaseEventID(vToMte2Event_); + pipe_->ReleaseEventID(vToMte3Event_); + pipe_->ReleaseEventID(mte3ToVEvent_); + pipe_->ReleaseEventID(mte2ToMte3Event_); + pipe_->ReleaseEventID(mte3ToMte2Event_); + vectorEventsAllocated_ = false; + } + + __aicore__ inline uint64_t QOffset(uint64_t b, uint64_t h, uint64_t t, uint64_t d) const + { + if (inputSequenceMajor_) { + return ((b * T_ + t) * H_ + h) * K_ + d; + } + return ((b * H_ + h) * T_ + t) * K_ + d; + } + + __aicore__ inline uint64_t KVOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t d, uint64_t dim) const + { + return ((b * HV_ + hv) * T_ + t) * dim + d; + } + + __aicore__ inline uint64_t BetaOffset(uint64_t b, uint64_t hv, uint64_t t) const { return (b * HV_ + hv) * T_ + t; } + + __aicore__ inline uint64_t AOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t j) const + { + return ((b * HV_ + hv) * T_ + t) * BT_ + j; + } + + __aicore__ inline uint64_t HOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t d, uint64_t r) const + { + return (((b * HV_ + hv) * NT_ + chunkIdx) * K_ + d) * V_ + r; + } + + __aicore__ inline uint64_t WScratchOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t t, uint64_t d) const + { + return (((b * HV_ + hv) * NT_ + chunkIdx) * BT_ + t) * K_ + d; + } + + __aicore__ inline uint64_t SolveScratchOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t slot) const + { + (void)b; + (void)hv; + (void)chunkIdx; + uint64_t matrixElements = BT_ * BT_; + return solveCoreIdx_ * KDA_SOLVE_SCRATCH_SLOTS * matrixElements + slot * matrixElements; + } + + __aicore__ inline uint64_t ScoreScratchOffset(uint64_t slot, uint64_t plane, uint64_t t = 0, uint64_t d = 0) const + { + return (((solveCoreIdx_ * KDA_SCORE_QUEUE_DEPTH + slot) * KDA_SCORE_SCRATCH_PLANES + plane) * BT_ + t) * K_ + d; + } + + __aicore__ inline uint64_t ScoreRefBlockSize() const { return KDA_SCORE_REF_BC; } + + __aicore__ inline uint64_t ScoreRowBlockCount(uint64_t curT, uint64_t rowBegin) const + { + uint64_t blockSize = ScoreRefBlockSize(); + uint64_t rowCount = curT - rowBegin; + if (rowCount > blockSize) { + rowCount = blockSize; + } + return rowCount; + } + + __aicore__ inline uint64_t ScoreRefToken(uint64_t start, uint64_t curT, uint64_t rowBegin, uint64_t rowCount) const + { + uint64_t ref = rowBegin + rowCount / 2; + if (ref >= curT) { + ref = curT - 1; + } + return start + ref; + } + + __aicore__ inline void RunExp2(LocalTensor &tensor, uint32_t count) + { + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + ClampExpInput(tensor, count); + Exp(tensor, tensor, count); + PipeBarrier(); + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + } + + __aicore__ inline void ClampExpInput(LocalTensor &tensor, uint32_t count) + { + Mins(tensor, tensor, KDA_EXP_INPUT_MAX, count); + PipeBarrier(); + Maxs(tensor, tensor, KDA_EXP_INPUT_MIN, count); + PipeBarrier(); + } + + __aicore__ inline void ClampFp32ToOutputType(LocalTensor &tensor, uint32_t count) + { + if constexpr (IsSameType::value) { + Mins(tensor, tensor, KDA_FP16_MAX, count); + PipeBarrier(); + Maxs(tensor, tensor, -KDA_FP16_MAX, count); + PipeBarrier(); + } + } + + template + __aicore__ inline void CopyVectorIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset, + uint64_t count) + { + uint64_t rowBytes = count * static_cast(sizeof(CopyT)); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst, src[offset], static_cast(count)); + return; + } + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPadParams padParams{false, 0, 0, 0}; + DataCopyPad(dst, src[offset], params, padParams); + } + + template + __aicore__ inline void CopyVectorOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src, + uint64_t count) + { + uint64_t rowBytes = count * static_cast(sizeof(CopyT)); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst[offset], src, static_cast(count)); + return; + } + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPad(dst[offset], src, params); + } + + template + __aicore__ inline void CopyRowsIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset, uint64_t rows, + uint64_t cols, uint64_t rowStride) + { + if (rows == 0 || cols == 0) { + return; + } + if (rowStride == cols) { + CopyVectorIn(dst, src, offset, rows * cols); + return; + } + DataCopyExtParams params{static_cast(rows), static_cast(cols * sizeof(CopyT)), + static_cast((rowStride - cols) * sizeof(CopyT)), 0, 0}; + DataCopyPadExtParams padParams{false, 0, 0, 0}; + DataCopyPad(dst, src[offset], params, padParams); + } + + template + __aicore__ inline void CopyRowIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset) + { + CopyVectorIn(dst, src, offset, K_); + } + + template + __aicore__ inline void CopyRowOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src) + { + CopyVectorOut(dst, offset, src, K_); + } + + __aicore__ inline LocalTensor VecScratch(uint64_t slot) + { + return vecBuf_.Get()[slot * EXP2_UB_ELEMENTS]; + } + + template + __aicore__ inline void LoadAsFloatRow(GlobalTensor &src, uint64_t srcOffset, LocalTensor &dst, + uint64_t count) + { + if constexpr (IsSameType::value) { + CopyVectorIn(dst, src, srcOffset, count); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Adds(dst, dst, 0.0f, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } else { + constexpr uint32_t typedOffset = EXP2_UB_ELEMENTS * sizeof(float) / sizeof(CopyT); + LocalTensor rowLocal = exp2Buf_.Get()[typedOffset]; + CopyVectorIn(rowLocal, src, srcOffset, count); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Cast(dst, rowLocal, RoundMode::CAST_NONE, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } + PipeBarrier(); + } + + template + __aicore__ inline void LoadAsFloatVector(GlobalTensor &src, uint64_t srcOffset, LocalTensor &dst, + LocalTensor &typedScratch, uint64_t count) + { + if constexpr (IsSameType::value) { + CopyVectorIn(dst, src, srcOffset, count); + } else { + CopyVectorIn(typedScratch, src, srcOffset, count); + } + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + if constexpr (!IsSameType::value) { + Cast(dst, typedScratch, RoundMode::CAST_NONE, static_cast(count)); + PipeBarrier(); + } + } + + template + __aicore__ inline void StoreFloatRow(GlobalTensor &dst, uint64_t dstOffset, LocalTensor &src, + uint64_t count) + { + if constexpr (IsSameType::value) { + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(dst, dstOffset, src, count); + } else { + constexpr uint32_t typedOffset = EXP2_UB_ELEMENTS * sizeof(float) / sizeof(CopyT); + LocalTensor rowLocal = exp2Buf_.Get()[typedOffset]; + Cast(rowLocal, src, RoundMode::CAST_RINT, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(dst, dstOffset, rowLocal, count); + } + SetFlag(mte3ToMte2Event_); + WaitFlag(mte3ToMte2Event_); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline LocalTensor Exp2NegG(uint64_t b, uint64_t hv, uint64_t t) + { + LocalTensor exp2Local = exp2Buf_.Get(); + LoadAsFloatRow(gk_, KVOffset(b, hv, t, 0, K_), exp2Local, K_); + Muls(exp2Local, exp2Local, -LN2, static_cast(K_)); + PipeBarrier(); + RunExp2(exp2Local, static_cast(K_)); + return exp2Local; + } + + __aicore__ inline uint64_t ScoreVectorMaxRows(uint64_t bytesPerElem) const + { + constexpr uint64_t arenaBytes = static_cast(KDA_VEC_ARENA_ELEMENTS) * sizeof(float); + uint64_t maxRows = (arenaBytes / bytesPerElem) / K_; + if (K_ >= 128 && maxRows > 32) { + maxRows = 32; + } + return maxRows; + } + __aicore__ inline bool UsePostWuCube(uint64_t curT) const + { + return curT > 0 && curT <= BT_ && (BT_ == 64 || BT_ == 128) && K_ >= 16 && V_ >= 16 && V_ <= 256 && + K_ % 16 == 0 && V_ % 16 == 0; + } + + __aicore__ inline void ComputePostWuCube(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t curT) + { + using ElementA = AKK_T; + using ElementB = T; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::RowMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using WTileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using UTileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using PostL1TileShape128 = tla::Shape; + using PostL0TileShape128 = tla::Shape; + using PostL1TileShape256 = tla::Shape; + using PostL0TileShape256 = tla::Shape; + using WBlockMmad = Catlass::Gemm::Block::BlockMmadTla; + using UBlockMmad128 = + Catlass::Gemm::Block::BlockMmadTla; + using UBlockMmad256 = + Catlass::Gemm::Block::BlockMmadTla; + LayoutTagA tagA = LayoutTagA::template MakeLayout(BT_, BT_); + auto layoutA = tla::MakeLayoutFromTag(tagA); + auto tensorA = tla::MakeTensor(preparedAqk_[AOffset(b, hv, start, 0)], layoutA, Catlass::Arch::PositionGM{}); + + { + LayoutTagB tagB = LayoutTagB::template MakeLayout(BT_, K_); + LayoutTagC tagC = LayoutTagC::template MakeLayout(BT_, K_); + auto layoutB = tla::MakeLayoutFromTag(tagB); + auto layoutC = tla::MakeLayoutFromTag(tagC); + Catlass::GemmCoord shape{static_cast(curT), static_cast(K_), + static_cast(curT)}; + auto tensorB = + tla::MakeTensor(preparedQG_[KVOffset(b, hv, start, 0, K_)], layoutB, Catlass::Arch::PositionGM{}); + auto tensorC = + tla::MakeTensor(h_[WScratchOffset(b, hv, chunkIdx, 0, 0)], layoutC, Catlass::Arch::PositionGM{}); + auto blockA = GetTile(tensorA, tla::MakeCoord(0, 0), tla::MakeShape(shape.m(), shape.k())); + auto blockB = GetTile(tensorB, tla::MakeCoord(0, 0), tla::MakeShape(shape.k(), shape.n())); + auto blockC = GetTile(tensorC, tla::MakeCoord(0, 0), tla::MakeShape(shape.m(), shape.n())); + Catlass::Arch::Resource wResource; + WBlockMmad wBlockMmad(wResource); + wBlockMmad(blockA, blockB, blockC, shape); + PipeBarrier(); + } + + { + LayoutTagB tagB = LayoutTagB::template MakeLayout(BT_, V_); + LayoutTagC tagC = LayoutTagC::template MakeLayout(BT_, V_); + auto layoutB = tla::MakeLayoutFromTag(tagB); + auto layoutC = tla::MakeLayoutFromTag(tagC); + Catlass::GemmCoord shape{static_cast(curT), static_cast(V_), + static_cast(curT)}; + auto tensorB = + tla::MakeTensor(propagatedVNew_[KVOffset(b, hv, start, 0, V_)], layoutB, Catlass::Arch::PositionGM{}); + auto tensorC = tla::MakeTensor(u_[KVOffset(b, hv, start, 0, V_)], layoutC, Catlass::Arch::PositionGM{}); + auto blockA = GetTile(tensorA, tla::MakeCoord(0, 0), tla::MakeShape(shape.m(), shape.k())); + auto blockB = GetTile(tensorB, tla::MakeCoord(0, 0), tla::MakeShape(shape.k(), shape.n())); + auto blockC = GetTile(tensorC, tla::MakeCoord(0, 0), tla::MakeShape(shape.m(), shape.n())); + Catlass::Arch::Resource uResource; + if (V_ <= 128) { + UBlockMmad128 uBlockMmad(uResource); + uBlockMmad(blockA, blockB, blockC, shape); + } else { + UBlockMmad256 uBlockMmad(uResource); + uBlockMmad(blockA, blockB, blockC, shape); + } + PipeBarrier(); + } + } + + __aicore__ inline void CopyScratchWAndFinalizeKg(uint64_t b, uint64_t h, uint64_t hv, uint64_t chunkIdx, + uint64_t start, uint64_t curT, uint64_t subBlockIdx, + uint64_t subBlockNum, bool copyScratchW) + { + constexpr uint64_t typedOffsetFloats = 20480; + constexpr uint64_t typedOffset = typedOffsetFloats * sizeof(float) / sizeof(T); + constexpr uint64_t kgFp32Planes = 4; + uint64_t rowBegin = (curT * subBlockIdx) / subBlockNum; + uint64_t rowEnd = (curT * (subBlockIdx + 1)) / subBlockNum; + if (rowBegin >= rowEnd) { + return; + } + uint64_t maxRows = (typedOffsetFloats / kgFp32Planes) / K_; + if (maxRows > 32) { + maxRows = 32; + } + if (maxRows == 0) { + return; + } + + uint64_t last = start + curT - 1; + LocalTensor arena = vecBuf_.Get(); + LocalTensor gateLast = exp2Buf_.Get(); + LocalTensor typedLocal = vecBuf_.Get()[typedOffset]; + LoadAsFloatRow(gk_, KVOffset(b, hv, last, 0, K_), gateLast, K_); + + for (uint64_t tileRow = rowBegin; tileRow < rowEnd; tileRow += maxRows) { + uint64_t tileRows = rowEnd - tileRow; + if (tileRows > maxRows) { + tileRows = maxRows; + } + uint64_t elemCount = tileRows * K_; + uint64_t token = start + tileRow; + + if (copyScratchW) { + uint64_t scratchBase = WScratchOffset(b, hv, chunkIdx, tileRow, 0); + DataCopy(arena, h_[scratchBase], static_cast(elemCount)); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Cast(typedLocal, arena, RoundMode::CAST_RINT, static_cast(elemCount)); + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(w_[KVOffset(b, hv, token, 0, K_)], typedLocal, static_cast(elemCount)); + SetFlag(mte3ToMte2Event_); + WaitFlag(mte3ToMte2Event_); + } + + LocalTensor kLocal = arena; + LocalTensor gLocal = arena[elemCount]; + LocalTensor expLocal = arena[2 * elemCount]; + LocalTensor outLocal = arena[3 * elemCount]; + const uint64_t gateOffsetBytes = (typedOffset + elemCount) * sizeof(T); + LocalTensor gateTyped = vecBuf_.Get()[(gateOffsetBytes + sizeof(GK_T) - 1) / sizeof(GK_T)]; + CopyRowsIn(typedLocal, k_, QOffset(b, h, token, 0), tileRows, K_, inputSequenceMajor_ ? H_ * K_ : K_); + LoadAsFloatVector(gk_, KVOffset(b, hv, token, 0, K_), gLocal, gateTyped, elemCount); + Cast(kLocal, typedLocal, RoundMode::CAST_NONE, static_cast(elemCount)); + PipeBarrier(); + + for (uint64_t row = 0; row < tileRows; ++row) { + Sub(expLocal[row * K_], gateLast, gLocal[row * K_], static_cast(K_)); + } + PipeBarrier(); + Muls(expLocal, expLocal, LN2, static_cast(elemCount)); + PipeBarrier(); + ClampExpInput(expLocal, static_cast(elemCount)); + Exp(expLocal, expLocal, static_cast(elemCount)); + PipeBarrier(); + Mul(outLocal, kLocal, expLocal, static_cast(elemCount)); + PipeBarrier(); + ClampFp32ToOutputType(outLocal, static_cast(elemCount)); + Cast(typedLocal, outLocal, RoundMode::CAST_RINT, static_cast(elemCount)); + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(kg_, KVOffset(b, hv, token, 0, K_), typedLocal, elemCount); + SetFlag(mte3ToMte2Event_); + WaitFlag(mte3ToMte2Event_); + } + if (rowEnd == curT) { + CopyVectorIn(typedLocal, k_, QOffset(b, h, last, 0), K_); + SetFlag(mte2ToMte3Event_); + WaitFlag(mte2ToMte3Event_); + CopyVectorOut(kg_, KVOffset(b, hv, last, 0, K_), typedLocal, K_); + SetFlag(mte3ToMte2Event_); + WaitFlag(mte3ToMte2Event_); + } + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + template + __aicore__ inline void ComputeTailWuRow(GlobalTensor &src, GlobalTensor &dst, + uint64_t akkBase, uint64_t srcBase, uint64_t dstBase, uint64_t curT, + uint64_t dim) + { + LocalTensor acc = vecBuf_.Get(); + LocalTensor value = vecBuf_.Get()[512]; + LocalTensor typed = vecBuf_.Get()[4096]; + LocalTensor coefficientTyped = exp2Buf_.Get(); + LocalTensor coefficients = exp2Buf_.Get()[128]; + CopyVectorIn(coefficientTyped, preparedAqk_, akkBase, curT); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Cast(coefficients, coefficientTyped, RoundMode::CAST_NONE, static_cast(curT)); + PipeBarrier(); + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + for (uint64_t j = 0; j < curT; ++j) { + LoadAsFloatVector(src, srcBase + j * dim, value, typed, dim); + float coefficient = coefficients.GetValue(j); + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + Muls(value, value, coefficient, static_cast(dim)); + PipeBarrier(); + if (j == 0) { + Adds(acc, value, 0.0f, static_cast(dim)); + } else { + Add(acc, acc, value, static_cast(dim)); + } + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + ClampFp32ToOutputType(acc, static_cast(dim)); + StoreFloatRow(dst, dstBase, acc, dim); + } + + __aicore__ inline void ComputeTailWuVector(uint64_t b, uint64_t hv, uint64_t start, uint64_t curT, + uint64_t subBlockIdx, uint64_t subBlockNum) + { + uint64_t rowBegin = (curT * subBlockIdx) / subBlockNum; + uint64_t rowEnd = (curT * (subBlockIdx + 1)) / subBlockNum; + for (uint64_t row = rowBegin; row < rowEnd; ++row) { + ComputeTailWuRow(preparedQG_, w_, AOffset(b, hv, start + row, 0), KVOffset(b, hv, start, 0, K_), + KVOffset(b, hv, start + row, 0, K_), curT, K_); + ComputeTailWuRow(propagatedVNew_, u_, AOffset(b, hv, start + row, 0), KVOffset(b, hv, start, 0, V_), + KVOffset(b, hv, start + row, 0, V_), curT, V_); + } + } + + __aicore__ inline bool ResolveFlatChunk(uint64_t task, uint64_t &seq, uint64_t &b, uint64_t &h, uint64_t &hv, + uint64_t &chunkIdx, uint64_t &start, uint64_t &end) + { + hv = task % HV_; + uint64_t flatChunk = task / HV_; + if (!isVarLen_) { + seq = flatChunk / NT_; + b = seq; + chunkIdx = flatChunk % NT_; + start = chunkIdx * BT_; + end = start + BT_; + if (end > T_) { + end = T_; + } + } else { + if (!KdaVarlen::ResolveChunkRange(cuSeqlensAddr_, chunkIndicesAddr_, N_, T_, BT_, flatChunk, seq, start, + end)) { + return false; + } + b = 0; + chunkIdx = flatChunk; + } + h = hv / (HV_ / H_); + return start < end; + } + + __aicore__ inline void ProcessChunkPostAiv(uint64_t b, uint64_t h, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t end, uint64_t subBlockIdx, uint64_t subBlockNum) + { + uint64_t curT = end - start; + if (curT == 0 || !UsePostWuCube(curT)) { + return; + } + if (curT < BT_) { + ComputeTailWuVector(b, hv, start, curT, subBlockIdx, subBlockNum); + CopyScratchWAndFinalizeKg(b, h, hv, chunkIdx, start, curT, subBlockIdx, subBlockNum, false); + return; + } + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE2>(syncDoneFlag_); + CopyScratchWAndFinalizeKg(b, h, hv, chunkIdx, start, curT, subBlockIdx, subBlockNum, true); + } + + __aicore__ inline void ProcessChunkPostAic(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t end) + { + if constexpr (IsSameType::value) { + ProcessChunkPostAicTyped(b, hv, chunkIdx, start, end); + } + } + + __aicore__ inline void ProcessChunkPostAicTyped(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t end) + { + uint64_t curT = end - start; + if (curT == 0 || !UsePostWuCube(curT)) { + return; + } + if (curT < BT_) { + return; + } + ComputePostWuCube(b, hv, chunkIdx, start, curT); + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(syncDoneFlag_); + } + + __aicore__ inline void ProcessPostAiv() + { + if constexpr (IsSameType::value) { + return; + } + uint64_t subBlockNum = static_cast(GetSubBlockNum()); + if (subBlockNum == 0) { + return; + } + uint64_t subBlockIdx = static_cast(GetSubBlockIdx()); + uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + uint64_t coreIdx = static_cast(GetBlockIdx()) / subBlockNum; + uint64_t taskNum = static_cast((isVarLen_ ? NT_ : B_ * NT_) * HV_); + for (uint64_t task = coreIdx; task < taskNum; task += coreNum) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveFlatChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + (void)seq; + ProcessChunkPostAiv(b, h, hv, chunkIdx, start, end, subBlockIdx, subBlockNum); + } + } + } + + __aicore__ inline void ProcessPostAic() + { + if constexpr (IsSameType::value) { + return; + } + uint64_t taskNum = static_cast((isVarLen_ ? NT_ : B_ * NT_) * HV_); + uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + for (uint64_t task = GetBlockIdx(); task < taskNum; task += coreNum) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveFlatChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + (void)seq; + (void)h; + ProcessChunkPostAic(b, hv, chunkIdx, start, end); + } + } + } + +private: + GlobalTensor q_; + GlobalTensor k_; + GlobalTensor v_; + GlobalTensor gk_; + GlobalTensor initialState_; + GlobalTensor o_; + GlobalTensor finalState_; + GlobalTensor aqk_; + GlobalTensor akk_; + GlobalTensor w_; + GlobalTensor u_; + GlobalTensor qg_; + GlobalTensor kg_; + GlobalTensor vNew_; + GlobalTensor h_; + GlobalTensor preparedQG_; + GlobalTensor preparedAqk_; + GlobalTensor propagatedVNew_; + GlobalTensor propagatedH_; + GlobalTensor solveWorkspace_; + GlobalTensor scoreWorkspace_; + TPipe *pipe_ = nullptr; + TBuf exp2Buf_; + TBuf vecBuf_; + TBuf gateWritebackBuf_; + TEventID mte2ToVEvent_ = 0; + TEventID vToMte2Event_ = 0; + TEventID vToMte3Event_ = 0; + TEventID mte3ToVEvent_ = 0; + TEventID mte2ToMte3Event_ = 0; + TEventID mte3ToMte2Event_ = 0; + bool vectorEventsAllocated_ = false; + Catlass::Arch::CrossCoreFlagWithReverse scoreReadyFlag_{KDA_SCORE_READY_FLAG0, + KDA_SCORE_READY_FLAG1}; + Catlass::Arch::CrossCoreFlagWithReverse scoreDoneFlag_{KDA_SCORE_DONE_FLAG0, + KDA_SCORE_DONE_FLAG1}; + // Score production is fully drained before solve starts, so the solve handshake can safely reuse + // the existing score flags without consuming additional hardware flag IDs. + Catlass::Arch::CrossCoreFlagWithReverse syncReadyFlag_{KDA_SCORE_READY_FLAG0, + KDA_SCORE_READY_FLAG1}; + Catlass::Arch::CrossCoreFlagWithReverse syncDoneFlag_{KDA_SCORE_DONE_FLAG0, + KDA_SCORE_DONE_FLAG1}; + uint64_t B_ = 0; + uint64_t N_ = 0; + uint64_t H_ = 0; + uint64_t HV_ = 0; + uint64_t T_ = 0; + uint64_t K_ = 0; + uint64_t V_ = 0; + uint64_t BT_ = 0; + uint64_t NT_ = 0; + float scale_ = 1.0f; + bool hasInitial_ = false; + bool isVarLen_ = false; + bool isAivOnly_ = false; + bool inputSequenceMajor_ = false; + uint64_t usedCoreNum_ = 1; + uint64_t solveCoreIdx_ = 0; + __gm__ int64_t *chunkIndicesAddr_ = nullptr; + __gm__ int64_t *cuSeqlensAddr_ = nullptr; +}; +} // namespace + +template +__aicore__ inline void RunChunkKdaPostWu(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR beta, + GM_ADDR initialState, GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR wSeed, + GM_ADDR akk, GM_ADDR uSeed, GM_ADDR w, GM_ADDR u, GM_ADDR kg, GM_ADDR vNew, + GM_ADDR userWorkspace, const TilingData &tiling, TPipe &pipe) +{ + GM_ADDR postScratch = userWorkspace + tiling.postWuScratchOffset; + if ASCEND_IS_AIC { + ChunkKdaFwdPostWuKernel op; + op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, wSeed, akk, uSeed, nullptr, userWorkspace, + userWorkspace, userWorkspace, akk, w, u, userWorkspace, kg, vNew, postScratch, postScratch, tiling, + &pipe, false); + op.ProcessAic(); + } + if ASCEND_IS_AIV { + ChunkKdaFwdPostWuKernel op; + op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, wSeed, akk, uSeed, nullptr, userWorkspace, + userWorkspace, userWorkspace, akk, w, u, userWorkspace, kg, vNew, postScratch, postScratch, tiling, + &pipe); + op.ProcessAiv(); + } +} + +} // namespace KdaPostWu diff --git a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_prepare.h b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_prepare.h new file mode 100644 index 0000000000..ee95145a5a --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_prepare.h @@ -0,0 +1,2535 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#pragma once + +#ifndef CATLASS_ARCH +#define CATLASS_ARCH 2201 +#endif + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/arch/kda_cross_core_sync.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_catlass.hpp" +#include "catlass/gemm/block/kda_block_mmad.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "kernel_utils/block/block_mmad_pingpong_tla_multi.hpp" +#include "catlass/layout/kda_layout.hpp" +#include "kernel_operator.h" +#include "chunk_kda_fwd_varlen.h" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +using namespace AscendC; + +namespace KdaPrepare { +namespace { +using KdaInt64 = tla::Int<64>; +using KdaInt128 = tla::Int<128>; +constexpr float LN2 = 0.69314718055994530942f; +constexpr float KDA_EXP2_CLAMP = 80.0f; +constexpr float KDA_EXP_INPUT_MAX = KDA_EXP2_CLAMP * LN2; +constexpr float KDA_EXP_INPUT_MIN = -KDA_EXP2_CLAMP * LN2; +constexpr float KDA_SCORE_EXP2_CLAMP = 120.0f; +constexpr float KDA_SCORE_EXP2_MIN_CLAMP = 126.0f; +constexpr float KDA_SCORE_EXP_INPUT_MAX = KDA_SCORE_EXP2_CLAMP * LN2; +constexpr float KDA_SCORE_EXP_INPUT_MIN = -KDA_SCORE_EXP2_MIN_CLAMP * LN2; +constexpr float KDA_FP16_MAX = 65504.0f; +constexpr uint32_t EXP2_UB_ELEMENTS = 256; +constexpr uint32_t EXP2_UB_BYTES = EXP2_UB_ELEMENTS * (sizeof(float) + sizeof(uint16_t)); +constexpr uint32_t EXP2_EVENT_ID = 0; +constexpr uint32_t KDA_SOLVE_BT = 64; +constexpr uint32_t KDA_SOLVE_MATRIX_ELEMENTS = KDA_SOLVE_BT * KDA_SOLVE_BT; +constexpr uint32_t KDA_SOLVE_SCRATCH_X = 0; +constexpr uint32_t KDA_SOLVE_SCRATCH_Y0 = 1; +constexpr uint32_t KDA_SOLVE_SCRATCH_TMP = 2; +constexpr uint32_t KDA_SOLVE_SCRATCH_Y1 = 3; +constexpr uint32_t KDA_SOLVE_SCRATCH_IDENTITY = 4; +constexpr uint32_t KDA_SOLVE_SCRATCH_SLOTS = 5; +constexpr uint32_t KDA_SOLVE_PIPELINE_DEPTH = 4; +constexpr uint32_t KDA_SOLVE_DIAG_BT = 16; +constexpr uint32_t KDA_SOLVE_DIAG_BLOCKS = KDA_SOLVE_BT / KDA_SOLVE_DIAG_BT; +constexpr uint32_t KDA_SOLVE_DIAG_MCH_ITERS = 3; +// Keep the local safe-gate exponent span within the BF16 score range while +// reducing repeated gate-factor work and AIV/AIC handshakes. +constexpr uint32_t KDA_SCORE_REF_BC = 32; +constexpr uint32_t KDA_SAFE_SCORE_REF_BC = 32; +constexpr uint32_t KDA_VEC_ARENA_ELEMENTS = 32768; +constexpr uint32_t KDA_BITS_PER_MASK_BYTE = 8; +constexpr uint32_t KDA_SELECT_COL_BLOCKS = 2; +constexpr uint32_t KDA_SELECT_COL_MASK_BYTES = KDA_SOLVE_MATRIX_ELEMENTS / KDA_BITS_PER_MASK_BYTE; +constexpr uint32_t KDA_SELECT_MASK_BYTES = KDA_SELECT_COL_BLOCKS * KDA_SELECT_COL_MASK_BYTES; +constexpr uint32_t KDA_SELECT_AQK_MASK_BYTE_OFFSET = 120 * 1024; +constexpr uint32_t KDA_SELECT_AKK_MASK_BYTE_OFFSET = KDA_SELECT_AQK_MASK_BYTE_OFFSET + KDA_SELECT_MASK_BYTES; +constexpr uint32_t KDA_SELECT_ZERO_BYTE_OFFSET = KDA_SELECT_AKK_MASK_BYTE_OFFSET + KDA_SELECT_MASK_BYTES; +constexpr uint32_t KDA_SELECT_ZERO_FLOAT_OFFSET = KDA_SELECT_ZERO_BYTE_OFFSET / sizeof(float); +constexpr uint8_t KDA_SCORE_DONE_FLAG0 = 2; +constexpr uint8_t KDA_SCORE_DONE_FLAG1 = 3; +constexpr uint8_t KDA_SCORE_READY_FLAG0 = 4; +constexpr uint8_t KDA_SCORE_READY_FLAG1 = 5; +constexpr uint8_t KDA_SOLVE_DONE_FLAG = 6; +constexpr uint8_t KDA_SOLVE_READY_FLAG = 7; +constexpr uint32_t KDA_SCORE_QUEUE_DEPTH = 2; +constexpr uint32_t KDA_SCORE_LANES = 2; +constexpr uint32_t KDA_SCORE_SCRATCH_SLOTS = KDA_SCORE_QUEUE_DEPTH * KDA_SCORE_LANES; +constexpr uint32_t KDA_SYNC_REVERSE_DEPTH = 1; +constexpr uint32_t KDA_SCORE_SCRATCH_PLANES = 3; +constexpr uint32_t KDA_SCORE_SCRATCH_QG = 0; +constexpr uint32_t KDA_SCORE_SCRATCH_W = 1; +constexpr uint32_t KDA_SCORE_SCRATCH_KG = 2; +constexpr uint64_t KDA_WORKSPACE_ALIGN = 512; +constexpr uint32_t KDA_GATE_TILE_ROWS = 16; +constexpr uint32_t KDA_GATE_PIPELINE_DEPTH = 3; +constexpr uint32_t KDA_AIV_UB_BUDGET_BYTES = 192 * 1024; +using KdaArchTag = Catlass::Arch::AtlasA2; +using KdaDispatchPolicy = Catlass::Gemm::MmadPingpong; +using KdaScoreDispatchPolicy = Catlass::Gemm::MmadPingpongTlaMulti; +static_assert(KdaScoreDispatchPolicy::ENABLE_L1_RESIDENT, + "KDA Aqk/Akk score MMAD must keep the shared right matrix resident in L1"); +static_assert(KdaScoreDispatchPolicy::L1B_STAGES == 1, + "KDA Aqk/Akk score MMAD needs one L1 B slot so the second MMAD reuses it"); +using KdaSolveDispatchPolicy = Catlass::Gemm::MmadPingpong; +static_assert(!KdaSolveDispatchPolicy::USE_HF32_MODE, "KDA triangular solve must use IEEE FP32 Cube mode"); +using KdaL1TileShape = tla::Shape; +using KdaL0TileShape = KdaL1TileShape; +using KdaSolveL1TileShape = tla::Shape; +using KdaSolveL0TileShape = KdaSolveL1TileShape; + +__aicore__ inline uint32_t FloatToBits(float value) +{ + union Bits { + __aicore__ Bits() {} + float f; + uint32_t u; + } bits; + bits.f = value; + return bits.u; +} + +__aicore__ inline float BitsToFloat(uint32_t value) +{ + union Bits { + __aicore__ Bits() {} + uint32_t u; + float f; + } bits; + bits.u = value; + return bits.f; +} + +__aicore__ inline uint16_t Bf16ToBits(bfloat16_t value) +{ + union Bits { + __aicore__ Bits() {} + bfloat16_t f; + uint16_t u; + } bits; + bits.f = value; + return bits.u; +} + +__aicore__ inline bfloat16_t BitsToBf16(uint16_t value) +{ + union Bits { + __aicore__ Bits() {} + uint16_t u; + bfloat16_t f; + } bits; + bits.u = value; + return bits.f; +} + +template +__aicore__ inline T FloatToType(float value) +{ + if constexpr (IsSameType::value) { + uint32_t bits = FloatToBits(value); + uint32_t bias = 0x7FFFu + ((bits >> 16) & 1u); + return BitsToBf16(static_cast((bits + bias) >> 16)); + } + return static_cast(value); +} + +template +class ChunkKdaFwdPrepareKernel { +public: + using OUT_T = T; + using AKK_T = float; + using SCORE_T = std::conditional_t::value, bfloat16_t, T>; + template + __aicore__ inline void Init(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR beta, GM_ADDR initialState, + GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR preparedQG, GM_ADDR preparedAqk, + GM_ADDR propagatedVNew, GM_ADDR propagatedH, GM_ADDR o, GM_ADDR finalState, GM_ADDR aqk, + GM_ADDR akk, GM_ADDR w, GM_ADDR u, GM_ADDR qg, GM_ADDR kg, GM_ADDR vNew, GM_ADDR h, + GM_ADDR workspace, const TilingData &tiling, TPipe *pipe, bool initVecBuffers = true) + { + pipe_ = pipe; + q_.SetGlobalBuffer((__gm__ T *)q); + k_.SetGlobalBuffer((__gm__ T *)k); + v_.SetGlobalBuffer((__gm__ T *)v); + gk_.SetGlobalBuffer((__gm__ GK_T *)gk); + betaFp32_.SetGlobalBuffer((__gm__ float *)beta); + betaBf16_.SetGlobalBuffer((__gm__ bfloat16_t *)beta); + if (initialState != nullptr) { + initialState_.SetGlobalBuffer((__gm__ float *)initialState); + } + cuSeqlensAddr_ = reinterpret_cast<__gm__ int64_t *>(cuSeqlens); + if (preparedQG != nullptr) { + preparedQG_.SetGlobalBuffer((__gm__ T *)preparedQG); + } + if (preparedAqk != nullptr) { + preparedAqk_.SetGlobalBuffer((__gm__ T *)preparedAqk); + } + if (propagatedVNew != nullptr) { + propagatedVNew_.SetGlobalBuffer((__gm__ T *)propagatedVNew); + } + if (propagatedH != nullptr) { + propagatedH_.SetGlobalBuffer((__gm__ T *)propagatedH); + } + chunkIndicesAddr_ = reinterpret_cast<__gm__ int64_t *>(chunkIndices); + o_.SetGlobalBuffer((__gm__ OUT_T *)o); + finalState_.SetGlobalBuffer((__gm__ float *)finalState); + aqk_.SetGlobalBuffer((__gm__ float *)aqk); + akk_.SetGlobalBuffer((__gm__ AKK_T *)akk); + w_.SetGlobalBuffer((__gm__ T *)w); + u_.SetGlobalBuffer((__gm__ OUT_T *)u); + qg_.SetGlobalBuffer((__gm__ T *)qg); + kg_.SetGlobalBuffer((__gm__ T *)kg); + vNew_.SetGlobalBuffer((__gm__ T *)vNew); + h_.SetGlobalBuffer((__gm__ float *)h); + solveWorkspace_.SetGlobalBuffer((__gm__ float *)workspace); + + B_ = tiling.batch; + N_ = tiling.seqNum; + H_ = tiling.qHeadNum; + HV_ = tiling.vHeadNum; + T_ = tiling.seqlen; + K_ = tiling.kHeadDim; + V_ = tiling.vHeadDim; + BT_ = tiling.chunkSize; + NT_ = tiling.totalChunks; + scale_ = tiling.scale; + hasInitial_ = tiling.hasInitialState; + isVarLen_ = tiling.isVarLen; + inputSequenceMajor_ = tiling.inputSequenceMajor; + usedCoreNum_ = tiling.prepareUsedCoreNum; + betaIsFp32_ = tiling.betaDataType == 0; + constexpr uint64_t solvePipelineDepth = SAFE_GATE ? KDA_SOLVE_PIPELINE_DEPTH : 1; + const uint64_t solveBytes = + usedCoreNum_ * solvePipelineDepth * KDA_SOLVE_SCRATCH_SLOTS * BT_ * BT_ * sizeof(float); + const uint64_t alignedSolveBytes = + (solveBytes + KDA_WORKSPACE_ALIGN - 1) / KDA_WORKSPACE_ALIGN * KDA_WORKSPACE_ALIGN; + scoreWorkspace_.SetGlobalBuffer((__gm__ SCORE_T *)(workspace + alignedSolveBytes)); + if ASCEND_IS_AIV { + uint64_t subBlockNum = static_cast(GetSubBlockNum()); + solveCoreIdx_ = subBlockNum == 0 ? 0 : static_cast(GetBlockIdx()) / subBlockNum; + } else { + solveCoreIdx_ = static_cast(GetBlockIdx()); + } + if (pipe_ != nullptr && initVecBuffers) { + pipe_->InitBuffer(exp2Buf_, EXP2_UB_BYTES); + pipe_->InitBuffer(vecBuf_, KDA_VEC_ARENA_ELEMENTS * sizeof(float)); + const uint64_t gateStageElems = GatePipelineRows() * K_; + const uint64_t gateInputSlotBytes = gateStageElems * (2 * sizeof(T) + sizeof(GK_T)); + const uint64_t gatePipelineBytes = + KDA_GATE_PIPELINE_DEPTH * (gateInputSlotBytes + gateStageElems * sizeof(T)); + pipe_->InitBuffer(gateWritebackBuf_, static_cast(gatePipelineBytes)); + AllocVectorEvents(); + } + } + __aicore__ inline void ProcessAivOnly() + { + isAivOnly_ = true; + ProcessPreAiv(); + ReleaseVectorEvents(); + } + + __aicore__ inline void ProcessAiv() + { + ProcessPreAiv(); + ReleaseVectorEvents(); + } + + __aicore__ inline void ProcessAic() { ProcessPreAic(); } + +private: + __aicore__ inline void AllocVectorEvents() + { + mte2ToVEvent_ = pipe_->AllocEventID(); + vToMte2Event_ = pipe_->AllocEventID(); + vToMte3Event_ = pipe_->AllocEventID(); + mte3ToVEvent_ = pipe_->AllocEventID(); + mte2ToMte3Event_ = pipe_->AllocEventID(); + for (uint32_t slot = 0; slot < KDA_GATE_PIPELINE_DEPTH; ++slot) { + mte3ToMte2Events_[slot] = pipe_->AllocEventID(); + } + vectorEventsAllocated_ = true; + } + + __aicore__ inline void ReleaseVectorEvents() + { + if (!vectorEventsAllocated_) { + return; + } + pipe_->ReleaseEventID(mte2ToVEvent_); + pipe_->ReleaseEventID(vToMte2Event_); + pipe_->ReleaseEventID(vToMte3Event_); + pipe_->ReleaseEventID(mte3ToVEvent_); + pipe_->ReleaseEventID(mte2ToMte3Event_); + for (uint32_t slot = 0; slot < KDA_GATE_PIPELINE_DEPTH; ++slot) { + pipe_->ReleaseEventID(mte3ToMte2Events_[slot]); + } + vectorEventsAllocated_ = false; + } + + __aicore__ inline uint64_t QOffset(uint64_t b, uint64_t h, uint64_t t, uint64_t d) const + { + if (inputSequenceMajor_) { + return ((b * T_ + t) * H_ + h) * K_ + d; + } + return ((b * H_ + h) * T_ + t) * K_ + d; + } + + __aicore__ inline uint64_t VInputOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t d) const + { + if (inputSequenceMajor_) { + return ((b * T_ + t) * HV_ + hv) * V_ + d; + } + return ((b * HV_ + hv) * T_ + t) * V_ + d; + } + + __aicore__ inline uint64_t KVOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t d, uint64_t dim) const + { + return ((b * HV_ + hv) * T_ + t) * dim + d; + } + + __aicore__ inline uint64_t BetaOffset(uint64_t b, uint64_t hv, uint64_t t) const { return (b * HV_ + hv) * T_ + t; } + + __aicore__ inline uint64_t AOffset(uint64_t b, uint64_t hv, uint64_t t, uint64_t j) const + { + return ((b * HV_ + hv) * T_ + t) * BT_ + j; + } + + __aicore__ inline uint64_t HOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t d, uint64_t r) const + { + return (((b * HV_ + hv) * NT_ + chunkIdx) * K_ + d) * V_ + r; + } + + __aicore__ inline uint64_t WScratchOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t t, uint64_t d) const + { + return (((b * HV_ + hv) * NT_ + chunkIdx) * BT_ + t) * K_ + d; + } + + __aicore__ inline uint64_t SolveScratchOffset(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t slot) const + { + (void)b; + (void)hv; + (void)chunkIdx; + constexpr uint64_t solvePipelineDepth = SAFE_GATE ? KDA_SOLVE_PIPELINE_DEPTH : 1; + uint64_t matrixElements = BT_ * BT_; + return ((solveCoreIdx_ * solvePipelineDepth + activeSolveSlot_) * KDA_SOLVE_SCRATCH_SLOTS + slot) * + matrixElements; + } + + __aicore__ inline uint64_t ScoreScratchOffset(uint64_t slot, uint64_t plane, uint64_t t = 0, uint64_t d = 0) const + { + return (((solveCoreIdx_ * KDA_SCORE_SCRATCH_SLOTS + slot) * KDA_SCORE_SCRATCH_PLANES + plane) * BT_ + t) * K_ + + d; + } + + __aicore__ inline uint64_t ScoreScratchSlot(uint64_t queueSlot, uint64_t lane, bool pairHeads) const + { + return pairHeads ? queueSlot * KDA_SCORE_LANES + lane : queueSlot; + } + + __aicore__ inline uint64_t ScoreRefBlockSize() const + { + if constexpr (SAFE_GATE) { + return KDA_SAFE_SCORE_REF_BC; + } + return KDA_SCORE_REF_BC; + } + + __aicore__ inline uint64_t ScoreRowBlockCount(uint64_t curT, uint64_t rowBegin) const + { + uint64_t blockSize = ScoreRefBlockSize(); + uint64_t rowCount = curT - rowBegin; + if (rowCount > blockSize) { + rowCount = blockSize; + } + return rowCount; + } + + __aicore__ inline uint64_t ScoreRefToken(uint64_t start, uint64_t curT, uint64_t rowBegin, uint64_t rowCount) const + { + uint64_t ref = rowBegin + rowCount / 2; + if (ref >= curT) { + ref = curT - 1; + } + return start + ref; + } + + __aicore__ inline void RunExp2(LocalTensor &tensor, uint32_t count) + { + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + ClampExpInput(tensor, count); + Exp(tensor, tensor, count); + PipeBarrier(); + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + } + + __aicore__ inline void ClampExpInput(LocalTensor &tensor, uint32_t count) + { + Mins(tensor, tensor, KDA_EXP_INPUT_MAX, count); + PipeBarrier(); + Maxs(tensor, tensor, KDA_EXP_INPUT_MIN, count); + PipeBarrier(); + } + + __aicore__ inline void ClampScoreExpInput(LocalTensor &tensor, uint32_t count) + { + constexpr float expInputMax = + IsSameType::value ? KDA_SCORE_EXP_INPUT_MAX : KDA_EXP_INPUT_MAX; + constexpr float expInputMin = + IsSameType::value ? KDA_SCORE_EXP_INPUT_MIN : KDA_EXP_INPUT_MIN; + Mins(tensor, tensor, expInputMax, count); + PipeBarrier(); + Maxs(tensor, tensor, expInputMin, count); + PipeBarrier(); + } + + template + __aicore__ inline void ClampFp32ForCast(LocalTensor &tensor, uint32_t count) + { + if constexpr (IsSameType::value) { + Mins(tensor, tensor, KDA_FP16_MAX, count); + PipeBarrier(); + Maxs(tensor, tensor, -KDA_FP16_MAX, count); + PipeBarrier(); + } + } + + __aicore__ inline void ClampFp32ToOutputType(LocalTensor &tensor, uint32_t count) + { + ClampFp32ForCast(tensor, count); + } + + template + __aicore__ inline void CopyVectorIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset, + uint64_t count) + { + uint64_t rowBytes = count * static_cast(sizeof(CopyT)); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst, src[offset], static_cast(count)); + return; + } + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPadParams padParams{false, 0, 0, 0}; + DataCopyPad(dst, src[offset], params, padParams); + } + + template + __aicore__ inline void CopyRowsIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset, uint64_t rows, + uint64_t cols, uint64_t rowStride) + { + if (rows == 0 || cols == 0) { + return; + } + if (rowStride == cols) { + CopyVectorIn(dst, src, offset, rows * cols); + return; + } + DataCopyExtParams params{static_cast(rows), static_cast(cols * sizeof(CopyT)), + static_cast((rowStride - cols) * sizeof(CopyT)), 0, 0}; + DataCopyPadExtParams padParams{false, 0, 0, 0}; + DataCopyPad(dst, src[offset], params, padParams); + } + + template + __aicore__ inline void CopyVectorOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src, + uint64_t count) + { + uint64_t rowBytes = count * static_cast(sizeof(CopyT)); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst[offset], src, static_cast(count)); + return; + } + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPad(dst[offset], src, params); + } + + template + __aicore__ inline void CopyRowIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset) + { + CopyVectorIn(dst, src, offset, K_); + } + + template + __aicore__ inline void CopyRowOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src) + { + CopyVectorOut(dst, offset, src, K_); + } + + __aicore__ inline LocalTensor VecScratch(uint64_t slot) + { + return vecBuf_.Get()[slot * EXP2_UB_ELEMENTS]; + } + + __aicore__ inline uint64_t GateStageElems() const { return GatePipelineRows() * K_; } + + __aicore__ inline uint64_t GatePipelineRows() const + { + constexpr uint64_t fixedBytes = static_cast(KDA_VEC_ARENA_ELEMENTS) * sizeof(float) + EXP2_UB_BYTES; + constexpr uint64_t availableBytes = KDA_AIV_UB_BUDGET_BYTES - fixedBytes; + uint64_t bytesPerRow = K_ * KDA_GATE_PIPELINE_DEPTH * (3 * sizeof(T) + sizeof(GK_T)); + uint64_t rows = bytesPerRow == 0 ? 0 : availableBytes / bytesPerRow; + return rows < KDA_GATE_TILE_ROWS ? rows : KDA_GATE_TILE_ROWS; + } + + __aicore__ inline uint64_t GateInputSlotBytes() const { return GateStageElems() * (2 * sizeof(T) + sizeof(GK_T)); } + + __aicore__ inline LocalTensor GateQTyped(uint64_t slot) + { + uint64_t byteOffset = slot * GateInputSlotBytes(); + return gateWritebackBuf_.Get()[byteOffset / sizeof(T)]; + } + + __aicore__ inline LocalTensor GateKTyped(uint64_t slot) + { + uint64_t byteOffset = slot * GateInputSlotBytes() + GateStageElems() * sizeof(T); + return gateWritebackBuf_.Get()[byteOffset / sizeof(T)]; + } + + __aicore__ inline LocalTensor GateGTyped(uint64_t slot) + { + uint64_t byteOffset = slot * GateInputSlotBytes() + 2 * GateStageElems() * sizeof(T); + return gateWritebackBuf_.Get()[byteOffset / sizeof(GK_T)]; + } + + __aicore__ inline LocalTensor GateKgTyped(uint64_t slot) + { + uint64_t byteOffset = KDA_GATE_PIPELINE_DEPTH * GateInputSlotBytes() + slot * GateStageElems() * sizeof(T); + return gateWritebackBuf_.Get()[byteOffset / sizeof(T)]; + } + + __aicore__ inline void PrefetchQKGate(uint64_t slot, uint64_t b, uint64_t h, uint64_t hv, uint64_t token, + uint64_t elems) + { + const uint64_t rows = elems / K_; + LocalTensor qTyped = GateQTyped(slot); + LocalTensor kTyped = GateKTyped(slot); + LocalTensor gateTyped = GateGTyped(slot); + CopyRowsIn(qTyped, q_, QOffset(b, h, token, 0), rows, K_, inputSequenceMajor_ ? H_ * K_ : K_); + CopyRowsIn(kTyped, k_, QOffset(b, h, token, 0), rows, K_, inputSequenceMajor_ ? H_ * K_ : K_); + CopyVectorIn(gateTyped, gk_, KVOffset(b, hv, token, 0, K_), elems); + SetFlag(mte2ToVEvent_); + } + + __aicore__ inline void PrefetchKGate(uint64_t slot, uint64_t b, uint64_t h, uint64_t hv, uint64_t token, + uint64_t elems) + { + const uint64_t rows = elems / K_; + LocalTensor kTyped = GateQTyped(slot); + LocalTensor gateTyped = GateGTyped(slot); + CopyRowsIn(kTyped, k_, QOffset(b, h, token, 0), rows, K_, inputSequenceMajor_ ? H_ * K_ : K_); + CopyVectorIn(gateTyped, gk_, KVOffset(b, hv, token, 0, K_), elems); + SetFlag(mte2ToVEvent_); + } + + __aicore__ inline void WaitGateInputReady() { WaitFlag(mte2ToVEvent_); } + + __aicore__ inline void WaitGateOutputForMte2(uint64_t slot = 0) + { + WaitFlag(mte3ToMte2Events_[slot]); + } + + __aicore__ inline void WaitGateOutputForVector() { WaitFlag(mte3ToVEvent_); } + + __aicore__ inline void SignalGateOutputDone() + { + SetFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + } + + __aicore__ inline void SignalGateOutputDoneForMte2(uint64_t slot) + { + SetFlag(mte3ToMte2Events_[slot]); + } + + template + __aicore__ inline void LoadAsFloatRow(GlobalTensor &src, uint64_t srcOffset, LocalTensor &dst, + uint64_t count) + { + if constexpr (IsSameType::value) { + CopyVectorIn(dst, src, srcOffset, count); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Adds(dst, dst, 0.0f, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } else { + constexpr uint32_t typedOffset = EXP2_UB_ELEMENTS * sizeof(float) / sizeof(CopyT); + LocalTensor rowLocal = exp2Buf_.Get()[typedOffset]; + CopyVectorIn(rowLocal, src, srcOffset, count); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Cast(dst, rowLocal, RoundMode::CAST_NONE, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + } + PipeBarrier(); + } + + // beta arrives as fp32 or bf16 (tiling.betaDataType). Resolving that at + // runtime keeps the whole prepare kernel independent of the beta dtype; + // templating on it doubled every Prepare instantiation for one vector cast. + __aicore__ inline void LoadBetaAsFloatRow(uint64_t srcOffset, LocalTensor &dst, uint64_t count) + { + if (betaIsFp32_) { + LoadAsFloatRow(betaFp32_, srcOffset, dst, count); + } else { + LoadAsFloatRow(betaBf16_, srcOffset, dst, count); + } + } + + template + __aicore__ inline void LoadAsFloatVector(GlobalTensor &src, uint64_t srcOffset, LocalTensor &dst, + LocalTensor &typedScratch, uint64_t count) + { + if constexpr (IsSameType::value) { + CopyVectorIn(dst, src, srcOffset, count); + } else { + CopyVectorIn(typedScratch, src, srcOffset, count); + } + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + if constexpr (!IsSameType::value) { + Cast(dst, typedScratch, RoundMode::CAST_NONE, static_cast(count)); + PipeBarrier(); + } + } + + template + __aicore__ inline void StoreFloatRow(GlobalTensor &dst, uint64_t dstOffset, LocalTensor &src, + uint64_t count) + { + if constexpr (IsSameType::value) { + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(dst, dstOffset, src, count); + } else { + constexpr uint32_t typedOffset = EXP2_UB_ELEMENTS * sizeof(float) / sizeof(CopyT); + LocalTensor rowLocal = exp2Buf_.Get()[typedOffset]; + Cast(rowLocal, src, RoundMode::CAST_RINT, static_cast(count)); + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(dst, dstOffset, rowLocal, count); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline LocalTensor Exp2NegG(uint64_t b, uint64_t hv, uint64_t t) + { + LocalTensor exp2Local = exp2Buf_.Get(); + LoadAsFloatRow(gk_, KVOffset(b, hv, t, 0, K_), exp2Local, K_); + Muls(exp2Local, exp2Local, -LN2, static_cast(K_)); + PipeBarrier(); + RunExp2(exp2Local, static_cast(K_)); + return exp2Local; + } + + __aicore__ inline void PrepareScoreFactorsBulk(uint64_t b, uint64_t h, uint64_t hv, uint64_t start, + uint64_t subBlockIdx, uint64_t subBlockNum, uint64_t refToken, + uint64_t scoreRowBegin, uint64_t scoreRowCount, uint64_t validColEnd, + uint64_t scoreSlot) + { + LocalTensor refFp32 = exp2Buf_.Get(); + LoadAsFloatRow(gk_, KVOffset(b, hv, refToken, 0, K_), refFp32, K_); + + uint64_t qwBegin = scoreRowBegin + (scoreRowCount * subBlockIdx) / subBlockNum; + uint64_t qwEnd = scoreRowBegin + (scoreRowCount * (subBlockIdx + 1)) / subBlockNum; + uint64_t qwMaxRows = GatePipelineRows(); + bool qwOutputPending = false; + uint64_t qwSlot = 0; + if (qwBegin < qwEnd && qwMaxRows > 0) { + uint64_t firstRows = qwEnd - qwBegin; + if (firstRows > qwMaxRows) { + firstRows = qwMaxRows; + } + PrefetchQKGate(qwSlot, b, h, hv, start + qwBegin, firstRows * K_); + } + for (uint64_t tileRow = qwBegin; tileRow < qwEnd && qwMaxRows > 0; tileRow += qwMaxRows) { + uint64_t tileRows = qwEnd - tileRow; + if (tileRows > qwMaxRows) { + tileRows = qwMaxRows; + } + uint64_t elems = tileRows * K_; + LocalTensor qTyped = GateQTyped(qwSlot); + LocalTensor kTyped = GateKTyped(qwSlot); + LocalTensor qScore = qTyped.template ReinterpretCast(); + LocalTensor kScore = kTyped.template ReinterpretCast(); + LocalTensor gateTyped = GateGTyped(qwSlot); + LocalTensor arena = vecBuf_.Get(); + LocalTensor qFp32 = arena; + LocalTensor kFp32 = arena[elems]; + LocalTensor gFp32 = arena[2 * elems]; + LocalTensor expFp32 = arena[3 * elems]; + LocalTensor outFp32 = arena[4 * elems]; + + WaitGateInputReady(); + Cast(qFp32, qTyped, RoundMode::CAST_NONE, static_cast(elems)); + Cast(kFp32, kTyped, RoundMode::CAST_NONE, static_cast(elems)); + if constexpr (IsSameType::value) { + gFp32 = gateTyped; + } else { + Cast(gFp32, gateTyped, RoundMode::CAST_NONE, static_cast(elems)); + } + if (qwOutputPending) { + WaitGateOutputForMte2(); + } + uint64_t nextTileRow = tileRow + qwMaxRows; + if (nextTileRow < qwEnd) { + uint64_t nextRows = qwEnd - nextTileRow; + if (nextRows > qwMaxRows) { + nextRows = qwMaxRows; + } + PrefetchQKGate(qwSlot ^ 1, b, h, hv, start + nextTileRow, nextRows * K_); + } + PipeBarrier(); + for (uint64_t row = 0; row < tileRows; ++row) { + Sub(expFp32[row * K_], gFp32[row * K_], refFp32, static_cast(K_)); + } + PipeBarrier(); + Muls(expFp32, expFp32, LN2, static_cast(elems)); + PipeBarrier(); + ClampScoreExpInput(expFp32, static_cast(elems)); + Exp(expFp32, expFp32, static_cast(elems)); + PipeBarrier(); + + Mul(outFp32, qFp32, expFp32, static_cast(elems)); + PipeBarrier(); + ClampFp32ForCast(outFp32, static_cast(elems)); + Cast(qScore, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + PipeBarrier(); + Mul(outFp32, kFp32, expFp32, static_cast(elems)); + PipeBarrier(); + ClampFp32ForCast(outFp32, static_cast(elems)); + Cast(kScore, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + PipeBarrier(); + + if (qwOutputPending) { + WaitGateOutputForVector(); + } + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(scoreWorkspace_, ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_QG, tileRow), qScore, elems); + CopyVectorOut(scoreWorkspace_, ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_W, tileRow), kScore, elems); + SignalGateOutputDone(); + qwOutputPending = true; + qwSlot ^= 1; + } + if (qwOutputPending) { + WaitGateOutputForMte2(); + WaitGateOutputForVector(); + } + + uint64_t kgBegin = (validColEnd * subBlockIdx) / subBlockNum; + uint64_t kgEnd = (validColEnd * (subBlockIdx + 1)) / subBlockNum; + uint64_t kgMaxRows = GatePipelineRows(); + bool kgOutputPending = false; + uint64_t kgSlot = 0; + if (kgBegin < kgEnd && kgMaxRows > 0) { + uint64_t firstRows = kgEnd - kgBegin; + if (firstRows > kgMaxRows) { + firstRows = kgMaxRows; + } + PrefetchKGate(kgSlot, b, h, hv, start + kgBegin, firstRows * K_); + } + for (uint64_t tileRow = kgBegin; tileRow < kgEnd && kgMaxRows > 0; tileRow += kgMaxRows) { + uint64_t tileRows = kgEnd - tileRow; + if (tileRows > kgMaxRows) { + tileRows = kgMaxRows; + } + uint64_t elems = tileRows * K_; + LocalTensor kTyped = GateQTyped(kgSlot); + LocalTensor kgScore = kTyped.template ReinterpretCast(); + LocalTensor gateTyped = GateGTyped(kgSlot); + LocalTensor arena = vecBuf_.Get(); + LocalTensor kFp32 = arena; + LocalTensor gFp32 = arena[elems]; + LocalTensor expFp32 = arena[2 * elems]; + LocalTensor outFp32 = arena[3 * elems]; + + WaitGateInputReady(); + Cast(kFp32, kTyped, RoundMode::CAST_NONE, static_cast(elems)); + if constexpr (IsSameType::value) { + gFp32 = gateTyped; + } else { + Cast(gFp32, gateTyped, RoundMode::CAST_NONE, static_cast(elems)); + } + if (kgOutputPending) { + WaitGateOutputForMte2(); + } + uint64_t nextTileRow = tileRow + kgMaxRows; + if (nextTileRow < kgEnd) { + uint64_t nextRows = kgEnd - nextTileRow; + if (nextRows > kgMaxRows) { + nextRows = kgMaxRows; + } + PrefetchKGate(kgSlot ^ 1, b, h, hv, start + nextTileRow, nextRows * K_); + } + PipeBarrier(); + for (uint64_t row = 0; row < tileRows; ++row) { + Sub(expFp32[row * K_], refFp32, gFp32[row * K_], static_cast(K_)); + } + PipeBarrier(); + Muls(expFp32, expFp32, LN2, static_cast(elems)); + PipeBarrier(); + ClampScoreExpInput(expFp32, static_cast(elems)); + Exp(expFp32, expFp32, static_cast(elems)); + PipeBarrier(); + Mul(outFp32, kFp32, expFp32, static_cast(elems)); + PipeBarrier(); + ClampFp32ForCast(outFp32, static_cast(elems)); + Cast(kgScore, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + PipeBarrier(); + + if (kgOutputPending) { + WaitGateOutputForVector(); + } + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(scoreWorkspace_, ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_KG, tileRow), kgScore, + elems); + SignalGateOutputDone(); + kgOutputPending = true; + kgSlot ^= 1; + } + if (kgOutputPending) { + WaitGateOutputForMte2(); + WaitGateOutputForVector(); + } + } + + __aicore__ inline void PrepareGateProductsBulk(uint64_t b, uint64_t h, uint64_t hv, uint64_t start, uint64_t curT, + uint64_t subBlockIdx, uint64_t subBlockNum, bool useRef, + uint64_t refToken, uint64_t validColEnd, bool writeScoreScratch, + uint64_t scoreSlot) + { + if constexpr (IsSameType::value) { + return; + } + if (subBlockNum == 0 || subBlockIdx >= subBlockNum || K_ == 0) { + return; + } + uint64_t rowBegin = (curT * subBlockIdx) / subBlockNum; + uint64_t rowEnd = (curT * (subBlockIdx + 1)) / subBlockNum; + if (rowBegin >= rowEnd) { + return; + } + + uint64_t maxRows = GatePipelineRows(); + if (maxRows == 0) { + return; + } + LocalTensor refFp32 = exp2Buf_.Get(); + if (useRef) { + LoadAsFloatRow(gk_, KVOffset(b, hv, refToken, 0, K_), refFp32, K_); + } + + bool outputPending = false; + uint64_t gateSlot = 0; + uint64_t firstRows = rowEnd - rowBegin; + if (firstRows > maxRows) { + firstRows = maxRows; + } + PrefetchQKGate(gateSlot, b, h, hv, start + rowBegin, firstRows * K_); + for (uint64_t tileRow = rowBegin; tileRow < rowEnd; tileRow += maxRows) { + uint64_t tileRows = rowEnd - tileRow; + if (tileRows > maxRows) { + tileRows = maxRows; + } + uint64_t elems = tileRows * K_; + LocalTensor qTyped = GateQTyped(gateSlot); + LocalTensor kTyped = GateKTyped(gateSlot); + LocalTensor kgTyped = GateKgTyped(gateSlot); + LocalTensor qScore = qTyped.template ReinterpretCast(); + LocalTensor wScore = kTyped.template ReinterpretCast(); + LocalTensor kgScore = kgTyped.template ReinterpretCast(); + LocalTensor gateTyped = GateGTyped(gateSlot); + LocalTensor arena = vecBuf_.Get(); + LocalTensor qFp32 = arena; + LocalTensor kFp32 = arena[elems]; + LocalTensor gFp32 = arena[2 * elems]; + LocalTensor expFp32 = arena[3 * elems]; + LocalTensor outFp32 = arena[4 * elems]; + + uint64_t token = start + tileRow; + WaitGateInputReady(); + Cast(qFp32, qTyped, RoundMode::CAST_NONE, static_cast(elems)); + Cast(kFp32, kTyped, RoundMode::CAST_NONE, static_cast(elems)); + if constexpr (IsSameType::value) { + gFp32 = gateTyped; + } else { + Cast(gFp32, gateTyped, RoundMode::CAST_NONE, static_cast(elems)); + } + uint64_t nextTileRow = tileRow + maxRows; + if (outputPending) { + WaitGateOutputForMte2(); + } + if (nextTileRow < rowEnd) { + uint64_t nextRows = rowEnd - nextTileRow; + if (nextRows > maxRows) { + nextRows = maxRows; + } + PrefetchQKGate(gateSlot ^ 1, b, h, hv, start + nextTileRow, nextRows * K_); + } + PipeBarrier(); + + if (useRef) { + for (uint64_t row = 0; row < tileRows; ++row) { + Sub(expFp32[row * K_], gFp32[row * K_], refFp32, static_cast(K_)); + } + } else { + Adds(expFp32, gFp32, 0.0f, static_cast(elems)); + } + PipeBarrier(); + Muls(expFp32, expFp32, LN2, static_cast(elems)); + PipeBarrier(); + if (writeScoreScratch) { + ClampScoreExpInput(expFp32, static_cast(elems)); + } else { + ClampExpInput(expFp32, static_cast(elems)); + } + Exp(expFp32, expFp32, static_cast(elems)); + PipeBarrier(); + + Mul(outFp32, qFp32, expFp32, static_cast(elems)); + PipeBarrier(); + if (writeScoreScratch) { + ClampFp32ForCast(outFp32, static_cast(elems)); + Cast(qScore, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + } else { + ClampFp32ToOutputType(outFp32, static_cast(elems)); + Cast(qTyped, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + } + PipeBarrier(); + + Mul(outFp32, kFp32, expFp32, static_cast(elems)); + PipeBarrier(); + if (writeScoreScratch) { + ClampFp32ForCast(outFp32, static_cast(elems)); + Cast(wScore, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + } else { + ClampFp32ToOutputType(outFp32, static_cast(elems)); + Cast(kTyped, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + } + PipeBarrier(); + + if (useRef) { + for (uint64_t row = 0; row < tileRows; ++row) { + Sub(expFp32[row * K_], refFp32, gFp32[row * K_], static_cast(K_)); + } + } else { + Muls(expFp32, gFp32, -1.0f, static_cast(elems)); + } + PipeBarrier(); + Muls(expFp32, expFp32, LN2, static_cast(elems)); + PipeBarrier(); + if (writeScoreScratch) { + ClampScoreExpInput(expFp32, static_cast(elems)); + } else { + ClampExpInput(expFp32, static_cast(elems)); + } + Exp(expFp32, expFp32, static_cast(elems)); + PipeBarrier(); + Mul(outFp32, kFp32, expFp32, static_cast(elems)); + PipeBarrier(); + if (useRef && tileRow + tileRows > validColEnd) { + for (uint64_t row = 0; row < tileRows; ++row) { + if (tileRow + row >= validColEnd) { + Duplicate(outFp32[row * K_], 0.0f, static_cast(K_)); + } + } + PipeBarrier(); + } + if (writeScoreScratch) { + ClampFp32ForCast(outFp32, static_cast(elems)); + } else { + ClampFp32ToOutputType(outFp32, static_cast(elems)); + } + if (outputPending) { + WaitGateOutputForVector(); + } + if (writeScoreScratch) { + Cast(kgScore, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + } else { + Cast(kgTyped, outFp32, RoundMode::CAST_RINT, static_cast(elems)); + } + PipeBarrier(); + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + if (writeScoreScratch) { + CopyVectorOut(scoreWorkspace_, ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_QG, tileRow), qScore, + elems); + CopyVectorOut(scoreWorkspace_, ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_W, tileRow), wScore, + elems); + CopyVectorOut(scoreWorkspace_, ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_KG, tileRow), kgScore, + elems); + } else { + CopyVectorOut(qg_, KVOffset(b, hv, token, 0, K_), qTyped, elems); + CopyVectorOut(w_, KVOffset(b, hv, token, 0, K_), kTyped, elems); + CopyVectorOut(kg_, KVOffset(b, hv, token, 0, K_), kgTyped, elems); + } + SignalGateOutputDone(); + outputPending = true; + gateSlot ^= 1; + } + if (outputPending) { + WaitGateOutputForMte2(); + WaitGateOutputForVector(); + } + return; + } + + __aicore__ inline void PrepareGateProducts(uint64_t b, uint64_t h, uint64_t hv, uint64_t start, uint64_t curT, + uint64_t subBlockIdx, uint64_t subBlockNum, bool useRef = false, + uint64_t refToken = 0, uint64_t validColEnd = 0, + bool writeScoreScratch = false, uint64_t scoreSlot = 0, + uint64_t scoreRowBegin = 0, uint64_t scoreRowCount = 0) + { + if (subBlockNum == 0 || subBlockIdx >= subBlockNum) { + return; + } + if (validColEnd == 0 || validColEnd > curT) { + validColEnd = curT; + } + if (writeScoreScratch) { + PrepareScoreFactorsBulk(b, h, hv, start, subBlockIdx, subBlockNum, refToken, scoreRowBegin, scoreRowCount, + validColEnd, scoreSlot); + return; + } + PrepareGateProductsBulk(b, h, hv, start, curT, subBlockIdx, subBlockNum, useRef, refToken, validColEnd, + writeScoreScratch, scoreSlot); + } + + __aicore__ inline void ComputeRawAqkAkkCube(uint64_t b, uint64_t hv, uint64_t start, uint64_t curT) + { + ComputeRawAqkAkkCubeBlock(b, hv, start, curT, 0, curT); + } + + __aicore__ inline void ComputeRawAqkAkkCubeBlock(uint64_t b, uint64_t hv, uint64_t start, uint64_t curT, + uint64_t rowBegin, uint64_t rowCount, + bool readScoreScratch = false, uint64_t scoreSlot = 0, + uint64_t colCount = 0) + { + using ElementA = SCORE_T; + using ElementB = SCORE_T; + using ElementC = float; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::ColumnMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using BlockMmad = Catlass::Gemm::Block::BlockMmadTla; + + Catlass::Arch::Resource resource; + BlockMmad blockMmad(resource); + auto layoutA = tla::MakeLayout(BT_, K_); + auto layoutB = tla::MakeLayout(K_, BT_); + auto layoutC = tla::MakeLayout(BT_, BT_); + if (colCount == 0 || colCount > curT) { + colCount = curT; + } + Catlass::GemmCoord shape{static_cast(rowCount), static_cast(colCount), + static_cast(K_)}; + + (void)readScoreScratch; + auto tensorQPos = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_QG)], layoutA, + Catlass::Arch::PositionGM{}); + auto tensorKPos = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_W)], layoutA, + Catlass::Arch::PositionGM{}); + auto tensorKNeg = tla::MakeTensor(scoreWorkspace_[ScoreScratchOffset(scoreSlot, KDA_SCORE_SCRATCH_KG)], layoutB, + Catlass::Arch::PositionGM{}); + auto tensorAqk = tla::MakeTensor(aqk_[AOffset(b, hv, start, 0)], layoutC, Catlass::Arch::PositionGM{}); + auto tensorAkk = tla::MakeTensor(akk_[AOffset(b, hv, start, 0)], layoutC, Catlass::Arch::PositionGM{}); + + auto blockQPos = GetTile(tensorQPos, tla::MakeCoord(rowBegin, 0), tla::MakeShape(shape.m(), shape.k())); + auto blockKPos = GetTile(tensorKPos, tla::MakeCoord(rowBegin, 0), tla::MakeShape(shape.m(), shape.k())); + auto blockKNeg = GetTile(tensorKNeg, tla::MakeCoord(0, 0), tla::MakeShape(shape.k(), shape.n())); + auto blockAqk = GetTile(tensorAqk, tla::MakeCoord(rowBegin, 0), tla::MakeShape(shape.m(), shape.n())); + auto blockAkk = GetTile(tensorAkk, tla::MakeCoord(rowBegin, 0), tla::MakeShape(shape.m(), shape.n())); + + blockMmad.preSetFlags(); + blockMmad(blockQPos, blockKNeg, blockAqk, shape); + blockMmad(blockKPos, blockKNeg, blockAkk, shape); + blockMmad.finalWaitFlags(); + } + + __aicore__ inline bool UseAkkCubeSolve(uint64_t curT) const + { + return curT > 0 && curT <= BT_ && (BT_ == 64 || BT_ == 128) && K_ >= 16 && V_ >= 16 && V_ <= 256 && + K_ % 16 == 0 && V_ % 16 == 0; + } + + __aicore__ inline bool UsePostWuCube(uint64_t curT) const + { + return curT > 0 && curT <= BT_ && (BT_ == 64 || BT_ == 128) && K_ >= 16 && V_ >= 16 && V_ <= 256 && + K_ % 16 == 0 && V_ % 16 == 0; + } + + __aicore__ inline void CopyLocalFloat(LocalTensor dst, LocalTensor src, uint64_t count) + { + if (count == 0) { + return; + } + Adds(dst, src, 0.0f, static_cast(count)); + PipeBarrier(); + } + + __aicore__ inline void FillLocalFloat(LocalTensor dst, float value, uint64_t count) + { + if (count == 0) { + return; + } + Duplicate(dst, value, static_cast(count)); + PipeBarrier(); + } + + __aicore__ inline void ForwardSubDiag16(LocalTensor diag, LocalTensor row, LocalTensor prod, + LocalTensor rowBrcb, LocalTensor reduced, uint64_t valid) + { + constexpr uint32_t brcbStride = 8; + constexpr uint32_t diagSize = KDA_SOLVE_DIAG_BT; + constexpr uint8_t rowBlk = diagSize * sizeof(float) / 32; + + for (uint64_t i = 2; i < valid; ++i) { + uint32_t rowOffset = static_cast(i * diagSize); + DataCopy(row, diag[rowOffset], diagSize); + PipeBarrier(); + + Brcb(rowBrcb, row, diagSize / brcbStride, {1, 8}); + PipeBarrier(); + for (uint32_t col = 0; col < diagSize; col += brcbStride) { + Mul(prod[col], diag[col], rowBrcb, brcbStride, static_cast(diagSize), + {1, 1, 0, rowBlk, rowBlk, 1}); + } + PipeBarrier(); + + uint32_t remain = diagSize; + while (remain > 1) { + uint32_t calcCount = (remain / 2) * diagSize; + remain = (remain + 1) / 2; + Add(prod, prod, prod[remain * diagSize], calcCount); + PipeBarrier(); + } + DataCopy(reduced, prod, diagSize); + PipeBarrier(); + Add(row, row, reduced, diagSize); + PipeBarrier(); + DataCopy(diag[rowOffset], row, diagSize); + PipeBarrier(); + } + + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + for (uint32_t i = 0; i < diagSize; ++i) { + uint32_t diagOffset = i * diagSize + i; + if (i < valid) { + diag.SetValue(diagOffset, diag.GetValue(diagOffset) + 1.0f); + } else { + diag.SetValue(diagOffset, 1.0f); + } + } + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + } + + __aicore__ inline void SolveDiagonalBlocksInRows(LocalTensor akkMat, LocalTensor xMat, + LocalTensor arena, uint64_t scratchBase, uint64_t curT, + uint64_t rowBegin, uint64_t rowCount) + { + constexpr uint32_t diagSize = KDA_SOLVE_DIAG_BT; + constexpr uint32_t diagElements = diagSize * diagSize; + constexpr uint32_t brcbElements = diagSize * 8; + + LocalTensor diag = arena[scratchBase]; + LocalTensor row = diag[diagElements]; + LocalTensor prod = row[diagSize]; + LocalTensor rowBrcb = prod[diagElements]; + LocalTensor reduced = rowBrcb[brcbElements]; + + uint64_t rowEnd = rowBegin + rowCount; + for (uint64_t blockBegin = 0; blockBegin < BT_; blockBegin += diagSize) { + if (blockBegin < rowBegin || blockBegin + diagSize > rowEnd) { + continue; + } + Duplicate(diag, 0.0f, diagElements); + PipeBarrier(); + + uint64_t localBlockRow = blockBegin - rowBegin; + uint64_t valid = blockBegin < curT ? curT - blockBegin : 0; + if (valid > diagSize) { + valid = diagSize; + } + for (uint32_t rowIdx = 0; rowIdx < diagSize; ++rowIdx) { + uint64_t srcOffset = (localBlockRow + rowIdx) * BT_ + blockBegin; + Muls(diag[rowIdx * diagSize], akkMat[srcOffset], -1.0f, diagSize); + } + PipeBarrier(); + + ForwardSubDiag16(diag, row, prod, rowBrcb, reduced, valid); + for (uint32_t rowIdx = 0; rowIdx < diagSize; ++rowIdx) { + uint64_t dstOffset = (localBlockRow + rowIdx) * BT_ + blockBegin; + Adds(xMat[dstOffset], diag[rowIdx * diagSize], 0.0f, diagSize); + } + PipeBarrier(); + } + } + + __aicore__ inline void BuildPrefixMask(LocalTensor dst, uint64_t prefix, uint64_t count) + { + if (prefix > count) { + prefix = count; + } + Duplicate(dst, 0.0f, static_cast(count)); + if (prefix > 0) { + Duplicate(dst, 1.0f, static_cast(prefix)); + } + PipeBarrier(); + } + + __aicore__ inline uint64_t BuildCausalMask(uint64_t threshold, uint64_t colBegin) const + { + if (threshold <= colBegin) { + return ~0ULL; + } + if (threshold >= colBegin + KDA_SOLVE_BT) { + return 0ULL; + } + return ~0ULL << (threshold - colBegin); + } + + __aicore__ inline void BuildCausalSelectMasks(LocalTensor aqkMask, LocalTensor akkMask, + uint64_t rowBegin, uint64_t rowCount, uint64_t colBegin) + { + __ubuf__ uint64_t *aqkMaskPtr = reinterpret_cast<__ubuf__ uint64_t *>(aqkMask.GetPhyAddr()); + __ubuf__ uint64_t *akkMaskPtr = reinterpret_cast<__ubuf__ uint64_t *>(akkMask.GetPhyAddr()); + for (uint32_t localRow = 0; localRow < rowCount; ++localRow) { + uint32_t row = static_cast(rowBegin + localRow); + aqkMaskPtr[localRow] = BuildCausalMask(static_cast(row) + 1, colBegin); + akkMaskPtr[localRow] = BuildCausalMask(static_cast(row), colBegin); + } + } + + __aicore__ inline void SelectCausalRows(LocalTensor aqkMat, LocalTensor akkMat, uint64_t rowBegin, + uint64_t rowCount) + { + LocalTensor aqkMask = vecBuf_.Get()[KDA_SELECT_AQK_MASK_BYTE_OFFSET]; + LocalTensor akkMask = vecBuf_.Get()[KDA_SELECT_AKK_MASK_BYTE_OFFSET]; + LocalTensor zeroLocal = vecBuf_.Get()[KDA_SELECT_ZERO_FLOAT_OFFSET]; + Duplicate(zeroLocal, 0.0f, 8); + PipeBarrier(); + + uint64_t colBlockCount = (BT_ + KDA_SOLVE_BT - 1) / KDA_SOLVE_BT; + for (uint64_t colBlock = 0; colBlock < colBlockCount; ++colBlock) { + uint64_t maskOffset = colBlock * KDA_SELECT_COL_MASK_BYTES; + uint64_t colBegin = colBlock * KDA_SOLVE_BT; + BuildCausalSelectMasks(aqkMask[maskOffset], akkMask[maskOffset], rowBegin, rowCount, colBegin); + } + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + + uint8_t rowStride = static_cast(BT_ * sizeof(float) / 32); + BinaryRepeatParams repeatParams = {1, 0, 1, rowStride, 0, rowStride}; + for (uint64_t colBlock = 0; colBlock < colBlockCount; ++colBlock) { + uint64_t maskOffset = colBlock * KDA_SELECT_COL_MASK_BYTES; + uint64_t colBegin = colBlock * KDA_SOLVE_BT; + Select(aqkMat[colBegin], aqkMask[maskOffset], zeroLocal, aqkMat[colBegin], SELMODE::VSEL_TENSOR_TENSOR_MODE, + KDA_SOLVE_BT, static_cast(rowCount), repeatParams); + Select(akkMat[colBegin], akkMask[maskOffset], zeroLocal, akkMat[colBegin], SELMODE::VSEL_TENSOR_TENSOR_MODE, + KDA_SOLVE_BT, static_cast(rowCount), repeatParams); + } + PipeBarrier(); + SetFlag(EXP2_EVENT_ID); + WaitFlag(EXP2_EVENT_ID); + } + + __aicore__ inline void PrepareAqkAkkSolveInput64(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start) + { + LocalTensor arena = vecBuf_.Get(); + LocalTensor aqkMat = arena; + LocalTensor akkMat = arena[KDA_SOLVE_MATRIX_ELEMENTS]; + LocalTensor xMat = arena[2 * KDA_SOLVE_MATRIX_ELEMENTS]; + LocalTensor betaLocal = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS]; + LocalTensor betaBrcb = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS + KDA_SOLVE_BT]; + LocalTensor maskLocal = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS + KDA_SOLVE_BT + 512]; + LocalTensor oneHotLocal = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS + KDA_SOLVE_BT + 512 + KDA_SOLVE_BT]; + + LoadBetaAsFloatRow(BetaOffset(b, hv, start), betaLocal, KDA_SOLVE_BT); + Brcb(betaBrcb, betaLocal, 8, {1, 8}); + PipeBarrier(); + + DataCopy(aqkMat, aqk_[AOffset(b, hv, start, 0)], KDA_SOLVE_MATRIX_ELEMENTS); + DataCopy(akkMat, akk_[AOffset(b, hv, start, 0)], KDA_SOLVE_MATRIX_ELEMENTS); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + + for (uint64_t col = 0; col < KDA_SOLVE_BT; col += 8) { + Mul(akkMat[col], akkMat[col], betaBrcb, 8, KDA_SOLVE_BT, {1, 1, 1, 8, 8, 1}); + PipeBarrier(); + } + SelectCausalRows(aqkMat, akkMat, 0, KDA_SOLVE_BT); + + Muls(xMat, akkMat, -1.0f, KDA_SOLVE_MATRIX_ELEMENTS); + PipeBarrier(); + for (uint64_t row = 0; row < KDA_SOLVE_BT; ++row) { + BuildPrefixMask(maskLocal, row + 1, KDA_SOLVE_BT); + BuildPrefixMask(oneHotLocal, row, KDA_SOLVE_BT); + Sub(maskLocal, maskLocal, oneHotLocal, KDA_SOLVE_BT); + PipeBarrier(); + Add(xMat[row * KDA_SOLVE_BT], xMat[row * KDA_SOLVE_BT], maskLocal, KDA_SOLVE_BT); + PipeBarrier(); + } + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(aqk_[AOffset(b, hv, start, 0)], aqkMat, KDA_SOLVE_MATRIX_ELEMENTS); + DataCopy(akk_[AOffset(b, hv, start, 0)], akkMat, KDA_SOLVE_MATRIX_ELEMENTS); + DataCopy(h_[SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X)], xMat, KDA_SOLVE_MATRIX_ELEMENTS); + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void PrepareAqkAkkSolveInputTail(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t curT) + { + uint64_t elemCount = curT * KDA_SOLVE_BT; + LocalTensor arena = vecBuf_.Get(); + LocalTensor aqkMat = arena; + LocalTensor akkMat = arena[KDA_SOLVE_MATRIX_ELEMENTS]; + LocalTensor xMat = arena[2 * KDA_SOLVE_MATRIX_ELEMENTS]; + LocalTensor betaLocal = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS]; + LocalTensor betaBrcb = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS + KDA_SOLVE_BT]; + LocalTensor maskLocal = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS + KDA_SOLVE_BT + 512]; + LocalTensor oneHotLocal = arena[3 * KDA_SOLVE_MATRIX_ELEMENTS + KDA_SOLVE_BT + 512 + KDA_SOLVE_BT]; + + FillLocalFloat(betaLocal, 0.0f, KDA_SOLVE_BT); + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + LoadBetaAsFloatRow(BetaOffset(b, hv, start), betaLocal, curT); + Brcb(betaBrcb, betaLocal, 8, {1, 8}); + PipeBarrier(); + + DataCopy(aqkMat, aqk_[AOffset(b, hv, start, 0)], static_cast(elemCount)); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + if (elemCount < KDA_SOLVE_MATRIX_ELEMENTS) { + FillLocalFloat(aqkMat[elemCount], 0.0f, KDA_SOLVE_MATRIX_ELEMENTS - elemCount); + } + DataCopy(akkMat, akk_[AOffset(b, hv, start, 0)], static_cast(elemCount)); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + if (elemCount < KDA_SOLVE_MATRIX_ELEMENTS) { + FillLocalFloat(akkMat[elemCount], 0.0f, KDA_SOLVE_MATRIX_ELEMENTS - elemCount); + } + + for (uint64_t col = 0; col < KDA_SOLVE_BT; col += 8) { + Mul(akkMat[col], akkMat[col], betaBrcb, 8, KDA_SOLVE_BT, {1, 1, 1, 8, 8, 1}); + PipeBarrier(); + } + SelectCausalRows(aqkMat, akkMat, 0, KDA_SOLVE_BT); + + Muls(xMat, akkMat, -1.0f, KDA_SOLVE_MATRIX_ELEMENTS); + PipeBarrier(); + for (uint64_t row = 0; row < KDA_SOLVE_BT; ++row) { + BuildPrefixMask(maskLocal, row + 1, KDA_SOLVE_BT); + BuildPrefixMask(oneHotLocal, row, KDA_SOLVE_BT); + Sub(maskLocal, maskLocal, oneHotLocal, KDA_SOLVE_BT); + PipeBarrier(); + Add(xMat[row * KDA_SOLVE_BT], xMat[row * KDA_SOLVE_BT], maskLocal, KDA_SOLVE_BT); + PipeBarrier(); + } + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(aqk_[AOffset(b, hv, start, 0)], aqkMat, static_cast(elemCount)); + DataCopy(h_[SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X)], xMat, KDA_SOLVE_MATRIX_ELEMENTS); + DataCopy(h_[SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_Y0)], akkMat, KDA_SOLVE_MATRIX_ELEMENTS); + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void GetSolveRowRange(uint64_t curT, uint64_t subBlockIdx, uint64_t subBlockNum, + uint64_t &rowBegin, uint64_t &rowEnd) const + { + if (subBlockNum == 0 || subBlockIdx >= subBlockNum) { + rowBegin = 0; + rowEnd = 0; + return; + } + rowBegin = (curT * subBlockIdx) / subBlockNum; + rowEnd = (curT * (subBlockIdx + 1)) / subBlockNum; + } + + __aicore__ inline void PrepareAqkAkkSolveInputRows(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t curT, uint64_t rowBegin, uint64_t rowEnd, + bool storeLToAkk, bool storeLToScratch) + { + uint64_t rowCount = rowEnd - rowBegin; + if (rowCount == 0) { + return; + } + uint64_t validRowCount = rowBegin < curT ? curT - rowBegin : 0; + if (validRowCount > rowCount) { + validRowCount = rowCount; + } + uint64_t elemCount = rowCount * BT_; + uint64_t validElemCount = validRowCount * BT_; + LocalTensor arena = vecBuf_.Get(); + LocalTensor aqkMat = arena; + LocalTensor akkMat = arena[elemCount]; + LocalTensor xMat = arena[2 * elemCount]; + LocalTensor betaLocal = arena[3 * elemCount]; + LocalTensor betaBrcb = arena[3 * elemCount + BT_]; + LocalTensor maskLocal = arena[3 * elemCount + BT_ + 512]; + LocalTensor oneHotLocal = arena[3 * elemCount + BT_ + 512 + BT_]; + + uint64_t token = start + rowBegin; + + if (validRowCount < rowCount) { + FillLocalFloat(aqkMat, 0.0f, elemCount); + FillLocalFloat(akkMat, 0.0f, elemCount); + FillLocalFloat(betaLocal, 0.0f, rowCount); + } + SetFlag(vToMte2Event_); + WaitFlag(vToMte2Event_); + if (validRowCount > 0) { + LoadBetaAsFloatRow(BetaOffset(b, hv, token), betaLocal, validRowCount); + DataCopy(aqkMat, aqk_[AOffset(b, hv, token, 0)], static_cast(validElemCount)); + DataCopy(akkMat, akk_[AOffset(b, hv, token, 0)], static_cast(validElemCount)); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + } + Brcb(betaBrcb, betaLocal, static_cast((rowCount + 7) / 8), {1, 8}); + PipeBarrier(); + uint8_t rowStride = static_cast(BT_ * sizeof(float) / 32); + for (uint64_t col = 0; col < BT_; col += 8) { + Mul(akkMat[col], akkMat[col], betaBrcb, 8, static_cast(rowCount), + {1, 1, 0, rowStride, rowStride, 1}); + } + PipeBarrier(); + if (validRowCount > 0) { + SelectCausalRows(aqkMat, akkMat, rowBegin, validRowCount); + } + + Muls(xMat, akkMat, -1.0f, static_cast(elemCount)); + PipeBarrier(); + if constexpr (SAFE_GATE) { + uint64_t scratchBase = 3 * elemCount + BT_ + 512 + 2 * BT_; + SolveDiagonalBlocksInRows(akkMat, xMat, arena, scratchBase, curT, rowBegin, rowCount); + } else { + for (uint64_t localRow = 0; localRow < rowCount; ++localRow) { + uint64_t row = rowBegin + localRow; + BuildPrefixMask(maskLocal, row + 1, BT_); + BuildPrefixMask(oneHotLocal, row, BT_); + Sub(maskLocal, maskLocal, oneHotLocal, static_cast(BT_)); + PipeBarrier(); + Add(xMat[localRow * BT_], xMat[localRow * BT_], maskLocal, static_cast(BT_)); + PipeBarrier(); + } + } + + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X) + rowBegin * BT_; + uint64_t lBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_Y0) + rowBegin * BT_; + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + if (validRowCount > 0) { + DataCopy(aqk_[AOffset(b, hv, token, 0)], aqkMat, static_cast(validElemCount)); + if (storeLToAkk) { + DataCopy(akk_[AOffset(b, hv, token, 0)], akkMat, static_cast(validElemCount)); + } + } + DataCopy(solveWorkspace_[xBase], xMat, static_cast(elemCount)); + if (storeLToScratch) { + DataCopy(solveWorkspace_[lBase], akkMat, static_cast(elemCount)); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void CubeGemmSolveSub(GlobalTensor &tensorA, uint64_t baseA, uint64_t rowA, uint64_t colA, + GlobalTensor &tensorB, uint64_t baseB, uint64_t rowB, uint64_t colB, + GlobalTensor &tensorC, uint64_t baseC, uint64_t rowC, uint64_t colC, + uint32_t m, uint32_t n, uint32_t k) + { + using ElementA = float; + using ElementB = float; + using ElementC = float; + using LayoutTagA = Catlass::layout::RowMajor; + using LayoutTagB = Catlass::layout::RowMajor; + using LayoutTagC = Catlass::layout::RowMajor; + using TileCopy = Catlass::Gemm::Tile::PackedTileCopyTla; + using BlockMmad = + Catlass::Gemm::Block::BlockMmadTla; + Catlass::Arch::Resource resource; + auto layoutA = tla::MakeLayout(BT_, BT_); + auto layoutB = tla::MakeLayout(BT_, BT_); + auto layoutC = tla::MakeLayout(BT_, BT_); + auto tensorLayoutA = tla::MakeTensor(tensorA[baseA], layoutA, Catlass::Arch::PositionGM{}); + auto tensorLayoutB = tla::MakeTensor(tensorB[baseB], layoutB, Catlass::Arch::PositionGM{}); + auto tensorLayoutC = tla::MakeTensor(tensorC[baseC], layoutC, Catlass::Arch::PositionGM{}); + Catlass::GemmCoord shape{m, n, k}; + auto blockA = GetTile(tensorLayoutA, tla::MakeCoord(rowA, colA), tla::MakeShape(shape.m(), shape.k())); + auto blockB = GetTile(tensorLayoutB, tla::MakeCoord(rowB, colB), tla::MakeShape(shape.k(), shape.n())); + auto blockC = GetTile(tensorLayoutC, tla::MakeCoord(rowC, colC), tla::MakeShape(shape.m(), shape.n())); + BlockMmad blockMmad(resource); + blockMmad(blockA, blockB, blockC, shape); + PipeBarrier(); + } + + __aicore__ inline void AddSolveTmpToX(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, bool storeAkk) + { + LocalTensor arena = vecBuf_.Get(); + LocalTensor xLocal = arena; + LocalTensor tmpLocal = arena[KDA_SOLVE_MATRIX_ELEMENTS]; + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X); + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP); + + DataCopy(xLocal, h_[xBase], KDA_SOLVE_MATRIX_ELEMENTS); + DataCopy(tmpLocal, h_[tmpBase], KDA_SOLVE_MATRIX_ELEMENTS); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + + Add(xLocal, xLocal, tmpLocal, KDA_SOLVE_MATRIX_ELEMENTS); + PipeBarrier(); + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(h_[xBase], xLocal, KDA_SOLVE_MATRIX_ELEMENTS); + if (storeAkk) { + DataCopy(akk_[AOffset(b, hv, start, 0)], xLocal, KDA_SOLVE_MATRIX_ELEMENTS); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void AddSolveTmpToXTail(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t curT, + bool storeAkk) + { + uint64_t elemCount = curT * KDA_SOLVE_BT; + LocalTensor arena = vecBuf_.Get(); + LocalTensor xLocal = arena; + LocalTensor tmpLocal = arena[KDA_SOLVE_MATRIX_ELEMENTS]; + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X); + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP); + + DataCopy(xLocal, h_[xBase], KDA_SOLVE_MATRIX_ELEMENTS); + DataCopy(tmpLocal, h_[tmpBase], KDA_SOLVE_MATRIX_ELEMENTS); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + + Add(xLocal, xLocal, tmpLocal, KDA_SOLVE_MATRIX_ELEMENTS); + PipeBarrier(); + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(h_[xBase], xLocal, KDA_SOLVE_MATRIX_ELEMENTS); + if (storeAkk) { + DataCopy(akk_[AOffset(b, hv, start, 0)], xLocal, static_cast(elemCount)); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void AddSolveTmpToXRows(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t curT, + uint64_t rowBegin, uint64_t rowEnd, bool storeAkk) + { + uint64_t rowCount = rowEnd - rowBegin; + if (rowCount == 0) { + return; + } + uint64_t validRowCount = rowBegin < curT ? curT - rowBegin : 0; + if (validRowCount > rowCount) { + validRowCount = rowCount; + } + uint64_t elemCount = rowCount * BT_; + uint64_t validElemCount = validRowCount * BT_; + LocalTensor arena = vecBuf_.Get(); + LocalTensor xLocal = arena; + LocalTensor tmpLocal = arena[elemCount]; + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X) + rowBegin * BT_; + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP) + rowBegin * BT_; + uint64_t token = start + rowBegin; + + DataCopy(xLocal, solveWorkspace_[xBase], static_cast(elemCount)); + DataCopy(tmpLocal, solveWorkspace_[tmpBase], static_cast(elemCount)); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + + Add(xLocal, xLocal, tmpLocal, static_cast(elemCount)); + PipeBarrier(); + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(solveWorkspace_[xBase], xLocal, static_cast(elemCount)); + if (storeAkk && validRowCount > 0) { + DataCopy(akk_[AOffset(b, hv, token, 0)], xLocal, static_cast(validElemCount)); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void AddSolveTmpToXDiagRows(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t rowBegin, uint64_t rowEnd, bool storeAkk) + { + uint64_t rowCount = rowEnd - rowBegin; + if (rowCount == 0) { + return; + } + uint64_t elemCount = rowCount * BT_; + LocalTensor arena = vecBuf_.Get(); + LocalTensor xLocal = arena; + LocalTensor tmpLocal = arena[elemCount]; + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X) + rowBegin * BT_; + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP) + rowBegin * BT_; + uint64_t token = start + rowBegin; + + DataCopy(xLocal, solveWorkspace_[xBase], static_cast(elemCount)); + DataCopy(tmpLocal, solveWorkspace_[tmpBase], static_cast(elemCount)); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + + for (uint64_t localRow = 0; localRow < rowCount; ++localRow) { + uint64_t row = rowBegin + localRow; + uint64_t col = (row / KDA_SOLVE_DIAG_BT) * KDA_SOLVE_DIAG_BT; + uint64_t offset = localRow * BT_ + col; + Add(xLocal[offset], xLocal[offset], tmpLocal[offset], KDA_SOLVE_DIAG_BT); + PipeBarrier(); + } + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(solveWorkspace_[xBase], xLocal, static_cast(elemCount)); + if (storeAkk) { + DataCopy(akk_[AOffset(b, hv, token, 0)], xLocal, static_cast(elemCount)); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void StoreSolveXRowsToAkk(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t curT, uint64_t rowBegin, uint64_t rowEnd) + { + uint64_t validRowCount = rowBegin < curT ? curT - rowBegin : 0; + uint64_t rowCount = rowEnd - rowBegin; + if (validRowCount > rowCount) { + validRowCount = rowCount; + } + if (validRowCount == 0) { + return; + } + uint64_t elemCount = validRowCount * BT_; + LocalTensor xLocal = vecBuf_.Get(); + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X) + rowBegin * BT_; + + DataCopy(xLocal, solveWorkspace_[xBase], static_cast(elemCount)); + SetFlag(mte2ToMte3Event_); + WaitFlag(mte2ToMte3Event_); + DataCopy(akk_[AOffset(b, hv, start + rowBegin, 0)], xLocal, static_cast(elemCount)); + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + } + + __aicore__ inline void ComputeAkkMergeCube(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start) + { + uint64_t aiBase = AOffset(b, hv, start, 0); + uint64_t negABase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X); + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP); + + for (uint32_t mergeSize = 2 * KDA_SOLVE_DIAG_BT; mergeSize <= BT_; mergeSize *= 2) { + uint32_t half = mergeSize / 2; + for (uint32_t block = 0; block < BT_; block += mergeSize) { + uint32_t lower = block + half; + CubeGemmSolveSub(akk_, aiBase, lower, lower, solveWorkspace_, negABase, lower, block, solveWorkspace_, + tmpBase, 0, 0, half, half, half); + CubeGemmSolveSub(solveWorkspace_, tmpBase, 0, 0, akk_, aiBase, block, block, akk_, aiBase, lower, block, + half, half, half); + } + } + } + + __aicore__ inline void ComputeAkkMergeCubeWorkspace(uint64_t b, uint64_t hv, uint64_t chunkIdx) + { + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X); + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP); + + for (uint32_t mergeSize = 2 * KDA_SOLVE_DIAG_BT; mergeSize <= BT_; mergeSize *= 2) { + uint32_t half = mergeSize / 2; + for (uint32_t block = 0; block < BT_; block += mergeSize) { + uint32_t lower = block + half; + CubeGemmSolveSub(solveWorkspace_, xBase, lower, lower, solveWorkspace_, xBase, lower, block, + solveWorkspace_, tmpBase, 0, 0, half, half, half); + CubeGemmSolveSub(solveWorkspace_, tmpBase, 0, 0, solveWorkspace_, xBase, block, block, solveWorkspace_, + xBase, lower, block, half, half, half); + } + } + } + + __aicore__ inline void ComputeAkkInverseMchFull(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start) + { + uint64_t aBase = AOffset(b, hv, start, 0); + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X); + uint64_t yBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_Y0); + uint64_t yNextBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_Y1); + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP); + + uint32_t diagBlocks = static_cast(BT_ / KDA_SOLVE_DIAG_BT); + for (uint32_t block = 0; block < diagBlocks; ++block) { + uint32_t off = block * KDA_SOLVE_DIAG_BT; + CubeGemmSolveSub(akk_, aBase, off, off, akk_, aBase, off, off, solveWorkspace_, yBase, off, off, + KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT); + } + for (uint32_t iter = 0; iter < KDA_SOLVE_DIAG_MCH_ITERS; ++iter) { + for (uint32_t block = 0; block < diagBlocks; ++block) { + uint32_t off = block * KDA_SOLVE_DIAG_BT; + CubeGemmSolveSub(solveWorkspace_, xBase, off, off, solveWorkspace_, yBase, off, off, solveWorkspace_, + tmpBase, off, off, KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT); + } + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(mchSyncDoneFlag_); + if (iter + 1 < KDA_SOLVE_DIAG_MCH_ITERS) { + for (uint32_t block = 0; block < diagBlocks; ++block) { + uint32_t off = block * KDA_SOLVE_DIAG_BT; + CubeGemmSolveSub(solveWorkspace_, yBase, off, off, solveWorkspace_, yBase, off, off, + solveWorkspace_, yNextBase, off, off, KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT, + KDA_SOLVE_DIAG_BT); + } + } + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_FIX>(mchSyncReadyFlag_); + if (iter + 1 < KDA_SOLVE_DIAG_MCH_ITERS) { + uint64_t oldYBase = yBase; + yBase = yNextBase; + yNextBase = oldYBase; + } + } + ComputeAkkMergeCube(b, hv, chunkIdx, start); + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(mchSyncDoneFlag_); + } + + __aicore__ inline void ComputeAkkInverseMchTail(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t curT) + { + uint64_t xBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_X); + uint64_t lBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_Y0); + uint64_t yBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_Y1); + uint64_t yNextBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_Y0); + uint64_t tmpBase = SolveScratchOffset(b, hv, chunkIdx, KDA_SOLVE_SCRATCH_TMP); + (void)start; + (void)curT; + + uint32_t diagBlocks = static_cast(BT_ / KDA_SOLVE_DIAG_BT); + for (uint32_t block = 0; block < diagBlocks; ++block) { + uint32_t off = block * KDA_SOLVE_DIAG_BT; + CubeGemmSolveSub(solveWorkspace_, lBase, off, off, solveWorkspace_, lBase, off, off, solveWorkspace_, yBase, + off, off, KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT); + } + for (uint32_t iter = 0; iter < KDA_SOLVE_DIAG_MCH_ITERS; ++iter) { + for (uint32_t block = 0; block < diagBlocks; ++block) { + uint32_t off = block * KDA_SOLVE_DIAG_BT; + CubeGemmSolveSub(solveWorkspace_, xBase, off, off, solveWorkspace_, yBase, off, off, solveWorkspace_, + tmpBase, off, off, KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT); + } + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(mchSyncDoneFlag_); + if (iter + 1 < KDA_SOLVE_DIAG_MCH_ITERS) { + for (uint32_t block = 0; block < diagBlocks; ++block) { + uint32_t off = block * KDA_SOLVE_DIAG_BT; + CubeGemmSolveSub(solveWorkspace_, yBase, off, off, solveWorkspace_, yBase, off, off, + solveWorkspace_, yNextBase, off, off, KDA_SOLVE_DIAG_BT, KDA_SOLVE_DIAG_BT, + KDA_SOLVE_DIAG_BT); + } + } + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_FIX>(mchSyncReadyFlag_); + if (iter + 1 < KDA_SOLVE_DIAG_MCH_ITERS) { + uint64_t oldYBase = yBase; + yBase = yNextBase; + yNextBase = oldYBase; + } + } + ComputeAkkMergeCubeWorkspace(b, hv, chunkIdx); + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(mchSyncDoneFlag_); + } + + __aicore__ inline void ScaleRowsByBeta(GlobalTensor &src, GlobalTensor &dst, uint64_t b, uint64_t hv, + uint64_t start, uint64_t rowBegin, uint64_t rowCount, uint64_t dim, + LocalTensor &betaLocal, LocalTensor &betaBrcb, + LocalTensor &matrixLocal, bool sourceSequenceMajor = false) + { + constexpr uint64_t vecElemsPerRepeat = 64; + constexpr uint64_t typedOffsetFloats = 20480; + constexpr uint64_t typedOffset = typedOffsetFloats * sizeof(float) / sizeof(T); + uint64_t elemCount = rowCount * dim; + uint64_t baseOffset = KVOffset(b, hv, start + rowBegin, 0, dim); + uint64_t sourceOffset = sourceSequenceMajor ? VInputOffset(b, hv, start + rowBegin, 0) : baseOffset; + uint64_t sourceStride = sourceSequenceMajor ? HV_ * dim : dim; + + if constexpr (IsSameType::value) { + CopyRowsIn(matrixLocal, src, sourceOffset, rowCount, dim, sourceStride); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + } else { + LocalTensor matrixTyped = vecBuf_.Get()[typedOffset]; + CopyRowsIn(matrixTyped, src, sourceOffset, rowCount, dim, sourceStride); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + Cast(matrixLocal, matrixTyped, RoundMode::CAST_NONE, static_cast(elemCount)); + PipeBarrier(); + } + + uint8_t repeatStride = static_cast(dim * sizeof(float) / 32); + for (uint64_t col = 0; col < dim; col += vecElemsPerRepeat) { + uint64_t mask = dim - col; + if (mask > vecElemsPerRepeat) { + mask = vecElemsPerRepeat; + } + Mul(matrixLocal[col], matrixLocal[col], betaBrcb, mask, static_cast(rowCount), + {1, 1, 0, repeatStride, repeatStride, 1}); + } + PipeBarrier(); + + if constexpr (IsSameType::value) { + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(dst[baseOffset], matrixLocal, static_cast(elemCount)); + } else { + LocalTensor matrixTyped = vecBuf_.Get()[typedOffset]; + Cast(matrixTyped, matrixLocal, RoundMode::CAST_RINT, static_cast(elemCount)); + PipeBarrier(); + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + DataCopy(dst[baseOffset], matrixTyped, static_cast(elemCount)); + } + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + + __aicore__ inline void PrepareWuCubeInputs(uint64_t b, uint64_t hv, uint64_t start, uint64_t curT, + uint64_t subBlockIdx, uint64_t subBlockNum) + { + uint64_t rowsPerSubBlock = (curT + subBlockNum - 1) / subBlockNum; + uint64_t rowBegin = subBlockIdx * rowsPerSubBlock; + if (rowBegin >= curT) { + return; + } + uint64_t rowCount = curT - rowBegin; + if (rowCount > rowsPerSubBlock) { + rowCount = rowsPerSubBlock; + } + LocalTensor arena = vecBuf_.Get(); + LocalTensor betaLocal = arena; + LocalTensor betaBrcb = arena[KDA_SOLVE_BT]; + LocalTensor matrixLocal = arena[KDA_SOLVE_BT + 512]; + LoadBetaAsFloatRow(BetaOffset(b, hv, start + rowBegin), betaLocal, rowCount); + Brcb(betaBrcb, betaLocal, static_cast((rowCount + 7) / 8), {1, 8}); + PipeBarrier(); + ScaleRowsByBeta(w_, w_, b, hv, start, rowBegin, rowCount, K_, betaLocal, betaBrcb, matrixLocal); + ScaleRowsByBeta(v_, vNew_, b, hv, start, rowBegin, rowCount, V_, betaLocal, betaBrcb, matrixLocal, + inputSequenceMajor_); + } + + __aicore__ inline void FinalizePrepareIntermediates(uint64_t b, uint64_t hv, uint64_t start, uint64_t curT, + uint64_t subBlockIdx, uint64_t subBlockNum) + { + constexpr uint64_t tileRows = 32; + // Keep tail rows on the same AIV that owns their padded solve rows. Splitting by curT would + // move short-tail export to AIV1 while AIV0 is still writing the solved matrix. + const uint64_t rowBegin = (BT_ * subBlockIdx) / subBlockNum; + uint64_t rowEnd = (BT_ * (subBlockIdx + 1)) / subBlockNum; + if (rowEnd > curT) { + rowEnd = curT; + } + if (rowBegin >= rowEnd) { + return; + } + for (uint64_t tileRow = rowBegin; tileRow < rowEnd; tileRow += tileRows) { + const uint64_t rows = (rowEnd - tileRow) > tileRows ? tileRows : (rowEnd - tileRow); + const uint64_t matrixElems = rows * BT_; + const uint64_t qgElems = rows * K_; + LocalTensor arena = vecBuf_.Get(); + LocalTensor aqkLocal = arena; + LocalTensor akkLocal = arena[matrixElems]; + const uint64_t typedOffset = (2 * matrixElems + qgElems) * sizeof(float) / sizeof(T); + LocalTensor typedBase = vecBuf_.Get()[typedOffset]; + LocalTensor aqkTyped = typedBase; + LocalTensor akkTyped = typedBase[matrixElems]; + LocalTensor qgTyped = typedBase[2 * matrixElems]; + + CopyVectorIn(aqkLocal, aqk_, AOffset(b, hv, start + tileRow, 0), matrixElems); + CopyVectorIn(akkLocal, akk_, AOffset(b, hv, start + tileRow, 0), matrixElems); + CopyVectorIn(qgTyped, qg_, KVOffset(b, hv, start + tileRow, 0, K_), qgElems); + SetFlag(mte2ToVEvent_); + WaitFlag(mte2ToVEvent_); + + Muls(aqkLocal, aqkLocal, scale_, static_cast(matrixElems)); + PipeBarrier(); + ClampFp32ToOutputType(aqkLocal, static_cast(matrixElems)); + ClampFp32ToOutputType(akkLocal, static_cast(matrixElems)); + Cast(aqkTyped, aqkLocal, RoundMode::CAST_RINT, static_cast(matrixElems)); + Cast(akkTyped, akkLocal, RoundMode::CAST_RINT, static_cast(matrixElems)); + PipeBarrier(); + + SetFlag(vToMte3Event_); + WaitFlag(vToMte3Event_); + CopyVectorOut(o_, AOffset(b, hv, start + tileRow, 0), aqkTyped, matrixElems); + CopyVectorOut(u_, AOffset(b, hv, start + tileRow, 0), akkTyped, matrixElems); + CopyVectorOut(kg_, KVOffset(b, hv, start + tileRow, 0, K_), qgTyped, qgElems); + SetFlag(mte3ToMte2Events_[0]); + WaitFlag(mte3ToMte2Events_[0]); + SetFlag(mte3ToVEvent_); + WaitFlag(mte3ToVEvent_); + } + } + + __aicore__ inline bool ResolveFlatChunk(uint64_t task, uint64_t &seq, uint64_t &b, uint64_t &h, uint64_t &hv, + uint64_t &chunkIdx, uint64_t &start, uint64_t &end) + { + hv = task % HV_; + uint64_t flatChunk = task / HV_; + if (!isVarLen_) { + seq = flatChunk / NT_; + b = seq; + chunkIdx = flatChunk % NT_; + start = chunkIdx * BT_; + end = start + BT_; + if (end > T_) { + end = T_; + } + } else { + if (!KdaVarlen::ResolveChunkRange(cuSeqlensAddr_, chunkIndicesAddr_, N_, T_, BT_, flatChunk, seq, start, + end)) { + return false; + } + b = 0; + chunkIdx = flatChunk; + } + h = hv / (HV_ / H_); + return start < end; + } + + __aicore__ inline void ProcessChunkPreAiv(uint64_t b, uint64_t h, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t end, uint64_t subBlockIdx, uint64_t subBlockNum) + { + if constexpr (IsSameType::value) { + ProcessChunkPreAivFp32(b, h, hv, chunkIdx, start, end, subBlockIdx, subBlockNum); + } + } + + template + __aicore__ inline void JoinAivMte3() + { + if constexpr (CORE_TYPE == AscendC::AIV) { + } + } + + template + __aicore__ inline void RunAicAfterBothAivReady(uint64_t subBlockIdx, uint64_t subBlockNum) + { + if constexpr (CORE_TYPE == AscendC::AIV) { + (void)subBlockIdx; + (void)subBlockNum; + JoinAivMte3(); + if constexpr (SAFE_GATE) { + Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(syncReadyFlag_); + Catlass::Arch::CrossCoreWaitFlag(syncDoneFlag_); + } else { + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_MTE3>(mchSyncReadyFlag_); + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE2>(mchSyncDoneFlag_); + } + } + } + + template + __aicore__ inline void SignalAicSolveReady() + { + if constexpr (CORE_TYPE == AscendC::AIV) { + JoinAivMte3(); + Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(syncReadyFlag_); + } + } + + template + __aicore__ inline void WaitAicSolveDone() + { + if constexpr (CORE_TYPE == AscendC::AIV) { + Catlass::Arch::CrossCoreWaitFlag(syncDoneFlag_); + } + } + + __aicore__ inline void ProcessChunkPreAivFp32(uint64_t b, uint64_t h, uint64_t hv, uint64_t chunkIdx, + uint64_t start, uint64_t end, uint64_t subBlockIdx, + uint64_t subBlockNum, bool deferSafeSolve = false, + bool waitPendingSafeSolve = false, uint64_t scoreLane = 0, + bool pairHeads = false) + { + uint64_t curT = end - start; + if (curT == 0) { + return; + } + if constexpr (IsSameType::value) { + return; + } + + if (K_ < 16) { + return; + } + bool usePostWuCube = UsePostWuCube(curT); + bool useAkkCubeSolve = UseAkkCubeSolve(curT); + uint64_t solveRowBegin = 0; + uint64_t solveRowEnd = 0; + GetSolveRowRange(BT_, subBlockIdx, subBlockNum, solveRowBegin, solveRowEnd); + uint64_t scoreBlockSize = ScoreRefBlockSize(); + uint64_t scoreBlockCount = (curT + scoreBlockSize - 1) / scoreBlockSize; + uint64_t pipelineBlockCount = + (scoreBlockCount + KDA_SCORE_QUEUE_DEPTH - 1) / KDA_SCORE_QUEUE_DEPTH * KDA_SCORE_QUEUE_DEPTH; + for (uint64_t block = 0; block < pipelineBlockCount; ++block) { + if (block < scoreBlockCount) { + uint64_t rowBegin = block * scoreBlockSize; + uint64_t rowCount = ScoreRowBlockCount(curT, rowBegin); + uint64_t refToken = ScoreRefToken(start, curT, rowBegin, rowCount); + uint64_t scoreSlot = ScoreScratchSlot(block % KDA_SCORE_QUEUE_DEPTH, scoreLane, pairHeads); + PrepareGateProducts(b, h, hv, start, curT, subBlockIdx, subBlockNum, true, refToken, + rowBegin + rowCount, true, scoreSlot, rowBegin, rowCount); + } + JoinAivMte3(); + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_MTE3>(scoreReadyFlag_); + if (block > 0) { + if constexpr (SAFE_GATE) { + if (waitPendingSafeSolve) { + WaitAicSolveDone(); + waitPendingSafeSolve = false; + } + } + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE2>(scoreDoneFlag_); + } + } + bool fusedScoreWriteback = false; + if (!fusedScoreWriteback) { + // The final score MMAD only consumes scoreWorkspace_. Run the + // independent gate writeback while AIC drains its MMAD/Fixpipe path. + PrepareGateProducts(b, h, hv, start, curT, subBlockIdx, subBlockNum); + } + if (pipelineBlockCount > 0) { + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE2>(scoreDoneFlag_); + } + if constexpr (SAFE_GATE) { + if (waitPendingSafeSolve) { + WaitAicSolveDone(); + waitPendingSafeSolve = false; + } + } + if (useAkkCubeSolve) { + bool fullChunk = curT == BT_; + if constexpr (SAFE_GATE) { + if (pairHeads) { + for (uint64_t rowPart = 0; rowPart < KDA_SCORE_LANES; ++rowPart) { + uint64_t pairRowBegin = 0; + uint64_t pairRowEnd = 0; + GetSolveRowRange(BT_, rowPart, KDA_SCORE_LANES, pairRowBegin, pairRowEnd); + PrepareAqkAkkSolveInputRows(b, hv, chunkIdx, start, curT, pairRowBegin, pairRowEnd, false, + false); + } + } else { + PrepareAqkAkkSolveInputRows(b, hv, chunkIdx, start, curT, solveRowBegin, solveRowEnd, false, false); + } + if (deferSafeSolve) { + SignalAicSolveReady(); + return; + } + RunAicAfterBothAivReady(subBlockIdx, subBlockNum); + StoreSolveXRowsToAkk(b, hv, chunkIdx, start, curT, solveRowBegin, solveRowEnd); + } else { + PrepareAqkAkkSolveInputRows(b, hv, chunkIdx, start, curT, solveRowBegin, solveRowEnd, fullChunk, + !fullChunk); + uint32_t solveIters = KDA_SOLVE_DIAG_MCH_ITERS; + RunAicAfterBothAivReady(subBlockIdx, subBlockNum); + for (uint32_t iter = 0; iter < solveIters; ++iter) { + AddSolveTmpToXDiagRows(b, hv, chunkIdx, start, solveRowBegin, solveRowEnd, + fullChunk && iter + 1 == solveIters); + RunAicAfterBothAivReady(subBlockIdx, subBlockNum); + } + if (!fullChunk) { + StoreSolveXRowsToAkk(b, hv, chunkIdx, start, curT, solveRowBegin, solveRowEnd); + } + } + } + // Host validation guarantees every accepted shape has enough workspace for this cube path. + PrepareWuCubeInputs(b, hv, start, curT, subBlockIdx, subBlockNum); + FinalizePrepareIntermediates(b, hv, start, curT, subBlockIdx, subBlockNum); + } + + __aicore__ inline void FinishDeferredSafeChunk(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t end, uint64_t subBlockIdx, uint64_t subBlockNum) + { + uint64_t curT = end - start; + uint64_t solveRowBegin = 0; + uint64_t solveRowEnd = 0; + GetSolveRowRange(BT_, subBlockIdx, subBlockNum, solveRowBegin, solveRowEnd); + StoreSolveXRowsToAkk(b, hv, chunkIdx, start, curT, solveRowBegin, solveRowEnd); + PrepareWuCubeInputs(b, hv, start, curT, subBlockIdx, subBlockNum); + FinalizePrepareIntermediates(b, hv, start, curT, subBlockIdx, subBlockNum); + } + + __aicore__ inline void FinishDeferredSafeChunkPair(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t end) + { + FinishDeferredSafeChunk(b, hv, chunkIdx, start, end, 0, 1); + } + + __aicore__ inline void ProcessChunkPreAic(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, uint64_t end) + { + if constexpr (IsSameType::value) { + ProcessChunkPreAicFp32(b, hv, chunkIdx, start, end); + } + } + + __aicore__ inline void ProcessChunkPreAicFp32(uint64_t b, uint64_t hv, uint64_t chunkIdx, uint64_t start, + uint64_t end) + { + uint64_t curT = end - start; + if (curT == 0 || K_ < 16) { + return; + } + uint64_t scoreBlockSize = ScoreRefBlockSize(); + uint64_t scoreBlockCount = (curT + scoreBlockSize - 1) / scoreBlockSize; + uint64_t pipelineBlockCount = + (scoreBlockCount + KDA_SCORE_QUEUE_DEPTH - 1) / KDA_SCORE_QUEUE_DEPTH * KDA_SCORE_QUEUE_DEPTH; + for (uint64_t block = 0; block < pipelineBlockCount; ++block) { + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_FIX>(scoreReadyFlag_); + if (block < scoreBlockCount) { + uint64_t rowBegin = block * scoreBlockSize; + uint64_t rowCount = ScoreRowBlockCount(curT, rowBegin); + ComputeRawAqkAkkCubeBlock(b, hv, start, curT, rowBegin, rowCount, true, block % KDA_SCORE_QUEUE_DEPTH, + rowBegin + rowCount); + } + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(scoreDoneFlag_); + } + bool usePostWuCube = UsePostWuCube(curT); + bool useAkkCubeSolve = UseAkkCubeSolve(curT); + if (useAkkCubeSolve) { + if constexpr (SAFE_GATE) { + Catlass::Arch::CrossCoreWaitFlag(syncReadyFlag_); + ComputeAkkMergeCubeWorkspace(b, hv, chunkIdx); + Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(syncDoneFlag_); + } else { + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_FIX>(mchSyncReadyFlag_); + if (curT == BT_) { + ComputeAkkInverseMchFull(b, hv, chunkIdx, start); + } else { + ComputeAkkInverseMchTail(b, hv, chunkIdx, start, curT); + } + } + } + (void)usePostWuCube; + (void)chunkIdx; + } + + __aicore__ inline void ProcessChunkPreAicHeadPairFp32(uint64_t b, uint64_t hvBase, uint64_t chunkIdx, + uint64_t start, uint64_t end, uint64_t localTaskIdx) + { + uint64_t curT = end - start; + if (curT == 0 || K_ < 16) { + return; + } + uint64_t scoreBlockSize = ScoreRefBlockSize(); + uint64_t scoreBlockCount = (curT + scoreBlockSize - 1) / scoreBlockSize; + uint64_t pipelineBlockCount = + (scoreBlockCount + KDA_SCORE_QUEUE_DEPTH - 1) / KDA_SCORE_QUEUE_DEPTH * KDA_SCORE_QUEUE_DEPTH; + for (uint64_t block = 0; block < pipelineBlockCount; ++block) { + Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_FIX>(scoreReadyFlag_); + if (block < scoreBlockCount) { + uint64_t rowBegin = block * scoreBlockSize; + uint64_t rowCount = ScoreRowBlockCount(curT, rowBegin); + for (uint64_t lane = 0; lane < KDA_SCORE_LANES; ++lane) { + uint64_t hv = hvBase + lane; + uint64_t scoreSlot = ScoreScratchSlot(block % KDA_SCORE_QUEUE_DEPTH, lane, true); + ComputeRawAqkAkkCubeBlock(b, hv, start, curT, rowBegin, rowCount, true, scoreSlot, + rowBegin + rowCount); + } + } + Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(scoreDoneFlag_); + } + + if (UseAkkCubeSolve(curT)) { + Catlass::Arch::CrossCoreWaitFlag(syncReadyFlag_); + for (uint64_t lane = 0; lane < KDA_SCORE_LANES; ++lane) { + activeSolveSlot_ = + (localTaskIdx % (KDA_SOLVE_PIPELINE_DEPTH / KDA_SCORE_LANES)) * KDA_SCORE_LANES + lane; + ComputeAkkMergeCubeWorkspace(b, hvBase + lane, chunkIdx); + } + Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(syncDoneFlag_); + } + } + + __aicore__ inline bool ResolveFlatChunkForHv(uint64_t flatChunk, uint64_t hv, uint64_t &seq, uint64_t &b, + uint64_t &h, uint64_t &chunkIdx, uint64_t &start, uint64_t &end) + { + if (!isVarLen_) { + seq = flatChunk / NT_; + b = seq; + chunkIdx = flatChunk % NT_; + start = chunkIdx * BT_; + end = start + BT_; + if (end > T_) { + end = T_; + } + } else { + if (!KdaVarlen::ResolveChunkRange(cuSeqlensAddr_, chunkIndicesAddr_, N_, T_, BT_, flatChunk, seq, start, + end)) { + return false; + } + b = 0; + chunkIdx = flatChunk; + } + h = hv / (HV_ / H_); + return start < end; + } + + __aicore__ inline void ProcessPreAivHeadPair() + { + const uint64_t subBlockIdx = static_cast(GetSubBlockIdx()); + const uint64_t subBlockNum = static_cast(GetSubBlockNum()); + const uint64_t coreNum = usedCoreNum_; + const uint64_t coreIdx = static_cast(GetBlockIdx()) / subBlockNum; + const uint64_t chunkCount = isVarLen_ ? NT_ : B_ * NT_; + const uint64_t headWindows = HV_ / KDA_SCORE_LANES; + const uint64_t taskNum = chunkCount * headWindows; + bool pendingValid = false; + uint64_t pendingB = 0; + uint64_t pendingHv = 0; + uint64_t pendingChunkIdx = 0; + uint64_t pendingStart = 0; + uint64_t pendingEnd = 0; + uint64_t pendingSlot = 0; + uint64_t localTaskIdx = 0; + + for (uint64_t task = coreIdx; task < taskNum; task += coreNum, ++localTaskIdx) { + uint64_t flatChunk = task / headWindows; + uint64_t hv = (task % headWindows) * KDA_SCORE_LANES + subBlockIdx; + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (!ResolveFlatChunkForHv(flatChunk, hv, seq, b, h, chunkIdx, start, end)) { + continue; + } + (void)seq; + uint64_t currentSlot = + (localTaskIdx % (KDA_SOLVE_PIPELINE_DEPTH / KDA_SCORE_LANES)) * KDA_SCORE_LANES + subBlockIdx; + activeSolveSlot_ = currentSlot; + bool deferSolve = UseAkkCubeSolve(end - start); + ProcessChunkPreAivFp32(b, h, hv, chunkIdx, start, end, 0, 1, deferSolve, pendingValid, subBlockIdx, true); + if (pendingValid) { + activeSolveSlot_ = pendingSlot; + FinishDeferredSafeChunkPair(pendingB, pendingHv, pendingChunkIdx, pendingStart, pendingEnd); + } + pendingValid = deferSolve; + if (pendingValid) { + pendingB = b; + pendingHv = hv; + pendingChunkIdx = chunkIdx; + pendingStart = start; + pendingEnd = end; + pendingSlot = currentSlot; + } + } + if (pendingValid) { + WaitAicSolveDone(); + activeSolveSlot_ = pendingSlot; + FinishDeferredSafeChunkPair(pendingB, pendingHv, pendingChunkIdx, pendingStart, pendingEnd); + } + } + + __aicore__ inline void ProcessPreAicHeadPair() + { + const uint64_t chunkCount = isVarLen_ ? NT_ : B_ * NT_; + const uint64_t headWindows = HV_ / KDA_SCORE_LANES; + const uint64_t taskNum = chunkCount * headWindows; + const uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + uint64_t localTaskIdx = 0; + for (uint64_t task = GetBlockIdx(); task < taskNum; task += coreNum, ++localTaskIdx) { + uint64_t flatChunk = task / headWindows; + uint64_t hvBase = (task % headWindows) * KDA_SCORE_LANES; + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveFlatChunkForHv(flatChunk, hvBase, seq, b, h, chunkIdx, start, end)) { + (void)seq; + (void)h; + ProcessChunkPreAicHeadPairFp32(b, hvBase, chunkIdx, start, end, localTaskIdx); + } + } + } + + __aicore__ inline void ProcessPreAiv() + { + if constexpr (IsSameType::value) { + isAivOnly_ = true; + } + uint64_t subBlockNum = isAivOnly_ ? 1 : static_cast(GetSubBlockNum()); + if (subBlockNum == 0) { + return; + } + uint64_t subBlockIdx = isAivOnly_ ? 0 : static_cast(GetSubBlockIdx()); + uint64_t coreNum = isAivOnly_ ? static_cast(GetBlockNum()) : usedCoreNum_; + uint64_t coreIdx = + isAivOnly_ ? static_cast(GetBlockIdx()) : static_cast(GetBlockIdx()) / subBlockNum; + uint64_t taskNum = static_cast((isVarLen_ ? NT_ : B_ * NT_) * HV_); + if constexpr (SAFE_GATE && !IsSameType::value) { + bool pendingValid = false; + uint64_t pendingB = 0; + uint64_t pendingHv = 0; + uint64_t pendingChunkIdx = 0; + uint64_t pendingStart = 0; + uint64_t pendingEnd = 0; + uint64_t pendingSlot = 0; + uint64_t localTaskIdx = 0; + for (uint64_t task = coreIdx; task < taskNum; task += coreNum, ++localTaskIdx) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (!ResolveFlatChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + continue; + } + (void)seq; + uint64_t currentSlot = localTaskIdx % KDA_SOLVE_PIPELINE_DEPTH; + activeSolveSlot_ = currentSlot; + bool deferSolve = UseAkkCubeSolve(end - start); + if (!deferSolve && pendingValid) { + WaitAicSolveDone(); + activeSolveSlot_ = pendingSlot; + FinishDeferredSafeChunk(pendingB, pendingHv, pendingChunkIdx, pendingStart, pendingEnd, subBlockIdx, + subBlockNum); + pendingValid = false; + activeSolveSlot_ = currentSlot; + } + ProcessChunkPreAivFp32(b, h, hv, chunkIdx, start, end, subBlockIdx, subBlockNum, deferSolve, + pendingValid); + if (pendingValid) { + activeSolveSlot_ = pendingSlot; + FinishDeferredSafeChunk(pendingB, pendingHv, pendingChunkIdx, pendingStart, pendingEnd, subBlockIdx, + subBlockNum); + } + pendingValid = deferSolve; + if (pendingValid) { + pendingB = b; + pendingHv = hv; + pendingChunkIdx = chunkIdx; + pendingStart = start; + pendingEnd = end; + pendingSlot = currentSlot; + } + } + if (pendingValid) { + WaitAicSolveDone(); + activeSolveSlot_ = pendingSlot; + FinishDeferredSafeChunk(pendingB, pendingHv, pendingChunkIdx, pendingStart, pendingEnd, subBlockIdx, + subBlockNum); + } + return; + } + for (uint64_t task = coreIdx; task < taskNum; task += coreNum) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveFlatChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + (void)seq; + ProcessChunkPreAiv(b, h, hv, chunkIdx, start, end, subBlockIdx, subBlockNum); + } + } + } + + __aicore__ inline void ProcessPreAic() + { + if constexpr (IsSameType::value) { + return; + } + uint64_t taskNum = static_cast((isVarLen_ ? NT_ : B_ * NT_) * HV_); + uint64_t coreNum = usedCoreNum_ == 0 ? 1 : usedCoreNum_; + uint64_t localTaskIdx = 0; + for (uint64_t task = GetBlockIdx(); task < taskNum; task += coreNum, ++localTaskIdx) { + uint64_t seq = 0; + uint64_t b = 0; + uint64_t h = 0; + uint64_t hv = 0; + uint64_t chunkIdx = 0; + uint64_t start = 0; + uint64_t end = 0; + if (ResolveFlatChunk(task, seq, b, h, hv, chunkIdx, start, end)) { + if constexpr (SAFE_GATE) { + activeSolveSlot_ = localTaskIdx % KDA_SOLVE_PIPELINE_DEPTH; + } + (void)seq; + (void)h; + ProcessChunkPreAic(b, hv, chunkIdx, start, end); + } + } + } + +private: + GlobalTensor q_; + GlobalTensor k_; + GlobalTensor v_; + GlobalTensor gk_; + GlobalTensor betaFp32_; + GlobalTensor betaBf16_; + GlobalTensor initialState_; + GlobalTensor o_; + GlobalTensor finalState_; + GlobalTensor aqk_; + GlobalTensor akk_; + GlobalTensor w_; + GlobalTensor u_; + GlobalTensor qg_; + GlobalTensor kg_; + GlobalTensor vNew_; + GlobalTensor h_; + GlobalTensor preparedQG_; + GlobalTensor preparedAqk_; + GlobalTensor propagatedVNew_; + GlobalTensor propagatedH_; + GlobalTensor solveWorkspace_; + GlobalTensor scoreWorkspace_; + TPipe *pipe_ = nullptr; + TBuf exp2Buf_; + TBuf vecBuf_; + TBuf gateWritebackBuf_; + TEventID mte2ToVEvent_ = 0; + TEventID vToMte2Event_ = 0; + TEventID vToMte3Event_ = 0; + TEventID mte3ToVEvent_ = 0; + TEventID mte2ToMte3Event_ = 0; + TEventID mte3ToMte2Events_[KDA_GATE_PIPELINE_DEPTH] = {0, 0, 0}; + bool vectorEventsAllocated_ = false; + Catlass::Arch::CrossCoreFlagWithReverse scoreReadyFlag_{KDA_SCORE_READY_FLAG0, + KDA_SCORE_READY_FLAG1}; + Catlass::Arch::CrossCoreFlagWithReverse scoreDoneFlag_{KDA_SCORE_DONE_FLAG0, + KDA_SCORE_DONE_FLAG1}; + // Solve has one outstanding task per core. Reuse the primary score IDs as an ordered token stream; + // score credits remain on the reverse IDs, so no additional hardware flag IDs are consumed. + Catlass::Arch::CrossCoreFlag syncReadyFlag_{KDA_SOLVE_READY_FLAG}; + Catlass::Arch::CrossCoreFlag syncDoneFlag_{KDA_SOLVE_DONE_FLAG}; + Catlass::Arch::CrossCoreFlagWithReverse mchSyncReadyFlag_{KDA_SCORE_READY_FLAG0, + KDA_SCORE_READY_FLAG1}; + Catlass::Arch::CrossCoreFlagWithReverse mchSyncDoneFlag_{KDA_SCORE_DONE_FLAG0, + KDA_SCORE_DONE_FLAG1}; + uint64_t B_ = 0; + uint64_t N_ = 0; + uint64_t H_ = 0; + uint64_t HV_ = 0; + uint64_t T_ = 0; + uint64_t K_ = 0; + uint64_t V_ = 0; + uint64_t BT_ = 0; + uint64_t NT_ = 0; + float scale_ = 1.0f; + bool hasInitial_ = false; + bool isVarLen_ = false; + bool inputSequenceMajor_ = false; + bool betaIsFp32_ = true; + bool isAivOnly_ = false; + uint64_t usedCoreNum_ = 1; + uint64_t solveCoreIdx_ = 0; + uint64_t activeSolveSlot_ = 0; + __gm__ int64_t *chunkIndicesAddr_ = nullptr; + __gm__ int64_t *cuSeqlensAddr_ = nullptr; +}; +} // namespace + +template +__aicore__ inline void RunChunkKdaPrepare(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR beta, + GM_ADDR initialState, GM_ADDR cuSeqlens, GM_ADDR chunkIndices, GM_ADDR aqk, + GM_ADDR akk, GM_ADDR qg, GM_ADDR qgScaled, GM_ADDR wSeed, GM_ADDR uSeed, + GM_ADDR userWorkspace, const TilingData &tiling, TPipe &pipe) +{ + GM_ADDR aqkFp32 = userWorkspace + tiling.prepareAqkFp32Offset; + GM_ADDR akkFp32 = userWorkspace + tiling.prepareAkkFp32Offset; + GM_ADDR prepareScratch = userWorkspace + tiling.prepareScratchOffset; + + if ASCEND_IS_AIC { + ChunkKdaFwdPrepareKernel op; + op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, nullptr, nullptr, nullptr, nullptr, aqk, + userWorkspace, aqkFp32, akkFp32, wSeed, akk, qg, qgScaled, uSeed, userWorkspace, prepareScratch, tiling, + &pipe, false); + op.ProcessAic(); + } + if ASCEND_IS_AIV { + ChunkKdaFwdPrepareKernel op; + op.Init(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, nullptr, nullptr, nullptr, nullptr, aqk, + userWorkspace, aqkFp32, akkFp32, wSeed, akk, qg, qgScaled, uSeed, userWorkspace, prepareScratch, tiling, + &pipe); + op.ProcessAiv(); + } +} + +template +__aicore__ inline void RunChunkKdaPrepare(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR gk, GM_ADDR, GM_ADDR, GM_ADDR, + GM_ADDR beta, GM_ADDR initialState, GM_ADDR cuSeqlens, GM_ADDR chunkIndices, + GM_ADDR aqk, GM_ADDR akk, GM_ADDR qg, GM_ADDR qgScaled, GM_ADDR wSeed, + GM_ADDR uSeed, GM_ADDR, GM_ADDR userWorkspace, const TilingData &tiling, + TPipe &pipe, bool = true) +{ + RunChunkKdaPrepare(q, k, v, gk, beta, initialState, cuSeqlens, chunkIndices, aqk, akk, qg, + qgScaled, wSeed, uSeed, userWorkspace, tiling, pipe); +} + +} // namespace KdaPrepare diff --git a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_tiling_data.h b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_tiling_data.h new file mode 100644 index 0000000000..a3d284b49e --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_tiling_data.h @@ -0,0 +1,99 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +/*! + * \file chunk_kda_fwd_tiling_data.h + * \brief direct-launch tiling data of chunk_kda_fwd + * + * Plain-struct mirror of the aclnn ChunkKdaFwdTilingData (register/tilingdata_base.h + * TILING_DATA_FIELD_DEF layout). Field order/type must be kept in sync with the host + * side filling routine so that the packed blob copied GM->stack in the kernel entry + * matches exactly what the host serializes. + */ + +#ifndef CHUNK_KDA_FWD_TILING_DATA_H +#define CHUNK_KDA_FWD_TILING_DATA_H + +#include + +namespace ChunkKdaFwd { + +#pragma pack(push, 8) +struct ChunkKdaFwdTilingData { + int64_t batch; + int64_t seqNum; + int64_t qHeadNum; + int64_t vHeadNum; + int64_t seqlen; + int64_t kHeadDim; + int64_t vHeadDim; + int64_t chunkSize; + int64_t totalChunks; + int64_t inputRank; + + float scale; + float lowerBound; + + bool hasInitialState; + bool isVarLen; + bool safeGate; + bool inputSequenceMajor; + bool useGateInKernel; + bool hasALog; + bool hasDtBias; + bool computeGateInPrepare; + bool fusePostWu; + bool fusePostWuIntoFwdH; + bool useDenseFwdH; + bool storeFinalState; + bool storeGk; + bool storeW; + bool storeU; + bool storeQG; + bool storeKg; + bool storeVNew; + bool storeH; + + int64_t gateDataType; + int64_t gateUsedCoreNum; + int64_t prepareUsedCoreNum; + int64_t postWuUsedCoreNum; + int64_t outputUsedCoreNum; + + int64_t gkStorageOffset; + int64_t finalStateStorageOffset; + int64_t wStorageOffset; + int64_t uStorageOffset; + int64_t qgStorageOffset; + int64_t kgStorageOffset; + int64_t vNewStorageOffset; + int64_t hStorageOffset; + int64_t qgScaledOffset; + int64_t prepareAqkFp32Offset; + int64_t prepareAkkFp32Offset; + int64_t prepareScratchOffset; + int64_t postWuScratchOffset; + int64_t outputScratchOffset; + int64_t fwdHWorkspaceBaseOffset; + int64_t vWorkspaceOffset; + int64_t vUpdateWorkspaceOffset; + int64_t kDecayWorkspaceOffset; + int64_t hWorkspaceOffset; + int64_t numSeqWorkspaceOffset; + int64_t numChunksWorkspaceOffset; + + // direct-launch only: beta storage dtype (0 = fp32, 1 = bf16) + int64_t betaDataType; +}; +#pragma pack(pop) + +} // namespace ChunkKdaFwd + +#endif // CHUNK_KDA_FWD_TILING_DATA_H diff --git a/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_varlen.h b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_varlen.h new file mode 100644 index 0000000000..69fabfb14f --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/chunk_kda_fwd_varlen.h @@ -0,0 +1,78 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#pragma once + +#include "kernel_operator.h" + +namespace KdaVarlen { + +__aicore__ inline bool ResolveChunkRange( + const __gm__ int64_t *cuSeqlens, const __gm__ int64_t *chunkIndices, + uint64_t seqNum, uint64_t totalTokens, uint64_t chunkSize, + uint64_t flatChunk, uint64_t &seq, uint64_t &start, uint64_t &end) +{ + if (cuSeqlens == nullptr || chunkSize == 0) { + return false; + } + + int64_t seqValue = -1; + int64_t localChunkValue = -1; + if (chunkIndices != nullptr) { + const uint64_t metadataOffset = flatChunk * 2; + seqValue = chunkIndices[metadataOffset]; + localChunkValue = chunkIndices[metadataOffset + 1]; + } else { + uint64_t chunkPrefix = 0; + for (uint64_t seqIdx = 0; seqIdx < seqNum; ++seqIdx) { + const int64_t seqStart = cuSeqlens[seqIdx]; + const int64_t seqEnd = cuSeqlens[seqIdx + 1]; + if (seqStart < 0 || seqEnd < seqStart) { + return false; + } + const uint64_t seqLength = static_cast(seqEnd - seqStart); + const uint64_t seqChunks = (seqLength + chunkSize - 1) / chunkSize; + if (flatChunk < chunkPrefix + seqChunks) { + seqValue = static_cast(seqIdx); + localChunkValue = static_cast(flatChunk - chunkPrefix); + break; + } + chunkPrefix += seqChunks; + } + } + + if (seqValue < 0 || localChunkValue < 0 || + static_cast(seqValue) >= seqNum) { + return false; + } + const int64_t seqStartValue = cuSeqlens[seqValue]; + const int64_t seqEndValue = cuSeqlens[seqValue + 1]; + if (seqStartValue < 0 || seqEndValue < seqStartValue) { + return false; + } + const uint64_t seqStart = static_cast(seqStartValue); + const uint64_t seqEnd = static_cast(seqEndValue); + const uint64_t localChunk = static_cast(localChunkValue); + start = seqStart + localChunk * chunkSize; + if (start >= seqEnd || start >= totalTokens) { + return false; + } + end = start + chunkSize; + if (end > seqEnd) { + end = seqEnd; + } + if (end > totalTokens) { + return false; + } + seq = static_cast(seqValue); + return start < end; +} + +} // namespace KdaVarlen diff --git a/csrc/chunk_kda_fwd/op_kernel/kda_gate_cumsum_kernel.h b/csrc/chunk_kda_fwd/op_kernel/kda_gate_cumsum_kernel.h new file mode 100644 index 0000000000..886322d9db --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/kda_gate_cumsum_kernel.h @@ -0,0 +1,749 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#pragma once + +#include "kernel_operator.h" +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +#ifndef FLA_NPU_REGBASE_HPP_INCLUDED +#define FLA_NPU_REGBASE_HPP_INCLUDED +#include "kernel_utils/vector/regbase.hpp" +#endif +#endif + +namespace KdaGateCumsum { + +using namespace AscendC; + +constexpr float RCP_LN2 = 1.4426950408889634f; +constexpr uint32_t GATE_ROW_ELEMENTS = 256; +constexpr uint32_t GATE_PIPELINE_DEPTH = 2; +constexpr uint32_t GATE_BULK_ROWS = 64; +constexpr uint32_t GATE_BULK_COLS = 128; +constexpr uint32_t GATE_BULK_ELEMENTS = GATE_BULK_ROWS * GATE_BULK_COLS; + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 +static __simd_vf__ inline void AccumulateGateRowRegbase(__ubuf__ float *input, __ubuf__ float *acc, + __ubuf__ float *output, uint16_t count) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t FLOAT_ELEMENTS_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(float); + constexpr uint16_t ELEMENTS_PER_PAIR = 2 * FLOAT_ELEMENTS_PER_REG; + + for (uint16_t offset = 0; offset < count; offset += ELEMENTS_PER_PAIR) { + RegTensor inputZeroReg; + RegTensor inputOneReg; + RegTensor accZeroReg; + RegTensor accOneReg; + MaskReg floatMask = CreateMask(); + + LoadAlign(inputZeroReg, input + offset); + LoadAlign(inputOneReg, input + offset + FLOAT_ELEMENTS_PER_REG); + LoadAlign(accZeroReg, acc + offset); + LoadAlign(accOneReg, acc + offset + FLOAT_ELEMENTS_PER_REG); + + Muls(inputZeroReg, inputZeroReg, RCP_LN2, floatMask); + Muls(inputOneReg, inputOneReg, RCP_LN2, floatMask); + Add(accZeroReg, accZeroReg, inputZeroReg, floatMask); + Add(accOneReg, accOneReg, inputOneReg, floatMask); + + StoreAlign(acc + offset, accZeroReg, floatMask); + StoreAlign(acc + offset + FLOAT_ELEMENTS_PER_REG, accOneReg, floatMask); + StoreAlign(output + offset, accZeroReg, floatMask); + StoreAlign(output + offset + FLOAT_ELEMENTS_PER_REG, accOneReg, floatMask); + } +} + +static __simd_vf__ inline void AccumulateGateChunk128Regbase(__ubuf__ float *input, + __ubuf__ float *output, + uint16_t rows) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t FLOAT_ELEMENTS_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(float); + constexpr uint16_t ROW_ELEMENTS = 2 * FLOAT_ELEMENTS_PER_REG; + + MaskReg floatMask = CreateMask(); + RegTensor accZeroReg; + RegTensor accOneReg; + Duplicate(accZeroReg, 0.0f, floatMask); + Duplicate(accOneReg, 0.0f, floatMask); + for (uint16_t row = 0; row < rows; ++row) { + uint32_t rowOffset = static_cast(row) * ROW_ELEMENTS; + RegTensor inputZeroReg; + RegTensor inputOneReg; + LoadAlign(inputZeroReg, input + rowOffset); + LoadAlign( + inputOneReg, input + rowOffset + FLOAT_ELEMENTS_PER_REG); + Muls(inputZeroReg, inputZeroReg, RCP_LN2, floatMask); + Muls(inputOneReg, inputOneReg, RCP_LN2, floatMask); + Add(accZeroReg, accZeroReg, inputZeroReg, floatMask); + Add(accOneReg, accOneReg, inputOneReg, floatMask); + StoreAlign(output + rowOffset, accZeroReg, floatMask); + StoreAlign(output + rowOffset + FLOAT_ELEMENTS_PER_REG, accOneReg, floatMask); + } +} + +template +static __simd_vf__ inline void AccumulateSafeGateChunk128Regbase( + __ubuf__ float *input, __ubuf__ float *bias, __ubuf__ float *output, + uint16_t rows, float expA, float lowerBound) +{ + using namespace AscendC::MicroAPI; + constexpr uint16_t FLOAT_ELEMENTS_PER_REG = AscendC::VECTOR_REG_WIDTH / sizeof(float); + constexpr uint16_t ROW_ELEMENTS = 2 * FLOAT_ELEMENTS_PER_REG; + + MaskReg floatMask = CreateMask(); + RegTensor accZeroReg; + RegTensor accOneReg; + RegTensor oneZeroReg; + RegTensor oneOneReg; + RegTensor biasZeroReg; + RegTensor biasOneReg; + Duplicate(accZeroReg, 0.0f, floatMask); + Duplicate(accOneReg, 0.0f, floatMask); + Duplicate(oneZeroReg, 1.0f, floatMask); + Duplicate(oneOneReg, 1.0f, floatMask); + if constexpr (HAS_BIAS) { + LoadAlign(biasZeroReg, bias); + LoadAlign(biasOneReg, bias + FLOAT_ELEMENTS_PER_REG); + } + + const float gateScale = lowerBound * RCP_LN2; + for (uint16_t row = 0; row < rows; ++row) { + uint32_t rowOffset = static_cast(row) * ROW_ELEMENTS; + RegTensor gateZeroReg; + RegTensor gateOneReg; + RegTensor sigmoidZeroReg; + RegTensor sigmoidOneReg; + LoadAlign(gateZeroReg, input + rowOffset); + LoadAlign( + gateOneReg, input + rowOffset + FLOAT_ELEMENTS_PER_REG); + if constexpr (HAS_BIAS) { + Add(gateZeroReg, gateZeroReg, biasZeroReg, floatMask); + Add(gateOneReg, gateOneReg, biasOneReg, floatMask); + } + Muls(gateZeroReg, gateZeroReg, -expA, floatMask); + Muls(gateOneReg, gateOneReg, -expA, floatMask); + Exp(gateZeroReg, gateZeroReg, floatMask); + Exp(gateOneReg, gateOneReg, floatMask); + Adds(gateZeroReg, gateZeroReg, 1.0f, floatMask); + Adds(gateOneReg, gateOneReg, 1.0f, floatMask); + Div(sigmoidZeroReg, oneZeroReg, gateZeroReg, floatMask); + Div(sigmoidOneReg, oneOneReg, gateOneReg, floatMask); + Muls(sigmoidZeroReg, sigmoidZeroReg, gateScale, floatMask); + Muls(sigmoidOneReg, sigmoidOneReg, gateScale, floatMask); + Add(accZeroReg, accZeroReg, sigmoidZeroReg, floatMask); + Add(accOneReg, accOneReg, sigmoidOneReg, floatMask); + StoreAlign(output + rowOffset, accZeroReg, floatMask); + StoreAlign(output + rowOffset + FLOAT_ELEMENTS_PER_REG, accOneReg, floatMask); + } +} + +#endif + +template +class KdaGateCumsumKernel { +public: + template + __aicore__ inline void Init(GM_ADDR g, GM_ADDR aLog, GM_ADDR dtBias, GM_ADDR cuSeqlens, GM_ADDR gk, + const TilingData &tiling, TPipe *pipe) + { + g_.SetGlobalBuffer(reinterpret_cast<__gm__ T *>(g)); + aLog_.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(aLog)); + dtBias_.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(dtBias)); + cuSeqlens_.SetGlobalBuffer(reinterpret_cast<__gm__ int64_t *>(cuSeqlens)); + gk_.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(gk)); + pipe_ = pipe; + batch_ = static_cast(tiling.batch); + t_ = static_cast(tiling.t); + hv_ = static_cast(tiling.hv); + k_ = static_cast(tiling.k); + rank_ = static_cast(tiling.rank); + chunkSize_ = static_cast(tiling.chunkSize); + seqNum_ = static_cast(tiling.seqNum); + hasCuSeqlens_ = tiling.hasCuSeqlens != 0; + hasALog_ = tiling.hasALog != 0; + hasDtBias_ = tiling.hasDtBias != 0; + inputSequenceMajor_ = tiling.inputSequenceMajor != 0; + lowerBound_ = tiling.lowerBound; + usedCoreNum_ = static_cast(tiling.usedCoreNum); + maxChunks_ = (t_ + chunkSize_ - 1) / chunkSize_; + pipe_->InitBuffer(rowBuf_, GATE_ROW_ELEMENTS * sizeof(float)); + pipe_->InitBuffer(accBuf_, GATE_ROW_ELEMENTS * sizeof(float)); + pipe_->InitBuffer(outBuf_, GATE_PIPELINE_DEPTH * GATE_ROW_ELEMENTS * sizeof(float)); + pipe_->InitBuffer(tmpBuf_, GATE_ROW_ELEMENTS * sizeof(float)); + pipe_->InitBuffer(oneBuf_, GATE_ROW_ELEMENTS * sizeof(float)); + pipe_->InitBuffer(biasBuf_, GATE_ROW_ELEMENTS * sizeof(float)); + pipe_->InitBuffer(inBuf_, GATE_PIPELINE_DEPTH * GATE_ROW_ELEMENTS * sizeof(T)); + pipe_->InitBuffer(chunkBuf_, 2 * GATE_BULK_ELEMENTS * sizeof(float)); + pipe_->InitBuffer(scalarBuf_, 32); + pipe_->InitBuffer(scalarI64Buf_, 32); + AllocEvents(); + } + + __aicore__ inline void Process() + { + uint64_t taskCount = hasCuSeqlens_ ? seqNum_ * hv_ : batch_ * hv_ * maxChunks_; + uint64_t coreIdx = static_cast(GetBlockIdx()); + for (uint64_t task = coreIdx; task < taskCount; task += usedCoreNum_) { + ProcessTask(task); + } + ReleaseEvents(); + } + +private: + __aicore__ inline void AllocEvents() + { + for (uint32_t slot = 0; slot < GATE_PIPELINE_DEPTH; ++slot) { + inputMte2ToVEvent_[slot] = pipe_->AllocEventID(); + inputVToMte2Event_[slot] = pipe_->AllocEventID(); + outputVToMte3Event_[slot] = pipe_->AllocEventID(); + outputMte3ToVEvent_[slot] = pipe_->AllocEventID(); + } + auxMte2ToVEvent_ = pipe_->AllocEventID(); + scalarVToSEvent_ = pipe_->AllocEventID(); + bulkMte3ToMte2Event_ = pipe_->AllocEventID(); + } + + __aicore__ inline void ReleaseEvents() + { + for (uint32_t slot = 0; slot < GATE_PIPELINE_DEPTH; ++slot) { + pipe_->ReleaseEventID(inputMte2ToVEvent_[slot]); + pipe_->ReleaseEventID(inputVToMte2Event_[slot]); + pipe_->ReleaseEventID(outputVToMte3Event_[slot]); + pipe_->ReleaseEventID(outputMte3ToVEvent_[slot]); + } + pipe_->ReleaseEventID(auxMte2ToVEvent_); + pipe_->ReleaseEventID(scalarVToSEvent_); + pipe_->ReleaseEventID(bulkMte3ToMte2Event_); + } + + __aicore__ inline uint64_t InputOffset(uint64_t b, uint64_t t, uint64_t hv, uint64_t k) const + { + if (inputSequenceMajor_) { + if (rank_ == 4) { + return ((b * t_ + t) * hv_ + hv) * k_ + k; + } + return (t * hv_ + hv) * k_ + k; + } + if (rank_ == 4) { + return ((b * hv_ + hv) * t_ + t) * k_ + k; + } + return (hv * t_ + t) * k_ + k; + } + + __aicore__ inline uint64_t OutputOffset(uint64_t b, uint64_t t, uint64_t hv, uint64_t k) const + { + if (rank_ == 4) { + return ((b * hv_ + hv) * t_ + t) * k_ + k; + } + return (hv * t_ + t) * k_ + k; + } + + __aicore__ inline void CopyVectorIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset, uint64_t count) + { + uint64_t rowBytes = count * static_cast(sizeof(T)); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst, src[offset], static_cast(count)); + } else { + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPadParams padParams{false, 0, 0, 0}; + DataCopyPad(dst, src[offset], params, padParams); + } + } + + __aicore__ inline void CopyFloatVectorIn(LocalTensor &dst, GlobalTensor &src, uint64_t offset, + uint64_t count) + { + uint64_t rowBytes = count * sizeof(float); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst, src[offset], static_cast(count)); + } else { + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPadParams padParams{false, 0, 0, 0}; + DataCopyPad(dst, src[offset], params, padParams); + } + } + + __aicore__ inline void CopyFloatVectorOut(GlobalTensor &dst, uint64_t offset, LocalTensor &src, + uint64_t count) + { + uint64_t rowBytes = count * sizeof(float); + if (rowBytes >= 32 && rowBytes % 32 == 0) { + DataCopy(dst[offset], src, static_cast(count)); + } else { + DataCopyParams params{1, static_cast(rowBytes), 0, 0}; + DataCopyPad(dst[offset], src, params); + } + } + + __aicore__ inline void CopyGateRowsIn(LocalTensor &dst, uint64_t b, uint64_t start, + uint64_t hv, uint64_t rows) + { + if (!inputSequenceMajor_) { + CopyVectorIn(dst, g_, InputOffset(b, start, hv, 0), rows * k_); + return; + } + DataCopyExtParams params{ + static_cast(rows), + static_cast(k_ * sizeof(T)), + static_cast((hv_ - 1) * k_ * sizeof(T)), + 0, + 0}; + DataCopyPadExtParams padParams{false, 0, 0, 0}; + DataCopyPad(dst, g_[InputOffset(b, start, hv, 0)], params, padParams); + } + + __aicore__ inline void PrefetchGateRow(uint64_t offset, uint32_t slot) + { + LocalTensor input = inBuf_.Get()[slot * GATE_ROW_ELEMENTS]; + CopyVectorIn(input, g_, offset, k_); + SetFlag(inputMte2ToVEvent_[slot]); + } + + __aicore__ inline void MaterializeGateRow(uint32_t slot, LocalTensor &row) + { + LocalTensor input = inBuf_.Get()[slot * GATE_ROW_ELEMENTS]; + if constexpr (IsSameType::value) { + Adds(row, input, 0.0f, static_cast(k_)); + } else { + Cast(row, input, RoundMode::CAST_NONE, static_cast(k_)); + } + PipeBarrier(); + } + + __aicore__ inline float ReadFloat(GlobalTensor &tensor, uint64_t offset) + { + LocalTensor scalar = scalarBuf_.Get(); + DataCopyParams params{1, static_cast(sizeof(float)), 0, 0}; + DataCopyPadParams padParams{false, 0, 0, 0}; + DataCopyPad(scalar, tensor[offset], params, padParams); + SetFlag(auxMte2ToVEvent_); + WaitFlag(auxMte2ToVEvent_); + Adds(scalar, scalar, 0.0f, 1); + PipeBarrier(); + SetFlag(scalarVToSEvent_); + WaitFlag(scalarVToSEvent_); + __ubuf__ float *ptr = (__ubuf__ float *)scalar.GetPhyAddr(); + return ptr[0]; + } + + __aicore__ inline int64_t ReadInt64(GlobalTensor &tensor, uint64_t offset) + { + LocalTensor scalar = scalarI64Buf_.Get(); + DataCopyParams params{1, static_cast(sizeof(int64_t)), 0, 0}; + DataCopyPadParams padParams{false, 0, 0, 0}; + DataCopyPad(scalar, tensor[offset], params, padParams); + SetFlag(auxMte2ToVEvent_); + WaitFlag(auxMte2ToVEvent_); + SetFlag(scalarVToSEvent_); + WaitFlag(scalarVToSEvent_); + __ubuf__ int64_t *ptr = (__ubuf__ int64_t *)scalar.GetPhyAddr(); + return ptr[0]; + } + + __aicore__ inline float ExpScalar(float x) + { + LocalTensor scalar = scalarBuf_.Get(); + Duplicate(scalar, x, 1); + PipeBarrier(); + Exp(scalar, scalar, 1); + PipeBarrier(); + SetFlag(scalarVToSEvent_); + WaitFlag(scalarVToSEvent_); + __ubuf__ float *ptr = (__ubuf__ float *)scalar.GetPhyAddr(); + return ptr[0]; + } + + __aicore__ inline void PrepareGate(uint64_t hv) + { + if constexpr (USE_GATE_IN_KERNEL) { + expA_ = ExpScalar(ReadFloat(aLog_, hv)); + if (hasDtBias_) { + LocalTensor bias = biasBuf_.Get(); + CopyFloatVectorIn(bias, dtBias_, hv * k_, k_); + SetFlag(auxMte2ToVEvent_); + WaitFlag(auxMte2ToVEvent_); + } + } + } + + __aicore__ inline void ApplyGate(LocalTensor &row) + { + if constexpr (USE_GATE_IN_KERNEL) { + if (hasDtBias_) { + LocalTensor bias = biasBuf_.Get(); + Add(row, row, bias, static_cast(k_)); + PipeBarrier(); + } + if constexpr (SAFE_GATE) { + Muls(row, row, expA_, static_cast(k_)); + PipeBarrier(); + + LocalTensor tmp = tmpBuf_.Get(); + Muls(tmp, row, -1.0f, static_cast(k_)); + PipeBarrier(); + Exp(tmp, tmp, static_cast(k_)); + PipeBarrier(); + Adds(tmp, tmp, 1.0f, static_cast(k_)); + PipeBarrier(); + + LocalTensor one = oneBuf_.Get(); + Duplicate(one, 1.0f, static_cast(k_)); + PipeBarrier(); + Div(row, one, tmp, static_cast(k_)); + PipeBarrier(); + Muls(row, row, lowerBound_, static_cast(k_)); + PipeBarrier(); + } else { + LocalTensor positive = oneBuf_.Get(); + LocalTensor tmp = tmpBuf_.Get(); + Maxs(positive, row, 0.0f, static_cast(k_)); + Abs(tmp, row, static_cast(k_)); + PipeBarrier(); + Muls(tmp, tmp, -1.0f, static_cast(k_)); + PipeBarrier(); + Exp(tmp, tmp, static_cast(k_)); + PipeBarrier(); + Adds(tmp, tmp, 1.0f, static_cast(k_)); + PipeBarrier(); + Ln(tmp, tmp, static_cast(k_)); + PipeBarrier(); + Add(row, positive, tmp, static_cast(k_)); + PipeBarrier(); + Muls(row, row, -expA_, static_cast(k_)); + PipeBarrier(); + } + } + } + + __aicore__ inline void ProcessTask(uint64_t task) + { + uint64_t hv = hasCuSeqlens_ ? task % hv_ : (task / maxChunks_) % hv_; + PrepareGate(hv); + if (!hasCuSeqlens_) { + uint64_t chunk = task % maxChunks_; + uint64_t b = task / (maxChunks_ * hv_); + uint64_t start = chunk * chunkSize_; + uint64_t end = start + chunkSize_; + if (end > t_) { + end = t_; + } + ProcessChunk(b, hv, start, end); + return; + } + uint64_t seq = task / hv_; + uint64_t seqStart = static_cast(ReadInt64(cuSeqlens_, seq)); + uint64_t seqEnd = static_cast(ReadInt64(cuSeqlens_, seq + 1)); + for (uint64_t start = seqStart; start < seqEnd; start += chunkSize_) { + uint64_t end = start + chunkSize_; + if (end > seqEnd) { + end = seqEnd; + } + ProcessChunk(0, hv, start, end); + } + } + + __aicore__ inline void ProcessChunk(uint64_t b, uint64_t hv, uint64_t start, uint64_t end) + { +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 220 + if constexpr (USE_GATE_IN_KERNEL && SAFE_GATE && IsSameType::value) { + if (k_ == GATE_BULK_COLS && chunkSize_ == GATE_BULK_ROWS && end - start == GATE_BULK_ROWS) { + ProcessChunkBulkFp32SafeVec(b, hv, start); + return; + } + } +#endif +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (USE_GATE_IN_KERNEL && SAFE_GATE && IsSameType::value) { + if (k_ == GATE_BULK_COLS && chunkSize_ == GATE_BULK_ROWS) { + ProcessChunkBulkFp32Safe(b, hv, start, end); + return; + } + } + if constexpr (!USE_GATE_IN_KERNEL && IsSameType::value) { + if (k_ == GATE_BULK_COLS && chunkSize_ == GATE_BULK_ROWS) { + ProcessChunkBulkFp32(b, hv, start, end); + return; + } + } +#endif + LocalTensor acc = accBuf_.Get(); + LocalTensor row = rowBuf_.Get(); + Duplicate(acc, 0.0f, static_cast(k_)); + PipeBarrier(); + uint64_t rows = end - start; + if (rows == 0) { + return; + } + + PrefetchGateRow(InputOffset(b, start, hv, 0), 0); + + for (uint64_t rowIdx = 0; rowIdx < rows; ++rowIdx) { + uint64_t token = start + rowIdx; + uint32_t slot = static_cast(rowIdx & 1); + WaitFlag(inputMte2ToVEvent_[slot]); + + if (rowIdx + 1 < rows) { + uint32_t nextSlot = slot ^ 1; + if (rowIdx >= 1) { + WaitFlag(inputVToMte2Event_[nextSlot]); + } + PrefetchGateRow(InputOffset(b, token + 1, hv, 0), nextSlot); + } + + uint32_t outputSlot = slot; + if (rowIdx >= GATE_PIPELINE_DEPTH) { + WaitFlag(outputMte3ToVEvent_[outputSlot]); + } + LocalTensor output = + outBuf_.Get()[outputSlot * GATE_ROW_ELEMENTS]; + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + if constexpr (!USE_GATE_IN_KERNEL && IsSameType::value) { + if ((k_ % 128) == 0) { + LocalTensor input = inBuf_.Get()[slot * GATE_ROW_ELEMENTS]; + AccumulateGateRowRegbase( + (__ubuf__ float *)input.GetPhyAddr(), (__ubuf__ float *)acc.GetPhyAddr(), + (__ubuf__ float *)output.GetPhyAddr(), static_cast(k_)); + PipeBarrier(); + } else { + MaterializeGateRow(slot, row); + Muls(row, row, RCP_LN2, static_cast(k_)); + PipeBarrier(); + Add(acc, acc, row, static_cast(k_)); + PipeBarrier(); + Adds(output, acc, 0.0f, static_cast(k_)); + } + } else { + MaterializeGateRow(slot, row); + ApplyGate(row); + Muls(row, row, RCP_LN2, static_cast(k_)); + PipeBarrier(); + Add(acc, acc, row, static_cast(k_)); + PipeBarrier(); + Adds(output, acc, 0.0f, static_cast(k_)); + } +#else + MaterializeGateRow(slot, row); + ApplyGate(row); + Muls(row, row, RCP_LN2, static_cast(k_)); + PipeBarrier(); + Add(acc, acc, row, static_cast(k_)); + PipeBarrier(); + Adds(output, acc, 0.0f, static_cast(k_)); +#endif + + PipeBarrier(); + SetFlag(inputVToMte2Event_[slot]); + SetFlag(outputVToMte3Event_[outputSlot]); + WaitFlag(outputVToMte3Event_[outputSlot]); + CopyFloatVectorOut(gk_, OutputOffset(b, token, hv, 0), output, k_); + SetFlag(outputMte3ToVEvent_[outputSlot]); + } + + uint64_t drainStart = rows > GATE_PIPELINE_DEPTH ? rows - GATE_PIPELINE_DEPTH : 0; + for (uint64_t rowIdx = drainStart; rowIdx < rows; ++rowIdx) { + uint32_t outputSlot = static_cast(rowIdx & 1); + WaitFlag(inputVToMte2Event_[outputSlot]); + WaitFlag(outputMte3ToVEvent_[outputSlot]); + } + } + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 220 + __aicore__ inline void ScanGateChunk64x128(LocalTensor input, LocalTensor output) + { + constexpr uint32_t rowElements = GATE_BULK_COLS; + Adds(input, output, 0.0f, rowElements); + Add(input[rowElements], output[rowElements], output, (GATE_BULK_ROWS - 1) * rowElements); + PipeBarrier(); + + Adds(output, input, 0.0f, 2 * rowElements); + Add(output[2 * rowElements], input[2 * rowElements], input, (GATE_BULK_ROWS - 2) * rowElements); + PipeBarrier(); + + Adds(input, output, 0.0f, 4 * rowElements); + Add(input[4 * rowElements], output[4 * rowElements], output, (GATE_BULK_ROWS - 4) * rowElements); + PipeBarrier(); + + Adds(output, input, 0.0f, 8 * rowElements); + Add(output[8 * rowElements], input[8 * rowElements], input, (GATE_BULK_ROWS - 8) * rowElements); + PipeBarrier(); + + Adds(input, output, 0.0f, 16 * rowElements); + Add(input[16 * rowElements], output[16 * rowElements], output, (GATE_BULK_ROWS - 16) * rowElements); + PipeBarrier(); + + Adds(output, input, 0.0f, 32 * rowElements); + Add(output[32 * rowElements], input[32 * rowElements], input, (GATE_BULK_ROWS - 32) * rowElements); + PipeBarrier(); + } + + __aicore__ inline void ProcessChunkBulkFp32SafeVec(uint64_t b, uint64_t hv, uint64_t start) + { + constexpr uint32_t elems = GATE_BULK_ELEMENTS; + LocalTensor input = chunkBuf_.Get(); + LocalTensor output = chunkBuf_.Get()[GATE_BULK_ELEMENTS]; + CopyGateRowsIn(input, b, start, hv, GATE_BULK_ROWS); + SetFlag(inputMte2ToVEvent_[0]); + WaitFlag(inputMte2ToVEvent_[0]); + + Duplicate(output, 1.0f, elems); + if (hasDtBias_) { + LocalTensor bias = biasBuf_.Get(); + for (uint32_t row = 0; row < GATE_BULK_ROWS; ++row) { + Add(input[row * GATE_BULK_COLS], input[row * GATE_BULK_COLS], bias, GATE_BULK_COLS); + } + PipeBarrier(); + } + Muls(input, input, -expA_, elems); + PipeBarrier(); + Exp(input, input, elems); + PipeBarrier(); + Adds(input, input, 1.0f, elems); + PipeBarrier(); + Div(output, output, input, elems); + PipeBarrier(); + Muls(output, output, lowerBound_ * RCP_LN2, elems); + PipeBarrier(); + ScanGateChunk64x128(input, output); + + SetFlag(outputVToMte3Event_[0]); + WaitFlag(outputVToMte3Event_[0]); + CopyFloatVectorOut(gk_, OutputOffset(b, start, hv, 0), output, elems); + SetFlag(bulkMte3ToMte2Event_); + WaitFlag(bulkMte3ToMte2Event_); + SetFlag(outputMte3ToVEvent_[0]); + WaitFlag(outputMte3ToVEvent_[0]); + } +#endif + +#if defined(__CCE_AICORE__) && __CCE_AICORE__ == 310 + __aicore__ inline void ProcessChunkBulkFp32Safe(uint64_t b, uint64_t hv, uint64_t start, uint64_t end) + { + uint64_t rows = end - start; + if (rows == 0) { + return; + } + uint32_t elems = static_cast(rows * k_); + LocalTensor input = chunkBuf_.Get(); + LocalTensor output = chunkBuf_.Get()[GATE_BULK_ELEMENTS]; + CopyGateRowsIn(input, b, start, hv, rows); + SetFlag(inputMte2ToVEvent_[0]); + WaitFlag(inputMte2ToVEvent_[0]); + LocalTensor bias = biasBuf_.Get(); + if (hasDtBias_) { + AccumulateSafeGateChunk128Regbase( + (__ubuf__ float *)input.GetPhyAddr(), (__ubuf__ float *)bias.GetPhyAddr(), + (__ubuf__ float *)output.GetPhyAddr(), static_cast(rows), expA_, lowerBound_); + } else { + AccumulateSafeGateChunk128Regbase( + (__ubuf__ float *)input.GetPhyAddr(), (__ubuf__ float *)bias.GetPhyAddr(), + (__ubuf__ float *)output.GetPhyAddr(), static_cast(rows), expA_, lowerBound_); + } + PipeBarrier(); + SetFlag(outputVToMte3Event_[0]); + WaitFlag(outputVToMte3Event_[0]); + CopyFloatVectorOut(gk_, OutputOffset(b, start, hv, 0), output, elems); + SetFlag(bulkMte3ToMte2Event_); + WaitFlag(bulkMte3ToMte2Event_); + SetFlag(outputMte3ToVEvent_[0]); + WaitFlag(outputMte3ToVEvent_[0]); + } + + __aicore__ inline void ProcessChunkBulkFp32(uint64_t b, uint64_t hv, uint64_t start, uint64_t end) + { + uint64_t rows = end - start; + if (rows == 0) { + return; + } + uint32_t elems = static_cast(rows * k_); + LocalTensor buffer0 = chunkBuf_.Get(); + CopyGateRowsIn(buffer0, b, start, hv, rows); + SetFlag(inputMte2ToVEvent_[0]); + WaitFlag(inputMte2ToVEvent_[0]); + AccumulateGateChunk128Regbase( + (__ubuf__ float *)buffer0.GetPhyAddr(), (__ubuf__ float *)buffer0.GetPhyAddr(), + static_cast(rows)); + PipeBarrier(); + SetFlag(outputVToMte3Event_[0]); + WaitFlag(outputVToMte3Event_[0]); + CopyFloatVectorOut(gk_, OutputOffset(b, start, hv, 0), buffer0, elems); + SetFlag(bulkMte3ToMte2Event_); + WaitFlag(bulkMte3ToMte2Event_); + SetFlag(outputMte3ToVEvent_[0]); + WaitFlag(outputMte3ToVEvent_[0]); + } +#endif + + GlobalTensor g_; + GlobalTensor aLog_; + GlobalTensor dtBias_; + GlobalTensor cuSeqlens_; + GlobalTensor gk_; + TPipe *pipe_ = nullptr; + TBuf rowBuf_; + TBuf accBuf_; + TBuf outBuf_; + TBuf tmpBuf_; + TBuf oneBuf_; + TBuf biasBuf_; + TBuf inBuf_; + TBuf chunkBuf_; + TBuf scalarBuf_; + TBuf scalarI64Buf_; + TEventID inputMte2ToVEvent_[GATE_PIPELINE_DEPTH]; + TEventID inputVToMte2Event_[GATE_PIPELINE_DEPTH]; + TEventID outputVToMte3Event_[GATE_PIPELINE_DEPTH]; + TEventID outputMte3ToVEvent_[GATE_PIPELINE_DEPTH]; + TEventID auxMte2ToVEvent_; + TEventID scalarVToSEvent_; + TEventID bulkMte3ToMte2Event_; + uint64_t batch_ = 0; + uint64_t t_ = 0; + uint64_t hv_ = 0; + uint64_t k_ = 0; + uint64_t rank_ = 0; + uint64_t chunkSize_ = 0; + uint64_t seqNum_ = 0; + uint64_t maxChunks_ = 0; + bool hasCuSeqlens_ = false; + bool hasALog_ = false; + bool hasDtBias_ = false; + bool inputSequenceMajor_ = false; + float expA_ = 1.0f; + float lowerBound_ = -5.0f; + uint64_t usedCoreNum_ = 1; +}; + +template +__aicore__ inline void RunKdaGateCumsum(GM_ADDR g, GM_ADDR aLog, GM_ADDR dtBias, GM_ADDR cuSeqlens, GM_ADDR gk, + const TilingData &tilingData, TPipe *pipe) +{ + KdaGateCumsumKernel op; + op.Init(g, aLog, dtBias, cuSeqlens, gk, tilingData, pipe); + op.Process(); +} + +template +__aicore__ inline void DispatchKdaGateCumsum(GM_ADDR g, GM_ADDR aLog, GM_ADDR dtBias, GM_ADDR cuSeqlens, + GM_ADDR gk, const TilingData &tilingData, TPipe *pipe) +{ + if (tilingData.useGateInKernel != 0) { + if (tilingData.safeGate != 0) { + RunKdaGateCumsum(g, aLog, dtBias, cuSeqlens, gk, tilingData, pipe); + } else { + RunKdaGateCumsum(g, aLog, dtBias, cuSeqlens, gk, tilingData, pipe); + } + } else { + RunKdaGateCumsum(g, aLog, dtBias, cuSeqlens, gk, tilingData, pipe); + } +} +} // namespace KdaGateCumsum diff --git a/csrc/chunk_kda_fwd/op_kernel/kernel_utils/block/block_mmad_pingpong_tla_kernel_utils.hpp b/csrc/chunk_kda_fwd/op_kernel/kernel_utils/block/block_mmad_pingpong_tla_kernel_utils.hpp new file mode 100644 index 0000000000..f099eb38d3 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/kernel_utils/block/block_mmad_pingpong_tla_kernel_utils.hpp @@ -0,0 +1,1014 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef COMMON_BLOCK_MMAD_PINGPONG_TLA_KERNEL_UTILS_HPP +#define COMMON_BLOCK_MMAD_PINGPONG_TLA_KERNEL_UTILS_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" +#include "../tile/copy_l0c_to_ub.hpp" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +namespace Common { + +template , + class TileMmad = Catlass::Gemm::Tile::TileMmadTla > +struct BlockMmadTla { + static_assert(DEPENDENT_FALSE, "BlockMmadTla is not implemented for this DispatchPolicy"); +}; + +// Now ENABLE_UNIT_FLAG_ must be false when intput element is int8 +template +struct MmadPingpong : public Catlass::Gemm::MmadBase { + static constexpr uint32_t L1A_STAGES = L1A_STAGES_; + static constexpr uint32_t L1B_STAGES = L1B_STAGES_; + static constexpr uint32_t L0A_STAGES = L0A_STAGES_; + static constexpr uint32_t L0B_STAGES = L0B_STAGES_; + static constexpr uint32_t L0C_STAGES = L0C_STAGES_; + static constexpr uint32_t UB_STAGES = UB_STAGES_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool USE_HF32_MODE = USE_HF32_MODE_; + static constexpr bool ENABLE_L1_RESIDENT = ENABLE_L1_RESIDENT_; +}; + +template +struct BlockMmadTla, + L1TileShape_, L0TileShape_, ElementA_, ElementB_, ElementC_, ElementBias_, TileCopy_, TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadPingpong; + using ArchTag = typename DispatchPolicy::ArchTag; + using TileCopy = TileCopy_; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = ElementA_; + using LayoutA = typename TileCopy::LayoutA; + using ElementB = ElementB_; + using LayoutB = typename TileCopy::LayoutB; + using ElementC = ElementC_; + using LayoutC = typename TileCopy::LayoutC; + using ElementBias = ElementBias_; + + using TileMmad = TileMmad_; + + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; + using CopyL1ToBT = typename TileCopy::CopyL1ToBT; + + using ElementAccumulator = typename TileCopy::ElementAccumulator; + + static constexpr bool HAS_BIAS = TileCopy::HAS_BIAS; + + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + + using L1AAlignHelper = typename TileCopy_::L1AAlignHelper; + using L1BAlignHelper = typename TileCopy_::L1BAlignHelper; + + static_assert(tla::is_tuple::value && tla::is_static::value, + "L1TileShape must be tla::tuple and static!"); + static_assert(tla::is_tuple::value && tla::is_static::value, + "L0TileShape must be tla::tuple and static!"); + + static constexpr uint64_t FLAG_ID_MAX = 16; + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr bool USE_HF32_MODE = DispatchPolicy::USE_HF32_MODE; + static constexpr bool ENABLE_L1_RESIDENT = DispatchPolicy::ENABLE_L1_RESIDENT; + static constexpr uint32_t L1A_STAGES = DispatchPolicy::L1A_STAGES; + static constexpr uint32_t L1B_STAGES = DispatchPolicy::L1B_STAGES; + static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES; + static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES; + static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; + static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); + static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); + static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); + static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); + static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); + static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); + static constexpr uint32_t UB_STAGES = UB_STAGES_; + static constexpr uint32_t MAX_CUBE_VEC_SYNC_NUM = 5; + + // L1 tile size + static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA); + static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); + // L0 tile size + static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L1_TILE_M * L1_TILE_N * sizeof(ElementAccumulator); + + // Check HF32_MODE + static_assert(!USE_HF32_MODE || + (USE_HF32_MODE && std::is_same_v && std::is_same_v), + "HF32 MODE only supports in float!"); + + // Check L0C_STAGES + static_assert(!(ENABLE_UNIT_FLAG && L0C_STAGES != 1), "L0C_STAGES must be 1 when UnitFlag is true!"); + + // Check LayoutC + static_assert(tla::detail::isRowMajor::value || + ((std::is_same_v || std::is_same_v || + std::is_same_v) && + tla::detail::iszN::value), + "LayoutC only supports zN in half or bfloat16 or float, RowMajor in all dtype yet!"); + + // Check L1TileShape + static_assert(L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static_assert(L0A_TILE_SIZE * L0A_STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert(L0B_TILE_SIZE * L0B_STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE * L0C_STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); + + static constexpr uint32_t _32B = 32 * 8; // in bits + static_assert(L1_TILE_M == L0_TILE_M && L1_TILE_N == L0_TILE_N, + "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); + static_assert(L0_TILE_K <= L1_TILE_K, "L0TileShape::K cannot exceed L1TileShape::K"); +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) + static_assert(L1_TILE_M * SizeOfBits::value % _32B == 0, "L1TileShape::M must be 32B aligned."); + static_assert(L1_TILE_K * SizeOfBits::value % _32B == 0, "L1TileShape::K must be 32B aligned."); + static_assert(L1_TILE_K * SizeOfBits::value % _32B == 0, "L1TileShape::K must be 32B aligned."); + static_assert(L1_TILE_N * SizeOfBits::value % _32B == 0, "L1TileShape::N must be 32B aligned."); + static_assert(L0_TILE_K * SizeOfBits::value % _32B == 0, "L0TileShape::K must be 32B aligned."); +#endif + + static_assert((!HAS_BIAS && (L1A_STAGES + L1B_STAGES) <= 8) || (HAS_BIAS && (L1A_STAGES + L1B_STAGES) <= 7), + "L1 Buffer overflow: Exceeds the supported range of EVENT(0~7)"); + + static_assert((!HAS_BIAS && (L0A_STAGES + L0B_STAGES) <= 8) || (HAS_BIAS && (L0A_STAGES + L0B_STAGES) <= 7), + "L0 Buffer overflow: Exceeds the supported range of EVENT_ID(0~7)"); + + static constexpr auto L1A_LAYOUT = + tla::MakeLayout(tla::Int{}, tla::Int{}); + static constexpr auto L1B_LAYOUT = + tla::MakeLayout(tla::Int{}, tla::Int{}); + static constexpr auto L1BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); + static constexpr auto L0BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); + + // When enabling L1 resident mode, restore the pointer and coordinates that record the last state + // to the initial state. if two blockmmad instances need to be consecutively invoked at the kernel layer, + // RestoreStatus() must be inserted between them. + CATLASS_DEVICE + void RestoreStatus() + { + for (int i = 0; i < L1A_STAGES; ++i) { + lastAddrA[i] = nullptr; + lastCoordA[i] = Catlass::MatrixCoord{0U, 0U}; + } + for (int i = 0; i < L1B_STAGES; ++i) { + lastAddrB[i] = nullptr; + lastCoordB[i] = Catlass::MatrixCoord{0U, 0U}; + } + } + + /// Construct + CATLASS_DEVICE + BlockMmadTla(Catlass::Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + if ASCEND_IS_AIC { + // use HF32 when USE_HF32_MODE is true + if constexpr (USE_HF32_MODE) { + AscendC::SetHF32Mode(true); + } else { + AscendC::SetHF32Mode(false); + } + if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { + AscendC::SetMMLayoutTransform(true); + } + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * L1A_STAGES; + // Init buffers + for (uint32_t i = 0; i < L1A_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_TILE_SIZE * i); + // Assign event ID for each stages + l1AEventList[i] = i; + // The event id that needs to be set before the loop + AscendC::SetFlag(l1AEventList[i]); + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); + // Assign event ID for each stages + l1BEventList[i] = i + L1A_STAGES; + // The event id that needs to be set before the loop + AscendC::SetFlag(l1BEventList[i]); + } + for (uint32_t i = 0; i < L0A_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); + // Assign event ID for each stages + l0AEventList[i] = i; + // The event id that needs to be set before the loop + AscendC::SetFlag(l0AEventList[i]); + } + for (uint32_t i = 0; i < L0B_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); + // Assign event ID for each stages + l0BEventList[i] = i + L0A_STAGES; + // The event id that needs to be set before the loop + AscendC::SetFlag(l0BEventList[i]); + } + if constexpr (!ENABLE_UNIT_FLAG) { + for (uint32_t i = 0; i < L0C_STAGES; i++) { + l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte(L0C_TILE_SIZE * i); + l0CEventList[i] = i; + AscendC::SetFlag(l0CEventList[i]); + } + } else { + l0CTensorList[0] = resource.l0CBuf.template GetBufferByByte(0); + } + if constexpr (HAS_BIAS) { + uint32_t l1BiasOffset = l1BOffset + L1B_TILE_SIZE * L1B_STAGES; + l1BiasTensor = resource.l1Buf.template GetBufferByByte(l1BiasOffset); + l0BiasTensor = resource.btBuf.template GetBufferByByte(0); + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + AscendC::SetFlag(L0A_STAGES + L0B_STAGES); + } + + if constexpr (ENABLE_L1_RESIDENT) { + RestoreStatus(); + } + } + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmadTla() + { + if ASCEND_IS_AIC { + if constexpr (USE_HF32_MODE) { + AscendC::SetHF32Mode(false); + } + if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { + AscendC::SetMMLayoutTransform(false); + } + for (uint32_t i = 0; i < L1A_STAGES; i++) { + AscendC::WaitFlag(l1AEventList[i]); + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + AscendC::WaitFlag(l1BEventList[i]); + } + for (uint32_t i = 0; i < L0A_STAGES; i++) { + AscendC::WaitFlag(l0AEventList[i]); + } + for (uint32_t i = 0; i < L0B_STAGES; i++) { + AscendC::WaitFlag(l0BEventList[i]); + } + if constexpr (!ENABLE_UNIT_FLAG) { + for (uint32_t i = 0; i < L0C_STAGES; i++) { + AscendC::WaitFlag(l0CEventList[i]); + } + } + if constexpr (HAS_BIAS) { + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + AscendC::WaitFlag(L0A_STAGES + L0B_STAGES); + } + } + } + + /// Perform a block-scoped matrix multiply-accumulate + template + CATLASS_DEVICE void operator()(TensorA &tensorA, TensorB &tensorB, TensorC &tensorC, + Catlass::GemmCoord const &actualShape, TensorBias const &tensorBias = {}) + { + // Check L1TileShape + if constexpr (HAS_BIAS) { + static constexpr uint32_t BIAS_BUF_SIZE = L0_TILE_N * sizeof(ElementAccumulator); + static constexpr uint32_t L1BIAS_SIZE = L1_TILE_N * sizeof(ElementBias); + static_assert(BIAS_BUF_SIZE <= ArchTag::BIAS_SIZE, + "BIAS_BUF_SIZE exceeding the BT space! Reduce L0_TILE_N"); + static_assert(L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES + L1BIAS_SIZE <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + } + + using CopyGmToL1A = typename TileCopy_::template CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::template CopyGmToL1B; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) + using CopyL0CToGm = typename TileCopy_::template CopyL0CToGm; + CopyL0CToGm copyL0CToDst; +#endif +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) + using CopyL0CToDst = typename TileCopy_::template CopyL0CToDst; + CopyL0CToDst copyL0CToDst; +#endif + + uint32_t mBlockActual = actualShape.m(); + uint32_t kBlockActual = actualShape.k(); + uint32_t nBlockActual = actualShape.n(); + + uint32_t mL1Actual = mBlockActual; + if constexpr (std::is_same_v) { + // Avoid using the gemv mode in mmad + if (mL1Actual == 1) { + mL1Actual = 16; + } + } + uint32_t nL1Actual = nBlockActual; + + auto layoutInL0C = tla::MakeLayoutL0C(mL1Actual, nL1Actual); + auto tensorL0C = tla::MakeTensor(l0CTensorList[l0CListId], layoutInL0C, Catlass::Arch::PositionL0C{}); + auto tensorL0Bias = tla::MakeTensor(l0BiasTensor, L0BIAS_LAYOUT, Catlass::Arch::PositionBias{}); + + uint32_t kL1Actual = min(kBlockActual, L1_TILE_K); + // load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1AListId]); + auto tensorL1A = tla::MakeTensor(l1ATensorList[l1AListId], L1A_LAYOUT, Catlass::Arch::PositionL1{}); + auto tensorTileA = GetTileA(tensorA, 0, 0, mBlockActual, kL1Actual); + if constexpr (ENABLE_L1_RESIDENT) { + // If the currently loaded GM pointer and block coordinates are the same as the last loaded ones, + // skip this loading. + if (lastAddrA[l1AListId] != tensorTileA.data().GetPhyAddr() || + tla::get<0>(tensorTileA.coord()) != lastCoordA[l1AListId].row() || + tla::get<1>(tensorTileA.coord()) != lastCoordA[l1AListId].column()) { + copyGmToL1A(tensorL1A, tensorTileA); + lastCoordA[l1AListId] = + Catlass::MatrixCoord{tla::get<0>(tensorTileA.coord()), tla::get<1>(tensorTileA.coord())}; + lastAddrA[l1AListId] = const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileA.data().GetPhyAddr()); + } + } else { + copyGmToL1A(tensorL1A, tensorTileA); + } + AscendC::SetFlag(l1AEventList[l1AListId]); + + // load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListId]); + auto tensorL1B = tla::MakeTensor(l1BTensorList[l1BListId], L1B_LAYOUT, Catlass::Arch::PositionL1{}); + auto tensorTileB = GetTile(tensorB, tla::MakeCoord(0, 0), tla::MakeShape(kL1Actual, nBlockActual)); + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrB[l1BListId] != tensorTileB.data().GetPhyAddr() || + tla::get<0>(tensorTileB.coord()) != lastCoordB[l1BListId].row() || + tla::get<1>(tensorTileB.coord()) != lastCoordB[l1BListId].column()) { + copyGmToL1B(tensorL1B, tensorTileB); + lastCoordB[l1BListId] = + Catlass::MatrixCoord{tla::get<0>(tensorTileB.coord()), tla::get<1>(tensorTileB.coord())}; + lastAddrB[l1BListId] = const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileB.data().GetPhyAddr()); + } + } else { + copyGmToL1B(tensorL1B, tensorTileB); + } + AscendC::SetFlag(l1BEventList[l1BListId]); + + if constexpr (HAS_BIAS && !std::is_same_v) { + using CopyGmToL1Bias = typename TileCopy::template CopyGmToL1Bias; + CopyGmToL1Bias copyGmToL1Bias; + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + auto l1Bias = l1BiasTensor.template ReinterpretCast(); + auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Catlass::Arch::PositionL1{}); + copyGmToL1Bias(tensorL1Bias, tensorBias); + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + } + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(l0CEventList[l0CListId]); + } + + uint32_t mL0Loop = CeilDiv(mL1Actual); + uint32_t nL0Loop = CeilDiv(nL1Actual); + + // main loop + uint32_t kL1Loop = CeilDiv(kBlockActual); + for (uint32_t kL1Idx = 0; kL1Idx < kL1Loop; kL1Idx++) { + uint32_t l1AListIdNext = (l1AListId + 1 < L1A_STAGES) ? (l1AListId + 1) : 0; + uint32_t l1BListIdNext = (l1BListId + 1 < L1B_STAGES) ? (l1BListId + 1) : 0; + uint32_t kL1ActualNext{0}; + // preload next tile from GM to L1 + if (kL1Idx < kL1Loop - 1) { + uint32_t kL1IdxNext = kL1Idx + 1; + kL1ActualNext = (kL1IdxNext < kL1Loop - 1) ? L1_TILE_K : (kBlockActual - kL1IdxNext * L1_TILE_K); + + // Get L1 tensor for next stage + auto l1ATensor = l1ATensorList[l1AListIdNext]; + auto l1BTensor = l1BTensorList[l1BListIdNext]; + auto tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Catlass::Arch::PositionL1{}); + auto tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Catlass::Arch::PositionL1{}); + // Get GM tile for next stage + auto tensorTileA = GetTileA(tensorA, 0, kL1IdxNext * L1_TILE_K, mBlockActual, kL1ActualNext); + auto tensorTileB = GetTile(tensorB, tla::MakeCoord(kL1IdxNext * L1_TILE_K, 0), + tla::MakeShape(kL1ActualNext, nBlockActual)); + + // load next matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1AListIdNext]); + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrA[l1AListIdNext] != tensorTileA.data().GetPhyAddr() || + tla::get<0>(tensorTileA.coord()) != lastCoordA[l1AListIdNext].row() || + tla::get<1>(tensorTileA.coord()) != lastCoordA[l1AListIdNext].column()) { + copyGmToL1A(tensorL1A, tensorTileA); + lastCoordA[l1AListIdNext] = + Catlass::MatrixCoord{tla::get<0>(tensorTileA.coord()), tla::get<1>(tensorTileA.coord())}; + lastAddrA[l1AListIdNext] = + const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileA.data().GetPhyAddr()); + } + } else { + copyGmToL1A(tensorL1A, tensorTileA); + } + AscendC::SetFlag(l1AEventList[l1AListIdNext]); + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListIdNext]); + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrB[l1BListIdNext] != tensorTileB.data().GetPhyAddr() || + tla::get<0>(tensorTileB.coord()) != lastCoordB[l1BListIdNext].row() || + tla::get<1>(tensorTileB.coord()) != lastCoordB[l1BListIdNext].column()) { + copyGmToL1B(tensorL1B, tensorTileB); + lastCoordB[l1BListIdNext] = + Catlass::MatrixCoord{tla::get<0>(tensorTileB.coord()), tla::get<1>(tensorTileB.coord())}; + lastAddrB[l1BListIdNext] = + const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileB.data().GetPhyAddr()); + } + } else { + copyGmToL1B(tensorL1B, tensorTileB); + } + AscendC::SetFlag(l1BEventList[l1BListIdNext]); + } + + // Get L1 tensor for current stage + auto l1ATensor = l1ATensorList[l1AListId]; + auto l1BTensor = l1BTensorList[l1BListId]; + tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Catlass::Arch::PositionL1{}); + tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Catlass::Arch::PositionL1{}); + // Get the loop nums on L0 + uint32_t kL0Loop = CeilDiv(kL1Actual); + + for (int mL0Idx = 0; mL0Idx < mL0Loop; mL0Idx++) { + uint32_t mL0Actual = (mL0Idx < mL0Loop - 1) ? L0_TILE_M : (mL1Actual - mL0Idx * L0_TILE_M); + + for (int kL0Idx = 0; kL0Idx < kL0Loop; kL0Idx++) { + uint32_t kL0Actual = (kL0Idx < kL0Loop - 1) ? L0_TILE_K : (kL1Actual - kL0Idx * L0_TILE_K); + + // Locate the current tile on L0A + auto l0ATile = l0ATensorList[l0AListId]; + auto layoutAInL0 = tla::MakeLayout(mL0Actual, kL0Actual); + auto tensorL0A = tla::MakeTensor(l0ATile, layoutAInL0, Catlass::Arch::PositionL0A{}); + // Locate the current tile of matrix A on L1 + auto tensorTileL1A = + GetTileA(tensorL1A, mL0Idx * L0_TILE_M, kL0Idx * L0_TILE_K, mL0Actual, kL0Actual); + + AscendC::WaitFlag(l0AEventList[l0AListId]); + if ((mL0Idx == 0) && (kL0Idx == 0)) { + AscendC::WaitFlag(l1AEventList[l1AListId]); + } + + // Load current tile from L1 to L0A + copyL1ToL0A(tensorL0A, tensorTileL1A); + + if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1)) { + AscendC::SetFlag(l1AEventList[l1AListId]); + } + + bool initC = ((kL1Idx == 0) && (kL0Idx == 0)); + for (int nL0Idx = 0; nL0Idx < nL0Loop; nL0Idx++) { + uint32_t nL0Actual = (nL0Idx < nL0Loop - 1) ? L0_TILE_N : (nL1Actual - nL0Idx * L0_TILE_N); + + // Locate the current tile on L0B + auto l0BTile = l0BTensorList[l0BListId]; + auto layoutBInL0 = tla::MakeLayout(kL0Actual, nL0Actual); + auto tensorL0B = tla::MakeTensor(l0BTile, layoutBInL0, Catlass::Arch::PositionL0B{}); + // Locate the current tile of matrix B on L1 + auto tensorTileL1B = GetTile(tensorL1B, tla::MakeCoord(kL0Idx * L0_TILE_K, nL0Idx * L0_TILE_N), + tla::MakeShape(kL0Actual, nL0Actual)); + + // Wait for mmad finished + AscendC::WaitFlag(l0BEventList[l0BListId]); + // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 + if ((mL0Idx == 0) && (kL0Idx == 0) && (nL0Idx == 0)) { + AscendC::WaitFlag(l1BEventList[l1BListId]); + } + + // Load current tile from L1 to L0B + copyL1ToL0B(tensorL0B, tensorTileL1B); + + // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 + if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1) && (nL0Idx == nL0Loop - 1)) { + AscendC::SetFlag(l1BEventList[l1BListId]); + } + + if constexpr (HAS_BIAS && !std::is_same_v) { + if (initC) { + if (nL0Idx == 0) { + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + } + AscendC::WaitFlag(L0A_STAGES + L0B_STAGES); + auto l1Bias = l1BiasTensor.template ReinterpretCast(); + auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Catlass::Arch::PositionL1{}); + auto tensorTileL1Bias = GetTile(tensorL1Bias, tla::MakeCoord(nL0Idx * L0_TILE_N), + tla::MakeShape(nL0Actual)); + // Load bias to l0 biasTable + copyL1ToBT(tensorL0Bias, tensorTileL1Bias); + if (nL0Idx == nL0Loop - 1) { + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + } + } + } + + // Notify to do mmad + AscendC::SetFlag(l0CEventList[l0CListId]); + + // Locate the current tile on L0C + auto tensorTileL0C = GetTile(tensorL0C, tla::MakeCoord(mL0Idx * L0_TILE_M, nL0Idx * L0_TILE_N), + tla::MakeShape(mL0Actual, nL0Actual)); + + // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator + // Wait for loading L0B + AscendC::WaitFlag(l0CEventList[l0CListId]); + + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if ((kL1Idx == kL1Loop - 1) && (mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1) && + (nL0Idx == nL0Loop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + + if constexpr (HAS_BIAS && !std::is_same_v) { + if (initC) { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, tensorL0Bias, mL0Actual, nL0Actual, + kL0Actual, initC, unitFlag); + AscendC::SetFlag(L0A_STAGES + L0B_STAGES); + } else { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, mL0Actual, nL0Actual, kL0Actual, initC, + unitFlag); + } + } else { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, mL0Actual, nL0Actual, kL0Actual, initC, + unitFlag); + } + + // Notify to move the next L0B tile + AscendC::SetFlag(l0BEventList[l0BListId]); + l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0; + } + AscendC::SetFlag(l0AEventList[l0AListId]); + l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0; + } + } + l1AListId = l1AListIdNext; + l1BListId = l1BListIdNext; + kL1Actual = kL1ActualNext; + } + + // copy block out + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(l0CEventList[l0CListId]); + AscendC::WaitFlag(l0CEventList[l0CListId]); + copyL0CToDst(tensorC, tensorL0C); + AscendC::SetFlag(l0CEventList[l0CListId]); + l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; + } else { + copyL0CToDst(tensorC, tensorL0C, 0b11); + } + } + + /// Perform a block-scoped matrix multiply-accumulate + template + CATLASS_DEVICE void operator()(TensorA &tensorA, TensorB &tensorB, TensorC tensorCList[MAX_CUBE_VEC_SYNC_NUM], + Catlass::GemmCoord const &actualShape, uint32_t ubRowNum, uint8_t beginSubBlockIdx, + uint64_t l0C2UBAIVAICFlag, uint64_t l0C2UBAICAIVFlag, uint32_t &ubListId, + uint8_t sendVecNum, uint64_t ubNum = 2, TensorBias const &tensorBias = {}) + { + // Check L1TileShape + if constexpr (HAS_BIAS) { + static constexpr uint32_t BIAS_BUF_SIZE = L0_TILE_N * sizeof(ElementAccumulator); + static constexpr uint32_t L1BIAS_SIZE = L1_TILE_N * sizeof(ElementBias); + static_assert(BIAS_BUF_SIZE <= ArchTag::BIAS_SIZE, + "BIAS_BUF_SIZE exceeding the BT space! Reduce L0_TILE_N"); + static_assert(L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES + L1BIAS_SIZE <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + } + + using CopyGmToL1A = typename TileCopy_::template CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::template CopyGmToL1B; + using CopyL0CToDst = typename TileCopy_::template CopyL0CToDst; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; + CopyL0CToDst copyL0CToDst; + + uint32_t mBlockActual = actualShape.m(); + uint32_t kBlockActual = actualShape.k(); + uint32_t nBlockActual = actualShape.n(); + + uint32_t mL1Actual = mBlockActual; + if constexpr (std::is_same_v) { + // Avoid using the gemv mode in mmad + if (mL1Actual == 1) { + mL1Actual = 16; + } + } + uint32_t nL1Actual = nBlockActual; + + auto layoutInL0C = tla::MakeLayoutL0C(mL1Actual, nL1Actual); + auto tensorL0C = tla::MakeTensor(l0CTensorList[l0CListId], layoutInL0C, Catlass::Arch::PositionL0C{}); + auto tensorL0Bias = tla::MakeTensor(l0BiasTensor, L0BIAS_LAYOUT, Catlass::Arch::PositionBias{}); + + uint32_t kL1Actual = min(kBlockActual, L1_TILE_K); + // load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1AListId]); + auto tensorL1A = tla::MakeTensor(l1ATensorList[l1AListId], L1A_LAYOUT, Catlass::Arch::PositionL1{}); + auto tensorTileA = GetTileA(tensorA, 0, 0, mBlockActual, kL1Actual); + if constexpr (ENABLE_L1_RESIDENT) { + // If the currently loaded GM pointer and block coordinates are the same as the last loaded ones, + // skip this loadding. + if (lastAddrA[l1AListId] != tensorTileA.data().GetPhyAddr() || + tla::get<0>(tensorTileA.coord()) != lastCoordA[l1AListId].row() || + tla::get<1>(tensorTileA.coord()) != lastCoordA[l1AListId].column()) { + copyGmToL1A(tensorL1A, tensorTileA); + lastCoordA[l1AListId] = + Catlass::MatrixCoord{tla::get<0>(tensorTileA.coord()), tla::get<1>(tensorTileA.coord())}; + lastAddrA[l1AListId] = const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileA.data().GetPhyAddr()); + } + } else { + copyGmToL1A(tensorL1A, tensorTileA); + } + AscendC::SetFlag(l1AEventList[l1AListId]); + + // load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListId]); + auto tensorL1B = tla::MakeTensor(l1BTensorList[l1BListId], L1B_LAYOUT, Catlass::Arch::PositionL1{}); + auto tensorTileB = GetTile(tensorB, tla::MakeCoord(0, 0), tla::MakeShape(kL1Actual, nBlockActual)); + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrB[l1BListId] != tensorTileB.data().GetPhyAddr() || + tla::get<0>(tensorTileB.coord()) != lastCoordB[l1BListId].row() || + tla::get<1>(tensorTileB.coord()) != lastCoordB[l1BListId].column()) { + copyGmToL1B(tensorL1B, tensorTileB); + lastCoordB[l1BListId] = + Catlass::MatrixCoord{tla::get<0>(tensorTileB.coord()), tla::get<1>(tensorTileB.coord())}; + lastAddrB[l1BListId] = const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileB.data().GetPhyAddr()); + } + } else { + copyGmToL1B(tensorL1B, tensorTileB); + } + AscendC::SetFlag(l1BEventList[l1BListId]); + + if constexpr (HAS_BIAS && !std::is_same_v) { + using CopyGmToL1Bias = typename TileCopy::template CopyGmToL1Bias; + CopyGmToL1Bias copyGmToL1Bias; + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + auto l1Bias = l1BiasTensor.template ReinterpretCast(); + auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Catlass::Arch::PositionL1{}); + copyGmToL1Bias(tensorL1Bias, tensorBias); + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + } + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(l0CEventList[l0CListId]); + } + + uint32_t mL0Loop = CeilDiv(mL1Actual); + uint32_t nL0Loop = CeilDiv(nL1Actual); + + // main loop + uint32_t kL1Loop = CeilDiv(kBlockActual); + for (uint32_t kL1Idx = 0; kL1Idx < kL1Loop; kL1Idx++) { + uint32_t l1AListIdNext = (l1AListId + 1 < L1A_STAGES) ? (l1AListId + 1) : 0; + uint32_t l1BListIdNext = (l1BListId + 1 < L1B_STAGES) ? (l1BListId + 1) : 0; + uint32_t kL1ActualNext{0}; + // preload next tile from GM to L1 + if (kL1Idx < kL1Loop - 1) { + uint32_t kL1IdxNext = kL1Idx + 1; + kL1ActualNext = (kL1IdxNext < kL1Loop - 1) ? L1_TILE_K : (kBlockActual - kL1IdxNext * L1_TILE_K); + + // Get L1 tensor for next stage + auto l1ATensor = l1ATensorList[l1AListIdNext]; + auto l1BTensor = l1BTensorList[l1BListIdNext]; + auto tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Catlass::Arch::PositionL1{}); + auto tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Catlass::Arch::PositionL1{}); + // Get GM tile for next stage + auto tensorTileA = GetTileA(tensorA, 0, kL1IdxNext * L1_TILE_K, mBlockActual, kL1ActualNext); + auto tensorTileB = GetTile(tensorB, tla::MakeCoord(kL1IdxNext * L1_TILE_K, 0), + tla::MakeShape(kL1ActualNext, nBlockActual)); + + // load next matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1AListIdNext]); + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrA[l1AListIdNext] != tensorTileA.data().GetPhyAddr() || + tla::get<0>(tensorTileA.coord()) != lastCoordA[l1AListIdNext].row() || + tla::get<1>(tensorTileA.coord()) != lastCoordA[l1AListIdNext].column()) { + copyGmToL1A(tensorL1A, tensorTileA); + lastCoordA[l1AListIdNext] = + Catlass::MatrixCoord{tla::get<0>(tensorTileA.coord()), tla::get<1>(tensorTileA.coord())}; + lastAddrA[l1AListIdNext] = + const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileA.data().GetPhyAddr()); + } + } else { + copyGmToL1A(tensorL1A, tensorTileA); + } + AscendC::SetFlag(l1AEventList[l1AListIdNext]); + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListIdNext]); + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrB[l1BListIdNext] != tensorTileB.data().GetPhyAddr() || + tla::get<0>(tensorTileB.coord()) != lastCoordB[l1BListIdNext].row() || + tla::get<1>(tensorTileB.coord()) != lastCoordB[l1BListIdNext].column()) { + copyGmToL1B(tensorL1B, tensorTileB); + lastCoordB[l1BListIdNext] = + Catlass::MatrixCoord{tla::get<0>(tensorTileB.coord()), tla::get<1>(tensorTileB.coord())}; + lastAddrB[l1BListIdNext] = + const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileB.data().GetPhyAddr()); + } + } else { + copyGmToL1B(tensorL1B, tensorTileB); + } + AscendC::SetFlag(l1BEventList[l1BListIdNext]); + } + + // Get L1 tensor for current stage + auto l1ATensor = l1ATensorList[l1AListId]; + auto l1BTensor = l1BTensorList[l1BListId]; + tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Catlass::Arch::PositionL1{}); + tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Catlass::Arch::PositionL1{}); + // Get the loop nums on L0 + uint32_t kL0Loop = CeilDiv(kL1Actual); + + for (int mL0Idx = 0; mL0Idx < mL0Loop; mL0Idx++) { + uint32_t mL0Actual = (mL0Idx < mL0Loop - 1) ? L0_TILE_M : (mL1Actual - mL0Idx * L0_TILE_M); + + for (int kL0Idx = 0; kL0Idx < kL0Loop; kL0Idx++) { + uint32_t kL0Actual = (kL0Idx < kL0Loop - 1) ? L0_TILE_K : (kL1Actual - kL0Idx * L0_TILE_K); + + // Locate the current tile on L0A + auto l0ATile = l0ATensorList[l0AListId]; + auto layoutAInL0 = tla::MakeLayout(mL0Actual, kL0Actual); + auto tensorL0A = tla::MakeTensor(l0ATile, layoutAInL0, Catlass::Arch::PositionL0A{}); + // Locate the current tile of matrix A on L1 + auto tensorTileL1A = + GetTileA(tensorL1A, mL0Idx * L0_TILE_M, kL0Idx * L0_TILE_K, mL0Actual, kL0Actual); + + AscendC::WaitFlag(l0AEventList[l0AListId]); + if ((mL0Idx == 0) && (kL0Idx == 0)) { + AscendC::WaitFlag(l1AEventList[l1AListId]); + } + + // Load current tile from L1 to L0A + copyL1ToL0A(tensorL0A, tensorTileL1A); + + if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1)) { + AscendC::SetFlag(l1AEventList[l1AListId]); + } + + bool initC = ((kL1Idx == 0) && (kL0Idx == 0)); + for (int nL0Idx = 0; nL0Idx < nL0Loop; nL0Idx++) { + uint32_t nL0Actual = (nL0Idx < nL0Loop - 1) ? L0_TILE_N : (nL1Actual - nL0Idx * L0_TILE_N); + + // Locate the current tile on L0B + auto l0BTile = l0BTensorList[l0BListId]; + auto layoutBInL0 = tla::MakeLayout(kL0Actual, nL0Actual); + auto tensorL0B = tla::MakeTensor(l0BTile, layoutBInL0, Catlass::Arch::PositionL0B{}); + // Locate the current tile of matrix B on L1 + auto tensorTileL1B = GetTile(tensorL1B, tla::MakeCoord(kL0Idx * L0_TILE_K, nL0Idx * L0_TILE_N), + tla::MakeShape(kL0Actual, nL0Actual)); + + // Wait for mmad finished + AscendC::WaitFlag(l0BEventList[l0BListId]); + // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 + if ((mL0Idx == 0) && (kL0Idx == 0) && (nL0Idx == 0)) { + AscendC::WaitFlag(l1BEventList[l1BListId]); + } + + // Load current tile from L1 to L0B + copyL1ToL0B(tensorL0B, tensorTileL1B); + + // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 + if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1) && (nL0Idx == nL0Loop - 1)) { + AscendC::SetFlag(l1BEventList[l1BListId]); + } + + if constexpr (HAS_BIAS && !std::is_same_v) { + if (initC) { + if (nL0Idx == 0) { + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + } + AscendC::WaitFlag(L0A_STAGES + L0B_STAGES); + auto l1Bias = l1BiasTensor.template ReinterpretCast(); + auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Catlass::Arch::PositionL1{}); + auto tensorTileL1Bias = GetTile(tensorL1Bias, tla::MakeCoord(nL0Idx * L0_TILE_N), + tla::MakeShape(nL0Actual)); + // Load bias to l0 biastable + copyL1ToBT(tensorL0Bias, tensorTileL1Bias); + if (nL0Idx == nL0Loop - 1) { + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + } + } + } + + // Notify to do mmad + AscendC::SetFlag(l0CEventList[l0CListId]); + + // Locate the current tile on L0C + auto tensorTileL0C = GetTile(tensorL0C, tla::MakeCoord(mL0Idx * L0_TILE_M, nL0Idx * L0_TILE_N), + tla::MakeShape(mL0Actual, nL0Actual)); + + // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator + // Wait for loading L0B + AscendC::WaitFlag(l0CEventList[l0CListId]); + + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if ((kL1Idx == kL1Loop - 1) && (mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1) && + (nL0Idx == nL0Loop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + + if constexpr (HAS_BIAS && !std::is_same_v) { + if (initC) { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, tensorL0Bias, mL0Actual, nL0Actual, + kL0Actual, initC, unitFlag); + AscendC::SetFlag(L0A_STAGES + L0B_STAGES); + } else { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, mL0Actual, nL0Actual, kL0Actual, initC, + unitFlag); + } + } else { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, mL0Actual, nL0Actual, kL0Actual, initC, + unitFlag); + } + + // Notify to move the next L0B tile + AscendC::SetFlag(l0BEventList[l0BListId]); + l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0; + } + AscendC::SetFlag(l0AEventList[l0AListId]); + l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0; + } + } + l1AListId = l1AListIdNext; + l1BListId = l1BListIdNext; + kL1Actual = kL1ActualNext; + } + + // copy block out + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(l0CEventList[l0CListId]); + AscendC::WaitFlag(l0CEventList[l0CListId]); + uint32_t leftRowNum = mL1Actual; + uint32_t twoVecRowNum = ubRowNum * sendVecNum; + uint32_t curRowNum = min(twoVecRowNum, leftRowNum); + uint32_t rowIdx = 0; + uint32_t tileNum = CeilDiv(mL1Actual, twoVecRowNum); + for (uint32_t i = 0; i < tileNum; i++) { + auto tensorTileL0C = + GetTile(tensorL0C, tla::MakeCoord(rowIdx, 0), tla::MakeShape(curRowNum, nL1Actual)); + if (sendVecNum == 2) { + AscendC::CrossCoreWaitFlag<0x4, PIPE_FIX>(l0C2UBAIVAICFlag + ubListId); + AscendC::CrossCoreWaitFlag<0x4, PIPE_FIX>(l0C2UBAIVAICFlag + FLAG_ID_MAX + ubListId); + copyL0CToDst(tensorCList[ubListId], tensorTileL0C, ubRowNum, beginSubBlockIdx, sendVecNum); + AscendC::CrossCoreSetFlag<0x4, PIPE_FIX>(l0C2UBAICAIVFlag + ubListId); + AscendC::CrossCoreSetFlag<0x4, PIPE_FIX>(l0C2UBAICAIVFlag + FLAG_ID_MAX + ubListId); + rowIdx += curRowNum; + leftRowNum -= curRowNum; + curRowNum = min(twoVecRowNum, leftRowNum); + ubListId = (ubListId + 1 < ubNum) ? (ubListId + 1) : 0; + } else if (sendVecNum == 1) { + if (beginSubBlockIdx == 0) { + AscendC::CrossCoreWaitFlag<0x4, PIPE_FIX>(l0C2UBAIVAICFlag + ubListId); + } else if (beginSubBlockIdx == 1) { + AscendC::CrossCoreWaitFlag<0x4, PIPE_FIX>(l0C2UBAIVAICFlag + FLAG_ID_MAX + ubListId); + } + copyL0CToDst(tensorCList[ubListId], tensorTileL0C, ubRowNum, beginSubBlockIdx, sendVecNum); + if (beginSubBlockIdx == 0) { + AscendC::CrossCoreSetFlag<0x4, PIPE_FIX>(l0C2UBAICAIVFlag + ubListId); + } else if (beginSubBlockIdx == 1) { + AscendC::CrossCoreSetFlag<0x4, PIPE_FIX>(l0C2UBAICAIVFlag + FLAG_ID_MAX + ubListId); + } + rowIdx += curRowNum; + leftRowNum -= curRowNum; + curRowNum = min(twoVecRowNum, leftRowNum); + ubListId = (ubListId + 1 < ubNum) ? (ubListId + 1) : 0; + } + } + AscendC::SetFlag(l0CEventList[l0CListId]); + l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; + } else { + uint32_t leftRowNum = mL1Actual; + uint32_t twoVecRowNum = ubRowNum * sendVecNum; + uint32_t curRowNum = min(twoVecRowNum, leftRowNum); + uint32_t rowIdx = 0; + uint32_t tileNum = CeilDiv(mL1Actual, twoVecRowNum); + for (uint32_t i = 0; i < tileNum; i++) { + auto tensorTileL0C = + GetTile(tensorL0C, tla::MakeCoord(rowIdx, 0), tla::MakeShape(curRowNum, nL1Actual)); + if (sendVecNum == 2) { + AscendC::CrossCoreWaitFlag<0x4, PIPE_FIX>(l0C2UBAIVAICFlag + ubListId); + AscendC::CrossCoreWaitFlag<0x4, PIPE_FIX>(l0C2UBAIVAICFlag + FLAG_ID_MAX + ubListId); + copyL0CToDst(tensorCList[ubListId], tensorTileL0C, ubRowNum, beginSubBlockIdx, sendVecNum); + AscendC::CrossCoreSetFlag<0x4, PIPE_FIX>(l0C2UBAICAIVFlag + ubListId); + AscendC::CrossCoreSetFlag<0x4, PIPE_FIX>(l0C2UBAICAIVFlag + FLAG_ID_MAX + ubListId); + rowIdx += curRowNum; + leftRowNum -= curRowNum; + curRowNum = min(twoVecRowNum, leftRowNum); + ubListId = (ubListId + 1 < ubNum) ? (ubListId + 1) : 0; + } else if (sendVecNum == 1) { + if (beginSubBlockIdx == 0) { + AscendC::CrossCoreWaitFlag<0x4, PIPE_FIX>(l0C2UBAIVAICFlag + ubListId); + } else if (beginSubBlockIdx == 1) { + AscendC::CrossCoreWaitFlag<0x4, PIPE_FIX>(l0C2UBAIVAICFlag + FLAG_ID_MAX + ubListId); + } + copyL0CToDst(tensorCList[ubListId], tensorTileL0C, ubRowNum, beginSubBlockIdx, sendVecNum); + if (beginSubBlockIdx == 0) { + AscendC::CrossCoreSetFlag<0x4, PIPE_FIX>(l0C2UBAICAIVFlag + ubListId); + } else if (beginSubBlockIdx == 1) { + AscendC::CrossCoreSetFlag<0x4, PIPE_FIX>(l0C2UBAICAIVFlag + FLAG_ID_MAX + ubListId); + } + rowIdx += curRowNum; + leftRowNum -= curRowNum; + curRowNum = min(twoVecRowNum, leftRowNum); + ubListId = (ubListId + 1 < ubNum) ? (ubListId + 1) : 0; + } + } + } + } + +protected: + template + CATLASS_DEVICE auto GetTileA(TensorA &tensorA, uint32_t mIndex, uint32_t kIndex, uint32_t mSize, uint32_t kSize) + { + if constexpr (tla::detail::isVector::value) { + return GetTile(tensorA, tla::MakeCoord(kIndex), tla::MakeShape(kSize)); + } else { + return GetTile(tensorA, tla::MakeCoord(mIndex, kIndex), tla::MakeShape(mSize, kSize)); + } + } + + // Multi-stage tensors list + AscendC::LocalTensor l1ATensorList[L1A_STAGES]; + AscendC::LocalTensor l1BTensorList[L1B_STAGES]; + AscendC::LocalTensor l0ATensorList[L0A_STAGES]; + AscendC::LocalTensor l0BTensorList[L0B_STAGES]; + AscendC::LocalTensor l0CTensorList[L0C_STAGES]; + AscendC::LocalTensor l1BiasTensor; + AscendC::LocalTensor l0BiasTensor; + + // Multi-stage event id list + int32_t l1AEventList[L1A_STAGES]; + int32_t l1BEventList[L1B_STAGES]; + int32_t l0AEventList[L0A_STAGES]; + int32_t l0BEventList[L0B_STAGES]; + int32_t l0CEventList[L0C_STAGES]; + + __gm__ typename AscendC::GlobalTensor::PrimType *lastAddrA[L1A_STAGES]; + __gm__ typename AscendC::GlobalTensor::PrimType *lastAddrB[L1B_STAGES]; + Catlass::MatrixCoord lastCoordA[L1A_STAGES]; + Catlass::MatrixCoord lastCoordB[L1B_STAGES]; + + // The id of current stage + uint32_t l1AListId{0}; + uint32_t l1BListId{0}; + uint32_t l0AListId{0}; + uint32_t l0BListId{0}; + uint32_t l0CListId{0}; + + TileMmad tileMmad; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL1ToBT copyL1ToBT; +}; + +} // namespace Common + +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_KERNEL_UTILS_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/kernel_utils/block/block_mmad_pingpong_tla_multi.hpp b/csrc/chunk_kda_fwd/op_kernel/kernel_utils/block/block_mmad_pingpong_tla_multi.hpp new file mode 100644 index 0000000000..08dd67ff11 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/kernel_utils/block/block_mmad_pingpong_tla_multi.hpp @@ -0,0 +1,677 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_MULTI_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_MULTI_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm { + +template +struct MmadPingpongTlaMulti : public MmadBase { + static constexpr uint32_t L1A_STAGES = L1A_STAGES_; + static constexpr uint32_t L1B_STAGES = L1B_STAGES_; + static constexpr uint32_t L0A_STAGES = L0A_STAGES_; + static constexpr uint32_t L0B_STAGES = L0B_STAGES_; + static constexpr uint32_t L0C_STAGES = L0C_STAGES_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; + static constexpr bool USE_HF32_MODE = USE_HF32_MODE_; + static constexpr bool ENABLE_L1_RESIDENT = ENABLE_L1_RESIDENT_; +}; + +} // namespace Catlass::Gemm + +namespace Catlass::Gemm::Block { + +template +struct BlockMmadTla, + L1TileShape_, L0TileShape_, ElementA_, ElementB_, ElementC_, ElementBias_, TileCopy_, TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = + MmadPingpongTlaMulti; + using ArchTag = typename DispatchPolicy::ArchTag; + using TileCopy = TileCopy_; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = ElementA_; + using LayoutA = typename TileCopy::LayoutA; + using ElementB = ElementB_; + using LayoutB = typename TileCopy::LayoutB; + using ElementC = ElementC_; + using LayoutC = typename TileCopy::LayoutC; + using ElementBias = ElementBias_; + + using TileMmad = TileMmad_; + + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; + using CopyL1ToBT = typename TileCopy::CopyL1ToBT; + + using ElementAccumulator = typename TileCopy::ElementAccumulator; + + static constexpr bool HAS_BIAS = TileCopy::HAS_BIAS; + + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + + using L1AAlignHelper = typename TileCopy_::L1AAlignHelper; + using L1BAlignHelper = typename TileCopy_::L1BAlignHelper; + + static_assert(tla::is_tuple::value && tla::is_static::value, + "L1TileShape must be tla::tuple and static!"); + static_assert(tla::is_tuple::value && tla::is_static::value, + "L0TileShape must be tla::tuple and static!"); + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr bool USE_HF32_MODE = DispatchPolicy::USE_HF32_MODE; + static constexpr bool ENABLE_L1_RESIDENT = DispatchPolicy::ENABLE_L1_RESIDENT; + static constexpr uint32_t L1A_STAGES = DispatchPolicy::L1A_STAGES; + static constexpr uint32_t L1B_STAGES = DispatchPolicy::L1B_STAGES; + static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES; + static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES; + static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; + static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); + static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); + static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); + static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); + static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); + static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); + + // L1 tile size + static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA); + static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); + // L0 tile size + static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L1_TILE_M * L1_TILE_N * sizeof(ElementAccumulator); + + // Check HF32_MODE + static_assert(!USE_HF32_MODE || + (USE_HF32_MODE && std::is_same_v && std::is_same_v), + "HF32 MODE only supports in float!"); + + // Check L0C_STAGES + static_assert(!(ENABLE_UNIT_FLAG && L0C_STAGES != 1), "L0C_STAGES must be 1 when UnitFlag is true!"); + + // Check LayoutC + static_assert(tla::detail::isRowMajor::value || + ((std::is_same_v || std::is_same_v || + std::is_same_v) && + tla::detail::iszN::value), + "LayoutC only supports zN in half or bfloat16 or float, RowMajor in all dtype yet!"); + + // Check L1TileShape + static_assert(L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static_assert(L0A_TILE_SIZE * L0A_STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert(L0B_TILE_SIZE * L0B_STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE * L0C_STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); + + static constexpr uint32_t _32B = 32 * 8; // in bits + static_assert(L1_TILE_M == L0_TILE_M && L1_TILE_N == L0_TILE_N, + "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); + static_assert(L0_TILE_K <= L1_TILE_K, "L0TileShape::K cannot exceed L1TileShape::K"); +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) + static_assert(L1_TILE_M * SizeOfBits::value % _32B == 0, "L1TileShape::M must be 32B aligned."); + static_assert(L1_TILE_K * SizeOfBits::value % _32B == 0, "L1TileShape::K must be 32B aligned."); + static_assert(L1_TILE_K * SizeOfBits::value % _32B == 0, "L1TileShape::K must be 32B aligned."); + static_assert(L1_TILE_N * SizeOfBits::value % _32B == 0, "L1TileShape::N must be 32B aligned."); + static_assert(L0_TILE_K * SizeOfBits::value % _32B == 0, "L0TileShape::K must be 32B aligned."); +#endif + + static_assert((!HAS_BIAS && (L1A_STAGES + L1B_STAGES) <= 8) || (HAS_BIAS && (L1A_STAGES + L1B_STAGES) <= 7), + "L1 Buffer overflow: Exceeds the supported range of EVENT(0~7)"); + + static_assert((!HAS_BIAS && (L0A_STAGES + L0B_STAGES) <= 8) || (HAS_BIAS && (L0A_STAGES + L0B_STAGES) <= 7), + "L0 Buffer overflow: Exceeds the supported range of EVENT_ID(0~7)"); + + static constexpr auto L1A_LAYOUT = + tla::MakeLayout(tla::Int{}, tla::Int{}); + static constexpr auto L1B_LAYOUT = + tla::MakeLayout(tla::Int{}, tla::Int{}); + static constexpr auto L1BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); + static constexpr auto L0BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); + + // When enabling L1 resident mode, restore the pointer and coordinates that record the last state + // to the initial state. if two blockmmad instances need to be consecutively invoked at the kernel layer, + // RestoreStatus() must be inserted between them. + CATLASS_DEVICE + void RestoreStatus() + { + for (int i = 0; i < L1A_STAGES; ++i) { + lastAddrA[i] = nullptr; + lastCoordA[i] = MatrixCoord{0U, 0U}; + } + for (int i = 0; i < L1B_STAGES; ++i) { + lastAddrB[i] = nullptr; + lastCoordB[i] = MatrixCoord{0U, 0U}; + } + } + + /// Construct + CATLASS_DEVICE + BlockMmadTla(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + if ASCEND_IS_AIC { + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * L1A_STAGES; + // Init buffers + for (uint32_t i = 0; i < L1A_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_TILE_SIZE * i); + // Assign event ID for each stages + l1AEventList[i] = i; + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); + // Assign event ID for each stages + l1BEventList[i] = i + L1A_STAGES; + } + for (uint32_t i = 0; i < L0A_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); + // Assign event ID for each stages + l0AEventList[i] = i; + } + for (uint32_t i = 0; i < L0B_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); + // Assign event ID for each stages + l0BEventList[i] = i + L0A_STAGES; + } + if constexpr (!ENABLE_UNIT_FLAG) { + for (uint32_t i = 0; i < L0C_STAGES; i++) { + l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte(L0C_TILE_SIZE * i); + l0CEventList[i] = i; + } + } else { + l0CTensorList[0] = resource.l0CBuf.template GetBufferByByte(0); + } + if constexpr (HAS_BIAS) { + uint32_t l1BiasOffset = l1BOffset + L1B_TILE_SIZE * L1B_STAGES; + l1BiasTensor = resource.l1Buf.template GetBufferByByte(l1BiasOffset); + l0BiasTensor = resource.btBuf.template GetBufferByByte(0); + } + } + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmadTla() {} + + CATLASS_DEVICE + void preSetFlags() + { + if ASCEND_IS_AIC { + // use HF32 when USE_HF32_MODE is true + if constexpr (USE_HF32_MODE) { + AscendC::SetHF32Mode(true); + } else { + AscendC::SetHF32Mode(false); + } + if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { + AscendC::SetMMLayoutTransform(true); + } + for (uint32_t i = 0; i < L1A_STAGES; i++) { + AscendC::SetFlag(l1AEventList[i]); + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + AscendC::SetFlag(l1BEventList[i]); + } + for (uint32_t i = 0; i < L0A_STAGES; i++) { + AscendC::SetFlag(l0AEventList[i]); + } + for (uint32_t i = 0; i < L0B_STAGES; i++) { + AscendC::SetFlag(l0BEventList[i]); + } + if constexpr (!ENABLE_UNIT_FLAG) { + for (uint32_t i = 0; i < L0C_STAGES; i++) { + AscendC::SetFlag(l0CEventList[i]); + } + } + if constexpr (HAS_BIAS) { + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + AscendC::SetFlag(L0A_STAGES + L0B_STAGES); + } + + if constexpr (ENABLE_L1_RESIDENT) { + RestoreStatus(); + } + } + } + + CATLASS_DEVICE + void finalWaitFlags() + { + if ASCEND_IS_AIC { + if constexpr (USE_HF32_MODE) { + AscendC::SetHF32Mode(false); + } + if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { + AscendC::SetMMLayoutTransform(false); + } + for (uint32_t i = 0; i < L1A_STAGES; i++) { + AscendC::WaitFlag(l1AEventList[i]); + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + AscendC::WaitFlag(l1BEventList[i]); + } + for (uint32_t i = 0; i < L0A_STAGES; i++) { + AscendC::WaitFlag(l0AEventList[i]); + } + for (uint32_t i = 0; i < L0B_STAGES; i++) { + AscendC::WaitFlag(l0BEventList[i]); + } + if constexpr (!ENABLE_UNIT_FLAG) { + for (uint32_t i = 0; i < L0C_STAGES; i++) { + AscendC::WaitFlag(l0CEventList[i]); + } + } + if constexpr (HAS_BIAS) { + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + AscendC::WaitFlag(L0A_STAGES + L0B_STAGES); + } + } + } + + /// Perform a block-scoped matrix multiply-accumulate + template + CATLASS_DEVICE void operator()(TensorA &tensorA, TensorB &tensorB, TensorC &tensorC, GemmCoord const &actualShape, + TensorBias const &tensorBias = {}, bool clearL1Padding = false) + { + // Check L1TileShape + if constexpr (HAS_BIAS) { + static constexpr uint32_t BIAS_BUF_SIZE = L0_TILE_N * sizeof(ElementAccumulator); + static constexpr uint32_t L1BIAS_SIZE = L1_TILE_N * sizeof(ElementBias); + static_assert(BIAS_BUF_SIZE <= ArchTag::BIAS_SIZE, + "BIAS_BUF_SIZE exceeding the BT space! Reduce L0_TILE_N"); + static_assert(L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES + L1BIAS_SIZE <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + } + + using CopyGmToL1A = typename TileCopy_::template CopyGmToL1A; + using CopyGmToL1B = typename TileCopy_::template CopyGmToL1B; + CopyGmToL1A copyGmToL1A; + CopyGmToL1B copyGmToL1B; +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) + using CopyL0CToGm = typename TileCopy_::template CopyL0CToGm; + CopyL0CToGm copyL0CToDst; +#endif +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) + using CopyL0CToDst = typename TileCopy_::template CopyL0CToDst; + CopyL0CToDst copyL0CToDst; +#endif + + uint32_t mBlockActual = actualShape.m(); + uint32_t kBlockActual = actualShape.k(); + uint32_t nBlockActual = actualShape.n(); + + uint32_t mL1Actual = mBlockActual; + if constexpr (std::is_same_v) { + // Avoid using the gemv mode in mmad + if (mL1Actual == 1) { + mL1Actual = 16; + } + } + uint32_t nL1Actual = nBlockActual; + + auto layoutInL0C = tla::MakeLayoutL0C(mL1Actual, nL1Actual); + auto tensorL0C = tla::MakeTensor(l0CTensorList[l0CListId], layoutInL0C, Arch::PositionL0C{}); + auto tensorL0Bias = tla::MakeTensor(l0BiasTensor, L0BIAS_LAYOUT, Arch::PositionBias{}); + + uint32_t kL1Actual = min(kBlockActual, L1_TILE_K); + // load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1AListId]); + if (clearL1Padding) { + AscendC::InitConstValueParams clearParams(1, static_cast(L1A_TILE_SIZE / 32), 0, + static_cast(0)); + AscendC::InitConstValue(l1ATensorList[l1AListId], clearParams); + } + auto tensorL1A = tla::MakeTensor(l1ATensorList[l1AListId], L1A_LAYOUT, Arch::PositionL1{}); + auto tensorTileA = GetTileA(tensorA, 0, 0, mBlockActual, kL1Actual); + if constexpr (ENABLE_L1_RESIDENT) { + // If the currently loaded GM pointer and block coordinates are the same as the last loaded ones, + // skip this loading. + if (lastAddrA[l1AListId] != tensorTileA.data().GetPhyAddr() || + tla::get<0>(tensorTileA.coord()) != lastCoordA[l1AListId].row() || + tla::get<1>(tensorTileA.coord()) != lastCoordA[l1AListId].column()) { + copyGmToL1A(tensorL1A, tensorTileA); + lastCoordA[l1AListId] = MatrixCoord{static_cast(tla::get<0>(tensorTileA.coord())), + static_cast(tla::get<1>(tensorTileA.coord()))}; + lastAddrA[l1AListId] = const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileA.data().GetPhyAddr()); + } + } else { + copyGmToL1A(tensorL1A, tensorTileA); + } + AscendC::SetFlag(l1AEventList[l1AListId]); + + // load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListId]); + if (clearL1Padding) { + AscendC::InitConstValueParams clearParams(1, static_cast(L1B_TILE_SIZE / 32), 0, + static_cast(0)); + AscendC::InitConstValue(l1BTensorList[l1BListId], clearParams); + } + auto tensorL1B = tla::MakeTensor(l1BTensorList[l1BListId], L1B_LAYOUT, Arch::PositionL1{}); + auto tensorTileB = GetTile(tensorB, tla::MakeCoord(0, 0), tla::MakeShape(kL1Actual, nBlockActual)); + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrB[l1BListId] != tensorTileB.data().GetPhyAddr() || + tla::get<0>(tensorTileB.coord()) != lastCoordB[l1BListId].row() || + tla::get<1>(tensorTileB.coord()) != lastCoordB[l1BListId].column()) { + copyGmToL1B(tensorL1B, tensorTileB); + lastCoordB[l1BListId] = MatrixCoord{static_cast(tla::get<0>(tensorTileB.coord())), + static_cast(tla::get<1>(tensorTileB.coord()))}; + lastAddrB[l1BListId] = const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileB.data().GetPhyAddr()); + } + } else { + copyGmToL1B(tensorL1B, tensorTileB); + } + AscendC::SetFlag(l1BEventList[l1BListId]); + + if constexpr (HAS_BIAS && !std::is_same_v) { + using CopyGmToL1Bias = typename TileCopy::template CopyGmToL1Bias; + CopyGmToL1Bias copyGmToL1Bias; + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + auto l1Bias = l1BiasTensor.template ReinterpretCast(); + auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Arch::PositionL1{}); + copyGmToL1Bias(tensorL1Bias, tensorBias); + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + } + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(l0CEventList[l0CListId]); + } + + uint32_t mL0Loop = CeilDiv(mL1Actual); + uint32_t nL0Loop = CeilDiv(nL1Actual); + + // main loop + uint32_t kL1Loop = CeilDiv(kBlockActual); + for (uint32_t kL1Idx = 0; kL1Idx < kL1Loop; kL1Idx++) { + uint32_t l1AListIdNext = (l1AListId + 1 < L1A_STAGES) ? (l1AListId + 1) : 0; + uint32_t l1BListIdNext = (l1BListId + 1 < L1B_STAGES) ? (l1BListId + 1) : 0; + uint32_t kL1ActualNext{0}; + // preload next tile from GM to L1 + if (kL1Idx < kL1Loop - 1) { + uint32_t kL1IdxNext = kL1Idx + 1; + kL1ActualNext = (kL1IdxNext < kL1Loop - 1) ? L1_TILE_K : (kBlockActual - kL1IdxNext * L1_TILE_K); + + // Get L1 tensor for next stage + auto l1ATensor = l1ATensorList[l1AListIdNext]; + auto l1BTensor = l1BTensorList[l1BListIdNext]; + auto tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Arch::PositionL1{}); + auto tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Arch::PositionL1{}); + // Get GM tile for next stage + auto tensorTileA = GetTileA(tensorA, 0, kL1IdxNext * L1_TILE_K, mBlockActual, kL1ActualNext); + auto tensorTileB = GetTile(tensorB, tla::MakeCoord(kL1IdxNext * L1_TILE_K, 0), + tla::MakeShape(kL1ActualNext, nBlockActual)); + + // load next matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1AListIdNext]); + if (clearL1Padding) { + AscendC::InitConstValueParams clearParams(1, static_cast(L1A_TILE_SIZE / 32), 0, + static_cast(0)); + AscendC::InitConstValue(l1ATensorList[l1AListIdNext], clearParams); + } + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrA[l1AListIdNext] != tensorTileA.data().GetPhyAddr() || + tla::get<0>(tensorTileA.coord()) != lastCoordA[l1AListIdNext].row() || + tla::get<1>(tensorTileA.coord()) != lastCoordA[l1AListIdNext].column()) { + copyGmToL1A(tensorL1A, tensorTileA); + lastCoordA[l1AListIdNext] = + MatrixCoord{static_cast(tla::get<0>(tensorTileA.coord())), + static_cast(tla::get<1>(tensorTileA.coord()))}; + lastAddrA[l1AListIdNext] = + const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileA.data().GetPhyAddr()); + } + } else { + copyGmToL1A(tensorL1A, tensorTileA); + } + AscendC::SetFlag(l1AEventList[l1AListIdNext]); + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListIdNext]); + if (clearL1Padding) { + AscendC::InitConstValueParams clearParams(1, static_cast(L1B_TILE_SIZE / 32), 0, + static_cast(0)); + AscendC::InitConstValue(l1BTensorList[l1BListIdNext], clearParams); + } + if constexpr (ENABLE_L1_RESIDENT) { + if (lastAddrB[l1BListIdNext] != tensorTileB.data().GetPhyAddr() || + tla::get<0>(tensorTileB.coord()) != lastCoordB[l1BListIdNext].row() || + tla::get<1>(tensorTileB.coord()) != lastCoordB[l1BListIdNext].column()) { + copyGmToL1B(tensorL1B, tensorTileB); + lastCoordB[l1BListIdNext] = + MatrixCoord{static_cast(tla::get<0>(tensorTileB.coord())), + static_cast(tla::get<1>(tensorTileB.coord()))}; + lastAddrB[l1BListIdNext] = + const_cast<__gm__ typename AscendC::GlobalTensor::PrimType *>( + tensorTileB.data().GetPhyAddr()); + } + } else { + copyGmToL1B(tensorL1B, tensorTileB); + } + AscendC::SetFlag(l1BEventList[l1BListIdNext]); + } + + // Get L1 tensor for current stage + auto l1ATensor = l1ATensorList[l1AListId]; + auto l1BTensor = l1BTensorList[l1BListId]; + tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Arch::PositionL1{}); + tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Arch::PositionL1{}); + // Get the loop nums on L0 + uint32_t kL0Loop = CeilDiv(kL1Actual); + + for (int mL0Idx = 0; mL0Idx < mL0Loop; mL0Idx++) { + uint32_t mL0Actual = (mL0Idx < mL0Loop - 1) ? L0_TILE_M : (mL1Actual - mL0Idx * L0_TILE_M); + + for (int kL0Idx = 0; kL0Idx < kL0Loop; kL0Idx++) { + uint32_t kL0Actual = (kL0Idx < kL0Loop - 1) ? L0_TILE_K : (kL1Actual - kL0Idx * L0_TILE_K); + + // Locate the current tile on L0A + auto l0ATile = l0ATensorList[l0AListId]; + auto layoutAInL0 = tla::MakeLayout(mL0Actual, kL0Actual); + auto tensorL0A = tla::MakeTensor(l0ATile, layoutAInL0, Arch::PositionL0A{}); + // Locate the current tile of matrix A on L1 + auto tensorTileL1A = + GetTileA(tensorL1A, mL0Idx * L0_TILE_M, kL0Idx * L0_TILE_K, mL0Actual, kL0Actual); + + AscendC::WaitFlag(l0AEventList[l0AListId]); + if ((mL0Idx == 0) && (kL0Idx == 0)) { + AscendC::WaitFlag(l1AEventList[l1AListId]); + } + + // Load current tile from L1 to L0A + copyL1ToL0A(tensorL0A, tensorTileL1A); + + if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1)) { + AscendC::SetFlag(l1AEventList[l1AListId]); + } + + bool initC = ((kL1Idx == 0) && (kL0Idx == 0)); + for (int nL0Idx = 0; nL0Idx < nL0Loop; nL0Idx++) { + uint32_t nL0Actual = (nL0Idx < nL0Loop - 1) ? L0_TILE_N : (nL1Actual - nL0Idx * L0_TILE_N); + + // Locate the current tile on L0B + auto l0BTile = l0BTensorList[l0BListId]; + auto layoutBInL0 = tla::MakeLayout(kL0Actual, nL0Actual); + auto tensorL0B = tla::MakeTensor(l0BTile, layoutBInL0, Arch::PositionL0B{}); + // Locate the current tile of matrix B on L1 + auto tensorTileL1B = GetTile(tensorL1B, tla::MakeCoord(kL0Idx * L0_TILE_K, nL0Idx * L0_TILE_N), + tla::MakeShape(kL0Actual, nL0Actual)); + + // Wait for mmad finished + AscendC::WaitFlag(l0BEventList[l0BListId]); + // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 + if ((mL0Idx == 0) && (kL0Idx == 0) && (nL0Idx == 0)) { + AscendC::WaitFlag(l1BEventList[l1BListId]); + } + + // Load current tile from L1 to L0B + copyL1ToL0B(tensorL0B, tensorTileL1B); + + // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 + if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1) && (nL0Idx == nL0Loop - 1)) { + AscendC::SetFlag(l1BEventList[l1BListId]); + } + + if constexpr (HAS_BIAS && !std::is_same_v) { + if (initC) { + if (nL0Idx == 0) { + AscendC::WaitFlag(L1A_STAGES + L1B_STAGES); + } + AscendC::WaitFlag(L0A_STAGES + L0B_STAGES); + auto l1Bias = l1BiasTensor.template ReinterpretCast(); + auto tensorL1Bias = tla::MakeTensor(l1Bias, L1BIAS_LAYOUT, Arch::PositionL1{}); + auto tensorTileL1Bias = GetTile(tensorL1Bias, tla::MakeCoord(nL0Idx * L0_TILE_N), + tla::MakeShape(nL0Actual)); + // Load bias to l0 biasTable + copyL1ToBT(tensorL0Bias, tensorTileL1Bias); + if (nL0Idx == nL0Loop - 1) { + AscendC::SetFlag(L1A_STAGES + L1B_STAGES); + } + } + } + + // Notify to do mmad + AscendC::SetFlag(l0CEventList[l0CListId]); + + // Locate the current tile on L0C + auto tensorTileL0C = GetTile(tensorL0C, tla::MakeCoord(mL0Idx * L0_TILE_M, nL0Idx * L0_TILE_N), + tla::MakeShape(mL0Actual, nL0Actual)); + + // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator + // Wait for loading L0B + AscendC::WaitFlag(l0CEventList[l0CListId]); + + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if ((kL1Idx == kL1Loop - 1) && (mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1) && + (nL0Idx == nL0Loop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + + if constexpr (HAS_BIAS && !std::is_same_v) { + if (initC) { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, tensorL0Bias, mL0Actual, nL0Actual, + kL0Actual, initC, unitFlag); + AscendC::SetFlag(L0A_STAGES + L0B_STAGES); + } else { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, mL0Actual, nL0Actual, kL0Actual, initC, + unitFlag); + } + } else { + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, mL0Actual, nL0Actual, kL0Actual, initC, + unitFlag); + } + + // Notify to move the next L0B tile + AscendC::SetFlag(l0BEventList[l0BListId]); + l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0; + } + AscendC::SetFlag(l0AEventList[l0AListId]); + l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0; + } + } + l1AListId = l1AListIdNext; + l1BListId = l1BListIdNext; + kL1Actual = kL1ActualNext; + } + + // copy block out + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(l0CEventList[l0CListId]); + AscendC::WaitFlag(l0CEventList[l0CListId]); + copyL0CToDst(tensorC, tensorL0C); + AscendC::SetFlag(l0CEventList[l0CListId]); + l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; + } else { + copyL0CToDst(tensorC, tensorL0C, 0b11); + } + } + +protected: + template + CATLASS_DEVICE auto GetTileA(TensorA &tensorA, uint32_t mIndex, uint32_t kIndex, uint32_t mSize, uint32_t kSize) + { + if constexpr (tla::detail::isVector::value) { + return GetTile(tensorA, tla::MakeCoord(kIndex), tla::MakeShape(kSize)); + } else { + return GetTile(tensorA, tla::MakeCoord(mIndex, kIndex), tla::MakeShape(mSize, kSize)); + } + } + + // Multi-stage tensors list + AscendC::LocalTensor l1ATensorList[L1A_STAGES]; + AscendC::LocalTensor l1BTensorList[L1B_STAGES]; + AscendC::LocalTensor l0ATensorList[L0A_STAGES]; + AscendC::LocalTensor l0BTensorList[L0B_STAGES]; + AscendC::LocalTensor l0CTensorList[L0C_STAGES]; + AscendC::LocalTensor l1BiasTensor; + AscendC::LocalTensor l0BiasTensor; + + // Multi-stage event id list + int32_t l1AEventList[L1A_STAGES]; + int32_t l1BEventList[L1B_STAGES]; + int32_t l0AEventList[L0A_STAGES]; + int32_t l0BEventList[L0B_STAGES]; + int32_t l0CEventList[L0C_STAGES]; + + __gm__ typename AscendC::GlobalTensor::PrimType *lastAddrA[L1A_STAGES]; + __gm__ typename AscendC::GlobalTensor::PrimType *lastAddrB[L1B_STAGES]; + MatrixCoord lastCoordA[L1A_STAGES]; + MatrixCoord lastCoordB[L1B_STAGES]; + + // The id of current stage + uint32_t l1AListId{0}; + uint32_t l1BListId{0}; + uint32_t l0AListId{0}; + uint32_t l0BListId{0}; + uint32_t l0CListId{0}; + + TileMmad tileMmad; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL1ToBT copyL1ToBT; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_MULTI_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/kernel_utils/block/block_mmad_pingpong_tla_preloadA_l1B.hpp b/csrc/chunk_kda_fwd/op_kernel/kernel_utils/block/block_mmad_pingpong_tla_preloadA_l1B.hpp new file mode 100644 index 0000000000..6b1341c9d9 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/kernel_utils/block/block_mmad_pingpong_tla_preloadA_l1B.hpp @@ -0,0 +1,498 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_PRELOADA_L1B_HPP +#define CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_PRELOADA_L1B_HPP + +#include "catlass/kda_catlass.hpp" +#include "catlass/arch/kda_resource.hpp" +#include "catlass/kda_coord.hpp" +#include "catlass/kda_gemm_coord.hpp" +#include "catlass/gemm/kda_gemm_dispatch_policy.hpp" +#include "catlass/gemm/kda_gemm_helper.hpp" +#include "catlass/gemm/tile/kda_gemm_tile_copy.hpp" +#include "catlass/gemm/tile/kda_tile_mmad.hpp" +#include "tla/kda_tla_layout.hpp" +#include "tla/kda_tensor.hpp" + +namespace Catlass::Gemm { + +template +struct MmadPingpongTlaPreloadAL1B : public MmadBase { + static constexpr uint32_t L1A_STAGES = L1A_STAGES_; + static constexpr uint32_t L1B_STAGES = L1B_STAGES_; + static constexpr uint32_t L0A_STAGES = L0A_STAGES_; + static constexpr uint32_t L0B_STAGES = L0B_STAGES_; + static constexpr uint32_t L0C_STAGES = L0C_STAGES_; + static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_; +}; + +} // namespace Catlass::Gemm + +namespace Catlass::Gemm::Block { + +template +struct BlockMmadTla, + L1TileShape_, L0TileShape_, ElementA_, ElementB_, ElementC_, ElementBias_, TileCopy_, TileMmad_> { +public: + // Type Aliases + using DispatchPolicy = MmadPingpongTlaPreloadAL1B; + using ArchTag = typename DispatchPolicy::ArchTag; + using TileCopy = TileCopy_; + using L1TileShape = L1TileShape_; + using L0TileShape = L0TileShape_; + using ElementA = ElementA_; + using LayoutA = typename TileCopy::LayoutA; + using ElementB = ElementB_; + using LayoutB = typename TileCopy::LayoutB; + using ElementC = ElementC_; + using LayoutC = typename TileCopy::LayoutC; + + using TileMmad = TileMmad_; + + using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A; + using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B; + using CopyL1ToBT = typename TileCopy::CopyL1ToBT; + + using ElementAccumulator = typename TileCopy::ElementAccumulator; + + using LayoutTagL1A = typename TileCopy::LayoutTagL1A; + using LayoutTagL1B = typename TileCopy::LayoutTagL1B; + using LayoutTagL0A = typename TileCopy::LayoutTagL0A; + using LayoutTagL0B = typename TileCopy::LayoutTagL0B; + + using L1AAlignHelper = typename TileCopy_::L1AAlignHelper; + using L1BAlignHelper = typename TileCopy_::L1BAlignHelper; + + static_assert(tla::is_tuple::value && tla::is_static::value, + "L1TileShape must be tla::tuple and static!"); + static_assert(tla::is_tuple::value && tla::is_static::value, + "L0TileShape must be tla::tuple and static!"); + + static constexpr bool ENABLE_UNIT_FLAG = DispatchPolicy::ENABLE_UNIT_FLAG; + static constexpr uint32_t L1A_STAGES = DispatchPolicy::L1A_STAGES; + static constexpr uint32_t L1B_STAGES = DispatchPolicy::L1B_STAGES; + static constexpr uint32_t L0A_STAGES = DispatchPolicy::L0A_STAGES; + static constexpr uint32_t L0B_STAGES = DispatchPolicy::L0B_STAGES; + static constexpr uint32_t L0C_STAGES = DispatchPolicy::L0C_STAGES; + static constexpr uint32_t L1_TILE_M = tla::get<0>(L1TileShape{}); + static constexpr uint32_t L1_TILE_N = tla::get<1>(L1TileShape{}); + static constexpr uint32_t L1_TILE_K = tla::get<2>(L1TileShape{}); + static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape{}); + static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape{}); + static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape{}); + + // L1 tile size + static constexpr uint32_t L1A_TILE_SIZE = L1_TILE_M * L1_TILE_K * sizeof(ElementA); + static constexpr uint32_t L1B_TILE_SIZE = L1_TILE_N * L1_TILE_K * sizeof(ElementB); + // L0 tile size + static constexpr uint32_t L0A_TILE_SIZE = L0_TILE_M * L0_TILE_K * sizeof(ElementA); + static constexpr uint32_t L0B_TILE_SIZE = L0_TILE_K * L0_TILE_N * sizeof(ElementB); + static constexpr uint32_t L0C_TILE_SIZE = L1_TILE_M * L1_TILE_N * sizeof(ElementAccumulator); + + // Check L0C_STAGES + static_assert(!(ENABLE_UNIT_FLAG && L0C_STAGES != 1), "L0C_STAGES must be 1 when UnitFlag is true!"); + + // Check LayoutC + static_assert(tla::detail::isRowMajor::value || + ((std::is_same_v || std::is_same_v || + std::is_same_v) && + tla::detail::iszN::value), + "LayoutC only supports zN in half or bfloat16 or float, RowMajor in all dtype yet!"); + + // Check L1TileShape + static_assert(L1A_TILE_SIZE * L1A_STAGES + L1B_TILE_SIZE * L1B_STAGES <= ArchTag::L1_SIZE, + "L1TileShape exceeding the L1 space!"); + + // Check L0TileShape + static_assert(L0A_TILE_SIZE * L0A_STAGES <= ArchTag::L0A_SIZE, "L0TileShape exceeding the L0A space!"); + static_assert(L0B_TILE_SIZE * L0B_STAGES <= ArchTag::L0B_SIZE, "L0TileShape exceeding the L0B space!"); + static_assert(L0C_TILE_SIZE * L0C_STAGES <= ArchTag::L0C_SIZE, "L0TileShape exceeding the L0C space!"); + + static constexpr uint32_t _32B = 32 * 8; // in bits + static_assert(L1_TILE_M == L0_TILE_M && L1_TILE_N == L0_TILE_N, + "The situation where the basic blocks of L1 and L0 differ on the m and n axes is not supported yet"); + static_assert(L0_TILE_K <= L1_TILE_K, "L0TileShape::K cannot exceed L1TileShape::K"); +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) + static_assert(L1_TILE_M * SizeOfBits::value % _32B == 0, "L1TileShape::M must be 32B aligned."); + static_assert(L1_TILE_K * SizeOfBits::value % _32B == 0, "L1TileShape::K must be 32B aligned."); + static_assert(L1_TILE_K * SizeOfBits::value % _32B == 0, "L1TileShape::K must be 32B aligned."); + static_assert(L1_TILE_N * SizeOfBits::value % _32B == 0, "L1TileShape::N must be 32B aligned."); + static_assert(L0_TILE_K * SizeOfBits::value % _32B == 0, "L0TileShape::K must be 32B aligned."); +#endif + + static constexpr auto L1A_LAYOUT = + tla::MakeLayout(tla::Int{}, tla::Int{}); + static constexpr auto L1B_LAYOUT = + tla::MakeLayout(tla::Int{}, tla::Int{}); + static constexpr auto L1BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); + static constexpr auto L0BIAS_LAYOUT = tla::MakeLayout(tla::Int{}); + + // When enabling L1 resident mode, restore the pointer and coordinates that record the last state + // to the initial state. if two blockmmad instances need to be consecutively invoked at the kernel layer, + // RestoreStatus() must be inserted between them. + CATLASS_DEVICE + void RestoreStatus() + { + for (int i = 0; i < L1A_STAGES; ++i) { + lastAddrA[i] = nullptr; + lastCoordA[i] = MatrixCoord{0U, 0U}; + } + for (int i = 0; i < L1B_STAGES; ++i) { + lastAddrB[i] = nullptr; + lastCoordB[i] = MatrixCoord{0U, 0U}; + } + } + + /// Construct + CATLASS_DEVICE + BlockMmadTla(Arch::Resource &resource, uint32_t l1BufAddrStart = 0) + { + if ASCEND_IS_AIC { + uint32_t l1AOffset = l1BufAddrStart; + uint32_t l1BOffset = l1BufAddrStart + L1A_TILE_SIZE * L1A_STAGES; + // Init buffers + for (uint32_t i = 0; i < L1A_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1ATensorList[i] = resource.l1Buf.template GetBufferByByte(l1AOffset + L1A_TILE_SIZE * i); + // Assign event ID for each stages + l1AEventList[i] = i; + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l1BTensorList[i] = resource.l1Buf.template GetBufferByByte(l1BOffset + L1B_TILE_SIZE * i); + // Assign event ID for each stages + l1BEventList[i] = i + L1A_STAGES; + } + for (uint32_t i = 0; i < L0A_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l0ATensorList[i] = resource.l0ABuf.template GetBufferByByte(L0A_TILE_SIZE * i); + // Assign event ID for each stages + l0AEventList[i] = i; + } + for (uint32_t i = 0; i < L0B_STAGES; i++) { + // Assign L1/L0A/L0B space for each stages + l0BTensorList[i] = resource.l0BBuf.template GetBufferByByte(L0B_TILE_SIZE * i); + // Assign event ID for each stages + l0BEventList[i] = i + L0A_STAGES; + } + if constexpr (!ENABLE_UNIT_FLAG) { + for (uint32_t i = 0; i < L0C_STAGES; i++) { + l0CTensorList[i] = resource.l0CBuf.template GetBufferByByte(L0C_TILE_SIZE * i); + l0CEventList[i] = i; + } + } else { + l0CTensorList[0] = resource.l0CBuf.template GetBufferByByte(0); + } + } + } + + /// Destructor + CATLASS_DEVICE + ~BlockMmadTla() {} + + CATLASS_DEVICE + void preSetFlags() + { + if ASCEND_IS_AIC { + if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { + AscendC::SetMMLayoutTransform(true); + } + for (uint32_t i = 0; i < L1A_STAGES; i++) { + AscendC::SetFlag(l1AEventList[i]); + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + AscendC::SetFlag(l1BEventList[i]); + } + for (uint32_t i = 0; i < L0A_STAGES; i++) { + AscendC::SetFlag(l0AEventList[i]); + } + for (uint32_t i = 0; i < L0B_STAGES; i++) { + AscendC::SetFlag(l0BEventList[i]); + } + if constexpr (!ENABLE_UNIT_FLAG) { + for (uint32_t i = 0; i < L0C_STAGES; i++) { + AscendC::SetFlag(l0CEventList[i]); + } + } + } + } + + CATLASS_DEVICE + void finalWaitFlags() + { + if ASCEND_IS_AIC { + if constexpr (ENABLE_UNIT_FLAG && tla::detail::isRowMajor::value) { + AscendC::SetMMLayoutTransform(false); + } + for (uint32_t i = 0; i < L1A_STAGES; i++) { + AscendC::WaitFlag(l1AEventList[i]); + } + for (uint32_t i = 0; i < L1B_STAGES; i++) { + AscendC::WaitFlag(l1BEventList[i]); + } + for (uint32_t i = 0; i < L0A_STAGES; i++) { + AscendC::WaitFlag(l0AEventList[i]); + } + for (uint32_t i = 0; i < L0B_STAGES; i++) { + AscendC::WaitFlag(l0BEventList[i]); + } + if constexpr (!ENABLE_UNIT_FLAG) { + for (uint32_t i = 0; i < L0C_STAGES; i++) { + AscendC::WaitFlag(l0CEventList[i]); + } + } + } + } + + /// Perform a block-scoped matrix multiply-accumulate + template + CATLASS_DEVICE void operator()(TensorA &tensorA, TensorB &tensorB, TensorC &tensorC, GemmCoord const &actualShape) + { + using CopyGmToL1A = typename TileCopy_::template CopyGmToL1A; + CopyGmToL1A copyGmToL1A; +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) + using CopyL0CToGm = typename TileCopy_::template CopyL0CToGm; + CopyL0CToGm copyL0CToDst; +#endif +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) + using CopyL0CToDst = typename TileCopy_::template CopyL0CToDst; + CopyL0CToDst copyL0CToDst; +#endif + + uint32_t mBlockActual = actualShape.m(); + uint32_t kBlockActual = actualShape.k(); + uint32_t nBlockActual = actualShape.n(); + + uint32_t mL1Actual = mBlockActual; + if constexpr (std::is_same_v) { + // Avoid using the gemv mode in mmad + if (mL1Actual == 1) { + mL1Actual = 16; + } + } + uint32_t nL1Actual = nBlockActual; + + auto layoutInL0C = tla::MakeLayoutL0C(mL1Actual, nL1Actual); + auto tensorL0C = tla::MakeTensor(l0CTensorList[l0CListId], layoutInL0C, Arch::PositionL0C{}); + auto tensorL0Bias = tla::MakeTensor(l0BiasTensor, L0BIAS_LAYOUT, Arch::PositionBias{}); + + uint32_t kL1Actual = min(kBlockActual, L1_TILE_K); + // load first matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1AListId]); + auto tensorL1A = tla::MakeTensor(l1ATensorList[l1AListId], L1A_LAYOUT, Arch::PositionL1{}); + auto tensorTileA = GetTileA(tensorA, 0, 0, mBlockActual, kL1Actual); + copyGmToL1A(tensorL1A, tensorTileA); + AscendC::SetFlag(l1AEventList[l1AListId]); + + // load first matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListId]); + AscendC::SetFlag(l1BEventList[l1BListId]); + + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::WaitFlag(l0CEventList[l0CListId]); + } + + uint32_t mL0Loop = CeilDiv(mL1Actual); + uint32_t nL0Loop = CeilDiv(nL1Actual); + + // main loop + uint32_t kL1Loop = CeilDiv(kBlockActual); + for (uint32_t kL1Idx = 0; kL1Idx < kL1Loop; kL1Idx++) { + uint32_t l1AListIdNext = (l1AListId + 1 < L1A_STAGES) ? (l1AListId + 1) : 0; + uint32_t l1BListIdNext = (l1BListId + 1 < L1B_STAGES) ? (l1BListId + 1) : 0; + uint32_t kL1ActualNext{0}; + // preload next tile from GM to L1 + if (kL1Idx < kL1Loop - 1) { + uint32_t kL1IdxNext = kL1Idx + 1; + kL1ActualNext = (kL1IdxNext < kL1Loop - 1) ? L1_TILE_K : (kBlockActual - kL1IdxNext * L1_TILE_K); + + // Get L1 tensor for next stage + auto l1ATensor = l1ATensorList[l1AListIdNext]; + auto l1BTensor = l1BTensorList[l1BListIdNext]; + auto tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Arch::PositionL1{}); + auto tensorL1B = tla::MakeTensor(l1BTensor, L1B_LAYOUT, Arch::PositionL1{}); + // Get GM tile for next stage + auto tensorTileA = GetTileA(tensorA, 0, kL1IdxNext * L1_TILE_K, mBlockActual, kL1ActualNext); + auto tensorTileB = GetTile(tensorB, tla::MakeCoord(kL1IdxNext * L1_TILE_K, 0), + tla::MakeShape(kL1ActualNext, nBlockActual)); + + // load next matrix A tile from GM to L1 + AscendC::WaitFlag(l1AEventList[l1AListIdNext]); + copyGmToL1A(tensorL1A, tensorTileA); + AscendC::SetFlag(l1AEventList[l1AListIdNext]); + + // load next matrix B tile from GM to L1 + AscendC::WaitFlag(l1BEventList[l1BListIdNext]); + // copyGmToL1B(tensorL1B, tensorTileB); + AscendC::SetFlag(l1BEventList[l1BListIdNext]); + } + + // Get L1 tensor for current stage + auto l1ATensor = l1ATensorList[l1AListId]; + tensorL1A = tla::MakeTensor(l1ATensor, L1A_LAYOUT, Arch::PositionL1{}); + // Get the loop nums on L0 + uint32_t kL0Loop = CeilDiv(kL1Actual); + + for (int mL0Idx = 0; mL0Idx < mL0Loop; mL0Idx++) { + uint32_t mL0Actual = (mL0Idx < mL0Loop - 1) ? L0_TILE_M : (mL1Actual - mL0Idx * L0_TILE_M); + + for (int kL0Idx = 0; kL0Idx < kL0Loop; kL0Idx++) { + uint32_t kL0Actual = (kL0Idx < kL0Loop - 1) ? L0_TILE_K : (kL1Actual - kL0Idx * L0_TILE_K); + + // Locate the current tile on L0A + auto l0ATile = l0ATensorList[l0AListId]; + auto layoutAInL0 = tla::MakeLayout(mL0Actual, kL0Actual); + auto tensorL0A = tla::MakeTensor(l0ATile, layoutAInL0, Arch::PositionL0A{}); + // Locate the current tile of matrix A on L1 + auto tensorTileL1A = + GetTileA(tensorL1A, mL0Idx * L0_TILE_M, kL0Idx * L0_TILE_K, mL0Actual, kL0Actual); + + AscendC::WaitFlag(l0AEventList[l0AListId]); + if ((mL0Idx == 0) && (kL0Idx == 0)) { + AscendC::WaitFlag(l1AEventList[l1AListId]); + } + + // Load current tile from L1 to L0A + copyL1ToL0A(tensorL0A, tensorTileL1A); + + if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1)) { + AscendC::SetFlag(l1AEventList[l1AListId]); + } + + bool initC = ((kL1Idx == 0) && (kL0Idx == 0)); + for (int nL0Idx = 0; nL0Idx < nL0Loop; nL0Idx++) { + uint32_t nL0Actual = (nL0Idx < nL0Loop - 1) ? L0_TILE_N : (nL1Actual - nL0Idx * L0_TILE_N); + + // Locate the current tile on L0B + auto l0BTile = l0BTensorList[l0BListId]; + auto layoutBInL0 = tla::MakeLayout(kL0Actual, nL0Actual); + auto tensorL0B = tla::MakeTensor(l0BTile, layoutBInL0, Arch::PositionL0B{}); + // Locate the current tile of matrix B on L1 + auto tensorTileL1B = GetTile(tensorB, tla::MakeCoord(kL0Idx * L0_TILE_K, nL0Idx * L0_TILE_N), + tla::MakeShape(kL0Actual, nL0Actual)); + + // Wait for mmad finished + AscendC::WaitFlag(l0BEventList[l0BListId]); + // If the current tile is the first one on the k&n axis, wait for loading matrix B from GM to L1 + if ((mL0Idx == 0) && (kL0Idx == 0) && (nL0Idx == 0)) { + AscendC::WaitFlag(l1BEventList[l1BListId]); + } + + // Load current tile from L1 to L0B + copyL1ToL0B(tensorL0B, tensorTileL1B); + + // If the current tile is the last one on the k&n axis, notify to load matrix B from GM to L1 + if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1) && (nL0Idx == nL0Loop - 1)) { + AscendC::SetFlag(l1BEventList[l1BListId]); + } + + // Notify to do mmad + AscendC::SetFlag(l0CEventList[l0CListId]); + + // Locate the current tile on L0C + auto tensorTileL0C = GetTile(tensorL0C, tla::MakeCoord(mL0Idx * L0_TILE_M, nL0Idx * L0_TILE_N), + tla::MakeShape(mL0Actual, nL0Actual)); + + // Compute the matrix multiplication on L0A and L0B and write the result to the accumulator + // Wait for loading L0B + AscendC::WaitFlag(l0CEventList[l0CListId]); + + // If the unit flag is enabled, the unit flag is set according to the calculation progress + uint8_t unitFlag = 0b00; + if constexpr (ENABLE_UNIT_FLAG) { + if ((kL1Idx == kL1Loop - 1) && (mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1) && + (nL0Idx == nL0Loop - 1)) { + unitFlag = 0b11; + } else { + unitFlag = 0b10; + } + } + + tileMmad(tensorTileL0C, tensorL0A, tensorL0B, mL0Actual, nL0Actual, kL0Actual, initC, unitFlag); + + // Notify to move the next L0B tile + AscendC::SetFlag(l0BEventList[l0BListId]); + l0BListId = (l0BListId + 1 < L0B_STAGES) ? (l0BListId + 1) : 0; + } + AscendC::SetFlag(l0AEventList[l0AListId]); + l0AListId = (l0AListId + 1 < L0A_STAGES) ? (l0AListId + 1) : 0; + } + } + l1AListId = l1AListIdNext; + l1BListId = l1BListIdNext; + kL1Actual = kL1ActualNext; + } + + // copy block out + if constexpr (!ENABLE_UNIT_FLAG) { + AscendC::SetFlag(l0CEventList[l0CListId]); + AscendC::WaitFlag(l0CEventList[l0CListId]); + copyL0CToDst(tensorC, tensorL0C); + AscendC::SetFlag(l0CEventList[l0CListId]); + l0CListId = (l0CListId + 1 < L0C_STAGES) ? (l0CListId + 1) : 0; + } else { + copyL0CToDst(tensorC, tensorL0C, 0b11); + } + } + +protected: + template + CATLASS_DEVICE auto GetTileA(TensorA &tensorA, uint32_t mIndex, uint32_t kIndex, uint32_t mSize, uint32_t kSize) + { + if constexpr (tla::detail::isVector::value) { + return GetTile(tensorA, tla::MakeCoord(kIndex), tla::MakeShape(kSize)); + } else { + return GetTile(tensorA, tla::MakeCoord(mIndex, kIndex), tla::MakeShape(mSize, kSize)); + } + } + + // Multi-stage tensors list + AscendC::LocalTensor l1ATensorList[L1A_STAGES]; + AscendC::LocalTensor l1BTensorList[L1B_STAGES]; + AscendC::LocalTensor l0ATensorList[L0A_STAGES]; + AscendC::LocalTensor l0BTensorList[L0B_STAGES]; + AscendC::LocalTensor l0CTensorList[L0C_STAGES]; + AscendC::LocalTensor l1BiasTensor; + AscendC::LocalTensor l0BiasTensor; + + // Multi-stage event id list + int32_t l1AEventList[L1A_STAGES]; + int32_t l1BEventList[L1B_STAGES]; + int32_t l0AEventList[L0A_STAGES]; + int32_t l0BEventList[L0B_STAGES]; + int32_t l0CEventList[L0C_STAGES]; + + __gm__ typename AscendC::GlobalTensor::PrimType *lastAddrA[L1A_STAGES]; + __gm__ typename AscendC::GlobalTensor::PrimType *lastAddrB[L1B_STAGES]; + MatrixCoord lastCoordA[L1A_STAGES]; + MatrixCoord lastCoordB[L1B_STAGES]; + + // The id of current stage + uint32_t l1AListId{0}; + uint32_t l1BListId{0}; + uint32_t l0AListId{0}; + uint32_t l0BListId{0}; + uint32_t l0CListId{0}; + + TileMmad tileMmad; + CopyL1ToL0A copyL1ToL0A; + CopyL1ToL0B copyL1ToL0B; + CopyL1ToBT copyL1ToBT; +}; + +} // namespace Catlass::Gemm::Block + +#endif // CATLASS_GEMM_BLOCK_BLOCK_MMAD_PINGPONG_TLA_PRELOADA_L1B_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/kernel_utils/tile/copy_l0c_to_ub.hpp b/csrc/chunk_kda_fwd/op_kernel/kernel_utils/tile/copy_l0c_to_ub.hpp new file mode 100644 index 0000000000..516e290c9d --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/kernel_utils/tile/copy_l0c_to_ub.hpp @@ -0,0 +1,381 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef COMMON_TILE_ASCEND950_COPY_L0C_TO_UB_950_HPP +#define COMMON_TILE_ASCEND950_COPY_L0C_TO_UB_950_HPP + +#include "catlass/arch/kda_arch.hpp" +#include "catlass/kda_catlass.hpp" +#if defined(CATLASS_ARCH) && CATLASS_ARCH == 3510 +#include "catlass/gemm/tile/ascend950/kda_ascend950_copy_l0c_to_dst.hpp" +#endif +#include "tla/kda_tensor.hpp" +#include "catlass/detail/kda_tag_to_layout.hpp" + +namespace Common::Tile { + +#if defined(CATLASS_ARCH) && CATLASS_ARCH == 3510 + +template +struct CopyL0CToUBTla { + static_assert(DEPENDENT_FALSE, "Unsupported copy l0c to ub, can not find the specialization."); +}; + +#endif + +#if defined(CATLASS_ARCH) && CATLASS_ARCH == 3510 + +template < + /// Tag indicating architecture + class ArchTag, class ElementA_, class LayoutTagA_, class ElementB_, class LayoutTagB_, class ElementC_, + class LayoutTagC_, class ElementBias = void, bool ReluEnable_ = false, + Catlass::Gemm::Tile::ScaleGranularity DEQUANT_GRANULARITY_ = Catlass::Gemm::Tile::ScaleGranularity::NO_QUANT, + class L0CCopyMode = Catlass::Gemm::Tile::CopyToGM> +struct PackedTileCopyTla { + using ElementA = ElementA_; + using ElementB = ElementB_; + using LayoutTagA = LayoutTagA_; + using LayoutTagB = LayoutTagB_; + using LayoutTagC = LayoutTagC_; + using ElementAccumulator = + typename Catlass::Gemm::helper::ElementAccumulatorSelector::ElementAccumulator; + static constexpr bool ReluEnable = ReluEnable_; + static constexpr Catlass::Gemm::Tile::ScaleGranularity DEQUANT_GRANULARITY = DEQUANT_GRANULARITY_; + + static constexpr bool HAS_BIAS = !std::is_void_v; + static constexpr bool HAS_QUANT_TENSOR = + (DEQUANT_GRANULARITY == Catlass::Gemm::Tile::ScaleGranularity::PER_CHANNEL); + + using LayoutTagL1A = + typename Catlass::Gemm::helper::L1ATypeSelector>::L1AType::Layout; + using LayoutTagL1B = + typename Catlass::Gemm::helper::L1BTypeSelector>::L1BType::Layout; + using LayoutTagL0A = typename Catlass::Gemm::helper::L0ALayoutSelector::Layout; + using LayoutTagL0B = Catlass::layout::nZ; + using LayoutTagL0C = Catlass::layout::L0C; + + using LayoutA = Catlass::detail::TagToLayout_t; + using LayoutB = Catlass::detail::TagToLayout_t; + using LayoutC = Catlass::detail::TagToLayout_t; + + using LayoutL1A = Catlass::detail::TagToLayout_t; + using LayoutL1B = Catlass::detail::TagToLayout_t; + using LayoutL0A = Catlass::detail::TagToLayout_t; + using LayoutL0B = Catlass::detail::TagToLayout_t; + using LayoutL0C = typename Catlass::detail::LayoutL0C; + + using TensorL1AVectorLayout = + tla::Tensor, LayoutL1A, tla::Coord, AscendC::TPosition::A1>; + using TensorL1ALayout = + tla::Tensor, LayoutL1A, tla::Coord, AscendC::TPosition::A1>; + + using TensorL1A = + std::conditional_t::value, TensorL1AVectorLayout, TensorL1ALayout>; + using TensorL1B = + tla::Tensor, LayoutL1B, tla::Coord, AscendC::TPosition::A1>; + using TensorL0A = + tla::Tensor, LayoutL0A, tla::Coord, AscendC::TPosition::A2>; + using TensorL0B = + tla::Tensor, LayoutL0B, tla::Coord, AscendC::TPosition::B2>; + using TensorL0C = tla::Tensor, LayoutL0C, tla::Coord, + AscendC::TPosition::CO1>; + using TensorL1Bias = + std::conditional_t, + Catlass::detail::TagToLayout_t, + tla::Coord, AscendC::TPosition::A1>, + Catlass::EmptyClass>; + using TensorL0Bias = tla::Tensor, + Catlass::detail::TagToLayout_t, + tla::Coord, AscendC::TPosition::C2>; + using TensorL1Quant = + std::conditional_t, + Catlass::detail::TagToLayout_t, + tla::Coord, AscendC::TPosition::A1>, + Catlass::EmptyClass>; + + using L1AAlignHelper = Catlass::Gemm::helper::L1AlignHelper; + using L1BAlignHelper = Catlass::Gemm::helper::L1AlignHelper; + + template + using CopyGmToL1A = Catlass::Gemm::Tile::TileCopyTla; + + template + using CopyGmToL1B = Catlass::Gemm::Tile::TileCopyTla; + + template + using CopyGmToL1Bias = + std::conditional_t, + Catlass::EmptyClass>; + + template + using CopyGmToL1Scale = + std::conditional_t, + Catlass::EmptyClass>; + + using CopyL1ToL0A = Catlass::Gemm::Tile::TileCopyTla; + using CopyL1ToL0B = Catlass::Gemm::Tile::TileCopyTla; + using CopyL1ToBT = + std::conditional_t, + Catlass::EmptyClass>; + +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 3510) + template + using CopyL0CToDst = + Catlass::Gemm::Tile::CopyL0CToGmTla; +#endif +#if (defined(CATLASS_ARCH) && CATLASS_ARCH == 2201) + template + using CopyL0CToGm = + Catlass::Gemm::Tile::CopyL0CToGmTla; +#endif +}; + +#if defined(CATLASS_ARCH) && CATLASS_ARCH == 3510 + +template +struct CopyL0CToUBTla< + Catlass::Arch::Ascend950, TensorSrc_, + tla::Tensor, LayoutDst_, CoordDst_, AscendC::TPosition::VECCALC>, + Catlass::Gemm::Tile::CopyL0CToUBMode::NO_SPLIT, Catlass::Gemm::Tile::ScaleGranularity::NO_QUANT, ReluEnable_, + std::enable_if_t::value>> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = typename TensorSrc_::Element; + static constexpr auto quantPre = + Catlass::Gemm::Tile::CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t unitFlag = 0) + { + static_assert( + tla::detail::isRowMajor::value && + TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::VECCALC, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and RowMajor"); + + AscendC::FixpipeParamsC310 intriParams; + + // Fixpipe layout information + intriParams.nSize = tla::get<1>(dstTensor.originShape()); + intriParams.mSize = tla::get<0>(dstTensor.originShape()); + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + intriParams.dstStride = tla::get<0>(dstTensor.stride()); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + // Call AscendC Fixpipe + AscendC::Fixpipe(dstTensor.data()[dstOffset], + srcTensor.data()[srcOffset], intriParams); + } + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint32_t ubRowNum, + uint8_t beginSubBlockIdx, uint8_t sendVecNum, uint8_t unitFlag = 0) + { + static_assert( + tla::detail::isRowMajor::value && + TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::VECCALC, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and RowMajor"); + + AscendC::FixpipeParamsC310 intriParams; + + uint32_t dstTensorRowNum = tla::get<0>(dstTensor.originShape()); + // Fixpipe layout information + intriParams.nSize = tla::get<1>(dstTensor.originShape()); + intriParams.mSize = min(dstTensorRowNum, ubRowNum); + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + intriParams.dstStride = tla::get<0>(dstTensor.stride()); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + intriParams.dualDstCtl = 0; + intriParams.subBlockId = beginSubBlockIdx; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + // Call AscendC Fixpipe + AscendC::Fixpipe(dstTensor.data()[dstOffset], + srcTensor.data()[srcOffset], intriParams); + if (sendVecNum > 1) { + if (dstTensorRowNum > ubRowNum) { + intriParams.mSize = dstTensorRowNum - ubRowNum; + intriParams.subBlockId = (beginSubBlockIdx + 1) % 2; + srcOffset += ubRowNum * tla::get<0, 0>(srcTensor.stride()); + AscendC::Fixpipe(dstTensor.data()[dstOffset], + srcTensor.data()[srcOffset], intriParams); + } + } + } + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t subBlockIdx, + uint8_t unitFlag = 0) + { + static_assert( + tla::detail::isRowMajor::value && + TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::VECCALC, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and RowMajor"); + + AscendC::FixpipeParamsC310 intriParams; + + // Fixpipe layout information + intriParams.nSize = tla::get<1>(dstTensor.shape()); + intriParams.mSize = tla::get<0>(dstTensor.shape()); + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + intriParams.dstStride = tla::get<0>(dstTensor.stride()); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + intriParams.dualDstCtl = 0; + intriParams.subBlockId = subBlockIdx; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + // Call AscendC Fixpipe + AscendC::Fixpipe(dstTensor.data()[dstOffset], + srcTensor.data()[srcOffset], intriParams); + } +}; + +template +struct CopyL0CToUBTla< + Catlass::Arch::Ascend950, TensorSrc_, + tla::Tensor, LayoutDst_, CoordDst_, AscendC::TPosition::VECCALC>, + Catlass::Gemm::Tile::CopyL0CToUBMode::SPLIT_M, Catlass::Gemm::Tile::ScaleGranularity::NO_QUANT, ReluEnable_, + std::enable_if_t::value>> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = typename TensorSrc_::Element; + static constexpr auto quantPre = + Catlass::Gemm::Tile::CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t unitFlag = 0) + { + static_assert( + tla::detail::isRowMajor::value && + TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::VECCALC, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and RowMajor"); + + AscendC::FixpipeParamsC310 intriParams; + + // Fixpipe layout information + intriParams.nSize = tla::get<1>(dstTensor.originShape()); + intriParams.mSize = RoundUp(tla::get<0>(dstTensor.originShape()), 2); // m must be even when spilt m + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + intriParams.dstStride = tla::get<0>(dstTensor.stride()); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + intriParams.dualDstCtl = 1; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + // Call AscendC Fixpipe + AscendC::Fixpipe(dstTensor.data()[dstOffset], + srcTensor.data()[srcOffset], intriParams); + } +}; + +template +struct CopyL0CToUBTla< + Catlass::Arch::Ascend950, TensorSrc_, + tla::Tensor, LayoutDst_, CoordDst_, AscendC::TPosition::VECCALC>, + Catlass::Gemm::Tile::CopyL0CToUBMode::SPLIT_N, Catlass::Gemm::Tile::ScaleGranularity::NO_QUANT, ReluEnable_, + std::enable_if_t::value>> { + using ArchTag = Catlass::Arch::Ascend950; + using ElementDst = ElementDst_; + using ElementSrc = typename TensorSrc_::Element; + static constexpr auto quantPre = + Catlass::Gemm::Tile::CopyL0CToDstQuantMode::VALUE; + static constexpr auto reluEn = ReluEnable_; + + template + CATLASS_DEVICE void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t unitFlag = 0) + { + static_assert( + tla::detail::isRowMajor::value && + TensorSrc::position == AscendC::TPosition::CO1 && TensorDst::position == AscendC::TPosition::VECCALC, + "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and RowMajor"); + + AscendC::FixpipeParamsC310 intriParams; + + // Fixpipe layout information + intriParams.nSize = RoundUp(tla::get<1>(dstTensor.originShape()), 32); + intriParams.mSize = tla::get<0>(dstTensor.originShape()); // m must be even when spilt m + intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride()); + intriParams.dstStride = tla::get<0>(dstTensor.stride()); + + // Fixpipe auxiliary arguments + intriParams.quantPre = quantPre; + intriParams.reluEn = reluEn; + intriParams.unitFlag = unitFlag; + intriParams.dualDstCtl = 2; + + auto dstOffset = dstTensor.layout()(dstTensor.coord()); + auto srcOffset = srcTensor.layout()(srcTensor.coord()); + + // Call AscendC Fixpipe + AscendC::Fixpipe(dstTensor.data()[dstOffset], + srcTensor.data()[srcOffset], intriParams); + } +}; + +#endif + +template < + /// Tag indicating architecture + class ArchTag, class ElementA_, class LayoutTagA, class ElementB_, class LayoutTagB, class ElementC_, + class LayoutTagC, class ElementBias = void, + Catlass::Gemm::Tile::CopyL0CToUBMode CopyMode_ = Catlass::Gemm::Tile::CopyL0CToUBMode::NO_SPLIT, + bool ReluEnable = false, + Catlass::Gemm::Tile::ScaleGranularity DEQUANT_GRANULARITY = Catlass::Gemm::Tile::ScaleGranularity::NO_QUANT> +struct PackedTileCopyTlaToUB : public PackedTileCopyTla { + static constexpr Catlass::Gemm::Tile::CopyL0CToUBMode CopyMode = CopyMode_; + // 重写 CopyL0CToDst + using TensorL0C = typename PackedTileCopyTla::TensorL0C; + + template + using CopyL0CToDst = CopyL0CToUBTla; +}; + +#endif + +///////////////////////////////////////////////////////////////////////////////////////////////////////////// + +} // namespace Common::Tile + +#endif // CATLASS_GEMM_TILE_ASCEND950_COPY_L0C_TO_UB_950_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/kernel_utils/vector/regbase.hpp b/csrc/chunk_kda_fwd/op_kernel/kernel_utils/vector/regbase.hpp new file mode 100644 index 0000000000..0bd2391742 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/kernel_utils/vector/regbase.hpp @@ -0,0 +1,150 @@ +/** + * Copyright (c) 2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +using namespace AscendC::MicroAPI; + +#pragma once + +constexpr static CastTrait ctHalf2Fp32Zero = { + RegLayout::ZERO, + SatMode::SAT, + MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_NONE, +}; +constexpr static CastTrait ctHalf2Fp32One = { + RegLayout::ONE, + SatMode::SAT, + MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_NONE, +}; +constexpr static CastTrait ctFp322HalfZero = {RegLayout::ZERO, SatMode::NO_SAT, MaskMergeMode::MERGING, + AscendC::RoundMode::CAST_ROUND}; +constexpr static CastTrait ctFp322HalfOne = {RegLayout::ONE, SatMode::NO_SAT, MaskMergeMode::ZEROING, + AscendC::RoundMode::CAST_ROUND}; + +constexpr uint16_t PRELOAD_NUM = 2; + +template +__simd_callee__ inline void LoadIn(RegTensor &dstReg, __ubuf__ TType *srcUb) +{ + if constexpr (BroatCast) { + if constexpr (!std::is_same()) { + LoadAlign(dstReg, srcUb); + } else { + LoadAlign(dstReg, srcUb); + } + } else { + LoadAlign(dstReg, srcUb); + } +} + +template +__simd_callee__ inline void CastHalf2Float(RegTensor &dstZeroReg, RegTensor &dstOneReg, + RegTensor &srcReg, MaskReg &mask) +{ + Cast(dstZeroReg, srcReg, mask); + Cast(dstOneReg, srcReg, mask); +} + +template +__simd_callee__ inline void CastFloat2Half(RegTensor &dstReg, RegTensor &srcZeroReg, + RegTensor &srcOneReg, MaskReg &mask) +{ + Cast(dstReg, srcOneReg, mask); + Cast(dstReg, srcZeroReg, mask); +} + +template +__simd_callee__ inline void HalfOrFloat2Float(RegTensor &dstReg, RegTensor &srcReg, MaskReg &maskHalf, + MaskReg &maskFloat) +{ + if constexpr (!std::is_same()) { + Cast(dstReg, srcReg, maskHalf); + } else { + Duplicate(dstReg, srcReg, maskFloat); + } +} + +__simd_callee__ inline void MulFloatTwoReg(RegTensor &dstZeroReg, RegTensor &dstOneReg, + RegTensor &src1ZeroReg, RegTensor &src1OneReg, + RegTensor &src2ZeroReg, RegTensor &src2OneReg, + MaskReg &maskFloat) +{ + Mul(dstZeroReg, src1ZeroReg, src2ZeroReg, maskFloat); + Mul(dstOneReg, src1OneReg, src2OneReg, maskFloat); +} + +__simd_callee__ inline void MinsFloatTwoReg(RegTensor &dstZeroReg, RegTensor &dstOneReg, + RegTensor &srcZeroReg, RegTensor &srcOneReg, + float scalarValue, MaskReg &maskFloat) +{ + Mins(dstZeroReg, srcZeroReg, scalarValue, maskFloat); + Mins(dstOneReg, srcOneReg, scalarValue, maskFloat); +} + +__simd_callee__ inline void ExpFloatTwoReg(RegTensor &dstZeroReg, RegTensor &dstOneReg, + RegTensor &srcZeroReg, RegTensor &srcOneReg, + MaskReg &maskFloat) +{ + Exp(dstZeroReg, srcZeroReg, maskFloat); + Exp(dstOneReg, srcOneReg, maskFloat); +} + +__simd_callee__ inline void SubFloatTwoReg(RegTensor &dstZeroReg, RegTensor &dstOneReg, + RegTensor &src1ZeroReg, RegTensor &src1OneReg, + RegTensor &src2ZeroReg, RegTensor &src2OneReg, + MaskReg &maskFloat) +{ + Sub(dstZeroReg, src1ZeroReg, src2ZeroReg, maskFloat); + Sub(dstOneReg, src1OneReg, src2OneReg, maskFloat); +} + +__simd_callee__ inline void AddFloatTwoReg(RegTensor &dstZeroReg, RegTensor &dstOneReg, + RegTensor &src1ZeroReg, RegTensor &src1OneReg, + RegTensor &src2ZeroReg, RegTensor &src2OneReg, + MaskReg &maskFloat) +{ + Add(dstZeroReg, src1ZeroReg, src2ZeroReg, maskFloat); + Add(dstOneReg, src1OneReg, src2OneReg, maskFloat); +} + +template +__simd_callee__ inline void CompareTwoReg(MaskReg &maskZeroReg, MaskReg &maskOneReg, RegTensor &cmp1ZeroReg, + RegTensor &cmp1OneReg, RegTensor &cmp2ZeroReg, + RegTensor &cmp2OneReg, MaskReg &mask) +{ + Compare(maskZeroReg, cmp1ZeroReg, cmp2ZeroReg, mask); + Compare(maskOneReg, cmp1OneReg, cmp2OneReg, mask); +} + +template +__simd_callee__ inline void SelectTwoReg(RegTensor &dstZeroReg, RegTensor &dstOneReg, + RegTensor &src1ZeroReg, RegTensor &src1OneReg, + RegTensor &src2ZeroReg, RegTensor &src2OneReg, MaskReg &mask1, + MaskReg &mask2) +{ + Select(dstZeroReg, src1ZeroReg, src2ZeroReg, mask1); + Select(dstOneReg, src1OneReg, src2OneReg, mask2); +} + +template +__simd_callee__ inline void StoreUnAlignOut(__ubuf__ TType *dstUb, RegTensor &srcReg, MaskReg &mask, + UnalignRegForStore &uStore, uint32_t copyNum) +{ + RegTensor outZeroReg; + if constexpr (!std::is_same()) { + Cast(outZeroReg, srcReg, mask); + StoreUnAlign(dstUb, outZeroReg, uStore, copyNum); + StoreUnAlignPost(dstUb, uStore, 0); + } else { + StoreUnAlign(dstUb, srcReg, uStore, copyNum); + StoreUnAlignPost(dstUb, uStore, 0); + } +} diff --git a/csrc/chunk_kda_fwd/op_kernel/tla/int_tuple.hpp b/csrc/chunk_kda_fwd/op_kernel/tla/int_tuple.hpp new file mode 100644 index 0000000000..b35081452c --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/tla/int_tuple.hpp @@ -0,0 +1,239 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef TLA_INT_TUPLE_HPP +#define TLA_INT_TUPLE_HPP + +#include "tla/type_traits.hpp" +#include "tla/tuple.hpp" +#include "tla/numeric/integral_constant.hpp" +#include "tla/numeric/integer_sequence.hpp" + +namespace tla { +// +// Apply (Unpack) +// (t, f) => f(t_0,t_1,...,t_n) +// + +namespace detail { + +template +CATLASS_HOST_DEVICE constexpr auto apply(T &&t, F &&f, seq) +{ + return f(get(static_cast(t))...); +} + +template +CATLASS_HOST_DEVICE constexpr auto tapply(T &&t, F &&f, G &&g, seq) +{ + return g(f(get(static_cast(t)))...); +} + +template +CATLASS_HOST_DEVICE constexpr auto tapply(T0 &&t0, T1 &&t1, F &&f, G &&g, seq) +{ + return g(f(get(static_cast(t0)), + get(static_cast(t1)))...); +} + +} // end namespace detail + +template +CATLASS_HOST_DEVICE constexpr auto apply(T &&t, F &&f) +{ + return detail::apply(static_cast(t), f, tuple_seq{}); +} + +template +CATLASS_HOST_DEVICE constexpr auto transform_apply(T &&t, F &&f, G &&g) +{ + if constexpr (is_tuple>::value) { + return detail::tapply(static_cast(t), f, g, tuple_seq{}); + } else { + return g(f(static_cast(t))); + } +} + +struct UnpackedMakeTuple { + template + CATLASS_HOST_DEVICE constexpr auto operator()(T const &...a) const + { + return tla::MakeTuple(a...); + } +}; + +template +CATLASS_HOST_DEVICE constexpr auto transform(T0 const &t0, T1 const &t1, F &&f) +{ + if constexpr (is_tuple::value) { + static_assert(tuple_size::value == tuple_size::value, "Mismatched tuple_size"); + return detail::tapply(t0, t1, f, UnpackedMakeTuple{}, tuple_seq{}); + } else { + return f(t0, t1); + } +} + +template >::value)> +CATLASS_HOST_DEVICE constexpr decltype(auto) get(T &&t) noexcept +{ + static_assert(I == 0, "Index out of range"); + return static_cast(t); +} + +template +CATLASS_HOST_DEVICE constexpr decltype(auto) get(T &&t) noexcept +{ + return get(get(static_cast(t))); +} + +// max +template +CATLASS_HOST_DEVICE constexpr auto max(T0 const &t0, Ts const &...ts); + +struct UnpackedMax { + template + CATLASS_HOST_DEVICE constexpr auto operator()(T const &...v) const + { + return tla::max(v...); + } +}; + +template +CATLASS_HOST_DEVICE constexpr auto max(T0 const &t0, Ts const &...ts) +{ + if constexpr (is_tuple::value) { + return tla::max(tla::apply(t0, UnpackedMax{}), ts...); + } else if constexpr (sizeof...(Ts) == 0) { + return t0; + } else { + return tla::max(t0, tla::max(ts...)); + } +} + +// rank +template +CATLASS_HOST_DEVICE constexpr auto rank(Tuple const &t) +{ + if constexpr (sizeof...(Is) == 0) { + if constexpr (is_tuple::value) { + return Int::value>{}; + } else { + return Int<1>{}; + } + } else { + return rank(get(t)); + } +} + +template +using rank_t = decltype(rank(std::declval())); + +template +constexpr auto rank_v = rank_t::value; + +// depth +template +CATLASS_HOST_DEVICE constexpr auto depth(Tuple const &t); + +struct UnpackedDepth { + template + CATLASS_HOST_DEVICE constexpr auto operator()(T const &...v) const + { + return tla::max(depth(v)...); + } +}; + +template +CATLASS_HOST_DEVICE constexpr auto depth(Tuple const &t) +{ + if constexpr (sizeof...(Is) == 0) { + if constexpr (is_tuple::value) { + return Int<1>{} + tla::apply(t, UnpackedDepth{}); + } else { + return Int<0>{}; + } + } else { + return depth(get(t)); + } +} + +template +using depth_t = decltype(depth(std::declval())); + +template +constexpr auto depth_v = depth_t::value; + +struct MultipliesUnaryLfold { + template + CATLASS_HOST_DEVICE constexpr auto operator()(T const &...v) const + { + return (... * v); + } +}; + +// Implementation of product as a function object +struct Product { + template + CATLASS_HOST_DEVICE constexpr auto operator()(IntTuple const &a) const + { + if constexpr (is_tuple::value) { + if constexpr (tuple_size::value == 0) { + return Int<1>{}; + } else { + return tla::transform_apply(a, Product{}, MultipliesUnaryLfold{}); + } + } else if constexpr (tla::is_integral::value) { + return a; + } + } +}; + +namespace detail { + +template +struct MakeZeroTupleImpl; + +template +struct MakeZeroTupleImpl> { + using type = tla::tuple...>; +}; + +template +using MakeZeroTuple = typename MakeZeroTupleImpl>::type; + +} // end namespace detail + +// Add +template +CATLASS_HOST_DEVICE constexpr auto Add(IntTupleA const &a, IntTupleB const &b); + +struct UnpackedAdd { + template + CATLASS_HOST_DEVICE constexpr auto operator()(IntTupleA const &x, IntTupleB const &y) const + { + return Add(x, y); + } +}; + +template +CATLASS_HOST_DEVICE constexpr auto Add(IntTupleA const &a, IntTupleB const &b) +{ + if constexpr (is_tuple::value && is_tuple::value) { + static_assert(tuple_size::value == tuple_size::value, "Mismatched ranks"); + return transform(a, b, UnpackedAdd{}); + } else { + return tla::add(a, b); + } +} + +} // end namespace tla + +#endif // TLA_INT_TUPLE_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/tla/kda_tensor.hpp b/csrc/chunk_kda_fwd/op_kernel/tla/kda_tensor.hpp new file mode 100644 index 0000000000..3214d1da02 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/tla/kda_tensor.hpp @@ -0,0 +1,413 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef TLA_KDA_TENSOR_HPP +#define TLA_KDA_TENSOR_HPP + +#include "catlass/arch/kda_arch.hpp" +#include "tla/kda_tla_layout.hpp" // tla::Shape +#include "tla/numeric/integral_constant.hpp" // tla::is_integral +#include "tla/int_tuple.hpp" + +namespace tla { + +// +// Underscore slicing utilities +// +namespace detail { + +// Safe "element type at index I" helper that avoids instantiating get for invalid indices. +template +struct coord_elem_type { + using type = void; +}; +template +struct coord_elem_type= 0) && (I < (int)tla::tuple_size>::value)>> { + using type = tla::remove_cvref_t(std::declval()))>; +}; + +template +struct coord_elem_is_underscore : tla::is_underscore::type> {}; + +// Count underscores. +template +struct underscore_count_from : tla::integral_constant {}; +template +struct underscore_count_from= 0)>> { + static constexpr int value = + (coord_elem_is_underscore::value ? 1 : 0) + underscore_count_from::value; +}; +template +struct underscore_count + : tla::integral_constant>::value - 1>::value> {}; + +// Build index sequences for underscore dims (stable 0..R-1 recursion). +template +struct underscore_indices_impl; +template +struct underscore_indices_impl { + using type = seq; +}; +template +struct underscore_indices_impl + : std::conditional_t::value, + underscore_indices_impl, + underscore_indices_impl> {}; + +template +using underscore_indices = + typename underscore_indices_impl>::value>::type; + +// Replace every tla::_ with 0 for offset computation. +template +CATLASS_HOST_DEVICE constexpr decltype(auto) underscore_to_zero(T const &x) +{ + if constexpr (tla::is_underscore::value) { + return tla::_0{}; + } else { + return x; + } +} + +template +CATLASS_HOST_DEVICE constexpr auto replace_underscore_with_zero_impl(Coord const &c, seq) +{ + return tla::MakeCoord(underscore_to_zero(tla::get(c))...); +} + +template +CATLASS_HOST_DEVICE constexpr auto replace_underscore_with_zero(Coord const &c) +{ + static_assert(tla::is_tuple>::value, + "Coord must be tla::tuple for underscore slicing."); + return replace_underscore_with_zero_impl(c, tuple_seq{}); +} + +// Build a layout from selected top-level indices of `layout`. +template +CATLASS_HOST_DEVICE constexpr auto select_layout(Layout const &layout, seq) +{ + auto shape_new = tla::MakeTuple(tla::get(layout.shape())...); + auto stride_new = tla::MakeTuple(tla::get(layout.stride())...); + auto origin_new = tla::MakeTuple(tla::get(layout.originShape())...); + return tla::MakeLayout(shape_new, stride_new, origin_new); +} + +} // namespace detail + +// +// slice_and_offset +// +// A lightweight helper that factors underscore slicing into: +// - a projected layout (keeping only underscored dimensions, in-order), and +// - a base offset computed at the fixed indices (underscored dims treated as 0). +// +// Notes: +// - `coord_arg` must be a one-level `tla::tuple` and contain at least one `tla::_`. +// - This function does not perform runtime bounds checks against originShape(); out-of-bounds is undefined behavior. +// - Returned `offset` is an element offset intended for `BuiltinTensor::operator[](offset)` view creation. +template +CATLASS_HOST_DEVICE constexpr auto slice_and_offset(CoordArg const &coord_arg, + Layout const &layout, + BaseCoord const &base_coord) +{ + static_assert(tla::is_tuple>::value, "slice_and_offset expects a tuple CoordArg."); + static_assert(depth_v == 1, "slice_and_offset only supports one-level CoordArg (no nested tuples)."); + static_assert((int)tla::tuple_size>::value == (int)Layout::rank, + "slice_and_offset requires CoordArg rank == Layout::rank."); + static_assert(tla::is_tuple>::value, "slice_and_offset expects a tuple BaseCoord."); + static_assert((int)tla::tuple_size>::value == (int)Layout::rank, + "slice_and_offset requires BaseCoord rank == Layout::rank."); + + constexpr int k = detail::underscore_count::value; + static_assert(k > 0, "slice_and_offset requires at least one underscore."); + static_assert(k <= Layout::rank, "Invalid underscore count."); + + // Compute base offset using zeros for underscores + auto coord0 = detail::replace_underscore_with_zero(coord_arg); + auto full0 = Add(base_coord, coord0); + auto offset = (int64_t)layout(full0); + + // Determine output dims (underscored dims) and build projected layout. + using Us = detail::underscore_indices; + auto layout_proj = detail::select_layout(layout, Us{}); + + return tla::MakeTuple(layout_proj, offset); +} + +// Convenience overload (no base_coord): assume base_coord == 0. +template +CATLASS_HOST_DEVICE constexpr auto slice_and_offset(CoordArg const &coord_arg, Layout const &layout) +{ + using Z = detail::MakeZeroTuple; + return slice_and_offset(coord_arg, layout, Z{}); +} + +// +// Tensor +// + +namespace detail { + +template +CATLASS_DEVICE constexpr auto HadamardU32(A const &a, B const &b, seq) +{ + return MakeCoord((static_cast(get(a)) * static_cast(get(b)))...); +} + +template +CATLASS_DEVICE constexpr auto GetTileImpl(TensorT const &tensor, CoordT const &coord, ShapeT const &shape, Int) +{ + static_assert(is_tuple::value && depth_v == 1 && rank_v == R, "Coord rank mismatch."); + static_assert(is_tuple::value && depth_v == 1 && rank_v == R, "Shape rank mismatch."); + + auto layoutNew = GetTileLayout(tensor.layout(), shape, coord); + auto coordNew = Add(tensor.coord(), coord); + return MakeTensor(tensor.data(), layoutNew, coordNew, Catlass::Arch::PositionType{}); +} + +template +CATLASS_DEVICE constexpr auto TileViewImpl(TensorT const &tensor, TileCoord const &tileCoord, TileShape const &tileShape, Int) +{ + static_assert(is_tuple::value && depth_v == 1 && rank_v == R, "TileCoord rank mismatch."); + static_assert(is_tuple::value && depth_v == 1 && rank_v == R, "TileShape rank mismatch."); + + auto elementOffset = HadamardU32(tileCoord, tileShape, tuple_seq{}); + auto layoutNew = GetTileLayout(tensor.layout(), tileShape, elementOffset); + auto coordNew = Add(tensor.coord(), elementOffset); + return MakeTensor(tensor.data(), layoutNew, coordNew, Catlass::Arch::PositionType{}); +} + +} // namespace detail + +template +struct Tensor { + using Element = typename BuiltinTensor::PrimType; + using Layout = Layout_; + using Coord = Coord_; + static constexpr AscendC::TPosition position = Position; + + CATLASS_HOST_DEVICE constexpr Tensor() {} + + CATLASS_HOST_DEVICE constexpr Tensor(BuiltinTensor const &builtinTensor, Layout const &layout, Coord const &coord = {}) + : rep_(builtinTensor, layout, coord) + {} + + // + // Accessors + // + + static constexpr int rank = Layout::rank; + + CATLASS_HOST_DEVICE constexpr decltype(auto) tensor() const + { + return *this; + } + + CATLASS_HOST_DEVICE constexpr decltype(auto) data() const + { + return get<0>(rep_); + } + + CATLASS_HOST_DEVICE constexpr decltype(auto) data() + { + return get<0>(rep_); + } + + CATLASS_HOST_DEVICE constexpr decltype(auto) layout() const + { + return get<1>(rep_); + } + + CATLASS_HOST_DEVICE constexpr decltype(auto) coord() const + { + return get<2>(rep_); + } + + CATLASS_HOST_DEVICE constexpr decltype(auto) shape() const + { + return layout().shape(); + } + + CATLASS_HOST_DEVICE constexpr decltype(auto) stride() const + { + return layout().stride(); + } + + CATLASS_HOST_DEVICE constexpr decltype(auto) originShape() const + { + return layout().originShape(); + } + + // + // Indexing / slicing + // + // - No underscore: returns `data()[layout()(coord()+coord_arg)]` + // - Underscores (0..rank, one-level coord): returns a subtensor view over the underscored dimensions (kept in-order) + // Notes: + // - Coord must be one-level (no nested tuples in coord elements). + // - Fixed (non-underscore) indices are expected to be within originShape(). This implementation does not + // generally perform runtime bounds checks/cropping; out-of-bounds indices result in undefined behavior. + template + CATLASS_HOST_DEVICE constexpr decltype(auto) operator()(CoordArg const &coord_arg) const + { + if constexpr (tla::is_tuple>::value) { + static_assert(depth_v == 1, "Underscore slicing only supports one-level Coord (no nested tuples)."); + static_assert(tla::tuple_size>::value == Layout::rank, + "Tensor::operator()(coord): Coord rank must equal tensor rank (Layout::rank)."); + + constexpr int k = detail::underscore_count::value; + if constexpr (k > 0) { + static_assert(k <= Layout::rank, "Invalid underscore count."); + + auto sliced = tla::slice_and_offset(coord_arg, layout(), coord()); + auto layout_proj = tla::get<0>(sliced); + auto offset = (int64_t)tla::get<1>(sliced); + + using CoordZ = detail::MakeZeroTuple<(size_t)k>; + auto data_new = data()[static_cast(offset)]; + return Tensor(data_new, layout_proj, CoordZ{}); + } else { + // No underscore: point view at coord() + coord_arg + auto full = Add(coord(), coord_arg); + return data()[layout()(full)]; + } + } else { + // Scalar coordinate convenience (rank-1): treat it as a 1D coord tuple. + static_assert(Layout::rank == 1, "Tensor::operator()(scalar) is only supported for rank-1 tensors."); + auto full = Add(coord(), MakeCoord(coord_arg)); + return data()[layout()(full)]; + } + } + + template + CATLASS_HOST_DEVICE constexpr decltype(auto) operator()(Coord0 const &c0, Coord1 const &c1, Coords const &...cs) const + { + return operator()(MakeCoord(c0, c1, cs...)); + } + + tla::tuple rep_; +}; + +template +CATLASS_HOST_DEVICE constexpr auto MakeTensor(BuiltinTensor const &builtinTensor, Layout const &layout, PositionType) +{ + using Coord = detail::MakeZeroTuple; + return Tensor(builtinTensor, layout); +} + +template +CATLASS_HOST_DEVICE constexpr auto MakeTensor(BuiltinTensor const &builtinTensor, Layout const &layout, Coord const &coord, PositionType) +{ + return Tensor(builtinTensor, layout, coord); +} + +// Get a tile from tensor, automatically handling boundary cases (tail tile). +// coord is element coordinate; shape is tile size for memory layout. +// The returned Tensor's layout.shape() is used for memory layout calculation, +// and layout.originShape() is the actual logical size (may be smaller than shape). +// Supports tensors of any rank (rank >= 1). +template +CATLASS_DEVICE constexpr auto GetTile(Tensor const &tensor, Coord const &coord, Shape const &shape) +{ + static_assert(Tensor::rank >= 1, "GetTile requires tensor rank >= 1."); + static_assert(Tensor::rank == rank_v && Tensor::rank == rank_v, + "GetTile: coord and shape must have the same rank as the tensor."); + return detail::GetTileImpl(tensor, coord, shape, Int{}); +} + +// 从 tensor 中获取一个 tile，自动处理边界情况（tail tile）。 +// tileCoord 是 tile 单位坐标；tileShape 是用于内存布局的 tile 尺寸。 +// 返回的 Tensor 的 layout.shape() 用于内存布局计算，layout.originShape() 是实际逻辑尺寸（可能小于 tileShape）。 +// TileView(tensor, tileCoord, tileShape) = GetTile(tensor, tileCoord ⊙ tileShape, tileShape) +// Supports tensors of any rank (rank >= 1). +template +CATLASS_DEVICE constexpr auto TileView(TensorT const &tensor, TileCoord const &tileCoord, TileShape const &tileShape) +{ + static_assert(TensorT::rank >= 1, "TileView requires tensor rank >= 1."); + static_assert(TensorT::rank == rank_v && TensorT::rank == rank_v, + "TileView: tileCoord and tileShape must have the same rank as the tensor."); + return detail::TileViewImpl(tensor, tileCoord, tileShape, Int{}); +} + +// 创建一个与另一个 Tensor 类似的 Tensor： +// 目标 layout 根据 LayoutTagDst 构造，从 LikeTensor::Element 推断 ElementDst，从 likeTensor 的 originShape 提取尺寸。 +template +CATLASS_HOST_DEVICE constexpr auto MakeTensorLike(BuiltinTensor const &builtinTensor, + LikeTensor const &likeTensor, + PositionType) +{ + using ElementDst = typename LikeTensor::Element; + static_assert(std::is_same_v, + "BuiltinTensor element type must match LikeTensor element type"); + return MakeTensorLike(builtinTensor, likeTensor, PositionType{}); +} + +// 创建一个与另一个 Tensor 类似的 Tensor： +// 使用 layoutBase 的 shape/stride，但继承 likeTensor 的 originShape，从 LikeTensor::Element 推断 ElementDst。 +template +CATLASS_HOST_DEVICE constexpr auto MakeTensorLike(BuiltinTensor const &builtinTensor, + LikeTensor const &likeTensor, + PositionType, + LayoutBase const &layoutBase) +{ + using ElementDst = typename LikeTensor::Element; + static_assert(std::is_same_v, + "BuiltinTensor element type must match LikeTensor element type"); + return MakeTensorLike(builtinTensor, likeTensor, PositionType{}, layoutBase); +} + +// 创建一个与另一个 Tensor 类似的 Tensor： +// 目标 layout 根据 LayoutTagDst 构造，从 LikeTensor::Element 推断 ElementDst，从 likeTensor 的 originShape 提取尺寸。(调用MakeLayout，可能会因分型布局合法要求对shape进行以分型为粒度的向上取整) +// 允许 BuiltinTensor 的元素类型与 LikeTensor 的元素类型不同（例如 L0C 使用 ElementAccumulator 而不是 ElementC）。 +template +CATLASS_HOST_DEVICE constexpr auto MakeTensorLike(BuiltinTensor const &builtinTensor, + LikeTensor const &likeTensor, + PositionType) +{ + static_assert(LikeTensor::rank == 1 || LikeTensor::rank == 2, + "MakeTensorLike(..., likeTensor) expects rank-1 or rank-2 likeTensor."); + static_assert(std::is_same_v, + "BuiltinTensor element type must match specified ElementDst type"); + // 根据目标布局格式（LayoutTagDst）和指定的元素类型构造 layout + if constexpr (LikeTensor::rank == 1) { + auto layoutNominal = MakeLayout(get<0>(likeTensor.layout().originShape())); + using Coord0 = detail::MakeZeroTuple; + return Tensor(builtinTensor, layoutNominal); + } else { + static_assert(LikeTensor::rank == 2, "MakeTensorLike(..., likeTensor) expects rank-1 or rank-2 likeTensor."); + auto layoutNominal = MakeLayout(get<0>(likeTensor.layout().originShape()), get<1>(likeTensor.layout().originShape())); + using Coord0 = detail::MakeZeroTuple; + return Tensor(builtinTensor, layoutNominal); + } +} + +// 创建一个与另一个 Tensor 类似的 Tensor： +// 使用 layoutBase 的 shape/stride，但继承 likeTensor 的 originShape。允许 BuiltinTensor 的元素类型与 LikeTensor 的元素类型不同。 +template +CATLASS_HOST_DEVICE constexpr auto MakeTensorLike(BuiltinTensor const &builtinTensor, + LikeTensor const &likeTensor, + PositionType, + LayoutBase const &layoutBase) +{ + static_assert(LikeTensor::rank == 1 || LikeTensor::rank == 2, "MakeTensorLike(..., likeTensor, layoutBase) expects rank-1 or rank-2 likeTensor."); + static_assert(std::is_same_v, + "BuiltinTensor element type must match specified ElementDst type"); + + auto layoutFixedStride = MakeLayout(layoutBase.shape(), layoutBase.stride(), likeTensor.originShape()); + using Coord0 = detail::MakeZeroTuple; + return Tensor(builtinTensor, layoutFixedStride); +} + +} // end namespace tla + +#endif // TLA_KDA_TENSOR_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/tla/kda_tla_layout.hpp b/csrc/chunk_kda_fwd/op_kernel/tla/kda_tla_layout.hpp new file mode 100644 index 0000000000..c68c52bfe4 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/tla/kda_tla_layout.hpp @@ -0,0 +1,580 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef KDA_TLA_LAYOUT_HPP +#define KDA_TLA_LAYOUT_HPP + +#include "catlass/kda_catlass.hpp" +#include "tla/numeric/integral_constant.hpp" +#include "tla/numeric/math.hpp" +#include "tla/tuple.hpp" +#include "tla/int_tuple.hpp" +#include "catlass/layout/kda_layout.hpp" + +namespace tla { + +// Aliases + +template +using Shape = tla::tuple; + +template +using Stride = tla::tuple; + +template +using Coord = tla::tuple; + +template +CATLASS_HOST_DEVICE constexpr Shape MakeShape(Ts const &...t) +{ + return {t...}; +} +template +CATLASS_HOST_DEVICE constexpr Stride MakeStride(Ts const &...t) +{ + return {t...}; +} +template +CATLASS_HOST_DEVICE constexpr Coord MakeCoord(Ts const &...t) +{ + return {t...}; +} + +namespace detail { + +// Type trait to generate OriginShape type with rank=R, depth=1, element type=uint32_t +template +struct MakeOriginShapeTypeImpl; + +template +struct MakeOriginShapeTypeImpl> { + template + using repeat_type = uint32_t; + using type = Shape...>; +}; + +template +using MakeOriginShapeType = typename MakeOriginShapeTypeImpl, tla::make_index_sequence>>::type; + +struct UnpackedMakeOriginShapeU32 { + template + CATLASS_HOST_DEVICE constexpr auto operator()(T const &...a) const + { + return MakeShape(static_cast(a)...); + } +}; + +} // namespace detail + +// +// Layout +// +// 自动推导 OriginShape 类型为 Shape，rank=rank_v +template > +struct Layout : private tla::tuple { + // NOTE: This defaults static Shapes/Strides correctly, but not dynamic + CATLASS_HOST_DEVICE constexpr Layout(Shape const &shape = {}, Stride const &stride = {}, OriginShape const &originShape = {}) + : tla::tuple(shape, stride, originShape) + {} + + // + // Accessors + // + + static constexpr int rank = rank_v; + static constexpr int depth = depth_v; + + template + CATLASS_HOST_DEVICE constexpr decltype(auto) shape() + { + return get<0, I...>(static_cast &>(*this)); + } + + template + CATLASS_HOST_DEVICE constexpr decltype(auto) shape() const + { + return get<0, I...>(static_cast const &>(*this)); + } + + template + CATLASS_HOST_DEVICE constexpr decltype(auto) stride() + { + return get<1, I...>(static_cast &>(*this)); + } + + template + CATLASS_HOST_DEVICE constexpr decltype(auto) stride() const + { + return get<1, I...>(static_cast const &>(*this)); + } + + template + CATLASS_HOST_DEVICE constexpr decltype(auto) originShape() + { + return get<2, I...>(static_cast &>(*this)); + } + + template + CATLASS_HOST_DEVICE constexpr decltype(auto) originShape() const + { + return get<2, I...>(static_cast const &>(*this)); + } + + template + CATLASS_HOST_DEVICE constexpr auto operator()(Coord const &coord) const + { + return crd2offset(coord, shape(), stride()); + } +}; + +// Layout construction + +template +CATLASS_HOST_DEVICE constexpr auto MakeLayout(Shape const &shape, Stride const &stride, OriginShape const &originShape) +{ + static_assert(is_tuple::value || is_integral::value); + static_assert(is_tuple::value || is_integral::value); + static_assert(is_tuple::value || is_integral::value); + return Layout(shape, stride, originShape); +} + +template +CATLASS_HOST_DEVICE constexpr auto MakeLayout(Shape const &shape, Stride const &stride) +{ + static_assert(is_tuple::value || is_integral::value); + static_assert(is_tuple::value || is_integral::value); + // 计算默认的 originShape：将 shape 扁平化为 depth=1，并将每个维度归一化为 uint32_t + return MakeLayout(shape, stride, tla::transform_apply(shape, Product{}, detail::UnpackedMakeOriginShapeU32{})); +} + +// Convenience tags for common layouts + +template +CATLASS_HOST_DEVICE constexpr auto MakeLayoutFromTag(LayoutTag const &tag) +{ + static_assert(std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v, + "Unsupported LayoutTag for MakeLayoutFromTag, only support Catlass::layout::RowMajor or" + "Catlass::layout::ColumnMajor or Catlass::layout::VectorLayout or Catlass::layout::zN or Catlass::layout::nZ or Catlass::layout::L0C"); + + if constexpr (std::is_same_v) { + return MakeLayout(MakeShape(tag.shape(0)), + MakeStride(tag.stride(0)), + MakeShape(tag.shape(0))); + } else if constexpr (std::is_same_v) { + return MakeLayout(MakeShape(tag.shape(0), tag.shape(1)), + MakeStride(tag.stride(0), Int<1>{}), + MakeShape(tag.shape(0), tag.shape(1))); + } else if constexpr (std::is_same_v) { + return MakeLayout(MakeShape(tag.shape(0), tag.shape(1)), + MakeStride(Int<1>{}, tag.stride(1)), + MakeShape(tag.shape(0), tag.shape(1))); + } else { // zN or nZ or L0C + return MakeLayout(MakeShape(MakeShape(tag.shape(0), tag.shape(1)), MakeShape(tag.shape(2), tag.shape(3))), + MakeStride(MakeStride(tag.stride(0), tag.stride(1)), MakeStride(tag.stride(2), tag.stride(3))), + MakeShape(tag.orgShape(0), tag.orgShape(1))); + } +} + +// Return the shape of a mode +template +CATLASS_HOST_DEVICE constexpr decltype(auto) shape(Layout &layout) +{ + return layout.template shape(); +} + +template +CATLASS_HOST_DEVICE constexpr decltype(auto) shape(Layout const &layout) +{ + return layout.template shape(); +} + +// Return the stride of a mode +template +CATLASS_HOST_DEVICE constexpr decltype(auto) stride(Layout &layout) +{ + return layout.template stride(); +} + +template +CATLASS_HOST_DEVICE constexpr decltype(auto) stride(Layout const &layout) +{ + return layout.template stride(); +} + +template +CATLASS_HOST_DEVICE constexpr decltype(auto) originShape(Layout &layout) +{ + return layout.template originShape(); +} + +template +CATLASS_HOST_DEVICE constexpr decltype(auto) originShape(Layout const &layout) +{ + return layout.template originShape(); +} + +// Return the rank of layout +template +CATLASS_HOST_DEVICE constexpr auto rank(Layout const &layout) +{ + return rank(shape(layout)); +} + +// Return the depth of the layout +template +CATLASS_HOST_DEVICE constexpr auto depth(Layout const &layout) +{ + return depth(shape(layout)); +} + +// Return the offset of coord +template +CATLASS_HOST_DEVICE constexpr auto crd2offset(Coord const &coord, Shape const &shape, Stride const &stride); + +namespace detail { + +template +CATLASS_HOST_DEVICE constexpr auto crd2offset_ttt(Coord const &coord, Shape const &shape, Stride const &stride, seq) +{ + return (... + crd2offset(get(coord), get(shape), get(stride))); +} + +template +CATLASS_HOST_DEVICE constexpr auto crd2offset_itt(CInt const &coord, STuple const &shape, DTuple const &stride, seq) +{ + if constexpr (sizeof...(Is) == 0) { // Avoid recursion and mod on single/last iter + return crd2offset(coord, get(shape), get(stride)); + } else if constexpr (is_constant<0, CInt>::value) { + return crd2offset(_0{}, get(shape), get(stride)) + + (_0{} + ... + crd2offset(_0{}, get(shape), get(stride))); + } else { // General case + return crd2offset(coord % Product{}(get(shape)), get(shape), get(stride)) + + crd2offset_itt(coord / Product{}(get(shape)), shape, stride, seq{}); + } +} + +} // end namespace detail + +template +CATLASS_HOST_DEVICE constexpr auto crd2offset(Coord const &coord, Shape const &shape, Stride const &stride) +{ + if constexpr (is_tuple::value) { + if constexpr (is_tuple::value) { // tuple tuple tuple + static_assert(tuple_size::value == tuple_size::value, "Mismatched Ranks"); + static_assert(tuple_size::value == tuple_size::value, "Mismatched Ranks"); + return detail::crd2offset_ttt(coord, shape, stride, tuple_seq{}); + } else { // tuple "int" "int" + static_assert(sizeof(Coord) == 0, "Invalid parameters"); + } + } else { + if constexpr (is_tuple::value) { // "int" tuple tuple + static_assert(tuple_size::value == tuple_size::value, "Mismatched Ranks"); + return detail::crd2offset_itt(coord, shape, stride, tuple_seq{}); + } else { // "int" "int" "int" + return coord * stride; + } + } +} + +template +struct is_layout : false_type {}; +template +struct is_layout> : true_type {}; + +// Layout Check +namespace detail { + +template +struct isVector { + static bool const value = false; +}; + +template +struct isVector> { + static bool const value = (stride<0>(Layout{}) == 1); +}; + +template +struct isRowMajor { + static bool const value = false; +}; + +template +struct isRowMajor> { + static bool const value = (stride<1>(Layout{}) == 1); +}; + +template +struct isColumnMajor { + static bool const value = false; +}; + +template +struct isColumnMajor> { + static bool const value = (stride<0>(Layout{}) == 1); +}; + +template +struct iszN { + static bool const value = false; +}; + +template +struct iszN> { + static constexpr uint32_t ELE_NUM_PER_C0 = Catlass::BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = Catlass::BYTE_PER_FRACTAL / sizeof(Element); + static bool const value = (shape<0, 0>(Layout{}) == Catlass::C0_NUM_PER_FRACTAL && + shape<1, 0>(Layout{}) == ELE_NUM_PER_C0 && + stride<1, 0>(Layout{}) == 1 && + stride<0, 1>(Layout{}) == ELE_NUM_PER_FRACTAL); +}; + +/* +For matmul m axis is not c0 Aligned. +Exp: oriShape(m, k) : (127, 256) +zNUnAlign shape:((127, 1), (16, 256/16)) zN shape: ((16, Ceil(127/16)), (16, 256/16)) +*/ +template +struct iszNUnAlign { + static bool const value = false; +}; + +template +struct iszNUnAlign, std::enable_if_t(Layout{}))> == 2 && rank_v(Layout{}))> == 2>> { + static constexpr uint32_t ELE_NUM_PER_C0 = Catlass::BYTE_PER_C0 / sizeof(Element); + static bool const value = (shape<0, 1>(Layout{}) == 1 && + shape<1, 0>(Layout{}) == ELE_NUM_PER_C0 && + stride<0, 0>(Layout{}) == ELE_NUM_PER_C0 && + stride<1, 0>(Layout{}) == 1); +}; + +template +struct iszZ { + static bool const value = false; +}; + +template +struct iszZ> { + static constexpr uint32_t ELE_NUM_PER_C0 = Catlass::BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = Catlass::BYTE_PER_FRACTAL / sizeof(Element); + static bool const value = (shape<0, 0>(Layout{}) == Catlass::C0_NUM_PER_FRACTAL && + shape<1, 0>(Layout{}) == ELE_NUM_PER_C0 && + stride<1, 0>(Layout{}) == 1 && + stride<1, 1>(Layout{}) == ELE_NUM_PER_FRACTAL); +}; + +template +struct isnZ { + static bool const value = false; +}; + +template +struct isnZ> { + static constexpr uint32_t ELE_NUM_PER_C0 = Catlass::BYTE_PER_C0 / sizeof(Element); + static constexpr uint32_t ELE_NUM_PER_FRACTAL = Catlass::BYTE_PER_FRACTAL / sizeof(Element); + static bool const value = (shape<0, 0>(Layout{}) == ELE_NUM_PER_C0 && + shape<1, 0>(Layout{}) == Catlass::C0_NUM_PER_FRACTAL && + stride<0, 0>(Layout{}) == 1 && + stride<1, 1>(Layout{}) == ELE_NUM_PER_FRACTAL); +}; + +} // end namespace detail + +// Advanced Layout constructions + +// Make a vector layout. +template +CATLASS_HOST_DEVICE constexpr auto MakeLayout(T const &len) +{ + return MakeLayout(MakeShape(len), MakeStride(Int<1>{}), MakeShape(len)); +} + +// Make a inner layout with Rows and Cols. +template +CATLASS_HOST_DEVICE constexpr auto MakeLayout(T const &rows, U const &cols) +{ + static_assert(std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v || + std::is_same_v, + "Unsupported LayoutTag for MakeLayoutFromTag, only support Catlass::layout::RowMajor or" + "Catlass::layout::ColumnMajor or Catlass::layout::zN or Catlass::layout::nZ or Catlass::layout::zZ or Catlass::layout::L0C"); + + constexpr uint32_t ELE_NUM_PER_C0 = Catlass::BYTE_PER_C0 / sizeof(Element); + constexpr uint32_t ELE_NUM_PER_FRACTAL = Catlass::BYTE_PER_FRACTAL / sizeof(Element); + + if constexpr (std::is_same_v) { + return MakeLayout(MakeShape(cols), MakeStride(Int<1>{}), MakeShape(cols)); + } else if constexpr (std::is_same_v) { + return MakeLayout(MakeShape(rows, cols), + MakeStride((int64_t)cols, Int<1>{}), + MakeShape(rows, cols)); + } else if constexpr (std::is_same_v) { + return MakeLayout(MakeShape(rows, cols), + MakeStride(Int<1>{}, (int64_t)rows), + MakeShape(rows, cols)); + } else if constexpr (std::is_same_v) { + return MakeLayout( + MakeShape(MakeShape(Int{}, CeilDiv(rows, Int{})), + MakeShape(Int{}, CeilDiv(cols, Int{}))), + MakeStride(MakeStride(Int{}, Int{}), + MakeStride(Int<1>{}, RoundUp((int64_t)rows, Int{}) * ELE_NUM_PER_C0)), + MakeShape(rows, cols)); + } else if constexpr (std::is_same_v) { + return MakeLayout( + MakeShape(MakeShape(Int{}, CeilDiv(rows, Int{})), + MakeShape(Int{}, CeilDiv(cols, Int{}))), + MakeStride(MakeStride(Int{}, + RoundUp((int64_t)cols, Int{}) * Catlass::C0_NUM_PER_FRACTAL), + MakeStride(Int<1>{}, Int{})), + MakeShape(rows, cols)); + } else if constexpr (std::is_same_v) { + constexpr uint32_t ELE_NUM_PER_FRACTAL = 256; + return MakeLayout( + MakeShape(MakeShape(Int{}, CeilDiv(rows, Int{})), + MakeShape(Int{}, CeilDiv(cols, Int{}))), + MakeStride(MakeStride(Int{}, Int{}), + MakeStride(Int<1>{}, RoundUp((int64_t)rows, Int{}) * Catlass::C0_NUM_PER_FRACTAL)), + MakeShape(rows, cols)); + } else { + return MakeLayout( + MakeShape(MakeShape(Int{}, CeilDiv(rows, Int{})), + MakeShape(Int{}, CeilDiv(cols, Int{}))), + MakeStride( + MakeStride(Int<1>{}, RoundUp((int64_t)cols, Int{}) * ELE_NUM_PER_C0), + MakeStride(Int{}, Int{})), + MakeShape(rows, cols)); + } +} + +namespace detail { + +template +CATLASS_HOST_DEVICE constexpr auto CropOriginShape(OriginBase const &originBase, TileShape const &tileShape, Coord const &coord, seq) +{ + return MakeShape( + tla::min( + static_cast(get(tileShape)), + (static_cast(get(coord)) < static_cast(get(originBase))) ? (static_cast(get(originBase)) - static_cast(get(coord))) : 0u)...); +} + +} // namespace detail + +/// 创建 tile layout：使用指定的 tile 尺寸用于内存布局计算，同时携带实际逻辑尺寸（origin_shape）。 +/// coord 是元素坐标，用于计算实际的 originShape（处理边界情况）。 +/// Supports layouts of any rank (rank >= 1) for depth==1 layouts. +/// For depth>1 (fractal) layouts, currently only rank-2 is supported. +template +CATLASS_HOST_DEVICE constexpr auto GetTileLayout(Layout const &layout, TileShape const &tileShape, Coord const &coord) +{ + static_assert(is_tuple::value && depth_v == 1 && rank_v >= 1, + "GetTileLayout: TileShape must be a flat tuple with rank >= 1."); + static_assert(is_tuple::value && depth_v == 1 && rank_v == rank_v, + "GetTileLayout: Coord must have the same rank as TileShape."); + + // 统一计算 tail tile 的逻辑尺寸（originShape 裁剪） + auto tileOriginShape = detail::CropOriginShape(layout.originShape(), tileShape, coord, tuple_seq{}); + + // depth==1 的布局（vector/matrix/tensor）：tile shape 直接作为 memory-layout shape + // 支持任意 rank >= 1（但必须与 layout.rank 匹配） + if constexpr (Layout::depth == 1) { + static_assert(Layout::rank == rank_v, + "GetTileLayout: for depth==1 layouts, TileShape rank must match layout rank."); + return MakeLayout(tileShape, layout.stride(), tileOriginShape); + } else { + // depth>1 的布局（fractal layout）：目前只支持 rank=2 + // 因为 fractal layout 通常用于矩阵（rank-2），需要把 (rows, cols) 转为同结构嵌套 shape + static_assert(rank_v == 2, + "GetTileLayout: for depth>1 (fractal) layouts, TileShape must be rank-2 (rows, cols)."); + + if constexpr (is_static(layout))>::value && + is_static(layout))>::value) { + const uint32_t rows = get<0>(tileShape); + const uint32_t cols = get<1>(tileShape); + constexpr uint32_t dstInnerShapeRow = decltype(shape<0, 0>(layout))::value; + constexpr uint32_t dstInnerShapeCol = decltype(shape<1, 0>(layout))::value; + return MakeLayout( + MakeShape(MakeShape(Int{}, CeilDiv(rows)), + MakeShape(Int{}, CeilDiv(cols))), + layout.stride(), + tileOriginShape); + } else { + const uint32_t rows = get<0>(tileShape); + const uint32_t cols = get<1>(tileShape); + const uint32_t dstInnerShapeRow = shape<0, 0>(layout); + const uint32_t dstInnerShapeCol = shape<1, 0>(layout); + return MakeLayout( + MakeShape(MakeShape(dstInnerShapeRow, CeilDiv(rows, dstInnerShapeRow)), + MakeShape(dstInnerShapeCol, CeilDiv(cols, dstInnerShapeCol))), + layout.stride(), + tileOriginShape); + } + } +} + +template +CATLASS_HOST_DEVICE constexpr auto MakeLayoutL0C(T const &rows, U const &cols) +{ + constexpr uint32_t ELE_NUM_PER_FRACTAL = 256; + return MakeLayout( + MakeShape(MakeShape(Int{}, CeilDiv(rows, Int{})), + MakeShape(Int{}, CeilDiv(cols, Int{}))), + MakeStride(MakeStride(Int{}, Int{}), + MakeStride(Int<1>{}, RoundUp((int64_t)rows, Int{}) * Catlass::C0_NUM_PER_FRACTAL)), + MakeShape(rows, cols)); +} + +// +// Layout transforms +// +namespace detail { + +// Prepend one leading dimension to a layout type +// This is the general form of "make batched layout": +// - Given a base layout type of rank R, create a new layout of rank R+1 by +// prefixing (batchShape, batchStride, batchOrigin) to the existing shape/stride/originShape. +// intentionally preserve the *types* of each stride element from the original layout +template > +struct PrependDimLayout; + +template +struct PrependDimLayout> { + using ShapeOld = tla::remove_cvref_t().shape())>; + using StrideOld = tla::remove_cvref_t().stride())>; + using OriginOld = tla::remove_cvref_t().originShape())>; + + using ShapeNew = tla::Shape(std::declval()))>...>; + using StrideNew = tla::Stride(std::declval()))>...>; + using OriginNew = tla::Shape(std::declval()))>...>; + + using type = tla::Layout; +}; + +} // namespace detail + +template +using PrependDimLayout_t = typename detail::PrependDimLayout::type; + +template +using MakeBatchedLayout_t = PrependDimLayout_t; + +} // end namespace tla + +#endif // KDA_TLA_LAYOUT_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/tla/numeric/integer_sequence.hpp b/csrc/chunk_kda_fwd/op_kernel/tla/numeric/integer_sequence.hpp new file mode 100644 index 0000000000..2bd81f0fce --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/tla/numeric/integer_sequence.hpp @@ -0,0 +1,67 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef TLA_NUMERIC_INTEGER_SEQUENCE_HPP +#define TLA_NUMERIC_INTEGER_SEQUENCE_HPP + +#include "tla/numeric/integral_constant.hpp" +#include "tla/type_traits.hpp" + +namespace tla { + +template +struct IntegerSequence { + using value_type = T; + static constexpr size_t size() { return sizeof...(Ns); } +}; + +template +struct MakeIntegerSequenceImpl; + +template +struct MakeIntegerSequenceImpl, T, 0> { + typedef IntegerSequence type; +}; + +template +struct MakeIntegerSequenceImpl, T, N> { + typedef typename MakeIntegerSequenceImpl, T, N - 1>::type type; +}; + +template +using MakeIntegerSequence = typename MakeIntegerSequenceImpl, T, N>::type; + +// index_sequence +template +using index_sequence = IntegerSequence; + +template +using make_index_sequence = MakeIntegerSequence; + +// int_sequence +template +using int_sequence = IntegerSequence; + +template +using make_int_sequence = MakeIntegerSequence; + +// Shortcuts +template +using seq = int_sequence; + +template +using make_seq = make_int_sequence; + +template +using tuple_seq = make_seq>::value>; + +} // end namespace tla + +#endif // TLA_NUMERIC_INTEGER_SEQUENCE_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/tla/numeric/integral_constant.hpp b/csrc/chunk_kda_fwd/op_kernel/tla/numeric/integral_constant.hpp new file mode 100644 index 0000000000..ea18b263d4 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/tla/numeric/integral_constant.hpp @@ -0,0 +1,196 @@ +/** + * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef TLA_NUMERIC_INTEGER_CONSTANT_HPP +#define TLA_NUMERIC_INTEGER_CONSTANT_HPP + +#include "catlass/detail/kda_macros.hpp" +#include "tla/type_traits.hpp" +#include "tla/numeric/math.hpp" + +namespace tla { + +// A constant value: short name and type-deduction for fast compilation +template +struct C { + using type = C; + static constexpr auto value = v; + using value_type = decltype(v); + CATLASS_HOST_DEVICE constexpr operator value_type() const noexcept { return value; } + CATLASS_HOST_DEVICE constexpr value_type operator()() const noexcept { return value; } +}; + +// Deprecate +template +using constant = C; + +template +using bool_constant = C; + +using true_type = bool_constant; +using false_type = bool_constant; + +template +using is_std_integral = std::is_integral; + +// A more std:: conforming integral_constant that enforces type but interops with C +template +struct integral_constant : C { + using type = integral_constant; + static constexpr T value = v; + using value_type = T; + CATLASS_HOST_DEVICE constexpr value_type operator()() const noexcept { return value; } +}; + +// Use tla::is_std_integral to match built-in integral types (int, int64_t, unsigned, etc) +// Use tla::is_integral to match both built-in integral types AND static integral types. + +template +struct is_integral : bool_constant::value> {}; +template +struct is_integral> : true_type {}; +template +struct is_integral> : true_type {}; + +// is_static detects if an (abstract) value is defined completely by its type (no members) +template +struct is_static : bool_constant>::value> {}; + +// is_constant detects if a type is a static integral type and if v is equal to a value + +template +struct is_constant : false_type {}; +template +struct is_constant : is_constant {}; +template +struct is_constant : is_constant {}; +template +struct is_constant : is_constant {}; +template +struct is_constant : is_constant {}; +template +struct is_constant> : bool_constant {}; +template +struct is_constant> : bool_constant {}; + +// +// Specializations +// + +template +using Int = C; +using _0 = Int<0>; +using _64 = Int<64>; +using _128 = Int<128>; +using _256 = Int<256>; +using _512 = Int<512>; + +// +// Underscore placeholder (for slicing semantics) +// +// Usage: +// - `tla::_` is an empty tag value that can be used inside `tla::Coord` / tensor indexing +// to indicate "take the whole dimension" (full slice). +struct Underscore { + using type = Underscore; +}; + +CATLASS_HOST_DEVICE constexpr Underscore _{}; + +template +struct is_underscore : false_type {}; +template <> +struct is_underscore : true_type {}; +template +struct is_underscore : is_underscore {}; +template +struct is_underscore : is_underscore {}; +template +struct is_underscore : is_underscore {}; +template +struct is_underscore : is_underscore {}; + +/***************/ +/** Operators **/ +/***************/ + +#define TLA_LEFT_UNARY_OP(OP) \ + template \ + CATLASS_HOST_DEVICE constexpr C<(OP t)> operator OP(C) \ + { \ + return {}; \ + } +#define TLA_BINARY_OP(OP) \ + template \ + CATLASS_HOST_DEVICE constexpr C<(t OP u)> operator OP(C, C) \ + { \ + return {}; \ + } + +TLA_LEFT_UNARY_OP(+); +TLA_LEFT_UNARY_OP(-); +TLA_LEFT_UNARY_OP(~); +TLA_LEFT_UNARY_OP(!); +TLA_LEFT_UNARY_OP(*); + +TLA_BINARY_OP(+); +TLA_BINARY_OP(-); +TLA_BINARY_OP(*); +TLA_BINARY_OP(/); +TLA_BINARY_OP(%); +TLA_BINARY_OP(&); +TLA_BINARY_OP(|); +TLA_BINARY_OP(^); +TLA_BINARY_OP(<<); +TLA_BINARY_OP(>>); + +#undef TLA_BINARY_OP +#undef TLA_LEFT_UNARY_OP +#undef TLA_RIGHT_UNARY_OP + +// +// Named functions from math.hpp +// + +#define TLA_NAMED_UNARY_FN(OP) \ + template \ + CATLASS_HOST_DEVICE constexpr auto OP(C) \ + { \ + return C{}; \ + } +#define TLA_NAMED_BINARY_FN(OP) \ + template \ + CATLASS_HOST_DEVICE constexpr auto OP(C, C) \ + { \ + return C{}; \ + } \ + template ::value)> \ + CATLASS_HOST_DEVICE constexpr auto OP(C, U u) \ + { \ + return OP(t, u); \ + } \ + template ::value)> \ + CATLASS_HOST_DEVICE constexpr auto OP(T t, C) \ + { \ + return OP(t, u); \ + } + +TLA_NAMED_BINARY_FN(max); +TLA_NAMED_BINARY_FN(min); +TLA_NAMED_BINARY_FN(add); + +#undef TLA_NAMED_UNARY_FN +#undef TLA_NAMED_BINARY_FN + +} // end namespace tla + +#endif // TLA_NUMERIC_INTEGER_CONSTANT_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/tla/numeric/math.hpp b/csrc/chunk_kda_fwd/op_kernel/tla/numeric/math.hpp new file mode 100644 index 0000000000..00ff91957b --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/tla/numeric/math.hpp @@ -0,0 +1,52 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef TLA_NUMERIC_MATH_HPP +#define TLA_NUMERIC_MATH_HPP + +#include "catlass/detail/kda_macros.hpp" +#include "tla/type_traits.hpp" + +namespace tla { + +// +// Common Operations +// + +template ::value && + std::is_arithmetic::value)> +CATLASS_HOST_DEVICE constexpr auto +max(T const &t, U const &u) +{ + return t < u ? u : t; +} + +template ::value && + std::is_arithmetic::value)> +CATLASS_HOST_DEVICE constexpr auto +min(T const &t, U const &u) +{ + return t < u ? t : u; +} + +template ::value && + std::is_arithmetic::value)> +CATLASS_HOST_DEVICE constexpr auto +add(T const &t, U const &u) +{ + return t + u; +} + +} // namespace tla + +#endif // TLA_NUMERIC_MATH_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/tla/tuple.hpp b/csrc/chunk_kda_fwd/op_kernel/tla/tuple.hpp new file mode 100644 index 0000000000..db8a39ec0a --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/tla/tuple.hpp @@ -0,0 +1,141 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef TLA_TUPLE_HPP +#define TLA_TUPLE_HPP + +#include "tla/numeric/integral_constant.hpp" +#include "tla/numeric/integer_sequence.hpp" + +namespace tla { + +namespace detail { + +// EBO stands for "empty base optimization." +template ::value> +struct EBO; + +// Specialization for types T that are empty; +template +struct EBO { + CATLASS_HOST_DEVICE constexpr EBO() {} + + CATLASS_HOST_DEVICE constexpr EBO(T const &) {} +}; + +template +CATLASS_HOST_DEVICE constexpr T getv(EBO const &) +{ + return {}; +} + +// Specialization for types T that are not empty; +template +struct EBO { + CATLASS_HOST_DEVICE constexpr EBO() + : t_{} + {} + + CATLASS_HOST_DEVICE constexpr EBO(T const &t) + : t_{t} + {} + + T t_; +}; + +template +CATLASS_HOST_DEVICE constexpr T const &getv(EBO const &x) +{ + return x.t_; +} + +template +CATLASS_HOST_DEVICE constexpr T &getv(EBO &x) +{ + return x.t_; +} + +// TupleBase +template +struct TupleBase; + +template +struct TupleBase, T...> : EBO... { + CATLASS_HOST_DEVICE constexpr TupleBase() {} + + CATLASS_HOST_DEVICE constexpr TupleBase(T const &...t) + : EBO(t)... + {} +}; + +} // end namespace detail + +// tla::tuple class. +template +struct tuple : detail::TupleBase, T...> { + CATLASS_HOST_DEVICE constexpr tuple() {} + + CATLASS_HOST_DEVICE constexpr tuple(T const &...t) + : detail::TupleBase, T...>(t...) + {} +}; + +// get for tla::tuple +template +CATLASS_HOST_DEVICE constexpr decltype(auto) get(tuple const &t) noexcept +{ + static_assert(I < sizeof...(T), "Index out of range"); + return detail::getv(t); +} + +template +CATLASS_HOST_DEVICE constexpr decltype(auto) get(tuple &t) noexcept +{ + static_assert(I < sizeof...(T), "Index out of range"); + return detail::getv(t); +} + +template +CATLASS_HOST_DEVICE constexpr decltype(auto) get(tuple &&t) noexcept +{ + static_assert(I < sizeof...(T), "Index out of range"); + return detail::getv(static_cast &&>(t)); +} + +namespace detail { + +template +auto has_tuple_size(T *) -> bool_constant<(0 <= tuple_size::value)>; +auto has_tuple_size(...) -> false_type; + +} // end namespace detail + +template +struct is_tuple : decltype(detail::has_tuple_size((T *)0)){}; + +template +struct tuple_size> + : std::integral_constant {}; + +template +struct tuple_size> + : std::integral_constant {}; + +// make_tuple +template +CATLASS_HOST_DEVICE constexpr tuple +MakeTuple(T const &...t) +{ + return {t...}; +} + +} // end namespace tla + +#endif // TLA_TUPLE_HPP diff --git a/csrc/chunk_kda_fwd/op_kernel/tla/type_traits.hpp b/csrc/chunk_kda_fwd/op_kernel/tla/type_traits.hpp new file mode 100644 index 0000000000..81ac1427c2 --- /dev/null +++ b/csrc/chunk_kda_fwd/op_kernel/tla/type_traits.hpp @@ -0,0 +1,43 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This program is free software, you can redistribute it and/or modify it under the terms and conditions of + * CANN Open Software License Agreement Version 2.0 (the "License"). + * Please refer to the License for details. You may not use this file except in compliance with the License. + * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, + * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. + * See LICENSE in the root of the software repository for the full text of the License. + */ + +#ifndef TLA_UTIL_TYPE_TRAITS_HPP +#define TLA_UTIL_TYPE_TRAITS_HPP + +#include + +#define TLA_REQUIRES(...) typename std::enable_if<(__VA_ARGS__)>::type * = nullptr + +namespace tla { + +// using std::remove_cvref; +template +struct remove_cvref { + using type = std::remove_cv_t>; +}; + +// using std::remove_cvref_t; +template +using remove_cvref_t = typename remove_cvref::type; + +// tuple_size, tuple_element +template +struct tuple_size; + +template +struct tuple_size::type>> + : std::integral_constant::value> {}; + +template +constexpr size_t tuple_size_v = tuple_size::value; + +} // end namespace tla + +#endif // TLA_UTIL_TYPE_TRAITS_HPP diff --git a/csrc/pytorch_extensions.cpp b/csrc/pytorch_extensions.cpp index f259c064a3..b368b52e66 100644 --- a/csrc/pytorch_extensions.cpp +++ b/csrc/pytorch_extensions.cpp @@ -102,7 +102,7 @@ TORCH_LIBRARY_FRAGMENT(npu, m) "causal_conv1d_update(Tensor x, Tensor weight, Tensor(a!) conv_state, " "Tensor conv_state_indices, Tensor? bias=None, Tensor? num_accepted_tokens=None, " "Tensor? query_start_loc=None, bool activation_mode=False, int pad_slot_id=-1) -> Tensor"); - + m.def( "causal_conv1d(Tensor x, Tensor weight, Tensor(a!) conv_states, Tensor? bias=None, " "Tensor? query_start_loc=None, Tensor? cache_indices=None, Tensor? has_initial_state=None, " @@ -198,6 +198,16 @@ TORCH_LIBRARY_FRAGMENT(npu, m) "bool has_ori_kv=True, bool has_cmp_kv=True) -> Tensor"); #endif + m.def( + "chunk_kda_fwd(Tensor q, Tensor k, Tensor v, Tensor g, Tensor beta, " + "Tensor? a_log=None, Tensor? dt_bias=None, Tensor? initial_state=None, " + "Tensor? cu_seqlens=None, Tensor? chunk_indices=None, " + "str layout='BSND', float scale=1.0, int chunk_size=64, " + "bool safe_gate=False, float lower_bound=-5.0, bool use_gate_in_kernel=False, " + "bool state_v_first=False, bool output_final_state=True, bool output_gk=False, " + "bool output_w=False, bool output_u=False, bool output_qg=False, " + "bool output_kg=False, bool output_v_new=False, bool output_h=False) " + "-> (Tensor, Tensor?, Tensor?, Tensor, Tensor, Tensor?, Tensor?, Tensor?, Tensor?, Tensor?, Tensor?)"); #ifdef SGL_KERNEL_ENABLE_A5_ONLY_OPS m.def( "kv_compress_epilog(Tensor(a!) kv_compress_cache, Tensor x, Tensor slot_mapping, " @@ -270,6 +280,10 @@ TORCH_LIBRARY_IMPL(npu, PrivateUse1, m) m.impl("sgemmc_shrink", TORCH_FN(sglang::npu_kernel::sgemmc_shrink)); +#ifdef SGL_KERNEL_ENABLE_A3_ONLY_OPS + m.impl("compressor", TORCH_FN(sglang::npu_kernel::compressor)); +#endif + m.impl("apply_token_bitmask", [](at::Tensor logits, at::Tensor bitmask, const c10::optional &indices) { auto indices_or_empty = indices.has_value() ? *indices : at::empty({0}, logits.options().dtype(at::kInt)); return sglang::npu_kernel::apply_token_bitmask(logits, bitmask, indices_or_empty); @@ -297,6 +311,7 @@ TORCH_LIBRARY_IMPL(npu, PrivateUse1, m) bias_or_empty, num_accepted_or_empty, query_loc_or_empty, activation_mode, pad_slot_id); }); + m.impl("causal_conv1d", [](const at::Tensor &x, const at::Tensor &weight, const at::Tensor &conv_states, const c10::optional &bias, const c10::optional &query_start_loc, @@ -344,6 +359,8 @@ TORCH_LIBRARY_IMPL(npu, PrivateUse1, m) m.impl("situ_mxfp8_quant", TORCH_FN(sglang::npu_kernel::situ_mxfp8_quant)); #endif + m.impl("chunk_kda_fwd", TORCH_FN(sglang::npu_kernel::chunk_kda_fwd)); + #ifdef BUILD_CATLASS_MODULE m.impl("catlass_matmul_basic", TORCH_FN(sglang::npu_kernel::catlass_matmul_basic)); diff --git a/include/sgl_kenel_npu_ops.h b/include/sgl_kenel_npu_ops.h index dfac314b04..1a88740fc2 100644 --- a/include/sgl_kenel_npu_ops.h +++ b/include/sgl_kenel_npu_ops.h @@ -239,6 +239,17 @@ at::Tensor sparse_attn_sharedkv_metadata_host( int64_t cmp_mask_mode, int64_t ori_win_left, int64_t ori_win_right, bool has_ori_kv, bool has_cmp_kv); +std::tuple, c10::optional, at::Tensor, at::Tensor, + c10::optional, c10::optional, c10::optional, + c10::optional, c10::optional, c10::optional> +chunk_kda_fwd(const at::Tensor &q, const at::Tensor &k, const at::Tensor &v, const at::Tensor &g, + const at::Tensor &beta, const c10::optional &aLog, + const c10::optional &dtBias, const c10::optional &initialState, + const c10::optional &cuSeqlens, + const c10::optional &chunkIndices, const std::string &layout, double scale, + int64_t chunkSize, bool safeGate, double lowerBound, bool useGateInKernel, + bool stateVFirst, bool outputFinalState, bool outputGk, bool outputW, bool outputU, + bool outputQG, bool outputKg, bool outputVNew, bool outputH); #ifdef SGL_KERNEL_ENABLE_A3_ONLY_OPS /** * @brief Sparse row copy: for each i where valid_mask[i] is true, diff --git a/python/sgl_kernel_npu/sgl_kernel_npu/fla/chunk_kda_fwd.py b/python/sgl_kernel_npu/sgl_kernel_npu/fla/chunk_kda_fwd.py new file mode 100644 index 0000000000..37dbefae01 --- /dev/null +++ b/python/sgl_kernel_npu/sgl_kernel_npu/fla/chunk_kda_fwd.py @@ -0,0 +1,93 @@ +from typing import Optional + +import torch + + +def chunk_kda_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + *, + a_log: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + initial_state: Optional[torch.Tensor] = None, + cu_seqlens: Optional[torch.Tensor] = None, + chunk_indices: Optional[torch.Tensor] = None, + layout: str = "BSND", + scale: float = 1.0, + chunk_size: int = 64, + safe_gate: bool = False, + lower_bound: float = -5.0, + use_gate_in_kernel: bool = False, + state_v_first: bool = False, + output_final_state: bool = True, + output_gk: bool = False, + output_w: bool = False, + output_u: bool = False, + output_qg: bool = False, + output_kg: bool = False, + output_v_new: bool = False, + output_h: bool = False, +) -> tuple: + r"""Fused chunk KDA forward kernel (direct-launch NPU implementation). + + Args: + q: (B, S, H, K) bfloat16 (layout-dependent, see ``layout``). + k: (B, S, H, K) bfloat16, same shape as q. + v: (B, S, HV, V) bfloat16. + g: (B, S, HV, K) float32 or bfloat16 raw gate. + beta: (B, S, HV) float32 or bfloat16 delta coefficient. + a_log: (HV,) float32 gate decay parameter, required when use_gate_in_kernel. + dt_bias: (HV*K,) float32 gate bias. + initial_state: (N, HV, K, V) [or (N, HV, V, K) when state_v_first] float32. + cu_seqlens: (N+1,) int64 varlen sequence boundaries. + chunk_indices: (2*NC,) int64 canonical (seq_id, chunk_id) pairs. + layout: one of "BSND", "BNSD", "TND", "NTD". + scale: query scaling factor, usually K^(-0.5). + chunk_size: 64 or 128. + safe_gate: whether to use a bounded gate. + lower_bound: safe gate lower bound in [-5, 0). + use_gate_in_kernel: whether to compute the activated gate in-kernel. + state_v_first: whether the state tensors have (V, K) last two dims. + output_*: whether to materialize the corresponding intermediate output. + + Returns: + A tuple of length 11: + (attn_out, final_state?, gk?, aqk_out, akk_out, w?, u?, qg?, kg?, v_new?, h?). + attn_out: (B, S, HV, V) [or (T, HV, V) for rank-3 layouts]. + final_state: (N, HV, K, V) [or (N, HV, V, K) when state_v_first] float32. + gk: (B, HV, S, K) [or (HV, S, K)] float32. + aqk_out/akk_out: (B, HV, S, C) [or (HV, S, C)]. + w/qg/kg: (B, HV, S, K) [or (HV, S, K)]. + u/v_new: (B, HV, S, V) [or (HV, S, V)]. + h: (B, NC, HV, K, V) [or (NC, HV, K, V)]. + """ + return torch.ops.npu.chunk_kda_fwd( + q, + k, + v, + g, + beta, + a_log, + dt_bias, + initial_state, + cu_seqlens, + chunk_indices, + layout, + scale, + chunk_size, + safe_gate, + lower_bound, + use_gate_in_kernel, + state_v_first, + output_final_state, + output_gk, + output_w, + output_u, + output_qg, + output_kg, + output_v_new, + output_h, + )