diff --git a/xllm_ops/moe/dequant_swiglu_quant/op_host/dequant_swiglu_quant_tiling_base.cpp b/xllm_ops/moe/dequant_swiglu_quant/op_host/dequant_swiglu_quant_tiling_base.cpp index 612ff81..3b60b51 100644 --- a/xllm_ops/moe/dequant_swiglu_quant/op_host/dequant_swiglu_quant_tiling_base.cpp +++ b/xllm_ops/moe/dequant_swiglu_quant/op_host/dequant_swiglu_quant_tiling_base.cpp @@ -62,6 +62,19 @@ constexpr uint32_t PERFORMANCE_ROW_LEN = 128; constexpr uint32_t MIN_CORE = 12; const int64_t DYNAMIC_INT_X_FLOAT32_BIAS_QUANT_D_PERFORMANCE = 30013; +// Vectorised-row fast path (see op_kernel/dequant_swiglu_quant_vecrow.hpp). +// It removes the per-row V_S sync + GetValue/SetValue round-trip the other +// dynamic classes pay, which only dominates once there are many rows: measured +// at colLen=256, rows<=512 is a wash or a small loss (bf16 256 rows +31%) while +// rows>=1024 wins steadily (int32 2048 rows -32%, bf16 8192 rows -21%). +// The win also shrinks as colLen grows, because the path spills the fp32 SwiGLU +// result to UB and reads it back -- by colLen=1024 it is a wash. +const int64_t DYNAMIC_VECROW_INT32 = 30020; +const int64_t DYNAMIC_VECROW_FLOAT16 = 30021; +const int64_t DYNAMIC_VECROW_BFLOAT16 = 30022; +constexpr uint32_t VECROW_COL_LEN = 512; +constexpr uint32_t VECROW_MIN_ROW_LEN = 1024; + // Tiling优选参数 struct GluSingleTilingOptParam { // Maximum amount of data that can be transferred by an operator UB at a time. Unit:element @@ -146,6 +159,8 @@ class DequantSwigluQuantTiling : public TilingBaseClass { bool isPerformanceBranch(); + bool isVecRowBranch(); + int64_t getTilingKeyStatic( const int32_t inputDtype, const ge::DataType biasType, const int64_t scaleSize) const; @@ -170,6 +185,7 @@ class DequantSwigluQuantTiling : public TilingBaseClass { ge::DataType xInputDataType; bool isPerfBranch = false; + bool isVecRow = false; ge::DataType biasDataType = ge::DT_FLOAT; uint64_t quantScaleShapeSize = 0; @@ -590,6 +606,7 @@ ge::graphStatus DequantSwigluQuantTiling::DoOpTiling() return ge::GRAPH_FAILED; } isPerfBranch = isPerformanceBranch(); + isVecRow = isVecRowBranch(); return ge::GRAPH_SUCCESS; } @@ -647,12 +664,18 @@ int64_t DequantSwigluQuantTiling::getTilingKeyDynamic( if (scaleSize == 1) { return DYNAMIC_FLOAT16_X; } else { + if (isVecRow) { + return DYNAMIC_VECROW_FLOAT16; + } return DYNAMIC_FLOAT16_XD; } } else { if (scaleSize == 1) { return DYNAMIC_BFLOAT16_X; } else { + if (isVecRow) { + return DYNAMIC_VECROW_BFLOAT16; + } return DYNAMIC_BFLOAT16_XD; } } @@ -671,6 +694,12 @@ int64_t DequantSwigluQuantTiling::getTilingKeyDynamic( if (biasType == ge::DT_INT32) { return DYNAMIC_INT_X_INT_BIAS_QUANT_D; } else if (biasType == ge::DT_FLOAT) { + // isVecRowBranch() checks the bias input directly, so the flag being + // set already means no bias is present (biasType defaults to + // DT_FLOAT when the optional input is absent). + if (isVecRow) { + return DYNAMIC_VECROW_INT32; + } if(isPerfBranch) { return DYNAMIC_INT_X_FLOAT32_BIAS_QUANT_D_PERFORMANCE; } @@ -695,6 +724,42 @@ bool DequantSwigluQuantTiling::isPerformanceBranch() { return false; } +// The vecrow kernel handles one whole row per VF iteration and keeps the +// per-row max in a vector register, so it needs the row fully loadable and no +// bias / quant_offset / group_index to fold in. Beyond that it is gated on +// shape: colLen must be small enough that spilling the fp32 SwiGLU result to UB +// still pays, and there must be enough rows to amortise its setup. +bool DequantSwigluQuantTiling::isVecRowBranch() { + if (tilingData.get_is32BAligned() != 1) { + return false; + } + // biasIsEmpty is only populated on the int32 branch (checkWeightBiasActivate + // runs under `xDataType == DT_INT32`), so it reads 0 for bf16/fp16 whether or + // not a bias exists. Ask the context directly instead of trusting the field. + if (context_->GetOptionalInputShape(INDEX_IN_BIAS) != nullptr) { + return false; + } + // quant_offset is not folded in by this kernel. + if (context_->GetOptionalInputShape(INDEX_IN_QUANT_OFFSET) != nullptr) { + return false; + } + // Row must be fully loadable: the kernel indexes a row in one shot. + if (tilingData.get_baseColLen() != tilingData.get_colLen()) { + return false; + } + uint64_t colLen = tilingData.get_colLen(); + if (colLen > VECROW_COL_LEN) { + return false; + } + // Below this row count the fast path's fixed setup is not amortised and it + // loses to the row-wise kernel (measured at colLen=256: 256 rows +31%, + // 1024 rows -10%, 8192 rows -21%). + if (tilingData.get_rowLen() < VECROW_MIN_ROW_LEN) { + return false; + } + return true; +} + uint64_t DequantSwigluQuantTiling::GetTilingKey() const { if (quantMode == 0) { // static diff --git a/xllm_ops/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant.cpp b/xllm_ops/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant.cpp index fb645f4..a52d391 100644 --- a/xllm_ops/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant.cpp +++ b/xllm_ops/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant.cpp @@ -26,6 +26,7 @@ #include "dequant_swiglu_quant_dynamic_bias_int32.hpp" #include "dequant_swiglu_quant_dynamic_bias_float.hpp" #include "dequant_swiglu_quant_dynamic_performance.hpp" +#include "dequant_swiglu_quant_vecrow.hpp" using namespace AscendC; @@ -58,6 +59,11 @@ using namespace AscendC; #define DEQUANT_SWIGLU_QUANT_WITH_GROUP_FP16_QS_GR 110000100 #define DEQUANT_SWIGLU_QUANT_WITH_GROUP_BF16_QS_GR 110000200 +// Vectorised-row fast path; keys assigned in dequant_swiglu_quant_tiling_base.cpp. +#define DEQUANT_SWIGLU_QUANT_VECROW_INT32 30020 +#define DEQUANT_SWIGLU_QUANT_VECROW_FLOAT16 30021 +#define DEQUANT_SWIGLU_QUANT_VECROW_BFLOAT16 30022 + extern "C" __global__ __aicore__ void dequant_swiglu_quant(GM_ADDR xGM, GM_ADDR weightSscaleGM, GM_ADDR activationScaleGM, GM_ADDR biasGM, GM_ADDR quantScaleGM, GM_ADDR quantOffsetGM, @@ -312,6 +318,12 @@ extern "C" __global__ __aicore__ void dequant_swiglu_quant(GM_ADDR xGM, GM_ADDR op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, userspace, tilingData, &(pipe)); op.Process(); + } else if (TILING_KEY_IS(DEQUANT_SWIGLU_QUANT_VECROW_INT32)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantVecRow op; + op.Init(xGM, weightSscaleGM, activationScaleGM, quantScaleGM, yGM, scaleGM, tilingData, &(pipe)); + op.Process(); } #if !(defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3003 || __NPU_ARCH__ == 3113)) // ORIG_DTYPE_BIAS == DT_BF16 @@ -375,6 +387,12 @@ extern "C" __global__ __aicore__ void dequant_swiglu_quant(GM_ADDR xGM, GM_ADDR op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, userspace, tilingData, &(pipe)); op.Process(); + } else if (TILING_KEY_IS(DEQUANT_SWIGLU_QUANT_VECROW_FLOAT16)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantVecRow op; + op.Init(xGM, weightSscaleGM, activationScaleGM, quantScaleGM, yGM, scaleGM, tilingData, &(pipe)); + op.Process(); } #endif #if !(defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3003 || __NPU_ARCH__ == 3113)) && (ORIG_DTYPE_X == DT_BF16) @@ -428,6 +446,12 @@ extern "C" __global__ __aicore__ void dequant_swiglu_quant(GM_ADDR xGM, GM_ADDR op.Init(xGM, weightSscaleGM, activationScaleGM, biasGM, quantScaleGM, quantOffsetGM, yGM, scaleGM, userspace, tilingData, &(pipe)); op.Process(); + } else if (TILING_KEY_IS(DEQUANT_SWIGLU_QUANT_VECROW_BFLOAT16)) { + GET_TILING_DATA_WITH_STRUCT(SwiGluTilingData, tilingDataIn, tiling); + const SwiGluTilingData* __restrict__ tilingData = &tilingDataIn; + DequantSwigluQuant::DequantSwigluQuantVecRow op; + op.Init(xGM, weightSscaleGM, activationScaleGM, quantScaleGM, yGM, scaleGM, tilingData, &(pipe)); + op.Process(); } #endif } \ No newline at end of file diff --git a/xllm_ops/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_vecrow.hpp b/xllm_ops/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_vecrow.hpp new file mode 100644 index 0000000..a4279bd --- /dev/null +++ b/xllm_ops/moe/dequant_swiglu_quant/op_kernel/dequant_swiglu_quant_vecrow.hpp @@ -0,0 +1,479 @@ +/** + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * Copyright 2026 The xLLM Authors. All Rights Reserved. + * 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 dequant_swiglu_quant_vecrow.hpp + * \brief Vectorised-row fast path for dynamic per-token quant. + * + * The existing dynamic classes process one row at a time and, per row, pay a + * ReduceMax + V_S sync + GetValue/SetValue + Muls(scalar) round-trip (see + * DequantSwigluQuantDynamicBase::dynamicMultiColMax, which ends in + * maxTempLocal.GetValue(rowId)). With rows in the tens of thousands that scalar + * traffic dominates: three shapes moving identical bytes took 47/70/123us purely + * because their row counts differ. + * + * Here the whole dequant -> SwiGLU -> abs-max -> quant chain runs in MicroAPI + * inside two __VEC_SCOPE__ blocks per tile. The per-row max never leaves a + * vector register: ReduceMax lands it in lane 0 and a broadcast load feeds it + * straight back into the divide. Input tiles are double-buffered so MTE2 + * overlaps the VF work. + * + * Scope: dynamic quant, no bias / quant_offset / group_index. Everything else + * keeps using the original classes. + */ + +#ifndef CANN_DEQUANT_SWIGLU_QUANT_VECROW_HPP +#define CANN_DEQUANT_SWIGLU_QUANT_VECROW_HPP + +#include "kernel_operator.h" + +namespace DequantSwigluQuant { +using namespace AscendC; + +namespace VecRow { + +constexpr uint32_t UB_BLOCK_BYTES = 32; +// Row strides are padded to a whole vector register so every chunk load in the +// VF loops sits on a natural boundary. +constexpr uint32_t VREG_BYTES = 256; +constexpr uint32_t VL_FP32 = 64; +constexpr uint32_t MAX_TILE_ROWS = 64; +constexpr uint32_t UB_BUDGET_BYTES = 184 * 1024; + +// resolves to the round-only cast; CAST_NONE is the plain +// widening the original int32 path uses. +constexpr Reg::CastTrait castS32ToF32 = { + Reg::RegLayout::UNKNOWN, Reg::SatMode::UNKNOWN, Reg::MaskMergeMode::ZEROING, RoundMode::CAST_NONE}; +constexpr Reg::CastTrait castB16ToF32 = { + Reg::RegLayout::ZERO, Reg::SatMode::UNKNOWN, Reg::MaskMergeMode::ZEROING, RoundMode::UNKNOWN}; +// fp32 -> int8 goes fp32 -> s16 -> fp16 -> s8, the same chain adv_api's +// TransRegForS8 uses (quantize_impl.h:145). Casting straight to int8 rounds and +// packs lanes wrongly -- it yields a correct scale but a garbage y. +constexpr Reg::CastTrait castF32ToS16 = { + Reg::RegLayout::ZERO, Reg::SatMode::SAT, Reg::MaskMergeMode::ZEROING, RoundMode::CAST_RINT}; + +__aicore__ inline uint32_t CeilDivU(uint32_t a, uint32_t b) +{ + return (b == 0) ? 0 : (a + b - 1) / b; +} + +// Pad `count` elements of `sz` bytes up to a whole `toBytes` boundary. +__aicore__ inline uint32_t AlignElTo(uint32_t count, uint32_t sz, uint32_t toBytes) +{ + uint32_t bytes = count * sz; + return (bytes + toBytes - 1) / toBytes * toBytes / sz; +} + +} // namespace VecRow + +/** + * InType: int32_t (dequant path) or half/bfloat16_t (direct path). + * hasWeightScale: int32 path multiplies by weight_scale[2H]. + */ +template +class DequantSwigluQuantVecRow { +public: + __aicore__ inline DequantSwigluQuantVecRow() {} + __aicore__ inline ~DequantSwigluQuantVecRow() {} + + __aicore__ inline void Init(GM_ADDR x_gm, GM_ADDR weight_scale_gm, GM_ADDR activation_scale_gm, + GM_ADDR quant_scale_gm, GM_ADDR y_gm, GM_ADDR scale_gm, + const SwiGluTilingData* tilingData, TPipe* pipe_) + { + using namespace VecRow; + pipe = pipe_; + curBlockIdx = GetBlockIdx(); + + H = static_cast(tilingData->colLen); + rowNum = static_cast(tilingData->rowLen); + actLeft = (tilingData->activateLeft != 0); + hasQs = (tilingData->quantScaleIsEmpty == 0); + hasActScale = (tilingData->activateScaleIsEmpty == 0); + + useCoreNum = tilingData->usedCoreNum; + if (rowNum < useCoreNum) { + useCoreNum = rowNum; + } + + // Same even split the other classes use, so core assignment matches. + uint32_t perRoundCnt = useCoreNum == 0 ? 0 : rowNum / useCoreNum; + uint32_t remainCnt = rowNum - useCoreNum * perRoundCnt; + numRound = perRoundCnt; + if (curBlockIdx < remainCnt) { + numRound = perRoundCnt + 1; + biasOffset = curBlockIdx * (perRoundCnt + 1); + } else { + biasOffset = (perRoundCnt + 1) * remainCnt + (curBlockIdx - remainCnt) * perRoundCnt; + } + + xGm.SetGlobalBuffer((__gm__ InType*)x_gm); + yGm.SetGlobalBuffer((__gm__ int8_t*)y_gm); + scaleGm.SetGlobalBuffer((__gm__ float*)scale_gm); + if constexpr (isInt32) { + weightScaleGm.SetGlobalBuffer((__gm__ float*)weight_scale_gm); + if (hasActScale) { + actScaleGm.SetGlobalBuffer((__gm__ float*)activation_scale_gm); + } + } + if (hasQs) { + quantScaleGm.SetGlobalBuffer((__gm__ float*)quant_scale_gm); + } + + strideIn = AlignElTo(H, sizeof(InType), VREG_BYTES); + strideF = AlignElTo(H, sizeof(float), VREG_BYTES); + strideI8 = AlignElTo(H, 1, VREG_BYTES); + + uint32_t perRowBytes = 2 * strideIn * sizeof(InType) + strideF * sizeof(float) + strideI8; + uint32_t sharedBytes = 2 * strideF * sizeof(float) + (hasQs ? strideF * sizeof(float) : 0) + + 3 * MAX_TILE_ROWS * sizeof(float); + uint32_t avail = (UB_BUDGET_BYTES > sharedBytes) ? (UB_BUDGET_BYTES - sharedBytes) : 0; + tileRows = (perRowBytes > 0) ? (avail / perRowBytes) : 1; + if (tileRows > MAX_TILE_ROWS) tileRows = MAX_TILE_ROWS; + if (tileRows > numRound) tileRows = numRound; + if (tileRows < 1) tileRows = 1; + // Two input sets so the next tile's MTE2 overlaps this tile's VEC work. + if (tileRows > 1) tileRows /= 2; + + pipe->InitBuffer(bufA, 2 * tileRows * strideIn * sizeof(InType)); + pipe->InitBuffer(bufB, 2 * tileRows * strideIn * sizeof(InType)); + pipe->InitBuffer(bufSwi, tileRows * strideF * sizeof(float)); + pipe->InitBuffer(bufY, tileRows * strideI8); + pipe->InitBuffer(bufMax, MAX_TILE_ROWS * sizeof(float)); + pipe->InitBuffer(bufScale, MAX_TILE_ROWS * sizeof(float)); + pipe->InitBuffer(bufWsA, strideF * sizeof(float)); + pipe->InitBuffer(bufWsB, strideF * sizeof(float)); + pipe->InitBuffer(bufAct, 2 * MAX_TILE_ROWS * sizeof(float)); + if (hasQs) { + pipe->InitBuffer(bufQs, strideF * sizeof(float)); + } + } + + __aicore__ inline void Process() + { + using namespace VecRow; + if (curBlockIdx >= useCoreNum || numRound == 0) { + return; + } + + LocalTensor aLocal = bufA.template Get(); + LocalTensor bLocal = bufB.template Get(); + LocalTensor swiLocal = bufSwi.Get(); + LocalTensor yLocal = bufY.Get(); + LocalTensor maxLocal = bufMax.Get(); + LocalTensor scaleLocal = bufScale.Get(); + LocalTensor wsALocal = bufWsA.Get(); + LocalTensor wsBLocal = bufWsB.Get(); + LocalTensor actLocal = bufAct.Get(); + + // weight_scale / quant_scale are row-invariant: load once per core. + if constexpr (isInt32) { + DataCopyExtParams cpWs{1, static_cast(H * sizeof(float)), 0, 0, 0}; + DataCopyPadExtParams ppWs{false, 0, 0, 0}; + DataCopyPad(wsALocal, weightScaleGm[0], cpWs, ppWs); + DataCopyPad(wsBLocal, weightScaleGm[H], cpWs, ppWs); + } + if (hasQs) { + DataCopyExtParams cpQs{1, static_cast(H * sizeof(float)), 0, 0, 0}; + DataCopyPadExtParams ppQs{false, 0, 0, 0}; + DataCopyPad(bufQs.Get(), quantScaleGm[0], cpQs, ppQs); + } + if (isInt32 || hasQs) { + event_t eid = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); + SetFlag(eid); + WaitFlag(eid); + } + + __local_mem__ InType* aAddr = (__local_mem__ InType*)aLocal.GetPhyAddr(); + __local_mem__ InType* bAddr = (__local_mem__ InType*)bLocal.GetPhyAddr(); + __local_mem__ float* swiAddr = (__local_mem__ float*)swiLocal.GetPhyAddr(); + __local_mem__ int8_t* yAddr = (__local_mem__ int8_t*)yLocal.GetPhyAddr(); + __local_mem__ float* maxAddr = (__local_mem__ float*)maxLocal.GetPhyAddr(); + __local_mem__ float* scaleAddr = (__local_mem__ float*)scaleLocal.GetPhyAddr(); + __local_mem__ float* wsAAddr = (__local_mem__ float*)wsALocal.GetPhyAddr(); + __local_mem__ float* wsBAddr = (__local_mem__ float*)wsBLocal.GetPhyAddr(); + __local_mem__ float* actAddr = (__local_mem__ float*)actLocal.GetPhyAddr(); + __local_mem__ float* qsAddr = hasQs + ? (__local_mem__ float*)bufQs.Get().GetPhyAddr() : nullptr; + + // activate_left picks which half drives the sigmoid. Matches + // DequantSwigluQuantDynamicBase::BaseProcess: activateLeft==0 -> gate is + // the second half. Resolved here so the VF loops stay branch-free. + __local_mem__ InType* gateAddr = actLeft ? aAddr : bAddr; + __local_mem__ InType* upAddr = actLeft ? bAddr : aAddr; + __local_mem__ float* wsGateAddr = actLeft ? wsAAddr : wsBAddr; + __local_mem__ float* wsUpAddr = actLeft ? wsBAddr : wsAAddr; + + uint32_t dstStrideIn = (strideIn - H) * sizeof(InType) / UB_BLOCK_BYTES; + uint32_t dstStrideY = (strideI8 - H) / UB_BLOCK_BYTES; + uint32_t inHalf = tileRows * strideIn; + uint32_t actHalf = VecRow::MAX_TILE_ROWS; + + DataCopyPadExtParams ppIn{false, 0, 0, 0}; + DataCopyPadExtParams ppAct{false, 0, 0, 0}; + + event_t eidM2V[2] = { + static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)), + static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V))}; + event_t eidM3M2 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE3_MTE2)); + event_t eidV2M3 = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3)); + +#define DSQ_VECROW_ISSUE_LOAD(rIdx, halfIdx) \ + do { \ + uint32_t _rows = tileRows; \ + if (_rows > numRound - (rIdx)) _rows = numRound - (rIdx); \ + uint32_t _start = biasOffset + (rIdx); \ + DataCopyExtParams _cpIn{static_cast(_rows), \ + static_cast(H * sizeof(InType)), \ + static_cast(H * sizeof(InType)), \ + static_cast(dstStrideIn), 0}; \ + DataCopyPad(aLocal[(halfIdx) * inHalf], xGm[_start * 2 * H], _cpIn, ppIn); \ + DataCopyPad(bLocal[(halfIdx) * inHalf], xGm[_start * 2 * H + H], _cpIn, ppIn); \ + if constexpr (isInt32) { \ + if (hasActScale) { \ + DataCopyExtParams _cpAct{1, static_cast(_rows * sizeof(float)), \ + 0, 0, 0}; \ + DataCopyPad(actLocal[(halfIdx) * actHalf], actScaleGm[_start], _cpAct, \ + ppAct); \ + } \ + } \ + } while (0) + + DSQ_VECROW_ISSUE_LOAD(0u, 0u); + SetFlag(eidM2V[0]); + + uint32_t half = 0; + for (uint32_t r = 0; r < numRound; r += tileRows) { + uint32_t curRows = tileRows; + if (curRows > numRound - r) curRows = numRound - r; + uint32_t rowStart = biasOffset + r; + uint32_t nextR = r + tileRows; + uint32_t nextHalf = half ^ 1; + + // Start the next tile's load before touching this one, so MTE2 runs + // underneath the VF work below. The other half is free: its compute + // finished an iteration ago. + if (nextR < numRound) { + DSQ_VECROW_ISSUE_LOAD(nextR, nextHalf); + SetFlag(eidM2V[nextHalf]); + } + + WaitFlag(eidM2V[half]); + + ComputeSwigluAndMax(gateAddr + half * inHalf, upAddr + half * inHalf, + wsGateAddr, wsUpAddr, qsAddr, actAddr + half * actHalf, + swiAddr, maxAddr, curRows); + + QuantizeRows(swiAddr, maxAddr, scaleAddr, yAddr, curRows); + + SetFlag(eidV2M3); + WaitFlag(eidV2M3); + + DataCopyExtParams cpY{static_cast(curRows), static_cast(H), + dstStrideY, 0, 0}; + DataCopyPad(yGm[rowStart * H], yLocal, cpY); + + DataCopyExtParams cpScale{1, static_cast(curRows * sizeof(float)), 0, 0, 0}; + DataCopyPad(scaleGm[rowStart], scaleLocal, cpScale); + + // swi/y/scale are single-buffered: the next iteration's VF writes + // must not start before these stores drain. + SetFlag(eidM3M2); + WaitFlag(eidM3M2); + + half = nextHalf; + } +#undef DSQ_VECROW_ISSUE_LOAD + } + +private: + /** + * Pass 1: dequant + SwiGLU for every row in the tile, spilling the fp32 + * result to `swi` and leaving each row's abs-max in maxBuf[i]. + * + * The abs-max accumulation deliberately uses a full mask while the Abs that + * feeds it uses the tail mask: ZEROING zeroes the inactive lanes of the Abs + * result, and 0 is the identity for a max over absolute values. Masking the + * Max itself would instead zero the lanes already holding the running max + * from earlier chunks. + */ + __aicore__ inline void ComputeSwigluAndMax( + __local_mem__ InType* gateIn, __local_mem__ InType* upIn, + __local_mem__ float* wsGate, __local_mem__ float* wsUp, + __local_mem__ float* qs, __local_mem__ float* actBuf, + __local_mem__ float* swi, __local_mem__ float* maxBuf, uint32_t rows) + { + using namespace VecRow; + const uint16_t chunks = static_cast(CeilDivU(H, VL_FP32)); + const uint16_t rowCnt = static_cast(rows); + const uint32_t hLocal = H; + const uint32_t strideInLocal = strideIn; + const uint32_t strideFLocal = strideF; + const bool hasQsLocal = hasQs; + const bool hasActLocal = hasActScale; + + __VEC_SCOPE__ + { + MicroAPI::RegTensor vGate, vUp, vScale, vTmp, vOne, vSig, vMax, vAbs; + MicroAPI::RegTensor vRaw; + MicroAPI::MaskReg m; + MicroAPI::MaskReg mAll = MicroAPI::CreateMask(); + MicroAPI::MaskReg mOne = MicroAPI::CreateMask(); + + for (uint16_t i = 0; i < rowCnt; i++) { + uint32_t sreg = hLocal; + MicroAPI::Duplicate(vMax, 0.0f); + MicroAPI::Duplicate(vOne, 1.0f, mAll); + + if constexpr (isInt32) { + // activation_scale is per row: one broadcast load feeds all + // chunks, replacing a per-row Muls(scalar). + if (hasActLocal) { + MicroAPI::DataCopy(vScale, actBuf + i); + } else { + MicroAPI::Duplicate(vScale, 1.0f, mAll); + } + } + + for (uint16_t c = 0; c < chunks; c++) { + m = MicroAPI::UpdateMask(sreg); + uint32_t inOff = i * strideInLocal + c * VL_FP32; + uint32_t fOff = i * strideFLocal + c * VL_FP32; + + if constexpr (isInt32) { + MicroAPI::DataCopy(vRaw, gateIn + inOff); + MicroAPI::Cast(vGate, vRaw, m); + MicroAPI::DataCopy(vTmp, wsGate + c * VL_FP32); + MicroAPI::Mul(vGate, vGate, vTmp, m); + MicroAPI::Mul(vGate, vGate, vScale, m); + + MicroAPI::DataCopy(vRaw, upIn + inOff); + MicroAPI::Cast(vUp, vRaw, m); + MicroAPI::DataCopy(vTmp, wsUp + c * VL_FP32); + MicroAPI::Mul(vUp, vUp, vTmp, m); + MicroAPI::Mul(vUp, vUp, vScale, m); + } else { + MicroAPI::DataCopy(vRaw, gateIn + inOff); + MicroAPI::Cast(vGate, vRaw, m); + MicroAPI::DataCopy(vRaw, upIn + inOff); + MicroAPI::Cast(vUp, vRaw, m); + } + + // SiLU(gate) = gate / (1 + exp(-gate)) -- the same op + // sequence adv_api's Sigmoid uses on this arch, so results + // match the original path bit for bit. + MicroAPI::Muls(vTmp, vGate, -1.0f, m); + MicroAPI::Exp(vTmp, vTmp, m); + MicroAPI::Adds(vTmp, vTmp, 1.0f, m); + MicroAPI::Div(vSig, vOne, vTmp, m); + MicroAPI::Mul(vGate, vGate, vSig, m); + + MicroAPI::Mul(vGate, vGate, vUp, m); + + if (hasQsLocal) { + MicroAPI::DataCopy(vTmp, qs + c * VL_FP32); + MicroAPI::Mul(vGate, vGate, vTmp, m); + } + + MicroAPI::DataCopy(swi + fOff, vGate, m); + + MicroAPI::Abs(vAbs, vGate, m); + MicroAPI::Max(vMax, vMax, vAbs, mAll); + } + + MicroAPI::ReduceMax(vMax, vMax, mAll); + MicroAPI::DataCopy(maxBuf + i, vMax, mOne); + } + } + } + + /** + * Pass 2: turn each row's max into its scale and quantise. The max is pulled + * back as a broadcast vector load, so scale and 1/scale are computed in-lane + * and the divide never touches the scalar unit. + */ + __aicore__ inline void QuantizeRows( + __local_mem__ float* swi, __local_mem__ float* maxBuf, + __local_mem__ float* scaleBuf, __local_mem__ int8_t* yOut, uint32_t rows) + { + using namespace VecRow; + const uint16_t chunks = static_cast(CeilDivU(H, VL_FP32)); + const uint16_t rowCnt = static_cast(rows); + const uint32_t hLocal = H; + const uint32_t strideFLocal = strideF; + const uint32_t strideI8Local = strideI8; + + __VEC_SCOPE__ + { + MicroAPI::RegTensor vVal, vScale, vInv, vOne; + MicroAPI::RegTensor vHalf; + MicroAPI::RegTensor vQ; + MicroAPI::MaskReg m; + MicroAPI::MaskReg mAll = MicroAPI::CreateMask(); + MicroAPI::MaskReg mOne = MicroAPI::CreateMask(); + + for (uint16_t i = 0; i < rowCnt; i++) { + uint32_t sreg = hLocal; + + MicroAPI::DataCopy(vScale, maxBuf + i); + MicroAPI::Muls(vScale, vScale, 1.0f / 127.0f, mAll); + MicroAPI::Maxs(vScale, vScale, 1e-12f, mAll); + MicroAPI::DataCopy(scaleBuf + i, vScale, mOne); + + MicroAPI::Duplicate(vOne, 1.0f, mAll); + MicroAPI::Div(vInv, vOne, vScale, mAll); + + for (uint16_t c = 0; c < chunks; c++) { + m = MicroAPI::UpdateMask(sreg); + MicroAPI::DataCopy( + vVal, swi + i * strideFLocal + c * VL_FP32); + MicroAPI::Mul(vVal, vVal, vInv, m); + // SAT on the s16 step clamps to [-128,127] once the value + // lands in int8, so the explicit clamp is free. + MicroAPI::Cast((MicroAPI::RegTensor&)vHalf, vVal, m); + MicroAPI::Cast(vHalf, (MicroAPI::RegTensor&)vHalf, m); + MicroAPI::Cast(vQ, vHalf, m); + MicroAPI::DataCopy( + yOut + i * strideI8Local + c * VL_FP32, vQ, m); + } + } + } + } + +private: + TPipe* pipe = nullptr; + uint32_t curBlockIdx = 0; + uint32_t H = 0; + uint32_t rowNum = 0; + uint32_t useCoreNum = 0; + uint32_t numRound = 0; + uint32_t biasOffset = 0; + uint32_t tileRows = 1; + uint32_t strideIn = 0; + uint32_t strideF = 0; + uint32_t strideI8 = 0; + bool actLeft = false; + bool hasQs = false; + bool hasActScale = false; + + GlobalTensor xGm; + GlobalTensor weightScaleGm; + GlobalTensor actScaleGm; + GlobalTensor quantScaleGm; + GlobalTensor yGm; + GlobalTensor scaleGm; + + TBuf bufA, bufB, bufSwi, bufY; + TBuf bufWsA, bufWsB, bufQs, bufAct, bufMax, bufScale; +}; + +} // namespace DequantSwigluQuant + +#endif // CANN_DEQUANT_SWIGLU_QUANT_VECROW_HPP diff --git a/xllm_ops/moe/dequant_swiglu_quant/tiling_base/error_log.h b/xllm_ops/moe/dequant_swiglu_quant/tiling_base/error_log.h index 1b09d20..9e2b4e4 100644 --- a/xllm_ops/moe/dequant_swiglu_quant/tiling_base/error_log.h +++ b/xllm_ops/moe/dequant_swiglu_quant/tiling_base/error_log.h @@ -10,6 +10,10 @@ #pragma once +// tiling_base.h calls CheckLogLevel(static_cast(OP), ...); the OP +// module id lives in base/log_types.h, which slog.h pulls in. log/log.h alone +// does not, so include it here as moe_gating_top_k/tiling_base/error_log.h does. +#include "toolchain/slog.h" #include "log/log.h" #ifndef OP_LOGE_FOR_INVALID_DTYPE diff --git a/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/BEFORE_AFTER_PERF.md b/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/BEFORE_AFTER_PERF.md new file mode 100644 index 0000000..af8ea51 --- /dev/null +++ b/xllm_ops/moe/hc_pre_sinkhorn/docs/perf-integrate/BEFORE_AFTER_PERF.md @@ -0,0 +1,47 @@ +# hc_pre_sinkhorn 20 case 优化前后性能 + +日期：2026-08-21 +口径：aclnn + torch_npu profiler，`HcPreSinkhorn` kernel 时间（µs） +优化前：Round 2 RegResident vendor（合入 SoA 前同机实测） +优化后：Round 3 SoA，源码在 `my-xllm-ops/xllm-ops/xllm_ops/moe/hc_pre_sinkhorn` + +## 结论 + +- 几何平均加速 **2.15x** +- 14/20 case 加速比 ≥ 1.10x；M≥6（9–12）未走 SoA，约 1.00x +- 最大：case 6（bs=16384, M=4）**204.48 → 12.88 µs，15.87x** +- 小 batch（1/2/18）基本持平；case 18 微回归 3.01 → 3.07 µs + +## 逐 case + +| case | bs | M | iters | 优化前 µs | 优化后 µs | 加速比 | 路径 | +|------|----|---|-------|-----------|-----------|--------|------| +| 1 | 1 | 4 | 20 | 3.155 | 3.142 | 1.00x | 小 batch | +| 2 | 64 | 4 | 20 | 4.602 | 4.546 | 1.01x | 小 batch | +| 3 | 512 | 4 | 20 | 10.780 | 5.512 | 1.96x | SoA | +| 4 | 1024 | 4 | 20 | 16.730 | 5.407 | 3.09x | SoA | +| 5 | 4096 | 4 | 20 | 53.710 | 6.978 | 7.70x | SoA | +| 6 | 16384 | 4 | 20 | 204.479 | 12.884 | **15.87x** | SoA | +| 7 | 1024 | 2 | 20 | 14.609 | 4.640 | 3.15x | SoA | +| 8 | 1024 | 3 | 20 | 16.120 | 5.323 | 3.03x | SoA | +| 9 | 1024 | 6 | 20 | 18.805 | 18.812 | 1.00x | AoS | +| 10 | 1024 | 8 | 20 | 20.854 | 20.786 | 1.00x | AoS | +| 11 | 1024 | 12 | 20 | 24.392 | 24.255 | 1.01x | AoS | +| 12 | 1024 | 16 | 20 | 29.433 | 29.423 | 1.00x | AoS | +| 13 | 1024 | 4 | 1 | 5.199 | 4.635 | 1.12x | SoA | +| 14 | 1024 | 4 | 5 | 6.498 | 4.521 | 1.44x | SoA | +| 15 | 1024 | 4 | 40 | 30.409 | 6.498 | 4.68x | SoA | +| 16 | 1024 | 4 | 20 | 17.030 | 5.760 | 2.96x | SoA | +| 17 | 1024 | 4 | 20 | 16.805 | 5.376 | 3.13x | SoA | +| 18 | 1 | 2 | 20 | 3.014 | 3.070 | 0.98x | 小 batch | +| 19 | 1024 | 4 | 20 | 16.809 | 5.524 | 3.04x | SoA | +| 20 | 1024 | 4 | 20 | 17.056 | 5.724 | 2.98x | SoA | +| **geomean** | | | | | | **2.15x** | | + +## 怎么读 + +- **SoA**（M≤4 且每核行数够）：CombFrag 改成 plane-major，行/列归一化变成逐元素运算，大 batch 和深迭代收益最大。 +- **AoS**（M≥6）：寄存器装不下 M² 个 plane，路径未改，时间不变。 +- **小 batch**：rowFactor<4 不走 SoA，时间在测量噪声内。 + +墙钟（含 aclnn 下发）约 260–400 µs，不能当 kernel 时间。 diff --git a/xllm_ops/moe/hc_pre_sinkhorn/op_host/CMakeLists.txt b/xllm_ops/moe/hc_pre_sinkhorn/op_host/CMakeLists.txt index 54d9747..74f3ca8 100644 --- a/xllm_ops/moe/hc_pre_sinkhorn/op_host/CMakeLists.txt +++ b/xllm_ops/moe/hc_pre_sinkhorn/op_host/CMakeLists.txt @@ -52,6 +52,7 @@ add_ops_compile_options( OP_NAME HcPreSinkhorn OPTIONS --cce-auto-sync=off -Wno-deprecated-declarations + -Wno-constant-conversion -Werror -mllvm -cce-aicore-hoist-movemask=false --op_relocatable_kernel_binary=true diff --git a/xllm_ops/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn_regbase_base.h b/xllm_ops/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn_regbase_base.h index 5261122..fa35ab5 100644 --- a/xllm_ops/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn_regbase_base.h +++ b/xllm_ops/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn_regbase_base.h @@ -26,6 +26,10 @@ using AscendC::MicroAPI::RegTensor; using AscendC::MicroAPI::UnalignReg; constexpr int32_t BLOCK_SIZE = 32; constexpr int32_t VL_FP32 = 64; +// Number of comb-matrix rows held in registers by the four-unfold fast path. +constexpr int32_t COMB_UNFOLD_NUM = 4; +// Below this many rows per tile the plane-major transpose overhead outweighs the gain. +constexpr int32_t COMB_SOA_MIN_ROWS = 4; __aicore__ inline int32_t CeilDiv(int32_t a, int b) { @@ -302,7 +306,350 @@ __aicore__ inline void VFProcessCombFragRLessVL( } } -__aicore__ inline void VFProcessIteration(RegTensor& sum0, RegTensor& sum1, RegTensor& mix, float eps, MaskReg pregLoop) +// [vec-05 + state_resident] Register-resident Sinkhorn iteration for generic M. +// Phase 1 (initial softmax + column norm) is identical to RLessVL (uses UB staging). +// Phase 2 (Sinkhorn iterations) loads M rows into RegTensor, iterates entirely in +// registers with no UB load/store and no LocalMemBar, then stores M rows back once. +// This eliminates iters * dim0 * M * 2 UB load/store + iters * dim0 LocalMemBar. +// +// Note: RegTensor arrays (RegTensor mix[M]) are NOT supported by bisheng backend +// ("Unsupported Inst must be hoisted"). Use individual variables + if constexpr chains, +// consistent with the FourUnfold path (mix1..mix4). The compiler eliminates unused +// variables and dead branches for each template instantiation. + +// Helper macros for compile-time unrolling over individual RegTensor variables +#define REG_LOAD_N(N) \ + if constexpr (M > N) { \ + LoadInputData(mix##N, combFragLocalAddr, pregLoop, \ + i * dim1 * dim2Align + (N) * dim2Align); \ + } + +#define REG_ROW_NORM_N(N) \ + if constexpr (M > N) { \ + ReduceSum(sumR, mix##N, pregLoop); \ + Duplicate(sumR, sumR, pregLoop); \ + Adds(sumR, sumR, eps, pregLoop); \ + Div(mix##N, mix##N, sumR, pregLoop); \ + Add(sumC, sumC, mix##N, pregLoop); \ + } + +#define REG_COL_NORM_N(N) \ + if constexpr (M > N) { \ + Div(mix##N, mix##N, sumC, pregLoop); \ + } + +#define REG_STORE_N(N) \ + if constexpr (M > N) { \ + StoreOutputData(combFragLocalAddr, mix##N, pregLoop, \ + i * dim1 * dim2Align + (N) * dim2Align); \ + } + +// Expand all 16 slots +#define REG_LOAD_ALL \ + REG_LOAD_N(0) REG_LOAD_N(1) REG_LOAD_N(2) REG_LOAD_N(3) \ + REG_LOAD_N(4) REG_LOAD_N(5) REG_LOAD_N(6) REG_LOAD_N(7) \ + REG_LOAD_N(8) REG_LOAD_N(9) REG_LOAD_N(10) REG_LOAD_N(11) \ + REG_LOAD_N(12) REG_LOAD_N(13) REG_LOAD_N(14) REG_LOAD_N(15) + +#define REG_ROW_NORM_ALL \ + REG_ROW_NORM_N(0) REG_ROW_NORM_N(1) REG_ROW_NORM_N(2) REG_ROW_NORM_N(3) \ + REG_ROW_NORM_N(4) REG_ROW_NORM_N(5) REG_ROW_NORM_N(6) REG_ROW_NORM_N(7) \ + REG_ROW_NORM_N(8) REG_ROW_NORM_N(9) REG_ROW_NORM_N(10) REG_ROW_NORM_N(11) \ + REG_ROW_NORM_N(12) REG_ROW_NORM_N(13) REG_ROW_NORM_N(14) REG_ROW_NORM_N(15) + +#define REG_COL_NORM_ALL \ + REG_COL_NORM_N(0) REG_COL_NORM_N(1) REG_COL_NORM_N(2) REG_COL_NORM_N(3) \ + REG_COL_NORM_N(4) REG_COL_NORM_N(5) REG_COL_NORM_N(6) REG_COL_NORM_N(7) \ + REG_COL_NORM_N(8) REG_COL_NORM_N(9) REG_COL_NORM_N(10) REG_COL_NORM_N(11) \ + REG_COL_NORM_N(12) REG_COL_NORM_N(13) REG_COL_NORM_N(14) REG_COL_NORM_N(15) + +#define REG_STORE_ALL \ + REG_STORE_N(0) REG_STORE_N(1) REG_STORE_N(2) REG_STORE_N(3) \ + REG_STORE_N(4) REG_STORE_N(5) REG_STORE_N(6) REG_STORE_N(7) \ + REG_STORE_N(8) REG_STORE_N(9) REG_STORE_N(10) REG_STORE_N(11) \ + REG_STORE_N(12) REG_STORE_N(13) REG_STORE_N(14) REG_STORE_N(15) + +template +__aicore__ inline void VFProcessCombFragRegResident( + const LocalTensor& combFragLocal, const LocalTensor& mixLocal, const LocalTensor& hcBaseLocal, + const LocalTensor& rsqrtLocal, float scale, float eps, uint16_t iters, uint16_t dim0, uint16_t dim1, + uint16_t dim2) +{ + __local_mem__ float* combFragLocalAddr = (__local_mem__ float*)combFragLocal.GetPhyAddr(); + __local_mem__ float* mixLocalAddr = (__local_mem__ float*)mixLocal.GetPhyAddr(); + __local_mem__ float* hcBaseLocalAddr = (__local_mem__ float*)hcBaseLocal.GetPhyAddr(); + __local_mem__ float* rsqrtLocalAddr = (__local_mem__ float*)rsqrtLocal.GetPhyAddr(); + uint32_t dim2Align = RoundUp(dim2); + __VEC_SCOPE__ + { + RegTensor base; + RegTensor mix; + RegTensor rsqrt; + RegTensor max; + RegTensor sum; + RegTensor sum1; + uint32_t sreg = dim2; + // [vec-09] mask creation done once, outside all loops + MaskReg pregLoop = UpdateMask(sreg); + + // Phase 1a: Initial softmax per row + accumulate column sum + for (uint16_t i = 0; i < dim0; i++) { + Duplicate(sum1, static_cast(0), pregLoop); + LoadInputDataWithBrc(rsqrt, rsqrtLocalAddr, pregLoop, i); + for (uint16_t j = 0; j < dim1; j++) { + LoadInputData(base, hcBaseLocalAddr, pregLoop, j * dim2Align); + LoadInputData(mix, mixLocalAddr, pregLoop, i * dim1 * dim2Align + j * dim2Align); + Mul(mix, mix, rsqrt, pregLoop); + Muls(mix, mix, scale, pregLoop); + Add(mix, mix, base, pregLoop); + ReduceMax(max, mix, pregLoop); + Duplicate(max, max, pregLoop); + Sub(mix, mix, max, pregLoop); + Exp(mix, mix, pregLoop); + ReduceSum(sum, mix, pregLoop); + Duplicate(sum, sum, pregLoop); + Div(mix, mix, sum, pregLoop); + Adds(mix, mix, eps, pregLoop); + Add(sum1, sum1, mix, pregLoop); + StoreOutputData(combFragLocalAddr, mix, pregLoop, i * dim1 * dim2Align + j * dim2Align); + } + // Phase 1b: Column normalization + LocalMemBar(); + Adds(sum1, sum1, eps, pregLoop); + for (uint16_t j = 0; j < dim1; j++) { + LoadInputData(mix, combFragLocalAddr, pregLoop, i * dim1 * dim2Align + j * dim2Align); + Div(mix, mix, sum1, pregLoop); + StoreOutputData(combFragLocalAddr, mix, pregLoop, i * dim1 * dim2Align + j * dim2Align); + } + } + + // Phase 2: Register-resident Sinkhorn iterations + // [vec-05] Load M rows into individual RegTensors, iterate without UB staging. + // Declare up to 16 individual RegTensor variables; compiler eliminates unused. + LocalMemBar(); + RegTensor mix0, mix1, mix2, mix3, mix4, mix5, mix6, mix7; + RegTensor mix8, mix9, mix10, mix11, mix12, mix13, mix14, mix15; + RegTensor sumR; // row sum (reused per row) + RegTensor sumC; // column sum accumulator + + for (uint16_t i = 0; i < dim0; i++) { + // Load M rows into registers (one-time UB load per dim0 iteration) + REG_LOAD_ALL + + // Sinkhorn iterations: all in registers, no UB load/store, no LocalMemBar + for (uint16_t iter = 0; iter < iters; iter++) { + Duplicate(sumC, static_cast(0), pregLoop); + // Row normalization + accumulate column sum + REG_ROW_NORM_ALL + // Column normalization + Adds(sumC, sumC, eps, pregLoop); + REG_COL_NORM_ALL + } + + // Store M rows back to UB (one-time UB store per dim0 iteration) + REG_STORE_ALL + } + } +} + +// =================================================================================== +// [round3 / vec-10] SoA (plane-major) comb_frag path for hcMult <= 4. +// +// Stores the M*M matrix entries as separate UB planes of length aAlign, so both +// Sinkhorn reductions become elementwise ops at full lane occupancy. +// MTE2 : NDDMA MultiCopy does (A, M*M) -> (M*M, A) on the fly. +// VEC : register-resident planes, no UB traffic per iteration. +// MTE3 : TransDataTo5HD (ldva address list) converts back to GM layout. +// =================================================================================== +constexpr int32_t COMB_SOA_MAX_MULT = 4; +constexpr int32_t TRANS_BLOCK_HALF = 8; +constexpr int32_t TRANS_BLOCK_FULL = 16; +// GM [A, hcMix] -> UB planes [planes, aAlign] via NDDMA MultiCopy. +// opc for this op is compiled with -Wno-constant-conversion so default +// NdDmaConfig::unsetPad survives -Werror. +__aicore__ inline void CopyInCombTransposed( + const GlobalTensor& mixesGm, const LocalTensor& combT, uint32_t curA, uint32_t aAlign, + uint32_t hcMix, uint32_t planes) +{ + MultiCopyLoopInfo<2> loopInfo; + loopInfo.loopSrcStride[0] = 1; + loopInfo.loopDstStride[0] = aAlign; + loopInfo.loopSize[0] = planes; + loopInfo.loopSrcStride[1] = hcMix; + loopInfo.loopDstStride[1] = 1; + loopInfo.loopSize[1] = curA; + MultiCopyParams params = {loopInfo, 0.0f}; + DataCopy(combT, mixesGm, params); +} + +#define SOA_HAS(R, C) ((R) < M && (C) < M) +#define SOA_PLANE(R, C) (((R) * M + (C)) * aAlign + off) + +#define SOA_AFFINE(R, C) \ + if constexpr (SOA_HAS(R, C)) { \ + DataCopy(c##R##C, combAddr + SOA_PLANE(R, C)); \ + Mul(c##R##C, c##R##C, rsq, preg); \ + Muls(c##R##C, c##R##C, scale, preg); \ + LoadInputDataWithBrc(bas, baseAddr, preg, (R) * M + (C)); \ + Add(c##R##C, c##R##C, bas, preg); \ + } + +#define SOA_AFFINE_ROW(R) SOA_AFFINE(R, 0) SOA_AFFINE(R, 1) SOA_AFFINE(R, 2) SOA_AFFINE(R, 3) + +#define SOA_REDUCE_ROW(OP, R) \ + OP(acc, c##R##0, c##R##1, preg); \ + if constexpr (M > 2) { OP(acc, acc, c##R##2, preg); } \ + if constexpr (M > 3) { OP(acc, acc, c##R##3, preg); } + +#define SOA_APPLY_ROW(OP, R) \ + OP(c##R##0, c##R##0, acc, preg); \ + OP(c##R##1, c##R##1, acc, preg); \ + if constexpr (M > 2) { OP(c##R##2, c##R##2, acc, preg); } \ + if constexpr (M > 3) { OP(c##R##3, c##R##3, acc, preg); } + +#define SOA_EXP_ROW(R) \ + Exp(c##R##0, c##R##0, preg); \ + Exp(c##R##1, c##R##1, preg); \ + if constexpr (M > 2) { Exp(c##R##2, c##R##2, preg); } \ + if constexpr (M > 3) { Exp(c##R##3, c##R##3, preg); } + +#define SOA_ADDS_ROW(R) \ + Adds(c##R##0, c##R##0, eps, preg); \ + Adds(c##R##1, c##R##1, eps, preg); \ + if constexpr (M > 2) { Adds(c##R##2, c##R##2, eps, preg); } \ + if constexpr (M > 3) { Adds(c##R##3, c##R##3, eps, preg); } + +#define SOA_SOFTMAX_ROW(R) \ + if constexpr ((R) < M) { \ + SOA_REDUCE_ROW(Max, R) \ + SOA_APPLY_ROW(Sub, R) \ + SOA_EXP_ROW(R) \ + SOA_REDUCE_ROW(Add, R) \ + SOA_APPLY_ROW(Div, R) \ + SOA_ADDS_ROW(R) \ + } + +#define SOA_ROW_NORM(R) \ + if constexpr ((R) < M) { \ + SOA_REDUCE_ROW(Add, R) \ + Adds(acc, acc, eps, preg); \ + SOA_APPLY_ROW(Div, R) \ + } + +#define SOA_COL_NORM(C) \ + if constexpr ((C) < M) { \ + Add(acc, c0##C, c1##C, preg); \ + if constexpr (M > 2) { Add(acc, acc, c2##C, preg); } \ + if constexpr (M > 3) { Add(acc, acc, c3##C, preg); } \ + Adds(acc, acc, eps, preg); \ + Div(c0##C, c0##C, acc, preg); \ + Div(c1##C, c1##C, acc, preg); \ + if constexpr (M > 2) { Div(c2##C, c2##C, acc, preg); } \ + if constexpr (M > 3) { Div(c3##C, c3##C, acc, preg); } \ + } + +#define SOA_STORE(R, C) \ + if constexpr (SOA_HAS(R, C)) { \ + DataCopy(combAddr + SOA_PLANE(R, C), c##R##C, preg); \ + } + +#define SOA_STORE_ROW(R) SOA_STORE(R, 0) SOA_STORE(R, 1) SOA_STORE(R, 2) SOA_STORE(R, 3) + +template +__aicore__ inline void VFProcessCombFragSoA( + const LocalTensor& combT, const LocalTensor& hcBaseLocal, const LocalTensor& rsqrtLocal, + float scale, float eps, uint16_t iters, uint16_t curA, uint32_t aAlign) +{ + __local_mem__ float* combAddr = (__local_mem__ float*)combT.GetPhyAddr(); + __local_mem__ float* baseAddr = (__local_mem__ float*)hcBaseLocal.GetPhyAddr(); + __local_mem__ float* rsqrtAddr = (__local_mem__ float*)rsqrtLocal.GetPhyAddr(); + uint16_t chunkCount = CeilDiv(curA, VL_FP32); + __VEC_SCOPE__ + { + RegTensor c00, c01, c02, c03; + RegTensor c10, c11, c12, c13; + RegTensor c20, c21, c22, c23; + RegTensor c30, c31, c32, c33; + RegTensor rsq; + RegTensor bas; + RegTensor acc; + uint32_t remain = curA; + MaskReg preg; + for (uint16_t chunk = 0; chunk < chunkCount; chunk++) { + preg = UpdateMask(remain); + uint32_t off = chunk * VL_FP32; + DataCopy(rsq, rsqrtAddr + off); + SOA_AFFINE_ROW(0) + SOA_AFFINE_ROW(1) + SOA_AFFINE_ROW(2) + SOA_AFFINE_ROW(3) + SOA_SOFTMAX_ROW(0) + SOA_SOFTMAX_ROW(1) + SOA_SOFTMAX_ROW(2) + SOA_SOFTMAX_ROW(3) + SOA_COL_NORM(0) + SOA_COL_NORM(1) + SOA_COL_NORM(2) + SOA_COL_NORM(3) + for (uint16_t iter = 0; iter < iters; iter++) { + SOA_ROW_NORM(0) + SOA_ROW_NORM(1) + SOA_ROW_NORM(2) + SOA_ROW_NORM(3) + SOA_COL_NORM(0) + SOA_COL_NORM(1) + SOA_COL_NORM(2) + SOA_COL_NORM(3) + } + SOA_STORE_ROW(0) + SOA_STORE_ROW(1) + SOA_STORE_ROW(2) + SOA_STORE_ROW(3) + } + } +} + +__aicore__ inline void TransposeCombSoAToAoS( + const LocalTensor& dst, const LocalTensor& src, const LocalTensor& vaAddr, + uint16_t curA, uint32_t aAlign, int32_t planes) +{ + int32_t groups = CeilDiv(planes, TRANS_BLOCK_HALF); + uint8_t aRepeat = static_cast(CeilDiv(curA, TRANS_BLOCK_FULL)); + TransDataTo5HDParams params; + params.repeatTimes = aRepeat; + params.srcRepStride = (aRepeat == 1) ? 0 : 2; + params.dstRepStride = (aRepeat == 1) ? 0 : TRANS_BLOCK_FULL * (TRANS_BLOCK_FULL / TRANS_BLOCK_HALF); + uint64_t dstBase = (uint64_t)(__ubuf__ float*)dst.GetPhyAddr(); + uint64_t srcBase = (uint64_t)(__ubuf__ float*)src.GetPhyAddr(); + LocalTensor dstAddr = vaAddr; + LocalTensor srcAddr = vaAddr[TRANS_BLOCK_FULL]; + // xllm-ops compiles with --cce-auto-sync=off. SetValue is scalar-pipe UB + // writes; TransDataTo5HD consumes the address list on the vector pipe. + PipeBarrier(); + event_t evSV = static_cast(GetTPipePtr()->FetchEventID(HardEvent::S_V)); + event_t evVS = static_cast(GetTPipePtr()->FetchEventID(HardEvent::V_S)); + for (int32_t i = 0; i < groups; i++) { + for (int32_t j = 0; j < TRANS_BLOCK_HALF; j++) { + int32_t plane = i * TRANS_BLOCK_HALF + j; + uint64_t srcOff = (uint64_t)(plane < planes ? plane : 0) * aAlign * sizeof(float); + srcAddr.SetValue(j, srcBase + srcOff); + srcAddr.SetValue(j + TRANS_BLOCK_HALF, srcBase + srcOff + TRANS_BLOCK_HALF * sizeof(float)); + dstAddr.SetValue( + j * 2, dstBase + (uint64_t)(i * TRANS_BLOCK_HALF + j * TRANS_BLOCK_FULL) * sizeof(float)); + dstAddr.SetValue( + j * 2 + 1, + dstBase + + (uint64_t)(i * TRANS_BLOCK_HALF + (j + TRANS_BLOCK_HALF) * TRANS_BLOCK_FULL) * sizeof(float)); + } + SetFlag(evSV); + WaitFlag(evSV); + AscendC::TransDataTo5HD(dstAddr, srcAddr, params); + SetFlag(evVS); + WaitFlag(evVS); + } +} + +__aicore__ inline void VFProcessIteration(RegTensor& sum0, RegTensor& sum1, RegTensor& mix, float eps, MaskReg pregLoop) { ReduceSum(sum1, mix, pregLoop); Duplicate(sum1, sum1, pregLoop); @@ -506,6 +853,23 @@ __aicore__ inline void CopyOut( DataCopyPad(outputGm, outputTensor, dataCopyParams); } +__aicore__ inline void CopyOutCombSoA( + const LocalTensor& outputTensor, const GlobalTensor& outputGm, uint32_t curA, int32_t planes) +{ + if (planes == TRANS_BLOCK_FULL) { + CopyOut(outputTensor, outputGm, 1, curA * TRANS_BLOCK_FULL); + return; + } + constexpr int32_t rowBytes = TRANS_BLOCK_FULL * sizeof(float); + int32_t burstBytes = planes * sizeof(float); + DataCopyExtParams dataCopyParams; + dataCopyParams.blockCount = static_cast(curA); + dataCopyParams.blockLen = burstBytes; + dataCopyParams.srcStride = (rowBytes - CeilAlign(burstBytes, ONE_BLK_SIZE)) / ONE_BLK_SIZE; + dataCopyParams.dstStride = 0; + DataCopyPad(outputGm, outputTensor, dataCopyParams); +} + } // namespace HCPreSinkhorn #endif \ No newline at end of file diff --git a/xllm_ops/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn_regbase_perf.h b/xllm_ops/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn_regbase_perf.h index 56e9c52..25091f1 100644 --- a/xllm_ops/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn_regbase_perf.h +++ b/xllm_ops/moe/hc_pre_sinkhorn/op_kernel/hc_pre_sinkhorn_regbase_perf.h @@ -43,12 +43,20 @@ class HcPreSinkhornPerf { postGm.SetGlobalBuffer((__gm__ float*)post); combFragGm.SetGlobalBuffer((__gm__ float*)combFrag); + useCombSoa = (tilingData->hcMult <= COMB_SOA_MAX_MULT) && (tilingData->rowFactor >= COMB_SOA_MIN_ROWS); + combPlanes = tilingData->hcMult * tilingData->hcMult; + combAAlign = CeilAlign(static_cast(tilingData->rowFactor), VL_FP32); + // InQue int64_t mixesQue01Size = tilingData->rowFactor * tilingData->hcMultAlign * 2 * sizeof(float); pipe->InitBuffer(mixesQue01, 2, mixesQue01Size); - pipe->InitBuffer( - mixesQue2, 2, tilingData->rowFactor * tilingData->hcMult * tilingData->hcMultAlign * sizeof(float)); - pipe->InitBuffer(rsqrtQue, 2, RoundUp(tilingData->rowFactor) * sizeof(float)); + int64_t combInSize = useCombSoa ? + (int64_t)combPlanes * combAAlign * sizeof(float) : + tilingData->rowFactor * tilingData->hcMult * tilingData->hcMultAlign * sizeof(float); + pipe->InitBuffer(mixesQue2, 2, combInSize); + int64_t rsqrtSize = useCombSoa ? combAAlign * sizeof(float) : + RoundUp(tilingData->rowFactor) * sizeof(float); + pipe->InitBuffer(rsqrtQue, 2, rsqrtSize); pipe->InitBuffer( xQue, 2, tilingData->rowFactor * tilingData->hcMult * RoundUp(tilingData->dFactor) * sizeof(T)); @@ -56,8 +64,8 @@ class HcPreSinkhornPerf { pipe->InitBuffer( yQue, 2, tilingData->rowFactor * RoundUp(tilingData->dFactor) * sizeof(T)); pipe->InitBuffer(postQue, 2, tilingData->rowFactor * tilingData->hcMultAlign * sizeof(float)); - pipe->InitBuffer( - combFragQue, 2, tilingData->rowFactor * tilingData->hcMult * tilingData->hcMultAlign * sizeof(float)); + int64_t combOutSize = useCombSoa ? (int64_t)TRANS_BLOCK_FULL * combAAlign * sizeof(float) : combInSize; + pipe->InitBuffer(combFragQue, 2, combOutSize); // TBuf pipe->InitBuffer(hcBaseBuf0, tilingData->hcMultAlign * sizeof(float)); @@ -67,6 +75,11 @@ class HcPreSinkhornPerf { hcBase0Local = hcBaseBuf0.Get(); hcBase1Local = hcBaseBuf1.Get(); hcBase2Local = hcBaseBuf2.Get(); + + if (useCombSoa) { + pipe->InitBuffer(vaAddrBuf, 2 * TRANS_BLOCK_FULL * sizeof(uint64_t)); + vaAddrLocal = vaAddrBuf.Get(); + } } __aicore__ inline void Process() @@ -81,7 +94,11 @@ class HcPreSinkhornPerf { CopyIn(hcBaseGm, hcBase0Local, 1, tilingData->hcMult); CopyIn(hcBaseGm[tilingData->hcMult], hcBase1Local, 1, tilingData->hcMult); - CopyIn(hcBaseGm[tilingData->hcMult * 2], hcBase2Local, tilingData->hcMult, tilingData->hcMult); + if (useCombSoa) { + CopyIn(hcBaseGm[tilingData->hcMult * 2], hcBase2Local, 1, tilingData->hcMult * tilingData->hcMult); + } else { + CopyIn(hcBaseGm[tilingData->hcMult * 2], hcBase2Local, tilingData->hcMult, tilingData->hcMult); + } event_t eventId = static_cast(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V)); SetFlag(eventId); WaitFlag(eventId); @@ -142,23 +159,101 @@ class HcPreSinkhornPerf { // combFrag mixes2Local = mixesQue2.AllocTensor(); - CopyInWithLoopMode( - mixesGm - [mixGmBaseOffset + rowOuterIdx * tilingData->rowFactor * tilingData->hcMix + tilingData->hcMult * 2], - mixes2Local, curRowFactor, tilingData->hcMult, tilingData->hcMult, tilingData->hcMix); + if (useCombSoa) { + CopyInCombTransposed( + mixesGm[mixGmBaseOffset + rowOuterIdx * tilingData->rowFactor * tilingData->hcMix + + tilingData->hcMult * 2], + mixes2Local, curRowFactor, combAAlign, tilingData->hcMix, combPlanes); + } else { + CopyInWithLoopMode( + mixesGm[mixGmBaseOffset + rowOuterIdx * tilingData->rowFactor * tilingData->hcMix + + tilingData->hcMult * 2], + mixes2Local, curRowFactor, tilingData->hcMult, tilingData->hcMult, tilingData->hcMix); + } mixesQue2.EnQue(mixes2Local); mixes2Local = mixesQue2.DeQue(); combFragLocal = combFragQue.AllocTensor(); - VFProcessCombFragRLessVLUseFourUnfold( - combFragLocal, mixes2Local, hcBase2Local, rsqrtLocal, hcScaleGm.GetValue(2), tilingData->eps, - tilingData->iterTimes - 1, curRowFactor, tilingData->hcMult, tilingData->hcMult); + // Branch-S (hcMult<=4, enough rows): plane-major SoA, full lane occupancy + // Branch-A (hcMult==4): FourUnfold fast path (already register-resident) + // Branch-B (hcMult in {2,3,6,8,12,16}): RLessVL register-resident via template + // Branch-C (other): fallback to original RLessVL UB-staging path + if (useCombSoa) { + switch (tilingData->hcMult) { + case 2: + VFProcessCombFragSoA<2>( + mixes2Local, hcBase2Local, rsqrtLocal, hcScaleGm.GetValue(2), tilingData->eps, + tilingData->iterTimes - 1, curRowFactor, combAAlign); + break; + case 3: + VFProcessCombFragSoA<3>( + mixes2Local, hcBase2Local, rsqrtLocal, hcScaleGm.GetValue(2), tilingData->eps, + tilingData->iterTimes - 1, curRowFactor, combAAlign); + break; + default: + VFProcessCombFragSoA<4>( + mixes2Local, hcBase2Local, rsqrtLocal, hcScaleGm.GetValue(2), tilingData->eps, + tilingData->iterTimes - 1, curRowFactor, combAAlign); + break; + } + TransposeCombSoAToAoS(combFragLocal, mixes2Local, vaAddrLocal, curRowFactor, combAAlign, combPlanes); + } else if (tilingData->hcMult == COMB_UNFOLD_NUM) { + VFProcessCombFragRLessVLUseFourUnfold( + combFragLocal, mixes2Local, hcBase2Local, rsqrtLocal, hcScaleGm.GetValue(2), tilingData->eps, + tilingData->iterTimes - 1, curRowFactor, tilingData->hcMult, tilingData->hcMult); + } else { + // [vec-05 + state_resident] Register-resident Sinkhorn iteration + switch (tilingData->hcMult) { + case 2: + VFProcessCombFragRegResident<2>( + combFragLocal, mixes2Local, hcBase2Local, rsqrtLocal, hcScaleGm.GetValue(2), tilingData->eps, + tilingData->iterTimes - 1, curRowFactor, tilingData->hcMult, tilingData->hcMult); + break; + case 3: + VFProcessCombFragRegResident<3>( + combFragLocal, mixes2Local, hcBase2Local, rsqrtLocal, hcScaleGm.GetValue(2), tilingData->eps, + tilingData->iterTimes - 1, curRowFactor, tilingData->hcMult, tilingData->hcMult); + break; + case 6: + VFProcessCombFragRegResident<6>( + combFragLocal, mixes2Local, hcBase2Local, rsqrtLocal, hcScaleGm.GetValue(2), tilingData->eps, + tilingData->iterTimes - 1, curRowFactor, tilingData->hcMult, tilingData->hcMult); + break; + case 8: + VFProcessCombFragRegResident<8>( + combFragLocal, mixes2Local, hcBase2Local, rsqrtLocal, hcScaleGm.GetValue(2), tilingData->eps, + tilingData->iterTimes - 1, curRowFactor, tilingData->hcMult, tilingData->hcMult); + break; + case 12: + VFProcessCombFragRegResident<12>( + combFragLocal, mixes2Local, hcBase2Local, rsqrtLocal, hcScaleGm.GetValue(2), tilingData->eps, + tilingData->iterTimes - 1, curRowFactor, tilingData->hcMult, tilingData->hcMult); + break; + case 16: + VFProcessCombFragRegResident<16>( + combFragLocal, mixes2Local, hcBase2Local, rsqrtLocal, hcScaleGm.GetValue(2), tilingData->eps, + tilingData->iterTimes - 1, curRowFactor, tilingData->hcMult, tilingData->hcMult); + break; + default: + VFProcessCombFragRLessVL( + combFragLocal, mixes2Local, hcBase2Local, rsqrtLocal, hcScaleGm.GetValue(2), tilingData->eps, + tilingData->iterTimes - 1, curRowFactor, tilingData->hcMult, tilingData->hcMult); + break; + } + } mixesQue2.FreeTensor(mixes2Local); rsqrtQue.FreeTensor(rsqrtLocal); combFragQue.EnQue(combFragLocal); combFragLocal = combFragQue.DeQue(); - CopyOut(combFragLocal, combFragGm[curBlockIdx * tilingData->rowOfFormerBlock * tilingData->hcMult * tilingData->hcMult + rowOuterIdx * tilingData->rowFactor * tilingData->hcMult * tilingData->hcMult], curRowFactor * tilingData->hcMult, tilingData->hcMult); + int64_t combGmOffset = curBlockIdx * tilingData->rowOfFormerBlock * tilingData->hcMult * tilingData->hcMult + + rowOuterIdx * tilingData->rowFactor * tilingData->hcMult * tilingData->hcMult; + if (useCombSoa) { + CopyOutCombSoA(combFragLocal, combFragGm[combGmOffset], curRowFactor, combPlanes); + } else { + CopyOut( + combFragLocal, combFragGm[combGmOffset], curRowFactor * tilingData->hcMult, tilingData->hcMult); + } combFragQue.FreeTensor(combFragLocal); } } @@ -166,6 +261,9 @@ class HcPreSinkhornPerf { private: TPipe* pipe; const HcPreSinkhornTilingData* tilingData; + bool useCombSoa = false; + int32_t combPlanes = 0; + int32_t combAAlign = 0; GlobalTensor mixesGm; GlobalTensor rsqrtGm; GlobalTensor hcScaleGm; @@ -186,6 +284,7 @@ class HcPreSinkhornPerf { TBuf hcBaseBuf0; TBuf hcBaseBuf1; TBuf hcBaseBuf2; + TBuf vaAddrBuf; LocalTensor mixes01Local; LocalTensor mixes2Local; @@ -197,6 +296,7 @@ class HcPreSinkhornPerf { LocalTensor hcBase0Local; LocalTensor hcBase1Local; LocalTensor hcBase2Local; + LocalTensor vaAddrLocal; }; } // namespace HCPreSinkhorn