Star137
243
代码
代码
Issues85
Pull Requests179
项目讨论
Wiki
分析
Star137
243
已合并
采样topk-topp用算子实现 #935
shenzhao创建于 4月22日
采样topk-topp用算子实现 #935
已合并
shenzhao创建于 4月22日
从shenzhao/MindIE-LLM_3169dev-0422合入到Ascend/MindIE-LLMdev
共 21 个文件变更+2223-20
Mmindie_llm/runtime/layers/sampling/sampler.py+25-2
@@ -14,6 +14,8 @@ from typing import Any, List
14 14 
15import numpy as np15import numpy as np
16import torch16import torch
17+import torch_npu as torch_npu
18+from mindie_llm.runtime.ops import mie_ops as mie_ops
17from mindie_llm.text_generator.samplers.logits_handlers.__init__ import get_handler_registry19from mindie_llm.text_generator.samplers.logits_handlers.__init__ import get_handler_registry
18from mindie_llm.text_generator.samplers.logits_handlers.logits_handler import LogitsHandlerList20from mindie_llm.text_generator.samplers.logits_handlers.logits_handler import LogitsHandlerList
19from mindie_llm.text_generator.samplers.sampler_params import HandlerParams, SelectorParams21from mindie_llm.text_generator.samplers.sampler_params import HandlerParams, SelectorParams
@@ -21,6 +23,7 @@ from mindie_llm.text_generator.samplers.token_selectors import get_selector_regi
21from mindie_llm.text_generator.utils.config import SamplerConfig, HandlingBackend23from mindie_llm.text_generator.utils.config import SamplerConfig, HandlingBackend
22from mindie_llm.text_generator.utils.sampling_output import SamplingOutput24from mindie_llm.text_generator.utils.sampling_output import SamplingOutput
23from mindie_llm.text_generator.utils.sampling_metadata import SamplingMetadata25from mindie_llm.text_generator.utils.sampling_metadata import SamplingMetadata
26+from mindie_llm.utils.log.logging import logger
24 27 
25 28 
26class SelectorType(IntEnum):29class SelectorType(IntEnum):
@@ -303,10 +306,30 @@ class Sampler:
303 306 
304 @staticmethod307 @staticmethod
305 def apply_top_k_top_p(logits: torch.Tensor, metadata: SamplingMetadata):308 def apply_top_k_top_p(logits: torch.Tensor, metadata: SamplingMetadata):
309+ has_top_k = metadata.top_k_idx is not None
310+ has_top_p = metadata.top_p is not None
311+ if has_top_k or has_top_p:
312+ p_tensor = None
313+ k_tensor = None
314+ if has_top_p:
315+ p_tensor = metadata.top_p.squeeze(-1).to(dtype=logits.dtype)
316+ else:
317+ p_tensor = torch.zeros(logits.shape[0], dtype=logits.dtype, device=logits.device)
318+ if has_top_k:
319+ k_tensor = (metadata.top_k_idx.squeeze(-1) + 1).to(torch.int32)
320+ if metadata.top_k_disabled_mask is not None:
321+ disabled = metadata.top_k_disabled_mask.squeeze(-1).to(torch.bool)
322+ k_tensor.masked_fill_(disabled, logits.size(-1))
323+ else:
324+ k_tensor = torch.full((logits.shape[0],), logits.size(-1), dtype=torch.int32, device=logits.device)
325+ try:
326+ return torch.ops.mie_ops.apply_top_k_top_p_custom(logits, p_tensor, k_tensor)
327+ except Exception as e:
328+ logger.warning("apply_top_k_top_p_custom failed, fallback to pytorch impl, error: %s", e)
306 probs = logits.softmax(dim=-1)329 probs = logits.softmax(dim=-1)
307 probs_sort, _ = probs.sort(dim=-1, descending=False)330 probs_sort, _ = probs.sort(dim=-1, descending=False)
308 masked_value = -float("inf")331 masked_value = -float("inf")
309- if metadata.top_k_idx is not None:332+ if has_top_k:
310 top_k_indices = probs_sort.size(1) - 1 - metadata.top_k_idx333 top_k_indices = probs_sort.size(1) - 1 - metadata.top_k_idx
311 if top_k_indices.dim() == 1:334 if top_k_indices.dim() == 1:
312 top_k_indices = top_k_indices.unsqueeze(-1)335 top_k_indices = top_k_indices.unsqueeze(-1)
@@ -318,7 +341,7 @@ class Sampler:
318 341 
319 indices_to_remove = probs < kth_probs342 indices_to_remove = probs < kth_probs
320 logits.masked_fill_(indices_to_remove, masked_value)343 logits.masked_fill_(indices_to_remove, masked_value)
321- if metadata.top_p is not None:344+ if has_top_p:
322 cumprob = torch.cumsum(probs_sort, dim=-1)345 cumprob = torch.cumsum(probs_sort, dim=-1)
323 top_p = 1 - metadata.top_p346 top_p = 1 - metadata.top_p
324 if top_p.dim() == 1:347 if top_p.dim() == 1:
Msrc/kernels/mie_ops/csrc/CMakeLists.txt+2-1
@@ -51,7 +51,7 @@ foreach(SOC_VERSION ${ASCEND_COMPUTE_UNIT})
51 list(GET ARCH_DIRECTORY_LIST ${INDEX} VAL)51 list(GET ARCH_DIRECTORY_LIST ${INDEX} VAL)
52 list(APPEND ARCH_DIRECTORY ${VAL})52 list(APPEND ARCH_DIRECTORY ${VAL})
53 else()53 else()
54- message(STATUS "unsupport chip type")54+ message(STATUS "unsupported chip type")
55 if ((NOT BUILD_OPS_RTY_KERNEL) AND (BUILD_OPEN_PROJECT))55 if ((NOT BUILD_OPS_RTY_KERNEL) AND (BUILD_OPEN_PROJECT))
56 include(cmake/build_empty_package.cmake)56 include(cmake/build_empty_package.cmake)
57 cpack_empty_package()57 cpack_empty_package()
@@ -261,6 +261,7 @@ add_subdirectory(common)
261if (NOT BUILD_OPS_RTY_KERNEL)261if (NOT BUILD_OPS_RTY_KERNEL)
262 add_subdirectory(mc2)262 add_subdirectory(mc2)
263 add_subdirectory(posembedding)263 add_subdirectory(posembedding)
264+ add_subdirectory(sampler)
264endif()265endif()
265 266 
266if(NOT DEFINED COMPILED_OPS)267if(NOT DEFINED COMPILED_OPS)
Msrc/kernels/mie_ops/csrc/build_aclnn.sh+2-0
@@ -12,6 +12,7 @@ if [[ "$SOC_VERSION" =~ ^ascend910b ]]; then
12 12 
13 CUSTOM_OPS_ARRAY=(13 CUSTOM_OPS_ARRAY=(
14 "lightning_indexer"14 "lightning_indexer"
15+ "apply_top_k_top_p_custom"
15 )16 )
16 17 
17 CUSTOM_OPS=$(IFS=';'; echo "${CUSTOM_OPS_ARRAY[*]}")18 CUSTOM_OPS=$(IFS=';'; echo "${CUSTOM_OPS_ARRAY[*]}")
@@ -24,6 +25,7 @@ elif [[ "$SOC_VERSION" =~ ^ascend910_93 ]]; then
24 CUSTOM_OPS_ARRAY=(25 CUSTOM_OPS_ARRAY=(
25 "dispatch_ffn_combine"26 "dispatch_ffn_combine"
26 "dispatch_gmm_combine_decode"27 "dispatch_gmm_combine_decode"
28+ "apply_top_k_top_p_custom"
27 )29 )
28 CUSTOM_OPS=$(IFS=';'; echo "${CUSTOM_OPS_ARRAY[*]}")30 CUSTOM_OPS=$(IFS=';'; echo "${CUSTOM_OPS_ARRAY[*]}")
29 SOC_ARG="ascend910_93"31 SOC_ARG="ascend910_93"
Msrc/kernels/mie_ops/csrc/cmake/func.cmake+3-2
@@ -86,6 +86,7 @@ function(op_add_subdirectory OP_LIST OP_DIR_LIST)
86 "${CMAKE_CURRENT_SOURCE_DIR}/moe/**/op_host/CMakeLists.txt"86 "${CMAKE_CURRENT_SOURCE_DIR}/moe/**/op_host/CMakeLists.txt"
87 "${CMAKE_CURRENT_SOURCE_DIR}/ffn/**/op_host/CMakeLists.txt"87 "${CMAKE_CURRENT_SOURCE_DIR}/ffn/**/op_host/CMakeLists.txt"
88 "${CMAKE_CURRENT_SOURCE_DIR}/mc2/**/op_host/CMakeLists.txt"88 "${CMAKE_CURRENT_SOURCE_DIR}/mc2/**/op_host/CMakeLists.txt"
89+ "${CMAKE_CURRENT_SOURCE_DIR}/sampler/**/op_host/CMakeLists.txt"
89 "${CMAKE_CURRENT_SOURCE_DIR}/posembedding/**/framework/CMakeLists.txt"90 "${CMAKE_CURRENT_SOURCE_DIR}/posembedding/**/framework/CMakeLists.txt"
90 "${CMAKE_CURRENT_SOURCE_DIR}/moe/**/framework/CMakeLists.txt"91 "${CMAKE_CURRENT_SOURCE_DIR}/moe/**/framework/CMakeLists.txt"
91 "${CMAKE_CURRENT_SOURCE_DIR}/ffn/**/framework/CMakeLists.txt"92 "${CMAKE_CURRENT_SOURCE_DIR}/ffn/**/framework/CMakeLists.txt"
@@ -815,9 +816,9 @@ function(add_aicpu_cust_kernel_modules op_name aicpu_sources aicpu_jsons)
815 -Wl,--no-whole-archive816 -Wl,--no-whole-archive
816 )817 )
817 if (NOT (UT_TEST_ALL OR OP_KERNEL_AICPU_UT))818 if (NOT (UT_TEST_ALL OR OP_KERNEL_AICPU_UT))
818- set_property(TARGET ${target_name} PROPERTY 819+ set_property(TARGET ${target_name} PROPERTY
819 CXX_COMPILER_LAUNCHER ${ASCEND_DIR}/toolkit/toolchain/hcc/bin/aarch64-target-linux-gnu-g++)820 CXX_COMPILER_LAUNCHER ${ASCEND_DIR}/toolkit/toolchain/hcc/bin/aarch64-target-linux-gnu-g++)
820- endif() 821+ endif()
821 target_sources(${target_name} PRIVATE ${aicpu_sources})822 target_sources(${target_name} PRIVATE ${aicpu_sources})
822 set_property(GLOBAL APPEND PROPERTY AICPU_JSON_FILES ${aicpu_jsons})823 set_property(GLOBAL APPEND PROPERTY AICPU_JSON_FILES ${aicpu_jsons})
823 if (NOT ${target_name} IN_LIST AICPU_CUST_OBJ_TARGETS)824 if (NOT ${target_name} IN_LIST AICPU_CUST_OBJ_TARGETS)
Asrc/kernels/mie_ops/csrc/sampler/CMakeLists.txt+6-0
@@ -0,0 +1,6 @@
1+file(GLOB SUBDIRECTORIES LIST_DIRECTORIES true RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
2+foreach(SUBDIR ${SUBDIRECTORIES})
3+ if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/${SUBDIR}/CMakeLists.txt)
4+ add_subdirectory(${SUBDIR})
5+ endif()
6+endforeach()
Asrc/kernels/mie_ops/csrc/sampler/apply_top_k_top_p_custom/CMakeLists.txt+18-0
@@ -0,0 +1,18 @@
1+file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
2+foreach(SUB_DIR ${CURRENT_DIRS})
3+ if (DEFINED ASCEND_OP_NAME AND NOT "${ASCEND_OP_NAME}" STREQUAL "")
4+ if (NOT "${ASCEND_OP_NAME}" STREQUAL "all" AND NOT "${ASCEND_OP_NAME}" STREQUAL "ALL")
5+ if (NOT ${SUB_DIR} IN_LIST ASCEND_OP_NAME)
6+ continue()
7+ endif ()
8+ endif ()
9+ endif ()
10+ 
11+ if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
12+ add_subdirectory(${SUB_DIR})
13+ else()
14+ if (EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/op_host/CMakeLists.txt")
15+ add_subdirectory(${SUB_DIR}/op_host)
16+ endif()
17+ endif()
18+endforeach()
Asrc/kernels/mie_ops/csrc/sampler/apply_top_k_top_p_custom/op_host/CMakeLists.txt+32-0
@@ -0,0 +1,32 @@
1+add_op_to_compiled_list()
2+ 
3+if (BUILD_OPEN_PROJECT)
4+ target_sources(op_host_aclnnInner PRIVATE
5+ apply_top_k_top_p_custom_def.cpp
6+ )
7+ add_ops_compile_options(
8+ OP_NAME ApplyTopKTopPCustom
9+ OPTIONS --cce-auto-sync=on
10+ -Wno-deprecated-declarations
11+ -Werror
12+ )
13+endif()
14+ 
15+if(NOT BUILD_OPS_RTY_KERNEL)
16+ add_modules_sources(OPTYPE apply_top_k_top_p_custom ACLNNTYPE aclnn_inner)
17+ 
18+ if (TARGET ${OPHOST_NAME}_opapi_obj)
19+ target_sources(${OPHOST_NAME}_opapi_obj PRIVATE
20+ ${CMAKE_CURRENT_SOURCE_DIR}/apply_top_k_top_p_custom.cpp
21+ )
22+ target_include_directories(${OPHOST_NAME}_opapi_obj PRIVATE
23+ ${CMAKE_CURRENT_SOURCE_DIR}
24+ )
25+ endif()
26+ 
27+ if (TARGET ${OPHOST_NAME}_tiling_obj)
28+ target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE
29+ ${CMAKE_CURRENT_SOURCE_DIR}
30+ )
31+ endif()
32+endif()
Asrc/kernels/mie_ops/csrc/sampler/apply_top_k_top_p_custom/op_host/apply_top_k_top_p_custom.cpp+45-0
@@ -0,0 +1,45 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file apply_top_k_top_p_custom.cpp
13+ * \brief
14+ */
15+#include "apply_top_k_top_p_custom.h"
16+ 
17+#include "opdev/data_type_utils.h"
18+#include "opdev/format_utils.h"
19+#include "opdev/make_op_executor.h"
20+#include "opdev/op_def.h"
21+#include "opdev/op_dfx.h"
22+#include "opdev/op_executor.h"
23+#include "opdev/op_log.h"
24+#include "opdev/shape_utils.h"
25+using namespace op;
26+ 
27+namespace l0op {
28+ 
29+OP_TYPE_REGISTER(ApplyTopKTopPCustom);
30+ 
31+const aclTensor* ApplyTopKTopPCustom(const aclTensor* sortedValue, const aclTensor* sortedIndices, const aclTensor* p,
32+ const aclTensor* k, aclOpExecutor* executor) {
33+ L0_DFX(ApplyTopKTopPCustom, sortedValue, sortedIndices, p, k);
34+ auto output = executor->AllocTensor(sortedValue->GetViewShape(), sortedValue->GetDataType());
35+ if (p == nullptr) {
36+ p = executor->AllocTensor(sortedValue->GetDataType(), Format::FORMAT_ND, Format::FORMAT_ND);
37+ }
38+ if (k == nullptr) {
39+ k = executor->AllocTensor(DataType::DT_INT32, Format::FORMAT_ND, Format::FORMAT_ND);
40+ }
41+ ADD_TO_LAUNCHER_LIST_AICORE(ApplyTopKTopPCustom, OP_INPUT(sortedValue, sortedIndices, p, k), OP_OUTPUT(output));
42+ 
43+ return output;
44+}
45+} // namespace l0op
Asrc/kernels/mie_ops/csrc/sampler/apply_top_k_top_p_custom/op_host/apply_top_k_top_p_custom.h+24-0
@@ -0,0 +1,24 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file apply_top_k_top_p_custom.h
13+ * \brief
14+ */
15+#ifndef OP_API_INC_LEVEL0_OP_APPLY_TOP_K_TOP_P_CUSTOM_OP_H_
16+#define OP_API_INC_LEVEL0_OP_APPLY_TOP_K_TOP_P_CUSTOM_OP_H_
17+ 
18+#include "opdev/op_executor.h"
19+ 
20+namespace l0op {
21+const aclTensor* ApplyTopKTopPCustom(const aclTensor* sortedValue, const aclTensor* sortedIndices, const aclTensor* p,
22+ const aclTensor* k, aclOpExecutor* executor);
23+}
24+#endif // OP_API_INC_LEVEL0_OP_APPLY_TOP_K_TOP_P_CUSTOM_OP_H_
Asrc/kernels/mie_ops/csrc/sampler/apply_top_k_top_p_custom/op_host/apply_top_k_top_p_custom_def.cpp+98-0
@@ -0,0 +1,98 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file apply_top_k_top_p_custom_def.cpp
13+ * \brief
14+ */
15+#include "register/op_def_registry.h"
16+ 
17+namespace ops {
18+class ApplyTopKTopPCustom : public OpDef {
19+ public:
20+ explicit ApplyTopKTopPCustom(const char *name) : OpDef(name) {
21+ this->Input("sorted_value")
22+ .ParamType(REQUIRED)
23+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16})
24+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
25+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
26+ this->Input("sorted_indices")
27+ .ParamType(REQUIRED)
28+ .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
29+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
30+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
31+ this->Input("p")
32+ .ParamType(OPTIONAL)
33+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16})
34+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
35+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
36+ this->Input("k")
37+ .ParamType(OPTIONAL)
38+ .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
39+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
40+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
41+ this->Output("out")
42+ .ParamType(REQUIRED)
43+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16})
44+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
45+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
46+ 
47+ OpAICoreConfig aicore_config;
48+ aicore_config.DynamicCompileStaticFlag(true)
49+ .DynamicFormatFlag(false)
50+ .DynamicRankSupportFlag(true)
51+ .DynamicShapeSupportFlag(true);
52+ this->AICore().AddConfig("ascend910b");
53+ this->AICore().AddConfig("ascend910_93");
54+ 
55+ OpAICoreConfig config_kirin = GetKirinCoreConfig();
56+ this->AICore().AddConfig("kirinx90", config_kirin);
57+ }
58+ 
59+ private:
60+ OpAICoreConfig GetKirinCoreConfig() const {
61+ OpAICoreConfig config_kirin;
62+ config_kirin.DynamicCompileStaticFlag(true)
63+ .DynamicFormatFlag(true)
64+ .DynamicRankSupportFlag(true)
65+ .DynamicShapeSupportFlag(true)
66+ .NeedCheckSupportFlag(false)
67+ .PrecisionReduceFlag(true);
68+ config_kirin.Input("sorted_value")
69+ .ParamType(REQUIRED)
70+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
71+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
72+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
73+ config_kirin.Input("sorted_indices")
74+ .ParamType(REQUIRED)
75+ .DataType({ge::DT_INT32, ge::DT_INT32})
76+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
77+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
78+ config_kirin.Input("p")
79+ .ParamType(OPTIONAL)
80+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
81+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
82+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
83+ config_kirin.Input("k")
84+ .ParamType(OPTIONAL)
85+ .DataType({ge::DT_INT32, ge::DT_INT32})
86+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
87+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
88+ config_kirin.Output("out")
89+ .ParamType(REQUIRED)
90+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
91+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
92+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
93+ return config_kirin;
94+ }
95+};
96+ 
97+OP_ADD(ApplyTopKTopPCustom);
98+} // namespace ops
Asrc/kernels/mie_ops/csrc/sampler/apply_top_k_top_p_custom/op_host/apply_top_k_top_p_custom_tiling.cpp+306-0
@@ -0,0 +1,306 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file apply_top_k_top_p_custom_tiling.cpp
13+ * \brief
14+ */
15+ 
16+#include "apply_top_k_top_p_custom_tiling.h"
17+ 
18+#include <iostream>
19+#include <map>
20+ 
21+#include "error_log.h"
22+#include "platform/platform_infos_def.h"
23+#include "register/op_def_registry.h"
24+#include "register/op_impl_registry.h"
25+#include "tiling/platform/platform_ascendc.h"
26+#include "tiling/tiling_api.h"
27+ 
28+namespace {
29+constexpr uint32_t SYS_RESERVED_UB = uint32_t(16 * 1024);
30+constexpr uint32_t SELECT_RESERVED_UB = uint32_t(8 * 1024);
31+constexpr uint32_t DIM_ONE = 1;
32+constexpr uint32_t DIM_TWO = 2;
33+constexpr int32_t SORTED_VALUE_INPUT_INDEX = 0;
34+constexpr int32_t SORTED_INDICES_INPUT_INDEX = 1;
35+constexpr int32_t P_INPUT_INDEX = 2;
36+constexpr int32_t K_INPUT_INDEX = 3;
37+constexpr uint32_t DIM_INDEX0 = 0;
38+constexpr uint32_t FLOAT_BYTES = 4;
39+static std::map<ge::DataType, uint32_t> DTYPE_MAP = {{ge::DT_BF16, 2}, {ge::DT_FLOAT16, 1}, {ge::DT_FLOAT, 0}};
40+static std::map<ge::DataType, uint32_t> DATATYPE_LEN_MAP = {{ge::DT_FLOAT16, 2}, {ge::DT_BF16, 2}, {ge::DT_FLOAT, 4}};
41+const static uint32_t SYS_WORKSPACESIZE = uint32_t(16 * 1024 * 1024);
42+ 
43+constexpr uint32_t DATA_PER_BLOCK_B32 = 8;
44+constexpr uint32_t BYTES_B32 = 4;
45+constexpr uint32_t BLOCK_BYTES = 32;
46+constexpr uint32_t K_VALUE_MAX = 1024;
47+constexpr uint32_t ONLY_TOP_P_KEY = 2;
48+constexpr uint32_t ONLY_TOP_K_KEY = 1;
49+constexpr uint32_t BATCH_MODE = 1;
50+} // namespace
51+ 
52+namespace optiling {
53+class ApplyTopKTopPCustomTiling {
54+ public:
55+ explicit ApplyTopKTopPCustomTiling(gert::TilingContext* context) : tilingcontext(context) {};
56+ ge::graphStatus Init();
57+ ge::graphStatus RunKernelTiling();
58+ 
59+ private:
60+ ApplyTopKTopPCustomTilingData tilingData;
61+ gert::TilingContext* tilingcontext = nullptr;
62+ ge::graphStatus CheckShape();
63+ void SetTilingKey();
64+ void GetUsedCore();
65+ void CalDataPerCore();
66+ void FillTilingData();
67+ void PrintTilingData();
68+ template <typename T1>
69+ inline auto CeilAlign(T1 a, T1 b) const -> T1 {
70+ return b == 0 ? a : (a + b - 1) / b * b;
71+ }
72+ template <typename T1>
73+ inline auto FloorAlign(T1 a, T1 b) const -> T1 {
74+ return b == 0 ? a : a / b * b;
75+ }
76+ 
77+ const char* opName_ = nullptr;
78+ uint32_t coreNum_ = 0;
79+ uint32_t calUbSize_ = 0;
80+ uint32_t batchSize_ = 0;
81+ uint32_t vocabSize_ = 0;
82+ uint32_t tilingKey_ = 0;
83+ uint32_t usedCoreNum_ = 0;
84+ uint32_t batchPerCore_ = 1;
85+ uint32_t tailBatch_ = 0;
86+ uint32_t dataNumInit_ = 0;
87+ uint32_t dataNumInitAligned_ = 0;
88+ uint32_t ubFactorElement_ = 0;
89+ uint32_t ubFactorElementAligned_ = 0;
90+ uint32_t tailUbFactorElement_ = 0;
91+ uint32_t tailUbFactorElementAligned_ = 0;
92+ uint32_t iterateTimes_ = 0;
93+ uint32_t onlyTopK_ = 0;
94+ uint32_t onlyTopP_ = 0;
95+ uint64_t platformUbSize_ = 0;
96+};
97+ 
98+ge::graphStatus ApplyTopKTopPCustomTiling::CheckShape() {
99+ auto sortedValueShapePtr = tilingcontext->GetInputShape(SORTED_VALUE_INPUT_INDEX);
100+ OP_CHECK_NULL_WITH_CONTEXT(tilingcontext, sortedValueShapePtr);
101+ auto sortedValueShape = sortedValueShapePtr->GetStorageShape();
102+ if (sortedValueShape.GetDimNum() != DIM_TWO) {
103+ OP_LOGE(opName_, "the dimNum of sorted_value should be 2, but got %u.", sortedValueShape.GetDimNum());
104+ return ge::GRAPH_FAILED;
105+ }
106+ auto sortedIndicesShapePtr = tilingcontext->GetInputShape(SORTED_INDICES_INPUT_INDEX);
107+ OP_CHECK_NULL_WITH_CONTEXT(tilingcontext, sortedIndicesShapePtr);
108+ auto sortedIndicesShape = sortedIndicesShapePtr->GetStorageShape();
109+ if (sortedIndicesShape.GetDimNum() != DIM_TWO) {
110+ OP_LOGE(opName_, "the dimNum of sorted_indices should be 2, but got %u.", sortedIndicesShape.GetDimNum());
111+ return ge::GRAPH_FAILED;
112+ }
113+ batchSize_ = sortedValueShape.GetDim(DIM_INDEX0);
114+ vocabSize_ = sortedValueShape.GetDim(DIM_ONE);
115+ if (sortedIndicesShape.GetDim(DIM_INDEX0) != batchSize_ || sortedIndicesShape.GetDim(DIM_ONE) != vocabSize_) {
116+ OP_LOGE(opName_, "the shape of sorted_indices should be equal to sorted_value.");
117+ return ge::GRAPH_FAILED;
118+ }
119+ 
120+ auto pShapePtr = tilingcontext->GetOptionalInputShape(P_INPUT_INDEX);
121+ OP_CHECK_NULL_WITH_CONTEXT(tilingcontext, pShapePtr);
122+ auto pShape = pShapePtr->GetStorageShape();
123+ auto pDimNum = pShape.GetDimNum();
124+ if (pDimNum != DIM_ONE && pDimNum != 0) {
125+ OP_LOGE(opName_, "the dimNum of p should be 1 or 0, but got %u.", pDimNum);
126+ return ge::GRAPH_FAILED;
127+ }
128+ if (pDimNum != 0 && batchSize_ != pShape.GetDim(DIM_INDEX0)) {
129+ OP_LOGE(opName_, "p.shape[0] should be equal to logits.shape[0].");
130+ return ge::GRAPH_FAILED;
131+ }
132+ 
133+ auto kShapePtr = tilingcontext->GetOptionalInputShape(K_INPUT_INDEX);
134+ OP_CHECK_NULL_WITH_CONTEXT(tilingcontext, kShapePtr);
135+ auto kShape = kShapePtr->GetStorageShape();
136+ auto kDimNum = kShape.GetDimNum();
137+ if (kDimNum != DIM_ONE && kDimNum != 0) {
138+ OP_LOGE(opName_, "the dimNum of k should be 1 or 0, but got %u.", kShape.GetDimNum());
139+ return ge::GRAPH_FAILED;
140+ }
141+ if (kDimNum != 0 && batchSize_ != kShape.GetDim(DIM_INDEX0)) {
142+ OP_LOGE(opName_, "k.shape[0] should be equal to logits.shape[0].");
143+ return ge::GRAPH_FAILED;
144+ }
145+ if (kDimNum == 0 && pDimNum == 0) {
146+ OP_LOGE(opName_, "the dimNum of q and k should be 0 at the same time.");
147+ return ge::GRAPH_FAILED;
148+ }
149+ onlyTopK_ = (kDimNum != 0 && pDimNum == 0) ? ONLY_TOP_K_KEY : 0;
150+ onlyTopP_ = (pDimNum != 0 && kDimNum == 0) ? ONLY_TOP_P_KEY : 0;
151+ return ge::GRAPH_SUCCESS;
152+}
153+ 
154+ge::graphStatus ApplyTopKTopPCustomTiling::Init() {
155+ opName_ = tilingcontext->GetNodeName();
156+ OP_LOGD(opName_, "TilingForApplyTopKTopPCustom init.");
157+ auto platformInfo = platform_ascendc::PlatformAscendC(tilingcontext->GetPlatformInfo());
158+ coreNum_ = platformInfo.GetCoreNumAiv();
159+ platformInfo.GetCoreMemSize(platform_ascendc::CoreMemType::UB, platformUbSize_);
160+ OP_LOGD(opName_, "platformUbSize: %lu.", platformUbSize_);
161+ uint32_t avaliableUb = static_cast<uint32_t>(platformUbSize_) - SYS_RESERVED_UB - SELECT_RESERVED_UB;
162+ calUbSize_ = FloorAlign(avaliableUb, BLOCK_BYTES);
163+ if (CheckShape() == ge::GRAPH_FAILED) {
164+ OP_LOGE(opName_, "check shape failed.");
165+ return ge::GRAPH_FAILED;
166+ }
167+ uint32_t tempValue = 1;
168+ while (tempValue < vocabSize_) {
169+ tempValue <<= 1;
170+ iterateTimes_++;
171+ } // ceil(log2(vocabSize_))
172+ return ge::GRAPH_SUCCESS;
173+}
174+ 
175+void ApplyTopKTopPCustomTiling::SetTilingKey() {
176+ tilingKey_ += onlyTopK_;
177+ tilingKey_ += onlyTopP_;
178+ tilingcontext->SetTilingKey(tilingKey_);
179+ if (tilingKey_ == ONLY_TOP_P_KEY) {
180+ tilingcontext->SetScheduleMode(BATCH_MODE);
181+ }
182+}
183+ 
184+void ApplyTopKTopPCustomTiling::GetUsedCore() {
185+ if (coreNum_ > 0) {
186+ batchPerCore_ = coreNum_ == uint32_t(0) ? batchSize_ : batchSize_ / coreNum_;
187+ tailBatch_ = batchSize_ % coreNum_;
188+ usedCoreNum_ = coreNum_;
189+ }
190+}
191+ 
192+void ApplyTopKTopPCustomTiling::CalDataPerCore() {
193+ uint32_t inputDataTypeByte = DATATYPE_LEN_MAP[tilingcontext->GetInputDesc(SORTED_VALUE_INPUT_INDEX)->GetDataType()];
194+ uint32_t dataPerBlock = BLOCK_BYTES / inputDataTypeByte;
195+ dataNumInit_ = vocabSize_ < K_VALUE_MAX ? vocabSize_ : K_VALUE_MAX;
196+ dataNumInitAligned_ = vocabSize_ < K_VALUE_MAX ? vocabSize_ : K_VALUE_MAX;
197+ ubFactorElement_ = vocabSize_ < K_VALUE_MAX ? vocabSize_ : K_VALUE_MAX;
198+ ubFactorElementAligned_ = CeilAlign(ubFactorElement_, dataPerBlock);
199+ tailUbFactorElement_ = vocabSize_ % ubFactorElement_;
200+ tailUbFactorElement_ = tailUbFactorElement_ == uint32_t(0) ? ubFactorElement_ : tailUbFactorElement_;
201+ tailUbFactorElementAligned_ = CeilAlign(tailUbFactorElement_, dataPerBlock);
202+ 
203+ uint32_t sortedValueBytes = ubFactorElementAligned_ * inputDataTypeByte + K_VALUE_MAX * inputDataTypeByte;
204+ uint32_t sortedIndicesBytes = ubFactorElementAligned_ * BYTES_B32 + K_VALUE_MAX * BYTES_B32;
205+ uint32_t pBytes = dataPerBlock * inputDataTypeByte;
206+ uint32_t kBytes = DATA_PER_BLOCK_B32 * BYTES_B32;
207+ uint32_t outTensorBytes = ubFactorElementAligned_ * inputDataTypeByte;
208+ 
209+ calUbSize_ = calUbSize_ - sortedValueBytes - sortedIndicesBytes - pBytes - kBytes - outTensorBytes;
210+ if (onlyTopP_ > 0) {
211+ calUbSize_ = static_cast<uint32_t>(platformUbSize_);
212+ }
213+}
214+ 
215+void ApplyTopKTopPCustomTiling::FillTilingData() {
216+ tilingData.set_batchSize(batchSize_);
217+ tilingData.set_vocabSize(vocabSize_);
218+ tilingData.set_batchPerCore(batchPerCore_);
219+ tilingData.set_tailBatch(tailBatch_);
220+ tilingData.set_blockNum(usedCoreNum_);
221+ tilingData.set_dataNumInit(dataNumInit_);
222+ tilingData.set_dataNumInitAligned(dataNumInitAligned_);
223+ tilingData.set_ubFactorElement(ubFactorElement_);
224+ tilingData.set_ubFactorElementAligned(ubFactorElementAligned_);
225+ tilingData.set_tailUbFactorElement(tailUbFactorElement_);
226+ tilingData.set_tailUbFactorElementAligned(tailUbFactorElementAligned_);
227+ tilingData.set_calUbSize(calUbSize_);
228+ tilingData.set_iterateTimes(iterateTimes_);
229+}
230+ 
231+void ApplyTopKTopPCustomTiling::PrintTilingData() {
232+ OP_LOGD(opName_, "batchSize: %u.", tilingData.get_batchSize());
233+ OP_LOGD(opName_, "vocabSize: %u.", tilingData.get_vocabSize());
234+ OP_LOGD(opName_, "batchPerCore: %u.", tilingData.get_batchPerCore());
235+ OP_LOGD(opName_, "tailBatch: %u.", tilingData.get_tailBatch());
236+ OP_LOGD(opName_, "usedCoreNum: %u.", tilingData.get_blockNum());
237+ OP_LOGD(opName_, "dataNumInit_: %u.", tilingData.get_dataNumInit());
238+ OP_LOGD(opName_, "dataNumInitAligned_: %u.", tilingData.get_dataNumInitAligned());
239+ OP_LOGD(opName_, "ubFactorElement: %u.", tilingData.get_ubFactorElement());
240+ OP_LOGD(opName_, "ubFactorElementAligned: %u.", tilingData.get_ubFactorElementAligned());
241+ OP_LOGD(opName_, "tailUbFactorElement: %u.", tilingData.get_tailUbFactorElement());
242+ OP_LOGD(opName_, "tailUbFactorElementAligned: %u.", tilingData.get_tailUbFactorElementAligned());
243+ OP_LOGD(opName_, "calUbSize: %u.", tilingData.get_calUbSize());
244+ OP_LOGD(opName_, "iterateTimes: %u.", tilingData.get_iterateTimes());
245+}
246+ 
247+ge::graphStatus ApplyTopKTopPCustomTiling::RunKernelTiling() {
248+ OP_LOGD(opName_, "TilingForApplyTopKTopPCustom start.");
249+ 
250+ SetTilingKey();
251+ GetUsedCore();
252+ CalDataPerCore();
253+ FillTilingData();
254+ PrintTilingData();
255+ 
256+ OP_LOGD(opName_, "tilingKey: %u.", tilingKey_);
257+ uint32_t syncWorkspaceSize = SYS_WORKSPACESIZE;
258+ size_t* currentWorkspace = tilingcontext->GetWorkspaceSizes(1);
259+ currentWorkspace[0] = onlyTopP_ > 0 ? syncWorkspaceSize + batchSize_ * vocabSize_ * FLOAT_BYTES : syncWorkspaceSize;
260+ 
261+ tilingData.SaveToBuffer(tilingcontext->GetRawTilingData()->GetData(),
262+ tilingcontext->GetRawTilingData()->GetCapacity());
263+ tilingcontext->GetRawTilingData()->SetDataSize(tilingData.GetDataSize());
264+ tilingcontext->SetBlockDim(usedCoreNum_);
265+ 
266+ OP_LOGD(opName_, "TilingForApplyTopKTopPCustom end.");
267+ return ge::GRAPH_SUCCESS;
268+}
269+ 
270+static ge::graphStatus TilingForApplyTopKTopPCustom(gert::TilingContext* context) {
271+ ApplyTopKTopPCustomTiling tilingObject(context);
272+ auto ret = tilingObject.Init();
273+ if (ret != ge::GRAPH_SUCCESS) {
274+ OP_LOGE(context->GetNodeName(), "tiling Init failed.");
275+ return ge::GRAPH_FAILED;
276+ }
277+ ret = tilingObject.RunKernelTiling();
278+ OP_LOGD(context->GetNodeName(), "TilingForApplyTopKTopPCustom end.");
279+ return ret;
280+}
281+ 
282+static ge::graphStatus TilingPrepareForApplyTopKTopPCustom(gert::TilingParseContext* context) {
283+ OP_LOGD(context->GetNodeName(), "TilingPrepareForApplyTopKTopPCustom start");
284+ auto compileInfo = context->GetCompiledInfo<TilingForApplyTopKTopPCustomCompileInfo>();
285+ OP_CHECK_NULL_WITH_CONTEXT(context, compileInfo);
286+ auto platformInfo = context->GetPlatformInfo();
287+ OP_CHECK_NULL_WITH_CONTEXT(context, platformInfo);
288+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
289+ compileInfo->totalCoreNum = ascendcPlatform.GetCoreNumAiv();
290+ uint64_t ubSizePlatForm;
291+ ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatForm);
292+ compileInfo->ubSizePlatForm = static_cast<int64_t>(ubSizePlatForm);
293+ OP_CHECK_IF(compileInfo->ubSizePlatForm <= 0, OP_LOGE(context->GetNodeName(), "Failed to get ub size"),
294+ return ge::GRAPH_FAILED);
295+ OP_LOGD(context->GetNodeName(), "ub_size_platform is %lu", compileInfo->ubSizePlatForm);
296+ uint64_t totalUbSize = 0;
297+ platformInfo->GetLocalMemSize(fe::LocalMemType::UB, totalUbSize);
298+ OP_LOGD(context->GetNodeName(), "total ub size is %lu", totalUbSize);
299+ OP_LOGD(context->GetNodeName(), "TilingPrepareForApplyTopKTopPCustom end");
300+ return ge::GRAPH_SUCCESS;
301+}
302+ 
303+IMPL_OP_OPTILING(ApplyTopKTopPCustom)
304+ .Tiling(TilingForApplyTopKTopPCustom)
305+ .TilingParse<TilingForApplyTopKTopPCustomCompileInfo>(TilingPrepareForApplyTopKTopPCustom);
306+} // namespace optiling
Asrc/kernels/mie_ops/csrc/sampler/apply_top_k_top_p_custom/op_host/apply_top_k_top_p_custom_tiling.h+51-0
@@ -0,0 +1,51 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file apply_top_k_top_p_custom_tiling.h
13+ * \brief
14+ * ATTENTION: MAKE SURE 'BEGIN_TILING_DATA_DEF' STAY IN THE SAME LINE (28) USING BLANK LINES.
15+ *
16+ *
17+ *
18+ *
19+ *
20+ */
21+#ifndef __APPLY_TOP_K_TOP_P_CUSTOM_TILINGDATA_H__
22+#define __APPLY_TOP_K_TOP_P_CUSTOM_TILINGDATA_H__
23+ 
24+#include "register/tilingdata_base.h"
25+ 
26+namespace optiling {
27+ 
28+BEGIN_TILING_DATA_DEF(ApplyTopKTopPCustomTilingData)
29+TILING_DATA_FIELD_DEF(uint32_t, batchSize);
30+TILING_DATA_FIELD_DEF(uint32_t, vocabSize);
31+TILING_DATA_FIELD_DEF(uint32_t, batchPerCore);
32+TILING_DATA_FIELD_DEF(uint32_t, tailBatch);
33+TILING_DATA_FIELD_DEF(uint32_t, blockNum);
34+TILING_DATA_FIELD_DEF(uint32_t, dataNumInit);
35+TILING_DATA_FIELD_DEF(uint32_t, dataNumInitAligned);
36+TILING_DATA_FIELD_DEF(uint32_t, ubFactorElement);
37+TILING_DATA_FIELD_DEF(uint32_t, ubFactorElementAligned);
38+TILING_DATA_FIELD_DEF(uint32_t, tailUbFactorElement);
39+TILING_DATA_FIELD_DEF(uint32_t, tailUbFactorElementAligned);
40+TILING_DATA_FIELD_DEF(uint32_t, calUbSize);
41+TILING_DATA_FIELD_DEF(uint32_t, iterateTimes);
42+END_TILING_DATA_DEF;
43+REGISTER_TILING_DATA_CLASS(ApplyTopKTopPCustom, ApplyTopKTopPCustomTilingData)
44+ 
45+struct TilingForApplyTopKTopPCustomCompileInfo {
46+ uint32_t totalCoreNum = 0;
47+ uint64_t ubSizePlatForm = 0;
48+};
49+ 
50+} // namespace optiling
51+#endif // __APPLY_TOP_K_TOP_P_CUSTOM_TILINGDATA_H__
Asrc/kernels/mie_ops/csrc/sampler/apply_top_k_top_p_custom/op_host/error_log.h+67-0
@@ -0,0 +1,67 @@
1+#ifndef OPS_BUILT_IN_OP_TILING_ERROR_LOG_H_
2+#define OPS_BUILT_IN_OP_TILING_ERROR_LOG_H_
3+ 
4+#include <string>
5+ 
6+#include "toolchain/slog.h"
7+ 
8+#define OP_LOGI(opname, ...)
9+#define OP_LOGW(opname, ...) \
10+ do { \
11+ printf("[WARN][%s] ", (opname), ##__VA_ARGS__); \
12+ printf("\n"); \
13+ } while (0)
14+ 
15+#define OP_LOGE_WITHOUT_REPORT(opname, ...) \
16+ do { \
17+ printf("[ERRORx][%s] ", (opname), ##__VA_ARGS__); \
18+ printf("\n"); \
19+ } while (0)
20+ 
21+#define OP_LOGE(opname, ...) \
22+ do { \
23+ printf("[ERROR][%s] ", (opname), ##__VA_ARGS__); \
24+ printf("\n"); \
25+ } while (0)
26+ 
27+#define OP_LOGD(opname, ...)
28+ 
29+namespace optiling {
30+ 
31+#define VECTOR_INNER_ERR_REPORT_TILIING(op_name, err_msg, ...) \
32+ do { \
33+ OP_LOGE_WITHOUT_REPORT(op_name, err_msg, ##__VA_ARGS__); \
34+ } while (0)
35+ 
36+#define OP_CHECK_IF(cond, log_func, expr) \
37+ do { \
38+ if (cond) { \
39+ log_func; \
40+ expr; \
41+ } \
42+ } while (0)
43+ 
44+#define OP_CHECK_NULL_WITH_CONTEXT(context, ptr) \
45+ do { \
46+ if ((ptr) == nullptr) { \
47+ OP_LOGE(context->GetNodeType(), "%s is null", #ptr); \
48+ return ge::GRAPH_FAILED; \
49+ } \
50+ } while (0)
51+ 
52+} // namespace optiling
53+ 
54+template <typename T>
55+T CeilAlign(T a, T b) {
56+ return (a + b - 1) / b * b;
57+}
58+ 
59+template <typename T>
60+T CeilDiv(T a, T b) {
61+ if (b == 0) {
62+ return a;
63+ }
64+ return (a + b - 1) / b;
65+}
66+ 
67+#endif // OPS_BUILT_IN_OP_TILING_ERROR_LOG_H_
Asrc/kernels/mie_ops/csrc/sampler/apply_top_k_top_p_custom/op_host/op_api/aclnn_apply_top_k_top_p_custom.cpp+198-0
@@ -0,0 +1,198 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file aclnn_apply_top_k_top_p_custom.cpp
13+ * \brief
14+ */
15+#include "aclnn_apply_top_k_top_p_custom.h"
16+ 
17+#include "../apply_top_k_top_p_custom.h"
18+#include "../sort.h"
19+#include "aclnn/aclnn_base.h"
20+#include "aclnn_kernels/common/op_error_check.h"
21+#include "aclnn_kernels/contiguous.h"
22+#include "opdev/common_types.h"
23+#include "opdev/data_type_utils.h"
24+#include "opdev/format_utils.h"
25+#include "opdev/make_op_executor.h"
26+#include "opdev/op_dfx.h"
27+#include "opdev/op_executor.h"
28+#include "opdev/op_log.h"
29+#include "opdev/platform.h"
30+#include "opdev/shape_utils.h"
31+#include "opdev/tensor_view_utils.h"
32+ 
33+using namespace op;
34+#ifdef __cplusplus
35+extern "C" {
36+#endif
37+ 
38+extern aclnnStatus aclnnInnerApplyTopKTopPCustomGetWorkspaceSize(const aclTensor* sortedValue,
39+ const aclTensor* sortedIndices, const aclTensor* p,
40+ const aclTensor* k, aclTensor* out,
41+ uint64_t* workspaceSize, aclOpExecutor** executor);
42+ 
43+extern aclnnStatus aclnnInnerApplyTopKTopPCustom(void* workspace, uint64_t workspaceSize, aclOpExecutor* executor,
44+ aclrtStream stream);
45+ 
46+namespace {
47+static const int64_t EXPECTED_DIM_ONE = 1;
48+static const int64_t EXPECTED_DIM_TWO = 2;
49+static constexpr size_t DIM_ONE = 1;
50+ 
51+static const std::initializer_list<op::DataType> DTYPE_SUPPORT_LIST = {op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16,
52+ op::DataType::DT_BF16};
53+ 
54+static const std::initializer_list<op::DataType> INT_DTYPE_SUPPORT_LIST = {op::DataType::DT_INT32};
55+ 
56+static bool CheckNotNull(const aclTensor* logits, const aclTensor* p, const aclTensor* k, const aclTensor* out) {
57+ OP_CHECK_NULL(logits, return false);
58+ if (p == nullptr && k == nullptr) {
59+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "The inputs, p and k, should not be nullptr at the same time.");
60+ }
61+ OP_CHECK_NULL(out, return false);
62+ return true;
63+}
64+ 
65+static bool CheckDtypeValid(const aclTensor* logits, const aclTensor* p, const aclTensor* k, const aclTensor* out) {
66+ OP_CHECK_DTYPE_NOT_SUPPORT(logits, DTYPE_SUPPORT_LIST, return false);
67+ if (p != nullptr) {
68+ OP_CHECK_DTYPE_NOT_SUPPORT(p, DTYPE_SUPPORT_LIST, return false);
69+ }
70+ if (k != nullptr) {
71+ OP_CHECK_DTYPE_NOT_SUPPORT(k, INT_DTYPE_SUPPORT_LIST, return false);
72+ }
73+ OP_CHECK_DTYPE_NOT_SUPPORT(out, DTYPE_SUPPORT_LIST, return false);
74+ 
75+ if (p != nullptr) {
76+ OP_CHECK_DTYPE_NOT_MATCH(p, logits->GetDataType(), return false);
77+ }
78+ OP_CHECK_DTYPE_NOT_MATCH(out, logits->GetDataType(), return false);
79+ return true;
80+}
81+ 
82+static bool CheckShapeValid(const aclTensor* logits, const aclTensor* p, const aclTensor* k, const aclTensor* out) {
83+ OP_CHECK_WRONG_DIMENSION(logits, EXPECTED_DIM_TWO, return false);
84+ OP_CHECK_SHAPE_NOT_EQUAL(out, logits, return false);
85+ if (p != nullptr) {
86+ OP_CHECK_WRONG_DIMENSION(p, EXPECTED_DIM_ONE, return false);
87+ }
88+ if (k != nullptr) {
89+ OP_CHECK_WRONG_DIMENSION(k, EXPECTED_DIM_ONE, return false);
90+ }
91+ if (p != nullptr && p->GetViewShape().GetDim(0) != logits->GetViewShape().GetDim(0)) {
92+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "expected p.size(0) is equal to logits.size(0), but got %ld.",
93+ p->GetViewShape().GetDim(0));
94+ return false;
95+ }
96+ if (k != nullptr && k->GetViewShape().GetDim(0) != logits->GetViewShape().GetDim(0)) {
97+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "expected k.size(0) is equal to logits.size(0), but got %ld.",
98+ k->GetViewShape().GetDim(0));
99+ return false;
100+ }
101+ return true;
102+}
103+ 
104+static bool CheckFormatValid(const aclTensor* logits, const aclTensor* p, const aclTensor* k, const aclTensor* out) {
105+ if (logits->GetStorageFormat() != Format::FORMAT_ND) {
106+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "logits format only support ND");
107+ return false;
108+ }
109+ if (p != nullptr && p->GetStorageFormat() != Format::FORMAT_ND) {
110+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "p format only support ND");
111+ return false;
112+ }
113+ if (k != nullptr && k->GetStorageFormat() != Format::FORMAT_ND) {
114+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "k format only support ND");
115+ return false;
116+ }
117+ if (out->GetStorageFormat() != Format::FORMAT_ND) {
118+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "out format only support ND");
119+ return false;
120+ }
121+ return true;
122+}
123+ 
124+static aclnnStatus CheckParams(const aclTensor* logits, const aclTensor* p, const aclTensor* k, const aclTensor* out) {
125+ CHECK_RET(CheckNotNull(logits, p, k, out), ACLNN_ERR_PARAM_NULLPTR);
126+ CHECK_RET(CheckDtypeValid(logits, p, k, out), ACLNN_ERR_PARAM_INVALID);
127+ CHECK_RET(CheckShapeValid(logits, p, k, out), ACLNN_ERR_PARAM_INVALID);
128+ CHECK_RET(CheckFormatValid(logits, p, k, out), ACLNN_ERR_PARAM_INVALID);
129+ return ACLNN_SUCCESS;
130+}
131+} // namespace
132+ 
133+__attribute__((visibility("default"))) aclnnStatus
134+aclnnApplyTopKTopPCustomGetWorkspaceSize(const aclTensor* logits, const aclTensor* p, const aclTensor* k,
135+ aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor) {
136+ OP_CHECK_COMM_INPUT(workspaceSize, executor);
137+ L2_DFX_PHASE_1(aclnnApplyTopKTopPCustom, DFX_IN(logits, p, k), DFX_OUT(out));
138+ auto uniqueExecutor = CREATE_EXECUTOR();
139+ CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR);
140+ 
141+ auto ret = CheckParams(logits, p, k, out);
142+ CHECK_RET(ret == ACLNN_SUCCESS, ret);
143+ bool pIsEmpty = false;
144+ bool kIsEmpty = false;
145+ if (p != nullptr) {
146+ pIsEmpty = p->IsEmpty();
147+ }
148+ if (k != nullptr) {
149+ kIsEmpty = k->IsEmpty();
150+ }
151+ if (logits->IsEmpty() || pIsEmpty || kIsEmpty) {
152+ *workspaceSize = 0;
153+ uniqueExecutor.ReleaseTo(executor);
154+ return ACLNN_SUCCESS;
155+ }
156+ auto logitsContiguous = l0op::Contiguous(logits, uniqueExecutor.get());
157+ CHECK_RET(logitsContiguous != nullptr, ACLNN_ERR_INNER_NULLPTR);
158+ const aclTensor* pContiguous = nullptr;
159+ const aclTensor* kContiguous = nullptr;
160+ if (p != nullptr) {
161+ pContiguous = l0op::Contiguous(p, uniqueExecutor.get());
162+ CHECK_RET(pContiguous != nullptr, ACLNN_ERR_INNER_NULLPTR);
163+ }
164+ if (k != nullptr) {
165+ kContiguous = l0op::Contiguous(k, uniqueExecutor.get());
166+ CHECK_RET(kContiguous != nullptr, ACLNN_ERR_INNER_NULLPTR);
167+ }
168+ bool isLastDimSizeOne = logits->GetViewShape()[DIM_ONE] == 1;
169+ auto viewCopyResult = logitsContiguous;
170+ if (isLastDimSizeOne) {
171+ viewCopyResult = l0op::ViewCopy(logitsContiguous, out, uniqueExecutor.get());
172+ } else {
173+ auto sortResult = l0op::Sort(logitsContiguous, -1, false, true, op::DataType::DT_INT32, uniqueExecutor.get());
174+ const aclTensor* sortedValue = std::get<0>(sortResult);
175+ CHECK_RET(sortedValue != nullptr, ACLNN_ERR_INNER_NULLPTR);
176+ const aclTensor* sortedIndices = std::get<1>(sortResult);
177+ CHECK_RET(sortedIndices != nullptr, ACLNN_ERR_INNER_NULLPTR);
178+ auto res =
179+ l0op::ApplyTopKTopPCustom(sortedValue, sortedIndices, pContiguous, kContiguous, uniqueExecutor.get());
180+ CHECK_RET(res != nullptr, ACLNN_ERR_INNER_NULLPTR);
181+ viewCopyResult = l0op::ViewCopy(res, out, uniqueExecutor.get());
182+ }
183+ CHECK_RET(viewCopyResult != nullptr, ACLNN_ERR_INNER_NULLPTR);
184+ *workspaceSize = uniqueExecutor->GetWorkspaceSize();
185+ uniqueExecutor.ReleaseTo(executor);
186+ return ACLNN_SUCCESS;
187+}
188+ 
189+__attribute__((visibility("default"))) aclnnStatus aclnnApplyTopKTopPCustom(void* workspace, uint64_t workspaceSize,
190+ aclOpExecutor* executor,
191+ aclrtStream stream) {
192+ L2_DFX_PHASE_2(aclnnApplyTopKTopPCustom);
193+ return CommonOpExecutorRun(workspace, workspaceSize, executor, stream);
194+}
195+ 
196+#ifdef __cplusplus
197+}
198+#endif
Asrc/kernels/mie_ops/csrc/sampler/apply_top_k_top_p_custom/op_host/op_api/aclnn_apply_top_k_top_p_custom.h+36-0
@@ -0,0 +1,36 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file aclnn_apply_top_k_top_p_custom.h
13+ * \brief
14+ */
15+#ifndef OP_API_INC_APPLY_TOP_K_TOP_P_CUSTOM_H_
16+#define OP_API_INC_APPLY_TOP_K_TOP_P_CUSTOM_H_
17+ 
18+#include "aclnn/aclnn_base.h"
19+ 
20+#ifdef __cplusplus
21+extern "C" {
22+#endif
23+ 
24+__attribute__((visibility("default"))) aclnnStatus
25+aclnnApplyTopKTopPCustomGetWorkspaceSize(const aclTensor* logits, const aclTensor* p, const aclTensor* k,
26+ aclTensor* out, uint64_t* workspaceSize, aclOpExecutor** executor);
27+ 
28+__attribute__((visibility("default"))) aclnnStatus aclnnApplyTopKTopPCustom(void* workspace, uint64_t workspaceSize,
29+ aclOpExecutor* executor,
30+ aclrtStream stream);
31+ 
32+#ifdef __cplusplus
33+}
34+#endif
35+ 
36+#endif // OP_API_INC_APPLY_TOP_K_TOP_P_CUSTOM_H_
Asrc/kernels/mie_ops/csrc/sampler/apply_top_k_top_p_custom/op_host/sort.h+26-0
@@ -0,0 +1,26 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file sort.h
13+ * \brief
14+ */
15+#ifndef PTA_NPU_OP_API_INC_LEVEL0_OP_SORT_OP_H_
16+#define PTA_NPU_OP_API_INC_LEVEL0_OP_SORT_OP_H_
17+ 
18+#include "opdev/fast_vector.h"
19+#include "opdev/op_executor.h"
20+ 
21+namespace l0op {
22+const std::tuple<aclTensor*, aclTensor*> Sort(const aclTensor* self, int64_t dim, bool descending, bool stable,
23+ op::DataType indicesType, aclOpExecutor* executor);
24+}
25+ 
26+#endif // PTA_NPU_OP_API_INC_LEVEL0_OP_SORT_OP_H_
Asrc/kernels/mie_ops/csrc/sampler/apply_top_k_top_p_custom/op_kernel/apply_top_k_top_p_custom.cpp+44-0
@@ -0,0 +1,44 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file apply_top_k_top_p_custom.cpp
13+ * \brief
14+ */
15+ 
16+#include "apply_top_k_top_p_custom.h"
17+ 
18+#include "apply_top_p_custom.h"
19+using namespace AscendC;
20+using namespace ApplyTopKTopPCustomOp;
21+using namespace ApplyTopPCustomOp;
22+ 
23+extern "C" __global__ __aicore__ void apply_top_k_top_p_custom(GM_ADDR sorted_value, GM_ADDR sorted_indices, GM_ADDR p,
24+ GM_ADDR k, GM_ADDR out, GM_ADDR workSpace,
25+ GM_ADDR tiling) {
26+ TPipe pipe;
27+ GET_TILING_DATA(tilingData, tiling);
28+ if (TILING_KEY_IS(0)) {
29+ ApplyTopKTopPCustomOp::ApplyTopKTopPCustom<DTYPE_OUT, float, DTYPE_OUT> op;
30+ op.InitTilingData(tilingData, sorted_value, sorted_indices, p, k, out);
31+ op.InitBuffer(&pipe);
32+ op.Process();
33+ } else if (TILING_KEY_IS(1)) {
34+ ApplyTopKTopPCustomOp::ApplyTopKTopPCustom<DTYPE_OUT, float, DTYPE_OUT> op;
35+ op.InitTilingData(tilingData, sorted_value, sorted_indices, p, k, out);
36+ op.InitBuffer(&pipe);
37+ op.ProcessTopK();
38+ } else if (TILING_KEY_IS(2)) {
39+ ApplyTopPCustomOp::ApplyTopPCustom<DTYPE_OUT, float, DTYPE_OUT> op;
40+ op.InitTilingData(tilingData, sorted_value, sorted_indices, p, k, out, workSpace);
41+ op.InitBuffer(&pipe);
42+ op.ProcessTopP();
43+ }
44+}
Asrc/kernels/mie_ops/csrc/sampler/apply_top_k_top_p_custom/op_kernel/apply_top_p_custom.h+466-0
@@ -0,0 +1,466 @@
1+/**
2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file apply_top_p_custom.h
13+ * \brief
14+ */
15+#ifndef APPLY_TOP_P_CUSTOM_H_KERNEL
16+#define APPLY_TOP_P_CUSTOM_H_KERNEL
17+ 
18+#include "kernel_operator.h"
19+ 
20+using namespace AscendC;
21+namespace ApplyTopPCustomOp {
22+constexpr uint16_t FLOAT16_NEG_INF = 0xFC00; // -inf 64512
23+constexpr uint16_t BF16_NEG_INF = 0xFF80; // -inf 65408
24+constexpr int32_t FLOAT32_NEG_INF = 0xFF800000; // -inf -2139095040
25+ 
26+constexpr uint32_t BLOCK_BYTES = 32;
27+constexpr uint32_t DATA_PER_BLOCK_B32 = 8;
28+constexpr uint32_t DATA_PER_REPEAT_B32 = 64;
29+constexpr uint32_t SCATTER_PART_LENGTH = 1024;
30+constexpr uint32_t RESERVED_UB = 1024;
31+constexpr uint32_t FLOAT_BYTES = 4;
32+constexpr uint32_t SOFTMAX_UB_NUM = 2;
33+ 
34+template <typename inputT, typename calT, typename outputT>
35+class ApplyTopPCustom {
36+ public:
37+ __aicore__ inline ApplyTopPCustom(){};
38+ __aicore__ inline void InitTilingData(const ApplyTopKTopPCustomTilingData &__restrict tilingData,
39+ GM_ADDR sorted_value, GM_ADDR sorted_indices, GM_ADDR p, GM_ADDR k,
40+ GM_ADDR out, GM_ADDR workspace);
41+ __aicore__ inline void InitBuffer(TPipe *inputPipe);
42+ __aicore__ inline void ProcessTopP();
43+ 
44+ private:
45+ __aicore__ inline void ReduceSumWithAddsAndExpImpl(uint32_t loopDataNum);
46+ // topp func
47+ __aicore__ inline void ProcessPreSingleBatch(uint32_t loopBatch);
48+ __aicore__ inline void GetSoftmaxSum(uint32_t loopBatch);
49+ __aicore__ inline void CumsumKoggleStone(uint32_t loopBatch);
50+ __aicore__ inline void GetPValue(uint32_t batchOffset);
51+ __aicore__ inline void CumsumParamCompute(uint32_t iterateTime);
52+ __aicore__ inline void GetMaxValue(int64_t baseGmIdx);
53+ __aicore__ inline void GetSoftMaxRes(uint32_t loopBatch);
54+ __aicore__ inline void ScatterSingleTask(uint32_t taskIndex);
55+ 
56+ __aicore__ inline void SToMTE3Sync() {
57+ event_t eventIDSToMTE3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_MTE3));
58+ SetFlag<HardEvent::S_MTE3>(eventIDSToMTE3);
59+ WaitFlag<HardEvent::S_MTE3>(eventIDSToMTE3);
60+ }
61+ __aicore__ inline void VToMTE3Sync() {
62+ event_t eventIDVToMTE3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
63+ SetFlag<HardEvent::V_MTE3>(eventIDVToMTE3);
64+ WaitFlag<HardEvent::V_MTE3>(eventIDVToMTE3);
65+ }
66+ __aicore__ inline void VToMTE2Sync() {
67+ event_t eventIDVToMTE2 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE2));
68+ SetFlag<HardEvent::V_MTE2>(eventIDVToMTE2);
69+ WaitFlag<HardEvent::V_MTE2>(eventIDVToMTE2);
70+ }
71+ __aicore__ inline void MTE3ToSSync() {
72+ event_t eventIDMTE3ToS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_S));
73+ SetFlag<HardEvent::MTE3_S>(eventIDMTE3ToS);
74+ WaitFlag<HardEvent::MTE3_S>(eventIDMTE3ToS);
75+ }
76+ __aicore__ inline void VToSSync() {
77+ event_t eventIdVToS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
78+ SetFlag<HardEvent::V_S>(eventIdVToS);
79+ WaitFlag<HardEvent::V_S>(eventIdVToS);
80+ }
81+ __aicore__ inline void SToVSync() {
82+ event_t eventIdSToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
83+ SetFlag<HardEvent::S_V>(eventIdSToV);
84+ WaitFlag<HardEvent::S_V>(eventIdSToV);
85+ }
86+ __aicore__ inline void MTE2ToVSync() {
87+ event_t eventIdMte2ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_V));
88+ SetFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
89+ WaitFlag<HardEvent::MTE2_V>(eventIdMte2ToV);
90+ }
91+ __aicore__ inline void MTE2ToSSync() {
92+ event_t eventIdMte2ToS = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S));
93+ SetFlag<HardEvent::MTE2_S>(eventIdMte2ToS);
94+ WaitFlag<HardEvent::MTE2_S>(eventIdMte2ToS);
95+ }
96+ 
97+ __aicore__ inline void MTE3ToMTE2Sync() {
98+ event_t eventIdMte3ToMte2 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_MTE2));
99+ SetFlag<HardEvent::MTE3_MTE2>(eventIdMte3ToMte2);
100+ WaitFlag<HardEvent::MTE3_MTE2>(eventIdMte3ToMte2);
101+ }
102+ 
103+ __aicore__ inline void MTE3ToVSync() {
104+ event_t eventIdMte3ToV = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V));
105+ SetFlag<HardEvent::MTE3_V>(eventIdMte3ToV);
106+ WaitFlag<HardEvent::MTE3_V>(eventIdMte3ToV);
107+ }
108+ 
109+ __aicore__ inline uint32_t CeilDiv(uint32_t x, uint32_t y) { return y == 0 ? x : (x + y - 1) / y; }
110+ 
111+ private:
112+ TPipe *pipe_;
113+ TBuf<TPosition::VECCALC> calBuf_;
114+ 
115+ // tilingData
116+ uint32_t batchSize_ = 0;
117+ uint32_t vocabSize_ = 0;
118+ uint32_t batchPerCore_ = 0;
119+ uint32_t tailBatch_ = 0;
120+ uint32_t blockNum_ = 0;
121+ uint32_t dataNumInitAligned_ = 0;
122+ uint32_t calUbSize_ = 0;
123+ uint32_t blockIdx_ = 0;
124+ uint32_t loopBatch_ = 0;
125+ uint32_t batchOffset_ = 0;
126+ uint32_t bufOffsetLoop = 0;
127+ int64_t baseGmIdx_ = 0;
128+ 
129+ // topp scalar
130+ uint32_t maxSoftmaxLength = 1;
131+ uint32_t softmaxLength = 1;
132+ uint32_t lineSfLoopTimes = 1;
133+ uint32_t softmaxLengthTail = 1;
134+ uint32_t scatterLength = 1;
135+ 
136+ uint32_t singleCoreB = 0;
137+ uint32_t singleCoreBTail = 0;
138+ uint32_t vCnt = 0;
139+ uint32_t bCnt = 0;
140+ uint32_t singleCoreV = 1;
141+ uint32_t singleCoreVTail = 1;
142+ uint32_t iterateTimes = 1;
143+ 
144+ GlobalTensor<inputT> mGmSortedValue_;
145+ GlobalTensor<int32_t> mGmSortedIndices_;
146+ GlobalTensor<inputT> mGmP_;
147+ GlobalTensor<int32_t> mGmK_;
148+ GlobalTensor<outputT> mGmOut_;
149+ GlobalTensor<float> softMaxGm;
150+ 
151+ LocalTensor<uint8_t> totalUb;
152+ 
153+ // softmax tensor
154+ LocalTensor<float> softMaxLocalFp32;
155+ LocalTensor<inputT> softMaxLocal;
156+ LocalTensor<float> softMaxResLocal;
157+ LocalTensor<float> reduceLocal;
158+ LocalTensor<inputT> outInfLocal;
159+ 
160+ // cumsum tensor
161+ LocalTensor<float> cumSumInput1Local;
162+ LocalTensor<float> cumSumInput2Local;
163+ 
164+ // scatter tensor
165+ LocalTensor<inputT> sortedValueLocal;
166+ LocalTensor<int32_t> sortedIndicesLocal;
167+ LocalTensor<float> sortedValueLocalFp32;
168+ LocalTensor<outputT> scatterLocal;
169+ LocalTensor<float> cumsumLocal;
170+ 
171+ float pValue = 0;
172+ float maxValue = 0;
173+ float reduceSumValueInvert = 0;
174+ float reduceSumValue = 0;
175+ BinaryRepeatParams repeatParams = {1, 0, 1, 8, 0, 8};
176+ DataCopyExtParams scatterCopyParams{1, (uint32_t)(sizeof(outputT)), 0, 0, 0};
177+};
178+ 
179+template <typename inputT, typename calT, typename outputT>
180+__aicore__ inline void ApplyTopPCustom<inputT, calT, outputT>::InitTilingData(
181+ const ApplyTopKTopPCustomTilingData &__restrict tilingData, GM_ADDR sorted_value, GM_ADDR sorted_indices, GM_ADDR p,
182+ GM_ADDR k, GM_ADDR out, GM_ADDR workspace) {
183+ batchSize_ = tilingData.batchSize;
184+ vocabSize_ = tilingData.vocabSize;
185+ batchPerCore_ = tilingData.batchPerCore;
186+ tailBatch_ = tilingData.tailBatch;
187+ blockNum_ = tilingData.blockNum;
188+ calUbSize_ = tilingData.calUbSize;
189+ iterateTimes = tilingData.iterateTimes;
190+ blockIdx_ = GetBlockIdx();
191+ if (blockIdx_ < tailBatch_) {
192+ loopBatch_ = batchPerCore_ + 1;
193+ batchOffset_ = blockIdx_ * loopBatch_;
194+ } else {
195+ loopBatch_ = batchPerCore_;
196+ batchOffset_ = blockIdx_ * batchPerCore_ + tailBatch_;
197+ }
198+ maxSoftmaxLength = (calUbSize_ - RESERVED_UB) / SOFTMAX_UB_NUM / FLOAT_BYTES;
199+ softmaxLength = maxSoftmaxLength < vocabSize_ ? maxSoftmaxLength : vocabSize_;
200+ 
201+ lineSfLoopTimes = (vocabSize_ + softmaxLength - 1) / softmaxLength;
202+ softmaxLengthTail = vocabSize_ - (lineSfLoopTimes - 1) * softmaxLength;
203+ scatterLength = (calUbSize_ - RESERVED_UB - BLOCK_BYTES) / (SOFTMAX_UB_NUM * FLOAT_BYTES + sizeof(inputT)) /
204+ SCATTER_PART_LENGTH * SCATTER_PART_LENGTH;
205+ singleCoreB = CeilDiv(batchSize_, blockNum_);
206+ vCnt = batchSize_ < blockNum_ ? blockNum_ / batchSize_ : 1;
207+ bCnt = batchSize_;
208+ singleCoreB = 1;
209+ singleCoreBTail = 1;
210+ singleCoreV = vocabSize_ / vCnt;
211+ singleCoreVTail = vocabSize_ - vCnt * singleCoreV;
212+ mGmSortedValue_.SetGlobalBuffer(reinterpret_cast<__gm__ inputT *>(sorted_value));
213+ mGmSortedIndices_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(sorted_indices));
214+ mGmP_.SetGlobalBuffer(reinterpret_cast<__gm__ inputT *>(p));
215+ mGmK_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(k));
216+ mGmOut_.SetGlobalBuffer(reinterpret_cast<__gm__ outputT *>(out));
217+ softMaxGm.SetGlobalBuffer((__gm__ float *)workspace, batchSize_ * vocabSize_);
218+}
219+ 
220+// init used buffer
221+template <typename inputT, typename calT, typename outputT>
222+__aicore__ inline void ApplyTopPCustom<inputT, calT, outputT>::InitBuffer(TPipe *inputPipe) {
223+ pipe_ = inputPipe;
224+ pipe_->InitBuffer(calBuf_, calUbSize_);
225+ totalUb = calBuf_.Get<uint8_t>();
226+ // softmax ub
227+ uint32_t softmaxLengthAligned = CeilDiv(softmaxLength, BLOCK_BYTES / sizeof(inputT)) * BLOCK_BYTES / sizeof(inputT);
228+ softMaxLocalFp32 = totalUb.ReinterpretCast<float>();
229+ softMaxLocal = totalUb[softmaxLengthAligned * sizeof(inputT)].ReinterpretCast<inputT>();
230+ softMaxResLocal = totalUb[softmaxLengthAligned * sizeof(float)].ReinterpretCast<float>();
231+ reduceLocal = totalUb[softmaxLengthAligned * sizeof(float) * 2].ReinterpretCast<float>(); // 32 bytes
232+ outInfLocal = totalUb.ReinterpretCast<inputT>(); // Take softmax ub
233+ 
234+ // cumsum ub
235+ cumSumInput1Local = totalUb.ReinterpretCast<float>(); // Take softmax local
236+ cumSumInput2Local = totalUb[softmaxLengthAligned * sizeof(float)].ReinterpretCast<float>(); // Take softmax res ub
237+ 
238+ // scatter ub
239+ sortedValueLocal = totalUb[0].ReinterpretCast<inputT>();
240+ sortedIndicesLocal = totalUb[scatterLength * sizeof(inputT)].ReinterpretCast<int32_t>();
241+ cumsumLocal = totalUb[scatterLength * (FLOAT_BYTES + sizeof(inputT))].ReinterpretCast<float>();
242+ scatterLocal = totalUb[calUbSize_ - RESERVED_UB + BLOCK_BYTES].ReinterpretCast<outputT>(); // 32 bytes
243+}
244+ 
245+template <typename inputT, typename calT, typename outputT>
246+__aicore__ inline void ApplyTopPCustom<inputT, calT, outputT>::GetMaxValue(int64_t baseGmIdx) {
247+ int64_t initGmIdx = baseGmIdx + vocabSize_ - 1;
248+ if constexpr (IsSameType<inputT, float>::value) {
249+ maxValue = -mGmSortedValue_[initGmIdx].GetValue(0);
250+ } else if constexpr (IsSameType<inputT, half>::value) {
251+ maxValue = -static_cast<float>(mGmSortedValue_[initGmIdx].GetValue(0));
252+ } else {
253+ maxValue = -ToFloat(mGmSortedValue_[initGmIdx].GetValue(0));
254+ }
255+}
256+ 
257+template <typename inputT, typename calT, typename outputT>
258+__aicore__ inline void ApplyTopPCustom<inputT, calT, outputT>::GetPValue(uint32_t batchOffset) {
259+ if constexpr (IsSameType<inputT, float>::value) {
260+ pValue = float(1.0) - mGmP_[batchOffset].GetValue(0);
261+ } else if constexpr (IsSameType<inputT, half>::value) {
262+ pValue = float(1.0) - static_cast<float>(mGmP_[batchOffset].GetValue(0));
263+ } else {
264+ pValue = float(1.0) - ToFloat(mGmP_[batchOffset].GetValue(0));
265+ }
266+}
267+ 
268+template <typename inputT, typename calT, typename outputT>
269+__aicore__ inline void ApplyTopPCustom<inputT, calT, outputT>::ProcessPreSingleBatch(uint32_t loopBatch) {
270+ reduceSumValue = 0;
271+ GetSoftmaxSum(loopBatch);
272+ GetSoftMaxRes(loopBatch);
273+ CumsumKoggleStone(loopBatch);
274+}
275+ 
276+template <typename inputT, typename calT, typename outputT>
277+__aicore__ inline void ApplyTopPCustom<inputT, calT, outputT>::ProcessTopP() {
278+ for (uint32_t loopBatch = 0; loopBatch < loopBatch_; loopBatch++) {
279+ baseGmIdx_ = batchOffset_ * vocabSize_ + loopBatch * vocabSize_;
280+ GetMaxValue(baseGmIdx_); // Get max value in softmax.
281+ ProcessPreSingleBatch(loopBatch); // Softmax and cumsum.
282+ }
283+ SyncAll();
284+ for (uint32_t taskIndex = 0; taskIndex < bCnt * vCnt; taskIndex++) {
285+ ScatterSingleTask(taskIndex);
286+ }
287+}
288+ 
289+template <typename inputT, typename calT, typename outputT>
290+__aicore__ inline void ApplyTopPCustom<inputT, calT, outputT>::ScatterSingleTask(uint32_t taskIndex) {
291+ if (GetBlockIdx() == taskIndex % blockNum_) {
292+ uint32_t bCntIndex = taskIndex / vCnt;
293+ uint32_t vCntIndex = taskIndex % vCnt;
294+ uint32_t vCurSingleCore = vCntIndex < singleCoreVTail ? (singleCoreV + 1) : singleCoreV;
295+ uint32_t copyTimes = CeilDiv(vCurSingleCore, scatterLength);
296+ uint32_t copyLength = scatterLength;
297+ uint32_t copyLengthTail = vCurSingleCore - (copyTimes - 1) * scatterLength;
298+ GetPValue(bCntIndex); // Get maxPValue.
299+ for (uint32_t cpIndex = 0; cpIndex < copyTimes; cpIndex++) {
300+ uint32_t curCopyLength = cpIndex == (copyTimes - 1) ? copyLengthTail : copyLength;
301+ int64_t gmOffset =
302+ vCntIndex < singleCoreVTail
303+ ? bCntIndex * vocabSize_ + vCntIndex * (singleCoreV + 1) + cpIndex * copyLength
304+ : bCntIndex * vocabSize_ + vCntIndex * singleCoreV + singleCoreVTail + cpIndex * copyLength;
305+ DataCopyPad(cumsumLocal, softMaxGm[gmOffset],
306+ {1, static_cast<uint32_t>(curCopyLength * sizeof(float)), 0, 0, 0}, {false, 0, 0, 0});
307+ DataCopyPad(sortedIndicesLocal, mGmSortedIndices_[gmOffset],
308+ {1, static_cast<uint32_t>(curCopyLength * sizeof(float)), 0, 0, 0}, {false, 0, 0, 0});
309+ DataCopyPad(sortedValueLocal, mGmSortedValue_[gmOffset],
310+ {1, static_cast<uint32_t>(curCopyLength * sizeof(inputT)), 0, 0, 0}, {false, 0, 0, 0});
311+ MTE2ToSSync();
312+ 
313+ if (cumsumLocal.GetValue(curCopyLength - 1) <= pValue) {
314+ continue;
315+ }
316+ uint32_t scatterLoop = CeilDiv(curCopyLength, SCATTER_PART_LENGTH);
317+ uint32_t scatterNumsTail = curCopyLength - (scatterLoop - 1) * SCATTER_PART_LENGTH;
318+ for (uint32_t scatterLoopIndex = 0; scatterLoopIndex < scatterLoop; scatterLoopIndex++) {
319+ uint32_t curScatterNums = scatterLoopIndex == (scatterLoop - 1) ? scatterNumsTail : SCATTER_PART_LENGTH;
320+ if (cumsumLocal.GetValue(scatterLoopIndex * SCATTER_PART_LENGTH + curScatterNums - 1) <= pValue) {
321+ continue;
322+ }
323+ for (uint32_t scatterIndex = 0; scatterIndex < curScatterNums; scatterIndex++) {
324+ int64_t scatterOffset = scatterLoopIndex * SCATTER_PART_LENGTH + scatterIndex;
325+ if (cumsumLocal.GetValue(scatterOffset) <= pValue) {
326+ continue;
327+ }
328+ scatterLocal.SetValue(0, sortedValueLocal.GetValue(scatterOffset));
329+ int32_t lineIndex = sortedIndicesLocal.GetValue(scatterOffset);
330+ SToMTE3Sync();
331+ DataCopyPad(mGmOut_[bCntIndex * vocabSize_ + lineIndex],
332+ scatterLocal.template ReinterpretCast<outputT>(),
333+ {1, (uint32_t)(1 * sizeof(outputT)), 0, 0, 0});
334+ MTE3ToSSync();
335+ }
336+ }
337+ }
338+ }
339+}
340+template <typename inputT, typename calT, typename outputT>
341+__aicore__ inline void ApplyTopPCustom<inputT, calT, outputT>::GetSoftMaxRes(uint32_t loopBatch) {
342+ uint32_t loopDataNum = softmaxLength;
343+ for (int32_t loopInner = 0; loopInner < lineSfLoopTimes; loopInner++) {
344+ int64_t currentGmIdx = baseGmIdx_ + loopInner * softmaxLength;
345+ if (loopInner == (lineSfLoopTimes - 1)) {
346+ loopDataNum = softmaxLengthTail;
347+ }
348+ if constexpr (!IsSameType<inputT, float>::value) {
349+ DataCopyPad(softMaxLocal, mGmSortedValue_[currentGmIdx],
350+ {1, static_cast<uint32_t>(loopDataNum * sizeof(inputT)), 0, 0, 0}, {false, 0, 0, 0});
351+ MTE2ToVSync();
352+ Cast(softMaxLocalFp32, softMaxLocal, RoundMode::CAST_NONE, loopDataNum);
353+ PipeBarrier<PIPE_V>();
354+ } else {
355+ DataCopyPad(softMaxLocalFp32, mGmSortedValue_[currentGmIdx],
356+ {1, static_cast<uint32_t>(loopDataNum * sizeof(float)), 0, 0, 0}, {false, 0, 0, 0});
357+ MTE2ToVSync();
358+ }
359+ Adds(softMaxResLocal, softMaxLocalFp32, maxValue, loopDataNum);
360+ VToMTE2Sync();
361+ PipeBarrier<PIPE_V>();
362+ Exp(softMaxResLocal, softMaxResLocal, loopDataNum);
363+ PipeBarrier<PIPE_V>();
364+ Muls(softMaxResLocal, softMaxResLocal, reduceSumValueInvert, loopDataNum);
365+ VToMTE3Sync();
366+ DataCopyPad(softMaxGm[currentGmIdx], softMaxResLocal,
367+ {1, static_cast<uint32_t>(loopDataNum * sizeof(float)), 0, 0, 0});
368+ MTE3ToMTE2Sync();
369+ }
370+}
371+ 
372+template <typename inputT, typename calT, typename outputT>
373+__aicore__ inline void ApplyTopPCustom<inputT, calT, outputT>::CumsumKoggleStone(uint32_t loopBatch) {
374+ uint32_t loopDataNum = softmaxLength;
375+ for (uint32_t iterateTime = 0; iterateTime < iterateTimes; iterateTime++) {
376+ int64_t iteratOffset = 1;
377+ for (uint32_t powerIdx = 0; powerIdx < iterateTime; powerIdx++) {
378+ iteratOffset = iteratOffset * 2;
379+ }
380+ uint32_t addLength = vocabSize_ - iteratOffset;
381+ uint32_t innerLoopNum = addLength / softmaxLength;
382+ uint32_t dataTail = addLength - innerLoopNum * softmaxLength;
383+ loopDataNum = softmaxLength;
384+ for (uint32_t innerLoopIdx = 0; innerLoopIdx < innerLoopNum; innerLoopIdx++) {
385+ // Copy data from right
386+ int64_t loopInnerOffset = dataTail + (innerLoopNum - 1 - innerLoopIdx) * softmaxLength;
387+ DataCopyPad(cumSumInput1Local, softMaxGm[baseGmIdx_ + loopInnerOffset],
388+ {1, static_cast<uint32_t>(loopDataNum * sizeof(float)), 0, 0, 0}, {false, 0, 0, 0});
389+ DataCopyPad(cumSumInput2Local, softMaxGm[baseGmIdx_ + loopInnerOffset + iteratOffset],
390+ {1, static_cast<uint32_t>(loopDataNum * sizeof(float)), 0, 0, 0}, {false, 0, 0, 0});
391+ MTE2ToVSync();
392+ Add(cumSumInput1Local, cumSumInput1Local, cumSumInput2Local, loopDataNum);
393+ VToMTE3Sync();
394+ DataCopyPad(softMaxGm[baseGmIdx_ + loopInnerOffset + iteratOffset], cumSumInput1Local,
395+ {1, static_cast<uint32_t>(loopDataNum * sizeof(float)), 0, 0, 0});
396+ MTE3ToMTE2Sync();
397+ }
398+ if (dataTail > 0) {
399+ loopDataNum = dataTail;
400+ DataCopyPad(cumSumInput1Local, softMaxGm[baseGmIdx_],
401+ {1, static_cast<uint32_t>(loopDataNum * sizeof(float)), 0, 0, 0}, {false, 0, 0, 0});
402+ DataCopyPad(cumSumInput2Local, softMaxGm[baseGmIdx_ + iteratOffset],
403+ {1, static_cast<uint32_t>(loopDataNum * sizeof(float)), 0, 0, 0}, {false, 0, 0, 0});
404+ MTE2ToVSync();
405+ Add(cumSumInput1Local, cumSumInput1Local, cumSumInput2Local, loopDataNum);
406+ VToMTE3Sync();
407+ DataCopyPad(softMaxGm[baseGmIdx_ + iteratOffset], cumSumInput1Local,
408+ {1, static_cast<uint32_t>(loopDataNum * sizeof(float)), 0, 0, 0});
409+ MTE3ToMTE2Sync();
410+ }
411+ }
412+ MTE3ToVSync();
413+}
414+ 
415+template <typename inputT, typename calT, typename outputT>
416+__aicore__ inline void ApplyTopPCustom<inputT, calT, outputT>::ReduceSumWithAddsAndExpImpl(uint32_t loopDataNum) {
417+ Adds(softMaxResLocal, softMaxLocalFp32, maxValue, loopDataNum);
418+ PipeBarrier<PIPE_V>();
419+ Exp(softMaxResLocal, softMaxResLocal, loopDataNum);
420+ PipeBarrier<PIPE_V>();
421+ ReduceSum(reduceLocal, softMaxResLocal, reduceLocal, loopDataNum);
422+}
423+ 
424+template <typename inputT, typename calT, typename outputT>
425+__aicore__ inline void ApplyTopPCustom<inputT, calT, outputT>::GetSoftmaxSum(uint32_t loopBatch) {
426+ uint32_t loopDataNum = softmaxLength;
427+ for (int32_t loopInner = 0; loopInner < lineSfLoopTimes; loopInner++) {
428+ int64_t currentGmIdx = baseGmIdx_ + loopInner * softmaxLength;
429+ if (loopInner == (lineSfLoopTimes - 1)) {
430+ loopDataNum = softmaxLengthTail;
431+ }
432+ if constexpr (IsSameType<inputT, float>::value) {
433+ Duplicate(outInfLocal.template ReinterpretCast<int32_t>(), FLOAT32_NEG_INF, loopDataNum);
434+ } else if constexpr (IsSameType<inputT, half>::value) {
435+ Duplicate(outInfLocal.template ReinterpretCast<uint16_t>(), FLOAT16_NEG_INF, loopDataNum);
436+ } else {
437+ Duplicate(outInfLocal.template ReinterpretCast<uint16_t>(), BF16_NEG_INF, loopDataNum);
438+ }
439+ VToMTE3Sync();
440+ DataCopyPad(mGmOut_[currentGmIdx], outInfLocal,
441+ {1, static_cast<uint32_t>(loopDataNum * sizeof(inputT)), 0, 0, 0});
442+ MTE3ToMTE2Sync();
443+ if constexpr (!IsSameType<inputT, float>::value) {
444+ DataCopyPad(softMaxLocal, mGmSortedValue_[currentGmIdx],
445+ {1, static_cast<uint32_t>(loopDataNum * sizeof(inputT)), 0, 0, 0}, {false, 0, 0, 0});
446+ MTE2ToVSync();
447+ Cast(softMaxLocalFp32, softMaxLocal, RoundMode::CAST_NONE, loopDataNum);
448+ PipeBarrier<PIPE_V>();
449+ } else {
450+ DataCopyPad(softMaxLocalFp32, mGmSortedValue_[currentGmIdx],
451+ {1, static_cast<uint32_t>(loopDataNum * sizeof(inputT)), 0, 0, 0}, {false, 0, 0, 0});
452+ MTE2ToVSync();
453+ }
454+ 
455+ ReduceSumWithAddsAndExpImpl(loopDataNum);
456+ VToSSync();
457+ // Sum up to obtain the sum of exp reduce for the first x loops in the row.
458+ reduceSumValue += reduceLocal.GetValue(0);
459+ SToVSync();
460+ }
461+ reduceSumValueInvert = 1 / reduceSumValue;
462+ SToVSync();
463+}
464+} // namespace ApplyTopPCustomOp
465+ 
466+#endif // APPLY_TOP_P_CUSTOM_H_KERNEL
Asrc/kernels/mie_ops/torch_ops_extension/npu_apply_top_k_top_p_custom.cpp+31-0
@@ -0,0 +1,31 @@
1+#include <torch/library.h>
2+ 
3+#include <iostream>
4+ 
5+#include "ops_common.h"
6+ 
7+namespace mie_ops {
8+using namespace at_npu::native;
9+ 
10+at::Tensor npu_apply_top_k_top_p_custom_npu(const at::Tensor& logits, const c10::optional<at::Tensor>& p,
11+ const c10::optional<at::Tensor>& k) {
12+ at::Tensor out = at::empty_like(logits);
13+ EXEC_NPU_CMD_V1(aclnnApplyTopKTopPCustom, logits, p, k, out);
14+ 
15+ return out;
16+}
17+ 
18+at::Tensor npu_apply_top_k_top_p_custom_meta(const at::Tensor& logits, const c10::optional<at::Tensor>& p,
19+ const c10::optional<at::Tensor>& k) {
20+ return at::empty_like(logits);
21+}
22+ 
23+} // namespace mie_ops
24+ 
25+TORCH_LIBRARY_IMPL(mie_ops, PrivateUse1, m) {
26+ m.impl("apply_top_k_top_p_custom", &mie_ops::npu_apply_top_k_top_p_custom_npu);
27+}
28+ 
29+TORCH_LIBRARY_IMPL(mie_ops, Meta, m) {
30+ m.impl("apply_top_k_top_p_custom", &mie_ops::npu_apply_top_k_top_p_custom_meta);
31+}
Msrc/kernels/mie_ops/torch_ops_extension/ops_def_registration.cpp+34-15
@@ -14,21 +14,40 @@
14// 在custom命名空间里注册add_custom和npu_selected_flash_attention和后续的XXX算子，每次新增自定义aten ir都需先增加定义14// 在custom命名空间里注册add_custom和npu_selected_flash_attention和后续的XXX算子，每次新增自定义aten ir都需先增加定义
15// step1, 为新增自定义算子添加定义15// step1, 为新增自定义算子添加定义
16TORCH_LIBRARY(mie_ops, m) {16TORCH_LIBRARY(mie_ops, m) {
17- m.def("npu_mla_process(Tensor input, Tensor gamma0, Tensor beta0, Tensor wdqkv, Tensor descale0, Tensor gamma1, Tensor beta1, Tensor wuq, Tensor descale1, Tensor gamma2, Tensor cos, Tensor sin, Tensor wuk, Tensor kv_cache, Tensor kv_cache_rope, Tensor slotmapping, *, Tensor? quant_scale0=None, Tensor? quant_offset0=None, Tensor? bias0=None, Tensor? quant_scale1=None, Tensor? quant_offset1=None, Tensor? bias1=None, Tensor? ctkv_scale=None, Tensor? q_nope_scale=None, str? cache_mode_opt=None, str? quant_mode_opt=None) -> (Tensor, Tensor, Tensor, Tensor)");17+ m.def(
18- m.def("npu_dispatch_gmm_combine_decode(Tensor x, Tensor expert_ids, Tensor[] gmm1_permuted_weight, Tensor[] gmm1_permuted_weight_scale, Tensor[] gmm2_weight, Tensor[] gmm2_weight_scale, Tensor expert_scales, Tensor? expert_smooth_scales=None, Tensor? x_active_mask=None, str group_ep='', int ep_rank_size=0, int ep_rank_id=0, int moe_expert_num=0, int shared_expert_num=1, int shared_expert_rank_num=0, int quant_mode=0, int global_bs=0) -> (Tensor output, Tensor expert_token_nums)");18+ "npu_mla_process(Tensor input, Tensor gamma0, Tensor beta0, Tensor wdqkv, Tensor descale0, Tensor gamma1, "
19- m.def("npu_lightning_indexer(Tensor query, Tensor key, Tensor weights, *, Tensor? actual_seq_lengths_query=None, Tensor? actual_seq_lengths_key=None, Tensor? block_table=None, str layout_query='BSND', str layout_key='PA_BSND', int selected_count=2048, int sparse_mode=3) -> Tensor");19+ "Tensor beta1, Tensor wuq, Tensor descale1, Tensor gamma2, Tensor cos, Tensor sin, Tensor wuk, Tensor "
20+ "kv_cache, Tensor kv_cache_rope, Tensor slotmapping, *, Tensor? quant_scale0=None, Tensor? quant_offset0=None, "
21+ "Tensor? bias0=None, Tensor? quant_scale1=None, Tensor? quant_offset1=None, Tensor? bias1=None, Tensor? "
22+ "ctkv_scale=None, Tensor? q_nope_scale=None, str? cache_mode_opt=None, str? quant_mode_opt=None) -> (Tensor, "
23+ "Tensor, Tensor, Tensor)");
24+ m.def(
25+ "npu_dispatch_gmm_combine_decode(Tensor x, Tensor expert_ids, Tensor[] gmm1_permuted_weight, Tensor[] "
26+ "gmm1_permuted_weight_scale, Tensor[] gmm2_weight, Tensor[] gmm2_weight_scale, Tensor expert_scales, Tensor? "
27+ "expert_smooth_scales=None, Tensor? x_active_mask=None, str group_ep='', int ep_rank_size=0, int ep_rank_id=0, "
28+ "int moe_expert_num=0, int shared_expert_num=1, int shared_expert_rank_num=0, int quant_mode=0, int "
29+ "global_bs=0) -> (Tensor output, Tensor expert_token_nums)");
30+ m.def(
31+ "npu_lightning_indexer(Tensor query, Tensor key, Tensor weights, *, Tensor? actual_seq_lengths_query=None, "
32+ "Tensor? actual_seq_lengths_key=None, Tensor? block_table=None, str layout_query='BSND', str "
33+ "layout_key='PA_BSND', int selected_count=2048, int sparse_mode=3) -> Tensor");
20 m.def(34 m.def(
21 "npu_dispatch_ffn_combine("35 "npu_dispatch_ffn_combine("
22- "Tensor x, "36+ "Tensor x, "
23- "Tensor[] weight1, "37+ "Tensor[] weight1, "
24- "Tensor[] weight2, "38+ "Tensor[] weight2, "
25- "Tensor expert_idx, "39+ "Tensor expert_idx, "
26- "Tensor[] scale1, "40+ "Tensor[] scale1, "
27- "Tensor[] scale2, "41+ "Tensor[] scale2, "
28- "Tensor probs, "42+ "Tensor probs, "
29- "str group, "43+ "str group, "
30- "int max_output_size, "44+ "int max_output_size, "
31- "Tensor! out"45+ "Tensor! out"
32- ") -> Tensor"46+ ") -> Tensor");
33- );47+ m.def(
48+ "apply_top_k_top_p_custom("
49+ "Tensor logits, "
50+ "Tensor? p, "
51+ "Tensor? k"
52+ ") -> Tensor");
34}53}