diff --git a/third_party/catlass b/third_party/catlass index c02a6e8..a527ccd 160000 --- a/third_party/catlass +++ b/third_party/catlass @@ -1 +1 @@ -Subproject commit c02a6e8d9055d79601bee66e9dbfab24fffc41ee +Subproject commit a527ccd238515c921e2f1a33de4eec387b05256b diff --git a/xllm_ops/build_aclnn.sh b/xllm_ops/build_aclnn.sh index c08f339..e9c54b6 100644 --- a/xllm_ops/build_aclnn.sh +++ b/xllm_ops/build_aclnn.sh @@ -150,6 +150,7 @@ elif [[ "$SOC_VERSION" =~ ^(ascend)?910b ]]; then "moe_init_routing_v3" "multi_latent_attention" "pp_matmul_opt" + "quant_matmul_nz_decode" "recurrent_gated_delta_rule" "replace_token" "mtp_prepare_next_draft" @@ -266,6 +267,7 @@ elif [[ "$SOC_VERSION" =~ ^ascend910_93 ]]; then "moe_grouped_matmul_swiglu_quant" "multi_latent_attention" "pp_matmul_opt" + "quant_matmul_nz_decode" "recurrent_gated_delta_rule" "replace_token" "mtp_prepare_next_draft" @@ -397,10 +399,39 @@ dump_selected_ops : "${SOC_VERSION:?SOC_VERSION is not set}" : "${SOC_ARG:?SOC_ARG is not set}" + if [[ -n "${BUILD_DIR_ARG}" && + ";${CUSTOM_OPS};" == *";quant_matmul_nz_decode;"* ]]; then + # The generated kernel targets do not depend on op_kernel sources. Force + # OPC to refresh this operator when reusing an incremental build tree. + quant_matmul_binary_dir="${BUILD_DIR_ARG}/binary/${SOC_ARG}" + rm -f -- \ + "${quant_matmul_binary_dir}/src/quant_matmul_nz_decode/quant_matmul_nz_decode_${SOC_ARG}_src_copy.done" \ + "${quant_matmul_binary_dir}/gen/quant_matmul_nz_decode_${SOC_ARG}_0.done" \ + "${quant_matmul_binary_dir}/gen/quant_matmul_nz_decode_${SOC_ARG}_1.done" \ + "${quant_matmul_binary_dir}/gen/quant_matmul_nz_decode_${SOC_ARG}_2.done" \ + "${quant_matmul_binary_dir}/gen/quant_matmul_nz_decode_${SOC_ARG}_3.done" + + quant_matmul_gen_dir="${quant_matmul_binary_dir}/gen" + if [[ -d "${quant_matmul_gen_dir}" ]]; then + while IFS= read -r -d '' kernel_meta_dir; do + find "${kernel_meta_dir}" -depth -delete + done < <(find "${quant_matmul_gen_dir}" -mindepth 1 -maxdepth 1 \ + -type d -name 'kernel_meta_QuantMatmulNzDecode_*' -print0) + fi + + quant_matmul_output_dir="${quant_matmul_binary_dir}/bin/quant_matmul_nz_decode" + if [[ -d "${quant_matmul_output_dir}" ]]; then + find "${quant_matmul_output_dir}" -mindepth 1 -maxdepth 1 \ + -type f \( -name 'QuantMatmulNzDecode_*.o' \ + -o -name 'QuantMatmulNzDecode_*.json' \) -delete + fi + fi + log "build command: bash build.sh --pkg --ops=\"${CUSTOM_OPS}\" --soc=\"${SOC_ARG}\"" log "building custom ops ${CUSTOM_OPS} for ${SOC_VERSION}" if [ -n "$BUILD_DIR_ARG" ]; then - bash build.sh --pkg --ops="${CUSTOM_OPS}" --soc="${SOC_ARG}" --build-dir="${BUILD_DIR_ARG}" + bash build.sh --pkg --ops="${CUSTOM_OPS}" --soc="${SOC_ARG}" \ + --build-dir="${BUILD_DIR_ARG}" else bash build.sh --pkg --ops="${CUSTOM_OPS}" --soc="${SOC_ARG}" fi diff --git a/xllm_ops/quant_matmul_nz_decode/CMakeLists.txt b/xllm_ops/quant_matmul_nz_decode/CMakeLists.txt new file mode 100644 index 0000000..cade2ed --- /dev/null +++ b/xllm_ops/quant_matmul_nz_decode/CMakeLists.txt @@ -0,0 +1,20 @@ +# Copyright 2026 The xLLM Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://gitcode.com/xLLM-AI/xllm_ops/blob/main/LICENSE +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) +foreach(SUB_DIR ${CURRENT_DIRS}) + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt") + add_subdirectory(${SUB_DIR}) + endif() +endforeach() diff --git a/xllm_ops/quant_matmul_nz_decode/op_host/CMakeLists.txt b/xllm_ops/quant_matmul_nz_decode/op_host/CMakeLists.txt new file mode 100644 index 0000000..e4e5de5 --- /dev/null +++ b/xllm_ops/quant_matmul_nz_decode/op_host/CMakeLists.txt @@ -0,0 +1,39 @@ +# Copyright 2026 The xLLM Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://gitcode.com/xLLM-AI/xllm_ops/blob/main/LICENSE +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +add_op_to_compiled_list() + +if(BUILD_OPEN_PROJECT) + target_sources(op_host_aclnn PRIVATE + quant_matmul_nz_decode_def.cpp + ) +endif() + +add_ops_tiling_keys( + OP_NAME QuantMatmulNzDecode + TILING_KEYS 1 3 2 0 4 5 6 7 8 9 10 11 12 13 14 15 +) + +add_ops_compile_options( + OP_NAME QuantMatmulNzDecode + OPTIONS --cce-auto-sync=off + -DCATLASS_ENABLE_L1_L0_MN_TILE_MISMATCH + -Wno-deprecated-declarations + -Wno-error + -I${CMAKE_SOURCE_DIR}/../third_party/catlass/include +) + +if(NOT BUILD_OPS_RTY_KERNEL) + add_modules_sources(OPTYPE quant_matmul_nz_decode ACLNNTYPE aclnn) +endif() diff --git a/xllm_ops/quant_matmul_nz_decode/op_host/quant_matmul_nz_decode_def.cpp b/xllm_ops/quant_matmul_nz_decode/op_host/quant_matmul_nz_decode_def.cpp new file mode 100644 index 0000000..53c9214 --- /dev/null +++ b/xllm_ops/quant_matmul_nz_decode/op_host/quant_matmul_nz_decode_def.cpp @@ -0,0 +1,53 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://gitcode.com/xLLM-AI/xllm_ops/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "register/op_def_registry.h" + +namespace ops { +class QuantMatmulNzDecode : public OpDef { + public: + explicit QuantMatmulNzDecode(const char* name) : OpDef(name) { + this->Input("x") + .ParamType(REQUIRED) + .DataType({ge::DT_INT8}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Input("weight") + .ParamType(REQUIRED) + .DataType({ge::DT_INT8}) + .Format({ge::FORMAT_FRACTAL_NZ}) + .UnknownShapeFormat({ge::FORMAT_FRACTAL_NZ}); + this->Input("scale") + .ParamType(REQUIRED) + .DataType({ge::DT_FLOAT}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Input("bias") + .ParamType(REQUIRED) + .DataType({ge::DT_INT32}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->Output("y") + .ParamType(REQUIRED) + .DataType({ge::DT_BF16}) + .Format({ge::FORMAT_ND}) + .UnknownShapeFormat({ge::FORMAT_ND}); + this->AICore().AddConfig("ascend910b"); + this->AICore().AddConfig("ascend910_93"); + } +}; + +OP_ADD(QuantMatmulNzDecode); +} // namespace ops diff --git a/xllm_ops/quant_matmul_nz_decode/op_host/quant_matmul_nz_decode_proto.cpp b/xllm_ops/quant_matmul_nz_decode/op_host/quant_matmul_nz_decode_proto.cpp new file mode 100644 index 0000000..592e028 --- /dev/null +++ b/xllm_ops/quant_matmul_nz_decode/op_host/quant_matmul_nz_decode_proto.cpp @@ -0,0 +1,38 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://gitcode.com/xLLM-AI/xllm_ops/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "quant_matmul_nz_decode_tiling.h" +#include "register/op_def_registry.h" + +namespace ge { +static graphStatus InferShape(gert::InferShapeContext* context) { + const gert::Shape* x_shape = context->GetInputShape(0); + const gert::Shape* weight_shape = context->GetInputShape(1); + gert::Shape* y_shape = context->GetOutputShape(0); + *y_shape = *x_shape; + y_shape->SetDim(y_shape->GetDimNum() - 1, + weight_shape->GetDim(weight_shape->GetDimNum() - 1)); + return GRAPH_SUCCESS; +} + +static graphStatus InferDataType(gert::InferDataTypeContext* context) { + context->SetOutputDataType(0, ge::DT_BF16); + return GRAPH_SUCCESS; +} + +IMPL_OP(QuantMatmulNzDecode) + .InferShape(InferShape) + .InferDataType(InferDataType); +} // namespace ge diff --git a/xllm_ops/quant_matmul_nz_decode/op_host/quant_matmul_nz_decode_tiling.cpp b/xllm_ops/quant_matmul_nz_decode/op_host/quant_matmul_nz_decode_tiling.cpp new file mode 100644 index 0000000..0d49f2f --- /dev/null +++ b/xllm_ops/quant_matmul_nz_decode/op_host/quant_matmul_nz_decode_tiling.cpp @@ -0,0 +1,122 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://gitcode.com/xLLM-AI/xllm_ops/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "quant_matmul_nz_decode_tiling.h" + +#include "register/op_def_registry.h" +#include "tiling/platform/platform_ascendc.h" + +namespace optiling { +namespace { +constexpr uint32_t kCoreCount = 20; +constexpr uint32_t kGateUpCoreCount = 20; +constexpr uint32_t kQkvCoreCount = 10; +constexpr uint32_t kTileM = 16; +constexpr uint32_t kMaxOptimizedM = 16; +constexpr uint32_t kTileN = 320; +constexpr uint32_t kWorkspaceStages = 2; +constexpr uint64_t kGateUpM1TilingKey = 6; +constexpr uint64_t kGateUpM2TilingKey = 7; +constexpr uint64_t kGateUpM4TilingKey = 8; +constexpr uint64_t kGateUpM8TilingKey = 9; +constexpr uint64_t kGateUpM16TilingKey = 10; +constexpr uint64_t kQkvM1TilingKey = 11; +constexpr uint64_t kQkvM2TilingKey = 12; +constexpr uint64_t kQkvM4TilingKey = 13; +constexpr uint64_t kQkvM8TilingKey = 14; +constexpr uint64_t kQkvM16TilingKey = 15; + +ge::graphStatus TilingFunc(gert::TilingContext* context) { + const auto x_shape = context->GetInputShape(0)->GetOriginShape(); + const auto weight_shape = context->GetInputShape(1)->GetOriginShape(); + const auto scale_shape = context->GetInputShape(2)->GetOriginShape(); + const auto bias_shape = context->GetInputShape(3)->GetOriginShape(); + const size_t x_dim = x_shape.GetDimNum(); + const size_t weight_dim = weight_shape.GetDimNum(); + if (x_dim != 2 || weight_dim != 2 || scale_shape.GetDimNum() != 1 || + bias_shape.GetDimNum() != 1) { + return ge::GRAPH_FAILED; + } + + const int64_t m_dim = x_shape.GetDim(0); + const int64_t k_dim = x_shape.GetDim(1); + const int64_t n_dim = weight_shape.GetDim(1); + + const bool is_gate_up_shape = k_dim == 5120 && n_dim == 6400; + const bool is_down_shape = k_dim == 3200 && n_dim == 5120; + const bool is_qkv_shape = k_dim == 5120 && n_dim == 1280; + if (m_dim <= 0 || m_dim > kMaxOptimizedM || + (!is_gate_up_shape && !is_down_shape && !is_qkv_shape) || + weight_shape.GetDim(0) != k_dim || scale_shape.GetDim(0) != n_dim || + bias_shape.GetDim(0) != n_dim) { + return ge::GRAPH_FAILED; + } + + const uint32_t m = static_cast(m_dim); + const uint32_t k = static_cast(k_dim); + const uint32_t n = static_cast(n_dim); + QuantMatmulNzDecodeTilingData tiling; + tiling.set_m(m); + tiling.set_k(k); + tiling.set_n(n); + tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), + context->GetRawTilingData()->GetCapacity()); + context->GetRawTilingData()->SetDataSize(tiling.GetDataSize()); + context->SetBlockDim( + is_qkv_shape ? kQkvCoreCount + : (is_gate_up_shape ? kGateUpCoreCount : kCoreCount)); + if (is_down_shape) { + context->SetTilingKey(m == 1 ? 1 : (m <= 4 ? 3 : 2)); + } else if (is_qkv_shape) { + if (m == 1) { + context->SetTilingKey(kQkvM1TilingKey); + } else if (m == 2) { + context->SetTilingKey(kQkvM2TilingKey); + } else if (m == 4) { + context->SetTilingKey(kQkvM4TilingKey); + } else if (m == 8) { + context->SetTilingKey(kQkvM8TilingKey); + } else if (m == 16) { + context->SetTilingKey(kQkvM16TilingKey); + } else { + context->SetTilingKey(4); + } + } else if (m == 1) { + context->SetTilingKey(kGateUpM1TilingKey); + } else if (m == 2) { + context->SetTilingKey(kGateUpM2TilingKey); + } else if (m == 4) { + context->SetTilingKey(kGateUpM4TilingKey); + } else if (m == 8) { + context->SetTilingKey(kGateUpM8TilingKey); + } else if (m == 16) { + context->SetTilingKey(kGateUpM16TilingKey); + } else if (m > 2) { + context->SetTilingKey(5); + } else { + context->SetTilingKey(0); + } + + auto platform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); + size_t* workspace_size = context->GetWorkspaceSizes(1); + workspace_size[0] = + platform.GetLibApiWorkSpaceSize() + + kTileM * kTileN * kGateUpCoreCount * kWorkspaceStages * sizeof(int32_t); + return ge::GRAPH_SUCCESS; +} +} // namespace + +IMPL_OP_OPTILING(QuantMatmulNzDecode).Tiling(TilingFunc); +} // namespace optiling diff --git a/xllm_ops/quant_matmul_nz_decode/op_host/quant_matmul_nz_decode_tiling.h b/xllm_ops/quant_matmul_nz_decode/op_host/quant_matmul_nz_decode_tiling.h new file mode 100644 index 0000000..1b84108 --- /dev/null +++ b/xllm_ops/quant_matmul_nz_decode/op_host/quant_matmul_nz_decode_tiling.h @@ -0,0 +1,29 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://gitcode.com/xLLM-AI/xllm_ops/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include "register/tilingdata_base.h" + +namespace optiling { +BEGIN_TILING_DATA_DEF(QuantMatmulNzDecodeTilingData) + TILING_DATA_FIELD_DEF(uint32_t, m); + TILING_DATA_FIELD_DEF(uint32_t, k); + TILING_DATA_FIELD_DEF(uint32_t, n); +END_TILING_DATA_DEF; + +REGISTER_TILING_DATA_CLASS(QuantMatmulNzDecode, + QuantMatmulNzDecodeTilingData) +} // namespace optiling diff --git a/xllm_ops/quant_matmul_nz_decode/op_kernel/quant_matmul_nz_decode.cpp b/xllm_ops/quant_matmul_nz_decode/op_kernel/quant_matmul_nz_decode.cpp new file mode 100644 index 0000000..8e03aa2 --- /dev/null +++ b/xllm_ops/quant_matmul_nz_decode/op_kernel/quant_matmul_nz_decode.cpp @@ -0,0 +1,211 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://gitcode.com/xLLM-AI/xllm_ops/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "catlass/arch/arch.hpp" +#include "catlass/epilogue/block/block_epilogue.hpp" +#include "catlass/epilogue/dispatch_policy.hpp" +#include "catlass/epilogue/tile/tile_broadcast_mul.hpp" +#include "catlass/epilogue/tile/tile_broadcast_one_blk.hpp" +#include "catlass/epilogue/tile/tile_swizzle.hpp" +#include "catlass/gemm/block/block_mmad.hpp" +#include "catlass/gemm/block/block_swizzle.hpp" +#include "catlass/gemm/dispatch_policy.hpp" +#include "catlass/gemm/gemm_type.hpp" +#include "catlass/layout/layout.hpp" +#include "kernel_operator.h" +#include "quant_matmul_nz_workspace.hpp" + +using namespace Catlass; + +template +__aicore__ inline void RunQuantMatmulNzDecode(GM_ADDR x, + GM_ADDR weight, + GM_ADDR scale, + GM_ADDR bias, + GM_ADDR y, + GM_ADDR workspace, + uint32_t m, + uint32_t k, + uint32_t n) { + using ArchTag = Arch::AtlasA2; + using L1TileShape = GemmShape<16, kL1TileN, kL1TileK>; + using L0TileShape = GemmShape<16, kL0TileN, kL0TileK>; + using AType = Gemm::GemmType; + using BType = Gemm::GemmType; + using CType = Gemm::GemmType; + using DispatchPolicy = + Gemm::MmadAtlasA2PreloadAsyncWithCallback<1, + 2, + 2, + 2, + 1, + false, + kEnableShuffleK>; + using BlockMmad = Gemm::Block:: + BlockMmad; + + using EpiloguePolicy = Epilogue::EpilogueAtlasA2PerTokenDequant<2>; + using ScaleType = Gemm::GemmType; + using PerTokenScaleType = Gemm::GemmType; + using DType = Gemm::GemmType; + using RowBroadcastMulType = Gemm::GemmType; + using BroadcastOneBlkType = Gemm::GemmType; + using ColumnBroadcastMulType = Gemm::GemmType; + using EpilogueTileShape = MatrixShape<16, 256>; + using TileRowBroadcastMul = Epilogue::Tile:: + TileRowBroadcastMul; + using TileBroadcastOneBlk = Epilogue::Tile:: + TileBroadcastOneBlk; + using TileColumnBroadcastMul = + Epilogue::Tile::TileOneBlkColumnBroadcastMul; + using TileCopy = Epilogue::Tile:: + TileCopy; + using EpilogueScheduler = Epilogue::Tile::EpilogueHorizontalTileSwizzle; + using BlockEpilogue = Epilogue::Block::BlockEpilogue; + using BlockScheduler = Gemm::Block::GemmIdentityBlockSwizzle<3, 1>; + using MatmulKernel = Gemm::Kernel:: + QuantMatmulNzWorkspace; + + const uint32_t problem_m = kProblemM == 0 ? m : kProblemM; + const uint32_t problem_k = kProblemK == 0 ? k : kProblemK; + const uint32_t problem_n = kProblemN == 0 ? n : kProblemN; + GemmCoord problem_shape{problem_m, problem_n, problem_k}; + layout::RowMajor layout_x{problem_m, problem_k}; + auto layout_weight = + layout::zN::MakeLayout(problem_k, problem_n); + layout::VectorLayout layout_scale{problem_n}; + layout::VectorLayout unused_layout{problem_m}; + layout::RowMajor layout_y{problem_m, problem_n}; + typename MatmulKernel::Params params{problem_shape, + x, + layout_x, + weight, + layout_weight, + scale, + layout_scale, + scale, + unused_layout, + bias, + y, + layout_y, + AscendC::GetUserWorkspace(workspace)}; + MatmulKernel matmul; + matmul(params); +} + +extern "C" __global__ __aicore__ void quant_matmul_nz_decode(GM_ADDR x, + GM_ADDR weight, + GM_ADDR scale, + GM_ADDR bias, + GM_ADDR y, + GM_ADDR workspace, + GM_ADDR tiling) { + KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_1); + GET_TILING_DATA(tiling_data, tiling); + + const uint32_t m = tiling_data.m; + const uint32_t k = tiling_data.k; + const uint32_t n = tiling_data.n; + if (TILING_KEY_IS(6)) { + RunQuantMatmulNzDecode<320, 512, 128, 256, false, 1, 5120, 6400>( + x, weight, scale, bias, y, workspace, m, k, n); + } else if (TILING_KEY_IS(7)) { + RunQuantMatmulNzDecode<320, 512, 128, 256, false, 2, 5120, 6400>( + x, weight, scale, bias, y, workspace, m, k, n); + } else if (TILING_KEY_IS(8)) { + RunQuantMatmulNzDecode<160, 512, 128, 256, false, 4, 5120, 6400>( + x, weight, scale, bias, y, workspace, m, k, n); + } else if (TILING_KEY_IS(9)) { + RunQuantMatmulNzDecode<160, 512, 128, 256, false, 8, 5120, 6400>( + x, weight, scale, bias, y, workspace, m, k, n); + } else if (TILING_KEY_IS(10)) { + RunQuantMatmulNzDecode<160, 512, 128, 256, false, 16, 5120, 6400>( + x, weight, scale, bias, y, workspace, m, k, n); + } else if (TILING_KEY_IS(11)) { + RunQuantMatmulNzDecode<128, 1536, 128, 256, true, 1, 5120, 1280>( + x, weight, scale, bias, y, workspace, m, k, n); + } else if (TILING_KEY_IS(12)) { + RunQuantMatmulNzDecode<128, 1536, 128, 256, true, 2, 5120, 1280>( + x, weight, scale, bias, y, workspace, m, k, n); + } else if (TILING_KEY_IS(13)) { + RunQuantMatmulNzDecode<128, 1536, 128, 256, true, 4, 5120, 1280>( + x, weight, scale, bias, y, workspace, m, k, n); + } else if (TILING_KEY_IS(14)) { + RunQuantMatmulNzDecode<128, 1536, 128, 256, true, 8, 5120, 1280>( + x, weight, scale, bias, y, workspace, m, k, n); + } else if (TILING_KEY_IS(15)) { + RunQuantMatmulNzDecode<128, 1536, 128, 256, true, 16, 5120, 1280>( + x, weight, scale, bias, y, workspace, m, k, n); + } else if (TILING_KEY_IS(0)) { + RunQuantMatmulNzDecode<320, 512>( + x, weight, scale, bias, y, workspace, m, k, n); + } else if (TILING_KEY_IS(5)) { + RunQuantMatmulNzDecode<160, 512>( + x, weight, scale, bias, y, workspace, m, k, n); + } else if (TILING_KEY_IS(1)) { + RunQuantMatmulNzDecode<64, 3200, 64, 512, true>( + x, weight, scale, bias, y, workspace, m, k, n); + } else if (TILING_KEY_IS(3)) { + RunQuantMatmulNzDecode<128, 1792, 128, 256, true>( + x, weight, scale, bias, y, workspace, m, k, n); + } else if (TILING_KEY_IS(2)) { + RunQuantMatmulNzDecode<256, 896, 128, 256, true>( + x, weight, scale, bias, y, workspace, m, k, n); + } else if (TILING_KEY_IS(4)) { + RunQuantMatmulNzDecode<128, 1536, 128, 256, true>( + x, weight, scale, bias, y, workspace, m, k, n); + } +} + +// The generated mixed-core wrapper clears workspace through matmul helpers. +#include "lib/matmul_intf.h" + +namespace AscendC { + +__aicore__ inline void PrepareQuantMatmulNzMixedCoreWorkspace(__gm__ uint8_t*) { +#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 2201) + SetAtomicNone(); + SetMaskNorm(); + SetLoadDataBoundary(static_cast(0)); + SetLoadDataPaddingValue(static_cast(0)); + NotifyEvent(WORKSPACE_SYNC_ID); +#endif +} + +} // namespace AscendC + +// CATLASS uses FFTS cross-core flags directly and does not use the KFC message +// queues. Preserve the wrapper's completion notification, but skip clearing +// the unused 30 KiB KFC region on every AIC. +#define clearWorkspace(workspace) \ + PrepareQuantMatmulNzMixedCoreWorkspace(workspace) diff --git a/xllm_ops/quant_matmul_nz_decode/op_kernel/quant_matmul_nz_workspace.hpp b/xllm_ops/quant_matmul_nz_decode/op_kernel/quant_matmul_nz_workspace.hpp new file mode 100644 index 0000000..c83a505 --- /dev/null +++ b/xllm_ops/quant_matmul_nz_decode/op_kernel/quant_matmul_nz_workspace.hpp @@ -0,0 +1,489 @@ +/* + * Copyright (c) 2025 Huawei Technologies Co., Ltd. + * This file is a part of the CANN Open Software. + * Licensed under CANN Open Software License Agreement Version 1.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 "catlass/arch/cross_core_sync.hpp" +#include "catlass/arch/resource.hpp" +#include "catlass/catlass.hpp" +#include "catlass/coord.hpp" +#include "catlass/detail/callback.hpp" +#include "catlass/gemm_coord.hpp" +#include "catlass/matrix_coord.hpp" + +namespace Catlass::Gemm::Kernel { + +template +class QuantMatmulNzWorkspace { + public: + using BlockMmad = BlockMmad_; + using ArchTag = typename BlockMmad::ArchTag; + using L1TileShape = typename BlockMmad::L1TileShape; + using ElementA = typename BlockMmad::ElementA; + using LayoutA = typename BlockMmad::LayoutA; + using ElementB = typename BlockMmad::ElementB; + using LayoutB = typename BlockMmad::LayoutB; + using ElementC = typename BlockMmad::ElementC; + using LayoutC = typename BlockMmad::LayoutC; + using ElementAccumulator = typename BlockMmad::ElementAccumulator; + + using BlockEpilogue = BlockEpilogue_; + using ElementScale = typename BlockEpilogue::ElementScale; + using LayoutScale = typename BlockEpilogue::LayoutScale; + using ElementPerTokenScale = typename BlockEpilogue::ElementPerTokenScale; + using LayoutPerTokenScale = typename BlockEpilogue::LayoutPerTokenScale; + using ElementD = typename BlockEpilogue::ElementD; + using LayoutD = typename BlockEpilogue::LayoutD; + using EpilogueParams = typename BlockEpilogue::Params; + + using BlockScheduler = BlockScheduler_; + static constexpr uint32_t WORKSPACE_STAGES = WORKSPACE_STAGES_; + + /// Parameters structure + struct Params { + // Data members + GemmCoord problemShape; + __gm__ ElementA* ptrA; + LayoutA layoutA; + __gm__ ElementB* ptrB; + LayoutB layoutB; + __gm__ ElementScale* ptrScale; + LayoutScale layoutScale; + __gm__ ElementPerTokenScale* ptrPerTokenScale; + LayoutPerTokenScale layoutPerTokenScale; + __gm__ int32_t* ptrBias; + __gm__ ElementD* ptrD; + LayoutD layoutD; + GM_ADDR ptrWorkspace; + + // Methods + CATLASS_DEVICE + Params() {} + + CATLASS_DEVICE + Params(GemmCoord problemShape_, + GM_ADDR ptrA_, + LayoutA layoutA_, + GM_ADDR ptrB_, + LayoutB layoutB_, + GM_ADDR ptrScale_, + LayoutScale layoutScale_, + GM_ADDR ptrPerTokenScale_, + LayoutPerTokenScale layoutPerTokenScale_, + GM_ADDR ptrBias_, + GM_ADDR ptrD_, + LayoutD layoutD_, + GM_ADDR ptrWorkspace_) + : problemShape(problemShape_), + ptrA(reinterpret_cast<__gm__ ElementA*>(ptrA_)), + layoutA(layoutA_), + ptrB(reinterpret_cast<__gm__ ElementB*>(ptrB_)), + layoutB(layoutB_), + ptrScale(reinterpret_cast<__gm__ ElementScale*>(ptrScale_)), + layoutScale(layoutScale_), + ptrPerTokenScale(reinterpret_cast<__gm__ ElementPerTokenScale*>( + ptrPerTokenScale_)), + layoutPerTokenScale(layoutPerTokenScale_), + ptrBias(reinterpret_cast<__gm__ int32_t*>(ptrBias_)), + ptrD(reinterpret_cast<__gm__ ElementD*>(ptrD_)), + layoutD(layoutD_), + ptrWorkspace(ptrWorkspace_) {} + }; + + // Methods + CATLASS_DEVICE + QuantMatmulNzWorkspace() { + Arch::FlagID flagId = 0; + for (uint32_t stageId = 0; stageId < WORKSPACE_STAGES; ++stageId) { + flagAicFinishStoreList[stageId] = Arch::CrossCoreFlag(flagId++); + flagAivFinishComputeList[stageId] = Arch::CrossCoreFlag(flagId++); + aicWaitFuncList[stageId] = {this, stageId}; + aicSetFuncList[stageId] = {this, stageId}; + } + } + + template + CATLASS_DEVICE void operator()(Params const& params); + + template <> + CATLASS_DEVICE void operator()(Params const& params) { + BlockScheduler blockScheduler; + blockScheduler.Update(params.problemShape, + MakeCoord(L1TileShape::M, L1TileShape::N)); + uint32_t coreLoops = blockScheduler.GetCoreLoops(); + + BlockMmad blockMmad(resource); + + // Represent the full gm + AscendC::GlobalTensor gmA; + gmA.SetGlobalBuffer(params.ptrA); + AscendC::GlobalTensor gmB; + gmB.SetGlobalBuffer(params.ptrB); + if (params.problemShape.n() != 1280) { + gmB.SetL2CacheHint(AscendC::CacheMode::CACHE_MODE_DISABLE); + } + + uint32_t coreIdx = AscendC::GetBlockIdx(); + uint32_t coreNum = AscendC::GetBlockNum(); + + AscendC::GlobalTensor gmC; + gmC.SetGlobalBuffer( + reinterpret_cast<__gm__ ElementC*>(params.ptrWorkspace)); + auto layoutC = layout::RowMajor{L1TileShape::M * coreNum * WORKSPACE_STAGES, + L1TileShape::N}; + + uint32_t stageId = 0; + uint32_t stageUsed = 0; + + // Loop through the matmul of each groupIdx + for (uint32_t loopIdx = coreIdx; loopIdx < coreLoops; loopIdx += coreNum) { + // Compute block location + GemmCoord blockCoord = blockScheduler.GetBlockCoord(loopIdx); + GemmCoord actualBlockShape = + blockScheduler.GetActualBlockShape(blockCoord); + + Callback callbackBeforeFixpipe{}; + if (stageUsed == WORKSPACE_STAGES) { + callbackBeforeFixpipe = MakeCallback(&aicWaitFuncList[stageId]); + } else { + ++stageUsed; + } + Callback callbackAfterFixpipe = MakeCallback(&aicSetFuncList[stageId]); + + // Compute initial location in logical coordinates + MatrixCoord offsetA{blockCoord.m() * L1TileShape::M, + blockCoord.k() * L1TileShape::K}; + MatrixCoord offsetB{blockCoord.k() * L1TileShape::K, + blockCoord.n() * L1TileShape::N}; + MatrixCoord offsetC{(stageId * coreNum + coreIdx) * L1TileShape::M, 0}; + int64_t gmOffsetA = params.layoutA.GetOffset(offsetA); + int64_t gmOffsetB = params.layoutB.GetOffset(offsetB); + int64_t gmOffsetC = layoutC.GetOffset(offsetC); + + // Compute block-scoped matrix multiply-add + if constexpr (BlockMmad::DispatchPolicy::ASYNC) { + blockMmad(gmA[gmOffsetA], + params.layoutA, + gmB[gmOffsetB], + params.layoutB, + gmC[gmOffsetC], + layoutC, + actualBlockShape, + callbackBeforeFixpipe, + callbackAfterFixpipe); + } else { + callbackBeforeFixpipe(); + blockMmad(gmA[gmOffsetA], + params.layoutA, + gmB[gmOffsetB], + params.layoutB, + gmC[gmOffsetC], + layoutC, + actualBlockShape); + callbackAfterFixpipe(); + } + + stageId = (stageId + 1 < WORKSPACE_STAGES) ? (stageId + 1) : 0; + } + + if constexpr (BlockMmad::DispatchPolicy::ASYNC) { + blockMmad.SynchronizeBlock(); + } + + while (stageUsed > 0) { + uint32_t aivComputeStageId = + (stageId >= stageUsed) ? (stageId - stageUsed) + : (stageId + WORKSPACE_STAGES - stageUsed); + Arch::CrossCoreWaitFlag(flagAivFinishComputeList[aivComputeStageId]); + --stageUsed; + } + } + + template <> + CATLASS_DEVICE void operator()(Params const& params) { + BlockScheduler blockScheduler; + uint32_t coreIdx = AscendC::GetBlockIdx() / AscendC::GetSubBlockNum(); + uint32_t coreNum = AscendC::GetBlockNum(); + + AscendC::GlobalTensor gmC; + gmC.SetGlobalBuffer( + reinterpret_cast<__gm__ ElementC*>(params.ptrWorkspace)); + auto layoutC = layout::RowMajor{L1TileShape::M * coreNum * WORKSPACE_STAGES, + L1TileShape::N}; + + uint32_t stageId = 0; + + blockScheduler.Update(params.problemShape, L1TileShape::ToCoordMN()); + uint32_t coreLoops = blockScheduler.GetCoreLoops(); + + GemmCoord blockShapeMNK = L1TileShape::ToCoord(); + for (uint32_t loopIdx = coreIdx; loopIdx < coreLoops; loopIdx += coreNum) { + GemmCoord blockCoordMNK = blockScheduler.GetBlockCoord(loopIdx); + GemmCoord actualBlockShapeMNK = + blockScheduler.GetActualBlockShape(blockCoordMNK); + + MatrixCoord offsetC{(stageId * coreNum + coreIdx) * L1TileShape::M, 0}; + int64_t gmOffsetC = layoutC.GetOffset(offsetC); + auto gmBlockC = gmC[gmOffsetC]; + + Arch::CrossCoreWaitFlag(flagAicFinishStoreList[stageId]); + static_epilogue(params, blockCoordMNK, actualBlockShapeMNK, gmBlockC); + Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(flagAivFinishComputeList[stageId]); + + stageId = (stageId + 1 < WORKSPACE_STAGES) ? (stageId + 1) : 0; + } + } + + private: + CATLASS_DEVICE + void static_epilogue(Params const& params, + GemmCoord const& block_coord, + GemmCoord const& actual_block_shape, + AscendC::GlobalTensor const& gm_block_c) { + constexpr bool kSupportsTwoDimensionalEpilogue = + L1TileShape::N == 128 || L1TileShape::N == 160 || L1TileShape::N == 256; + constexpr uint32_t kMinTwoDimensionalRows = L1TileShape::N == 256 ? 8 : 2; + if constexpr (kSupportsTwoDimensionalEpilogue) { + if (actual_block_shape.m() >= kMinTwoDimensionalRows && + actual_block_shape.n() == L1TileShape::N) { + if (AscendC::GetSubBlockIdx() == 0) { + static_epilogue_two_dimensional( + params, block_coord, actual_block_shape, gm_block_c); + } + return; + } + } + static_epilogue_row(params, block_coord, actual_block_shape, gm_block_c); + } + + CATLASS_DEVICE + void static_epilogue_row(Params const& params, + GemmCoord const& blockCoord, + GemmCoord const& actualBlockShape, + AscendC::GlobalTensor const& gmBlockC) { + constexpr uint32_t kEpilogueTileN = + L1TileShape::N == 320 ? 320 : 256; + uint32_t subblockIdx = AscendC::GetSubBlockIdx(); + uint32_t tileNOffset = subblockIdx * kEpilogueTileN; + if (tileNOffset >= actualBlockShape.n()) { + return; + } + uint32_t tileN = actualBlockShape.n() - tileNOffset; + tileN = tileN < kEpilogueTileN ? tileN : kEpilogueTileN; + uint32_t blockNOffset = blockCoord.n() * L1TileShape::N; + + size_t ubOffset = 0; + auto ubC = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += kEpilogueTileN * sizeof(int32_t); + auto ubBias = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += kEpilogueTileN * sizeof(int32_t); + auto ubScale = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += kEpilogueTileN * sizeof(float); + auto ubFloat = resource.ubBuf.template GetBufferByByte(ubOffset); + ubOffset += kEpilogueTileN * sizeof(float); + auto ubD = resource.ubBuf.template GetBufferByByte(ubOffset); + + AscendC::GlobalTensor gmBias; + gmBias.SetGlobalBuffer(params.ptrBias); + AscendC::GlobalTensor gmScale; + gmScale.SetGlobalBuffer(params.ptrScale); + AscendC::GlobalTensor gmD; + gmD.SetGlobalBuffer(params.ptrD); + + AscendC::DataCopy(ubBias, gmBias[blockNOffset + tileNOffset], tileN); + AscendC::DataCopy(ubScale, gmScale[blockNOffset + tileNOffset], tileN); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID0); + AscendC::SetFlag(EVENT_ID0); + + for (uint32_t row = 0; row < actualBlockShape.m(); ++row) { + AscendC::WaitFlag(EVENT_ID0); + AscendC::DataCopy( + ubC, gmBlockC[row * L1TileShape::N + tileNOffset], tileN); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::Add(ubC, ubC, ubBias, tileN); + AscendC::PipeBarrier(); + AscendC::Cast(ubFloat, ubC, AscendC::RoundMode::CAST_RINT, tileN); + AscendC::SetFlag(EVENT_ID0); + AscendC::PipeBarrier(); + AscendC::Mul(ubFloat, ubFloat, ubScale, tileN); + AscendC::PipeBarrier(); + AscendC::WaitFlag(EVENT_ID0); + AscendC::Cast(ubD, ubFloat, AscendC::RoundMode::CAST_RINT, tileN); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + AscendC::DataCopy( + gmD[row * params.problemShape.n() + blockNOffset + tileNOffset], + ubD, + tileN); + AscendC::SetFlag(EVENT_ID0); + } + AscendC::WaitFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + + CATLASS_DEVICE + void static_epilogue_two_dimensional( + Params const& params, + GemmCoord const& block_coord, + GemmCoord const& actual_block_shape, + AscendC::GlobalTensor const& gm_block_c) { + constexpr uint32_t kEpilogueTileN = 256; + constexpr uint32_t kMaxRows = 16; + constexpr uint32_t kMaxElements = kMaxRows * kEpilogueTileN; + constexpr uint32_t kDataBlockBytes = 32; + constexpr uint32_t kVectorRepeatElements = 64; + + const uint32_t row_count = actual_block_shape.m(); + const uint32_t tile_n = actual_block_shape.n(); + const uint32_t element_count = row_count * tile_n; + const uint32_t block_n_offset = block_coord.n() * L1TileShape::N; + + size_t ub_offset = 0; + auto ub_c = resource.ubBuf.template GetBufferByByte(ub_offset); + ub_offset += kMaxElements * sizeof(int32_t); + auto ub_bias = resource.ubBuf.template GetBufferByByte(ub_offset); + ub_offset += kEpilogueTileN * sizeof(int32_t); + auto ub_scale = resource.ubBuf.template GetBufferByByte(ub_offset); + ub_offset += kEpilogueTileN * sizeof(float); + auto ub_float = resource.ubBuf.template GetBufferByByte(ub_offset); + ub_offset += kMaxElements * sizeof(float); + auto ub_d = resource.ubBuf.template GetBufferByByte(ub_offset); + + AscendC::GlobalTensor gm_bias; + gm_bias.SetGlobalBuffer(params.ptrBias); + AscendC::GlobalTensor gm_scale; + gm_scale.SetGlobalBuffer(params.ptrScale); + AscendC::GlobalTensor gm_d; + gm_d.SetGlobalBuffer(params.ptrD); + + const uint16_t accumulator_burst_length = + static_cast(tile_n * sizeof(int32_t) / kDataBlockBytes); + const uint16_t accumulator_source_gap = static_cast( + (L1TileShape::N - tile_n) * sizeof(int32_t) / kDataBlockBytes); + AscendC::DataCopy(ub_bias, gm_bias[block_n_offset], tile_n); + AscendC::DataCopy(ub_scale, gm_scale[block_n_offset], tile_n); + AscendC::DataCopy(ub_c, + gm_block_c, + AscendC::DataCopyParams{static_cast(row_count), + accumulator_burst_length, + accumulator_source_gap, + 0}); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + + const uint8_t row_stride = static_cast(accumulator_burst_length); + const AscendC::BinaryRepeatParams broadcast_params{ + 1, 1, 1, row_stride, row_stride, 0}; + for (uint32_t column_offset = 0; column_offset < tile_n; + column_offset += kVectorRepeatElements) { + const uint32_t vector_elements = + tile_n - column_offset < kVectorRepeatElements + ? tile_n - column_offset + : kVectorRepeatElements; + AscendC::Add(ub_c[column_offset], + ub_c[column_offset], + ub_bias[column_offset], + vector_elements, + static_cast(row_count), + broadcast_params); + } + AscendC::PipeBarrier(); + AscendC::Cast(ub_float, ub_c, AscendC::RoundMode::CAST_RINT, element_count); + AscendC::PipeBarrier(); + for (uint32_t column_offset = 0; column_offset < tile_n; + column_offset += kVectorRepeatElements) { + const uint32_t vector_elements = + tile_n - column_offset < kVectorRepeatElements + ? tile_n - column_offset + : kVectorRepeatElements; + AscendC::Mul(ub_float[column_offset], + ub_float[column_offset], + ub_scale[column_offset], + vector_elements, + static_cast(row_count), + broadcast_params); + } + AscendC::PipeBarrier(); + AscendC::Cast(ub_d, ub_float, AscendC::RoundMode::CAST_RINT, element_count); + + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + const uint16_t output_burst_length = + static_cast(tile_n * sizeof(ElementD) / kDataBlockBytes); + const uint16_t output_destination_gap = + static_cast((params.problemShape.n() - tile_n) * + sizeof(ElementD) / kDataBlockBytes); + AscendC::DataCopy(gm_d[block_n_offset], + ub_d, + AscendC::DataCopyParams{static_cast(row_count), + output_burst_length, + 0, + output_destination_gap}); + AscendC::SetFlag(EVENT_ID0); + AscendC::WaitFlag(EVENT_ID0); + } + + friend struct AicWaitFunc; + friend struct AicSetFunc; + + struct AicWaitFunc { + using MatmulKernel = QuantMatmulNzWorkspace; + + CATLASS_DEVICE + AicWaitFunc() = default; + + CATLASS_DEVICE + void operator()() const { + Arch::CrossCoreWaitFlag(ptr->flagAivFinishComputeList[stageId]); + } + + MatmulKernel* ptr{nullptr}; + uint32_t stageId; + }; + + struct AicSetFunc { + using MatmulKernel = QuantMatmulNzWorkspace; + + CATLASS_DEVICE + AicSetFunc() = default; + + CATLASS_DEVICE + void operator()() const { + Arch::CrossCoreSetFlag<0x2, PIPE_FIX>( + ptr->flagAicFinishStoreList[stageId]); + } + + MatmulKernel* ptr{nullptr}; + uint32_t stageId; + }; + + Arch::CrossCoreFlag flagAicFinishStoreList[WORKSPACE_STAGES]; + Arch::CrossCoreFlag flagAivFinishComputeList[WORKSPACE_STAGES]; + + AicWaitFunc aicWaitFuncList[WORKSPACE_STAGES]; + AicSetFunc aicSetFuncList[WORKSPACE_STAGES]; + Arch::Resource resource; +}; + +} // namespace Catlass::Gemm::Kernel